Merge pull request #148 from VectorInstitute/bugfix/required_args_bon_echo

Move REQUIRED_ARGS from _client_vars to environment.yaml

Author: XkunW

README.md

@@ -44,7 +44,7 @@ You should see an output like the following:
 
 <img width="720" alt="launch_image" src="https://github.com/user-attachments/assets/c1e0c60c-cf7a-49ed-a426-fdb38ebf88ee" />
 
-**NOTE**: On Vector Killarney Cluster environment, the following fields are required:
+**NOTE**: You can set the required fields in the environment configuration (`environment.yaml`), it's a mapping between required arguments and their corresponding environment variables. On the Vector **Killarney** Cluster environment, the required fields are:
   * `--account`, `-A`: The Slurm account, this argument can be set to default by setting environment variable `VEC_INF_ACCOUNT`.
   * `--work-dir`, `-D`: A working directory other than your home directory, this argument can be set to default by seeting environment variable `VEC_INF_WORK_DIR`.
 

docs/user_guide.md

@@ -37,7 +37,7 @@ You should see an output like the following:
 └─────────────────────────┴───────────────────────────────────────────┘
 ```
 
-**NOTE**: On Vector Killarney Cluster environment, the following fields are required:
+**NOTE**: You can set the required fields in the environment configuration (`environment.yaml`), it's a mapping between required arguments and their corresponding environment variables. On the Vector **Killarney** Cluster environment, the required fields are:
   * `--account`, `-A`: The Slurm account, this argument can be set to default by setting environment variable `VEC_INF_ACCOUNT`.
   * `--work-dir`, `-D`: A working directory other than your home directory, this argument can be set to default by seeting environment variable `VEC_INF_WORK_DIR`.
 

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"