Remove SINGULARITY_IMAGE env var as it breaks reproducibility

Author: XkunW

vec_inf/client/_client_vars.py

@@ -144,10 +144,10 @@ class SlurmScriptTemplate(TypedDict):
     },
     "singularity_setup": [
         SINGULARITY_LOAD_CMD,
-        "singularity exec $SINGULARITY_IMAGE ray stop",
+        "singularity exec {singularity_image} ray stop",
     ],
     "imports": "source {src_dir}/find_port.sh",
-    "singularity_command": "singularity exec --nv --bind {model_weights_path}:{model_weights_path} --containall $SINGULARITY_IMAGE",
+    "singularity_command": "singularity exec --nv --bind {model_weights_path}:{model_weights_path} --containall {singularity_image}",
     "activate_venv": "source {venv}/bin/activate",
     "server_setup": {
         "single_node": [

vec_inf/client/_helper.py

@@ -37,7 +37,6 @@
 )
 from vec_inf.client.slurm_vars import (
     LD_LIBRARY_PATH,
-    SINGULARITY_IMAGE,
     VLLM_NCCL_SO_PATH,
 )
 
@@ -229,7 +228,6 @@ def _set_env_vars(self) -> None:
         """Set environment variables for the launch command."""
         os.environ["LD_LIBRARY_PATH"] = LD_LIBRARY_PATH
         os.environ["VLLM_NCCL_SO_PATH"] = VLLM_NCCL_SO_PATH
-        os.environ["SINGULARITY_IMAGE"] = SINGULARITY_IMAGE
 
     def _build_launch_command(self) -> str:
         """Generate the slurm script and construct the launch command.

vec_inf/client/_slurm_script_generator.py

@@ -12,6 +12,7 @@
     SLURM_JOB_CONFIG_ARGS,
     SLURM_SCRIPT_TEMPLATE,
 )
+from vec_inf.client.slurm_vars import SINGULARITY_IMAGE
 
 
 class SlurmScriptGenerator:
@@ -85,7 +86,11 @@ def _generate_server_setup(self) -> str:
         """
         server_script = ["\n"]
         if self.use_singularity:
-            server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["singularity_setup"]))
+            server_script.append(
+                "\n".join(SLURM_SCRIPT_TEMPLATE["singularity_setup"]).format(
+                    singularity_image=SINGULARITY_IMAGE,
+                )
+            )
         server_script.append(
             SLURM_SCRIPT_TEMPLATE["imports"].format(src_dir=self.params["src_dir"])
         )
@@ -97,7 +102,8 @@ def _generate_server_setup(self) -> str:
                 server_setup_str = server_setup_str.replace(
                     "SINGULARITY_PLACEHOLDER",
                     SLURM_SCRIPT_TEMPLATE["singularity_command"].format(
-                        model_weights_path=self.model_weights_path
+                        model_weights_path=self.model_weights_path,
+                        singularity_image=SINGULARITY_IMAGE,
                     ),
                 )
         else:
@@ -128,7 +134,8 @@ def _generate_launch_cmd(self) -> str:
         if self.use_singularity:
             launcher_script.append(
                 SLURM_SCRIPT_TEMPLATE["singularity_command"].format(
-                    model_weights_path=self.model_weights_path
+                    model_weights_path=self.model_weights_path,
+                    singularity_image=SINGULARITY_IMAGE,
                 )
                 + " \\"
             )