Architecture search¶

Training every model idea to convergence wastes compute: weak candidates often become distinguishable after only a small amount of training. Samudra's search runner uses successive halving to train all candidates with a small initial budget, retain the best fraction, and give progressively larger budgets only to the survivors. This is useful for comparing model architectures as well as ordinary hyperparameters such as learning rate and batch size.

Successive halving is the resource-allocation primitive used by Hyperband. Samudra currently runs one successive-halving bracket. It does not yet run Hyperband's collection of brackets with different exploration-versus-training tradeoffs.

Run a search¶

Copy the bundled example and Torch executor configuration from src/samudra/configs/search/, then edit the output paths, Slurm account, and candidate list. Build the configured code layer from that same committed revision. The runner verifies the layer's manifest before submitting work, and completed training results must report the same commit before promotion. All candidate models must therefore exist in one experiment branch before the search is submitted.

Run the entire search with one command:

python -m samudra.search path/to/search.yaml

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 or Colab notebook run, include search/local.yaml instead of search/torch.yaml; candidates then run sequentially in the current environment. Each local task gets a fresh multiton scope and releases cached CUDA memory before the next candidate, so W&B and normalization state cannot leak between models.

Clusters that expose the container runtime through environment modules should set executor.apptainer_module to the exact loadable module name (for example, singularity-ce/4.3.3 on Torch). The resolved value is exported to every worker and retained in the search configuration.

Every invocation gets a readable instance identifier such as perceiver-2deg--20260813T192612.123456Z. The stable name describes the experiment design; this generated run_id names its filesystem, object-store, Slurm, and W&B resources so repeated trials cannot collide. Set run_id explicitly only when an external system needs to allocate the identity. The runner still refuses to overwrite an existing local or published instance.

The uniquely named search directory contains:

• config.yaml: the fully resolved, validated search configuration;
• candidates/: one fully resolved training configuration per candidate;
• state.json: internal resumable scheduler state;
• results.csv and results.parquet: one analysis-ready row per candidate and rung, including all requested metrics, timing, checkpoint lineage, W&B identity, and job ID;
• epochs.parquet: full epoch-level training, validation, inference, timing, throughput, and variable/depth/channel metrics from every completed run;
• artifacts.parquet: hashes, sizes, media types, and public/queryable locations for every published artifact;
• analysis/report.md: a continuously updated leaderboard, rung history, outcome (once available), and explicit candidate failures;
• logs/: scheduler output when using Slurm.

Fail-fast worker probes¶

For costly Slurm searches, set executor.rung0_probe: true. Before releasing the first candidate array, the executor runs the first candidate through the real data loader, model forward/backward path, and one optimizer update with W&B disabled. The bulk array is submitted only when the probe records a finite training batch and at least one optimizer step. A missing or failed probe marks the search terminal, publishes its diagnostics, and consumes no candidate or fixed-anchor array allocation. Fixed anchors are submitted only after the same probe succeeds.

Every search-managed training process atomically maintains search_worker_status.json with its lifecycle history: launched, initialized, first_batch, first optimizer_step, and completed or failed. Events include batch/optimizer counts, loss and loader timings when available, Slurm identity, and an error type/message on failure. These files are copied into the public research record, and their latest stage is included in failed result rows and reports. A scheduler process merely starting is therefore not treated as evidence that training occurred.

results.csv is deliberately denormalized and readable directly with pandas:

import pandas as pd

results = pd.read_csv(
    "/scratch/USER/searches/my-search--20260813T192612.123456Z/results.csv"
)
print(results.sort_values(["rung", "validation_loss"]))

Publish an inspectable research record¶

Artifact publication is opt-in and independent of the executor. Local and Slurm searches first create the same record on their working filesystem, then an optional publisher mirrors it to another local directory or any S3-compatible object store. This keeps compute scheduling separate from where research results are retained.

For the public m2lines OSN pod, uncomment the following line in a copied search config:

artifacts: !include osn-artifacts.yaml

The packaged template publishes under s3://m2lines-pubs/Samudra/experiments/searches/<run-id>/. It does not contain credentials. On every machine that may run the search controller, provide write credentials through the normal environment:

export AWS_ACCESS_KEY_ID=...
export AWS_SECRET_ACCESS_KEY=...
# Only when the credentials are temporary:
export AWS_SESSION_TOKEN=...

Slurm submissions inherit these variables; secrets are neither embedded in the search config nor written to its artifacts. The S3 destination carries the OSN endpoint explicitly. Publication is required once configured: an upload failure stops promotion loudly instead of silently creating an incomplete public record.

The published record includes resolved search and training configs, Git provenance, scheduler state, per-rung outcomes (including failures), a Markdown summary report, full epoch histories, structured worker status, W&B identities, and an SHA-256 inventory. results.parquet is created and published with a stable empty schema before jobs are submitted, so its public URL can be queried immediately rather than returning 404 while rung zero is running. Search-level or per-run analysis hooks have a simple artifact contract: write their tables, reports, or figures beneath an analysis/ directory and the publisher includes them automatically with their hashes and locations. This keeps future diagnostics independent of Local versus Slurm execution. Checkpoint publication is configurable:

• none publishes no model weights;
• final (the default) publishes the best-validation checkpoint from each successful final-rung run and fixed baseline (falling back to its latest checkpoint when necessary);
• all publishes every saved checkpoint from every completed rung.

Raw scheduler stdout/stderr and experiment.log/error.log are excluded from artifact publication by default because jobs inherit credentials and arbitrary process output is not safe to mirror into a public bucket. Set artifacts.logs: all only for a destination with an appropriate access policy and after auditing the workload's logging. Structured worker errors redact credential-like environment values before publication.

The local filesystem always retains the checkpoints required for promotion. final is normally enough to reproduce a promising model and run deferred rollout diagnostics without paying object-storage costs for every eliminated candidate.

The Parquet tables can be queried in place by agents or collaborators. For a public HTTP endpoint, DuckDB needs no local download:

SELECT candidate, rung, epochs, validation_loss, error
FROM read_parquet(
  'https://nyu1.osn.mghpcc.org/m2lines-pubs/Samudra/experiments/searches/my-search--20260813T192612.123456Z/results.parquet'
)
ORDER BY rung, validation_loss;

epochs.parquet supports deeper questions such as which model learned fastest, which variable or depth stalled, whether throughput or optimizer-step counts differed, and where divergence began. artifacts.parquet lets an agent locate and verify the exact config, logs, or checkpoint behind any row.

These tables are intentionally single Parquet files, not shards. Their row counts grow with candidates times rungs or candidates times epochs, rather than with the ocean dataset, and are expected to remain small compared with one checkpoint. A future search large enough to need sharding can publish a partitioned Parquet dataset with the same columns and query interface; no current workload benefits from that complexity.

Deferred model diagnostics¶

Search currently records and publishes the inputs needed for later analysis rather than coupling training to a particular scientific metric suite. Future post-search jobs can load published checkpoints, generate matched validation rollouts, and write tables or figures beneath analysis/. Keeping that work separate preserves fast promotion decisions while leaving room to add the most informative diagnostics after experience with real searches.

Configure candidates and metrics¶

Search configuration uses the same Pydantic/YAML system as training and eval, including !include, packaged presets, command-line overrides, generated JSON schemas, and @config.yaml values.

name: perceiver-2deg

algorithm:
  type: successive_halving
  rungs: [1, 3, 6, 12]  # cumulative total epochs
  promotion_fraction: 0.5
  minimum_promoted: 1

objective: {metric: validation_loss, mode: min}
metrics: [validation_loss, train_loss, best_validation_loss]

executor: !include torch.yaml

candidates:
  - name: control
    fixed: true
    config: experiments/control/train.yaml
  - name: sdpa-perceiver
    config: experiments/sdpa/train.yaml
    args: [--batch_size=2, --learning_rate=0.0006]

The objective is the single metric used for promotion. metrics controls the additional columns collected in results.csv. Any finite scalar emitted by training, validation, or inference can be named here, including namespaced diagnostics such as val/seam/window_jump_ratio/zos. A missing or non-finite metric, an incomplete epoch budget, or a missing checkpoint makes that result ineligible.

minimum_promoted is a floor, so promotion can intentionally stop reducing the candidate pool once it reaches that size. It cannot exceed the number of non-fixed candidates configured for the search; if worker failures leave fewer eligible candidates, every eligible candidate advances.

A search is complete only when every scheduled candidate result across all rungs and fixed anchors is eligible. If the search can continue after one or more worker failures, its terminal status is partial; the report identifies the best completed finalist without presenting it as an uncontested winner. A rung with no eligible candidates is failed.

Candidates marked fixed are reference baselines. They run independently at the largest budget and do not consume promotion slots. Rung budgets are total epochs, not extra epochs: a survivor resumes its complete checkpoint from epoch 1 to epoch 3, then from epoch 3 to epoch 6. When a candidate's scheduler omits target_epochs, the search snapshots the largest rung as the common scheduler horizon for every rung and fixed anchor. This prevents a resumed cosine scheduler from restoring the first rung's short T_max. An explicitly configured target_epochs is preserved.

Controller jobs inherit the verified search commit explicitly rather than requiring an editable Git checkout in the controller environment. Their partition, CPU count, memory, and walltime are configurable with controller_partition, controller_cpus_per_task, controller_memory, and controller_time. Published-object hashes are recorded locally after each successful upload, so retries skip unchanged multi-gigabyte checkpoints.

The resumable state.json contract is versioned and validated on every read and write. A malformed or stale state therefore fails at the controller boundary with a schema error instead of producing a later key error.

W&B¶

Each candidate uses the unique search run ID as its W&B group and receives the tags search, the stable search name, the run ID, and the candidate name. This makes repeated trials separately filterable while preserving a stable tag for cross-run comparisons. Search identity, rung, objective, epoch budget, executor, job ID, parent checkpoint, and the public artifact root are stored under config.experiment.search. Promoted rungs resume the same W&B run from the checkpoint, preserving one continuous learning curve per candidate.

W&B is used for curves and interactive comparison. Promotion reads the local training summary so temporary W&B or network failures do not control scheduling.

Python API¶

SearchConfig.from_yaml_and_cli() loads and validates configuration, config.build() selects the configured search algorithm, and search.start() submits or runs it. Internal worker entry points are implementation details used by executors.

Compute-specific logic lives in samudra.search.executors. Built-in executor classes are selected by a small dictionary in successive_halving.py. Adding a future Empire AI executor requires implementing the same submit_anchors and submit_rung interface and adding one dictionary entry; no plugin registration system is imposed today. The search algorithm has a separate typed-config and factory boundary. A future Hyperband implementation can coordinate several successive-halving brackets through the existing executor submissions and use the algorithm-neutral artifact publisher. A scheduler that is not rung-based may require broadening the executor task interface; the current API does not pretend every possible explore/exploit strategy already fits it.

77
78

def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)

122
123
124
125
126

def build(self) -> "SuccessiveHalving":
    """Build the configured search algorithm."""
    from samudra.search.successive_halving import SuccessiveHalving

    return SuccessiveHalving(self)

106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

def __init__(self, config: SearchConfig) -> None:
    self.config = config
    self.slug = resource_slug(config.name)
    if config.run_id is None:
        config.run_id = _new_run_id(config.name)
    self.run_id = resource_slug(config.run_id)
    self.search_dir = config.executor.output_dir / self.run_id
    self.config_path = self.search_dir / "config.yaml"
    self.state_path = self.search_dir / "state.json"
    self.results_path = self.search_dir / "results.csv"
    self.results_parquet_path = self.search_dir / "results.parquet"
    self.checkpoint = CHECKPOINT
    self.rungs = config.algorithm.rungs
    self.executor = EXECUTORS[config.executor.type](self)
    self.publisher = (
        ArtifactPublisher(self, config.artifacts)
        if config.artifacts is not None
        else None
    )

373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430

def release_probe(self, rung: int) -> None:
    """Release a Slurm rung only after its probe proves an optimizer update."""
    state = self.read_state()
    controller_commit = os.environ.get("SAMUDRA_CODE_COMMIT")
    expected_commit = state.get("provenance", {}).get("commit")
    if controller_commit != expected_commit:
        raise ValueError(
            f"Search was started at {expected_commit}; probe controller "
            f"reports {controller_commit!r}"
        )
    current = state["rungs"][rung]
    probe = current.get("probe")
    if not isinstance(probe, dict):
        raise ValueError(f"Rung {rung} has no configured probe")
    name = probe["candidate"]
    status_path = (
        self.search_dir / "probe" / resource_slug(name) / SEARCH_WORKER_STATUS_NAME
    )
    try:
        status = json.loads(status_path.read_text(encoding="utf-8"))
        if status.get("stage") != "completed":
            raise ValueError(
                f"worker stopped at stage {status.get('stage')!r}: "
                f"{status.get('error', 'no worker error recorded')}"
            )
        if int(status.get("optimizer_steps", 0)) < 1:
            raise ValueError("worker completed without an optimizer update")
    except (OSError, TypeError, ValueError, json.JSONDecodeError) as error:
        state["status"] = "failed"
        state["failure"] = {
            "stage": "rung_probe",
            "type": type(error).__name__,
            "message": str(error),
            "candidate": name,
            "failed_at": datetime.datetime.now(datetime.UTC).isoformat(),
        }
        probe["status"] = "failed"
        self.write_state(state)
        self._write_results(state)
        self.publish(state)
        try:
            self._write_report(state)
            self.publish(state)
        except Exception as report_error:
            error.add_note(f"Failed to render or publish report: {report_error!r}")
        raise ValueError(f"Rung {rung} probe failed: {error}") from error
    probe.update(
        status="complete",
        optimizer_steps=status["optimizer_steps"],
        batches_seen=status.get("batches_seen"),
        validated_at=datetime.datetime.now(datetime.UTC).isoformat(),
    )
    self.write_state(state)
    if not isinstance(self.executor, SlurmExecutor):
        raise TypeError("Rung probes are only supported by the Slurm executor")
    self.executor.submit_anchors(state)
    self.executor.submit_validated_rung(state, rung)
    self.publish(self.read_state())

77
78

def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)

122
123
124
125
126

def build(self) -> "SuccessiveHalving":
    """Build the configured search algorithm."""
    from samudra.search.successive_halving import SuccessiveHalving

    return SuccessiveHalving(self)

20
21
22
23
24
25

def resource_slug(value: str) -> str:
    """Convert a display name to the identifier used for search resources."""
    slug = re.sub(r"[^a-zA-Z0-9_.-]+", "-", value).strip("-.")
    if not slug:
        raise ValueError(f"Name has no usable characters: {value!r}")
    return slug

106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

def __init__(self, config: SearchConfig) -> None:
    self.config = config
    self.slug = resource_slug(config.name)
    if config.run_id is None:
        config.run_id = _new_run_id(config.name)
    self.run_id = resource_slug(config.run_id)
    self.search_dir = config.executor.output_dir / self.run_id
    self.config_path = self.search_dir / "config.yaml"
    self.state_path = self.search_dir / "state.json"
    self.results_path = self.search_dir / "results.csv"
    self.results_parquet_path = self.search_dir / "results.parquet"
    self.checkpoint = CHECKPOINT
    self.rungs = config.algorithm.rungs
    self.executor = EXECUTORS[config.executor.type](self)
    self.publisher = (
        ArtifactPublisher(self, config.artifacts)
        if config.artifacts is not None
        else None
    )

373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430

def release_probe(self, rung: int) -> None:
    """Release a Slurm rung only after its probe proves an optimizer update."""
    state = self.read_state()
    controller_commit = os.environ.get("SAMUDRA_CODE_COMMIT")
    expected_commit = state.get("provenance", {}).get("commit")
    if controller_commit != expected_commit:
        raise ValueError(
            f"Search was started at {expected_commit}; probe controller "
            f"reports {controller_commit!r}"
        )
    current = state["rungs"][rung]
    probe = current.get("probe")
    if not isinstance(probe, dict):
        raise ValueError(f"Rung {rung} has no configured probe")
    name = probe["candidate"]
    status_path = (
        self.search_dir / "probe" / resource_slug(name) / SEARCH_WORKER_STATUS_NAME
    )
    try:
        status = json.loads(status_path.read_text(encoding="utf-8"))
        if status.get("stage") != "completed":
            raise ValueError(
                f"worker stopped at stage {status.get('stage')!r}: "
                f"{status.get('error', 'no worker error recorded')}"
            )
        if int(status.get("optimizer_steps", 0)) < 1:
            raise ValueError("worker completed without an optimizer update")
    except (OSError, TypeError, ValueError, json.JSONDecodeError) as error:
        state["status"] = "failed"
        state["failure"] = {
            "stage": "rung_probe",
            "type": type(error).__name__,
            "message": str(error),
            "candidate": name,
            "failed_at": datetime.datetime.now(datetime.UTC).isoformat(),
        }
        probe["status"] = "failed"
        self.write_state(state)
        self._write_results(state)
        self.publish(state)
        try:
            self._write_report(state)
            self.publish(state)
        except Exception as report_error:
            error.add_note(f"Failed to render or publish report: {report_error!r}")
        raise ValueError(f"Rung {rung} probe failed: {error}") from error
    probe.update(
        status="complete",
        optimizer_steps=status["optimizer_steps"],
        batches_seen=status.get("batches_seen"),
        validated_at=datetime.datetime.now(datetime.UTC).isoformat(),
    )
    self.write_state(state)
    if not isinstance(self.executor, SlurmExecutor):
        raise TypeError("Rung probes are only supported by the Slurm executor")
    self.executor.submit_anchors(state)
    self.executor.submit_validated_rung(state, rung)
    self.publish(self.read_state())