diff --git a/.claude/agents/exec-local-slurm.md b/.claude/agents/exec-local-slurm.md new file mode 100644 index 000000000000..795f5751441c --- /dev/null +++ b/.claude/agents/exec-local-slurm.md @@ -0,0 +1,31 @@ +--- +name: exec-local-slurm +description: > + Execute a TensorRT-LLM workload on a local Slurm cluster. Supports persistent + allocation (allocate once via nohup salloc, reuse across runs) and one-shot + sbatch. Workflow-agnostic — handles pytest, eval, benchmark, and custom + scripts identically. The orchestrator (typically trtllm-case-executor) writes + a job spec to /job_spec.json and invokes this agent to run it. +tools: ["Read", "Write", "Edit", "Bash", "Grep", "Glob"] +model: sonnet +license: Apache-2.0 +--- + +You are the local Slurm executor agent. Load the `exec-local-slurm` skill (`trtllm-agent-toolkit:exec-local-slurm`) and follow its procedure exactly. The skill is the single source of truth for the execution flow; this agent file only adds the contract with the caller and the invariants that must hold across every run. + +## Input + +The caller passes a path to a job spec — typically `/job_spec.json` — plus a short summary of the fields used in this run. **Read `job_spec.json` first.** The skill's "Input (from orchestrator prompt)" section enumerates every field it consumes (`script_path`, `work_dir`, `model_name`, `workflow_type`, `success_patterns`, `failure_patterns`, `log_file_pattern`, `monitor_timeout_seconds`, `persistent_mode`, `release_allocation`, `alloc_time_limit`, `docker_image`, `container_name`, `container_mounts`, `repo_root`, `slurm_params`). + +Do not re-derive any field that `trtllm-case-executor` already wrote into the spec. + +## Invariants + +- **Hang detection.** Poll the log periodically; on a case-insensitive `hang detected` match, kill the process group and report `HANG_DETECTED`. Implementation lives in the skill. +- **Wall-clock limit.** Honor `monitor_timeout_seconds` (default `3600`). On timeout, kill and report `TIMEOUT`. +- **Persistent allocation lifecycle.** Default `persistent_mode=true`; reuse an existing allocation when present and valid. Only release when `release_allocation=true` is explicitly set — never auto-release. +- **Single source of truth.** `node_count`, `job_name`, `container_image`, and `slurm_params` come from `job_spec.json`. Never recompute them or re-grep `current_image_tags.properties`. + +## Output + +Return a single report to the caller with: task type, status (`PASSED` / `FAILED` / `TIMEOUT` / `HANG_DETECTED` / `OUT_OF_MEMORY` / `CANCELLED` / `ERROR` / `BUILD_FAILED`), Slurm job id (and allocation id when persistent), log file path, summary, and any error excerpts (last ~100 lines on build/job failure). Do not perform follow-up actions beyond what the skill prescribes. diff --git a/.claude/agents/exec-remote-slurm.md b/.claude/agents/exec-remote-slurm.md new file mode 100644 index 000000000000..496493d0aa56 --- /dev/null +++ b/.claude/agents/exec-remote-slurm.md @@ -0,0 +1,40 @@ +--- +name: exec-remote-slurm +description: > + Execute a TensorRT-LLM workload on a remote Slurm cluster via SSH. Resolves + the cluster (explicit name or auto-select from device_type + + required_devices_per_node), handles MFA-aware SSH, seeds the remote checkout + from a local repo URL/branch, submits jobs with pyxis/enroot, tails logs, + and reports back. The orchestrator (typically trtllm-case-executor) writes a + job spec to /job_spec.json and invokes this agent to run it. +tools: ["Read", "Write", "Edit", "Bash", "Grep", "Glob"] +model: sonnet +license: Apache-2.0 +--- + +You are the remote Slurm executor agent. Load the `exec-remote-slurm` skill (`trtllm-agent-toolkit:exec-remote-slurm`) and follow its procedure exactly. The skill is the single source of truth for the execution flow; this agent file only adds the contract with the caller and the invariants that must hold across every run. + +## Input + +The caller passes a path to a job spec — typically `/job_spec.json` — plus a short summary of the fields used in this run. **Read `job_spec.json` first.** The skill's "Input" section enumerates every field it consumes (`script_path`, `script_name`, `work_dir`, `model_name`, `workflow_type`, `success_patterns`, `failure_patterns`, `log_file_pattern`, `slurm_cluster`, `ssh_host`, `slurm_user`, `remote_cwd`, `remote_work_dir`, `slurm_password`, `extra_files`, `repo_url`, `repo_branch`, `device_type`, `total_required_devices`, `required_devices_per_node`, `container_image`, `node_count`, `job_name`, `monitor_timeout_seconds`). + +Do not re-derive any field that `trtllm-case-executor` already wrote into the spec. + +## Cluster resolution + +**Precondition — optional dependency.** Before either bullet below, check whether the `skills/internal-env-info/` directory exists in the toolkit. If it does **not**, do not attempt to load any reference file or invoke the skill. Proceed using only the cluster fields the orchestrator wrote into `job_spec.json` (`ssh_host`, `slurm_user`, `remote_cwd`, `partition`, `account`, `container_image`, `mounts`, `gpus_per_node`, `mfa_style`). If a required field is also absent, stop and ask the user to supply it — do **not** report the missing skill as an error. + +- **Explicit cluster** (`slurm_cluster` is set) → invoke the `internal-env-info` skill in single-cluster mode to fetch per-cluster info (`mfa_style`, `default_models_repo`, `default_user_root_dir`, `gpus_per_node`). Connection fields (`ssh_host`, `slurm_user`, `remote_cwd`, `account`, `partition`, `mounts`, `container_image`) come from `job_spec.json` — there is no per-cluster connection-config file to parse. +- **Auto-select** (only hardware constraints are present) → invoke the `internal-env-info` skill in constraint-based mode, passing `device_type` and `required_devices_per_node` (and optionally `total_required_devices`). Never re-implement the constraint filter. + +## Invariants + +- **`mfa_style` decides the SSH path.** `false` → direct; `true` → MFA flow; `null` → probe direct first, then fall back, and ask the user to update the cluster reference. Do **not** probe-and-fall-back on the SSH error string when `mfa_style` is known. +- **Hang detection + `monitor_timeout_seconds`.** Same semantics as local execution; implementation lives in the skill. +- **Pass-through fields.** `container_image`, `node_count`, and `job_name` come from `job_spec.json` and are used verbatim. Apply transport-specific URL rewrites (e.g., enroot `/` → `#`) at use-time. Never re-grep `current_image_tags.properties` or reconstruct `job_name` from `account` / `subproject` / `detail`. +- **Remote repo bootstrap.** Use `repo_url` and `repo_branch` from `job_spec.json` to ensure the remote checkout matches the local one before submission. +- **`node_count` is authoritative.** Use it directly as `--nodes`; never recompute from totals. + +## Output + +Return a single report to the caller with: task type, status (`PASSED` / `FAILED` / `TIMEOUT` / `HANG_DETECTED` / `OUT_OF_MEMORY` / `CANCELLED` / `ERROR` / `BUILD_FAILED`), remote Slurm job id, remote log path (and a local copy when synced back), summary, and any error excerpts (last ~100 lines on build/job failure). Do not perform follow-up actions beyond what the skill prescribes. diff --git a/.claude/agents/trtllm-test-specialist.md b/.claude/agents/trtllm-test-specialist.md new file mode 100644 index 000000000000..0952a24f01c2 --- /dev/null +++ b/.claude/agents/trtllm-test-specialist.md @@ -0,0 +1,18 @@ +--- +name: trtllm-test-specialist +description: > + Runs model-level and module-level tests for TensorRT-LLM. Classifies the test + scope (module test or model test), builds the appropriate test commands, and + delegates execution to trtllm-case-executor. Supports functionality/smoke + tests, benchmarks, and evaluations. Writes structured test reports to a + caller-specified path. +tools: ["Read", "Write", "Edit", "Bash", "Grep", "Glob"] +license: Apache-2.0 +--- + +Role: dispatch TRT-LLM model-level and module-level test requests. + +Load the `trtllm-agent-toolkit:trtllm-test-specialist` skill, pass the caller's parameters through verbatim, and return its result. + +- If the caller supplies `report_file`, write the report to that exact path — do not substitute the skill's default (`./-auto-test-report.md`) or invent your own. +- Return the skill's status and final report path verbatim; do not re-summarize or rename. diff --git a/.claude/skills/exec-env-check/SKILL.md b/.claude/skills/exec-env-check/SKILL.md new file mode 100644 index 000000000000..3b91216a74b1 --- /dev/null +++ b/.claude/skills/exec-env-check/SKILL.md @@ -0,0 +1,142 @@ +--- +name: exec-env-check +description: >- + Check the local execution environment for GPU availability, Docker support, + and Slurm access. Returns the execution scenario (`satisfied, local, docker`, + `satisfied, local, direct`, `satisfied, slurm, local`, or `not_satisfied`), + the number of available GPUs, and the GPU type. On Slurm login nodes + without local GPUs, the cluster is identified by delegating the hostname + to internal-env-info (hostname-based mode), which owns the + hostname → cluster_name patterns; GPU type and gpus_per_node then come + from that skill's reference files. If internal-env-info is not + installed, the scenario falls back to `not_satisfied` without probing + compute nodes via srun. +tags: [infrastructure, slurm, environment] +license: Apache-2.0 +metadata: + author: NVIDIA Corporation +--- + +# TensorRT-LLM Environment Check + +Detect whether the current machine can run a GPU workload locally (Docker) or via Slurm, and report hardware details. + +## Input + +| Field | Description | Required | +|-------|-------------|----------| +| `required_devices` | Minimum number of GPUs needed | Yes | +| `account` | Slurm account (unused for GPU probing, kept for compatibility) | No | + +## Procedure + +### 1. Check local GPUs + +```bash +timeout 5 nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null +``` + +If this succeeds: +- Count the number of GPU lines → `available_gpus` +- Extract the GPU name from the first line → `device_type` (e.g., `NVIDIA B200`, `NVIDIA H100 80GB HBM3`) +- Normalize `device_type`: strip `NVIDIA ` prefix and trailing memory info to get the short name (e.g., `B200`, `H100`, `A100`, `L40S`, `RTX 6000`) + +If `nvidia-smi` fails or returns no GPUs → `available_gpus = 0`, `device_type = null`. + +### 2. Check if local GPUs are sufficient + +If `available_gpus >= required_devices`, continue to step 2a to determine whether Docker is available on this host. + +### 2a. Check Docker availability + +```bash +command -v docker >/dev/null 2>&1 && timeout 3 docker info >/dev/null 2>&1 +``` + +- If the command succeeds (Docker CLI exists **and** the daemon responds) → **Result**: `satisfied, local, docker` +- Otherwise (Docker CLI missing, daemon not running, or permission denied) → **Result**: `satisfied, local, direct` + +Include `available_gpus` in the result. Do NOT include `device_type` — local execution does not need it. + +### 3. Check Slurm availability + +If local GPUs are insufficient (or as additional detection), check for Slurm: + +```bash +which squeue 2>/dev/null && squeue --version 2>/dev/null +``` + +If Slurm is NOT available → go to step 5. + +### 4. Resolve GPU type + +**Optional dependency.** Before doing anything else in this step, check whether `skills/internal-env-info/` exists in the toolkit. If it does **not**, skip this step entirely and return `scenario: not_satisfied` with `available_gpus = 0`, `device_type = null`, `gpus_per_node = null`, `cluster_name = null`, `default_models_repo = null`, and `default_user_root_dir = null`. Do **not** report this as an error — the skill is an internal-only dependency. + +When Slurm is available but local `nvidia-smi` returned no GPUs or `device_type` is null (login nodes typically have no GPUs), capture the hostname and delegate cluster identification to `internal-env-info`: + +```bash +hostname -f 2>/dev/null || hostname +``` + +- Pass the captured hostname to `internal-env-info` in **hostname-based mode**. That skill owns the NVIDIA-internal login-host patterns and the hostname → `` mapping; do **not** parse the hostname or hard-code any cluster identifier here. +- It returns the standard output template (`cluster_name`, `device_type`, `gpus_per_node`, `default_models_repo`, `default_user_root_dir`) plus supplementary field (`mfa_style`). +- Set `available_gpus` = `gpus_per_node` (if resolved). +- Set `cluster_name` = the value returned (a placeholder `` token in this skill). The orchestrator uses this to fetch per-cluster info (`mfa_style`, `default_models_repo`, `default_user_root_dir`, `gpus_per_node`) from `internal-env-info`; connection fields (`mounts`, `ssh_host`, `partition`, etc.) come from caller-supplied inputs in `job_spec.json` (with `internal-env-info` default values / ask-the-user fallbacks). +- Set `default_user_root_dir` = the user root directory returned (with `` substituted with the actual SLURM username). Set to `null` if not found. +- If `internal-env-info` returns `null` for `cluster_name` (no pattern matched), set `device_type = null`, `cluster_name = null`, `default_user_root_dir = null` and fall through to Step 5 (`not_satisfied`). + +**Result**: `satisfied, slurm, local` with `device_type`, `gpus_per_node`, `cluster_name`, and `default_user_root_dir`. + +### 5. Not satisfied + +If neither local GPUs nor Slurm is available: +- **Result**: `not_satisfied` + +## Output + +Return a single structured result: + +``` +scenario: +available_gpus: +device_type: +gpus_per_node: +cluster_name: +default_models_repo: +default_user_root_dir: +``` + +**Examples:** + +``` +scenario: satisfied, local, docker +available_gpus: 4 +``` + +``` +scenario: satisfied, local, direct +available_gpus: 4 +``` + +``` +scenario: satisfied, slurm, local +available_gpus: 4 +device_type: B200 +gpus_per_node: 4 +cluster_name: +default_user_root_dir: / +``` + +``` +scenario: not_satisfied +available_gpus: 0 +device_type: null +cluster_name: null +``` + +## Rules + +- Never install drivers or modify the system +- If `nvidia-smi` hangs, use a 5-second timeout: `timeout 5 nvidia-smi ...` +- GPU type is derived from the hostname by matching the cluster name via the `internal-env-info` skill when it is installed — no `srun` allocation needed. If that skill is absent, the scenario falls back to `not_satisfied` (see Step 4); do not error. +- Report the GPU type exactly as found in the mapping (do not guess or fabricate) diff --git a/.claude/skills/exec-local-docker/SKILL.md b/.claude/skills/exec-local-docker/SKILL.md new file mode 100644 index 000000000000..cd73eed6e043 --- /dev/null +++ b/.claude/skills/exec-local-docker/SKILL.md @@ -0,0 +1,121 @@ +--- +name: exec-local-docker +description: >- + Execute a TensorRT-LLM workload locally in Docker. Runs a fully-resolved + Docker command in background, monitors completion, reads logs, and reports + results. Workflow-agnostic — does not need to know if the workload is pytest, + eval, benchmark, or a custom script. +tags: [docker, execution, infrastructure] +license: Apache-2.0 +metadata: + author: NVIDIA Corporation +--- + +# Local Docker Executor + +Run a Docker command locally, monitor it, and report results. + +## Input (from orchestrator prompt) + +The orchestrator passes these fields in the skill prompt: + +| Field | Description | +|-------|-------------| +| `docker_cmd` | Complete `docker run` command string, ready to execute | +| `work_dir` | Local work directory for logs and artifacts | +| `log_file` | Full path to the log file (output redirected here) | +| `model_name` | Short model name for reporting | +| `workflow_type` | `pytest`, `eval`, `custom`, or `benchmark` — for output parsing hints | +| `success_patterns` | Comma-separated patterns indicating success (e.g., `passed,accuracy:`) | +| `failure_patterns` | Comma-separated patterns indicating failure (e.g., `FAILED,Error,AssertionError`) | + +## Procedure + +### Step 0: Resolve Image and Build (when `build_project=true`) + +This executor owns image selection and the build for the local Docker target. Skip the entire step when `build_project=false`. + +1. **Detect the target GPU type.** Use `gpu_type` from `job_spec.json` if upstream env-check resolved it; otherwise probe locally with `nvidia-smi --query-gpu=name --format=csv,noheader | head -1`. +2. **Detect the host CPU arch** with `uname -m` (`x86_64` or `aarch64`). +3. **Resolve the container image.** Read `/jenkins/current_image_tags.properties` and pick the tag whose CPU-arch flavor matches the host. If the orchestrator already passed a `container_image` field in the job spec, use that and skip the lookup. +4. **Map GPU → build arch (`-a` flag):** `H100`/`H200` → `90-real`; `B200`/`GB200`/`B300`/`GB300` → `100-real`; `A100` → `80-real`; `L40S` → `89-real`. Default `100-real` when the GPU is unknown. +5. **Compile.** Invoke the `exec-local-compile` skill with `repo_dir=`, `image=`, `arch=`. Wait for completion. +6. **On failure**, do not launch the workload. Report `BUILD_FAILED` with the last 100 lines of the compile log. + +`build_project`, `gpu_type`, `repo_root`, and (optionally) `container_image` come from `job_spec.json`. + +### Step 1: Launch + +Run the Docker command in background using `run_in_background`: + +```bash + 2>&1 | tee +``` + +Report to the orchestrator: "Launched locally, log at ``" + +### Step 2: Monitor for Hangs + +While waiting for the background process to complete, actively monitor the log +file for hang indicators. Launch a monitoring loop using `run_in_background`: + +```bash +while true; do + sleep 60 + if [ -f "" ] && grep -qi "hang detected" ""; then + echo "HANG_DETECTED: Found 'hang detected' in log file" + docker ps --filter "ancestor=" -q | xargs -r docker kill 2>/dev/null + exit 1 + fi +done +``` + +- If the monitor detects a hang, it kills the Docker container and exits with + code 1. The main background process will also terminate. +- When the main process completes normally (background notification received), + kill the monitoring loop (it is no longer needed). +- If a hang is detected, skip to Step 4 and report `HANG_DETECTED` status + instead of proceeding to normal result collection. + +### Step 3: Wait for Completion + +The Bash tool's `run_in_background` will notify when the process finishes +(either normally or because the container was killed by the hang monitor). + +### Step 4: Read Results + +On completion: + +1. **Read exit code** from the background command result. +2. **Read the last 100 lines** of `` using the Read tool. +3. **If exit code != 0**, also read the first 50 lines to catch early errors (import failures, setup crashes). +4. **Search for patterns**: + - Grep `` for each `success_patterns` entry + - Grep `` for each `failure_patterns` entry + +### Step 5: Report + +Return a structured result: + +``` +Status: PASSED | FAILED | ERROR | HANG_DETECTED +Exit code: +Log file: +Work directory: +Summary: +Errors: +``` + +### Output Parsing by Workflow Type + +- **pytest**: Look for `X passed, Y failed in Zs` summary line +- **eval**: Look for `accuracy:` or `score:` lines; check for `Expected accuracy >= X, but got Y` assertion +- **custom**: No specific patterns — report last 10 lines of output +- **benchmark**: Look for throughput/latency numbers + +## Rules + +- Never run the Docker command in foreground — always use `run_in_background` +- Never `cat` the full log file — use Read with offset/limit or tail +- If the Docker command fails immediately (exit code within seconds), check if the image exists locally +- Report results even if the log file is empty (container may have failed to start) diff --git a/.claude/skills/exec-local-slurm/SKILL.md b/.claude/skills/exec-local-slurm/SKILL.md new file mode 100644 index 000000000000..ecfd3a860fbc --- /dev/null +++ b/.claude/skills/exec-local-slurm/SKILL.md @@ -0,0 +1,443 @@ +--- +name: exec-local-slurm +description: >- + Submit and monitor a Slurm job on a local cluster. Supports two modes: + (1) Persistent allocation (default) — allocates nodes once via nohup salloc, + imports the container once, installs once, and reuses across runs by setting + SLURM env vars and running the sbatch script via bash. (2) One-shot sbatch — + submits a fully-generated Slurm script via sbatch, polls job status, reads + logs on completion, and reports results. Workflow-agnostic — handles pytest, + eval, benchmark, and custom scripts identically. +tags: [slurm, execution, infrastructure] +license: Apache-2.0 +metadata: + author: NVIDIA Corporation +--- + +# Local Slurm Executor + +Submit a Slurm job locally, monitor it, and report results. Uses persistent +allocation by default to eliminate queue wait, container import, and install +overhead on repeated runs. + +## Input (from orchestrator prompt) + +The orchestrator passes these fields in the skill prompt: + +| Field | Description | +|-------|-------------| +| `script_path` | Full local path to the generated `.slurm` script | +| `work_dir` | Local work directory for logs and artifacts | +| `model_name` | Short model name for reporting | +| `workflow_type` | `pytest`, `eval`, `custom`, or `benchmark` — for output parsing hints | +| `success_patterns` | Comma-separated patterns indicating success | +| `failure_patterns` | Comma-separated patterns indicating failure | +| `log_file_pattern` | Log filename pattern with `%j` placeholder (e.g., `llama_auto_test_%j.out`) | +| `persistent_mode` | Default: `true` for all local slurm workflows. Set to `false` to force one-shot sbatch (opt-out). | +| `release_allocation` | `true` to release the current allocation and stop. Only set when the user explicitly says no more jobs are needed. Default: `false`. Never auto-release. | +| `alloc_time_limit` | Walltime for the persistent allocation. Default: `04:00:00`. | +| `docker_image` | Container image for persistent container import. From `job_spec.json`. | +| `container_name` | Container name used in the `.slurm` script (e.g., `llama_auto_test`). Must match exactly. From `job_spec.json`. | +| `container_mounts` | Comma-separated mount mappings. From `job_spec.json`. | +| `repo_root` | Repo root path (for locating state file and project path). | +| `slurm_params` | Slurm parameters object: `partition`, `account`, `nodes`, `ntasks`, `ntasks_per_node`, `gpus_per_node`. From `job_spec.json`. | + +## Procedure + +### Pre-step: Resolve Image and Build (when `build_project=true`) + +This executor owns image selection and the build for the local SLURM cluster. Skip the entire pre-step when `build_project=false`. + +1. **Read the cluster GPU type** from `device_type` in `job_spec.json` (resolved upstream by env-check / internal-env-info). If absent, probe from any compute node via `srun -p --ntasks-per-node=1 nvidia-smi --query-gpu=name --format=csv,noheader | head -1`. If `skills/internal-env-info/` is not installed, skip the upstream lookup silently and rely on the srun probe — do not report the missing skill as an error. +2. **Determine the partition's CPU arch.** Look up `job_spec.slurm_env.partitions[]` for the entry where `name == partition` and read its `arch` (`x86_64` / `aarch64`). Case-executor's Step 2.5 already detected this — do **not** run `scontrol show node` here. If `slurm_env` is absent or the entry's `arch` is null (case-executor ran without SLURM tools), fall back to inferring from `device_type`: Grace-based parts (`GB*` / `GH*`) → `aarch64`; everything else → `x86_64`. +3. **Resolve the container image.** If `docker_image` is already set in `job_spec.json`, use it directly. Otherwise read `/jenkins/current_image_tags.properties` and pick the tag matching the partition's CPU arch. +4. **Map GPU → build arch (`-a` flag):** `H100`/`H200` → `90-real`; `B200`/`GB200`/`B300`/`GB300` → `100-real`; `A100` → `80-real`; `L40S` → `89-real`. Default `100-real`. +5. **Compile.** Invoke the `exec-slurm-compile` skill with `repo_dir=`, `partition`, `account`, `container_image=`, `user_root_dir`, `arch`. Wait for completion. +6. **On failure**, do not allocate or run the workload. Report `BUILD_FAILED` with the last 100 lines of the build log. + +`build_project`, `device_type`, `repo_root`, `partition`, `account`, `user_root_dir`, and (optionally) `docker_image` come from `job_spec.json`. + +### Step 0: Allocation Management + +This step runs before any execution. It determines whether to reuse an existing +persistent allocation, create a new one, or fall through to one-shot sbatch. + +#### Step 0A — Release mode + +If `release_allocation=true`: + +1. Read `/work_dirs/.slurm_alloc.json` +2. If file exists and `job_id` is present: + ```bash + scancel + ``` +3. Delete the state file and `.salloc.log` +4. Report: "Allocation released." Stop. + +#### Step 0B — One-shot mode + +If `persistent_mode=false`, skip entirely to Step 1 (One-Shot Path). + +#### Step 0C — Validate existing allocation + +Read `/work_dirs/.slurm_alloc.json`. + +- If state file **does not exist** → go to Step 0D. +- If state file exists, validate the allocation: + ```bash + squeue -j -h -o "%T %L" + ``` + - **RUNNING + remaining > 5 min + params compatible** → **reuse** (skip to Step 0F) + - **RUNNING + remaining <= 5 min** → warn user ("Allocation expiring soon"), `scancel `, delete state file, go to Step 0D + - **RUNNING + params incompatible** → `scancel `, delete state file, go to Step 0D + - **Empty output / error** (job gone) → stale state file, delete it, go to Step 0D + +**Params compatibility check:** +- `partition` must match +- `nodes` in state must be **>=** requested nodes (a 2-node allocation can serve 1-node jobs) +- `docker_image` must match + +#### Step 0D — Justify and prepare allocation + +Before allocating, ensure no orphaned allocations exist and log the reason: + +1. Check for existing persistent jobs: + ```bash + squeue -u $(whoami) -h -o "%i %T %j" --name=-trtllm.persistent + ``` +2. If found → orphan (state file was missing/corrupt). Cancel it: + ```bash + scancel + ``` + Log: "Released orphaned allocation — state file was missing." +3. Log justification: "No reusable allocation found. Allocating node(s) on partition for ." + +#### Step 0E — Allocate new + +Convert `alloc_time_limit` from `HH:MM:SS` to seconds (e.g., `04:00:00` → `14400`). + +```bash +mkdir -p /work_dirs +nohup salloc --partition= --account= \ + --nodes= --time= \ + --job-name=-trtllm.persistent \ + sleep \ + > /work_dirs/.salloc.log 2>&1 & +``` + +Retrieve job ID: +```bash +squeue -u $(whoami) -h -o "%i" --name=-trtllm.persistent +``` + +If no job appears after 10 seconds, read `/work_dirs/.salloc.log` +for errors (bad partition, invalid account, etc.) and report to user. + +Otherwise, poll until RUNNING (every 10s, max 60 polls). Once RUNNING, get +the nodelist: +```bash +squeue -j -h -o "%N" +``` + +**Import container** on all nodes (using the **same container name** as the +`.slurm` script — critical for pyxis reuse): +```bash +srun --jobid= -N --ntasks-per-node=1 \ + --container-image= \ + --container-name= true +``` + +**Warm up filesystem mounts** on all nodes — `ls` each mounted path so +Lustre/NFS metadata is cached for later use: +```bash +srun --jobid= -N --ntasks-per-node=1 \ + --container-name= \ + --container-mounts= \ + bash -c 'for p in ...; do ls "$p" > /dev/null 2>&1; done' +``` +Parse mount targets from `container_mounts` — the right-hand side of each +`host:container` pair. + +**Check GPU status** on all nodes (no container needed — `nvidia-smi` is on +the host): +```bash +srun --jobid= -N --ntasks-per-node=1 \ + bash -c 'echo "=== $(hostname) ===" && nvidia-smi --query-compute-apps=pid,name,used_memory --format=csv,noheader' +``` +- If output shows no processes → GPUs are clean, proceed +- If unexpected processes found → warn user: "GPU processes found on + : . These may interfere with the job." + +**Run install** on all nodes with `--container-writable` (so packages persist +in the named container across srun calls). Install all common requirements +upfront so that subsequent jobs of any workflow type can skip install: +```bash +srun --jobid= -N --ntasks-per-node=1 \ + --container-name= --container-writable \ + --container-mounts= \ + bash -c 'cd && pip install -e . && pip install -r requirements-dev.txt && \ + if [ -f examples/trtllm-eval/requirements.txt ]; then pip install -r examples/trtllm-eval/requirements.txt; fi' +``` +For custom workflow with `skip_install=true`, skip this step entirely. + +**Write state file** `/work_dirs/.slurm_alloc.json`: +```json +{ + "job_id": "", + "container_name": "", + "nodelist": "", + "partition": "", + "account": "", + "nodes": , + "gpus_per_node": , + "docker_image": "", + "container_mounts": "", + "allocated_at": "", + "time_limit": "", + "installed": true +} +``` + +Proceed to Step 1 (Persistent Path). + +#### Step 0F — Reuse path validation + +Allocation is valid. Check if container name matches current request: + +- If `container_name` in state file **matches** `container_name` from + `job_spec.json` → proceed to Step 1 (Persistent Path). Container and install + are already set up. +- If **different** (model changed) → import the new container, warm up mounts, + and install: + ```bash + srun --jobid= -N --ntasks-per-node=1 \ + --container-image= \ + --container-name= true + srun --jobid= -N --ntasks-per-node=1 \ + --container-name= \ + --container-mounts= \ + bash -c 'for p in ; do ls "$p" > /dev/null 2>&1; done' + srun --jobid= -N --ntasks-per-node=1 \ + --container-name= --container-writable \ + --container-mounts= \ + bash -c 'cd && pip install -e . && pip install -r requirements-dev.txt' + ``` + Update `container_name` in state file. Proceed to Step 1 (Persistent Path). + +--- + +### Step 1: Execute + +Two execution paths depending on `persistent_mode`. + +#### Persistent Path (persistent_mode=true) + +**Pre-flight time check** — verify the allocation has enough remaining time: +```bash +squeue -j -h -o "%L" +``` +If remaining time < job's `time_limit` from `slurm_params`, warn user: +"Allocation has left but job expects . The job may +be killed early." + +**Pre-flight GPU check** — verify GPUs are not occupied by leftover processes: +```bash +srun --jobid= -N --ntasks-per-node=1 \ + bash -c 'procs=$(nvidia-smi --query-compute-apps=pid,name,used_memory --format=csv,noheader 2>/dev/null); [ -n "$procs" ] && echo "WARNING: GPU processes on $(hostname): $procs"' +``` +If unexpected processes found, warn user before proceeding. + +**Run the sbatch script** with SLURM env vars — the `#SBATCH` directives are +comments when run via `bash`; the inner `srun` inherits the env vars and uses +the persistent allocation: +```bash +export SLURM_JOB_ID= +export SLURM_JOB_NUM_NODES= +export SLURM_NNODES= +export SLURM_NTASKS= +export SLURM_NTASKS_PER_NODE= +export SLURM_NODELIST= +export SLURM_JOB_NODELIST= +bash 2>&1 | tee /_.log +``` + +Run with `run_in_background`. Container reuse is automatic (pyxis skips +re-import when `--container-name` already exists on the node). The install +step (Step 1 srun) always runs to keep the container up to date. + +**Step cancellation:** If a job hangs or the user wants to abort, cancel just +the srun step without killing the allocation: +```bash +# List active steps: +squeue -s -j +# Cancel a specific step (e.g., step 0): +scancel .0 +``` +The allocation stays RUNNING — new jobs can run immediately. The job name +`-trtllm.persistent` uniquely identifies allocations created by this +skill via `squeue --name=-trtllm.persistent`. + +#### One-Shot Path (persistent_mode=false) + +```bash +sbatch +``` + +Parse the job ID from output: `Submitted batch job `. + +If sbatch fails, report the error immediately and stop. + +### Step 2: Report Submission + +**One-shot mode:** +``` +Job ID: +Script: +Work directory: +Log files: / +Monitor: squeue -j +``` + +**Persistent mode:** +``` +Allocation: (persistent, remaining: