flashruntime 0.3.0__py3-none-any.whl

flashruntime/strategies/__init__.py

@@ -0,0 +1,156 @@
+"""Strategy compilers: a frozen `StrategyPlan` → concrete launch configuration.
+
+This is the seam that keeps the planner honest: the planner reasons about
+strategies *by name and number* (never importing frameworks), and a
+compiler here translates the chosen plan into the real thing — torchrun
+argv + env for DDP/FSDP2, a DeepSpeed JSON config file for ZeRO families,
+a lease-job JobSpec for `lease_tasks`. One compiler per strategy family;
+the registry dispatches on `supports()`.
+
+Design rules (ADR-0003, HANDBOOK §5):
+- Compilers may import framework *constants* knowledge but must not
+  require the framework installed to *compile* (emitting an argv string
+  needs no torch). Anything that must run framework code belongs in the
+  launcher or the task, not here.
+- Compilation is pure and deterministic: same plan ⇒ same LaunchSpec.
+  No environment inspection, no network. Preflight checks (is torchrun on
+  PATH? is the image pullable?) belong to `Launcher.healthy()`.
+- A compiler must refuse loudly (`CompileError` listing every problem)
+  rather than emit a config it cannot stand behind — the same fail-closed
+  stance as the executor allowlists.
+
+Status: interface complete (final surface, by design review); concrete
+compilers land per SPRINT_PLAN (torchrun/DDP first, with the LoRA recipe).
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import ClassVar
+
+from pydantic import BaseModel, Field
+
+from flashruntime.protocol.plan_v1alpha1 import StrategyPlan
+
+__all__ = [
+    "LaunchSpec",
+    "StrategyCompiler",
+    "CompileError",
+    "register_compiler",
+    "compiler_for",
+]
+
+
+class CompileError(Exception):
+    """The plan cannot be compiled by this family — carries *all* reasons
+    at once (like planner rejections: half the value is the explanation)."""
+
+
+class LaunchSpec(BaseModel):
+    """Everything a `Launcher` needs to start one worker group / job.
+
+    Backend-neutral on purpose: the same model describes a torchrun
+    process group, a DeepSpeed launch, or a Mode A lease job. Fields a
+    given launcher does not need are simply ignored by it.
+    """
+
+    argv: list[str] = Field(description="Command to execute per worker (rank 0's view)")
+    env: dict[str, str] = Field(
+        default_factory=dict,
+        description="Environment to merge over the launcher's base env "
+        "(e.g. WORLD_SIZE, precision flags). Never secrets — credentials "
+        "are the launcher's concern.",
+    )
+    files: dict[str, str] = Field(
+        default_factory=dict,
+        description="Config files to materialize into the attempt OUTPUT dir "
+        "(the one exported as FLASHML_OUTPUT_DIR) before launch, name → "
+        "content (e.g. 'ds_config.json' for DeepSpeed). Keeps compilation "
+        "pure while supporting file-driven frameworks. Note: files land in "
+        "the output dir, NOT the command's cwd — wiring cwd-reading "
+        "frameworks to find them is a flagged follow-up to settle before any "
+        "files-emitting compiler ships.",
+    )
+    world_size: int = Field(default=1, ge=1, description="Total processes across all nodes")
+    gpus_per_process: int = Field(default=0, ge=0)
+    rendezvous: dict[str, str] = Field(
+        default_factory=dict,
+        description="Rendezvous hints the launcher resolves at start time "
+        "(e.g. {'backend': 'c10d'}). Addresses/ports are ALWAYS resolved "
+        "by the launcher — a compiled plan must stay host-agnostic.",
+    )
+    workdir_hint: str = Field(
+        default="", description="Relative workdir layout the command expects, if any"
+    )
+    notes: list[str] = Field(
+        default_factory=list,
+        description="Human-readable compilation notes for the ledger "
+        "(mirrors the planner's explanation habit).",
+    )
+
+
+class StrategyCompiler(ABC):
+    """One strategy family's translator.
+
+    Lifecycle: `supports(plan)` → `validate(plan)` → `compile(plan)`.
+    The registry calls `supports` to route; callers may run `validate`
+    separately for dry-run UX (surface every problem before spending
+    anything); `compile` implies validation and raises `CompileError` on
+    any defect.
+    """
+
+    #: Family name this compiler owns — must match
+    #: `StrategyPlan.strategy_family` values (e.g. "ddp", "fsdp2",
+    #: "zero3_cpu_offload", "lease_tasks").
+    family: ClassVar[str]
+
+    @abstractmethod
+    def supports(self, plan: StrategyPlan) -> bool:
+        """Cheap routing predicate. Must not raise; unknown plans → False."""
+
+    def validate(self, plan: StrategyPlan) -> list[str]:
+        """Return every reason this plan cannot compile (empty = fine).
+
+        Default: no extra checks beyond `supports`. Override to verify
+        family-specific invariants (e.g. FSDP2 requires workers > 1;
+        DeepSpeed offload requires `plan.offload != 'none'`). Never check
+        the *environment* here — that is `Launcher.healthy()`.
+        """
+        return [] if self.supports(plan) else [f"{self.family}: unsupported plan"]
+
+    @abstractmethod
+    def compile(self, plan: StrategyPlan) -> LaunchSpec:
+        """Translate the plan into a LaunchSpec.
+
+        Input: a frozen, planner-emitted StrategyPlan (treat as immutable).
+        Output: a complete LaunchSpec — argv, env, config files.
+        Raises: CompileError with *all* problems when validation fails.
+        Determinism contract: identical plan ⇒ identical LaunchSpec
+        (the ledger records the spec hash next to the plan_id).
+        """
+
+
+_REGISTRY: list[StrategyCompiler] = []
+
+
+def register_compiler(compiler: StrategyCompiler) -> None:
+    """Register a compiler instance (idempotent per family: a re-register
+    of the same family replaces the previous one — supports test setups
+    and plugin reloads)."""
+    _REGISTRY[:] = [c for c in _REGISTRY if c.family != compiler.family]
+    _REGISTRY.append(compiler)
+
+
+def compiler_for(plan: StrategyPlan) -> StrategyCompiler:
+    """Dispatch: first registered compiler whose `supports(plan)` is True.
+
+    Raises LookupError naming the family — the caller surfaces this as a
+    422 (\"plan compiled by no installed strategy family\"), never a 500.
+    """
+    for compiler in _REGISTRY:
+        if compiler.supports(plan):
+            return compiler
+    raise LookupError(
+        f"no strategy compiler for family {plan.strategy_family!r} "
+        f"(registered: {[c.family for c in _REGISTRY] or 'none'})"
+    )

flashruntime/strategies/command.py

@@ -0,0 +1,56 @@
+"""Compile a CommandWorkload into the backend-neutral LaunchSpec.
+
+Pure and deterministic (same rules as StrategyCompiler): no environment
+inspection, no filesystem access — resolving the workdir and creating
+output directories is the launcher's job.
+
+Note: this is a module function, not a StrategyCompiler subclass — a
+StrategyPlan carries no argv, so a plan-driven compiler for commands is
+meaningless until flash.run() wiring lands (spec §10 follow-up).
+"""
+
+from __future__ import annotations
+
+from flashruntime.strategies import LaunchSpec
+from flashruntime.workloads.command import CommandWorkload
+
+
+def compile_workload(workload: CommandWorkload, params: dict | None = None) -> LaunchSpec:
+    argv = workload.argv(params)
+    env = {k: (v.format(**params) if params else v) for k, v in workload.env.items()}
+    world_size = 1
+    notes = [f"mode={workload.resolved_mode()}"]
+    if argv and argv[0] == "torchrun":
+        # torchrun's worker count flag has two spellings (hyphen/underscore)
+        # and two forms (--flag=N or --flag N). Values may also be
+        # symbolic (auto/gpu/cpu), resolved by torchrun at launch — those
+        # must not crash compilation: leave world_size=1 and note it.
+        nproc_flags = ("--nproc-per-node", "--nproc_per_node")
+        i = 1
+        while i < len(argv):
+            token = argv[i]
+            flag = value = None
+            if "=" in token:
+                head, _, tail = token.partition("=")
+                if head in nproc_flags:
+                    flag, value = head, tail
+            elif token in nproc_flags:
+                flag = token
+                value = argv[i + 1] if i + 1 < len(argv) else None
+                i += 1  # consume the value token
+            if flag is not None and value is not None:
+                try:
+                    world_size = int(value)
+                    notes.append(f"world_size from torchrun: {world_size}")
+                except ValueError:
+                    notes.append(
+                        f"world_size unresolved: {flag}={value} (resolved at launch time)"
+                    )
+            i += 1
+    return LaunchSpec(
+        argv=argv,
+        env=env,
+        world_size=world_size,
+        workdir_hint=workload.source.path,
+        notes=notes,
+    )

flashruntime/torch/__init__.py

@@ -0,0 +1,274 @@
+"""In-training-script helper: one import makes a PyTorch script
+launch-anywhere and fault-tolerant.
+
+    import flashruntime.torch as ft
+    model, opt, loader = ft.prepare(model, opt, loader)
+    ...
+    ft.checkpoint(model, opt, step=step, every=100)
+    ft.log_metrics({"loss": float(loss)})
+
+Launched by torchrun (WORLD_SIZE>1): prepare() wires torch's OWN DDP +
+DistributedSampler and restores the newest VALID checkpoint manifest.
+Launched as plain `python train.py`: prepare() is a no-op passthrough.
+
+GUARDRAIL (ADR-0003 — do not rebuild Accelerate): this module wires
+torch's own primitives and REPORTS launch facts — nothing more (that
+reporting includes a per-rank heartbeat file the run viewer reads). The
+surface is prepare / checkpoint / log_metrics plus read-only accessors
+(rank / world_size / is_main / start_step / device / backend). The
+boundary is capability, not count: no FSDP policies, no autocast, no
+DeepSpeed config, no strategy knobs — ever. Users wanting those use the
+real framework features, which the launcher still launches correctly.
+
+torch is imported inside functions only: flashruntime's core never
+depends on it.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from pathlib import Path
+
+__all__ = [
+    "prepare",
+    "checkpoint",
+    "log_metrics",
+    "rank",
+    "world_size",
+    "is_main",
+    "start_step",
+    "device",
+    "backend",
+]
+
+_restored_step = 0
+# Set by prepare() from launch facts so a script can REPORT what it ran on
+# (device string + torch.distributed backend) without re-deriving it. Both
+# stay at their launched-single-process defaults until prepare() runs.
+_device = "cpu"
+_backend: str | None = None
+
+# Heartbeat state: the run viewer draws machine → worker → rank from a small
+# per-rank JSON each process mirrors to ranks/rank-N.json. Throttled so a
+# tight training loop calling log_metrics()/checkpoint() every step costs at
+# most one small atomic write per second per rank.
+_last_beat = 0.0
+_last_step: int | None = None
+_BEAT_MIN_INTERVAL_S = 1.0
+
+
+def _write_heartbeat(step: int | None = None, force: bool = False) -> None:
+    """Mirror this rank's identity + progress to ranks/rank-<N>.json.
+
+    Atomic (tmp + os.replace, the run.json idiom) so the viewer never reads
+    a torn file. Best-effort by contract: observability must never be able
+    to crash training, so every failure is swallowed. `step=None` reuses the
+    last step this process reported (a refresh must not erase progress)."""
+    global _last_beat, _last_step
+    if step is not None:
+        _last_step = step
+    now = time.time()
+    if not force and now - _last_beat < _BEAT_MIN_INTERVAL_S:
+        return
+    try:
+        beat = {
+            "rank": rank(),
+            "local_rank": int(os.environ.get("LOCAL_RANK", "0")),
+            "pid": os.getpid(),
+            "device": _device,
+            "backend": _backend,
+            "world_size": world_size(),
+            "step": _last_step,
+            "ts": now,
+        }
+        ranks_dir = _output_dir() / "ranks"
+        ranks_dir.mkdir(parents=True, exist_ok=True)
+        tmp = ranks_dir / f".rank-{rank()}.{os.getpid()}.tmp"
+        tmp.write_text(json.dumps(beat))
+        os.replace(tmp, ranks_dir / f"rank-{rank()}.json")
+        _last_beat = now
+    except Exception:  # noqa: BLE001 — by contract, swallow everything
+        pass
+
+
+def _resolve_device(world_size: int, cuda_available: bool, local_rank: int) -> str:
+    """Decide the training device from launch facts alone — returns ``"cpu"``
+    or ``f"cuda:{local_rank}"``.
+
+    Kept pure (no torch import, no globals) so the device *decision* is
+    testable on this CPU-only box: the CUDA code paths in ``prepare`` that
+    consume the result cannot execute here, but the choice they depend on
+    can. ``world_size`` is part of the launch-facts signature; the device
+    string itself depends only on whether CUDA is present and which local
+    rank this process owns.
+    """
+    if cuda_available:
+        return f"cuda:{local_rank}"
+    return "cpu"
+
+
+def world_size() -> int:
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def rank() -> int:
+    return int(os.environ.get("RANK", "0"))
+
+
+def is_main() -> bool:
+    return rank() == 0
+
+
+def start_step() -> int:
+    """First step the training loop should run: 0 fresh, >0 after a
+    resume (set by prepare() when it restores a checkpoint)."""
+    return _restored_step
+
+
+def device() -> str:
+    """The device the last ``prepare()`` placed this rank's model on —
+    ``"cpu"`` or ``f"cuda:{local_rank}"``. Lets a script report where it
+    trained (e.g. into ``metrics.json``) without re-deriving the choice."""
+    return _device
+
+
+def backend() -> str | None:
+    """The ``torch.distributed`` backend the last ``prepare()`` initialized
+    (``"nccl"`` on GPU, ``"gloo"`` on CPU), or ``None`` when launched
+    single-process (no process group)."""
+    return _backend
+
+
+def _output_dir() -> Path:
+    return Path(os.environ.get("FLASHML_OUTPUT_DIR", "."))
+
+
+def _ckpt_root() -> Path:
+    # job-scoped (NOT attempt-scoped): a restarted attempt must find its
+    # predecessor's manifests — the launcher exports FLASHML_CKPT_DIR
+    root = os.environ.get("FLASHML_CKPT_DIR")
+    return Path(root) if root else _output_dir() / "ckpt"
+
+
+def prepare(model, optimizer=None, dataloader=None):
+    """Wire distributed execution (when launched distributed) and restore
+    the newest valid checkpoint (when one exists). Returns the possibly
+    wrapped/rebuilt (model, optimizer, dataloader) triple."""
+    global _restored_step, _device, _backend
+    import torch
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    device = _resolve_device(world_size(), torch.cuda.is_available(), local_rank)
+    on_cuda = device != "cpu"
+    _device = device  # remember for device()/metrics reporting
+
+    if on_cuda:
+        # place the model on this rank's GPU BEFORE the DDP wrap (and for a
+        # single-process GPU box too — the "just works on a GPU" path). This
+        # branch is CUDA-only and never runs on the CPU-only test machine.
+        torch.cuda.set_device(local_rank)
+        model = model.to(device)
+
+    if world_size() > 1:
+        import torch.distributed as dist
+
+        _backend = "nccl" if torch.cuda.is_available() else "gloo"
+        if not dist.is_initialized():
+            dist.init_process_group(backend=_backend)
+        if on_cuda:
+            # bind DDP to the device; gloo/CPU must keep the no-args wrap.
+            model = torch.nn.parallel.DistributedDataParallel(
+                model, device_ids=[local_rank], output_device=local_rank
+            )
+        else:
+            model = torch.nn.parallel.DistributedDataParallel(model)
+        if dataloader is not None:
+            from torch.utils.data import DataLoader
+            from torch.utils.data.distributed import DistributedSampler
+
+            dataloader = DataLoader(
+                dataloader.dataset,
+                batch_size=dataloader.batch_size,
+                sampler=DistributedSampler(dataloader.dataset),
+                collate_fn=dataloader.collate_fn,
+                num_workers=dataloader.num_workers,
+                drop_last=dataloader.drop_last,
+            )
+
+    from flashruntime.checkpoint.local import latest_valid_manifest
+
+    manifest = latest_valid_manifest(_ckpt_root())
+    if manifest is not None:
+        step_dir = Path(manifest.storage_prefix)
+        target = model.module if hasattr(model, "module") else model
+        # load straight onto the live device: CPU-saved (device-agnostic)
+        # parts map onto "cpu" (unchanged) or this rank's "cuda:N".
+        target.load_state_dict(torch.load(step_dir / "model.pt", map_location=device))
+        if optimizer is not None and (step_dir / "optimizer.pt").is_file():
+            optimizer.load_state_dict(torch.load(step_dir / "optimizer.pt", map_location=device))
+        _restored_step = manifest.step
+
+    # announce this rank to the run viewer the moment it is wired
+    _write_heartbeat(step=_restored_step, force=True)
+    return model, optimizer, dataloader
+
+
+def checkpoint(model, optimizer=None, *, step: int, every: int | None = None) -> None:
+    """Write a resumable checkpoint under the manifest contract (parts
+    first, manifest last). rank 0 writes; every rank synchronizes on the
+    barrier so no one races past a half-written checkpoint.
+
+    `every=None` (default) checkpoints unconditionally; `every=N` gates to
+    every Nth step; `every<=0` means "no periodic checkpointing" and is a
+    no-op — a helper whose contract is fault tolerance must never itself
+    crash training over a config value (found by the benchmark suite:
+    `every=0` used to raise ZeroDivisionError)."""
+    _write_heartbeat(step=step)  # progress signal — fires even on gated calls
+    if every is not None and (every <= 0 or step == 0 or step % every != 0):
+        return
+    import torch
+
+    if is_main():
+        from flashruntime.checkpoint.local import write_manifest
+
+        step_dir = _ckpt_root() / f"step-{step:06d}"
+        step_dir.mkdir(parents=True, exist_ok=True)
+        target = model.module if hasattr(model, "module") else model
+        # detach + move to CPU so the saved parts are device-agnostic: a
+        # checkpoint trained on cuda:3 must restore onto cpu or a box with
+        # fewer GPUs (map_location handles the rest at restore time).
+        model_state = {k: v.detach().cpu() for k, v in target.state_dict().items()}
+        torch.save(model_state, step_dir / "model.pt")
+        if optimizer is not None:
+            torch.save(optimizer.state_dict(), step_dir / "optimizer.pt")
+        write_manifest(
+            step_dir,
+            job_id=os.environ.get("FLASHML_JOB_ID", "local"),
+            attempt_id=os.environ.get("FLASHML_ATTEMPT_ID", "local"),
+            step=step,
+            world_size=world_size(),
+            framework=f"pytorch-{torch.__version__.split('+')[0]}",
+        )
+    if world_size() > 1:
+        import torch.distributed as dist
+
+        dist.barrier()
+
+
+def log_metrics(metrics: dict) -> None:
+    """Append one JSON record to FLASHML_OUTPUT_DIR/metrics.jsonl (rank 0
+    only). Never raises — metrics must never kill training."""
+    # every rank beats (identity/progress); only rank 0 appends metrics below
+    step = metrics.get("step") if isinstance(metrics, dict) else None
+    _write_heartbeat(step=step if isinstance(step, int) else None)
+    if not is_main():
+        return
+    try:
+        path = _output_dir() / "metrics.jsonl"
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "a") as f:
+            f.write(json.dumps(metrics) + "\n")
+    except Exception:  # noqa: BLE001 — by contract, swallow everything
+        pass

flashruntime/viewer/__init__.py

@@ -0,0 +1,20 @@
+"""The stdlib run viewer: a read-only window onto a live run's directory.
+
+Two pieces, both stdlib-only so `import flashruntime.viewer` stays clean-core
+importable (no numpy/torch/fastapi): `state.collect()` assembles the
+`/api/state` snapshot from disk, and `server.RunViewerServer` serves it (plus
+the page and the docs) over HTTP. The only flashruntime dependency is
+`flashruntime.checkpoint.local` for manifest hash-verification — itself
+core-safe (stdlib + pydantic).
+
+Contract source: the `viewer_v1` run.json written by `flashruntime.sdk.Run`.
+The viewer consumes that versioned contract and NOTHING else about the SDK's
+internals, so the two evolve independently (spec §2b, horizontal
+extensibility).
+"""
+
+from __future__ import annotations
+
+from flashruntime.viewer.state import collect
+
+__all__ = ["collect"]