Skip to content

DSE: single-element list dims ([1]) aren't scalarized before command-gen (no central invariant) #946

Description

@rutayan-nv
Prompt for AI agent — drop-in spec to implement the fix (canonicalization invariant + central boundary guard)

Where you are

Work in the OSS cloudai package: cloudai/src/cloudai/... (the workspace also holds cloudaix* worktrees — ignore those; cloudaix consumes cloudai as a pinned dep and inherits this fix). Run all quality gates from the cloudai repo. Conventions from the /cloudaix skill apply: Python 3.10, line length 120, Pydantic v2, ruff + pyright + pytest must pass, SPDX/copyright header on new files. Note the skill's rule: Union[T, List[T]] on a cmd_args field marks it DSE-able (scalar = fixed, list = swept); DSE auto-enables when any value is a list.

Problem

CloudAI overloads "a cmd_args leaf is a list" to mean "this is a DSE sweep axis" (TestRun.param_space / TestDefinition.is_dse_job). A single-element list [1] is semantically fixed (one candidate) but is structurally a list, so it survives as a degenerate 1-point "sweep" all the way to command generation. There it either:

  • leaks onto the CLI verbatim (e.g. aiconfig emits --p-pp [1]), or
  • is silently patched per-workload by ad-hoc unwrap/guard code.

There is no central invariant that command-gen only ever sees a scalar for a sweepable field, so every workload re-invents the defense (inconsistently), and two forgot.

Invariant to establish

After parsing + DSE resolution: fixed = scalar, sweep = list of >= 2 candidates. No single-element list survives on a sweepable field, and no unresolved sweep ever reaches command generation.

Decision (already made — do NOT redesign)

Implement TWO layers and delete the [1] band-aids. Do NOT attempt the deeper "make sweep an explicit declaration" redesign — that's a separate, breaking change and is OUT OF SCOPE. No TOML changes are required by this fix.


Layer 1 — parse-time canonicalization (TestDefinition)

File: cloudai/src/cloudai/models/workload.py. Add a @model_validator(mode="after") on TestDefinition that recursively walks cmd_args (nested BaseModels AND the extra="allow" bag) and extra_env_vars, collapsing [x] -> x. It must be TYPE-AWARE and SCALAR-ELEMENT-ONLY — collapse a length-1 list only when ALL hold:

  1. the single element is a scalar primitive (int | float | str | bool);
  2. the key is NOT env-sampled (is_env_sampled) and NOT dse-excluded (is_dse_excluded_arg);
  3. for a declared field, its annotation admits a scalar branch (the T | list[T] signature). For extra="allow" bag fields (typed Any), collapse freely.

Hard constraints (these are the bugs a naive version introduces — verified by review)

  • aiperf_phases: list[AIPerfPhase] in cloudai/src/cloudai/workloads/ai_dynamo/ai_dynamo.py (~line 403) is a pure list of models with NO scalar branch. Collapsing [{...}] would fail Pydantic re-validation at apply_params_set reconstruct and is semantically wrong. The scalar-element + type-aware rules above must exclude it — add a test that proves it.
  • Skipping env-sampled keys is MANDATORY, not optional: validate_env_params (models/workload.py ~line 213) requires every env-sampled cmd_args field to be a list (see the not isinstance(value, list) reject ~line 252). Collapsing one would break parse. Validators are all mode="after" and read self.env_params, so order is safe as long as canonicalization skips env-sampled keys.
  • Fold in the fix(configurator): make env_params first-class to fix the trajectory cache key #901 follow-up: a [1] env_param is a no-op (one candidate, nothing to sample) and must be rejected at parse, mirroring the existing scalar-annotation rejection. Tighten validate_env_params so an env-sampled candidate list with < 2 entries raises.
  • apply_params_set (_core/test_scenario.py ~line 190) reconstructs the model; confirm the new validator is a no-op there (post-overlay swept fields are already scalars; a multi-element [1,2,4] is never present at that point; env_params is popped before reconstruct). Do not break that path.

Layer 1b — num_nodes

