diff --git a/agents/completed.md b/agents/completed.md index a0094d4..74f07ea 100644 --- a/agents/completed.md +++ b/agents/completed.md @@ -8,6 +8,33 @@ --- +# #377: HF binding files= not honored on every download path (14GB instead of ~3GB) + +**Completed:** yes +**Status:** DONE (2026-07-04, Claude) — root-cause fix: `Executor.__init__` registers every endpoint binding by wire ref on the `ModelStore` (`register_binding`), and `ModelStore.ensure_local` falls back to that index when called with a bare ref — so `files=`/provider now apply on the ModelOp DOWNLOAD/LOAD and lifecycle startup-prefetch paths, not just setup. Also found + fixed the fp32-pull selector bug: `select_hf_files` treated root monolithic checkpoints (sd1.5's untagged `v1-5-pruned*.safetensors`, 12GB) as their own weight group and "untagged" won for that group even though component dirs had fp16 — diffusers-layout repos now exclude root single-file checkpoints entirely. Boundary tests (no network) assert allow_patterns/provider at the download layer: tests/test_executor_prefetch_binding.py, tests/test_hf_selection.py. + +## Metadata +- Category: models +- Status: done + +--- + +# #376: reserved-source materialization — ctx.source_path is never populated + +**Completed:** yes +**Status:** DONE (2026-07-04, Claude) — `Executor._run_job` now extracts the reserved `source`/`destination` structs from producer-kind (non-inference) payloads, populates `ctx.source`/`ctx.destination`, and materializes `payload.source` through the same `ModelStore.ensure_local` path as model bindings (identical retry + RETRYABLE/INVALID classification + ModelEvent emission) before the handler runs — `ctx._set_source_path` finally has a caller. Missing `source` is a no-op (dataset producers like generate_prompt_corpus carry none); present-but-empty ref → INVALID. Deferred: live cast-dtype run against the hub (needs hub infra + GPU; training-endpoints #37 deferral stands). Tests: tests/test_executor_source_materialization.py. + +## Tasks +- [x] Executor pre-invoke hook: detect reserved `source` in payload for producer kinds, ensure_local, set ctx.source_path +- [x] Failure classification + ModelEvent parity with model-binding downloads +- [x] Tests: producer endpoint with source ref gets a real local dir; missing/invalid source → INVALID not crash +- [ ] Verify training-endpoints cast-dtype runs live against the hub (their #37 deferral) — deferred, needs live hub + GPU + +## Acceptance +A transform endpoint (e.g. cast-dtype) invoked with `source: {repo: owner/name:tag}` receives a materialized local snapshot at ctx.source_path. + +--- + # #369: residency registry must own the executor's pipelines — honest vram_bytes, no load races **Completed:** yes diff --git a/agents/progress.md b/agents/progress.md index 4501108..9071c82 100644 --- a/agents/progress.md +++ b/agents/progress.md @@ -11,38 +11,6 @@ next_id: 378 --- -# #377: HF binding files= not honored on every download path (14GB instead of ~3GB) - -**Completed:** no -**Status:** OPEN (2026-07-04, filed from e2e #105 J4) — the sd15-image example binds `HF("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype="fp16", files=("*.json","*.txt","*.fp16.safetensors"))`, yet the worker's cold download fetched the FULL fp32 root checkpoints too (v1-5-pruned.safetensors 7.7GB + v1-5-pruned-emaonly.safetensors 4.3GB) — ~14GB on disk where ~3GB suffices. `Store.ensure_local` threads `allow_patterns=binding.files` on the setup path (executor.py:540→359), but at least one download path reaches `models/download.py` without the binding (ModelOp/prefetch path calls `store.ensure_local(ref, snap)` with no binding, executor.py:753) and falls back to `select_hf_files`' variant selector or the whole repo. Fix: the worker knows the binding for every ref in its own endpoint spec — resolve binding-by-ref once (registry) and apply its files/provider on ALL download paths. Also worth checking why the variant selector (if it ran) still took fp32 root weights when an fp16 group exists. - -## Metadata -- Category: models -- Status: planned - ---- - - - -# #376: reserved-source materialization — ctx.source_path is never populated - -**Completed:** no -**Status:** OPEN — filed 2026-07-04 from training-endpoints #37 rewrite. `RequestContext.source_path` + `_set_source_path()` exist (request_context/__init__.py:770,777) but no caller: grep shows zero call sites in src/ or packages/. training-endpoints transform/quant producers (cast-dtype, convert-gguf, convert-flashpack, repackage-singlefile, torchao/bnb/modelopt quant) read `ctx.source_path` expecting the library to have materialized the payload's reserved `source` repo locally before invoke. Today they get None and fail. - -Executor should, for producer-kind endpoints whose payload carries the reserved `source` struct: resolve the source ref → `ensure_local` (models/download.py, same snapshot path as model bindings) → `ctx._set_source_path(local_dir)` before the handler runs; download failures map to the same RETRYABLE/INVALID classification as model bindings. No new wire messages — the source ref resolves like any tensorhub ref (hub snapshot on RunJob or resolve-at-worker per CONTRACT §5). - -## Tasks -- [ ] Executor pre-invoke hook: detect reserved `source` in payload for producer kinds, ensure_local, set ctx.source_path -- [ ] Failure classification + ModelEvent parity with model-binding downloads -- [ ] Tests: producer endpoint with source ref gets a real local dir; missing/invalid source → INVALID not crash -- [ ] Verify training-endpoints cast-dtype runs live against the hub (their #37 deferral) - -## Acceptance -A transform endpoint (e.g. cast-dtype) invoked with `source: {repo: owner/name:tag}` receives a materialized local snapshot at ctx.source_path. - ---- - - # #362: docs teach a deleted API — every README quickstart raises ImportError **Completed:** no diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index 29aeb98..5078a92 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -80,6 +80,21 @@ def _sanitize(message: str) -> str: return out[:1024] +def _reserved_repo_info(payload: Any, field_name: str) -> Dict[str, Any]: + """``payload.source`` / ``payload.destination`` as a plain dict ({} when + absent). Producer payloads carry these reserved-name structs (#376).""" + obj = getattr(payload, field_name, None) + if obj is None: + return {} + if isinstance(obj, dict): + return dict(obj) + try: + out = msgspec.to_builtins(obj) + except Exception: + return {} + return out if isinstance(out, dict) else {} + + def _map_exception(exc: BaseException) -> Tuple[int, str]: """-> (JobStatus, safe_message).""" if isinstance(exc, (CanceledError, asyncio.CancelledError)): @@ -177,6 +192,7 @@ def __init__( on_event=self._on_residency_event, vram_budget_bytes=vram_budget_bytes, ) self._locks: Dict[str, asyncio.Lock] = {} + self._bindings: Dict[str, Any] = {} self.keep: set[str] = set() self._loop: Optional[asyncio.AbstractEventLoop] = None self._index = disk_gc.RefIndex(self._cache_dir) @@ -305,6 +321,12 @@ async def _ensure_disk_headroom(self, ref: str, needed_bytes: int) -> None: def _lock(self, ref: str) -> asyncio.Lock: return self._locks.setdefault(ref, asyncio.Lock()) + def register_binding(self, ref: str, binding: Any) -> None: + """Endpoint-spec binding for ``ref`` — supplies files/provider on + download paths that only carry the bare ref (ModelOp, startup + prefetch), so ``files=`` selections apply everywhere (#377).""" + self._bindings.setdefault(ref, binding) + async def ensure_local( self, ref: str, @@ -314,8 +336,11 @@ async def ensure_local( ) -> Path: """Materialize `ref` on disk. Transient failures retry with backoff; terminal (4xx-class) failures raise immediately. Emits ModelEvents. - ``binding`` (when known) supplies provider + file-selection metadata.""" + ``binding`` (when known) supplies provider + file-selection metadata; + bare-ref callers fall back to the registered endpoint binding.""" self.bind_loop() + if binding is None: + binding = self._bindings.get(ref) async with self._lock(ref): cached = self.residency.local_path(ref) if cached is not None and cached.exists(): @@ -434,6 +459,9 @@ def __init__( self._send = send self._settings = settings self.store = store or ModelStore(send) + for s in specs: + for b in s.models.values(): + self.store.register_binding(wire_ref(b), b) self._gpu_semaphore = asyncio.Semaphore(max(1, gpu_slots)) # Model loads/promotions serialize so allocator-delta measurements # and free-VRAM reads don't cross-contaminate (#369). @@ -934,6 +962,10 @@ async def _run_job(self, job: _Job, run: pb.RunJob) -> None: gpu_index = int(compute.gpu_index) if compute is not None else 0 timeout_ms = int(run.timeout_ms or 0) or int(spec.timeout_ms or 0) + producer = spec.kind != "inference" + source_info = _reserved_repo_info(payload, "source") if producer else {} + destination_info = _reserved_repo_info(payload, "destination") if producer else {} + ctx_cls = _CONTEXT_BY_KIND.get(spec.kind, RequestContext) ctx = ctx_cls( request_id=run.request_id, @@ -949,6 +981,8 @@ async def _run_job(self, job: _Job, run: pb.RunJob) -> None: gpu_count=int(compute.gpu_count) if compute is not None else 0, ), models={b.slot: b.ref for b in run.models}, + source_info=source_info, + destination_info=destination_info, execution_hints=( {"output_format": "inline"} if run.output_mode == pb.OUTPUT_MODE_INLINE else {} ), @@ -957,6 +991,8 @@ async def _run_job(self, job: _Job, run: pb.RunJob) -> None: job.ctx = ctx try: + if source_info: + await self._materialize_source(ctx, source_info, snapshots) instance = await self.ensure_setup(spec, snapshots) kwargs = await self._handler_kwargs(spec, snapshots) except asyncio.CancelledError: @@ -1017,6 +1053,18 @@ async def _run_job(self, job: _Job, run: pb.RunJob) -> None: self._gpu_semaphore.release() self._maybe_idle() + async def _materialize_source( + self, ctx: Any, info: Dict[str, Any], snapshots: Dict[str, pb.Snapshot] + ) -> None: + """Reserved-source contract (#376): materialize ``payload.source`` + locally before the handler runs. Same ModelStore path as model + bindings — identical retry/classification and ModelEvent emission.""" + ref = str(info.get("ref") or "").strip() + if not ref: + raise ValidationError("payload.source.ref must be a non-empty repo ref") + path = await self.store.ensure_local(ref, snapshots.get(ref)) + ctx._set_source_path(str(path)) + async def _handler_kwargs( self, spec: EndpointSpec, snapshots: Dict[str, pb.Snapshot] ) -> Dict[str, Any]: diff --git a/src/gen_worker/models/download.py b/src/gen_worker/models/download.py index 2c8a35a..22a4e1f 100644 --- a/src/gen_worker/models/download.py +++ b/src/gen_worker/models/download.py @@ -267,8 +267,9 @@ def select_hf_files( """Small variant selector: which repo files to download. - Diffusers-style repos (component dirs): every non-weight file (configs, - tokenizers) plus ONE weight set per directory — the requested ``flavor`` - when present, else bf16 > fp16 > untagged, safetensors preferred. + tokenizers) plus ONE weight set per component directory — the requested + ``flavor`` when present, else bf16 > fp16 > untagged, safetensors + preferred. Root monolithic checkpoints are excluded (redundant). - Root-weights repos: the root weight files (same variant rule) + sidecars. - Anything else: None (download the whole repo). """ @@ -295,6 +296,12 @@ def select_hf_files( if not diffusers_like and set(weights_by_dir) != {""}: return None # unrecognized layout: full repo + if diffusers_like and "" in weights_by_dir and len(weights_by_dir) > 1: + # Root single-file distributions (sd1.5's v1-5-pruned*.safetensors) + # duplicate the component-dir weights — diffusers loads from the + # component dirs, so never pull the monolithic root checkpoints. + del weights_by_dir[""] + def _pick(group: list[str]) -> list[str]: st = [f for f in group if f.lower().endswith(".safetensors")] pool = st or group diff --git a/tests/test_executor_prefetch_binding.py b/tests/test_executor_prefetch_binding.py new file mode 100644 index 0000000..6ba0e5b --- /dev/null +++ b/tests/test_executor_prefetch_binding.py @@ -0,0 +1,117 @@ +"""#377: binding files=/provider must apply on EVERY download path. + +The ModelOp/startup-prefetch paths only carry the bare ref; the executor +registers each endpoint binding by wire ref so ModelStore.ensure_local can +recover files/provider at the download-layer boundary (no network).""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any, Dict, List + +import msgspec + +import gen_worker.executor as executor_mod +from gen_worker.api.binding import HF +from gen_worker.executor import Executor +from gen_worker.pb import worker_scheduler_pb2 as pb +from gen_worker.registry import EndpointSpec + + +class _In(msgspec.Struct): + x: str + + +class _Out(msgspec.Struct): + y: str + + +_BINDING = HF( + "stable-diffusion-v1-5/stable-diffusion-v1-5", + dtype="fp16", + files=("*.json", "*.txt", "*.fp16.safetensors"), +) +_REF = "stable-diffusion-v1-5/stable-diffusion-v1-5" + + +def _spec() -> EndpointSpec: + class Endpoint: + def setup(self, model: str) -> None: # pragma: no cover + pass + + def run(self, ctx, payload: _In) -> _Out: # pragma: no cover + return _Out(y=payload.x) + + return EndpointSpec( + name="sd15", method=Endpoint.run, kind="inference", + payload_type=_In, output_mode="single", cls=Endpoint, + attr_name="run", models={"model": _BINDING}, + ) + + +async def _noop_send(msg: pb.WorkerMessage) -> None: + pass + + +def _capture_download(monkeypatch, tmp_path: Path) -> List[Dict[str, Any]]: + calls: List[Dict[str, Any]] = [] + + async def _fake_ensure_local(ref: str, **kwargs: Any) -> Path: + calls.append({"ref": ref, **kwargs}) + return tmp_path + + monkeypatch.setattr(executor_mod, "ensure_local", _fake_ensure_local) + return calls + + +def test_model_op_download_applies_binding_files(monkeypatch, tmp_path) -> None: + calls = _capture_download(monkeypatch, tmp_path) + + async def _run() -> None: + ex = Executor([_spec()], _noop_send) + await ex.handle_model_op(pb.ModelOp(op=pb.MODEL_OP_KIND_DOWNLOAD, ref=_REF)) + + asyncio.run(_run()) + assert len(calls) == 1 + assert calls[0]["ref"] == _REF + assert calls[0]["provider"] == "hf" + assert calls[0]["allow_patterns"] == _BINDING.files + + +def test_bare_ref_store_ensure_local_applies_binding_files(monkeypatch, tmp_path) -> None: + # The lifecycle startup-prefetch shape: store.ensure_local(ref) with + # nothing else. + calls = _capture_download(monkeypatch, tmp_path) + + async def _run() -> None: + ex = Executor([_spec()], _noop_send) + await ex.store.ensure_local(_REF) + + asyncio.run(_run()) + assert calls[0]["provider"] == "hf" + assert calls[0]["allow_patterns"] == _BINDING.files + + +def test_explicit_binding_still_wins(monkeypatch, tmp_path) -> None: + calls = _capture_download(monkeypatch, tmp_path) + override = HF(_REF, files=("only-this.safetensors",)) + + async def _run() -> None: + ex = Executor([_spec()], _noop_send) + await ex.store.ensure_local(_REF, binding=override) + + asyncio.run(_run()) + assert calls[0]["allow_patterns"] == ("only-this.safetensors",) + + +def test_unknown_ref_downloads_without_patterns(monkeypatch, tmp_path) -> None: + calls = _capture_download(monkeypatch, tmp_path) + + async def _run() -> None: + ex = Executor([_spec()], _noop_send) + await ex.store.ensure_local("acme/unbound") + + asyncio.run(_run()) + assert calls[0]["provider"] is None + assert calls[0]["allow_patterns"] == () diff --git a/tests/test_executor_source_materialization.py b/tests/test_executor_source_materialization.py new file mode 100644 index 0000000..e17c546 --- /dev/null +++ b/tests/test_executor_source_materialization.py @@ -0,0 +1,162 @@ +"""#376: reserved-source materialization. + +Producer-kind (conversion/dataset/training) payloads carrying the reserved +``source`` struct get a materialized local snapshot at ``ctx.source_path`` +(and ``ctx.source`` metadata) before the handler runs. Failures classify like +model-binding downloads; missing source is a no-op.""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import List, Optional + +import msgspec + +from gen_worker.api.errors import RetryableError +from gen_worker.api.types import SourceRepo +from gen_worker.executor import Executor +from gen_worker.pb import worker_scheduler_pb2 as pb +from gen_worker.registry import EndpointSpec + + +class _ConvIn(msgspec.Struct): + source: Optional[SourceRepo] = None + dtype: str = "bf16" + + +class _Out(msgspec.Struct): + source_path: str + source_ref: str + + +def _convert(ctx, payload: _ConvIn) -> _Out: + return _Out( + source_path=str(ctx.source_path or ""), + source_ref=str((ctx.source or {}).get("ref") or ""), + ) + + +def _spec() -> EndpointSpec: + return EndpointSpec( + name="cast-dtype", method=_convert, kind="conversion", + payload_type=_ConvIn, output_mode="single", + ) + + +class _Harness: + def __init__(self, tmp_path: Path, *, fail_with: Optional[BaseException] = None) -> None: + self.sent: List[pb.WorkerMessage] = [] + self.ensured: List[str] = [] + self.tmp_path = tmp_path + self.fail_with = fail_with + + async def _send(msg: pb.WorkerMessage) -> None: + self.sent.append(msg) + + self.executor = Executor([_spec()], _send) + + async def _fake_ensure_local(ref, snapshot=None, *, binding=None) -> Path: + self.ensured.append(ref) + if self.fail_with is not None: + raise self.fail_with + return self.tmp_path + + self.executor.store.ensure_local = _fake_ensure_local # type: ignore[method-assign] + + async def run(self, payload: _ConvIn) -> pb.JobResult: + await self.executor.handle_run_job(pb.RunJob( + request_id="r1", attempt=1, function_name="cast-dtype", + input_payload=msgspec.msgpack.encode(payload))) + job = self.executor.jobs[("r1", 1)] + assert job.task is not None + await job.task + results = [m.job_result for m in self.sent if m.WhichOneof("msg") == "job_result"] + assert results, f"no job_result; sent={self.sent}" + return results[-1] + + +def test_source_materialized_before_handler(tmp_path) -> None: + async def _run() -> None: + h = _Harness(tmp_path) + res = await h.run(_ConvIn(source=SourceRepo( + ref="acme/base-model:prod", attributes={"dtype": "fp32"}))) + assert res.status == pb.JOB_STATUS_OK + out = msgspec.msgpack.decode(res.inline, type=_Out) + assert out.source_path == str(tmp_path) + assert out.source_ref == "acme/base-model:prod" + assert h.ensured == ["acme/base-model:prod"] + + asyncio.run(_run()) + + +def test_missing_source_is_noop(tmp_path) -> None: + async def _run() -> None: + h = _Harness(tmp_path) + res = await h.run(_ConvIn()) + assert res.status == pb.JOB_STATUS_OK + out = msgspec.msgpack.decode(res.inline, type=_Out) + assert out.source_path == "" + assert h.ensured == [] + + asyncio.run(_run()) + + +def test_empty_source_ref_is_invalid(tmp_path) -> None: + async def _run() -> None: + h = _Harness(tmp_path) + res = await h.run(_ConvIn(source=SourceRepo(ref=" "))) + assert res.status == pb.JOB_STATUS_INVALID + assert h.ensured == [] + + asyncio.run(_run()) + + +def test_source_download_failure_classifies_like_model_bindings(tmp_path) -> None: + async def _run() -> None: + h = _Harness(tmp_path, fail_with=RetryableError("snapshot not provided")) + res = await h.run(_ConvIn(source=SourceRepo(ref="acme/base-model:prod"))) + assert res.status == pb.JOB_STATUS_RETRYABLE + + h2 = _Harness(tmp_path, fail_with=ValueError("unsupported model ref")) + res2 = await h2.run(_ConvIn(source=SourceRepo(ref="not-a-ref"))) + assert res2.status == pb.JOB_STATUS_INVALID + + asyncio.run(_run()) + + +def test_inference_kind_ignores_reserved_source(tmp_path) -> None: + class _InfOut(msgspec.Struct): + ok: bool + + def _infer(ctx, payload: _ConvIn) -> _InfOut: + return _InfOut(ok=not hasattr(ctx, "source_path")) + + spec = EndpointSpec( + name="infer", method=_infer, kind="inference", + payload_type=_ConvIn, output_mode="single", + ) + + async def _run() -> None: + sent: List[pb.WorkerMessage] = [] + + async def _send(msg: pb.WorkerMessage) -> None: + sent.append(msg) + + ex = Executor([spec], _send) + + async def _boom(ref, snapshot=None, *, binding=None) -> Path: + raise AssertionError("inference must not materialize source") + + ex.store.ensure_local = _boom # type: ignore[method-assign] + await ex.handle_run_job(pb.RunJob( + request_id="r1", attempt=1, function_name="infer", + input_payload=msgspec.msgpack.encode( + _ConvIn(source=SourceRepo(ref="acme/base-model:prod"))))) + job = ex.jobs[("r1", 1)] + assert job.task is not None + await job.task + results = [m.job_result for m in sent if m.WhichOneof("msg") == "job_result"] + assert results[-1].status == pb.JOB_STATUS_OK + + asyncio.run(_run()) diff --git a/tests/test_hf_selection.py b/tests/test_hf_selection.py index d303c39..e8a4de6 100644 --- a/tests/test_hf_selection.py +++ b/tests/test_hf_selection.py @@ -65,6 +65,40 @@ def test_safetensors_preferred_over_bin() -> None: assert "unet/diffusion_pytorch_model.bin" not in sel +def test_diffusers_repo_excludes_root_monolithic_checkpoints() -> None: + # sd1.5 shape (e2e #105 J4): untagged fp32 root checkpoints coexist with + # fp16 component weights. The root single-file distributions are redundant + # — picking them pulled 12GB of fp32 where ~2GB of fp16 suffices. + files = [ + "model_index.json", + "README.md", + "v1-5-pruned.ckpt", + "v1-5-pruned.safetensors", + "v1-5-pruned-emaonly.ckpt", + "v1-5-pruned-emaonly.safetensors", + "scheduler/scheduler_config.json", + "text_encoder/config.json", + "text_encoder/model.safetensors", + "text_encoder/model.fp16.safetensors", + "unet/config.json", + "unet/diffusion_pytorch_model.safetensors", + "unet/diffusion_pytorch_model.fp16.safetensors", + "unet/diffusion_pytorch_model.non_ema.safetensors", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", + "vae/diffusion_pytorch_model.fp16.safetensors", + ] + sel = select_hf_files(files) + assert sel is not None + assert "v1-5-pruned.safetensors" not in sel + assert "v1-5-pruned-emaonly.safetensors" not in sel + assert "v1-5-pruned.ckpt" not in sel + assert "unet/diffusion_pytorch_model.fp16.safetensors" in sel + assert "vae/diffusion_pytorch_model.fp16.safetensors" in sel + assert "text_encoder/model.fp16.safetensors" in sel + assert "model_index.json" in sel + + def test_root_weights_repo_selects_weights_and_sidecars() -> None: files = [ "config.json",