Add successive-halving Slurm architecture harness - #841
Conversation
alxmrs
left a comment
alxmrs
left a comment
There was a problem hiding this comment.
Could use some major improvement. That we run the train.py script with another script via a subprocess and don't just import what we need from train.py into a new runner in search is also a suprise (the train class should be importable! We have lots of examples setting these up in tests).
| metric: validation_loss | ||
| mode: min | ||
|
|
||
| runtime: |
There was a problem hiding this comment.
See below -- this should be a template.
| @@ -0,0 +1,17 @@ | |||
| # SPDX-FileCopyrightText: 2026 Samudra Authors | |||
There was a problem hiding this comment.
This test is unnecessary.
There was a problem hiding this comment.
Removed this test together with the top-level CLI integration it tested.
| """Launch and promote Samudra candidates through successive-halving rungs.""" | ||
|
|
||
| from __future__ import annotations | ||
|
|
There was a problem hiding this comment.
I was kind of hoping for a simple way to mix in a different searching algorithm. What if we wanted to use something else instead of successive halving rings? We shouldn't implement the something else now, but having a version of this that was easily refactorable to perform other forms of search later would be a big win.
There was a problem hiding this comment.
Like, I'm not asking for a pluggable system. I'm asking for the ability to easily make a pluggable system when we need it.
There was a problem hiding this comment.
Separated SuccessiveHalving from the executors and added a small build_search() boundary. A future algorithm can become another class and factory branch without committing to a plugin framework now.
| # SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| """Console-script entry point: ``samudra <train|eval|viz> CONFIG [OVERRIDES...]``. | ||
| """Console entry point for Samudra's training, evaluation, and search tasks. |
There was a problem hiding this comment.
Seeing that this script is more of a manager of the standard train.py script that we currently use, I don't think we should add it to our CLI.
There was a problem hiding this comment.
Removed search from the samudra console CLI. It is now an independent manager invoked with python -m samudra.search.
| "progress": self.train_progress.state_dict(), | ||
| "wandb_id": self.wandb_id, | ||
| "wandb_name": self.wandb_name, | ||
| "search": { |
There was a problem hiding this comment.
Driven by a config system, I think that this should only be enabled when we do a search with the tool (maybe, these should be added via fn call to be more succinct).
There was a problem hiding this comment.
Search summaries are now gated by typed experiment.search configuration, and construction moved into the concise _search_summary() helper. Ordinary training does not emit this file.
| "SLURM_ARRAY_TASK_ID" | ||
| ), | ||
| }, | ||
| "provenance": { |
There was a problem hiding this comment.
I like logging this all the time. Can we make it work for normal train? This can be done later.
There was a problem hiding this comment.
Also: don't we currently get most of this info in wandb today??
There was a problem hiding this comment.
Yes, most observability remains in W&B. The local summary is intentionally small and search-only: it supplies the promotion metric and completion/checkpoint/provenance checks without making scheduling depend on W&B availability.
There was a problem hiding this comment.
Agreed this can be useful later. I kept this PR scoped to the search consumer; ordinary training continues to rely on its existing checkpoint and W&B records.
3320a33 to
b0328d1
Compare
|
Implemented the agent-observability extension in The boundary I chose is search-level artifact publication, independent of Local/Slurm execution. An opt-in OSN template uses standard environment credentials; no secrets enter configs or artifacts. Each search now emits local CSV plus queryable Parquet results, consolidated epoch-level scalar histories, failure and scheduler logs, runtime/container/code-layer provenance, timestamps, W&B/public-record links, and a hashed artifact inventory. Finalists retain best-validation checkpoints by default. I did not automatically run PR #834's metrics/viz suite inside each rung. Those diagnostics require completed rollout datasets and would make the cheap inner loop expensive. Instead, A useful next increment after the first real Perceiver search is a finalist-analysis config that selects rollout length and a small diagnostic subset. The first runs should tell us whether variable/depth loss trajectories plus short-rollout bias/variance/spectra are the highest-value set, or whether Perceiver-specific probes such as latent utilization and attention entropy deserve priority. |
|
|
||
| The runner submits the first rung and fixed baselines. On Slurm, dependent | ||
| controller jobs automatically rank completed candidates and submit each later | ||
| rung; users do not manually advance the search. For a local laptop run, include |
There was a problem hiding this comment.
Could local also be used in colab? If so, maybe we could say "For a local run (on a laptop or a Colab notebook)"
There was a problem hiding this comment.
Updated the wording to explicitly include both laptops and Colab notebooks.
alxmrs
left a comment
alxmrs
left a comment
There was a problem hiding this comment.
Another round of review, then we can improve this by trying it out.
| `search/local.yaml` instead of `search/torch.yaml`; candidates then run | ||
| sequentially in the current environment. | ||
|
|
||
| The search directory contains: |
There was a problem hiding this comment.
Is the search directory namespace with a subdir? Can we make the names readable, but unique per experiment (for example, could we use the config name + timestamp, or the config name + hash of the config?).
There was a problem hiding this comment.
Implemented. Every invocation now defaults to <readable-name>--<microsecond-UTC-timestamp>, and that run_id consistently namespaces local directories, training output directories, Slurm jobs, W&B groups, result rows, and published objects. An explicit run_id is also supported for an external allocator, and existing destinations are still never overwritten.
| separate post-search job keeps cheap promotion decisions fast while still | ||
| producing maps, spectra, and time series that explain *why* a finalist worked | ||
| or failed. The publisher is the executor-independent boundary where those | ||
| analysis artifacts will be added once the first Perceiver searches establish |
There was a problem hiding this comment.
Talking about our specific searches is not helpful in these kinds of doc. We can just keep it vague and say it will be added later.
There was a problem hiding this comment.
Agreed. I removed the Perceiver-specific proposed analysis and kept this as a vague executor-independent extension point for future post-search diagnostics.
| return digest.hexdigest() | ||
|
|
||
|
|
||
| def atomic_parquet(frame: pd.DataFrame, path: Path) -> None: |
There was a problem hiding this comment.
Does this work when writing to a cloud optimized store? Or, is this for writing to a local store?
Why are these atomic?
There was a problem hiding this comment.
frame.to_parquet might have all the logic that we need to handle this (e.g. writing to an object store).
There was a problem hiding this comment.
Renamed this to atomic_local_parquet and made its one-line contract explicit. It is intentionally local: a search first assembles a complete local snapshot, atomically replacing catalogs so a controller cannot observe a truncated file, then uploads closed files to local or S3 destinations.
There was a problem hiding this comment.
I’m retaining the staged write here. DataFrame.to_parquet() can target S3, but a direct object write does not make the overall multi-file research record atomic, and it makes retry/debug behavior depend on the remote filesystem. We instead atomically close each local catalog, hash the exact bytes, and publish that immutable file through the destination adapter. This keeps local and S3 records byte-identical and retryable.
| @property | ||
| def root(self) -> str: | ||
| if self.config.public_url is not None: | ||
| return f"{self.config.public_url.rstrip('/')}/{self.search.slug}" |
There was a problem hiding this comment.
Are we enforcing unique slugs per experiment? That is more of the contention we have with our current training experiments. Please look to see what we do there for output_dirs and mimic that approach here.
There was a problem hiding this comment.
Implemented a unique run namespace throughout. The generated run_id is readable, collision-resistant, saved in config/state/results, and used by each candidate output directory. This automates the unique experiment-name discipline used by training while retaining the existing refusal to overwrite.
| config = SearchConfig.from_yaml_and_cli([str(args.config)]) | ||
| search = build_search(config) |
There was a problem hiding this comment.
Nit: I think our codebase convention is to add the build method to the config so this would become:
search = config.build()
There was a problem hiding this comment.
Agreed—updated the codebase to use the config-owned factory consistently: search = config.build(). I removed the standalone build_search helper, migrated the CLI, worker, tests, and docs, and kept the algorithm selection boundary inside SearchConfig.build().
|
@codex may I have your review? |
|
@fomo-bot may I also have your review? |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5e5cf24706
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| state["rungs"][rung]["job_id"] = job_id | ||
| state["status"] = "running" | ||
| self.search.write_state(state) |
There was a problem hiding this comment.
If the Slurm array submission succeeds but the following controller sbatch fails, these lines persist job_id without any controller_job_id. The retry path in SuccessiveHalving.advance() only checks whether the next rung has job_id, so a retry skips resubmission and leaves that rung with no dependent advance job; the search then stops after the array finishes instead of failing or recovering loudly.
AGENTS.md reference: AGENTS.md:L310-L310
Useful? React with 👍 / 👎.
|
Finding
I did not find other blocking issues. Live PR state at Validation I ran: Debug info
|
| Run the entire search with one command: | ||
|
|
||
| ```bash | ||
| python -m samudra.search path/to/search.yaml |
There was a problem hiding this comment.
Question for the reviewer -- should this be promoted to the top level Samudra CLI? My decision so far has been to wait.
alxmrs
left a comment
alxmrs
left a comment
There was a problem hiding this comment.
🤖 Critical review — successive-halving harness
I read this as a fresh adversary: 3.3k lines, a new subpackage, a new persistent-state machine, changes to the shared training loop and the shared Slurm harness. The provenance/artifact design is genuinely good — commit pinning across controller/code-layer/summary, atomic writes everywhere, the SHA-256 catalog, the fail-fast probe, and docs/search.md are all above the bar for research code. The EXECUTORS seam is the right size.
But I do not think this is landable yet. Three findings are blocking, and two of them mean an advertised path has never actually been run:
- Every promoted candidate trains under a corrupted LR schedule.
T_maxrides along in the scheduler state dict, so a rung-1 resume reverts the horizon to rung 0's. Measured on this branch: a promoted run's LR goes0.002 → 0 → 0.001against a configured1e-3. The bundledexample.yamlis exactly this case. Anchors get a clean schedule, so the one comparison the search exists to make is biased. - The local executor cannot run a second candidate.
Trainer.__init__callsWandBLogger.init_instance();WandBLoggeris aMultiton; the second call raisesValueError: already initialized.configs/search/local.yamland the docs both advertise this path. - A full-history JSON rewrite with
fsyncon every optimizer step, in the shared training loop, with unbounded growth — O(n²) synchronous I/O on shared scratch.
The reason (1) and (2) aren't caught is structural, not accidental: the search suite runs in 0.67 s because train_task is monkeypatched away in one test and Trainer is replaced by a FakeTrainer in the other. "371 passed" in the verification section doesn't speak to either. Any fix should come with a test that constructs two real Trainers in one process and one that asserts the scheduler horizon survives a rung boundary.
Beyond those, inline: a disabled overwrite guard, a rung double-submit on controller retry, status: "complete" after a mass rung-0 failure, allow_dirty being unusable on the executor it matters for, a hardcoded 10-minute/2 GB controller job that has to upload multi-GB checkpoints, and the search's own non-finite-metric handling being dead code.
Two cross-cutting things that didn't fit on a line:
- Public logs. The publisher uploads every Slurm stdout/stderr,
experiment.log, and up to 16 KB of raw stderr tail embedded inresults.csv— into a public OSN bucket, from jobs launched with--export=ALLcarryingAWS_SECRET_ACCESS_KEYandWANDB_API_KEY. The default should not be "mirror unfiltered process output to a public URL"; at minimum this needs a documented decision and a scrub pass. - PR description drift. It cites
build_search()as the algorithm boundary (the function isSearchConfig.build()) and says "W&B runs use the search name as their group" (the code usesrun_id, which the docs correctly describe). Worth a pass before merge — reviewers calibrate on this text.
None of this is a design objection to successive halving here; the shape is right. It's that the science-critical path (resume semantics) and the demo path (local) are both unexercised, and the training loop took on a cost that will show up at quarter-degree scale.
| write_search_worker_status( | ||
| self.output_dir, "first_batch", **status_metrics | ||
| ) | ||
| if batch_progress.optimizer_stepped: |
There was a problem hiding this comment.
🤖 Blocking (performance): a full-file JSON rewrite with fsync on every optimizer step, over an unbounded append-only history.
write_search_worker_status reads the existing file, appends to history, and re-serialises the entire history with indent=2, sort_keys=True followed by os.fsync (training_summary.py:25-65). Here it is called from inside the batch loop.
For the bundled example (gradient_accumulation_steps=16, 12-epoch finalist) that is thousands of events per run, and each write rewrites everything written so far — total bytes fsync'd grows as O(n^2). Hundreds of MB to GB of synchronous small writes on shared Lustre/GPFS scratch, per rank-0 worker, per candidate. The file then gets SHA-256'd and re-uploaded to the public bucket on every publish.
AGENTS.md explicitly calls out keeping the core train loop performant, and this lands in the hot path of every search-managed run — including the quarter-degree configs this repo is being scaled for.
Two options that keep the diagnostic value: hold the per-step counters in memory and flush on a wall-clock interval, or drop step events from history entirely and keep this file to lifecycle transitions only — which is exactly what docs/search.md:83 describes it as. search_metrics.parquet and W&B already carry the per-step curve.
| if probe | ||
| else self.output_dir(name, rung) | ||
| ) | ||
| if output.exists() and "RANK" not in os.environ: |
There was a problem hiding this comment.
🤖 This guard is disabled on precisely the path that needs it.
Under the Slurm harness every worker runs via torch.distributed.run, which sets RANK for all ranks including rank 0. So "RANK" not in os.environ is false and Refusing to overwrite never fires for any distributed run — it only protects the local executor, where collisions are least likely.
If the intent was "don't let non-main ranks race on this check", the predicate you want is is_main_process(), already imported on line 35. As written, it silently opts out of the one safety check standing between a resubmitted array and a half-overwritten run directory (see the double-submit path in advance).
| from samudra.search.config import SearchConfig | ||
|
|
||
|
|
||
| class SearchController(Protocol): |
There was a problem hiding this comment.
🤖 Three structural protocols now describe one concrete class: SearchController here, plus ArtifactSearch and ArtifactSearchConfig in artifacts.py. There is exactly one algorithm — SearchConfig.build() is annotated -> "SuccessiveHalving" with an unreachable raise AssertionError behind a Literal-typed discriminator, and AlgorithmConfig = SuccessiveHalvingConfig aliases a single class.
The abstraction isn't paying rent, it's charging it. Undoing it takes four assert isinstance calls: SlurmExecutor.config (slurm.py:19) to recover the config type it was handed, release_probe (successive_halving.py:400) because submit_validated_rung isn't on the Executor ABC, and two on Location in artifacts.py (92, 255). Those asserts vanish under python -O, so the type narrowing they stand in for is unenforced in exactly the deployment where a mistake is expensive.
The EXECUTORS dict is a real seam with two real implementations — keep it. I'd drop the algorithm and artifact protocols until a second implementation exists and let the concrete types flow; submit_validated_rung belonging to the base interface (or the probe living behind submit_rung) would remove one assert on its own.
| worker_batches_seen=worker_status.get("batches_seen"), | ||
| worker_error_type=worker_status.get("error_type"), | ||
| worker_error=worker_status.get("error"), | ||
| worker_status_log=str(status_path.relative_to(self.search_dir.parent)), |
There was a problem hiding this comment.
🤖 Path bases are inconsistent between adjacent columns of the same table.
worker_status_log is relative to search_dir.parent, yielding <run-dir>/search_worker_status.json, while scheduler_stdout_log / scheduler_stderr_log two methods down are relative to search_dir, yielding logs/r0-123_1.out. The published catalog keys the same worker-status file as runs/<run-dir>/search_worker_status.json.
So joining results.parquet.worker_status_log against artifacts.parquet.artifact — which is the stated point of publishing both — misses, while the scheduler-log columns join fine. Pick one base (search-dir-relative, matching the catalog) and derive all four the same way.
| CONFIG is a path to a YAML file or the name of a bundled preset such as | ||
| `samudra_om4/train.yaml`. Any config key can be overridden inline (e.g. | ||
| `--epochs 100`); run `samudra <command> --help` for a command's full options. | ||
| For train, eval, and viz, CONFIG is a YAML path or bundled preset such as |
There was a problem hiding this comment.
🤖 This hunk is leftover churn from a reverted decision. fd4d592 added search to _COMMANDS with a "Plan and run successive-halving architecture searches" line; b0328d1 took it back out. What survives in the diff is a reworded help paragraph and a dropped else: — leaving main() with three elifs and no fallback branch.
Meanwhile the commit is still titled "Expose architecture search through Samudra CLI" and the PR body opens with the user interface section, but samudra search does not exist — python -m samudra.search is the only entry point, which a pip install samudra user has no way to discover from samudra --help.
Either finish it (search in _COMMANDS plus a dispatch branch and a help line) or revert cli.py entirely, so the diff stops implying a feature that isn't there.
|
No blocking findings in the current code review. The prior slug-collision issue is fixed: candidate names are now validated after Validation run:
Live PR state at Debug info
|
|
Rechecked the PR at Debug info
|
cb9cf74 to
47201b7
Compare
adam-subel
left a comment
adam-subel
left a comment
There was a problem hiding this comment.
I think overall things seem to make sense, I went through it with the help of claude as it is quite dense and left some notes. Happy to talk this through on a call, it might be helpful for me to better give feedback on the larger structuring
|
|
||
| @staticmethod | ||
| def _kind(relative: str) -> str: | ||
| if relative.endswith("ckpt.pt"): |
There was a problem hiding this comment.
Did you intend for this to capture any checkpoints of the form ckpt_{epochs}.pt?
| self.write_state(state) | ||
| # Publish stable, queryable schemas immediately. Without these files, | ||
| # the public results URL returns 404 until the first rung completes. | ||
| self._write_results(state) |
There was a problem hiding this comment.
Noted by claude:
_write_results:519-544 writes results.parquet, results.csv, and then calls write_search_report(self, state). advance calls it between write_state and publish:
So any exception from report rendering — a formatting bug on an unusual value, a None where a number was expected — skips publication and skips the next rung's submission. A cosmetic failure becomes a silently stalled search.
There was a problem hiding this comment.
I think this is addressed in the new lines on 217 below.
f5bd9a3 to
cd532d9
Compare
3 participants
Goal
Use fragmented compute capacity to compare model architectures quickly, then allocate larger epoch budgets only to promising candidates.
This implements one successive-halving bracket, the resource-allocation primitive underlying Hyperband. It deliberately leaves multiple-bracket Hyperband and deferred rollout metrics for later work.
User interface
Run a complete search with:
python -m samudra.search path/to/search.yamlUsers do not manually advance rungs. The Local executor runs candidates sequentially in isolated multiton/logging scopes; the Slurm executor submits arrays and dependent promotion jobs automatically. Search configs use Samudra's Pydantic/YAML system, including
!include, packaged presets, CLI overrides, and generated schemas.Design and correctness
SuccessiveHalvingowns ranking and promotion;SearchConfig.build()constructs it.T_maxcorruption.completeonly if every scheduled result is eligible; recoverable worker failures producepartial, and a rung with no eligible candidate records a durable terminal failure.Reproducibility, observability, and W&B
run_idas the group; runs receive stable search-name, run-ID, and candidate tags.launched,initialized,first_batch, firstoptimizer_step, terminal stage), rather than rewriting unbounded per-step history.results.csv/results.parquet,epochs.parquet,artifacts.parquet, resolved configs, reports, state, provenance, W&B identity, and configurable checkpoints form the agent-observable record.artifacts.logs: all; structured error strings are scrubbed.See
docs/search.mdfor setup, schemas, safety policy, and DuckDB examples.Verification