Merge pull request #135 from VectorInstitute/add_env_arg

Added additional launch command options:
* --env option for adding environment variables
* --config option for path to custom config

Author: XkunW

tests/vec_inf/cli/test_cli.py

@@ -40,6 +40,7 @@ def test_launch_command_success(runner):
             "model_weights_parent_dir": "/model-weights",
             "vocab_size": "128000",
             "vllm_args": {"max_model_len": 8192},
+            "env": {"CACHE": "/cache"},
         }
         mock_client.launch_model.return_value = mock_response
 

tests/vec_inf/cli/test_helper.py

@@ -36,6 +36,7 @@ def test_format_table_output(self):
             "model_weights_parent_dir": "/model-weights",
             "log_dir": "/tmp/logs",
             "vllm_args": {"max_model_len": 8192, "enable_prefix_caching": True},
+            "env": {"CACHE": "/cache"},
         }
 
         formatter = LaunchResponseFormatter(model_name, params)
@@ -63,6 +64,7 @@ def test_format_table_output_with_minimal_params(self):
             "model_weights_parent_dir": "/weights",
             "log_dir": "/logs",
             "vllm_args": {},
+            "env": {},
         }
 
         formatter = LaunchResponseFormatter(model_name, params)

tests/vec_inf/client/test_helper.py

@@ -124,6 +124,22 @@ def test_process_vllm_args(self, mock_path_exists, mock_load_config, mock_config
         assert vllm_args["--quantization"] == "awq"
         assert vllm_args["--compilation-config"] == "3"
 
+    @patch("vec_inf.client._helper.utils.load_config")
+    @patch("pathlib.Path.exists")
+    def test_process_env_vars(self, mock_path_exists, mock_load_config, mock_configs):
+        """Test that vars from `--env` flag are parsed correctly."""
+        mock_load_config.return_value = mock_configs
+        mock_path_exists.return_value = True
+
+        # Get filepath of dummy env file
+        file_path = Path(__file__).parent / "test_vars.env"
+
+        launcher = ModelLauncher("test-model", {})
+        env_vars = launcher._process_env_vars(f"CACHE_DIR=/cache,{file_path}")
+        assert env_vars["CACHE_DIR"] == "/cache"
+        assert env_vars["MY_VAR"] == "5"
+        assert env_vars["VLLM_CACHE_ROOT"] == "/cache/vllm"
+
     @patch("vec_inf.client._helper.utils.load_config")
     def test_get_launch_params_merges_config_and_cli_args(
         self, mock_load_config, model_config

tests/vec_inf/client/test_slurm_script_generator.py

@@ -55,6 +55,11 @@ def singularity_params(self, basic_params):
             {
                 "venv": "singularity",
                 "bind": "/scratch:/scratch,/data:/data",
+                "env": {
+                    "CACHE_DIR": "/cache",
+                    "MY_VAR": "5",
+                    "VLLM_CACHE_ROOT": "/cache/vllm",
+                },
             }
         )
         return singularity
@@ -74,6 +79,7 @@ def test_init_single_node(self, basic_params):
         assert not generator.use_container
         assert generator.additional_binds == ""
         assert generator.model_weights_path == "/path/to/model_weights/test-model"
+        assert generator.env_str == ""
 
     def test_init_multinode(self, multinode_params):
         """Test initialization with multi-node configuration."""
@@ -84,6 +90,7 @@ def test_init_multinode(self, multinode_params):
         assert not generator.use_container
         assert generator.additional_binds == ""
         assert generator.model_weights_path == "/path/to/model_weights/test-model"
+        assert generator.env_str == ""
 
     def test_init_singularity(self, singularity_params):
         """Test initialization with Singularity configuration."""
@@ -94,6 +101,10 @@ def test_init_singularity(self, singularity_params):
         assert not generator.is_multinode
         assert generator.additional_binds == " --bind /scratch:/scratch,/data:/data"
         assert generator.model_weights_path == "/path/to/model_weights/test-model"
+        assert (
+            generator.env_str
+            == "--env CACHE_DIR=/cache,MY_VAR=5,VLLM_CACHE_ROOT=/cache/vllm"
+        )
 
     def test_init_singularity_no_bind(self, basic_params):
         """Test Singularity initialization without additional binds."""
@@ -106,6 +117,7 @@ def test_init_singularity_no_bind(self, basic_params):
         assert not generator.is_multinode
         assert generator.additional_binds == ""
         assert generator.model_weights_path == "/path/to/model_weights/test-model"
+        assert generator.env_str == ""
 
     def test_generate_shebang_single_node(self, basic_params):
         """Test shebang generation for single-node setup."""

tests/vec_inf/client/test_vars.env

@@ -0,0 +1,2 @@
+MY_VAR=5
+VLLM_CACHE_ROOT=/cache/vllm

uv.lock

@@ -131,6 +131,16 @@ def cli() -> None:
     is_flag=True,
     help="Output in JSON string",
 )
