flashruntime 0.3.0__py3-none-any.whl

flashruntime/planner/selector.py

@@ -0,0 +1,169 @@
+"""Deterministic selection: filter → rank by objective → mint the plan.
+
+Identical inputs always produce the identical plan (the report carries the
+planner version and a request digest so decisions are reproducible and
+auditable — same principle as the recovery policy engine).
+
+Ranking is lexicographic per objective mode. Hard constraints (deadline,
+budget, allow_* flags) were already applied as `rejected_policy` during
+evaluation — the ranker only orders survivors:
+
+  cheapest : (cost, time)         — unknown values sort last, never win by omission
+  fastest  : (time, cost)
+  balanced : cost first (deadline is already a hard gate), time tiebreak
+  reliable : fewest moving parts — smallest world size, no offload, no
+             quantization, most memory headroom
+
+When nothing survives, the report still answers usefully: the nearest-miss
+hint names the minimal relaxation that would unlock a plan.
+"""
+
+from __future__ import annotations
+
+import math
+
+from flashruntime.planner.candidates import Evaluated
+from flashruntime.protocol.plan_v1alpha1 import (
+    CheckpointPolicy,
+    PlanReport,
+    PlanRequest,
+    StrategyPlan,
+)
+
+_INF = math.inf
+
+
+def _cost(e: Evaluated) -> float:
+    return e.verdict.est_cost_usd.value if e.verdict.est_cost_usd else _INF
+
+
+def _time(e: Evaluated) -> float:
+    return e.verdict.est_time_min.value if e.verdict.est_time_min else _INF
+
+
+def _rank_key(e: Evaluated, mode: str):
+    if mode == "cheapest":
+        return (_cost(e), _time(e), e.verdict.workers)
+    if mode == "fastest":
+        return (_time(e), _cost(e), e.verdict.workers)
+    if mode == "reliable":
+        return (
+            e.verdict.workers,
+            0 if e.offload == "none" else 1,
+            0 if e.quantization is None else 1,
+            -e.headroom_gb,
+        )
+    # balanced
+    return (_cost(e), _time(e), e.verdict.workers)
+
+
+def select(
+    req: PlanRequest,
+    evaluated: list[Evaluated],
+    planner_version: str,
+    request_digest: str,
+    notes: list[str],
+) -> PlanReport:
+    feasible = [e for e in evaluated if e.verdict.status == "feasible"]
+    report = PlanReport(
+        planner_version=planner_version,
+        request_digest=request_digest,
+        candidates=[e.verdict for e in evaluated],
+        warnings=list(notes),
+    )
+
+    if not feasible:
+        report.no_valid_strategy_hint = _nearest_miss(req, evaluated)
+        return report
+
+    feasible.sort(key=lambda e: _rank_key(e, req.objective.mode))
+    winner = feasible[0]
+    winner.verdict.status = "selected"
+    report.selected = _mint_plan(req, winner, planner_version, request_digest, feasible)
+    return report
+
+
+def _mint_plan(
+    req: PlanRequest,
+    e: Evaluated,
+    planner_version: str,
+    request_digest: str,
+    feasible: list[Evaluated],
+) -> StrategyPlan:
+    v = e.verdict
+    because = list(v.reasons)
+    because.append(f"objective '{req.objective.mode}': best-ranked of {len(feasible)} feasible candidate(s)")
+    runner_up = feasible[1] if len(feasible) > 1 else None
+    if runner_up is not None:
+        because.append(
+            f"runner-up was '{runner_up.verdict.name}' "
+            f"(time {_fmt(runner_up.verdict.est_time_min and runner_up.verdict.est_time_min.value, 'min')}, "
+            f"cost {_fmt(runner_up.verdict.est_cost_usd and runner_up.verdict.est_cost_usd.value, 'usd')})"
+        )
+
+    checkpoint = None
+    if e.workload_mode == "coordinated_training":
+        # Static default; Young–Daly (τ* = √(2·C·MTBF)) needs a measured
+        # checkpoint duration and pool failure rate — profiling/ledger work.
+        checkpoint = CheckpointPolicy(
+            backend="pytorch_dcp",
+            interval_seconds=300,
+            note="static default; tune with Young–Daly once checkpoint duration is measured",
+        )
+
+    return StrategyPlan(
+        plan_id=f"pl_{request_digest}",
+        planner_version=planner_version,
+        workload_mode=e.workload_mode,  # type: ignore[arg-type]
+        strategy_family=v.strategy_family,
+        launcher=e.launcher,
+        workers=v.workers,
+        gpus_per_worker=e.gpus_per_worker,
+        colocated=e.colocated,
+        precision=e.precision,
+        quantization=e.quantization,
+        peft=e.peft,
+        micro_batch_per_gpu=e.micro_batch,
+        grad_accum=e.grad_accum,
+        activation_checkpointing=e.act_ckpt,
+        offload=e.offload,  # type: ignore[arg-type]
+        libraries=e.libraries,
+        checkpoint=checkpoint,
+        memory=v.memory,
+        est_time_min=v.est_time_min,
+        est_cost_usd=v.est_cost_usd,
+        scaling_efficiency=v.scaling_efficiency,
+        profiling_required=v.profiling_required,
+        selected_because=because,
+    )
+
+
+def _nearest_miss(req: PlanRequest, evaluated: list[Evaluated]) -> str:
+    """Turn a dead end into the minimal unlocking relaxation."""
+    hints: list[str] = []
+    mem_missed = [e for e in evaluated if e.vram_deficit_gb > 0]
+    if mem_missed:
+        closest = min(mem_missed, key=lambda e: e.vram_deficit_gb)
+        hints.append(
+            f"closest candidate '{closest.verdict.name}' missed VRAM by "
+            f"{closest.vram_deficit_gb:.1f} GB/GPU — add GPUs or larger GPUs"
+        )
+        if not req.objective.allow_quantization:
+            hints.append("allowing quantization (objective.allow_quantization) may unlock QLoRA")
+        if not req.objective.allow_cpu_offload:
+            hints.append("allowing CPU offload (objective.allow_cpu_offload) may unlock ZeRO-3 offload")
+    policy_missed = [e for e in evaluated if e.verdict.status == "rejected_policy"]
+    if policy_missed:
+        hints.append(
+            f"{len(policy_missed)} candidate(s) were feasible but violated the "
+            "deadline/budget — relaxing objective constraints would unlock them"
+        )
+    if not hints:
+        hints.append("no candidate came close — this workload is outside the supported envelope")
+    return "; ".join(hints)
+
+
+def _fmt(value: float | None, unit: str) -> str:
+    if value is None:
+        return "unknown"
+    return f"${value:.2f}" if unit == "usd" else f"{value:.0f} {unit}"

flashruntime/planner/timecost.py

@@ -0,0 +1,81 @@
+"""Static time and cost estimation.
+
+Transformer training throughput uses the standard FLOPs identity:
+    FLOPs per token (forward + backward) ≈ 6 · parameters
+    tokens/s/GPU ≈ MFU · peak_FLOPS / (6 · parameters)
+with an assumed MFU (catalog.ASSUMED_MFU). Multiplied by GPU count and
+scaling efficiency for the fleet rate; penalized for slow strategy knobs.
+LoRA/QLoRA save optimizer *memory*, not backward FLOPs — activation
+gradients still flow through every frozen layer, so 6·P stands.
+
+Cost = wall-clock × GPUs × hourly rate (when the user provided one).
+Everything here is `basis: static`; a profiling run replaces these numbers
+with measurements, and honest absence beats invented precision — when
+inputs are missing, estimates are omitted, never guessed silently.
+"""
+
+from __future__ import annotations
+
+from flashruntime.planner.catalog import ASSUMED_MFU
+from flashruntime.protocol.plan_v1alpha1 import Estimate
+
+# Multiplicative step-time penalties for strategy knobs. [assumption]
+PENALTY_ACT_CKPT = 1.30  # recompute in backward
+PENALTY_QLORA = 1.35  # dequantization overhead
+PENALTY_CPU_OFFLOAD = 4.0  # optimizer step across PCIe — order-of-magnitude, profile it
+
+
+def tokens_per_second_per_gpu(params: float, gpu_tflops: float) -> float:
+    """Static single-GPU throughput from the 6·P FLOPs identity."""
+    if gpu_tflops <= 0 or params <= 0:
+        return 0.0
+    return ASSUMED_MFU * gpu_tflops * 1e12 / (6.0 * params)
+
+
+def compute_time_per_micro_step_s(
+    params: float, seq_len: int, micro_batch: int, gpu_tflops: float
+) -> float:
+    """Seconds of pure compute for one micro-batch on one GPU (feeds the
+    communication overlap model)."""
+    tps = tokens_per_second_per_gpu(params, gpu_tflops)
+    if tps <= 0:
+        return 0.0
+    return seq_len * micro_batch / tps
+
+
+def training_time_minutes(
+    params: float,
+    train_tokens: float,
+    gpus: int,
+    gpu_tflops: float,
+    scaling_efficiency: float,
+    *,
+    activation_checkpointing: bool = False,
+    qlora: bool = False,
+    cpu_offload: bool = False,
+) -> Estimate | None:
+    """Wall-clock estimate for the whole run, or None when inputs are missing."""
+    tps = tokens_per_second_per_gpu(params, gpu_tflops)
+    if tps <= 0 or train_tokens <= 0:
+        return None
+    fleet = tps * max(1, gpus) * max(0.01, scaling_efficiency)
+    penalty = 1.0
+    notes = [f"MFU {ASSUMED_MFU} assumed"]
+    if activation_checkpointing:
+        penalty *= PENALTY_ACT_CKPT
+        notes.append("activation-checkpointing recompute +30%")
+    if qlora:
+        penalty *= PENALTY_QLORA
+        notes.append("QLoRA dequantization +35%")
+    if cpu_offload:
+        penalty *= PENALTY_CPU_OFFLOAD
+        notes.append("CPU-offload optimizer ×4 — profile before trusting")
+    minutes = train_tokens * penalty / fleet / 60.0
+    return Estimate(value=round(minutes, 1), unit="min", basis="static", note="; ".join(notes))
+
+
+def cost_usd(time_min: Estimate | None, gpus: int, hourly_per_gpu: float | None) -> Estimate | None:
+    if time_min is None or hourly_per_gpu is None:
+        return None
+    usd = time_min.value / 60.0 * max(1, gpus) * hourly_per_gpu
+    return Estimate(value=round(usd, 2), unit="usd", basis="static", note=time_min.note)

flashruntime/profiling/__init__.py

@@ -0,0 +1,113 @@
+"""Profiling: measured numbers that replace the planner's assumptions.
+
+The planner's tier-2 upgrade (FLASHRUNTIME_EVALUATION §F): run a few real
+steps of the real workload and replace the `[assumption]` constants in
+`planner/catalog.py` (activation bytes, MFU, offload penalty) with
+measurements. A `ProfileResult` is born with `basis="profiled"` — the
+whole honesty ladder (`static → profiled → ledger`) hangs on never faking
+that field.
+
+Isolation invariants (every implementation MUST honor all four — they are
+what makes profiling safe to run against production jobs):
+1. Separate namespace: profile runs use their own run-id/prefix; they
+   never write under a real job's artifact prefix.
+2. No commits: a profile run never creates ArtifactRecords, never commits
+   checkpoints to the catalog, never counts as accepted work.
+3. Independent RNG: profiling must not consume the training run's seed
+   stream (the sgd_trainer's step-indexed batching makes this trivial —
+   keep that property in future trainers).
+4. Bounded cost: hard wall-clock + step budgets; a hung profile is
+   cancelled, and "no measurement" is reported honestly rather than a
+   fabricated number.
+
+Measurement protocol (defaults from the evaluation): `warmup_steps=3`
+(skip compile/cache effects), `measure_steps=20`, peak memory via the
+framework's allocator stats plus process RSS, one checkpoint save/restore
+cycle — which doubles as free validation of the checkpoint contract
+before the real run depends on it.
+
+Status: interface complete (final surface); first concrete implementation
+targets the LoRA recipe (SPRINT_PLAN Days 8–10) — profile locally via the
+subprocess runner before wiring cluster profiling.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Literal, Protocol
+
+from pydantic import BaseModel, Field
+
+from flashruntime.protocol.plan_v1alpha1 import PlanRequest, StrategyPlan
+
+__all__ = ["ProfileResult", "ProfileCache", "Profiler", "ProfileError"]
+
+
+class ProfileError(Exception):
+    """The profile run could not produce trustworthy numbers (crashed,
+    exceeded budget, measured nonsense). Callers fall back to static
+    estimates — with the plan still labeled `basis: static`, never
+    upgraded on a failed profile."""
+
+
+class ProfileResult(BaseModel):
+    """What a profile run measured. Every field is optional except the
+    bookkeeping — a partial profile is still useful (peak memory without
+    throughput beats assumptions), but absent numbers stay absent."""
+
+    basis: Literal["profiled"] = "profiled"
+    peak_vram_gb: float | None = Field(default=None, ge=0)
+    host_ram_gb: float | None = Field(default=None, ge=0)
+    step_time_s: float | None = Field(default=None, gt=0)
+    tokens_per_s: float | None = Field(default=None, gt=0)
+    dataloader_wait_fraction: float | None = Field(default=None, ge=0, le=1)
+    checkpoint_save_s: float | None = Field(default=None, ge=0)
+    checkpoint_restore_s: float | None = Field(default=None, ge=0)
+    checkpoint_size_gb: float | None = Field(default=None, ge=0)
+    measured_steps: int = Field(ge=1)
+    warmup_steps: int = Field(ge=0)
+    notes: list[str] = Field(default_factory=list)
+
+
+class ProfileCache(Protocol):
+    """Profiles are expensive; identical questions must be answered once.
+
+    The cache key is the tuple that actually determines the numbers:
+    (model digest, strategy family + knobs, GPU class, framework versions,
+    seq-len bucket, batch bucket) — see FLASHRUNTIME_EVALUATION §F. Real
+    runs later back-feed the same table with `basis: ledger` at higher
+    trust; the cache interface stays identical.
+    """
+
+    def key(self, request: PlanRequest, plan: StrategyPlan) -> str: ...
+
+    def get(self, key: str) -> ProfileResult | None: ...
+
+    def put(self, key: str, result: ProfileResult) -> None: ...
+
+
+class Profiler(ABC):
+    """Run a bounded, isolated measurement of one candidate plan.
+
+    Skip policy lives in the *caller* (the planner), not here: the planner
+    skips profiling when a cache hit exists, when static margins are
+    comfortable (≤ 0.6·VRAM), or when the profile would cost more than
+    2–5% of the job budget. The profiler's only judgment is "did I measure
+    something trustworthy" — everything else is planning.
+    """
+
+    @abstractmethod
+    def profile(
+        self,
+        request: PlanRequest,
+        plan: StrategyPlan,
+        budget_seconds: float = 300.0,
+    ) -> ProfileResult:
+        """Execute warmup + measured steps of `plan` for `request`.
+
+        Inputs: the user's request (data/model identity), the candidate
+        plan (strategy + knobs to measure under), a hard wall-clock budget.
+        Output: a ProfileResult (partial allowed; see model notes).
+        Raises: ProfileError when nothing trustworthy was measured.
+        Must honor all four isolation invariants in the module docstring.
+        """

flashruntime/protocol/__init__.py

@@ -0,0 +1,18 @@
+"""Versioned public schemas shared by FlashNode and FlashML Cloud.
+
+Current version: v1alpha1 (`flashruntime.protocol.v1alpha1`) — job
+specification, job state, event vocabulary, artifact records, and node
+registration/heartbeat messages. Every schema carries an explicit version
+field; security-relevant fields fail closed on unknown values.
+
+Planning contract: `flashruntime.protocol.plan_v1alpha1` — PlanRequest
+(workload + resources + objective), StrategyPlan, PlanReport with candidate
+verdicts and explanations.
+
+Future (still scaffolds elsewhere): lease and attempt payloads, checkpoint
+manifests.
+"""
+
+from flashruntime.protocol import plan_v1alpha1, v1alpha1
+
+__all__ = ["v1alpha1", "plan_v1alpha1"]