num_nodes lives on TestRun (dataclass, _core/test_scenario.py ~line 83), NOT TestDefinition, and is built via multiple paths (models/scenario.py TestRunModel.num_nodes, test_scenario_parser.py ~line 201, and direct construction in code/tests). param_space/is_dse_job special-case isinstance(num_nodes, list). Normalize num_nodes=[1] -> 1 in TestRun.__post_init__ (covers the dataclass path uniformly), and also add a field_validator on TestRunModel.num_nodes for the parse model. No path requires a length-1 list (TestRun.nnodes raises on a list; single-sbatch unroll always sets a scalar NUM_NODES).

Layer 2 — central boundary guard (the point of the whole change)

Add a single guard so an unresolved sweep can NEVER silently reach command-gen. Recommended location: CommandGenStrategy.__init__ (cloudai/src/cloudai/_core/command_gen_strategy.py ~line 28) — assert that self.test_run.param_space is empty at command-gen time, raising a clear error naming the offending keys otherwise. Rationale: by command-gen time every sweep axis (real, or a value-list mis-detected as one) has been resolved to a scalar by the DSE/unroll machinery, so a non-empty param_space means an unresolved sweep — a real bug to surface loudly, in ONE place. This reuses the existing dse_excluded / env_sampled filtering, so it needs no TOML changes.

  • Verify before shipping: confirm no legitimate path constructs a command-gen strategy on an unresolved DSE test_run (especially dry-run). If one exists, resolve/representative-pick first, or scope the guard to the run path. Do not let the guard false-positive on dry-run.

Band-aids to DELETE (sweep-resolution workarounds, now subsumed)

  • NCCL _nccl_cmd_scalar + its call sites: cloudai/src/cloudai/workloads/nccl_test/slurm_command_gen_strategy.py (~lines 37-40, 121-146).
  • megatron_bridge generic flag emitter length-1 branch: cloudai/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py (~line 311), and the explicit 1 or [1] vp special-case (~lines 391-394).
  • nemo_run pipeline_model_parallel_size list->[0]: cloudai/src/cloudai/workloads/nemo_run/slurm_command_gen_strategy.py (~lines 44-45).
  • nixl_ep num_processes_per_node int-guard: cloudai/src/cloudai/workloads/nixl_ep/slurm_command_gen_strategy.py (~lines 64-66).
  • ai_dynamo num_nodes int asserts: cloudai/src/cloudai/workloads/ai_dynamo/slurm_command_gen_strategy.py (~lines 618-619).
  • Inverse band-aids (no defense -> active [1] leak; canonicalization fixes them, just verify output): aiconfig standalone (workloads/aiconfig/standalone_command_gen_strategy.py, the str(d.p_pp) block ~lines 60-124) and dynamo_mocker standalone (workloads/dynamo_mocker/standalone_command_gen_strategy.py ~line 58).
  • Before deleting NCCL/megatron/nemo_run unwraps, confirm a multi-element list can no longer reach command-gen (the Layer-2 guard guarantees this) — those helpers also coerced multi-element lists by silently taking [0]/comma-joining, which was itself a latent bug.

KEEP (genuine value-list handling — NOT sweep-resolution; do not touch)

  • ai_dynamo AIPerf arg list/dict reject (ai_dynamo/slurm_command_gen_strategy.py ~lines 140-144) — multi-value AIPerf needs CLI-string syntax; real semantic.
  • nemo_launcher Hydra list->comma (nemo_launcher/slurm_command_gen_strategy.py ~lines 204-205).
  • megatron _normalize_recompute_modules, --dataset_paths, --profiling_ranks.
  • fio _format_option, gpu_ids / CUDA_VISIBLE_DEVICES joining.

