Move REQUIRED_ARGS from _client_vars to environment.yaml

Author: XkunW

vec_inf/client/_client_vars.py

@@ -71,10 +71,3 @@
 
 # Required matching arguments for batch mode
 BATCH_MODE_REQUIRED_MATCHING_ARGS = ["venv", "log_dir"]
-
-# Required arguments for launching jobs that don't have a default value and their
-# corresponding environment variables
-REQUIRED_ARGS = {
-    "account": "VEC_INF_ACCOUNT",
-    "work_dir": "VEC_INF_WORK_DIR",
-}

vec_inf/client/_slurm_vars.py

@@ -78,5 +78,9 @@ def create_literal_type(values: list[str], fallback: str = "") -> Any:
     _config["allowed_values"]["resource_type"]
 )
 
+# Extract required arguments, for launching jobs that don't have a default value and their
+# corresponding environment variables
+REQUIRED_ARGS: dict[str, str] = _config["required_args"]
+
 # Extract default arguments
 DEFAULT_ARGS: dict[str, str] = _config["default_args"]

vec_inf/client/_utils.py

@@ -14,9 +14,9 @@
 import requests
 import yaml
 
-from vec_inf.client._client_vars import MODEL_READY_SIGNATURE, REQUIRED_ARGS
+from vec_inf.client._client_vars import MODEL_READY_SIGNATURE
 from vec_inf.client._exceptions import MissingRequiredFieldsError
-from vec_inf.client._slurm_vars import CACHED_CONFIG_DIR
+from vec_inf.client._slurm_vars import CACHED_CONFIG_DIR, REQUIRED_ARGS
 from vec_inf.client.config import ModelConfig
 from vec_inf.client.models import ModelStatus
 

vec_inf/config/environment.yaml

@@ -15,6 +15,10 @@ allowed_values:
   partition: []
   resource_type: ["l40s", "h100"]
 
+required_args:
+  account: "VEC_INF_ACCOUNT"
+  work_dir: "VEC_INF_WORK_DIR"
+
 default_args:
   cpus_per_task: "16"
   mem_per_node: "64G"