slurmforge 0.1.0__py3-none-any.whl

slurmforge/__init__.py

@@ -0,0 +1,7 @@
+"""Generic Slurm-oriented experiment orchestration package for slurmforge."""
+
+from __future__ import annotations
+
+from .identity import PACKAGE_NAME, __version__
+
+__all__ = ["PACKAGE_NAME", "__version__"]

slurmforge/cli/__init__.py

@@ -0,0 +1 @@
+"""CLI subcommands for slurmforge."""

slurmforge/cli/common.py

@@ -0,0 +1,95 @@
+"""Shared CLI argument builders and batch materialization helpers."""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import yaml
+
+from ..errors import ConfigContractError
+from ..identity import __version__, regenerate_after_upgrade_note
+from ..pipeline.materialization import MaterializationResult, materialize_batch, print_dry_run
+from ..pipeline.planning import PlannedBatch
+from ..sweep import deep_set, parse_override
+from ..templating import build_template_env
+
+
+def load_raw_cfg(config_path: Path) -> dict:
+    """Load a YAML config file as a plain dict, without applying any overrides."""
+    text = config_path.read_text(encoding="utf-8")
+    cfg = yaml.safe_load(text)
+    if not isinstance(cfg, dict):
+        raise ConfigContractError(f"Config must be a YAML mapping: {config_path}")
+    return cfg
+
+
+def load_effective_cfg(config_path: Path, cli_overrides: list[str]) -> dict:
+    """
+    Load YAML config and apply CLI overrides, producing the effective config.
+
+    This represents the user's declared intent — raw YAML merged with any
+    ``--set`` overrides — and is the single config view that all pre-compile
+    checks (completeness, correctness, advisory) should operate on.
+
+    The compiler receives ``(config_path, cli_overrides)`` separately because
+    it tracks override provenance for diagnostics and replay.
+    """
+    cfg = load_raw_cfg(config_path)
+    for override in cli_overrides:
+        key, value = parse_override(override)
+        deep_set(cfg, key, value)
+    return cfg
+
+
+def add_config_override_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--set",
+        action="append",
+        default=[],
+        help="Override config by dot-path, e.g. --set run.args.lr=0.004",
+    )
+    parser.add_argument(
+        "--project_root",
+        default=None,
+        help="Override project root used to resolve relative paths (default: config file directory)",
+    )
+
+
+def add_dry_run_arg(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--dry_run",
+        action="store_true",
+        help="Expand runs and print commands without writing files",
+    )
+
+
+def add_common_args(parser: argparse.ArgumentParser) -> None:
+    add_config_override_args(parser)
+    add_dry_run_arg(parser)
+
+
+def materialize_or_print_batch(
+    *,
+    planned_batch: PlannedBatch,
+    dry_run: bool,
+) -> MaterializationResult | None:
+    if dry_run:
+        print_dry_run(planned_run.plan for planned_run in planned_batch.planned_runs)
+        return None
+
+    env = build_template_env()
+    return materialize_batch(
+        planned_batch=planned_batch,
+        env=env,
+    )
+
+
+def print_batch_ready(*, dispatch: MaterializationResult, sbatch_dir: Path, retry: bool = False) -> None:
+    if retry:
+        print(f"[OK] Generated retry batch with {len(dispatch.array_groups_meta)} array sbatch file(s) in: {sbatch_dir}")
+    else:
+        print(f"[OK] Generated {len(dispatch.array_groups_meta)} array sbatch file(s) in: {sbatch_dir}")
+    print(f"[OK] Generated by: slurmforge {__version__}")
+    print(f"[OK] Submit all via: {dispatch.submit_script}")
+    print(f"[OK] Batch manifest: {dispatch.manifest_path}")
+    print(f"[NOTE] {regenerate_after_upgrade_note()}")

slurmforge/cli/examples.py

@@ -0,0 +1,50 @@
+"""``sforge examples`` -- list, show, or export shipped YAML reference examples."""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from ..example_configs import export_example, list_example_catalog, read_example_text
+
+
+def handle_examples_list(_args: argparse.Namespace) -> None:
+    for name, description in list_example_catalog():
+        if description:
+            print(f"{name:<26} {description}")
+        else:
+            print(name)
+
+
+def handle_examples_show(args: argparse.Namespace) -> None:
+    print(read_example_text(args.name), end="")
+
+
+def handle_examples_export(args: argparse.Namespace) -> None:
+    exported = export_example(args.name, Path(args.out), force=args.force)
+    print(f"[OK] Wrote example config: {exported}")
+
+
+def add_subparser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
+    examples_parser = subparsers.add_parser(
+        "examples",
+        help="List, show, or export shipped raw YAML reference examples",
+    )
+    examples_subparsers = examples_parser.add_subparsers(dest="examples_command")
+    examples_subparsers.required = True
+
+    list_parser = examples_subparsers.add_parser("list", help="List shipped raw YAML reference examples")
+    list_parser.set_defaults(handler=handle_examples_list)
+
+    show_parser = examples_subparsers.add_parser("show", help="Print one shipped raw YAML reference example")
+    show_parser.add_argument("name", help="Example name, with or without .yaml suffix")
+    show_parser.set_defaults(handler=handle_examples_show)
+
+    export_parser = examples_subparsers.add_parser("export", help="Copy one shipped raw YAML reference example to a file")
+    export_parser.add_argument("name", help="Example name, with or without .yaml suffix")
+    export_parser.add_argument("--out", required=True, help="Destination path for the exported raw YAML example")
+    export_parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Overwrite the destination file if it already exists",
+    )
+    export_parser.set_defaults(handler=handle_examples_export)

slurmforge/cli/generate.py

@@ -0,0 +1,102 @@
+"""``sforge generate`` -- expand an experiment config into sbatch arrays."""
+from __future__ import annotations
+
+import argparse
+import datetime
+from pathlib import Path
+
+from ..pipeline.compiler import AuthoringSourceRequest, BatchCompileError, compile_source, iter_compile_report_lines
+from ..pipeline.compiler.reports import require_success
+from ..pipeline.config.validation.advisory import check_advisory
+from ..pipeline.config.validation.completeness import assert_complete
+from ..pipeline.config.validation.correctness import check_correctness
+from ..pipeline.config.validation.messages import format_advisory_report, format_correctness_report
+from .common import (
+    add_common_args,
+    load_effective_cfg,
+    materialize_or_print_batch,
+    print_batch_ready,
+)
+
+
+def render_generate(
+    *,
+    config_path: Path,
+    cli_overrides: list[str],
+    dry_run: bool,
+    project_root_override: str | None,
+) -> None:
+    resolved_config_path = config_path.resolve()
+    project_root = (
+        Path(project_root_override).resolve()
+        if project_root_override is not None
+        else resolved_config_path.parent
+    )
+
+    effective_cfg = load_effective_cfg(resolved_config_path, cli_overrides)
+
+    # ── Level 0: completeness gate (hard — null sentinels block generation) ──
+    assert_complete(effective_cfg, config_path=resolved_config_path, project_root=project_root)
+
+    # ── Level 1: correctness gate (hard — format/logic errors block generation) ─
+    correctness_errors = check_correctness(effective_cfg, config_path=resolved_config_path)
+    if correctness_errors:
+        print(format_correctness_report(
+            correctness_errors,
+            config_path=resolved_config_path,
+            force_flag_available=False,
+        ))
+        raise SystemExit(1)
+
+    # ── Level 2: advisory (soft — warnings printed, generation continues) ────
+    advisory_warnings = check_advisory(effective_cfg, config_path=resolved_config_path)
+    if advisory_warnings:
+        print(format_advisory_report(advisory_warnings, config_path=resolved_config_path))
+        print()
+
+    # ── Compiler pipeline ─────────────────────────────────────────────────────
+    report = compile_source(
+        AuthoringSourceRequest(
+            config_path=resolved_config_path,
+            cli_overrides=tuple(cli_overrides),
+            project_root=None if project_root_override is None else Path(project_root_override),
+            default_batch_name=datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f"),
+        ),
+    )
+    for line in iter_compile_report_lines(report):
+        print(line)
+    try:
+        planned_batch = require_success(report)
+    except BatchCompileError:
+        raise
+
+    dispatch = materialize_or_print_batch(
+        planned_batch=planned_batch,
+        dry_run=dry_run,
+    )
+    if dispatch is None:
+        return
+    print_batch_ready(dispatch=dispatch, sbatch_dir=planned_batch.sbatch_dir)
+
+
+def handle_generate(args: argparse.Namespace) -> None:
+    render_generate(
+        config_path=Path(args.config),
+        cli_overrides=args.set,
+        dry_run=args.dry_run,
+        project_root_override=args.project_root,
+    )
+
+
+def add_subparser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
+    generate_parser = subparsers.add_parser(
+        "generate",
+        help="Expand config into a new batch and render sbatch arrays",
+    )
+    generate_parser.add_argument(
+        "--config",
+        required=True,
+        help="Path to experiment config yaml",
+    )
+    add_common_args(generate_parser)
+    generate_parser.set_defaults(handler=handle_generate)

slurmforge/cli/init.py

@@ -0,0 +1,148 @@
+"""
+`sforge init` — create a starter project scaffold.
+
+Decision tree (two orthogonal choices):
+
+  TRAINING TYPE  (how is your training code invoked?)
+    script    → train.py with CLI args            [most common]
+    command   → complete shell command
+    registry  → shared team model registry
+    adapter   → interface bridge script
+
+  PROFILE  (cluster complexity)
+    starter   → single GPU, minimal config        [default]
+    hpc       → multi-GPU, sweep, eval, artifact sync
+
+Examples
+--------
+  sforge init                          # interactive wizard
+  sforge init script                   # script · starter profile
+  sforge init script --profile hpc    # script · hpc profile
+  sforge init command
+  sforge init command  --profile hpc
+  sforge init registry
+  sforge init registry --profile hpc
+  sforge init adapter
+  sforge init adapter  --profile hpc
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from ..starter_catalog import PROFILES, TEMPLATE_TYPES, get_starter_spec
+from ..starter_projects import init_project
+from .init_wizard import run_wizard
+
+_DEFAULT_OUT = "./slurmforge_starter"
+
+_TYPE_DESCRIPTIONS = {
+    "script":   "Scaffold for a train.py-style script — slurmforge manages args and submission.",
+    "command":  "Scaffold that wraps a complete shell command in Slurm.",
+    "registry": "Scaffold using a shared team model registry.",
+    "adapter":  "Scaffold with an interface-bridge adapter script (advanced).",
+}
+
+
+# ---------------------------------------------------------------------------
+# Shared argument builder
+# ---------------------------------------------------------------------------
+
+def _add_common_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--profile",
+        default="starter",
+        choices=PROFILES,
+        metavar="PROFILE",
+        help=(
+            "Cluster complexity profile. "
+            "'starter' = single GPU, minimal config (default). "
+            "'hpc' = multi-GPU, sweep, eval, artifact sync."
+        ),
+    )
+    parser.add_argument(
+        "--out",
+        default=_DEFAULT_OUT,
+        metavar="DIR",
+        help="Destination directory for the project scaffold (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Overwrite existing files in the destination directory",
+    )
+
+
+# ---------------------------------------------------------------------------
+# Handlers
+# ---------------------------------------------------------------------------
+
+def _do_init(*, template_type: str, profile: str, out: str, force: bool) -> None:
+    spec = get_starter_spec(template_type, profile)
+    written = init_project(template_type, profile, Path(out), force=force)
+    out_dir = Path(out).expanduser().resolve()
+    print(f"[OK] Initialized '{template_type}' scaffold (profile: {profile}) in: {out_dir}")
+    print(f"[INFO] {spec.post_init_guidance}")
+    print()
+    print("  Files created:")
+    for path in written:
+        print(f"    {path}")
+    print()
+    print("  Next steps:")
+    print(f"    1. Open  {out_dir / 'experiment.yaml'}")
+    print("    2. Fill in every field marked with ~  (required — see STEP 1 comments)")
+    print(f"    3. Run:  sforge validate --config {out_dir / 'experiment.yaml'}")
+    print(f"    4. Run:  sforge generate --config {out_dir / 'experiment.yaml'}")
+
+
+def handle_init_template(args: argparse.Namespace) -> None:
+    _do_init(
+        template_type=args.template_type,
+        profile=args.profile,
+        out=args.out,
+        force=args.force,
+    )
+
+
+def handle_init_wizard(args: argparse.Namespace) -> None:
+    """Fallback handler when no TYPE subcommand is given — launches interactive wizard."""
+    template_type, profile, out = run_wizard(out=args.out, force=args.force)
+    _do_init(template_type=template_type, profile=profile, out=out, force=args.force)
+
+
+# ---------------------------------------------------------------------------
+# Parser registration
+# ---------------------------------------------------------------------------
+
+def add_subparser(subparsers: argparse._SubParsersAction) -> None:  # type: ignore[type-arg]
+    init_parser = subparsers.add_parser(
+        "init",
+        description=__doc__,
+        help="Create a starter project scaffold  (run 'sforge init' for interactive setup)",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    # Top-level --out/--force for wizard path (TYPE subcommand overrides these)
+    init_parser.add_argument(
+        "--out",
+        default=_DEFAULT_OUT,
+        metavar="DIR",
+        help="Output directory (wizard mode — overridden by TYPE subcommand flags)",
+    )
+    init_parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Overwrite existing files",
+    )
+    init_parser.set_defaults(handler=handle_init_wizard)
+
+    # TYPE subcommands: script / command / registry / adapter
+    type_subparsers = init_parser.add_subparsers(dest="template_type")
+
+    for ttype in TEMPLATE_TYPES:
+        tp = type_subparsers.add_parser(
+            ttype,
+            help=_TYPE_DESCRIPTIONS.get(ttype, ""),
+            description=_TYPE_DESCRIPTIONS.get(ttype, ""),
+        )
+        tp.set_defaults(template_type=ttype, handler=handle_init_template)
+        _add_common_args(tp)

slurmforge/cli/init_wizard.py

@@ -0,0 +1,80 @@
+"""
+Interactive 2-question wizard for `sforge init` (no arguments).
+Falls back gracefully when stdin is not a TTY (CI, piped input).
+"""
+from __future__ import annotations
+
+import sys
+
+
+_TRAINING_TYPES = [
+    ("script",   "I have a train.py — slurmforge manages my args and submission"),
+    ("command",  "I have a complete launch command — just wrap it in Slurm"),
+    ("registry", "My training code lives in a shared team model registry"),
+    ("adapter",  "I need an interface bridge between slurmforge and my training code"),
+]
+
+_PROFILES = [
+    ("starter", "Quick start — single GPU, minimal config (recommended for first run)"),
+    ("hpc",     "Full HPC — multi-GPU, sweep, eval, artifact sync"),
+]
+
+
+def _prompt_choice(
+    prompt: str,
+    options: list[tuple[str, str]],
+) -> str:
+    """Display numbered options and return the selected key."""
+    print(prompt)
+    for i, (key, description) in enumerate(options, 1):
+        print(f"  {i}) {key:<12} {description}")
+    print()
+    while True:
+        try:
+            raw = input("  Enter number: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print()
+            sys.exit(0)
+        if raw.isdigit():
+            idx = int(raw) - 1
+            if 0 <= idx < len(options):
+                chosen_key = options[idx][0]
+                print(f"  → {chosen_key}")
+                print()
+                return chosen_key
+        print(f"  Please enter a number between 1 and {len(options)}.")
+
+
+def run_wizard(*, out: str, force: bool) -> tuple[str, str, str]:
+    """
+    Run the interactive wizard and return (template_type, profile, out_dir).
+    Raises SystemExit if the user cancels.
+    """
+    if not sys.stdin.isatty():
+        print(
+            "[sforge init] No template type specified.\n"
+            "Usage: sforge init <TYPE> [--profile starter|hpc] [--out DIR]\n"
+            "\n"
+            "Available types: script, command, registry, adapter\n"
+            "Run 'sforge init --help' for full usage.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    print()
+    print("  ┌──────────────────────────────────────────────────────┐")
+    print("  │  sforge init · project setup wizard                  │")
+    print("  └──────────────────────────────────────────────────────┘")
+    print()
+
+    template_type = _prompt_choice(
+        "  How is your training code invoked?",
+        _TRAINING_TYPES,
+    )
+
+    profile = _prompt_choice(
+        "  Which cluster profile fits your setup?",
+        _PROFILES,
+    )
+
+    return template_type, profile, out

slurmforge/cli/replay.py

@@ -0,0 +1,102 @@
+"""``sforge replay`` -- regenerate a batch from persisted run snapshots."""
+from __future__ import annotations
+
+import argparse
+import datetime
+from pathlib import Path
+
+from ..pipeline.compiler import BatchCompileError, ReplaySourceRequest, compile_source, iter_compile_report_lines
+from ..pipeline.compiler.reports import report_total_runs, require_success
+from .common import add_common_args, materialize_or_print_batch, print_batch_ready
+
+
+def render_replay(
+    *,
+    source_run_dir: Path | None,
+    source_snapshot_path: Path | None,
+    source_batch_root: Path | None,
+    run_ids: list[str],
+    run_indices: list[int],
+    cli_overrides: list[str],
+    dry_run: bool,
+    project_root_override: str | None,
+) -> None:
+    default_batch_name = datetime.datetime.now().strftime("replay_%Y%m%d_%H%M%S_%f")
+    request = ReplaySourceRequest(
+        source_run_dir=source_run_dir,
+        source_snapshot_path=source_snapshot_path,
+        source_batch_root=source_batch_root,
+        run_ids=tuple(run_ids),
+        run_indices=tuple(run_indices),
+        cli_overrides=tuple(cli_overrides),
+        project_root=None if project_root_override is None else Path(project_root_override),
+        default_batch_name=default_batch_name,
+    )
+    report = compile_source(request)
+    source_summary = getattr(report, "source_summary", "") or "<missing replay source>"
+    print(f"[REPLAY] source={source_summary} selected_runs={report_total_runs(report)}")
+    for line in iter_compile_report_lines(report):
+        print(line)
+    try:
+        planned_batch = require_success(report)
+    except BatchCompileError:
+        raise
+
+    dispatch = materialize_or_print_batch(
+        planned_batch=planned_batch,
+        dry_run=dry_run,
+    )
+    if dispatch is None:
+        return
+    print_batch_ready(dispatch=dispatch, sbatch_dir=planned_batch.sbatch_dir)
+
+
+def handle_replay(args: argparse.Namespace) -> None:
+    render_replay(
+        source_run_dir=None if args.source_run_dir is None else Path(args.source_run_dir),
+        source_snapshot_path=None if args.source_snapshot_path is None else Path(args.source_snapshot_path),
+        source_batch_root=None if args.source_batch_root is None else Path(args.source_batch_root),
+        run_ids=args.run_id,
+        run_indices=args.run_index,
+        cli_overrides=args.set,
+        dry_run=args.dry_run,
+        project_root_override=args.project_root,
+    )
+
+
+def add_subparser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
+    replay_parser = subparsers.add_parser(
+        "replay",
+        help="Replay one or more persisted runs from a run dir, snapshot, or batch root",
+    )
+    source_group = replay_parser.add_mutually_exclusive_group(required=True)
+    source_group.add_argument(
+        "--from-run",
+        dest="source_run_dir",
+        help="Path to a persisted run directory containing meta/run_snapshot.json",
+    )
+    source_group.add_argument(
+        "--from-snapshot",
+        dest="source_snapshot_path",
+        help="Path to a persisted run_snapshot.json file",
+    )
+    source_group.add_argument(
+        "--from-batch",
+        dest="source_batch_root",
+        help="Path to an existing batch_root; replays all runs unless selectors are provided",
+    )
+    replay_parser.add_argument(
+        "--run_id",
+        action="append",
+        default=[],
+        help="Select specific run_id values when replaying from --from-batch",
+    )
+    replay_parser.add_argument(
+        "--run_index",
+        action="append",
+        type=int,
+        default=[],
+        help="Select specific run_index values when replaying from --from-batch",
+    )
+    add_common_args(replay_parser)
+    replay_parser.set_defaults(handler=handle_replay)

slurmforge/cli/rerun.py

@@ -0,0 +1,76 @@
+"""``sforge rerun`` -- rebuild a retry batch from an existing batch's run records."""
+from __future__ import annotations
+
+import argparse
+import datetime
+from pathlib import Path
+
+from ..pipeline.compiler import BatchCompileError, RetrySourceRequest, compile_source, iter_compile_report_lines
+from ..pipeline.compiler.reports import report_total_runs, require_success
+from .common import add_common_args, materialize_or_print_batch, print_batch_ready
+
+
+def render_rerun(
+    *,
+    source_batch_root: Path,
+    cli_overrides: list[str],
+    dry_run: bool,
+    project_root_override: str | None,
+    status_query: str,
+) -> None:
+    default_batch_name = datetime.datetime.now().strftime("retry_%Y%m%d_%H%M%S_%f")
+    report = compile_source(
+        RetrySourceRequest(
+            source_batch_root=source_batch_root,
+            status_query=status_query,
+            cli_overrides=tuple(cli_overrides),
+            project_root=None if project_root_override is None else Path(project_root_override),
+            default_batch_name=default_batch_name,
+        )
+    )
+    source_summary = getattr(report, "source_summary", "") or str(source_batch_root)
+    print(f"[RETRY] source={source_summary} selected_runs={report_total_runs(report)}")
+    for line in iter_compile_report_lines(report):
+        print(line)
+    try:
+        planned_batch = require_success(report)
+    except BatchCompileError:
+        raise
+
+    dispatch = materialize_or_print_batch(
+        planned_batch=planned_batch,
+        dry_run=dry_run,
+    )
+    if dispatch is None:
+        return
+    print_batch_ready(dispatch=dispatch, sbatch_dir=planned_batch.sbatch_dir, retry=True)
+
+
+def handle_rerun(args: argparse.Namespace) -> None:
+    render_rerun(
+        source_batch_root=Path(args.source_batch_root).resolve(),
+        cli_overrides=args.set,
+        dry_run=args.dry_run,
+        project_root_override=args.project_root,
+        status_query=args.status,
+    )
+
+
+def add_subparser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
+    rerun_parser = subparsers.add_parser("rerun", help="Rebuild a retry batch from an existing batch_root")
+    rerun_parser.add_argument(
+        "--from",
+        dest="source_batch_root",
+        required=True,
+        help="Path to an existing batch_root; rebuild and resubmit a filtered retry batch from its run records",
+    )
+    rerun_parser.add_argument(
+        "--status",
+        default="failed",
+        help=(
+            "Retry filter for existing batch runs: failed(non-success), success, "
+            "pending, running, oom, preempted, node_failure, script_error, eval_failed, all"
+        ),
+    )
+    add_common_args(rerun_parser)
+    rerun_parser.set_defaults(handler=handle_rerun)