Tests

  • Fix cloudai/tests/workloads/nixl_ep/test_command_gen_strategy_slurm.py test_gen_srun_command_rejects_process_list (~lines 535-556): it uses num_processes_per_node=[4], which now canonicalizes to 4 and won't raise. Change to a genuine multi-element [4, 8] and assert the Layer-2 guard raises.
  • Add parse-time tests: gpu_ids=["1"] -> "1"; nested disagg.p_pp=[1] -> 1 and absent from param_space / is_dse_job False; [1, 2] untouched; num_nodes=[1] -> 1; env_params with a 1-element candidate list rejected; aiperf_phases=[{...}] NOT collapsed.
  • Audit cloudai/tests/.../test_cloudaigym.py (~lines 237, 302 use a [3] param) and test_handlers.py (~lines 444-486 build the dataclass directly, bypassing parse) for the is_dse_job flip side-effect: an all-[1] config now becomes non-DSE, and if it also declares env_params, validate_dse_env_params (configurator/env_params.py ~lines 188-198) will now raise. Make sure no existing test/config regresses on this.

Validation / quality gates (must all pass)

uv sync --extra dev
uv run ruff check
uv run ruff format --check
uv run pyright
uv run pytest -vv

Out of scope

The deeper redesign that separates "sweep axis" from "genuine value" (explicit sweep declaration, removing the list == sweep overload that mis-detects gpu_ids / dataset_paths as sweeps) is intentionally NOT part of this change.


Summary

A DSE dim authored as a single-element list (e.g. p_pp = [1]) can reach command generation still as a list, producing an invalid CLI arg (--p-pp [1]). Nothing in core guarantees cmd_args is fully resolved to scalars before a workload stringifies it. #938 patches this inside the aiconfigurator command-gen, but the gap is generic.

Mechanism

  • Sweepable dims are typed Union[int, List[int]] (list = explore, scalar = fixed).
  • param_space / is_dse_job classify by isinstance(value, list)length-agnostic — so a single-element [1] is treated as an action dim even though it carries no real choice.
  • apply_params_set overlays only the keys present in the agent's action; it does not iterate param_space. So a degenerate dim is collapsed to a scalar only if the active agent happens to emit it:
    • grid_search: emits every dim (cartesian product) → always collapses [1]1.
    • GymnasiumAdapter (RL): explicitly splits tunable (len > 1) vs fixed (len == 1) and injects the fixed scalars on every step → also collapses.
    • Any other agent/path is under no obligation to, so cmd_args.p_pp stays [1].
  • Command-gen then does str(value)--p-pp [1], which simple_predictor.py (int("[1]")) rejects.

Why it surfaced now (not RL-specific)

  • Long latent: grid was the only DSE agent and it always collapses, and fixed dims used to be authored as scalars (p_pp = 1).
  • The auto-generated agent configs (bo, ga, nvopt, rl) now emit fixed dims uniformly as [1], and non-grid agents don't guarantee overlay — so the latent gap became reachable. It is config/agent-convention driven, not specific to RL.

Impact

  • Any workload whose cmd_args field is typed Union[scalar, list] and stringified into a CLI flag is exposed; aiconfigurator is just the first.
  • Failure is a hard crash at predictor parse time, not a silent wrong result.

Root cause

No single normalization point owns "resolve every dim to a concrete scalar before command-gen." The responsibility is implicitly spread across each agent, and apply_params_set is action-driven rather than param_space-driven. A single-element list is semantically identical to its scalar to a human, but not at the CLI boundary (str(1) vs str([1])).

Suggested fix

Lift the fixed-dim collapse into core so every agent/path benefits, instead of per-workload band-aids:

  • Option A: in param_space, treat len == 1 lists as fixed (exclude from the search space) and materialize them onto cmd_args; or
  • Option B: have apply_params_set (or a dedicated post-resolve step) always scalarize degenerate dims regardless of which keys the agent's action carries.
  • GymnasiumAdapter._fixed_params already implements exactly this tunable/fixed split — it is the reference behavior to hoist into core.
  • Keep fix(aiconfig): unwrap single-value list dims in standalone command-gen #938's command-gen guard as defense-in-depth (fail loud on a genuine multi-element leak, which means an unresolved sweep escaped the agent).

Severity

Medium — hard crash, workload-reachable, currently mitigated only for aiconfigurator (#938).

Code: src/cloudai/_core/test_scenario.py (param_space, apply_params_set); src/cloudai/configurator/gymnasium_adapter.py (reference impl); src/cloudai/workloads/aiconfig/standalone_command_gen_strategy.py (current local fix, #938).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions