Skip to content

Commit 69ea120

Browse files
XkunWXkunW
committed
Repalce singularity with container
1 parent d950d06 commit 69ea120

File tree

6 files changed

+58
-57
lines changed

6 files changed

+58
-57
lines changed

‎tests/vec_inf/client/test_slurm_script_generator.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ def test_init_single_node(self, basic_params):
7171

7272
assert generator.params == basic_params
7373
assert not generator.is_multinode
74-
assert not generator.use_singularity
74+
assert not generator.use_container
7575
assert generator.additional_binds == ""
7676
assert generator.model_weights_path == "/path/to/model_weights/test-model"
7777

@@ -81,7 +81,7 @@ def test_init_multinode(self, multinode_params):
8181

8282
assert generator.params == multinode_params
8383
assert generator.is_multinode
84-
assert not generator.use_singularity
84+
assert not generator.use_container
8585
assert generator.additional_binds == ""
8686
assert generator.model_weights_path == "/path/to/model_weights/test-model"
8787

@@ -90,7 +90,7 @@ def test_init_singularity(self, singularity_params):
9090
generator = SlurmScriptGenerator(singularity_params)
9191

9292
assert generator.params == singularity_params
93-
assert generator.use_singularity
93+
assert generator.use_container
9494
assert not generator.is_multinode
9595
assert generator.additional_binds == " --bind /scratch:/scratch,/data:/data"
9696
assert generator.model_weights_path == "/path/to/model_weights/test-model"
@@ -102,7 +102,7 @@ def test_init_singularity_no_bind(self, basic_params):
102102
generator = SlurmScriptGenerator(params)
103103

104104
assert generator.params == params
105-
assert generator.use_singularity
105+
assert generator.use_container
106106
assert not generator.is_multinode
107107
assert generator.additional_binds == ""
108108
assert generator.model_weights_path == "/path/to/model_weights/test-model"
@@ -307,7 +307,7 @@ def test_init_basic(self, batch_params):
307307
generator = BatchSlurmScriptGenerator(batch_params)
308308

309309
assert generator.params == batch_params
310-
assert not generator.use_singularity
310+
assert not generator.use_container
311311
assert len(generator.script_paths) == 0
312312
assert "model1" in generator.params["models"]
313313
assert "model2" in generator.params["models"]
@@ -316,7 +316,7 @@ def test_init_singularity(self, batch_singularity_params):
316316
"""Test initialization with Singularity configuration."""
317317
generator = BatchSlurmScriptGenerator(batch_singularity_params)
318318

