vec-inf 0.6.1__py3-none-any.whl → 0.7.1__py3-none-any.whl

vec_inf/client/_slurm_script_generator.py

@@ -1,6 +1,6 @@
-"""Class for generating SLURM scripts to run vLLM servers.
+"""Class for generating Slurm scripts to run vLLM servers.
 
-This module provides functionality to generate SLURM scripts for running vLLM servers
+This module provides functionality to generate Slurm scripts for running vLLM servers
 in both single-node and multi-node configurations.
 """
 
@@ -8,51 +8,56 @@ from datetime import datetime
 from pathlib import Path
 from typing import Any
 
-from vec_inf.client._client_vars import (
-    SLURM_JOB_CONFIG_ARGS,
+from vec_inf.client._client_vars import SLURM_JOB_CONFIG_ARGS
+from vec_inf.client._slurm_templates import (
+    BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE,
+    BATCH_SLURM_SCRIPT_TEMPLATE,
     SLURM_SCRIPT_TEMPLATE,
 )
 
 
 class SlurmScriptGenerator:
-    """A class to generate SLURM scripts for running vLLM servers.
+    """A class to generate Slurm scripts for running vLLM servers.
 
-    This class handles the generation of SLURM scripts for both single-node and
+    This class handles the generation of Slurm scripts for both single-node and
     multi-node configurations, supporting different virtualization environments
-    (venv or singularity).
+    (venv or singularity/apptainer).
 
     Parameters
     ----------
-    params : dict[str, Any]
-        Configuration parameters for the SLURM script. Contains settings for job
-        configuration, model parameters, and virtualization environment.
+        params : dict[str, Any]
+            Configuration parameters for the Slurm script.
     """
 
     def __init__(self, params: dict[str, Any]):
-        """Initialize the SlurmScriptGenerator with configuration parameters.
-
-        Parameters
-        ----------
-        params : dict[str, Any]
-            Configuration parameters for the SLURM script.
-        """
         self.params = params
         self.is_multinode = int(self.params["num_nodes"]) > 1