+@click.option(
+    "--env",
+    type=str,
+    help="Environment variables to be set. Seperate variables with commas. Can also include path to a file containing environment variables seperated by newlines. e.g. --env 'TRITON_CACHE_DIR=/scratch/.cache/triton,my_custom_vars_file.env'",
+)
+@click.option(
+    "--config",
+    type=str,
+    help="Path to a model config yaml file to use in place of the default",
+)
 def launch(
     model_name: str,
     **cli_kwargs: Optional[Union[str, int, float, bool]],
@@ -177,6 +187,10 @@ def launch(
             Path to model weights directory
         - vllm_args : str, optional
             vLLM engine arguments
+        - env : str, optional
+            Environment variables
+        - config : str, optional
+            Path to custom model config yaml file
         - json_mode : bool, optional
             Output in JSON format
 
@@ -212,7 +226,10 @@ def launch(
         raise click.ClickException(f"Launch failed: {str(e)}") from e
 
 
-@cli.command("batch-launch", help="Launch multiple models in a batch.")
+@cli.command(
+    "batch-launch",
+    help="Launch multiple models in a batch, separate model names with spaces.",
+)
 @click.argument("model-names", type=str, nargs=-1)
 @click.option(
     "--batch-config",

vec_inf/cli/_cli.py

@@ -87,6 +87,11 @@ def format_table_output(self) -> Table:
         for arg, value in self.params["vllm_args"].items():
             table.add_row(f"  {arg}:", str(value))
 
+        # Add Environment Variable Configuration Details
+        table.add_row("Environment Variables", style="magenta")
+        for arg, value in self.params["env"].items():
+            table.add_row(f"  {arg}:", str(value))
+
         return table
 
 

vec_inf/cli/_helper.py

@@ -61,7 +61,7 @@ def __init__(self, model_name: str, kwargs: Optional[dict[str, Any]]):
         self.kwargs = kwargs or {}
         self.slurm_job_id = ""
         self.slurm_script_path = Path("")
-        self.model_config = self._get_model_configuration()
+        self.model_config = self._get_model_configuration(self.kwargs.get("config"))
         self.params = self._get_launch_params()
 
     def _warn(self, message: str) -> None:
@@ -74,9 +74,14 @@ def _warn(self, message: str) -> None:
         """
         warnings.warn(message, UserWarning, stacklevel=2)
 
-    def _get_model_configuration(self) -> ModelConfig:
+    def _get_model_configuration(self, config_path: str | None = None) -> ModelConfig:
         """Load and validate model configuration.
 
+        Parameters
+        ----------
+        config_path : str | None, optional
+            Path to a yaml file with custom model config to use in place of the default
+
         Returns
         -------
         ModelConfig
@@ -89,7 +94,7 @@ def _get_model_configuration(self) -> ModelConfig:
         ModelConfigurationError
             If model configuration is not found and weights don't exist
         """
-        model_configs = utils.load_config()
+        model_configs = utils.load_config(config_path=config_path)
         config = next(
             (m for m in model_configs if m.model_name == self.model_name), None
         )
@@ -158,6 +163,38 @@ def _process_vllm_args(self, arg_string: str) -> dict[str, Any]:
                 vllm_args[arg.strip()] = True
         return vllm_args
 
+    def _process_env_vars(self, env_arg: str) -> dict[str, str]:
+        """Process the env string into a dictionary of environment variables.
+
+        Parameters
+        ----------
+        env_arg : str
+            String containing comma separated list of environment variable definitions
+            (eg. MY_VAR=1), file paths containing environment variable definitions
+            (separated by newlines), or a combination of both
+            (eg. 'MY_VAR=5,my_env.env')
+
+        Returns
+        -------
+        dict[str, str]
+            Processed environment variables as key-value pairs.
+        """
+        env_vars: dict[str, str] = {}
+        for arg in env_arg.split(","):
+            if "=" in arg:  # Arg is an env var definition
+                key, value = arg.split("=")
+                env_vars[key.strip()] = value.strip()
+            else:  # Arg is a path to a file
+                with open(arg, "r") as file:
+                    lines = [line.rstrip() for line in file]
+                for line in lines:
+                    if "=" in line:
+                        key, value = line.split("=")
+                        env_vars[key.strip()] = value.strip()
+                    else:
+                        print(f"WARNING: Could not parse env var: {line}")
+        return env_vars
+
     def _get_launch_params(self) -> dict[str, Any]:
         """Prepare launch parameters, set log dir, and validate required fields.
 
@@ -181,6 +218,12 @@ def _get_launch_params(self) -> dict[str, Any]:
                 params["vllm_args"][key] = value
             del self.kwargs["vllm_args"]
 
+        if self.kwargs.get("env"):
+            env_vars = self._process_env_vars(self.kwargs["env"])
+            for key, value in env_vars.items():
+                params["env"][key] = str(value)
+            del self.kwargs["env"]
+
         for key, value in self.kwargs.items():
             params[key] = value
 
@@ -233,7 +276,7 @@ def _get_launch_params(self) -> dict[str, Any]:
 
         # Convert path to string for JSON serialization
         for field in params:
-            if field == "vllm_args":
+            if field in ["vllm_args", "env"]:
                 continue
             params[field] = str(params[field])
 

vec_inf/client/_helper.py

@@ -41,6 +41,15 @@ def __init__(self, params: dict[str, Any]):
         self.model_weights_path = str(
             Path(self.params["model_weights_parent_dir"], self.params["model_name"])
         )
+        env_dict: dict[str, str] = self.params.get("env", {})
+        # Create string of environment variables
+        self.env_str = ""
+        for key, val in env_dict.items():
+            if len(self.env_str) == 0:
+                self.env_str = "--env "
+            else:
+                self.env_str += ","
+            self.env_str += key + "=" + val
 
     def _generate_script_content(self) -> str:
         """Generate the complete Slurm script content.
@@ -100,6 +109,7 @@ def _generate_server_setup(self) -> str:
                     SLURM_SCRIPT_TEMPLATE["container_command"].format(
                         model_weights_path=self.model_weights_path,
                         additional_binds=self.additional_binds,
+                        env_str=self.env_str,
                     ),
                 )
         else:
@@ -132,6 +142,7 @@ def _generate_launch_cmd(self) -> str:
                 SLURM_SCRIPT_TEMPLATE["container_command"].format(
                     model_weights_path=self.model_weights_path,
                     additional_binds=self.additional_binds,
+                    env_str=self.env_str,
                 )
             )
         else: