From 58223c0c41812ff3b1ce01d2fac7ca192550378d Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 16 Jun 2026 19:13:36 -0700 Subject: [PATCH] [feat] Register cross-model workflows as first-class named servables MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Answers "what's the right way to name/register custom pipelines like T2I→I2V": treat a Workflow like a card — a stable namespaced id in the same servable namespace, declared dependencies, and a two-level registry. No new concepts; mirrors how cards are registered (and vllm-omni's pipeline_registry). - program/workflow.py: Workflow gains `requires` (the cards it composes, derived from stages) and `validate(engine)` (fail-fast if a required card is absent, P7). New WorkflowRegistry: declarative workflow_id -> builder catalog for out-of-tree/ad hoc use. - runtime/engine.py: `_workflows` registry + register_workflow (validates deps, rejects id collision with a model_id) + serves(); engine.run routes a request whose model_id is a workflow to workflow.run — addressable exactly like a model. Single-model hot path untouched. - runtime/async_engine.py + serving/server.py: serves() and /models include workflows (discoverable as servables). - models/__init__.py: declarative `_WORKFLOWS` catalog (the cross-model analog of _BUILDERS) + register_workflows() helper; build_image_video_engine now registers the workflow too. Adding a custom pipeline = one catalog line. - Naming convention: dotted/namespaced workflow_id (`image_video.t2i_i2v`), distinct from kebab model ids, collision-checked. Renamed from `t2i_then_i2v`. - tests (+5, 12 total in the file): addressable by id, requires/validate, id collision, registry catalog, register_workflows helper. Full suite 114 passed. - designv4 §9.6: the naming & registration convention documented. --- designv4.md | 33 +++++++++++++++-- v2/models/__init__.py | 41 ++++++++++++++++----- v2/models/image_video/program.py | 2 +- v2/program/__init__.py | 5 ++- v2/program/workflow.py | 61 ++++++++++++++++++++++++++++++- v2/runtime/async_engine.py | 6 ++- v2/runtime/engine.py | 17 +++++++++ v2/serving/server.py | 3 +- v2/tests/test_t2i_i2v_workflow.py | 53 ++++++++++++++++++++++++++- 9 files changed, 199 insertions(+), 22 deletions(-) diff --git a/designv4.md b/designv4.md index 32cec5032..2b690a9bf 100644 --- a/designv4.md +++ b/designv4.md @@ -5,7 +5,7 @@ > together with **what was actually built and tested** in the `v2/` package, and with the one workload > that most stressed the design — **UniRL/PromptRL-style joint LM + generator reinforcement learning** > (arXiv 2510.17937; PromptRL, arXiv 2602.01382). v4 is not aspirational. Every structural claim below -> is backed by code in `v2/` and a passing test (109 tests, 17 files, two independent runners — `pytest` +> is backed by code in `v2/` and a passing test (114 tests, 17 files, two independent runners — `pytest` > and a zero-dependency `v2/run_tests.py`). The neural forwards are toy numpy stand-ins (no GPU/torch in > this environment, §12); the **control flow, contracts, scheduling, caching, parity gates, and training > math are real**, which is the whole point of the (recipe, runtime) separation: swap the component @@ -370,8 +370,33 @@ threaded stage→stage. The `image_video` package ships two distinct cards (`flu `ImagePart`; stage 2's program folds the image into the I2V conditioning (a program node — `WanDenoiseLoop` itself is unchanged) and denoises a video. Crucially, crossing instances is a *Workflow boundary*, not a program loop step, so **each model keeps its own interleave-parity guarantee** (proven per-stage). Cost on -the per-step scheduling path: zero. ✅ `v2/tests/test_t2i_i2v_workflow.py` (7 tests: the video provably -depends on the generated image; two distinct instances; each stage standalone + parity-gated). +the per-step scheduling path: zero. ✅ `v2/tests/test_t2i_i2v_workflow.py` (12 tests: the video provably +depends on the generated image; two distinct instances; each stage standalone + parity-gated; plus the +naming/registration tests below). + +**Naming & registering custom pipelines** (asked directly). A workflow is a *first-class servable*, named +and registered with the **same discipline as a card** — two levels, mirroring how cards work: + + * **Name = `workflow_id`, in the same namespace as `model_id`.** A workflow is addressed exactly like a + model: `engine.run(request)` with `request.model_id == "image_video.t2i_i2v"` routes to the workflow; + `engine.serves(name)` reports it; the OpenAI `/models` list includes it. Convention: a **dotted, + namespaced** id — `.` (`image_video.t2i_i2v`) — so it never collides with a card id + (kebab, e.g. `flux-t2i`); `register_workflow` rejects a collision outright. + * **Declared dependencies, validated.** A `Workflow.requires` is the list of cards it composes (derived + from its stages — `["flux-t2i", "wan-i2v"]`), and `register_workflow` calls `workflow.validate(engine)` + to fail fast if any required card is absent (design.md P7: declared, never inferred). + * **Two-level registry, exactly like cards.** A *declarative catalog* — `_WORKFLOWS` in + `models/__init__.py` (the cross-model analog of `_BUILDERS`/`_OMNI_BUILDERS`, and of vllm-omni's + `pipeline_registry`) — maps `workflow_id → (builder, required-card-builders)`; **adding a custom + pipeline is one line there.** `register_workflows(engine, only=[...])` brings up the cards and the + workflow together. A live engine then holds validated `Workflow`s in `engine._workflows`, just as it + holds instances in `_registry`. A standalone `WorkflowRegistry` class is available for out-of-tree/ad + hoc catalogs (the `register_pipeline`-style escape hatch). + +So "more custom ones" is: write the stage builders, add one `_WORKFLOWS` entry with a namespaced id, done — +it is then addressable, discoverable, and dependency-checked like any model. (Async/streaming dispatch of a +workflow servable through `AsyncEngine.submit` is the one piece not yet wired — the offline `engine.run` +path and `serves()`/listing are; it is the same routing hook, noted as the remaining step.) ### 9.7 N-way joint RL — arbitrary experts (asked: possible, or a rewrite?) @@ -464,7 +489,7 @@ v2/ methods/{finetune,dmd2,diffusion_nft,self_forcing,unified_rl,joint_multi_rl} serving/ AsyncEngine, pools, DisaggregatedRunner, connectors, OpenAI server deploy/ DeploymentCard, LocalFleet, DynamoWorkerAdapter - tests/ 17 files, 109 tests run via `pytest v2/tests/` OR `python3 v2/run_tests.py` + tests/ 17 files, 114 tests run via `pytest v2/tests/` OR `python3 v2/run_tests.py` ``` **Enforced boundaries:** `card/` imports no product/runtime; `runtime/` executes `card/` loops but defines diff --git a/v2/models/__init__.py b/v2/models/__init__.py index 3cf4604fe..76e034f79 100644 --- a/v2/models/__init__.py +++ b/v2/models/__init__.py @@ -31,7 +31,7 @@ __all__ = [ "build_qwen_omni_card", "build_qwen_omni_program", "build_unified_card", "build_unified_program", "build_flux_t2i_card", "build_flux_t2i_program", "build_wan_i2v_card", "build_wan_i2v_program", - "build_t2i_then_i2v_workflow", + "build_t2i_then_i2v_workflow", "register_workflows", "build_default_engine", "build_omni_engine", "build_unified_engine", "build_image_video_engine", ] @@ -80,19 +80,40 @@ def build_omni_engine(engine: Any = None) -> Any: return eng -def build_image_video_engine(engine: Any = None) -> Any: - """Register the T2I (``flux-t2i``) and I2V (``wan-i2v``) cards on one engine — two *separate* - models chained by ``build_t2i_then_i2v_workflow`` (design_v3 §13 cross-model composition).""" +# Declarative catalog of cross-model workflows: workflow_id -> (builder, required model cards). +# Adding a custom pipeline is one line here (the cross-model analog of _BUILDERS / _OMNI_BUILDERS; +# cf. vllm-omni's pipeline_registry). ``register_workflows`` registers each whose cards are present. +_WORKFLOWS: dict[str, tuple] = { + "image_video.t2i_i2v": (build_t2i_then_i2v_workflow, + [(build_flux_t2i_card, build_flux_t2i_program), + (build_wan_i2v_card, build_wan_i2v_program)]), +} + + +def register_workflows(engine: Any, *, only: list[str] | None = None) -> Any: + """Register catalog workflows (and the cards they require) onto ``engine``. ``only`` selects a + subset by workflow_id; default registers all whose cards aren't yet present.""" from ..cache import CacheManager from ..card import load_card + names = only if only is not None else list(_WORKFLOWS) + for wf_id in names: + build_workflow, card_builders = _WORKFLOWS[wf_id] + for build_card, build_program in card_builders: + card = build_card() + if not engine.serves(card.model_id): + inst = load_card(card, cache_manager=CacheManager.from_card(card)) + engine.register(card.model_id, inst, build_program()) + engine.register_workflow(build_workflow()) + return engine + + +def build_image_video_engine(engine: Any = None) -> Any: + """Register the T2I (``flux-t2i``) and I2V (``wan-i2v``) cards plus the ``image_video.t2i_i2v`` + workflow on one engine — two *separate* models chained by a cross-model workflow, addressable by + its workflow_id like any servable (design_v3 §13).""" from ..runtime import Engine eng = engine if engine is not None else Engine() - for build_card, build_program in [(build_flux_t2i_card, build_flux_t2i_program), - (build_wan_i2v_card, build_wan_i2v_program)]: - card = build_card() - inst = load_card(card, cache_manager=CacheManager.from_card(card)) - eng.register(card.model_id, inst, build_program()) - return eng + return register_workflows(eng, only=["image_video.t2i_i2v"]) def build_unified_engine(engine: Any = None) -> Any: diff --git a/v2/models/image_video/program.py b/v2/models/image_video/program.py index 1fa70b132..bec13a79d 100644 --- a/v2/models/image_video/program.py +++ b/v2/models/image_video/program.py @@ -95,7 +95,7 @@ def build_t2i_then_i2v_workflow(t2i_id: str = "flux-t2i", i2v_id: str = "wan-i2v diffusion=DiffusionParams(num_steps=4, num_frames=81, # → multi-frame video seed=state.get("seed", 0))) - return Workflow("t2i_then_i2v", [ + return Workflow("image_video.t2i_i2v", [ # namespaced id: . (a servable) WorkflowStage(t2i_id, t2i_stage, label="t2i"), WorkflowStage(i2v_id, i2v_stage, label="i2v"), ]) diff --git a/v2/program/__init__.py b/v2/program/__init__.py index 994d297a6..d18f8b010 100644 --- a/v2/program/__init__.py +++ b/v2/program/__init__.py @@ -13,7 +13,8 @@ from .specs import ( when_opt, when_task, ) -from .workflow import Workflow, WorkflowStage +from .workflow import Workflow, WorkflowRegistry, WorkflowStage __all__ = ["Program", "ProgramKind", "ProgramNode", "ComponentNode", "ModelLoopNode", - "Edge", "EdgeKind", "always", "when_task", "when_opt", "Workflow", "WorkflowStage"] + "Edge", "EdgeKind", "always", "when_task", "when_opt", + "Workflow", "WorkflowStage", "WorkflowRegistry"] diff --git a/v2/program/workflow.py b/v2/program/workflow.py index 8e4f3db58..635659179 100644 --- a/v2/program/workflow.py +++ b/v2/program/workflow.py @@ -34,11 +34,44 @@ class WorkflowStage: label: str = "" +def _engine_serves(engine: Any, model_id: str) -> bool: + if hasattr(engine, "serves"): + return engine.serves(model_id) # AsyncEngine + return model_id in getattr(engine, "_registry", {}) # Engine + + @dataclass class Workflow: + """A named, registrable cross-model pipeline. ``workflow_id`` is its servable name — it lives in + the SAME namespace as a card's ``model_id`` (so the engine, server, and fleet route to it the same + way), and by convention is dotted/namespaced to avoid colliding with a card id (e.g. + ``"image_video.t2i_i2v"``). ``requires`` is the set of cards it composes — declared, validated, + never inferred (design.md P7).""" workflow_id: str stages: list[WorkflowStage] = field(default_factory=list) + @property + def requires(self) -> list[str]: + """The model_ids this workflow composes (deduped, in first-use order). The engine must serve + all of them before the workflow can run.""" + seen: set[str] = set() + out: list[str] = [] + for s in self.stages: + if s.model_id not in seen: + seen.add(s.model_id) + out.append(s.model_id) + return out + + def validate(self, engine: Any) -> "Workflow": + """Fail-fast: every required card must already be registered on ``engine`` (no silent skips).""" + if not self.stages: + raise ValueError(f"workflow {self.workflow_id!r} has no stages") + missing = [m for m in self.requires if not _engine_serves(engine, m)] + if missing: + raise ValueError(f"workflow {self.workflow_id!r} requires unregistered models {missing} " + f"(register those cards before the workflow)") + return self + def run(self, engine: Any, **initial: Any) -> Any: """Execute the stages in order on ``engine``; return the final stage's Output. @@ -48,7 +81,7 @@ class Workflow: raise ValueError(f"workflow {self.workflow_id!r} has no stages") state: dict[str, Any] = dict(initial) out = None - for i, stage in enumerate(self.stages): + for stage in self.stages: label = stage.label or stage.model_id req = stage.make_request(state) if req.model_id != stage.model_id: @@ -59,3 +92,29 @@ class Workflow: for name, art in out.artifacts.items(): state[f"{label}:{name}"] = art return out + + +class WorkflowRegistry: + """Declarative ``workflow_id → builder`` catalog — the cross-model analog of the card builders in + ``models/__init__.py`` (cf. vllm-omni's ``pipeline_registry``). Adding a custom pipeline is one + ``register`` call; the builder is a zero/kw-arg factory returning a ``Workflow``.""" + + def __init__(self) -> None: + self._builders: dict[str, Any] = {} + + def register(self, workflow_id: str, builder: Any) -> Any: + if workflow_id in self._builders: + raise ValueError(f"workflow {workflow_id!r} already registered") + self._builders[workflow_id] = builder + return builder + + def build(self, workflow_id: str, **kw: Any) -> Workflow: + if workflow_id not in self._builders: + raise KeyError(f"no workflow {workflow_id!r} (have {list(self._builders)})") + return self._builders[workflow_id](**kw) + + def names(self) -> list[str]: + return list(self._builders) + + def __contains__(self, workflow_id: str) -> bool: + return workflow_id in self._builders diff --git a/v2/runtime/async_engine.py b/v2/runtime/async_engine.py index 36e003cd3..992eb8e1a 100644 --- a/v2/runtime/async_engine.py +++ b/v2/runtime/async_engine.py @@ -74,6 +74,9 @@ class AsyncEngine: def register_disaggregated(self, model_id: str, pools: Any, program: Any) -> None: self._disagg[model_id] = (pools, program) + def register_workflow(self, workflow: Any) -> Any: + return self.engine.register_workflow(workflow) + def is_disaggregated(self, model_id: str) -> bool: return model_id in self._disagg @@ -107,7 +110,8 @@ class AsyncEngine: return sum(1 for s in self._states.values() if s == RequestState.WAITING) def serves(self, model_id: str) -> bool: - return model_id in self.engine._registry or model_id in self._disagg + return (model_id in self.engine._registry or model_id in self._disagg + or model_id in self.engine._workflows) # --- the driver: one asyncio task per request --------------------------- # async def _run(self, request: Any) -> None: diff --git a/v2/runtime/engine.py b/v2/runtime/engine.py index a96d5de4c..9b739cc95 100644 --- a/v2/runtime/engine.py +++ b/v2/runtime/engine.py @@ -194,10 +194,23 @@ class Engine: admission: AdmissionController = field(default_factory=AdmissionController) metrics: SchedulerMetrics = field(default_factory=SchedulerMetrics) _registry: dict[str, tuple[Any, Program]] = field(default_factory=dict) + _workflows: dict[str, Any] = field(default_factory=dict) # workflow_id -> validated Workflow def register(self, model_id: str, instance: Any, program: Program) -> None: self._registry[model_id] = (instance, program) + def register_workflow(self, workflow: Any) -> Any: + """Register a cross-model Workflow as a servable, addressable by ``workflow.workflow_id`` in + the SAME namespace as model ids. Validates its required cards are already registered.""" + workflow.validate(self) + if workflow.workflow_id in self._registry: + raise ValueError(f"workflow id {workflow.workflow_id!r} collides with a registered model") + self._workflows[workflow.workflow_id] = workflow + return workflow + + def serves(self, name: str) -> bool: + return name in self._registry or name in self._workflows + def _resolve(self, request: Any) -> tuple[Any, Program]: if request.model_id not in self._registry: raise KeyError(f"no model registered for {request.model_id!r} (have {list(self._registry)})") @@ -209,6 +222,10 @@ class Engine: # --- offline single-request (the VideoGenerator path) -------------------- # def run(self, request: Any) -> Output: + if request.model_id in self._workflows: # a workflow servable → orchestrate stages + wf = self._workflows[request.model_id] + seed = request.diffusion.seed if request.diffusion.seed is not None else 0 + return wf.run(self, prompt=request.prompt(), seed=seed, request=request) r = self._make_runner(request) r.run_to_completion() return r.output() diff --git a/v2/serving/server.py b/v2/serving/server.py index ab8aaf922..9c6038735 100644 --- a/v2/serving/server.py +++ b/v2/serving/server.py @@ -113,7 +113,8 @@ class OmniOpenAIServer: "gpu_seconds": round(adm.gpu_seconds, 6), "by_kind": dict(adm.by_kind)} def _models(self) -> dict: - ids = list(self.engine.engine._registry) + list(self.engine._disagg) + ids = (list(self.engine.engine._registry) + list(self.engine._disagg) + + list(self.engine.engine._workflows)) # workflows are servables too return {"object": "list", "data": [{"id": m, "object": "model"} for m in ids]} # --- images ------------------------------------------------------------- # diff --git a/v2/tests/test_t2i_i2v_workflow.py b/v2/tests/test_t2i_i2v_workflow.py index 927dd38d1..3e3c42cc6 100644 --- a/v2/tests/test_t2i_i2v_workflow.py +++ b/v2/tests/test_t2i_i2v_workflow.py @@ -11,11 +11,12 @@ from __future__ import annotations import numpy as np import pytest -from v2.models import build_image_video_engine, build_t2i_then_i2v_workflow +from v2.models import build_image_video_engine, build_t2i_then_i2v_workflow, register_workflows from v2.parity import assert_interleave_parity -from v2.program import Workflow, WorkflowStage +from v2.program import Workflow, WorkflowRegistry, WorkflowStage from v2.request import DiffusionParams, TaskType, make_request from v2.request.modalpart import ImagePart, TextPart +from v2.runtime import Engine def _i2v(eng, image, *, prompt="a scene", seed=9): @@ -95,3 +96,51 @@ def test_workflow_rejects_model_mismatch(): lambda s: make_request(TaskType.I2V, "wan-i2v", "x"))]) # declares flux-t2i, builds wan-i2v with pytest.raises(ValueError): bad.run(eng, prompt="x") + + +# --- naming + registration: a workflow is a first-class servable ------------------ # + +def test_workflow_is_addressable_by_its_id_like_a_model(): + """The workflow_id is a servable name in the SAME namespace as model ids: serves() reports it and + engine.run routes to it — so the OpenAI server / fleet address it exactly like a card.""" + eng = build_image_video_engine() + assert eng.serves("image_video.t2i_i2v") # registered as a servable + assert eng.serves("flux-t2i") and eng.serves("wan-i2v") # its required cards too + out = eng.run(make_request(TaskType.T2V, "image_video.t2i_i2v", "a fox", + diffusion=DiffusionParams(seed=3))) + assert np.asarray(out.artifacts["video"].frames).shape[1] > 1 + + +def test_workflow_declares_and_validates_required_cards(): + """``requires`` is declared (the stage model_ids); registration fails fast if a card is absent.""" + wf = build_t2i_then_i2v_workflow() + assert wf.requires == ["flux-t2i", "wan-i2v"] + with pytest.raises(ValueError, match="requires unregistered"): + Engine().register_workflow(wf) # no cards registered → fail fast + + +def test_workflow_id_must_not_collide_with_a_model_id(): + eng = build_image_video_engine() + clash = Workflow("flux-t2i", [WorkflowStage("flux-t2i", lambda s: None)]) # id == a card id + with pytest.raises(ValueError, match="collides"): + eng.register_workflow(clash) + + +def test_workflow_registry_catalog(): + reg = WorkflowRegistry() + reg.register("image_video.t2i_i2v", build_t2i_then_i2v_workflow) + assert "image_video.t2i_i2v" in reg and reg.names() == ["image_video.t2i_i2v"] + assert reg.build("image_video.t2i_i2v").requires == ["flux-t2i", "wan-i2v"] + with pytest.raises(ValueError, match="already registered"): + reg.register("image_video.t2i_i2v", build_t2i_then_i2v_workflow) + with pytest.raises(KeyError): + reg.build("nonexistent.workflow") + + +def test_register_workflows_helper_brings_up_cards_and_workflow(): + """The declarative catalog: one call registers the workflow and the cards it needs.""" + eng = register_workflows(Engine(), only=["image_video.t2i_i2v"]) + assert eng.serves("image_video.t2i_i2v") + out = eng.run(make_request(TaskType.T2V, "image_video.t2i_i2v", "a city", + diffusion=DiffusionParams(seed=1))) + assert out.artifacts["video"].frames is not None