-        self.use_singularity = self.params["venv"] == "singularity"
+        self.use_container = (
+            self.params["venv"] == "singularity" or self.params["venv"] == "apptainer"
+        )
         self.additional_binds = self.params.get("bind", "")
         if self.additional_binds:
             self.additional_binds = f" --bind {self.additional_binds}"
         self.model_weights_path = str(
-            Path(params["model_weights_parent_dir"], params["model_name"])
+            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.
+        """Generate the complete Slurm script content.
 
         Returns
         -------
         str
-            The complete SLURM script as a string.
+            The complete Slurm script as a string.
         """
         script_content = []
         script_content.append(self._generate_shebang())
@@ -61,12 +66,12 @@ class SlurmScriptGenerator:
         return "\n".join(script_content)
 
     def _generate_shebang(self) -> str:
-        """Generate the SLURM script shebang with job specifications.
+        """Generate the Slurm script shebang with job specifications.
 
         Returns
         -------
         str
-            SLURM shebang containing job specifications.
+            Slurm shebang containing job specifications.
         """
         shebang = [SLURM_SCRIPT_TEMPLATE["shebang"]["base"]]
         for arg, value in SLURM_JOB_CONFIG_ARGS.items():
@@ -88,8 +93,8 @@ class SlurmScriptGenerator:
             Server initialization script content.
         """
         server_script = ["\n"]
-        if self.use_singularity:
-            server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["singularity_setup"]))
+        if self.use_container:
+            server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["container_setup"]))
         server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["env_vars"]))
         server_script.append(
             SLURM_SCRIPT_TEMPLATE["imports"].format(src_dir=self.params["src_dir"])
@@ -97,13 +102,14 @@ class SlurmScriptGenerator:
         if self.is_multinode:
             server_setup_str = "\n".join(
                 SLURM_SCRIPT_TEMPLATE["server_setup"]["multinode"]
-            )
-            if self.use_singularity:
+            ).format(gpus_per_node=self.params["gpus_per_node"])
+            if self.use_container:
                 server_setup_str = server_setup_str.replace(
-                    "SINGULARITY_PLACEHOLDER",
-                    SLURM_SCRIPT_TEMPLATE["singularity_command"].format(
+                    "CONTAINER_PLACEHOLDER",
+                    SLURM_SCRIPT_TEMPLATE["container_command"].format(
                         model_weights_path=self.model_weights_path,
                         additional_binds=self.additional_binds,
+                        env_str=self.env_str,
                     ),
                 )
         else:
@@ -123,7 +129,7 @@ class SlurmScriptGenerator:
         """Generate the vLLM server launch command.
 
         Creates the command to launch the vLLM server, handling different virtualization
-        environments (venv or singularity).
+        environments (venv or singularity/apptainer).
 
         Returns
         -------
@@ -131,13 +137,13 @@ class SlurmScriptGenerator:
             Server launch command.
         """
         launcher_script = ["\n"]
-        if self.use_singularity:
+        if self.use_container:
             launcher_script.append(
-                SLURM_SCRIPT_TEMPLATE["singularity_command"].format(
+                SLURM_SCRIPT_TEMPLATE["container_command"].format(
                     model_weights_path=self.model_weights_path,
                     additional_binds=self.additional_binds,
+                    env_str=self.env_str,
                 )
-                + " \\"
             )
         else:
             launcher_script.append(
@@ -158,21 +164,183 @@ class SlurmScriptGenerator:
         return "\n".join(launcher_script)
 
     def write_to_log_dir(self) -> Path:
-        """Write the generated SLURM script to the log directory.
+        """Write the generated Slurm script to the log directory.
 
         Creates a timestamped script file in the configured log directory.
 
         Returns
         -------
         Path
-            Path to the generated SLURM script file.
+            Path to the generated Slurm script file.
         """
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         script_path: Path = (
             Path(self.params["log_dir"])
-            / f"launch_{self.params['model_name']}_{timestamp}.slurm"
+            / f"launch_{self.params['model_name']}_{timestamp}.sbatch"
         )
 
         content = self._generate_script_content()
         script_path.write_text(content)
         return script_path
+
+
+class BatchSlurmScriptGenerator:
+    """A class to generate Slurm scripts for batch mode.
+
+    This class handles the generation of Slurm scripts for batch mode, which
+    launches multiple vLLM servers with different configurations in parallel.
+    """
+
+    def __init__(self, params: dict[str, Any]):
+        self.params = params
+        self.script_paths: list[Path] = []
+        self.use_container = (
+            self.params["venv"] == "singularity" or self.params["venv"] == "apptainer"
+        )
+        for model_name in self.params["models"]:
+            self.params["models"][model_name]["additional_binds"] = ""
+            if self.params["models"][model_name].get("bind"):
+                self.params["models"][model_name]["additional_binds"] = (
+                    f" --bind {self.params['models'][model_name]['bind']}"
+                )
+            self.params["models"][model_name]["model_weights_path"] = str(
+                Path(
+                    self.params["models"][model_name]["model_weights_parent_dir"],
+                    model_name,
+                )
+            )
+
+    def _write_to_log_dir(self, script_content: list[str], script_name: str) -> Path:
+        """Write the generated Slurm script to the log directory.
+
+        Returns
+        -------
+        Path
+            The Path object to the generated Slurm script file.
+        """
+        script_path = Path(self.params["log_dir"]) / script_name
+        script_path.touch(exist_ok=True)
+        script_path.write_text("\n".join(script_content))
+        return script_path
+
+    def _generate_model_launch_script(self, model_name: str) -> Path:
+        """Generate the bash script for launching individual vLLM servers.
+
+        Parameters
+        ----------
+        model_name : str
+            The name of the model to launch.
+
+        Returns
+        -------
+        Path
+            The bash script path for launching the vLLM server.
+        """
+        # Generate the bash script content
+        script_content = []
+        model_params = self.params["models"][model_name]
+        script_content.append(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["shebang"])
+        if self.use_container:
+            script_content.append(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["container_setup"])
+        script_content.append("\n".join(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["env_vars"]))
+        script_content.append(
+            "\n".join(
+                BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["server_address_setup"]
+            ).format(src_dir=self.params["src_dir"])
+        )
+        script_content.append(
+            "\n".join(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["write_to_json"]).format(
+                het_group_id=model_params["het_group_id"],
+                log_dir=self.params["log_dir"],
+                slurm_job_name=self.params["slurm_job_name"],
+                model_name=model_name,
+            )
+        )
+        if self.use_container:
+            script_content.append(
+                BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["container_command"].format(
+                    model_weights_path=model_params["model_weights_path"],
+                    additional_binds=model_params["additional_binds"],
+                )
+            )
+        script_content.append(
+            "\n".join(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["launch_cmd"]).format(
+                model_weights_path=model_params["model_weights_path"],
+                model_name=model_name,
+            )
+        )
+        for arg, value in model_params["vllm_args"].items():
+            if isinstance(value, bool):
+                script_content.append(f"    {arg} \\")
+            else:
+                script_content.append(f"    {arg} {value} \\")
+        script_content[-1] = script_content[-1].replace("\\", "")
+        # Write the bash script to the log directory
+        launch_script_path = self._write_to_log_dir(
+            script_content, f"launch_{model_name}.sh"
+        )
+        self.script_paths.append(launch_script_path)
+        return launch_script_path
+
+    def _generate_batch_slurm_script_shebang(self) -> str:
+        """Generate the shebang for batch mode Slurm script.
+
+        Returns
+        -------
+        str
+            The shebang for batch mode Slurm script.
+        """
+        shebang = [BATCH_SLURM_SCRIPT_TEMPLATE["shebang"]]
+
+        for arg, value in SLURM_JOB_CONFIG_ARGS.items():
+            if self.params.get(value):
+                shebang.append(f"#SBATCH --{arg}={self.params[value]}")
+        shebang.append("#SBATCH --ntasks=1")
+        shebang.append("\n")
+
+        for model_name in self.params["models"]:
+            shebang.append(f"# ===== Resource group for {model_name} =====")
+            for arg, value in SLURM_JOB_CONFIG_ARGS.items():
+                model_params = self.params["models"][model_name]
+                if model_params.get(value) and value not in ["out_file", "err_file"]:
+                    shebang.append(f"#SBATCH --{arg}={model_params[value]}")
+            shebang[-1] += "\n"
+            shebang.append(BATCH_SLURM_SCRIPT_TEMPLATE["hetjob"])
+        # Remove the last hetjob line
+        shebang.pop()
+        return "\n".join(shebang)
+
+    def generate_batch_slurm_script(self) -> Path:
+        """Generate the Slurm script for launching multiple vLLM servers in batch mode.
+
+        Returns
+        -------
+        Path
+            The Slurm script for launching multiple vLLM servers in batch mode.
+        """
+        script_content = []
+
+        script_content.append(self._generate_batch_slurm_script_shebang())
+
+        for model_name in self.params["models"]:
+            model_params = self.params["models"][model_name]
+            script_content.append(f"# ===== Launching {model_name} =====")
+            launch_script_path = str(self._generate_model_launch_script(model_name))
+            script_content.append(
+                BATCH_SLURM_SCRIPT_TEMPLATE["permission_update"].format(
+                    script_name=launch_script_path
+                )
+            )
+            script_content.append(
+                "\n".join(BATCH_SLURM_SCRIPT_TEMPLATE["launch_model_scripts"]).format(
+                    het_group_id=model_params["het_group_id"],
+                    out_file=model_params["out_file"],
+                    err_file=model_params["err_file"],
+                    script_name=launch_script_path,
+                )
+            )
+        script_content.append("wait")
+
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        script_name = f"{self.params['slurm_job_name']}_{timestamp}.sbatch"
+        return self._write_to_log_dir(script_content, script_name)

vec_inf/client/_slurm_templates.py

@@ -0,0 +1,248 @@
+"""SLURM script templates for Vector Inference.
+
+This module contains the SLURM script templates for Vector Inference, including
+single-node, multi-node, and batch mode templates.
+"""
+
+from typing import TypedDict
+
+from vec_inf.client._slurm_vars import (
+    CONTAINER_LOAD_CMD,
+    CONTAINER_MODULE_NAME,
+    IMAGE_PATH,
+)
+
+
+CONTAINER_MODULE_NAME_UPPER = CONTAINER_MODULE_NAME.upper()
+
+
+class ShebangConfig(TypedDict):
+    """TypedDict for SLURM script shebang configuration.
+
+    Parameters
+    ----------
+    base : str
+        Base shebang line for all SLURM scripts
+    multinode : list[str]
+        Additional SLURM directives for multi-node configurations
+    """
+
+    base: str
+    multinode: list[str]
+
+
+class ServerSetupConfig(TypedDict):
+    """TypedDict for server setup configuration.
+
+    Parameters
+    ----------
+    single_node : list[str]
+        Setup commands for single-node deployments
+    multinode : list[str]
+        Setup commands for multi-node deployments, including Ray initialization
+    """
+
+    single_node: list[str]
+    multinode: list[str]
+
+
+class SlurmScriptTemplate(TypedDict):
+    """TypedDict for complete SLURM script template configuration.
+
+    Parameters
+    ----------
+    shebang : ShebangConfig
+        Shebang and SLURM directive configuration
+    container_setup : list[str]
+        Commands for container setup
+    imports : str
+        Import statements and source commands
+    container_command : str
+        Template for container execution command
+    activate_venv : str
+        Template for virtual environment activation
+    server_setup : ServerSetupConfig
+        Server initialization commands for different deployment modes
+    find_vllm_port : list[str]
+        Commands to find available ports for vLLM server
+    write_to_json : list[str]
+        Commands to write server configuration to JSON
+    launch_cmd : list[str]
+        vLLM server launch commands
+    """
+
+    shebang: ShebangConfig
+    container_setup: list[str]
+    imports: str
+    env_vars: list[str]
+    container_command: str
+    activate_venv: str
+    server_setup: ServerSetupConfig
+    find_vllm_port: list[str]
+    write_to_json: list[str]
+    launch_cmd: list[str]
+
+
+SLURM_SCRIPT_TEMPLATE: SlurmScriptTemplate = {
+    "shebang": {
+        "base": "#!/bin/bash",
+        "multinode": [
+            "#SBATCH --exclusive",
+            "#SBATCH --tasks-per-node=1",
+        ],
+    },
+    "container_setup": [
+        CONTAINER_LOAD_CMD,
+        f"{CONTAINER_MODULE_NAME} exec {IMAGE_PATH} ray stop",
+    ],
+    "imports": "source {src_dir}/find_port.sh",
+    "env_vars": [
+        f"export {CONTAINER_MODULE_NAME}_BINDPATH=${CONTAINER_MODULE_NAME}_BINDPATH,$(echo /dev/infiniband* | sed -e 's/ /,/g')"
+    ],
+    "container_command": f"{CONTAINER_MODULE_NAME} exec --nv {{env_str}} --bind {{model_weights_path}}{{additional_binds}} --containall {IMAGE_PATH} \\",
+    "activate_venv": "source {venv}/bin/activate",
+    "server_setup": {
+        "single_node": [
+            "\n# Find available port",
+            "head_node_ip=${SLURMD_NODENAME}",
+        ],
+        "multinode": [
+            "\n# Get list of nodes",
+            'nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST")',
+            "nodes_array=($nodes)",
+            "head_node=${{nodes_array[0]}}",
+            'head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)',
+            "\n# Start Ray head node",
+            "head_node_port=$(find_available_port $head_node_ip 8080 65535)",
+            "ray_head=$head_node_ip:$head_node_port",
+            'echo "Ray Head IP: $ray_head"',
+            'echo "Starting HEAD at $head_node"',
+            'srun --nodes=1 --ntasks=1 -w "$head_node" \\',
+            "    CONTAINER_PLACEHOLDER",
+            '    ray start --head --node-ip-address="$head_node_ip" --port=$head_node_port \\',
+            '    --num-cpus "$SLURM_CPUS_PER_TASK" --num-gpus {gpus_per_node} --block &',
+            "sleep 10",
+            "\n# Start Ray worker nodes",
+            "worker_num=$((SLURM_JOB_NUM_NODES - 1))",
+            "for ((i = 1; i <= worker_num; i++)); do",
+            "    node_i=${{nodes_array[$i]}}",
+            '    echo "Starting WORKER $i at $node_i"',
+            '    srun --nodes=1 --ntasks=1 -w "$node_i" \\',
+            "        CONTAINER_PLACEHOLDER",
+            '        ray start --address "$ray_head" \\',
+            '        --num-cpus "$SLURM_CPUS_PER_TASK" --num-gpus {gpus_per_node} --block &',
+            "    sleep 5",
+            "done",
+        ],
+    },
+    "find_vllm_port": [
+        "\nvllm_port_number=$(find_available_port $head_node_ip 8080 65535)",
+        'server_address="http://${head_node_ip}:${vllm_port_number}/v1"',
+    ],
+    "write_to_json": [
+        '\njson_path="{log_dir}/{model_name}.$SLURM_JOB_ID/{model_name}.$SLURM_JOB_ID.json"',
+        'jq --arg server_addr "$server_address" \\',
+        "    '. + {{\"server_address\": $server_addr}}' \\",
+        '    "$json_path" > temp.json \\',
+        '    && mv temp.json "$json_path"',
+    ],
+    "launch_cmd": [
+        "vllm serve {model_weights_path} \\",
+        "    --served-model-name {model_name} \\",
+        '    --host "0.0.0.0" \\',
+        "    --port $vllm_port_number \\",
+    ],
+}
+
+
+class BatchSlurmScriptTemplate(TypedDict):
+    """TypedDict for batch SLURM script template configuration.
+
+    Parameters
+    ----------
+    shebang : str
+        Shebang line for the script
+    hetjob : str
+        SLURM directive for hetjob
+    permission_update : str
+        Command to update permissions of the script
+    launch_model_scripts : list[str]
+        Commands to launch the vLLM server
+    """
+
+    shebang: str
+    hetjob: str
+    permission_update: str
+    launch_model_scripts: list[str]
+
+
+BATCH_SLURM_SCRIPT_TEMPLATE: BatchSlurmScriptTemplate = {
+    "shebang": "#!/bin/bash",
+    "hetjob": "#SBATCH hetjob\n",
+    "permission_update": "chmod +x {script_name}",
+    "launch_model_scripts": [
+        "\nsrun --het-group={het_group_id} \\",
+        "    --output={out_file} \\",
+        "    --error={err_file} \\",
+        "    {script_name} &\n",
+    ],
+}
+
+
+class BatchModelLaunchScriptTemplate(TypedDict):
+    """TypedDict for batch model launch script template configuration.
+
+    Parameters
+    ----------
+    shebang : str
+        Shebang line for the script
+    container_setup : list[str]
+        Commands for container setup
+    env_vars : list[str]
+        Environment variables to set
+    server_address_setup : list[str]
+        Commands to setup the server address
+    launch_cmd : list[str]
+        Commands to launch the vLLM server
+    container_command : str
+        Commands to setup the container command
+    """
+
+    shebang: str
+    container_setup: str
+    env_vars: list[str]
+    server_address_setup: list[str]
+    write_to_json: list[str]
+    launch_cmd: list[str]
+    container_command: str
+
+
+BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE: BatchModelLaunchScriptTemplate = {
+    "shebang": "#!/bin/bash\n",
+    "container_setup": f"{CONTAINER_LOAD_CMD}\n",
+    "env_vars": [
+        f"export {CONTAINER_MODULE_NAME}_BINDPATH=${CONTAINER_MODULE_NAME}_BINDPATH,$(echo /dev/infiniband* | sed -e 's/ /,/g')"
+    ],
+    "server_address_setup": [
+        "source {src_dir}/find_port.sh",
+        "head_node_ip=${{SLURMD_NODENAME}}",
+        "vllm_port_number=$(find_available_port $head_node_ip 8080 65535)",
+        'server_address="http://${{head_node_ip}}:${{vllm_port_number}}/v1"\n',
+        "echo $server_address\n",
+    ],
+    "write_to_json": [
+        "het_job_id=$(($SLURM_JOB_ID+{het_group_id}))",
+        'json_path="{log_dir}/{slurm_job_name}.$het_job_id/{model_name}.$het_job_id.json"',
+        'jq --arg server_addr "$server_address" \\',
+        "    '. + {{\"server_address\": $server_addr}}' \\",
+        '    "$json_path" > temp_{model_name}.json \\',
+        '    && mv temp_{model_name}.json "$json_path"\n',
+    ],
+    "container_command": f"{CONTAINER_MODULE_NAME} exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {IMAGE_PATH} \\",
+    "launch_cmd": [
+        "vllm serve {model_weights_path} \\",
+        "    --served-model-name {model_name} \\",
+        '    --host "0.0.0.0" \\',
+        "    --port $vllm_port_number \\",
+    ],
+}

vec_inf/client/_slurm_vars.py

@@ -0,0 +1,86 @@
+"""Slurm cluster configuration variables."""
+
+import os
+import warnings
+from pathlib import Path
+from typing import Any, TypeAlias
+
+import yaml
+from typing_extensions import Literal
+
+
+CACHED_CONFIG_DIR = Path("/model-weights/vec-inf-shared")
+
+
+def load_env_config() -> dict[str, Any]:
+    """Load the environment configuration."""
+
+    def load_yaml_config(path: Path) -> dict[str, Any]:
+        """Load YAML config with error handling."""
+        try:
+            with path.open() as f:
+                return yaml.safe_load(f) or {}
+        except FileNotFoundError as err:
+            raise FileNotFoundError(f"Could not find config: {path}") from err
+        except yaml.YAMLError as err:
+            raise ValueError(f"Error parsing YAML config at {path}: {err}") from err
+
+    cached_config_path = CACHED_CONFIG_DIR / "environment.yaml"
+    default_path = (
+        cached_config_path
+        if cached_config_path.exists()
+        else Path(__file__).resolve().parent.parent / "config" / "environment.yaml"
+    )
+    config = load_yaml_config(default_path)
+
+    user_path = os.getenv("VEC_INF_CONFIG_DIR")
+    if user_path:
+        user_path_obj = Path(user_path, "environment.yaml")
+        if user_path_obj.exists():
+            user_config = load_yaml_config(user_path_obj)
+            config.update(user_config)
+        else:
+            warnings.warn(
+                f"WARNING: Could not find user config directory: {user_path}, revert to default config located at {default_path}",
+                UserWarning,
+                stacklevel=2,
+            )
+
+    return config
+
+
+_config = load_env_config()
+
+# Extract path values
+IMAGE_PATH = _config["paths"]["image_path"]
+
+# Extract containerization info
+CONTAINER_LOAD_CMD = _config["containerization"]["module_load_cmd"]
+CONTAINER_MODULE_NAME = _config["containerization"]["module_name"]
+
+# Extract limits
+MAX_GPUS_PER_NODE = _config["limits"]["max_gpus_per_node"]
+MAX_NUM_NODES = _config["limits"]["max_num_nodes"]
+MAX_CPUS_PER_TASK = _config["limits"]["max_cpus_per_task"]
+
+
+# Create dynamic Literal types
+def create_literal_type(values: list[str], fallback: str = "") -> Any:
+    """Create a Literal type from a list, with configurable fallback."""
+    if not values:
+        return Literal[fallback]
+    return Literal[tuple(values)]
+
+
+QOS: TypeAlias = create_literal_type(_config["allowed_values"]["qos"])  # type: ignore[valid-type]
+PARTITION: TypeAlias = create_literal_type(_config["allowed_values"]["partition"])  # type: ignore[valid-type]
+RESOURCE_TYPE: TypeAlias = create_literal_type(  # type: ignore[valid-type]
+    _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"]