[podkidyshev]

Summary

DSE trials can run in parallel. max_parallelism (default 1, set under agent_config) is the maximum number of trials CloudAI keeps outstanding at once.

• Trials run in waves: up to max_parallelism trials start, all of them finish, then the next wave starts. Proposals rejected by the constraint check, already in the trajectory, or identical to a trial of the same wave take no slot.
• Agents opt in through BaseAgent.select_batch, which returns every proposal the agent can make before it needs feedback; it may be an endless generator. The default wraps select_action, so existing agents stay serial. grid_search proposes the whole grid.
• CloudAIGymEnv.evaluate_batch owns the sweep bookkeeping (step allowance, trial numbers, constraint retries, skipped steps) and returns one result per proposal, in proposal order. step() is a one-proposal batch. The trajectory and the set of executed trials do not depend on max_parallelism.
• A failed trial does not stop its wave. Cancelling (Ctrl+C) now stops the whole DSE: the running wave is terminated and not recorded.
• BaseRunner separates job polling from scenario policy. poll_jobs yields finished jobs with their status; monitor_jobs applies iterations, dependencies and fail-fast on top; the DSE uses submit_test and poll_jobs directly. A job now costs one status query per pass instead of three, plus one more only while its run is published as pending or a start_post_init dependent waits for it.
• pin_nodes cannot be combined with max_parallelism > 1, and single-sbatch DSE rejects the setting.

Test Plan

CI passed.

The DSE cases (tests/dse) run a probe workload on an in-memory system where the duration of every trial, in polling passes, and its outcome are scripted, so the composition of each wave is asserted exactly.

• Waves: a grid runs in waves of max_parallelism; the next wave starts only after the slowest trial of the current one; the trajectory keeps proposal order when trials finish in reverse; every trial gets its own output directory; the trajectory is the same for max_parallelism 1, 2, 3, 5 and 16; env_params are sampled per trial whatever the parallelism.
• Constraint failures: a rejected proposal takes no slot and no trial number; a stream of rejected proposals ends through max_consecutive_rejected_proposals, and the skipped step consumes a trial number.
• Execution failures: a failed job, a job that reports no metric and a trial that cannot be prepared each fail only their own trial, the rest of the wave and the sweep go on; an error while waiting for jobs fails the trials still running and the sweep goes on.
• Cache and duplicates: earlier waves feed the cache of later ones; identical proposals in one wave run once.
• Agent contract: proposals are pulled lazily; the step allowance bounds an endless stream; an agent without select_batch stays serial; an agent gets the feedback of a whole batch before it proposes the next one; step() is a one-proposal batch.
• Cancel and submission errors: cancelling stops the sweep, terminates the running trials and drops the interrupted wave, including a trial submitted while the cancel arrives; an interrupt inside a wave terminates its trials; a submission error terminates the wave, aborts the sweep and the CLI exits with 1.
• Limits: pin_nodes is rejected with max_parallelism > 1 and works with 1; single-sbatch DSE rejects the setting.
• Runner (tests/test_base_runner.py): poll_jobs asks for completion once per job and pass, only reports a failed job, and lets the caller handle a finished job before the next one is checked; monitor_jobs stops the scenario on a failed job, resubmits per iteration, submits a start_post_init dependent once per iteration and does not iterate an end_post_comp dependent that finished with its source; nothing starts after shutdown, and shutdown terminates every job even if one cannot be terminated.
• Real processes on the standalone system: the processes of a wave overlap and waves do not; cancelling kills the running processes.

Manual: real cloudai run on the standalone system with the Aiconfigurator workload, grid of 9 trials: 36.0 s with max_parallelism = 1, 13.8 s with max_parallelism = 3, identical trajectory.csv. Nothing was run on Slurm; node allocation for several trials submitted at once is untested there.

Additional Notes

Compatibility changes:

• Cancelling a DSE stops it and returns a non-zero code. Before, only the current trial was killed and the sweep went on.
• A job submission error raises JobSubmissionError (now exported from cloudai.core) instead of calling exit(1): outstanding jobs are terminated, the CLI still exits with 1, cloudai.api.run_experiment raises it.
• No test is started once the runner is shutting down.
• In a DSE, an error while waiting for jobs fails the trials that were still running and terminates their jobs; before, their jobs were left running.
• BaseRunner.check_start_post_init_dependencies is removed; monitor_jobs covers it. handle_job_completion no longer removes the job from jobs and testrun_to_job_map; poll_jobs does.
• BaseGym.progress and CloudAIGymEnv.progress are removed. BaseAgent.run calls env.evaluate_batch; BaseGym provides a serial default built on step, so environments that only implement step keep working, but a wrapper that only overrides step is bypassed by the default loop.
• Feedback passed to update_policy keeps the existing keys and adds status, step, output_path and env_params.
• For a cached trial, env.test_run.output_path is the directory of the trial whose result is reused.
• Standalone jobs close their process pipes on completion.

🤖 Generated with Claude Code

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Pre-merge checks: 2

Failed checks: 2 inconclusive

Check name	Status	Explanation	Resolution
Title check		This assessment has not completed. The check will be retried on the next pre-merge check run.
Description check		This assessment has not completed. The check will be retried on the next pre-merge check run.

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts