added --config flag

Author: scarere

vec_inf/cli/_cli.py

@@ -124,6 +124,11 @@ def cli() -> None:
     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]],
@@ -166,6 +171,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
 

vec_inf/client/_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
         )

vec_inf/client/models.py

@@ -218,6 +218,8 @@ class LaunchOptions:
         Additional arguments for vLLM
     env : str, optional
         Environment variables to be set
+    config : str, optional
+        Path to custom model config yaml
     """
 
     model_family: Optional[str] = None
@@ -238,6 +240,7 @@ class LaunchOptions:
     model_weights_parent_dir: Optional[str] = None
     vllm_args: Optional[str] = None
     env: Optional[str] = None
+    config: Optional[str] = None
 
 
 @dataclass