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

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions