trtllm-test-script-builder

# TensorRT-LLM Script Builder Resolve parameters and generate execution scripts for all workflow types. ## Input (from orchestrator prompt) The orchestrator passes these fields. Not all fields are present for every workflow. | Field | Description | Required | |-------|-------------|----------| | `workflow_type` | `pytest`, `eval`, `bench`, `custom`, or `perf_sanity` | Always | | `work_dir` | Absolute path to local work directory (already created) | Always | | `skill_dir` | Absolute path to the test-script-builder skill directory (`.claude/skills/trtllm-test-script-builder`) | Always | | `repo_root` | Absolute path to the TensorRT-LLM repo root | Always | | `test_cmd` | Pytest command string | pytest | | `eval_cmd` | `trtllm-eval` command string, or `evaluation`, `accuracy`, `score`, `dataset` | eval | | `bench_cmd` | `trtllm-bench` command (may chain `prepare-dataset` and `throughput`/`latency` with `&&`). Auto-parses `--tp`/`--pp`/`--ep` from the last `trtllm-bench` invocation. | bench | | `custom_cmd` | Custom command string, `test command`, non-pytest command | custom | | `custom_script` | Path to custom script file | custom (optional) | | `custom_env` | Space-separated `KEY=VALUE` pairs | custom (optional) | | `perf_config_yaml` | Path to pre-built perf-sanity or benchmark config YAML | perf_sanity (required) | | `perf_test_name` | Server config name for aggregated mode | pytest (optional) | | `model_name` | Short model name (auto-derived if not provided) | Optional | | `required_devices` | Number of GPUs needed (auto-derived if not provided) | Optional | | `node_count` | Number of Slurm nodes (auto-derived if not provided) | Optional | | `partition` | Slurm partition. **For slurm scenarios, this MUST be the value resolved by `trtllm-case-executor` Step 2.5 against the cluster (validated to exist in `slurm_env.partitions[].name`).** Do not detect or default it here. | Required (slurm) | | `account` | Slurm account. **For slurm scenarios, this MUST be the value resolved by `trtllm-case-executor` Step 2.5 against the cluster (validated to exist in `slurm_env.accounts[]`).** Do not detect or default it here. | Required (slurm) | | `slurm_env_file` | Absolute path to the `slurm_env.json` file written by `trtllm-case-executor` Step 2.5 (canonically `<work_dir>/slurm_env.json`). Schema: `{user, default_account, accounts[], default_partition, partitions[{name,state,time_limit,nodes,arch,gres,gpus_per_node,requires_gres}], pmix:{available[],preferred}, errors[]}`. The script builder loads it once in Step 1, uses it to select the Docker image (arch), drive `--gres` / `--mpi=` flags, and embeds the parsed object into `job_spec.json` for downstream consumers. | Required (slurm) | | `time_limit` | Slurm time limit (default: `02:00:00`) | Optional | | `llm_models_root` | Models root **inside the container**. Resolution order: caller input → value already resolved by `trtllm-case-executor` and forwarded in `job_spec.json` → delegate to `internal-env-info` for the `default_llm_models_root` key when that skill is installed → otherwise **ask the user**. The script builder does not invent a default. | Optional | | `checkpoint_path` | Host path to model checkpoint | Optional | | `models_path` | Host path to LLM models directory | Optional | | `device_type` | Required GPU/device hardware type | Optional | | `execution_scenario` | `local_docker`, `local_slurm`, or `remote_slurm` | Always | | `persistent_mode` | `true` or `false`. **Only meaningful for `execution_scenario=local_slurm`** — when true, add persistent-mode fields to `job_spec.json` so `exec-local-slurm` can reuse a persistent allocation. Default: `true` for `local_slurm`. Ignored (no-op) for `local_docker` and `local_direct`. For `remote_slurm`, this field is **not used** — `exec-remote-slurm` owns its own persistence model (tmux-held `salloc` + named container; see remote-executor Recipe 6) and does not consume `persistent_mode` from `job_spec.json`. | Optional | > Per-partition hardware (`arch`, `gres`, `gpus_per_node`, `requires_gres`) lives under `slurm_env.partitions[]`. The chosen `partition` / `account` / `container_image` are passed as top-level inputs. ## Procedure ### Step 1: Resolve Docker Image 1. If the top-level `container_image` input is set (case-executor resolved it as the single source of truth in its Step 3), use it directly — skip the remaining steps. 2. Search `<repo_root>` for an image-tags properties file (e.g., under `jenkins/` or a similar CI config directory). Extract the relevant key based on arch (see step 3). 3. Determine compute node architecture: **Local Docker** (`execution_scenario == local_docker`): - Run `uname -m` via Bash tool on the local machine. **Local Slurm / Remote Slurm**: - Load the SLURM environment from the file at `slurm_env_file` (e.g. `python3 -c 'import json,sys; print(json.load(open(sys.argv[1])))' "$slurm_env_file"` or `jq` — pick whatever's available). The path is canonically `<work_dir>/slurm_env.json`, written by `trtllm-case-executor` Step 2.5. Do **not** invoke any SLURM detection command (`sinfo`, `scontrol`, `srun --mpi=list`) — case-executor already ran the one-pass probe. If `slurm_env_file` is missing or unreadable for a slurm scenario, stop and report that case-executor did not run Step 2.5 before invoking the script builder. The loaded shape is: ```json { "partitions": [ {"name": "<part>", "state": "up|down", "time_limit": "...", "nodes": "...", "arch": "x86_64|aarch64|null", "gres": "gpu:8|(null)", "gpus_per_node": 8|null, "requires_gres": true|false} ], "pmix": {"available": ["pmix","pmix_v3","pmix_v4","pmix_v5"], "preferred": "pmix_v5"}, "errors": ["..."] } ``` - Look up the entry where `name == partition` (the top-level `partition` input, resolved by case-executor Step 2.5); use its `arch` for the image selection below and pass `gres`, `gpus_per_node`, `requires_gres`, and the top-level `pmix.preferred` through to Step 3 / `job_spec.json` (consumed by the executors for the `--mpi=` flag). - If `slurm_env.errors[]` indicates `"SLURM client tools not found"` (the case-executor probe ran on a host without `sinfo`/`scontrol`) or the entry's `arch` is `null`, fall back to the `internal-env-info` skill for arch and leave `pmix.preferred` unset so Step 3 falls back to the bare `pmix` plugin name. If `internal-env-info` is also not installed, skip the lookup silently and use `uname -m` for `local_slurm` or default to `x86_64` for `remote_slurm` when otherwise unknowable. Do not report the missing skill as an error. **Arch → properties key mapping** (apply to the lookup result): | `Arch=` value | Properties key | |---------------|----------------| | `aarch64` | `LLM_SBSA_DOCKER_IMAGE` | | `x86_64` | `LLM_DOCKER_IMAGE` | Use the arch for the chosen `partition` (the target partition for this job) to select the Docker image. If `slurm_env` is missing or its entry's `arch` is null, fall back as described above. **GRES interpretation** — applied in Step 3: - `gres` starts with `gpu:` (e.g., `gpu:4`, `gpu:4(IDX:0-3)`) → partition requires `--gres` - `gres` is `(null)` → partition does not use GRES scheduling; omit `--gres` from the script ### Step 2: Resolve Container Mounts Build comma-separated `<CONTAINER_MOUNTS>`: **Resolve models host path** — evaluate rules in priority order and stop at the first match. All mount entries below use `<llm_models_root>` as the container-side target — the resolved input from Step 2.5 of `trtllm-case-executor` (see the `llm_models_root` row in [Input](#input-from-orchestrator-prompt)). 1. **`checkpoint_path` provided** → use its **parent directory** as the models host path. Construct one mount entry: `<parent_of_checkpoint_path>:<llm_models_root>`. This takes priority over `slurm_env.mounts`. 2. **`models_path` provided** → use that as the models host path. Construct: `<models_path>:<llm_models_root>`. 3. **`slurm_env.mounts` available** → use all entries verbatim (pre-configured for the cluster by `trtllm-case-executor` Step 2.5; includes the models directory and any symlink targets). Do **not** construct a separate models mount entry. 4. **Fallback** → resolve `<default_models_repo_host>` in this order: `env_check.default_models_repo` (forwarded by `trtllm-case-executor`) → `internal-env-info`'s `default_models_repo_host` key if the skill is installed → otherwise **stop and ask the user** for the host path. Do not hard-code a fallback in this skill. Once resolved, construct: `<default_models_repo_host>:<llm_models_root>`. **Home directory**: For local, `$HOME`. For remote, resolve from `slurm_env.remote_cwd`'s parent (the user's home on the cluster) or `ssh <host> 'echo $HOME'`. **Project path**: For local, `<repo_root>`. For remote, `remote_repo_path` (read from `job_spec.json`; constructed by `trtllm-case-executor` as `<default_user_root_dir>/<repo_dir_name>`, where `<default_user_root_dir>` is resolved by case-executor in this order: `internal-env-info` per cluster → fall back to `internal-env-info`'s `default_user_root_dir_example` key when only documentation-level guidance is needed → otherwise case-executor asks the user. The script builder consumes the resolved `remote_repo_path` verbatim and is never hard-coded in this skill). Do **not** use `slurm_env.remote_cwd` — that is the user's root, not the cloned repo path. **Remote work_dir**: For remote scenarios, `remote_work_dir` is read directly from `job_spec.json` — the case-executor is the single source of truth for that path (`<remote_repo_path>/work_dirs/<basename(local work_dir)>`). The script builder consumes it verbatim and never recomputes. **Additional mounts** (custom workflow only): If `custom_script` parent directory is not covered by existing mounts, add it. Assemble: `<home>:<home>,<project>:<project>,<models_mount or slurm_env.mounts entries>[,<extra>:<extra>]` ### Step 3: Resolve Parallelism & Slurm Parameters **By workflow type:** **Pytest**: - Detect sub-type: perf-sanity vs non-perf-sanity (check if `test_cmd` contains `perf-sanity` or a perf sanity test file path, or if `perf_config_yaml` is provided) - Non-perf: inspect test file for `skip_less_device(N)`, `@pytest.mark.parametrize` GPU markers → `required_devices`. If not derivable, use `required_devices` input (default: 1). - Non-perf: pytest is always single-node. If `required_devices` > `slurm_env.partitions[name=<partition>].gpus_per_node` (or `available_gpus` from env-check for non-slurm scenarios), **stop and report an error** — pytest does not support multi-node execution. - Non-perf: **Hard limit — `required_devices` must be `<= 8`.** If `required_devices > 8`, stop and report an error: `required_devices = N exceeds the maximum of 8 GPUs for pytest.` - Non-perf: **Hard limit — `node_count` must be `<= 1`.** Pytest is always single-node; `node_count` is always set to `1`. This is a redundant safety check — if any derivation produces `node_count > 1`, stop and report an error. - Perf-sanity: all Slurm parameters (nodes, ntasks-per-node, gpus_per_node) are derived by `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py` exclusively from the config YAML hardware section — do not use `slurm_env.partitions[].gpus_per_node`, user input, or any default value **Eval**: - Parse from `eval_cmd`: `--tp_size N` (default: 1), `--pp_size N` (default: 1), `--gpus_per_node N` (overrides cluster config) - `world_size` = `tp_size * pp_size`; `required_devices` = `world_size` - **Single-node only**: `trtllm-eval` does not support multi-node execution. If `world_size > gpus_per_node` (i.e., the job would need more than one node), **do not generate a Category 5 script** — fall back to the perf-sanity path (Category 4) by generating a benchmark config with `accuracy.enable_accuracy_test: true` and submitting via `<repo_root_or_remote_repo_path>/jenkins/scripts/perf/local/submit.py`. Same fallback applies to disaggregated accuracy testing. - Auto-derive `model_name` from last component of `--model` arg (e.g., `Llama-3.1-8B-Instruct` → `llama_3_1_8b_instruct`) - Eval task subcommand (e.g., `gsm8k`, `mmlu`) passes through as-is in `eval_cmd` - **Model path resolution**: If `--model` is a HuggingFace model ID (contains `/` but does NOT start with `/`), resolve it to a local path under `<llm_models_root>` (the resolved input — see the `llm_models_root` row in [Input](#input-from-orchestrator-prompt) for the priority chain). Search subdirectories (max depth 3) for a directory matching the model name component (case-insensitive). If found, rewrite `--model` in `eval_cmd` to the local absolute path. If not found, keep the original and warn that HF download will be attempted. **Bench** (trtllm-bench):