319-
assert generator.use_singularity
319+
assert generator.use_container
320320
assert (
321321
generator.params["models"]["model1"]["additional_binds"]
322322
== " --bind /scratch:/scratch,/data:/data"
@@ -335,7 +335,7 @@ def test_init_singularity_no_bind(self, batch_params):
335335

336336
generator = BatchSlurmScriptGenerator(params)
337337

338-
assert generator.use_singularity
338+
assert generator.use_container
339339
assert generator.params["models"]["model1"]["additional_binds"] == ""
340340
assert generator.params["models"]["model2"]["additional_binds"] == ""
341341

‎vec_inf/cli/_cli.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ def cli() -> None:
9999
@click.option(
100100
"--bind",
101101
type=str,
102-
help="Additional binds for the singularity container as a comma separated list of bind paths",
102+
help="Additional binds for the container as a comma separated list of bind paths",
103103
)
104104
@click.option(
105105
"--time",
@@ -166,7 +166,7 @@ def launch(
166166
- nodelist : str, optional
167167
Request a specific list of nodes for deployment
168168
- bind : str, optional
169-
Additional binds for the singularity container
169+
Additional binds for the container as a comma separated list of bind paths
170170
- time : str, optional
171171
Time limit for job
172172
- venv : str, optional

‎vec_inf/client/_slurm_script_generator.py‎

Lines changed: 19 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ class SlurmScriptGenerator:
2121
2222
This class handles the generation of Slurm scripts for both single-node and
2323
multi-node configurations, supporting different virtualization environments
24-
(venv or singularity).
24+
(venv or singularity/apptainer).
2525
2626
Parameters
2727
----------
@@ -32,7 +32,9 @@ class SlurmScriptGenerator:
3232
def __init__(self, params: dict[str, Any]):
3333
self.params = params
3434
self.is_multinode = int(self.params["num_nodes"]) > 1
35-
self.use_singularity = self.params["venv"] == "singularity"
35+
self.use_container = (
36+
self.params["venv"] == "singularity" or self.params["venv"] == "apptainer"
37+
)
3638
self.additional_binds = self.params.get("bind", "")
3739
if self.additional_binds:
3840
self.additional_binds = f" --bind {self.additional_binds}"
@@ -82,8 +84,8 @@ def _generate_server_setup(self) -> str:
8284
Server initialization script content.
8385
"""
8486
server_script = ["\n"]
85-
if self.use_singularity:
86-
server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["singularity_setup"]))
87+
if self.use_container:
88+
server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["container_setup"]))
8789
server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["env_vars"]))
8890
server_script.append(
8991
SLURM_SCRIPT_TEMPLATE["imports"].format(src_dir=self.params["src_dir"])
@@ -92,10 +94,10 @@ def _generate_server_setup(self) -> str:
9294
server_setup_str = "\n".join(
9395
SLURM_SCRIPT_TEMPLATE["server_setup"]["multinode"]
9496
).format(gpus_per_node=self.params["gpus_per_node"])
95-
if self.use_singularity:
97+
if self.use_container:
9698
server_setup_str = server_setup_str.replace(
97-
"SINGULARITY_PLACEHOLDER",
98-
SLURM_SCRIPT_TEMPLATE["singularity_command"].format(
99+
"CONTAINER_PLACEHOLDER",
100+
SLURM_SCRIPT_TEMPLATE["container_command"].format(
99101
model_weights_path=self.model_weights_path,
100102
additional_binds=self.additional_binds,
101103
),
@@ -117,17 +119,17 @@ def _generate_launch_cmd(self) -> str:
117119
"""Generate the vLLM server launch command.
118120
119121
Creates the command to launch the vLLM server, handling different virtualization
120-
environments (venv or singularity).
122+
environments (venv or singularity/apptainer).
121123
122124
Returns
123125
-------
124126
str
125127
Server launch command.
126128
"""
127129
launcher_script = ["\n"]
128-
if self.use_singularity:
130+
if self.use_container:
129131
launcher_script.append(
130-
SLURM_SCRIPT_TEMPLATE["singularity_command"].format(
132+
SLURM_SCRIPT_TEMPLATE["container_command"].format(
131133
model_weights_path=self.model_weights_path,
132134
additional_binds=self.additional_binds,
133135
)
@@ -181,7 +183,9 @@ class BatchSlurmScriptGenerator:
181183
def __init__(self, params: dict[str, Any]):
182184
self.params = params
183185
self.script_paths: list[Path] = []
184-
self.use_singularity = self.params["venv"] == "singularity"
186+
self.use_container = (
187+
self.params["venv"] == "singularity" or self.params["venv"] == "apptainer"
188+
)
185189
for model_name in self.params["models"]:
186190
self.params["models"][model_name]["additional_binds"] = ""
187191
if self.params["models"][model_name].get("bind"):
@@ -225,10 +229,8 @@ def _generate_model_launch_script(self, model_name: str) -> Path:
225229
script_content = []
226230
model_params = self.params["models"][model_name]
227231
script_content.append(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["shebang"])
228-
if self.use_singularity:
229-
script_content.append(
230-
BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["singularity_setup"]
231-
)
232+
if self.use_container:
233+
script_content.append(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["container_setup"])
232234
script_content.append("\n".join(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["env_vars"]))
233235
script_content.append(
234236
"\n".join(
@@ -243,9 +245,9 @@ def _generate_model_launch_script(self, model_name: str) -> Path:
243245
model_name=model_name,
244246
)
245247
)
246-
if self.use_singularity:
248+
if self.use_container:
247249
script_content.append(
248-
BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["singularity_command"].format(
250+
BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["container_command"].format(
249251
model_weights_path=model_params["model_weights_path"],
250252
additional_binds=model_params["additional_binds"],
251253
)

‎vec_inf/client/_slurm_templates.py‎

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,13 @@
77
from typing import TypedDict
88

99
from vec_inf.client._slurm_vars import (
10-
SINGULARITY_IMAGE,
11-
SINGULARITY_LOAD_CMD,
12-
SINGULARITY_MODULE_NAME,
10+
CONTAINER_LOAD_CMD,
11+
CONTAINER_MODULE_NAME,
12+
IMAGE_PATH,
1313
)
1414

1515

16-
SINGULARITY_MODULE_NAME_UPPER = SINGULARITY_MODULE_NAME.upper()
16+
CONTAINER_MODULE_NAME_UPPER = CONTAINER_MODULE_NAME.upper()
1717

1818

1919
class ShebangConfig(TypedDict):
@@ -53,12 +53,12 @@ class SlurmScriptTemplate(TypedDict):
5353
----------
5454
shebang : ShebangConfig
5555
Shebang and SLURM directive configuration
56-
singularity_setup : list[str]
57-
Commands for Singularity container setup
56+
container_setup : list[str]
57+
Commands for container setup
5858
imports : str
5959
Import statements and source commands
60-
singularity_command : str
61-
Template for Singularity execution command
60+
container_command : str
61+
Template for container execution command
6262
activate_venv : str
6363
Template for virtual environment activation
6464
server_setup : ServerSetupConfig
@@ -72,10 +72,10 @@ class SlurmScriptTemplate(TypedDict):
7272
"""
7373

7474
shebang: ShebangConfig
75-
singularity_setup: list[str]
75+
container_setup: list[str]
7676
imports: str
7777
env_vars: list[str]
78-
singularity_command: str
78+
container_command: str
7979
activate_venv: str
8080
server_setup: ServerSetupConfig
8181
find_vllm_port: list[str]
@@ -91,15 +91,15 @@ class SlurmScriptTemplate(TypedDict):
9191
"#SBATCH --tasks-per-node=1",
9292
],
9393
},
94-
"singularity_setup": [
95-
SINGULARITY_LOAD_CMD,
96-
f"{SINGULARITY_MODULE_NAME} exec {SINGULARITY_IMAGE} ray stop",
94+
"container_setup": [
95+
CONTAINER_LOAD_CMD,
96+
f"{CONTAINER_MODULE_NAME} exec {IMAGE_PATH} ray stop",
9797
],
9898
"imports": "source {src_dir}/find_port.sh",
9999
"env_vars": [
100-
f"export {SINGULARITY_MODULE_NAME}_BINDPATH=${SINGULARITY_MODULE_NAME}_BINDPATH,$(echo /dev/infiniband* | sed -e 's/ /,/g')"
100+
f"export {CONTAINER_MODULE_NAME}_BINDPATH=${CONTAINER_MODULE_NAME}_BINDPATH,$(echo /dev/infiniband* | sed -e 's/ /,/g')"
101101
],
102-
"singularity_command": f"{SINGULARITY_MODULE_NAME} exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {SINGULARITY_IMAGE} \\",
102+
"container_command": f"{CONTAINER_MODULE_NAME} exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {IMAGE_PATH} \\",
103103
"activate_venv": "source {venv}/bin/activate",
104104
"server_setup": {
105105
"single_node": [
@@ -118,7 +118,7 @@ class SlurmScriptTemplate(TypedDict):
118118
'echo "Ray Head IP: $ray_head"',
119119
'echo "Starting HEAD at $head_node"',
120120
'srun --nodes=1 --ntasks=1 -w "$head_node" \\',
121-
" SINGULARITY_PLACEHOLDER",
121+
" CONTAINER_PLACEHOLDER",
122122
' ray start --head --node-ip-address="$head_node_ip" --port=$head_node_port \\',
123123
' --num-cpus "$SLURM_CPUS_PER_TASK" --num-gpus {gpus_per_node} --block &',
124124
"sleep 10",
@@ -128,7 +128,7 @@ class SlurmScriptTemplate(TypedDict):
128128
" node_i=${{nodes_array[$i]}}",
129129
' echo "Starting WORKER $i at $node_i"',
130130
' srun --nodes=1 --ntasks=1 -w "$node_i" \\',
131-
" SINGULARITY_PLACEHOLDER",
131+
" CONTAINER_PLACEHOLDER",
132132
' ray start --address "$ray_head" \\',
133133
' --num-cpus "$SLURM_CPUS_PER_TASK" --num-gpus {gpus_per_node} --block &',
134134
" sleep 5",
@@ -196,32 +196,32 @@ class BatchModelLaunchScriptTemplate(TypedDict):
196196
----------
197197
shebang : str
198198
Shebang line for the script
199-
singularity_setup : list[str]
200-
Commands for Singularity container setup
199+
container_setup : list[str]
200+
Commands for container setup
201201
env_vars : list[str]
202202
Environment variables to set
203203
server_address_setup : list[str]
204204
Commands to setup the server address
205205
launch_cmd : list[str]
206206
Commands to launch the vLLM server
207-
singularity_command : str
208-
Commands to setup the singularity command
207+
container_command : str
208+
Commands to setup the container command
209209
"""
210210

211211
shebang: str
212-
singularity_setup: str
212+
container_setup: str
213213
env_vars: list[str]
214214
server_address_setup: list[str]
215215
write_to_json: list[str]
216216
launch_cmd: list[str]
217-
singularity_command: str
217+
container_command: str
218218

219219

220220
BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE: BatchModelLaunchScriptTemplate = {
221221
"shebang": "#!/bin/bash\n",
222-
"singularity_setup": f"{SINGULARITY_LOAD_CMD}\n",
222+
"container_setup": f"{CONTAINER_LOAD_CMD}\n",
223223
"env_vars": [
224-
f"export {SINGULARITY_MODULE_NAME}_BINDPATH=${SINGULARITY_MODULE_NAME}_BINDPATH,$(echo /dev/infiniband* | sed -e 's/ /,/g')"
224+
f"export {CONTAINER_MODULE_NAME}_BINDPATH=${CONTAINER_MODULE_NAME}_BINDPATH,$(echo /dev/infiniband* | sed -e 's/ /,/g')"
225225
],
226226
"server_address_setup": [
227227
"source {src_dir}/find_port.sh",
@@ -238,7 +238,7 @@ class BatchModelLaunchScriptTemplate(TypedDict):
238238
' "$json_path" > temp_{model_name}.json \\',
239239
' && mv temp_{model_name}.json "$json_path"\n',
240240
],
241-
"singularity_command": f"{SINGULARITY_MODULE_NAME} exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {SINGULARITY_IMAGE} \\",
241+
"container_command": f"{CONTAINER_MODULE_NAME} exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {IMAGE_PATH} \\",
242242
"launch_cmd": [
243243
"vllm serve {model_weights_path} \\",
244244
" --served-model-name {model_name} \\",

‎vec_inf/client/_slurm_vars.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -52,12 +52,11 @@ def load_yaml_config(path: Path) -> dict[str, Any]:
5252
_config = load_env_config()
5353

5454
# Extract path values
55-
SINGULARITY_IMAGE = _config["paths"]["image_path"]
56-
VLLM_NCCL_SO_PATH = _config["paths"]["vllm_nccl_so_path"]
55+
IMAGE_PATH = _config["paths"]["image_path"]
5756

5857
# Extract containerization info
59-
SINGULARITY_LOAD_CMD = _config["containerization"]["module_load_cmd"]
60-
SINGULARITY_MODULE_NAME = _config["containerization"]["module_name"]
58+
CONTAINER_LOAD_CMD = _config["containerization"]["module_load_cmd"]
59+
CONTAINER_MODULE_NAME = _config["containerization"]["module_name"]
6160

6261
# Extract limits
6362
MAX_GPUS_PER_NODE = _config["limits"]["max_gpus_per_node"]

‎vec_inf/client/models.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,7 @@ class LaunchOptions:
207207
node_list : str, optional
208208
Request a specific list of nodes for deployment
209209
bind : str, optional
210-
Additional binds for the singularity container
210+
Additional binds for the container as a comma separated list of bind paths
211211
vocab_size : int, optional
212212
Size of model vocabulary
213213
data_type : str, optional

0 commit comments

Comments
 (0)