diff --git a/components/src/dynamo/experimental/workflow/__init__.py b/components/src/dynamo/experimental/workflow/__init__.py index 822f190845c0..8b040539a2f2 100644 --- a/components/src/dynamo/experimental/workflow/__init__.py +++ b/components/src/dynamo/experimental/workflow/__init__.py @@ -4,7 +4,29 @@ """Experimental inference workflow APIs.""" from dynamo.experimental.workflow.builder import StageHandle, Workflow +from dynamo.experimental.workflow.compiler import DeploymentSpec, compile_workflow +from dynamo.experimental.workflow.endpoint import WorkflowEndpointHandler from dynamo.experimental.workflow.ir import StageIR, WorkflowIR +from dynamo.experimental.workflow.nixl import ( + NixlLeaseRegistry, + NixlTensorCarrier, + NixlTensorFanout, + NixlTensorRef, +) +from dynamo.experimental.workflow.orchestrator import WorkflowOrchestrator +from dynamo.experimental.workflow.plan import ( + ExecutionPlan, + GenerateEndpointBinding, + InlineBinding, + RemoteBinding, +) +from dynamo.experimental.workflow.remote import RemoteStageClient, RemoteStageServer +from dynamo.experimental.workflow.runtime import ( + StageContext, + StageRunner, + TensorCarrier, + WorkflowExecutionError, +) from dynamo.experimental.workflow.types import ( StageContract, ValueRef, @@ -12,11 +34,29 @@ ) __all__ = [ + "DeploymentSpec", "StageContract", "StageHandle", "StageIR", + "StageContext", + "StageRunner", + "TensorCarrier", "ValueRef", "Workflow", "WorkflowIR", + "WorkflowEndpointHandler", + "WorkflowExecutionError", + "WorkflowOrchestrator", "WorkflowValidationError", + "ExecutionPlan", + "GenerateEndpointBinding", + "InlineBinding", + "NixlLeaseRegistry", + "NixlTensorCarrier", + "NixlTensorFanout", + "NixlTensorRef", + "RemoteBinding", + "RemoteStageClient", + "RemoteStageServer", + "compile_workflow", ] diff --git a/components/src/dynamo/experimental/workflow/builder.py b/components/src/dynamo/experimental/workflow/builder.py index ed0488ef708b..72b0fa3cdf89 100644 --- a/components/src/dynamo/experimental/workflow/builder.py +++ b/components/src/dynamo/experimental/workflow/builder.py @@ -127,7 +127,7 @@ def build(self) -> WorkflowIR: return WorkflowIR( name=self._name, - inputs=self._inputs, + inputs=frozenset(self._inputs), stages=tuple(self._stages.values()), outputs=self._outputs, ) diff --git a/components/src/dynamo/experimental/workflow/compiler.py b/components/src/dynamo/experimental/workflow/compiler.py new file mode 100644 index 000000000000..1c72805b9cc8 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/compiler.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lower placement-neutral workflows into physical execution plans.""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Mapping, Optional, Union + +from dynamo.experimental.workflow.builder import Workflow +from dynamo.experimental.workflow.ir import WorkflowIR +from dynamo.experimental.workflow.plan import ( + Binding, + ExecutionPlan, + InlineBinding, + RemoteBinding, +) +from dynamo.experimental.workflow.types import WorkflowValidationError, validate_name + + +@dataclass(frozen=True) +class DeploymentSpec: + """Stage placement requested when compiling one WorkflowIR.""" + + bindings: Mapping[str, Binding] + + def __post_init__(self) -> None: + if not isinstance(self.bindings, Mapping): + raise WorkflowValidationError("deployment bindings must be a mapping") + bindings: dict[str, Binding] = {} + for stage_id, binding in sorted(self.bindings.items()): + validate_name(stage_id, "deployment stage id") + if not isinstance(binding, (InlineBinding, RemoteBinding)): + raise WorkflowValidationError( + f"binding for stage {stage_id!r} uses an unsupported type" + ) + bindings[stage_id] = binding + object.__setattr__(self, "bindings", MappingProxyType(bindings)) + + @staticmethod + def inline(**runner_keys: str) -> "DeploymentSpec": + """Build bindings to runners in the orchestrator process.""" + + return DeploymentSpec( + bindings={ + stage_id: InlineBinding(runner_key) + for stage_id, runner_key in runner_keys.items() + } + ) + + @staticmethod + def remote( + *, tensor_carrier: str | None = None, **endpoint_ids: str + ) -> "DeploymentSpec": + """Build round-robin bindings to discovered Dynamo endpoints.""" + + return DeploymentSpec( + bindings={ + stage_id: RemoteBinding(endpoint_id, tensor_carrier=tensor_carrier) + for stage_id, endpoint_id in endpoint_ids.items() + } + ) + + +def compile_workflow( + workflow: Union[Workflow, WorkflowIR], + deployment: Optional[DeploymentSpec] = None, +) -> ExecutionPlan: + """Compile one logical workflow, defaulting every stage to inline placement.""" + + workflow_ir = workflow.build() if isinstance(workflow, Workflow) else workflow + if not isinstance(workflow_ir, WorkflowIR): + raise TypeError("workflow must be a Workflow or WorkflowIR") + stage_ids = tuple(stage.id for stage in workflow_ir.stages) + if deployment is None: + deployment = DeploymentSpec.inline( + **{stage_id: stage_id for stage_id in stage_ids} + ) + if not isinstance(deployment, DeploymentSpec): + raise TypeError("deployment must use DeploymentSpec") + + expected = set(stage_ids) + actual = set(deployment.bindings) + if actual != expected: + raise WorkflowValidationError( + "deployment bindings differ from workflow stages; " + f"missing={sorted(expected - actual)}, extra={sorted(actual - expected)}" + ) + + return ExecutionPlan(workflow_ir, deployment.bindings) diff --git a/components/src/dynamo/experimental/workflow/dispatcher.py b/components/src/dynamo/experimental/workflow/dispatcher.py new file mode 100644 index 000000000000..77c34971e703 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/dispatcher.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runtime-bound workflow stage dispatch.""" + +from __future__ import annotations + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Any, Protocol, runtime_checkable + +from dynamo.experimental.workflow.generate import GenerateEndpointInvoker +from dynamo.experimental.workflow.nixl import NixlTensorFanout +from dynamo.experimental.workflow.plan import ( + NIXL_CARRIER, + ExecutionPlan, + GenerateEndpointBinding, + InlineBinding, + RemoteBinding, +) +from dynamo.experimental.workflow.remote import NixlCarriedValue, RemoteStageClient +from dynamo.experimental.workflow.runtime import ( + StageContext, + StageRunner, + WorkflowExecutionError, +) +from dynamo.experimental.workflow.types import ( + StageContract, + ValueRef, + WorkflowValidationError, +) + + +@runtime_checkable +class RemoteStageInvoker(Protocol): + """Internal transport boundary used by the dispatcher.""" + + async def run( + self, + stage_id: str, + contract: StageContract, + inputs: Mapping[str, Any], + context: StageContext, + output_transfers: Mapping[str, tuple[str, ...]], + *, + request_context: Any = None, + ) -> Mapping[str, Any]: + ... + + +class StageDispatcher: + """Validate and invoke stages through their compiled bindings.""" + + def __init__( + self, + plan: ExecutionPlan, + inline_runners: Mapping[str, StageRunner], + remote_clients: Mapping[str, Any] = MappingProxyType({}), + ) -> None: + if not isinstance(plan, ExecutionPlan): + raise TypeError("plan must use ExecutionPlan") + if not isinstance(inline_runners, Mapping): + raise TypeError("inline_runners must be a mapping") + if not isinstance(remote_clients, Mapping): + raise TypeError("remote_clients must be a mapping") + + expected_keys = { + binding.runner_key + for binding in plan.bindings.values() + if isinstance(binding, InlineBinding) + } + actual_keys = set(inline_runners) + if actual_keys != expected_keys: + raise WorkflowValidationError( + "inline runners differ from execution plan; " + f"missing={sorted(expected_keys - actual_keys)}, " + f"extra={sorted(actual_keys - expected_keys)}" + ) + + expected_endpoints = { + binding.endpoint_id + for binding in plan.bindings.values() + if isinstance(binding, RemoteBinding) + } + actual_endpoints = set(remote_clients) + if actual_endpoints != expected_endpoints: + raise WorkflowValidationError( + "remote clients differ from execution plan; " + f"missing={sorted(expected_endpoints - actual_endpoints)}, " + f"extra={sorted(actual_endpoints - expected_endpoints)}" + ) + + runners = dict(inline_runners) + for stage_id, contract in plan.stage_contracts.items(): + binding = plan.bindings[stage_id] + if isinstance(binding, RemoteBinding): + continue + runner = runners[binding.runner_key] + if not isinstance(runner, StageRunner): + raise WorkflowValidationError( + f"runner {binding.runner_key!r} must implement StageRunner" + ) + if runner.contract != contract: + raise WorkflowValidationError( + f"runner {binding.runner_key!r} for stage {stage_id!r} " + "does not match its authored contract" + ) + + output_transfers: dict[str, dict[str, list[str]]] = {} + for target in plan.workflow.stages: + target_binding = plan.bindings[target.id] + if not self._nixl_remote(target_binding): + continue + for target_port, source in target.inputs.items(): + if source.stage_id is None: + continue + source_binding = plan.bindings[source.stage_id] + if not self._nixl_remote(source_binding): + continue + source_output = source.output_name + assert source_output is not None + outputs = output_transfers.setdefault(source.stage_id, {}) + outputs.setdefault(source_output, []).append( + f"{target.id}.{target_port}" + ) + + self._plan = plan + self._inline_runners = MappingProxyType(runners) + self._remote_clients = MappingProxyType(dict(remote_clients)) + self._output_transfers = MappingProxyType( + { + stage_id: MappingProxyType( + { + output_name: tuple(sorted(transfer_ids)) + for output_name, transfer_ids in outputs.items() + } + ) + for stage_id, outputs in output_transfers.items() + } + ) + + @staticmethod + def _nixl_remote(binding: Any) -> bool: + return ( + isinstance(binding, RemoteBinding) + and binding.tensor_carrier == NIXL_CARRIER + ) + + @classmethod + async def bind( + cls, + plan: ExecutionPlan, + *, + runtime: Any = None, + inline_runners: Mapping[str, StageRunner] = MappingProxyType({}), + ) -> "StageDispatcher": + """Resolve remote endpoints once and bind all physical stage targets.""" + + endpoint_ids = { + binding.endpoint_id + for binding in plan.bindings.values() + if isinstance(binding, RemoteBinding) + } + if endpoint_ids and runtime is None: + raise WorkflowValidationError( + "runtime is required to bind remote workflow stages" + ) + + clients: dict[str, Any] = {} + for endpoint_id in sorted(endpoint_ids): + bindings = [ + binding + for binding in plan.bindings.values() + if isinstance(binding, RemoteBinding) + and binding.endpoint_id == endpoint_id + ] + protocols = { + GenerateEndpointBinding + if isinstance(binding, GenerateEndpointBinding) + else RemoteBinding + for binding in bindings + } + if len(protocols) != 1: + raise WorkflowValidationError( + f"remote endpoint {endpoint_id!r} cannot mix stage protocols" + ) + endpoint = runtime.endpoint(endpoint_id) + client = await endpoint.client() + await client.wait_for_instances() + clients[endpoint_id] = ( + GenerateEndpointInvoker(client) + if protocols == {GenerateEndpointBinding} + else RemoteStageClient(client) + ) + return cls(plan, inline_runners, clients) + + async def call( + self, + stage_id: str, + inputs: Mapping[str, Any], + context: StageContext, + *, + request_context: Any = None, + ) -> dict[str, Any]: + """Invoke one stage and validate its complete input/output contract.""" + + contract = self._plan.stage_contracts[stage_id] + expected_inputs = set(contract.inputs) + actual_inputs = set(inputs) + if actual_inputs != expected_inputs: + raise WorkflowExecutionError( + f"stage {stage_id!r} inputs differ from its contract; " + f"missing={sorted(expected_inputs - actual_inputs)}, " + f"extra={sorted(actual_inputs - expected_inputs)}" + ) + + binding = self._plan.bindings[stage_id] + frozen_inputs = MappingProxyType(dict(inputs)) + if isinstance(binding, InlineBinding): + result = await self._inline_runners[binding.runner_key].run( + frozen_inputs, context + ) + elif isinstance(binding, GenerateEndpointBinding): + result = await self._remote_clients[binding.endpoint_id].run( + stage_id, + contract, + frozen_inputs, + context, + request_context=request_context, + ) + else: + result = await self._remote_clients[binding.endpoint_id].run( + stage_id, + contract, + frozen_inputs, + context, + self._output_transfers.get(stage_id, MappingProxyType({})), + request_context=request_context, + ) + if not isinstance(result, Mapping): + raise WorkflowExecutionError( + f"stage {stage_id!r} returned a non-mapping result" + ) + expected_outputs = set(contract.outputs) + actual_outputs = set(result) + if actual_outputs != expected_outputs: + raise WorkflowExecutionError( + f"stage {stage_id!r} outputs differ from its contract; " + f"missing={sorted(expected_outputs - actual_outputs)}, " + f"extra={sorted(actual_outputs - expected_outputs)}" + ) + return dict(result) + + async def resolve_edge( + self, + reference: ValueRef, + target_stage: str, + target_port: str, + value: Any, + ) -> Any: + """Select a consumer-specific NIXL reference when a value was carried.""" + + if not isinstance(value, NixlCarriedValue): + return value + transfer_id = f"{target_stage}.{target_port}" + source_stage = reference.stage_id + if source_stage is None or not self._nixl_remote( + self._plan.bindings[source_stage] + ): + raise WorkflowExecutionError( + f"NIXL edge {transfer_id!r} must originate from a NIXL remote stage" + ) + if not self._nixl_remote(self._plan.bindings[target_stage]): + raise WorkflowExecutionError( + f"NIXL edge {transfer_id!r} must target a NIXL remote stage" + ) + reference_value = NixlTensorFanout.from_dict(value.value).for_transfer( + transfer_id + ) + return NixlCarriedValue(reference_value.to_dict()) + + def resolve_workflow_output(self, output_name: str, value: Any) -> Any: + """Reject carried tensors from the current unary workflow boundary.""" + + if isinstance(value, NixlCarriedValue): + raise WorkflowExecutionError( + f"workflow output {output_name!r} cannot expose a NIXL tensor" + ) + return value diff --git a/components/src/dynamo/experimental/workflow/endpoint.py b/components/src/dynamo/experimental/workflow/endpoint.py new file mode 100644 index 000000000000..450f386e9070 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/endpoint.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Serve a bound workflow through a Dynamo model endpoint.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Mapping +from typing import Any + +from dynamo.experimental.workflow.orchestrator import WorkflowOrchestrator +from dynamo.experimental.workflow.types import WorkflowValidationError + + +class WorkflowEndpointHandler: + """Adapt a request/chunk workflow to Dynamo's endpoint handler ABI.""" + + def __init__(self, orchestrator: WorkflowOrchestrator) -> None: + if not isinstance(orchestrator, WorkflowOrchestrator): + raise TypeError("orchestrator must use WorkflowOrchestrator") + _validate_endpoint_abi(orchestrator) + self._orchestrator = orchestrator + + async def generate( + self, request: Mapping[str, Any], context: Any + ) -> AsyncIterator[dict[str, Any]]: + """Execute one workflow attempt and yield its terminal chunk.""" + + if not isinstance(request, Mapping): + raise TypeError("workflow endpoint request must be a mapping") + + attempt_id = context.id() + execution = asyncio.create_task( + self._orchestrator.run( + {"request": request}, + attempt_id=attempt_id, + request_context=context, + ), + name=f"workflow-endpoint:{attempt_id}", + ) + cancellation = asyncio.ensure_future(context.async_killed_or_stopped()) + try: + done, _ = await asyncio.wait( + {execution, cancellation}, return_when=asyncio.FIRST_COMPLETED + ) + if execution not in done: + execution.cancel() + await asyncio.gather(execution, return_exceptions=True) + return + result = await execution + except BaseException: + if not execution.done(): + execution.cancel() + await asyncio.gather(execution, return_exceptions=True) + raise + finally: + if not cancellation.done(): + cancellation.cancel() + await asyncio.gather(cancellation, return_exceptions=True) + + chunk = result["chunk"] + if not isinstance(chunk, Mapping): + raise TypeError("workflow 'chunk' output must be a mapping") + yield dict(chunk) + + +def _validate_endpoint_abi(orchestrator: WorkflowOrchestrator) -> None: + workflow = orchestrator.plan.workflow + if workflow.inputs != frozenset({"request"}): + raise WorkflowValidationError( + "workflow endpoints require exactly one 'request' input" + ) + if set(workflow.outputs) != {"chunk"}: + raise WorkflowValidationError( + "workflow endpoints require exactly one 'chunk' output" + ) diff --git a/components/src/dynamo/experimental/workflow/generate.py b/components/src/dynamo/experimental/workflow/generate.py new file mode 100644 index 000000000000..67d2314cbfe5 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/generate.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Adapter from a workflow stage invocation to Dynamo's Generate endpoint.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Mapping +from typing import Any, Protocol + +from dynamo.experimental.workflow.runtime import StageContext, WorkflowExecutionError +from dynamo.experimental.workflow.types import StageContract + +GENERATE_REQUEST_PORT = "request" +GENERATE_OUTPUT_PORT = "completion" + + +class _DynamoClient(Protocol): + async def round_robin( + self, + request: Mapping[str, Any], + *, + annotated: bool, + context: Any = None, + ) -> AsyncIterator[Any]: + ... + + +class GenerateEndpointStream: + """Own one Generate transport stream and its cancellation context.""" + + def __init__(self, stream: AsyncIterator[Any], transport_context: Any) -> None: + self._stream = stream.__aiter__() + self._transport_context = transport_context + self._closed = False + + def __aiter__(self) -> "GenerateEndpointStream": + return self + + async def __anext__(self) -> Any: + return await self._stream.__anext__() + + async def aclose(self, *, cancel: bool = False) -> None: + if self._closed: + return + self._closed = True + if cancel and self._transport_context is not None: + self._transport_context.stop_generating() + close = getattr(self._stream, "aclose", None) + if callable(close): + await close() + + +class GenerateEndpointInvoker: + """Open or collect a stock token Generate endpoint invocation.""" + + def __init__(self, client: _DynamoClient) -> None: + self._client = client + + async def open( + self, + stage_id: str, + inputs: Mapping[str, Any], + context: StageContext, + *, + request_context: Any = None, + ) -> GenerateEndpointStream: + """Prepare the Generate request and return its owned token stream.""" + + request_value = inputs[GENERATE_REQUEST_PORT] + if not isinstance(request_value, Mapping): + raise WorkflowExecutionError("Generate endpoint request must be an object") + request = dict(request_value) + _validate_request_options(request) + + transport_context = None + if request_context is not None: + detach = getattr(request_context, "detached", None) + if not callable(detach): + raise WorkflowExecutionError( + "request context cannot create a detached child context" + ) + transport_context = detach(f"{context.attempt_id}:{stage_id}") + + try: + stream = await self._client.round_robin( + request, annotated=False, context=transport_context + ) + except BaseException: + if transport_context is not None: + transport_context.stop_generating() + raise + return GenerateEndpointStream(stream, transport_context) + + async def run( + self, + stage_id: str, + contract: StageContract, + inputs: Mapping[str, Any], + context: StageContext, + *, + request_context: Any = None, + ) -> Mapping[str, Any]: + del contract + stream = await self.open( + stage_id, + inputs, + context, + request_context=request_context, + ) + try: + completion = await collect_generation(stream, stage_id) + except BaseException: + await stream.aclose(cancel=True) + raise + await stream.aclose() + return {GENERATE_OUTPUT_PORT: completion} + + +def _validate_request_options(request: Mapping[str, Any]) -> None: + sampling_options = request.get("sampling_options", {}) + if not isinstance(sampling_options, Mapping): + raise WorkflowExecutionError("sampling_options must be an object") + # The Rust frontend represents an omitted OpenAI ``n`` as JSON null. + # Both omission forms mean the protocol default of one sequence. + if sampling_options.get("n") not in (None, 1): + raise WorkflowExecutionError("Generate endpoint workflow stage requires n=1") + + output_options = request.get("output_options", {}) + if not isinstance(output_options, Mapping): + raise WorkflowExecutionError("output_options must be an object") + if ( + output_options.get("logprobs") is not None + or output_options.get("prompt_logprobs") is not None + ): + raise WorkflowExecutionError( + "Generate endpoint workflow stage does not support logprobs" + ) + + +async def collect_generation( + stream: AsyncIterator[Any], stage_id: str +) -> dict[str, Any]: + """Fold one non-streaming workflow completion from Generate token deltas.""" + + token_ids: list[int] = [] + terminal: dict[str, Any] | None = None + async for value in stream: + if terminal is not None: + raise WorkflowExecutionError( + f"Generate endpoint stage {stage_id!r} returned data after terminal" + ) + if not isinstance(value, Mapping): + raise WorkflowExecutionError( + f"Generate endpoint stage {stage_id!r} returned a non-object chunk" + ) + chunk = dict(value) + if chunk.get("index") != 0: + raise WorkflowExecutionError( + f"Generate endpoint stage {stage_id!r} requires choice index 0" + ) + delta = chunk.get("token_ids") + if not isinstance(delta, list) or any( + isinstance(token_id, bool) or not isinstance(token_id, int) + for token_id in delta + ): + raise WorkflowExecutionError( + f"Generate endpoint stage {stage_id!r} returned invalid token_ids" + ) + if "log_probs" in chunk or "top_logprobs" in chunk: + raise WorkflowExecutionError( + f"Generate endpoint stage {stage_id!r} returned unsupported logprobs" + ) + token_ids.extend(delta) + finish_reason = chunk.get("finish_reason") + if finish_reason is not None: + if not isinstance(finish_reason, str) or not finish_reason: + raise WorkflowExecutionError( + f"Generate endpoint stage {stage_id!r} returned invalid finish_reason" + ) + terminal = chunk + + if terminal is None: + raise WorkflowExecutionError( + f"Generate endpoint stage {stage_id!r} returned no terminal chunk" + ) + terminal["token_ids"] = token_ids + return terminal diff --git a/components/src/dynamo/experimental/workflow/nixl.py b/components/src/dynamo/experimental/workflow/nixl.py new file mode 100644 index 000000000000..667a582649cf --- /dev/null +++ b/components/src/dynamo/experimental/workflow/nixl.py @@ -0,0 +1,435 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Explicit NIXL tensor carrier and producer-side lease ownership.""" + +from __future__ import annotations + +import asyncio +import logging +import math +import uuid +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Mapping, Optional + +from dynamo.experimental.workflow.runtime import WorkflowExecutionError + +logger = logging.getLogger(__name__) + +NIXL_TENSOR_SCHEMA = "dynamo.experimental.workflow.nixl_tensor" +NIXL_TENSOR_FANOUT_SCHEMA = "dynamo.experimental.workflow.nixl_tensor_fanout" +NIXL_TENSOR_VERSION = 0 +DEFAULT_NIXL_LEASE_TIMEOUT_S = 60.0 + + +def _check_keys(data: Mapping[str, Any], required: set[str]) -> None: + keys = set(data) + missing = required - keys + unknown = keys - required + if missing: + raise WorkflowExecutionError(f"NIXL metadata missing fields: {sorted(missing)}") + if unknown: + raise WorkflowExecutionError( + f"NIXL metadata has unknown fields: {sorted(unknown)}" + ) + + +def _validate_wire_version(data: Mapping[str, Any], schema: str) -> None: + if data["schema"] != schema: + raise WorkflowExecutionError(f"unsupported NIXL schema {data['schema']!r}") + version = data["version"] + if ( + not isinstance(version, int) + or isinstance(version, bool) + or version != NIXL_TENSOR_VERSION + ): + raise WorkflowExecutionError(f"unsupported NIXL version {version!r}") + + +@dataclass(frozen=True) +class NixlTensorRef: + """Serializable reference to one tensor-readable NIXL operation.""" + + transfer_id: str + lease_id: str + shape: tuple[int, ...] + dtype: str + device: str + rdma_metadata: Mapping[str, Any] + + def __post_init__(self) -> None: + for field_name, value in ( + ("transfer_id", self.transfer_id), + ("lease_id", self.lease_id), + ("dtype", self.dtype), + ("device", self.device), + ): + if not isinstance(value, str) or not value: + raise WorkflowExecutionError( + f"NIXL tensor {field_name} must be a non-empty string" + ) + shape = tuple(self.shape) + if any( + isinstance(dimension, bool) + or not isinstance(dimension, int) + or dimension < 0 + for dimension in shape + ): + raise WorkflowExecutionError( + "NIXL tensor shape must contain non-negative integers" + ) + if not isinstance(self.rdma_metadata, Mapping): + raise WorkflowExecutionError("NIXL rdma_metadata must be an object") + object.__setattr__(self, "shape", shape) + object.__setattr__( + self, "rdma_metadata", MappingProxyType(dict(self.rdma_metadata)) + ) + + def to_dict(self) -> dict[str, Any]: + return { + "schema": NIXL_TENSOR_SCHEMA, + "version": NIXL_TENSOR_VERSION, + "transfer_id": self.transfer_id, + "lease_id": self.lease_id, + "shape": list(self.shape), + "dtype": self.dtype, + "device": self.device, + "rdma_metadata": dict(self.rdma_metadata), + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "NixlTensorRef": + if not isinstance(data, Mapping): + raise WorkflowExecutionError("NIXL tensor reference must be an object") + _check_keys( + data, + { + "schema", + "version", + "transfer_id", + "lease_id", + "shape", + "dtype", + "device", + "rdma_metadata", + }, + ) + _validate_wire_version(data, NIXL_TENSOR_SCHEMA) + shape = data["shape"] + if not isinstance(shape, list): + raise WorkflowExecutionError("NIXL tensor shape must be an array") + return cls( + transfer_id=data["transfer_id"], + lease_id=data["lease_id"], + shape=tuple(shape), + dtype=data["dtype"], + device=data["device"], + rdma_metadata=data["rdma_metadata"], + ) + + +@dataclass(frozen=True) +class NixlTensorFanout: + """Per-consumer NIXL references for one logical tensor output.""" + + transfers: Mapping[str, NixlTensorRef] + + def __post_init__(self) -> None: + if not isinstance(self.transfers, Mapping) or not self.transfers: + raise WorkflowExecutionError("NIXL tensor fanout requires transfers") + transfers: dict[str, NixlTensorRef] = {} + for transfer_id, reference in sorted(self.transfers.items()): + if not isinstance(reference, NixlTensorRef): + raise WorkflowExecutionError( + "NIXL tensor fanout values must use NixlTensorRef" + ) + if transfer_id != reference.transfer_id: + raise WorkflowExecutionError( + "NIXL tensor fanout key does not match transfer id" + ) + transfers[transfer_id] = reference + object.__setattr__(self, "transfers", MappingProxyType(transfers)) + + def to_dict(self) -> dict[str, Any]: + return { + "schema": NIXL_TENSOR_FANOUT_SCHEMA, + "version": NIXL_TENSOR_VERSION, + "transfers": { + transfer_id: reference.to_dict() + for transfer_id, reference in self.transfers.items() + }, + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "NixlTensorFanout": + if not isinstance(data, Mapping): + raise WorkflowExecutionError("NIXL tensor fanout must be an object") + _check_keys(data, {"schema", "version", "transfers"}) + _validate_wire_version(data, NIXL_TENSOR_FANOUT_SCHEMA) + transfers = data["transfers"] + if not isinstance(transfers, Mapping): + raise WorkflowExecutionError( + "NIXL tensor fanout transfers must be an object" + ) + return cls( + { + transfer_id: NixlTensorRef.from_dict(reference) + for transfer_id, reference in transfers.items() + } + ) + + def for_transfer(self, transfer_id: str) -> NixlTensorRef: + try: + return self.transfers[transfer_id] + except KeyError as error: + raise WorkflowExecutionError( + f"NIXL tensor fanout has no transfer {transfer_id!r}" + ) from error + + +@dataclass +class _Lease: + operations: tuple[Any, ...] + value: Any + task: asyncio.Task[None] + + +class NixlLeaseRegistry: + """Keep producer memory registered until every read is confirmed complete.""" + + def __init__(self, timeout_s: float = DEFAULT_NIXL_LEASE_TIMEOUT_S) -> None: + if ( + isinstance(timeout_s, bool) + or not isinstance(timeout_s, (int, float)) + or not math.isfinite(timeout_s) + or timeout_s <= 0 + ): + raise ValueError("NIXL lease timeout must be a finite positive number") + self._timeout_s = float(timeout_s) + self._leases: dict[str, _Lease] = {} + + @property + def active_count(self) -> int: + return len(self._leases) + + def track(self, lease_id: str, operation: Any, value: Any) -> None: + self.track_fanout({lease_id: operation}, value) + + def track_fanout(self, operations: Mapping[str, Any], value: Any) -> None: + """Keep one shared tensor alive until every consumer read is terminal.""" + + if not isinstance(operations, Mapping) or not operations: + raise WorkflowExecutionError("NIXL fanout lease requires operations") + lease_ids = tuple(operations) + duplicate = next( + (lease_id for lease_id in lease_ids if lease_id in self._leases), None + ) + if duplicate is not None: + raise WorkflowExecutionError(f"duplicate NIXL lease {duplicate!r}") + operation_values = tuple(operations.values()) + task = asyncio.create_task( + self._wait_and_release(lease_ids, operation_values), + name=f"workflow-nixl-lease:{lease_ids[0]}", + ) + lease = _Lease(operation_values, value, task) + for lease_id in lease_ids: + self._leases[lease_id] = lease + + async def _wait_and_release( + self, lease_ids: tuple[str, ...], operations: tuple[Any, ...] + ) -> None: + async def wait_for_all() -> None: + results = await asyncio.gather( + *(operation.wait_for_completion() for operation in operations), + return_exceptions=True, + ) + failures = [ + result for result in results if isinstance(result, BaseException) + ] + if failures: + raise WorkflowExecutionError( + f"NIXL workflow leases {lease_ids!r} have uncertain read state" + ) from failures[0] + + completion = asyncio.create_task( + wait_for_all(), + name=f"workflow-nixl-completion:{lease_ids[0]}", + ) + try: + done, _ = await asyncio.wait({completion}, timeout=self._timeout_s) + if not done: + logger.warning( + "NIXL workflow leases %s were not all read within %.1fs; " + "retaining producer memory until late completion or process exit", + lease_ids, + self._timeout_s, + ) + await asyncio.shield(completion) + except asyncio.CancelledError: + logger.warning( + "NIXL workflow lease monitor for %s was cancelled; retaining " + "producer memory because read completion is uncertain", + lease_ids, + ) + raise + except Exception: + logger.warning( + "NIXL workflow leases %s entered quarantine because read completion " + "is uncertain; retaining producer memory until process exit", + lease_ids, + exc_info=True, + ) + return + + for lease_id in lease_ids: + self._leases.pop(lease_id, None) + for lease_id, operation in zip(lease_ids, operations): + try: + operation.__exit__(None, None, None) + except Exception: + logger.warning( + "failed to release completed NIXL workflow lease %s", + lease_id, + exc_info=True, + ) + + async def close(self) -> None: + if self._leases: + logger.warning( + "closing NIXL workflow registry with %d active leases; retaining " + "producer memory until late completion or process exit", + len(self._leases), + ) + await asyncio.sleep(0) + + +def _model_dump(value: Any) -> dict[str, Any]: + if hasattr(value, "model_dump"): + result = value.model_dump() + elif isinstance(value, Mapping): + result = dict(value) + else: + raise WorkflowExecutionError("NIXL metadata is not serializable") + if not isinstance(result, dict): + raise WorkflowExecutionError("NIXL metadata must serialize as an object") + return result + + +class NixlTensorCarrier: + """Export and import torch tensors with explicit NIXL READ operations.""" + + def __init__( + self, + *, + receive_device: Optional[str] = None, + lease_timeout_s: float = DEFAULT_NIXL_LEASE_TIMEOUT_S, + connector: Any = None, + nixl_module: Any = None, + torch_module: Any = None, + ) -> None: + if nixl_module is None: + try: + from dynamo import nixl_connect as default_nixl_module + except ImportError as error: + raise RuntimeError( + "NIXL tensor carrier requested but dynamo.nixl_connect is unavailable" + ) from error + nixl_module = default_nixl_module + if torch_module is None: + try: + import torch + except ImportError as error: + raise RuntimeError("NIXL tensor carrier requires PyTorch") from error + torch_module = torch + if receive_device is not None and ( + not isinstance(receive_device, str) or not receive_device + ): + raise ValueError("NIXL receive_device must be non-empty when set") + self._nixl = nixl_module + self._torch = torch_module + self._connector = connector or nixl_module.Connector() + self._receive_device = receive_device + self._leases = NixlLeaseRegistry(lease_timeout_s) + + @property + def active_leases(self) -> int: + return self._leases.active_count + + def can_export(self, value: Any) -> bool: + """Whether ``value`` is a complete top-level tensor port value.""" + + return isinstance(value, self._torch.Tensor) + + async def export_tensor(self, tensor: Any, transfer_id: str) -> Mapping[str, Any]: + references = await self.export_tensor_fanout(tensor, (transfer_id,)) + return references[transfer_id] + + async def export_tensor_fanout( + self, tensor: Any, transfer_ids: tuple[str, ...] + ) -> Mapping[str, Mapping[str, Any]]: + """Export one shared tensor through one independently addressed read per consumer.""" + + if not self.can_export(tensor): + raise WorkflowExecutionError("NIXL carrier can export torch.Tensor only") + if not isinstance(transfer_ids, tuple) or not transfer_ids: + raise WorkflowExecutionError("NIXL tensor export requires transfer ids") + if any( + not isinstance(transfer_id, str) or not transfer_id + for transfer_id in transfer_ids + ): + raise WorkflowExecutionError( + "NIXL tensor transfer ids must be non-empty strings" + ) + if len(set(transfer_ids)) != len(transfer_ids): + raise WorkflowExecutionError("NIXL tensor transfer ids must be unique") + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + descriptor = self._nixl.Descriptor(tensor) + readables: dict[str, Any] = {} + try: + references: dict[str, NixlTensorRef] = {} + for transfer_id in transfer_ids: + readable = await self._connector.create_readable(descriptor) + lease_id = uuid.uuid4().hex + readables[lease_id] = readable + references[transfer_id] = NixlTensorRef( + transfer_id=transfer_id, + lease_id=lease_id, + shape=tuple(tensor.shape), + dtype=str(tensor.dtype).rsplit(".", 1)[-1], + device=str(tensor.device), + rdma_metadata=_model_dump(readable.metadata()), + ) + self._leases.track_fanout(readables, tensor) + except BaseException: + for readable in readables.values(): + readable.__exit__(None, None, None) + raise + return { + transfer_id: reference.to_dict() + for transfer_id, reference in references.items() + } + + async def import_tensor(self, reference: Mapping[str, Any]) -> Any: + parsed = NixlTensorRef.from_dict(reference) + dtype = getattr(self._torch, parsed.dtype, None) + if dtype is None: + raise WorkflowExecutionError( + f"unsupported NIXL tensor dtype {parsed.dtype!r}" + ) + device = self._receive_device or parsed.device + tensor = self._torch.empty(parsed.shape, dtype=dtype, device=device) + descriptor = self._nixl.Descriptor(tensor) + rdma_metadata = self._nixl.RdmaMetadata.model_validate( + dict(parsed.rdma_metadata) + ) + operation = await self._connector.begin_read(rdma_metadata, descriptor) + try: + await operation.wait_for_completion() + finally: + operation.__exit__(None, None, None) + return tensor + + async def close(self) -> None: + await self._leases.close() diff --git a/components/src/dynamo/experimental/workflow/orchestrator.py b/components/src/dynamo/experimental/workflow/orchestrator.py new file mode 100644 index 000000000000..420af9ab85dc --- /dev/null +++ b/components/src/dynamo/experimental/workflow/orchestrator.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Execution lifecycle for compiled declarative workflows.""" + +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import Mapping +from types import MappingProxyType +from typing import Any + +from dynamo.experimental.workflow.dispatcher import StageDispatcher +from dynamo.experimental.workflow.plan import ExecutionPlan +from dynamo.experimental.workflow.runtime import StageRunner, WorkflowExecutionError +from dynamo.experimental.workflow.scheduler import GraphScheduler + + +class WorkflowOrchestrator: + """Own one compiled workflow's request and result lifecycle.""" + + # Declarative workflows execute through GraphScheduler. + # TODO: Support imperative authoring through WorkflowHandler. + + def __init__(self, plan: ExecutionPlan, dispatcher: StageDispatcher) -> None: + self._plan = plan + self._dispatcher = dispatcher + + @classmethod + async def bind( + cls, + plan: ExecutionPlan, + *, + runtime: Any = None, + inline_runners: Mapping[str, StageRunner] = MappingProxyType({}), + ) -> "WorkflowOrchestrator": + """Bind initialized resources to an immutable execution plan.""" + + dispatcher = await StageDispatcher.bind( + plan, + runtime=runtime, + inline_runners=inline_runners, + ) + return cls(plan, dispatcher) + + @property + def plan(self) -> ExecutionPlan: + return self._plan + + async def run( + self, + inputs: Mapping[str, Any], + *, + timeout: float | None = None, + attempt_id: str | None = None, + request_context: Any = None, + ) -> dict[str, Any]: + """Execute one graph request and return its named results.""" + + if timeout is not None and timeout <= 0: + raise ValueError("timeout must be positive") + workflow = self._plan.workflow + input_values = dict(inputs) + expected_inputs = set(workflow.inputs) + actual_inputs = set(input_values) + if actual_inputs != expected_inputs: + raise WorkflowExecutionError( + "workflow inputs differ from its definition; " + f"missing={sorted(expected_inputs - actual_inputs)}, " + f"extra={sorted(actual_inputs - expected_inputs)}" + ) + resolved_attempt_id = attempt_id or uuid.uuid4().hex + + async def execute() -> dict[str, Any]: + result = await GraphScheduler(workflow, self._dispatcher).run( + MappingProxyType(input_values), resolved_attempt_id, request_context + ) + return _validate_result(set(workflow.outputs), result) + + execution = asyncio.create_task( + execute(), name=f"workflow-attempt:{resolved_attempt_id}" + ) + try: + if timeout is None: + return await asyncio.shield(execution) + return await asyncio.wait_for(asyncio.shield(execution), timeout=timeout) + except BaseException: + if not execution.done(): + execution.cancel() + await asyncio.gather(execution, return_exceptions=True) + raise + + +def _validate_result(expected: set[str], result: Mapping[str, Any]) -> dict[str, Any]: + if not isinstance(result, Mapping): + raise WorkflowExecutionError("workflow returned a non-mapping result") + actual = set(result) + if actual != expected: + raise WorkflowExecutionError( + "workflow outputs differ from its definition; " + f"missing={sorted(expected - actual)}, extra={sorted(actual - expected)}" + ) + return dict(result) diff --git a/components/src/dynamo/experimental/workflow/plan.py b/components/src/dynamo/experimental/workflow/plan.py new file mode 100644 index 000000000000..8db2d707aa8e --- /dev/null +++ b/components/src/dynamo/experimental/workflow/plan.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""In-memory physical plans for Dynamo workflows.""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Mapping, Optional, Union + +from dynamo.experimental.workflow.ir import WorkflowIR +from dynamo.experimental.workflow.types import ( + StageContract, + WorkflowValidationError, + validate_name, +) + +NIXL_CARRIER = "nixl" + + +@dataclass(frozen=True) +class InlineBinding: + """Resolve one logical stage to a named in-process runner at bind time.""" + + runner_key: str + + def __post_init__(self) -> None: + validate_name(self.runner_key, "inline runner key") + + +def _validate_endpoint_id(endpoint_id: str) -> None: + if not isinstance(endpoint_id, str): + raise WorkflowValidationError("remote endpoint id must be a string") + parts = endpoint_id.split(".") + if len(parts) != 3: + raise WorkflowValidationError( + "remote endpoint id must use 'namespace.component.endpoint'" + ) + for kind, part in zip(("namespace", "component", "endpoint"), parts): + validate_name(part, f"remote {kind}") + + +@dataclass(frozen=True) +class RemoteBinding: + """Resolve one logical stage through a discovered Dynamo endpoint.""" + + endpoint_id: str + routing_policy: str = "round_robin" + tensor_carrier: Optional[str] = None + + def __post_init__(self) -> None: + _validate_endpoint_id(self.endpoint_id) + if self.routing_policy != "round_robin": + raise WorkflowValidationError( + f"unsupported remote routing policy {self.routing_policy!r}" + ) + if self.tensor_carrier not in {None, NIXL_CARRIER}: + raise WorkflowValidationError( + f"unsupported remote tensor carrier {self.tensor_carrier!r}" + ) + + +@dataclass(frozen=True) +class GenerateEndpointBinding(RemoteBinding): + """Bind a contracted stage to Dynamo's stock token Generate endpoint.""" + + +Binding = Union[InlineBinding, RemoteBinding] + + +def validate_binding_contract(binding: Binding, contract: StageContract) -> None: + """Validate protocol-specific stage ports against one physical binding.""" + + if not isinstance(binding, GenerateEndpointBinding): + return + if contract.inputs != {"request"}: + raise WorkflowValidationError("Generate endpoint stage input must be request") + if contract.outputs != {"completion"}: + raise WorkflowValidationError( + "Generate endpoint stage output must be completion" + ) + + +@dataclass(frozen=True) +class ExecutionPlan: + """A workflow plus immutable in-memory stage bindings.""" + + workflow: WorkflowIR + bindings: Mapping[str, Binding] + + def __post_init__(self) -> None: + if not isinstance(self.workflow, WorkflowIR): + raise WorkflowValidationError("execution plan requires WorkflowIR") + if not isinstance(self.bindings, Mapping): + raise WorkflowValidationError("execution plan bindings must be a mapping") + + bindings: dict[str, Binding] = {} + for stage_id, binding in sorted(self.bindings.items()): + validate_name(stage_id, "binding stage id") + if not isinstance(binding, (InlineBinding, RemoteBinding)): + raise WorkflowValidationError( + f"binding for stage {stage_id!r} uses an unsupported type" + ) + bindings[stage_id] = binding + + expected_stages = set(self.stage_contracts) + actual_stages = set(bindings) + if actual_stages != expected_stages: + raise WorkflowValidationError( + "execution plan bindings differ from workflow stages; " + f"missing={sorted(expected_stages - actual_stages)}, " + f"extra={sorted(actual_stages - expected_stages)}" + ) + + for stage_id, contract in self.stage_contracts.items(): + validate_binding_contract(bindings[stage_id], contract) + object.__setattr__(self, "bindings", MappingProxyType(bindings)) + + @property + def stage_contracts(self) -> Mapping[str, StageContract]: + """Return every stage contract keyed by its authored stage ID.""" + + return MappingProxyType( + {stage.id: stage.contract for stage in self.workflow.stages} + ) + + @property + def remote(self) -> bool: + """Whether every stage is bound to a remote endpoint.""" + + return bool(self.bindings) and all( + isinstance(binding, RemoteBinding) for binding in self.bindings.values() + ) diff --git a/components/src/dynamo/experimental/workflow/remote.py b/components/src/dynamo/experimental/workflow/remote.py new file mode 100644 index 000000000000..4efa7f6bdc4a --- /dev/null +++ b/components/src/dynamo/experimental/workflow/remote.py @@ -0,0 +1,451 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Adapters between unary workflow stages and streaming Dynamo endpoints.""" + +from __future__ import annotations + +import asyncio +import uuid +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, AsyncIterator, Mapping, Optional, Protocol + +from dynamo.experimental.workflow.nixl import NixlTensorFanout, NixlTensorRef +from dynamo.experimental.workflow.plan import NIXL_CARRIER +from dynamo.experimental.workflow.runtime import ( + StageContext, + StageRunner, + TensorCarrier, + WorkflowExecutionError, +) +from dynamo.experimental.workflow.types import ( + StageContract, + WorkflowValidationError, + validate_name, +) + +STAGE_REQUEST_SCHEMA = "dynamo.experimental.workflow.carrier_request" +STAGE_RESPONSE_SCHEMA = "dynamo.experimental.workflow.carrier_response" +STAGE_WIRE_VERSION = 1 + + +def _check_keys(data: Mapping[str, Any], required: set[str]) -> None: + actual = set(data) + missing = required - actual + unknown = actual - required + if missing: + raise WorkflowExecutionError( + f"remote carrier envelope missing fields: {sorted(missing)}" + ) + if unknown: + raise WorkflowExecutionError( + f"remote carrier envelope has unknown fields: {sorted(unknown)}" + ) + + +def _validate_header(data: Mapping[str, Any], schema: str) -> None: + if data["schema"] != schema: + raise WorkflowExecutionError( + f"unsupported remote carrier schema {data['schema']!r}" + ) + version = data["version"] + if isinstance(version, bool) or not isinstance(version, int): + raise WorkflowExecutionError("remote carrier version must be an integer") + if version != STAGE_WIRE_VERSION: + raise WorkflowExecutionError(f"unsupported remote carrier version {version!r}") + + +def _freeze_carriers( + carriers: Mapping[str, str], values: Mapping[str, Any] +) -> Mapping[str, str]: + if not isinstance(carriers, Mapping): + raise WorkflowExecutionError("remote carrier tags must be an object") + unknown = set(carriers) - set(values) + if unknown: + raise WorkflowExecutionError( + f"remote carrier tags reference unknown ports {sorted(unknown)}" + ) + normalized: dict[str, str] = {} + for name, carrier in carriers.items(): + validate_name(name, "remote carried port") + if carrier != NIXL_CARRIER: + raise WorkflowExecutionError( + f"unsupported remote value carrier {carrier!r}" + ) + normalized[name] = carrier + return MappingProxyType(normalized) + + +@dataclass(frozen=True) +class NixlCarriedValue: + """Internal marker that keeps NIXL metadata distinct from ordinary mappings.""" + + value: Mapping[str, Any] + + def __post_init__(self) -> None: + if not isinstance(self.value, Mapping): + raise WorkflowExecutionError("NIXL carried value must be an object") + object.__setattr__(self, "value", MappingProxyType(dict(self.value))) + + +@dataclass(frozen=True) +class StageRequestEnvelope: + """Internal carrier metadata sent to one remote stage.""" + + inputs: Mapping[str, Any] + input_carriers: Mapping[str, str] + output_transfers: Mapping[str, tuple[str, ...]] + + def __post_init__(self) -> None: + if not isinstance(self.inputs, Mapping): + raise WorkflowExecutionError("remote stage inputs must be an object") + inputs = MappingProxyType(dict(self.inputs)) + carriers = _freeze_carriers(self.input_carriers, inputs) + if not isinstance(self.output_transfers, Mapping): + raise WorkflowExecutionError("remote output transfers must be an object") + transfers: dict[str, tuple[str, ...]] = {} + for output_name, transfer_ids in self.output_transfers.items(): + validate_name(output_name, "remote output transfer port") + if not isinstance(transfer_ids, (list, tuple)): + raise WorkflowExecutionError( + "remote output transfer ids must be a list or tuple" + ) + normalized = tuple(transfer_ids) + if any( + not isinstance(transfer_id, str) or not transfer_id + for transfer_id in normalized + ): + raise WorkflowExecutionError( + "remote output transfer ids must be non-empty strings" + ) + if len(set(normalized)) != len(normalized): + raise WorkflowExecutionError( + f"remote output {output_name!r} has duplicate transfer ids" + ) + transfers[output_name] = normalized + object.__setattr__(self, "inputs", inputs) + object.__setattr__(self, "input_carriers", carriers) + object.__setattr__(self, "output_transfers", MappingProxyType(transfers)) + + def to_dict(self) -> dict[str, Any]: + return { + "schema": STAGE_REQUEST_SCHEMA, + "version": STAGE_WIRE_VERSION, + "inputs": dict(self.inputs), + "input_carriers": dict(self.input_carriers), + "output_transfers": { + name: list(transfer_ids) + for name, transfer_ids in self.output_transfers.items() + }, + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "StageRequestEnvelope": + if not isinstance(data, Mapping): + raise WorkflowExecutionError("remote stage request must be an object") + _check_keys( + data, + { + "schema", + "version", + "inputs", + "input_carriers", + "output_transfers", + }, + ) + _validate_header(data, STAGE_REQUEST_SCHEMA) + return cls( + inputs=data["inputs"], + input_carriers=data["input_carriers"], + output_transfers=data["output_transfers"], + ) + + +@dataclass(frozen=True) +class StageResponseEnvelope: + """Internal carrier metadata returned by one remote stage.""" + + outputs: Mapping[str, Any] + output_carriers: Mapping[str, str] + + def __post_init__(self) -> None: + if not isinstance(self.outputs, Mapping): + raise WorkflowExecutionError("remote stage outputs must be an object") + outputs = MappingProxyType(dict(self.outputs)) + object.__setattr__(self, "outputs", outputs) + object.__setattr__( + self, + "output_carriers", + _freeze_carriers(self.output_carriers, outputs), + ) + + def to_dict(self) -> dict[str, Any]: + return { + "schema": STAGE_RESPONSE_SCHEMA, + "version": STAGE_WIRE_VERSION, + "outputs": dict(self.outputs), + "output_carriers": dict(self.output_carriers), + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "StageResponseEnvelope": + if not isinstance(data, Mapping): + raise WorkflowExecutionError("remote stage response must be an object") + _check_keys( + data, + {"schema", "version", "outputs", "output_carriers"}, + ) + _validate_header(data, STAGE_RESPONSE_SCHEMA) + return cls( + outputs=data["outputs"], + output_carriers=data["output_carriers"], + ) + + +class _DynamoClient(Protocol): + async def round_robin( + self, + request: Mapping[str, Any], + *, + annotated: bool, + context: Any = None, + ) -> AsyncIterator[Any]: + ... + + +class RemoteStageClient: + """Adapt a unary workflow call to Dynamo's streaming endpoint API.""" + + def __init__(self, client: _DynamoClient) -> None: + self._client = client + + async def run( + self, + stage_id: str, + contract: StageContract, + inputs: Mapping[str, Any], + context: StageContext, + output_transfers: Mapping[str, tuple[str, ...]], + *, + request_context: Any = None, + ) -> Mapping[str, Any]: + stage_label = f"remote stage {stage_id!r} with contract {contract.id!r}" + wire_inputs: dict[str, Any] = {} + input_carriers: dict[str, str] = {} + for name, value in inputs.items(): + if isinstance(value, NixlCarriedValue): + wire_inputs[name] = dict(value.value) + input_carriers[name] = NIXL_CARRIER + else: + wire_inputs[name] = value + request = StageRequestEnvelope( + inputs=wire_inputs, + input_carriers=input_carriers, + output_transfers=output_transfers, + ) + + transport_context = None + if request_context is not None: + detach = getattr(request_context, "detached", None) + if not callable(detach): + raise WorkflowExecutionError( + "request context cannot create a detached child context" + ) + transport_context = detach(f"{context.attempt_id}:{context.stage_id}") + + try: + stream = await self._client.round_robin( + request.to_dict(), annotated=False, context=transport_context + ) + except asyncio.CancelledError: + raise + except Exception as error: + raise WorkflowExecutionError( + f"{stage_label} request failed at the transport boundary" + ) from error + + try: + try: + response = await stream.__anext__() + except StopAsyncIteration as error: + raise WorkflowExecutionError( + f"{stage_label} returned no response mapping" + ) from error + try: + await stream.__anext__() + except StopAsyncIteration: + pass + else: + raise WorkflowExecutionError( + f"{stage_label} returned multiple response mappings" + ) + envelope = StageResponseEnvelope.from_dict(response) + except BaseException as error: + if transport_context is not None: + transport_context.stop_generating() + close = getattr(stream, "aclose", None) + if callable(close): + await close() + if isinstance(error, (asyncio.CancelledError, WorkflowExecutionError)): + raise + if isinstance(error, Exception): + raise WorkflowExecutionError( + f"{stage_label} response failed at the transport boundary" + ) from error + raise + + outputs: dict[str, Any] = dict(envelope.outputs) + for name, carrier in envelope.output_carriers.items(): + if carrier == NIXL_CARRIER: + outputs[name] = NixlCarriedValue(envelope.outputs[name]) + return outputs + + +class RemoteStageServer: + """Adapt a unary ``StageRunner`` to Dynamo's streaming endpoint API.""" + + def __init__( + self, + stage_id: str, + runner: StageRunner, + tensor_carrier: Optional[TensorCarrier] = None, + ) -> None: + validate_name(stage_id, "remote stage id") + if not isinstance(runner, StageRunner): + raise WorkflowValidationError("remote runner must implement StageRunner") + if tensor_carrier is not None and not isinstance(tensor_carrier, TensorCarrier): + raise WorkflowValidationError( + "remote tensor_carrier must implement TensorCarrier" + ) + self._stage_id = stage_id + self._runner = runner + self._tensor_carrier = tensor_carrier + + async def generate( + self, request: Mapping[str, Any], context: Any = None + ) -> AsyncIterator[dict[str, Any]]: + envelope = StageRequestEnvelope.from_dict(request) + expected_inputs = set(self._runner.contract.inputs) + actual_inputs = set(envelope.inputs) + if actual_inputs != expected_inputs: + raise WorkflowExecutionError( + f"remote stage {self._stage_id!r} inputs differ from its contract; " + f"missing={sorted(expected_inputs - actual_inputs)}, " + f"extra={sorted(actual_inputs - expected_inputs)}" + ) + unknown_transfer_outputs = set(envelope.output_transfers) - set( + self._runner.contract.outputs + ) + if unknown_transfer_outputs: + raise WorkflowExecutionError( + f"remote stage {self._stage_id!r} has transfer requests for " + f"unknown outputs {sorted(unknown_transfer_outputs)}" + ) + + request_id = uuid.uuid4().hex + if context is not None: + get_request_id = getattr(context, "id", None) + if callable(get_request_id): + candidate = get_request_id() + if isinstance(candidate, str) and candidate: + request_id = candidate + stage_context = StageContext( + workflow_name=None, + stage_id=self._stage_id, + attempt_id=request_id, + ) + + async def invoke() -> tuple[dict[str, Any], dict[str, str]]: + runner_inputs = dict(envelope.inputs) + for name in envelope.input_carriers: + if self._tensor_carrier is None: + raise WorkflowExecutionError( + f"remote stage {self._stage_id!r} has no NIXL tensor carrier" + ) + runner_inputs[name] = await self._tensor_carrier.import_tensor( + envelope.inputs[name] + ) + + result = await self._runner.run( + MappingProxyType(runner_inputs), stage_context + ) + if not isinstance(result, Mapping): + raise WorkflowExecutionError( + f"remote stage {self._stage_id!r} returned a non-mapping result" + ) + expected_outputs = set(self._runner.contract.outputs) + actual_outputs = set(result) + if actual_outputs != expected_outputs: + raise WorkflowExecutionError( + f"remote stage {self._stage_id!r} outputs differ from its contract; " + f"missing={sorted(expected_outputs - actual_outputs)}, " + f"extra={sorted(actual_outputs - expected_outputs)}" + ) + + wire_outputs = dict(result) + output_carriers: dict[str, str] = {} + for name, value in result.items(): + if self._tensor_carrier is None or not self._tensor_carrier.can_export( + value + ): + continue + transfer_ids = envelope.output_transfers.get(name, ()) + if not transfer_ids: + raise WorkflowExecutionError( + f"remote tensor output {name!r} has no NIXL consumer transfers" + ) + references = await self._tensor_carrier.export_tensor_fanout( + value, transfer_ids + ) + if set(references) != set(transfer_ids): + raise WorkflowExecutionError( + f"remote tensor output {name!r} NIXL references differ " + "from requested consumer transfers" + ) + wire_outputs[name] = NixlTensorFanout( + { + transfer_id: NixlTensorRef.from_dict(reference) + for transfer_id, reference in references.items() + } + ).to_dict() + output_carriers[name] = NIXL_CARRIER + return wire_outputs, output_carriers + + invocation = asyncio.create_task(invoke(), name=f"workflow-remote:{request_id}") + transport_task: asyncio.Future[Any] | None = None + if context is not None: + wait_for_stop = getattr(context, "async_killed_or_stopped", None) + if callable(wait_for_stop): + transport_task = asyncio.ensure_future(wait_for_stop()) + try: + if transport_task is None: + wire_outputs, output_carriers = await invocation + else: + done, _ = await asyncio.wait( + {invocation, transport_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + if invocation in done: + wire_outputs, output_carriers = invocation.result() + else: + raise asyncio.CancelledError() + except BaseException: + if not invocation.done(): + invocation.cancel() + await asyncio.gather(invocation, return_exceptions=True) + raise + finally: + if transport_task is not None and not transport_task.done(): + transport_task.cancel() + await asyncio.gather(transport_task, return_exceptions=True) + + if context is not None and ( + bool(getattr(context, "is_stopped", lambda: False)()) + or bool(getattr(context, "is_killed", lambda: False)()) + ): + raise asyncio.CancelledError() + yield StageResponseEnvelope( + outputs=wire_outputs, + output_carriers=output_carriers, + ).to_dict() diff --git a/components/src/dynamo/experimental/workflow/runtime.py b/components/src/dynamo/experimental/workflow/runtime.py new file mode 100644 index 000000000000..6c65aff10447 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/runtime.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runtime binding and execution for compiled Dynamo workflows.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping, Protocol, runtime_checkable + +from dynamo.experimental.workflow.types import StageContract + + +class WorkflowExecutionError(RuntimeError): + """Raised when runtime values do not honor the authored workflow.""" + + +@dataclass(frozen=True) +class StageContext: + """Logical request identity available to one stage call. + + ``workflow_name`` is unavailable for remote stages until the transport + carries workflow lineage explicitly. + """ + + workflow_name: str | None + stage_id: str + attempt_id: str + + +@runtime_checkable +class StageRunner(Protocol): + """The small interface implemented by custom and Dynamo-provided workers.""" + + contract: StageContract + + async def run( + self, inputs: Mapping[str, Any], context: StageContext + ) -> Mapping[str, Any]: + """Run one stage attempt and return all declared outputs.""" + + ... + + +@runtime_checkable +class TensorCarrier(Protocol): + """Runtime-bound carrier for optional out-of-band tensor transport.""" + + def can_export(self, value: Any) -> bool: + """Whether this carrier can export the complete port value.""" + + ... + + async def export_tensor(self, tensor: Any, transfer_id: str) -> Mapping[str, Any]: + ... + + async def export_tensor_fanout( + self, tensor: Any, transfer_ids: tuple[str, ...] + ) -> Mapping[str, Mapping[str, Any]]: + ... + + async def import_tensor(self, reference: Mapping[str, Any]) -> Any: + ... diff --git a/components/src/dynamo/experimental/workflow/scheduler.py b/components/src/dynamo/experimental/workflow/scheduler.py new file mode 100644 index 000000000000..39ca8039e925 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/scheduler.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dependency scheduling for declarative workflow graphs.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Mapping +from typing import Any, cast + +from dynamo.experimental.workflow.dispatcher import StageDispatcher +from dynamo.experimental.workflow.ir import StageIR, WorkflowIR +from dynamo.experimental.workflow.runtime import StageContext +from dynamo.experimental.workflow.types import ValueRef + + +class GraphScheduler: + """Resolve static dependencies and run ready graph stages concurrently.""" + + def __init__(self, workflow: WorkflowIR, dispatcher: StageDispatcher) -> None: + self._workflow = workflow + self._dispatcher = dispatcher + + async def run( + self, + inputs: Mapping[str, Any], + attempt_id: str, + request_context: Any = None, + ) -> dict[str, Any]: + tasks: dict[str, asyncio.Task[dict[str, Any]]] = {} + + async def run_stage(stage: StageIR) -> dict[str, Any]: + stage_inputs = {} + for name, reference in stage.inputs.items(): + value = await resolve_raw(reference) + stage_inputs[name] = await self._dispatcher.resolve_edge( + reference, + stage.id, + name, + value, + ) + return await self._dispatcher.call( + stage.id, + stage_inputs, + StageContext( + workflow_name=self._workflow.name, + stage_id=stage.id, + attempt_id=attempt_id, + ), + request_context=request_context, + ) + + async def resolve_raw(reference: ValueRef) -> Any: + if reference.input_name is not None: + return inputs[reference.input_name] + stage_id = cast(str, reference.stage_id) + output_name = cast(str, reference.output_name) + stage_outputs = await asyncio.shield(tasks[stage_id]) + return stage_outputs[output_name] + + async def resolve_output(name: str, reference: ValueRef) -> Any: + value = await resolve_raw(reference) + return self._dispatcher.resolve_workflow_output(name, value) + + for stage in self._workflow.stages: + tasks[stage.id] = asyncio.create_task( + run_stage(stage), name=f"workflow:{stage.id}" + ) + + try: + output_values = await asyncio.gather( + *( + resolve_output(name, reference) + for name, reference in self._workflow.outputs.items() + ) + ) + return dict(zip(self._workflow.outputs, output_values)) + except BaseException: + for task in tasks.values(): + if not task.done(): + task.cancel() + await asyncio.gather(*tasks.values(), return_exceptions=True) + raise diff --git a/components/src/dynamo/experimental/workflow/tests/test_endpoint.py b/components/src/dynamo/experimental/workflow/tests/test_endpoint.py new file mode 100644 index 000000000000..4547923577d4 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/tests/test_endpoint.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import pytest + +from dynamo.experimental.workflow import ( + DeploymentSpec, + StageContract, + Workflow, + WorkflowEndpointHandler, + WorkflowOrchestrator, + WorkflowValidationError, + compile_workflow, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.gpu_0, + pytest.mark.core, +] + + +CONTRACT = StageContract( + id="token-stage", + inputs={"request"}, + outputs={"chunk"}, +) + + +class _Runner: + contract = CONTRACT + + async def run(self, inputs, context): + return { + "chunk": { + "token_ids": [inputs["request"]["token_ids"][0] + 1], + "index": 0, + "finish_reason": "stop", + } + } + + +async def _orchestrator(runner=None) -> WorkflowOrchestrator: + workflow = Workflow("endpoint-workflow") + request = workflow.input("request") + stage = workflow.stage("generate", CONTRACT, request=request) + workflow.output("chunk", stage.chunk) + return await WorkflowOrchestrator.bind( + compile_workflow(workflow, DeploymentSpec.inline(generate="generate")), + inline_runners={"generate": runner or _Runner()}, + ) + + +class _Context: + def __init__(self) -> None: + self.stopped = asyncio.Event() + + def id(self): + return "request-7" + + async def async_killed_or_stopped(self): + await self.stopped.wait() + return True + + +class _FutureContext(_Context): + def async_killed_or_stopped(self): + return asyncio.get_running_loop().create_future() + + +@pytest.mark.parametrize("context_type", [_Context, _FutureContext]) +async def test_endpoint_handler_runs_fixed_request_and_chunk_abi( + context_type, +) -> None: + context = context_type() + chunks = [ + chunk + async for chunk in WorkflowEndpointHandler(await _orchestrator()).generate( + {"token_ids": [41]}, context + ) + ] + + assert chunks == [{"token_ids": [42], "index": 0, "finish_reason": "stop"}] + + +async def test_endpoint_handler_rejects_noncanonical_workflow_abi() -> None: + workflow = Workflow("wrong-endpoint-abi") + request = workflow.input("payload") + stage = workflow.stage("generate", CONTRACT, request=request) + workflow.output("chunk", stage.chunk) + orchestrator = await WorkflowOrchestrator.bind( + compile_workflow(workflow), inline_runners={"generate": _Runner()} + ) + + with pytest.raises(WorkflowValidationError, match="'request' input"): + WorkflowEndpointHandler(orchestrator) + + +async def test_endpoint_handler_cancels_workflow_when_context_stops() -> None: + started = asyncio.Event() + cancelled = asyncio.Event() + + class BlockingRunner: + contract = CONTRACT + + async def run(self, inputs, context): + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + + context = _Context() + + async def stop(): + await started.wait() + context.stopped.set() + + stop_task = asyncio.create_task(stop()) + chunks = [ + chunk + async for chunk in WorkflowEndpointHandler( + await _orchestrator(BlockingRunner()) + ).generate({"token_ids": [1]}, context) + ] + await stop_task + + assert chunks == [] + assert cancelled.is_set() diff --git a/components/src/dynamo/experimental/workflow/tests/test_generate.py b/components/src/dynamo/experimental/workflow/tests/test_generate.py new file mode 100644 index 000000000000..244b1b937b09 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/tests/test_generate.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Mapping +from typing import Any + +import pytest + +from dynamo.experimental.workflow import ( + DeploymentSpec, + GenerateEndpointBinding, + StageContext, + StageContract, + Workflow, + WorkflowExecutionError, + WorkflowValidationError, + compile_workflow, +) +from dynamo.experimental.workflow.dispatcher import StageDispatcher +from dynamo.experimental.workflow.generate import ( + GenerateEndpointInvoker, + collect_generation, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.gpu_0, + pytest.mark.core, +] + + +GENERATOR = StageContract( + id="generator", + inputs={"request"}, + outputs={"completion"}, +) + + +def _workflow(generator_contract: StageContract = GENERATOR) -> Workflow: + workflow = Workflow("request-generator") + request = workflow.input("request") + generator = workflow.stage( + "generator", + generator_contract, + request=request, + ) + workflow.output( + "completion", generator.output(sorted(generator_contract.outputs)[0]) + ) + return workflow + + +def test_generate_binding_compiles_for_request_only_contract() -> None: + plan = compile_workflow( + _workflow(), + DeploymentSpec( + {"generator": GenerateEndpointBinding("models.decoder.generate")} + ), + ) + + assert plan.remote + assert plan.bindings == { + "generator": GenerateEndpointBinding("models.decoder.generate") + } + + +def test_generate_binding_rejects_a_non_generate_stage_contract() -> None: + incompatible = StageContract( + id="generator", + inputs=GENERATOR.inputs, + outputs={"chunk"}, + ) + + with pytest.raises(WorkflowValidationError, match="stage output"): + compile_workflow( + _workflow(incompatible), + DeploymentSpec( + {"generator": GenerateEndpointBinding("models.decoder.generate")} + ), + ) + + +class _Client: + def __init__(self, responses: list[Mapping[str, Any]]) -> None: + self.responses = responses + self.request: Mapping[str, Any] | None = None + self.context: Any = None + self.stream_closed = False + + async def wait_for_instances(self) -> None: + return None + + async def round_robin( + self, request: Mapping[str, Any], *, annotated: bool, context: Any = None + ) -> Any: + assert annotated is False + self.request = request + self.context = context + + async def stream(): + try: + for response in self.responses: + yield response + finally: + self.stream_closed = True + + return stream() + + +class _Endpoint: + def __init__(self, client: _Client) -> None: + self._client = client + + async def client(self) -> _Client: + return self._client + + +class _Runtime: + def __init__(self, clients: Mapping[str, _Client]) -> None: + self._clients = clients + + def endpoint(self, endpoint_id: str) -> _Endpoint: + return _Endpoint(self._clients[endpoint_id]) + + +def _context() -> StageContext: + return StageContext( + workflow_name="request-generator", + stage_id="generator", + attempt_id="request-1", + ) + + +class _TransportContext: + def __init__(self) -> None: + self.stopped = False + + def stop_generating(self) -> None: + self.stopped = True + + +class _ParentContext: + def __init__(self) -> None: + self.child = _TransportContext() + + def detached(self, context_id: str) -> _TransportContext: + assert context_id == "request-1:generator" + return self.child + + +def _request() -> dict[str, Any]: + return { + "token_ids": [1, 2], + "sampling_options": {"n": 1}, + "output_options": {}, + "multi_modal_data": {"image_url": [{"Url": "data:image/jpeg;base64,AA=="}]}, + "multi_modal_uuids": ["image-1"], + "mm_processor_kwargs": {"max_pixels": 1024}, + "mm_routing_info": {"mm_hashes": ["hash-1"]}, + } + + +async def test_generate_invoker_forwards_multimodal_request_unchanged() -> None: + transport = _Client( + [ + {"token_ids": [7], "index": 0}, + {"token_ids": [8, 9], "index": 0, "finish_reason": "stop"}, + ] + ) + request = _request() + + stream = await GenerateEndpointInvoker(transport).open( + "generator", + {"request": request}, + _context(), + ) + completion = await collect_generation(stream, "generator") + await stream.aclose() + + assert completion == { + "token_ids": [7, 8, 9], + "index": 0, + "finish_reason": "stop", + } + assert transport.request == request + + +async def test_generate_invoker_accepts_null_n_as_the_frontend_default() -> None: + transport = _Client([{"token_ids": [42], "index": 0, "finish_reason": "stop"}]) + + result = await GenerateEndpointInvoker(transport).run( + "generator", + GENERATOR, + { + "request": { + "token_ids": [1, 2], + "sampling_options": {"n": None}, + "output_options": {}, + } + }, + _context(), + ) + + assert result["completion"]["token_ids"] == [42] + + +async def test_generate_invoker_cancels_owned_stream_on_collection_error() -> None: + transport = _Client( + [ + {"token_ids": [42], "index": 0, "finish_reason": "stop"}, + {"token_ids": [43], "index": 0}, + ] + ) + parent = _ParentContext() + + with pytest.raises(WorkflowExecutionError, match="after terminal"): + await GenerateEndpointInvoker(transport).run( + "generator", + GENERATOR, + {"request": {"token_ids": [1], "output_options": {}}}, + _context(), + request_context=parent, + ) + + assert parent.child.stopped + assert transport.stream_closed + + +async def test_dispatcher_binds_generate_protocol_for_stock_endpoint() -> None: + plan = compile_workflow( + _workflow(), + DeploymentSpec( + {"generator": GenerateEndpointBinding("models.decoder.generate")} + ), + ) + generator_client = _Client( + [{"token_ids": [42], "index": 0, "finish_reason": "stop"}] + ) + dispatcher = await StageDispatcher.bind( + plan, + runtime=_Runtime({"models.decoder.generate": generator_client}), + ) + request = _request() + + result = await dispatcher.call( + "generator", + {"request": request}, + _context(), + ) + + assert result["completion"]["token_ids"] == [42] + assert generator_client.request == request + + +@pytest.mark.parametrize( + "request_value, message", + [ + ({"sampling_options": {"n": 2}}, "requires n=1"), + ({"output_options": {"logprobs": 0}}, "does not support logprobs"), + ], +) +async def test_generate_invoker_rejects_unsupported_frontend_options( + request_value: Mapping[str, Any], message: str +) -> None: + with pytest.raises(WorkflowExecutionError, match=message): + await GenerateEndpointInvoker(_Client([])).run( + "generator", + GENERATOR, + {"request": request_value}, + _context(), + ) diff --git a/components/src/dynamo/experimental/workflow/tests/test_nixl.py b/components/src/dynamo/experimental/workflow/tests/test_nixl.py new file mode 100644 index 000000000000..9309fe968d19 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/tests/test_nixl.py @@ -0,0 +1,379 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import pytest +import torch + +from dynamo.experimental.workflow import ( + DeploymentSpec, + InlineBinding, + NixlLeaseRegistry, + NixlTensorCarrier, + NixlTensorFanout, + NixlTensorRef, + RemoteBinding, + StageContract, + Workflow, + WorkflowExecutionError, + WorkflowOrchestrator, + compile_workflow, +) +from dynamo.experimental.workflow.dispatcher import StageDispatcher +from dynamo.experimental.workflow.remote import NixlCarriedValue + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.gpu_0, + pytest.mark.core, +] + + +class _Metadata: + def __init__(self, value): + self.value = value + + def model_dump(self): + return dict(self.value) + + +class _Descriptor: + def __init__(self, tensor): + self.tensor = tensor + + +class _Readable: + def __init__(self, key, descriptor): + self.key = key + self.descriptor = descriptor + self.completed = asyncio.Event() + self.released = False + + def metadata(self): + return _Metadata({"key": self.key}) + + async def wait_for_completion(self): + await self.completed.wait() + + def __exit__(self, exc_type, exc_value, traceback): + self.released = True + + +class _Read: + def __init__(self, source, destination): + self.source = source + self.destination = destination + self.released = False + + async def wait_for_completion(self): + self.destination.tensor.copy_(self.source.descriptor.tensor) + self.source.completed.set() + + def __exit__(self, exc_type, exc_value, traceback): + self.released = True + + +class _RdmaMetadata: + @classmethod + def model_validate(cls, value): + return dict(value) + + +class _NixlModule: + Descriptor = _Descriptor + RdmaMetadata = _RdmaMetadata + + +class _Connector: + def __init__(self): + self.readables = {} + self.reads = [] + + async def create_readable(self, descriptor): + key = f"transfer-{len(self.readables)}" + readable = _Readable(key, descriptor) + self.readables[key] = readable + return readable + + async def begin_read(self, metadata, descriptor): + read = _Read(self.readables[metadata["key"]], descriptor) + self.reads.append(read) + return read + + +async def _wait_for_no_leases(carrier): + for _ in range(20): + if carrier.active_leases == 0: + return + await asyncio.sleep(0) + raise AssertionError("NIXL lease did not complete") + + +async def _wait_for_registry_empty(registry: NixlLeaseRegistry) -> None: + for _ in range(20): + if registry.active_count == 0: + return + await asyncio.sleep(0) + raise AssertionError("NIXL registry did not release completed reads") + + +async def test_tensor_carrier_round_trip_keeps_per_transfer_lease() -> None: + connector = _Connector() + carrier = NixlTensorCarrier( + connector=connector, + nixl_module=_NixlModule, + torch_module=torch, + ) + source = torch.arange(12, dtype=torch.float32).reshape(3, 4) + + reference = await carrier.export_tensor(source, "classifier.embedding") + assert carrier.active_leases == 1 + received = await carrier.import_tensor(reference) + await _wait_for_no_leases(carrier) + + assert torch.equal(received, source) + assert connector.readables["transfer-0"].released + assert connector.reads[0].released + + +async def test_one_logical_tensor_can_have_independent_consumer_leases() -> None: + connector = _Connector() + carrier = NixlTensorCarrier( + connector=connector, + nixl_module=_NixlModule, + torch_module=torch, + ) + source = torch.ones((2, 8), dtype=torch.float16) + + references = await carrier.export_tensor_fanout( + source, ("classifier.embedding", "generator.embedding") + ) + classifier = NixlTensorRef.from_dict(references["classifier.embedding"]) + generator = NixlTensorRef.from_dict(references["generator.embedding"]) + fanout = NixlTensorFanout( + { + classifier.transfer_id: classifier, + generator.transfer_id: generator, + } + ) + + assert carrier.active_leases == 2 + assert NixlTensorFanout.from_dict(fanout.to_dict()) == fanout + assert ( + fanout.for_transfer("classifier.embedding").lease_id + != fanout.for_transfer("generator.embedding").lease_id + ) + + await carrier.import_tensor(classifier.to_dict()) + assert carrier.active_leases == 2 + assert not connector.readables["transfer-0"].released + assert not connector.readables["transfer-1"].released + await carrier.import_tensor(generator.to_dict()) + await _wait_for_no_leases(carrier) + assert connector.readables["transfer-0"].released + assert connector.readables["transfer-1"].released + + +async def test_lease_registry_retains_unread_operation_after_timeout() -> None: + registry = NixlLeaseRegistry(timeout_s=0.01) + operation = _Readable("never-read", _Descriptor(torch.ones(1))) + + registry.track("lease-1", operation, torch.ones(1024)) + assert registry.active_count == 1 + await asyncio.sleep(0.02) + + assert registry.active_count == 1 + assert not operation.released + + operation.completed.set() + await _wait_for_registry_empty(registry) + + assert operation.released + + +async def test_lease_registry_close_does_not_release_uncertain_read() -> None: + registry = NixlLeaseRegistry(timeout_s=1.0) + operation = _Readable("active-read", _Descriptor(torch.ones(1))) + + registry.track("lease-1", operation, torch.ones(1024)) + await registry.close() + + assert registry.active_count == 1 + assert not operation.released + + operation.completed.set() + await _wait_for_registry_empty(registry) + assert operation.released + + +def test_tensor_reference_rejects_unknown_wire_fields() -> None: + reference = NixlTensorRef( + transfer_id="stage.input", + lease_id="lease", + shape=(2, 3), + dtype="float32", + device="cpu", + rdma_metadata={"opaque": True}, + ).to_dict() + reference["fallback"] = "inline" + + with pytest.raises(WorkflowExecutionError, match="unknown fields"): + NixlTensorRef.from_dict(reference) + + +ENCODER = StageContract( + id="encoder", + inputs={"request"}, + outputs={"embedding"}, +) +CLASSIFIER = StageContract( + id="classifier", + inputs={"embedding"}, + outputs={"scores"}, +) +GENERATOR = StageContract( + id="generator", + inputs={"embedding"}, + outputs={"text"}, +) + + +def _tensor_workflow() -> Workflow: + workflow = Workflow("nixl-fanout") + request = workflow.input("request") + encoder = workflow.stage("encoder", ENCODER, request=request) + classifier = workflow.stage("classifier", CLASSIFIER, embedding=encoder.embedding) + generator = workflow.stage("generator", GENERATOR, embedding=encoder.embedding) + workflow.output("scores", classifier.scores) + workflow.output("text", generator.text) + return workflow + + +class _RemoteTensorInvoker: + def __init__(self, role: str) -> None: + self.role = role + self.calls: list[tuple[dict, dict]] = [] + + async def run( + self, + stage_id, + contract, + inputs, + context, + output_transfers, + *, + request_context=None, + ): + del request_context + self.calls.append((dict(inputs), dict(output_transfers))) + assert stage_id == self.role + if self.role == "encoder": + transfer_ids = output_transfers.get("embedding", ("classifier.embedding",)) + return { + "embedding": NixlCarriedValue( + NixlTensorFanout( + { + transfer_id: NixlTensorRef( + transfer_id=transfer_id, + lease_id=f"lease-{transfer_id}", + shape=(4, 8), + dtype="float32", + device="cuda:0", + rdma_metadata={"opaque": transfer_id}, + ) + for transfer_id in transfer_ids + } + ).to_dict() + ) + } + assert isinstance(inputs["embedding"], NixlCarriedValue) + reference = NixlTensorRef.from_dict(inputs["embedding"].value) + assert reference.transfer_id == f"{self.role}.embedding" + if self.role == "classifier": + return {"scores": {"ok": 1.0}} + return {"text": "generated"} + + +async def test_graph_dispatches_one_remote_nixl_reference_per_consumer() -> None: + endpoints = { + stage_id: f"workflows.{stage_id}.generate" + for stage_id in ("encoder", "classifier", "generator") + } + plan = compile_workflow( + _tensor_workflow(), + DeploymentSpec.remote(tensor_carrier="nixl", **endpoints), + ) + invokers = { + endpoints[stage_id]: _RemoteTensorInvoker(stage_id) for stage_id in endpoints + } + dispatcher = StageDispatcher(plan, {}, invokers) + executor = WorkflowOrchestrator(plan, dispatcher) + + assert await executor.run({"request": {"token_ids": [1]}}) == { + "scores": {"ok": 1.0}, + "text": "generated", + } + assert invokers[endpoints["encoder"]].calls[0][1] == { + "embedding": ("classifier.embedding", "generator.embedding") + } + + +async def test_carried_tensor_cannot_cross_inline_boundary() -> None: + workflow = Workflow("mixed-nixl") + request = workflow.input("request") + encoder = workflow.stage("encoder", ENCODER, request=request) + classifier = workflow.stage("classifier", CLASSIFIER, embedding=encoder.embedding) + workflow.output("scores", classifier.scores) + plan = compile_workflow( + workflow, + DeploymentSpec( + { + "encoder": RemoteBinding( + "workflows.encoder.generate", tensor_carrier="nixl" + ), + "classifier": InlineBinding("classifier"), + } + ), + ) + + class InlineClassifier: + contract = CLASSIFIER + + async def run(self, inputs, context): + raise AssertionError("inline runner must not receive a NIXL reference") + + dispatcher = StageDispatcher( + plan, + {"classifier": InlineClassifier()}, + {"workflows.encoder.generate": _RemoteTensorInvoker("encoder")}, + ) + executor = WorkflowOrchestrator(plan, dispatcher) + + with pytest.raises(WorkflowExecutionError, match="must target a NIXL remote"): + await executor.run({"request": {"token_ids": [1]}}) + + +async def test_carried_tensor_cannot_be_a_workflow_output() -> None: + workflow = Workflow("nixl-output") + request = workflow.input("request") + encoder = workflow.stage("encoder", ENCODER, request=request) + workflow.output("embedding", encoder.embedding) + endpoint = "workflows.encoder.generate" + plan = compile_workflow( + workflow, + DeploymentSpec.remote( + tensor_carrier="nixl", + encoder=endpoint, + ), + ) + dispatcher = StageDispatcher( + plan, + {}, + {endpoint: _RemoteTensorInvoker("encoder")}, + ) + executor = WorkflowOrchestrator(plan, dispatcher) + + with pytest.raises(WorkflowExecutionError, match="cannot expose a NIXL tensor"): + await executor.run({"request": {"token_ids": [1]}}) diff --git a/components/src/dynamo/experimental/workflow/tests/test_plan.py b/components/src/dynamo/experimental/workflow/tests/test_plan.py new file mode 100644 index 000000000000..28b31d6e68db --- /dev/null +++ b/components/src/dynamo/experimental/workflow/tests/test_plan.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from dynamo.experimental.workflow import ( + DeploymentSpec, + ExecutionPlan, + InlineBinding, + RemoteBinding, + StageContract, + Workflow, + WorkflowValidationError, + compile_workflow, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.gpu_0, + pytest.mark.core, +] + + +def _workflow(stage_id: str = "normalize") -> Workflow: + workflow = Workflow("physical-plan") + text = workflow.input("text") + stage = workflow.stage( + stage_id, + StageContract( + id="normalize", + inputs={"text"}, + outputs={"normalized"}, + ), + text=text, + ) + workflow.output("text", stage.normalized) + return workflow + + +def test_execution_plan_contains_only_in_memory_decisions() -> None: + plan = compile_workflow(_workflow(), DeploymentSpec.inline(normalize="normalizer")) + + assert plan.workflow == _workflow().build() + assert plan.bindings == {"normalize": InlineBinding(runner_key="normalizer")} + + +def test_compilation_defaults_to_stage_id_inline_bindings() -> None: + plan = compile_workflow(_workflow()) + + assert plan.bindings == {"normalize": InlineBinding(runner_key="normalize")} + + +def test_default_compilation_supports_stage_named_cls() -> None: + plan = compile_workflow(_workflow("cls")) + + assert plan.bindings == {"cls": InlineBinding(runner_key="cls")} + + +def test_execution_plan_rejects_missing_stage_bindings() -> None: + plan = compile_workflow(_workflow(), DeploymentSpec.inline(normalize="normalizer")) + + with pytest.raises(WorkflowValidationError, match=r"missing=\['normalize'\]"): + ExecutionPlan( + workflow=plan.workflow, + bindings={}, + ) + + +def test_remote_plan_records_remote_binding() -> None: + plan = compile_workflow( + _workflow(), + DeploymentSpec.remote(normalize="workflows.normalize.generate"), + ) + + assert plan.remote + assert plan.bindings == {"normalize": RemoteBinding("workflows.normalize.generate")} + + +def test_mixed_placement_records_each_stage_binding() -> None: + contract = StageContract( + id="text-stage", + inputs={"text"}, + outputs={"text"}, + ) + workflow = Workflow("mixed-placement") + value = workflow.input("text") + first = workflow.stage("first", contract, text=value) + second = workflow.stage("second", contract, text=first.text) + workflow.output("text", second.text) + + plan = compile_workflow( + workflow, + DeploymentSpec( + { + "first": InlineBinding("first"), + "second": RemoteBinding("workflows.second.generate"), + } + ), + ) + + assert plan.bindings == { + "first": InlineBinding("first"), + "second": RemoteBinding("workflows.second.generate"), + } + + +def test_remote_planning_does_not_assume_a_runtime_value_type() -> None: + contract = StageContract( + id="opaque-stage", + inputs={"value"}, + outputs={"result"}, + ) + workflow = Workflow("remote-opaque-value") + value = workflow.input("value") + result = workflow.stage("stage", contract, value=value) + workflow.output("result", result.result) + + plan = compile_workflow( + workflow, + DeploymentSpec.remote(stage="workflows.stage.generate"), + ) + + assert plan.bindings == {"stage": RemoteBinding("workflows.stage.generate")} + + +def test_remote_endpoint_id_is_a_stable_discovery_identity() -> None: + assert RemoteBinding("namespace.component.endpoint").endpoint_id == ( + "namespace.component.endpoint" + ) + + with pytest.raises(WorkflowValidationError, match="namespace.component.endpoint"): + RemoteBinding("component.endpoint") + + +def test_nixl_is_an_opt_in_remote_binding_capability() -> None: + encoder_contract = StageContract( + id="encoder", + inputs={"request"}, + outputs={"embedding"}, + ) + classifier_contract = StageContract( + id="classifier", + inputs={"embedding"}, + outputs={"scores"}, + ) + generator_contract = StageContract( + id="generator", + inputs={"embedding"}, + outputs={"text"}, + ) + workflow = Workflow("nixl-fanout") + request = workflow.input("request") + encoder = workflow.stage("encoder", encoder_contract, request=request) + classifier = workflow.stage( + "classifier", classifier_contract, embedding=encoder.embedding + ) + generator = workflow.stage( + "generator", generator_contract, embedding=encoder.embedding + ) + workflow.output("scores", classifier.scores) + workflow.output("text", generator.text) + + plan = compile_workflow( + workflow, + DeploymentSpec.remote( + tensor_carrier="nixl", + encoder="workflows.encoder.generate", + classifier="workflows.classifier.generate", + generator="workflows.generator.generate", + ), + ) + assert all(binding.tensor_carrier == "nixl" for binding in plan.bindings.values()) + assert not hasattr(plan, "edges") + + +def test_compilation_keeps_mixed_placement_value_opaque() -> None: + producer_contract = StageContract( + id="producer", + inputs={"request"}, + outputs={"embedding"}, + ) + consumer_contract = StageContract( + id="consumer", + inputs={"embedding"}, + outputs={"result"}, + ) + workflow = Workflow("unsupported-mixed-tensor") + request = workflow.input("request") + producer = workflow.stage("producer", producer_contract, request=request) + consumer = workflow.stage( + "consumer", consumer_contract, embedding=producer.embedding + ) + workflow.output("result", consumer.result) + + placements = ( + { + "producer": InlineBinding("producer"), + "consumer": RemoteBinding( + "workflows.consumer.generate", tensor_carrier="nixl" + ), + }, + { + "producer": RemoteBinding( + "workflows.producer.generate", tensor_carrier="nixl" + ), + "consumer": InlineBinding("consumer"), + }, + ) + for bindings in placements: + plan = compile_workflow(workflow, DeploymentSpec(bindings)) + assert plan.bindings == bindings diff --git a/components/src/dynamo/experimental/workflow/tests/test_remote.py b/components/src/dynamo/experimental/workflow/tests/test_remote.py new file mode 100644 index 000000000000..b563f952963d --- /dev/null +++ b/components/src/dynamo/experimental/workflow/tests/test_remote.py @@ -0,0 +1,492 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from collections.abc import Mapping +from typing import Any + +import pytest +import torch + +from dynamo.experimental.workflow import ( + DeploymentSpec, + InlineBinding, + NixlTensorFanout, + NixlTensorRef, + RemoteBinding, + StageContext, + StageContract, + Workflow, + WorkflowExecutionError, + WorkflowOrchestrator, + compile_workflow, +) +from dynamo.experimental.workflow.remote import ( + STAGE_REQUEST_SCHEMA, + NixlCarriedValue, + RemoteStageClient, + RemoteStageServer, + StageRequestEnvelope, + StageResponseEnvelope, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.gpu_0, + pytest.mark.core, +] + + +CONTRACT = StageContract( + id="normalize", + inputs={"text"}, + outputs={"normalized"}, +) + + +def _context(): + return StageContext( + workflow_name="remote-wire", + stage_id="normalize", + attempt_id="request-1", + ) + + +def _request(inputs, *, carriers=None, transfers=None): + return StageRequestEnvelope( + inputs=inputs, + input_carriers={} if carriers is None else carriers, + output_transfers={} if transfers is None else transfers, + ).to_dict() + + +def _response(outputs, *, carriers=None): + return StageResponseEnvelope( + outputs=outputs, + output_carriers={} if carriers is None else carriers, + ).to_dict() + + +def test_carrier_envelope_round_trip_is_strict_and_versioned() -> None: + envelope = StageRequestEnvelope( + inputs={"text": "HELLO"}, + input_carriers={}, + output_transfers={}, + ) + + assert StageRequestEnvelope.from_dict(envelope.to_dict()) == envelope + bad = envelope.to_dict() + bad["extra"] = True + with pytest.raises(WorkflowExecutionError, match="unknown fields"): + StageRequestEnvelope.from_dict(bad) + + bad = envelope.to_dict() + bad["schema"] = f"{STAGE_REQUEST_SCHEMA}.future" + with pytest.raises(WorkflowExecutionError, match="unsupported.*schema"): + StageRequestEnvelope.from_dict(bad) + + bad = envelope.to_dict() + bad["version"] = 0 + with pytest.raises(WorkflowExecutionError, match="unsupported.*version"): + StageRequestEnvelope.from_dict(bad) + + +class _Client: + def __init__(self, responses) -> None: + self.responses = responses + self.request = None + self.context = None + + async def round_robin(self, request, *, annotated, context=None): + assert annotated is False + self.request = request + self.context = context + + async def stream(): + for response in self.responses: + yield response + + return stream() + + +class _ChildContext: + def __init__(self, context_id: str) -> None: + self.context_id = context_id + self._stopped = asyncio.Event() + + def stop_generating(self) -> None: + self._stopped.set() + + def is_stopped(self) -> bool: + return self._stopped.is_set() + + def is_killed(self) -> bool: + return False + + def id(self) -> str: + return self.context_id + + async def async_killed_or_stopped(self) -> bool: + await self._stopped.wait() + return True + + +class _ParentContext: + def __init__(self) -> None: + self.children: list[_ChildContext] = [] + + def detached(self, context_id: str) -> _ChildContext: + child = _ChildContext(context_id) + self.children.append(child) + return child + + +async def test_remote_client_wraps_ordinary_values_for_the_endpoint() -> None: + transport = _Client([_response({"normalized": "hello"})]) + + result = await RemoteStageClient(transport).run( + "normalize", CONTRACT, {"text": "HELLO"}, _context(), {} + ) + + assert result == {"normalized": "hello"} + assert StageRequestEnvelope.from_dict(transport.request).inputs == {"text": "HELLO"} + + +async def test_remote_client_tags_only_internal_carried_values() -> None: + reference = {"transfer_id": "normalize.text", "opaque": True} + transport = _Client( + [_response({"normalized": reference}, carriers={"normalized": "nixl"})] + ) + + result = await RemoteStageClient(transport).run( + "normalize", + CONTRACT, + {"text": NixlCarriedValue(reference)}, + _context(), + {}, + ) + + request = StageRequestEnvelope.from_dict(transport.request) + assert request.inputs == {"text": reference} + assert request.input_carriers == {"text": "nixl"} + assert result == {"normalized": NixlCarriedValue(reference)} + + +async def test_remote_client_creates_invocation_scoped_transport_context() -> None: + transport = _Client([_response({"normalized": "hello"})]) + parent = _ParentContext() + + await RemoteStageClient(transport).run( + "normalize", + CONTRACT, + {"text": "HELLO"}, + _context(), + {}, + request_context=parent, + ) + + assert [child.context_id for child in parent.children] == ["request-1:normalize"] + assert transport.context is parent.children[0] + + +async def test_remote_client_enforces_one_response_mapping() -> None: + with pytest.raises(WorkflowExecutionError, match="no response mapping"): + await RemoteStageClient(_Client([])).run( + "normalize", CONTRACT, {"text": "HELLO"}, _context(), {} + ) + + response = _response({"normalized": "hello"}) + parent = _ParentContext() + with pytest.raises(WorkflowExecutionError, match="multiple response mappings"): + await RemoteStageClient(_Client([response, response])).run( + "normalize", + CONTRACT, + {"text": "HELLO"}, + _context(), + {}, + request_context=parent, + ) + assert parent.children[0].is_stopped() + + +async def test_remote_client_adds_stage_context_to_transport_failures() -> None: + class FailingClient: + async def round_robin(self, request, *, annotated, context=None): + raise TypeError("unsupported payload") + + with pytest.raises( + WorkflowExecutionError, + match="remote stage 'normalize'.*transport boundary", + ) as error: + await RemoteStageClient(FailingClient()).run( + "normalize", CONTRACT, {"text": object()}, _context(), {} + ) + + assert isinstance(error.value.__cause__, TypeError) + + +class _Runner: + contract = CONTRACT + + async def run(self, inputs, context): + assert context.workflow_name is None + assert context.stage_id == "normalize" + assert context.attempt_id == "request-1:normalize" + return {"normalized": inputs["text"].strip().lower()} + + +async def test_remote_server_runs_unary_stage_over_streaming_endpoint() -> None: + context = _ChildContext("request-1:normalize") + + responses = [ + response + async for response in RemoteStageServer("normalize", _Runner()).generate( + _request({"text": " HELLO "}), context=context + ) + ] + + assert dict(StageResponseEnvelope.from_dict(responses[0]).outputs) == { + "normalized": "hello" + } + + +async def test_ordinary_mapping_cannot_collide_with_carrier_metadata() -> None: + value = { + "schema": STAGE_REQUEST_SCHEMA, + "version": 1, + "input_carriers": {"value": "nixl"}, + } + + class OpaqueRunner: + contract = StageContract(id="opaque", inputs={"value"}, outputs={"result"}) + + async def run(self, inputs, context): + return {"result": inputs["value"]} + + response = ( + await RemoteStageServer("opaque", OpaqueRunner()) + .generate(_request({"value": value})) + .__anext__() + ) + + assert StageResponseEnvelope.from_dict(response).outputs["result"] == value + + +async def test_remote_server_cancels_runner_when_transport_stops() -> None: + class BlockingRunner: + contract = CONTRACT + + def __init__(self) -> None: + self.started = asyncio.Event() + self.cancelled = asyncio.Event() + + async def run(self, inputs, context): + self.started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled.set() + raise + + runner = BlockingRunner() + context = _ChildContext("request-1:normalize") + response = asyncio.create_task( + RemoteStageServer("normalize", runner) + .generate(_request({"text": "hello"}), context=context) + .__anext__() + ) + await runner.started.wait() + context.stop_generating() + + with pytest.raises(asyncio.CancelledError): + await response + assert runner.cancelled.is_set() + + +class _TensorCarrier: + def __init__(self) -> None: + self.exports: list[str] = [] + + def can_export(self, value): + return isinstance(value, torch.Tensor) + + async def import_tensor(self, reference): + assert reference == {"remote": "reference"} + return torch.ones((2, 4), dtype=torch.float32) + + async def export_tensor(self, tensor, transfer_id): + return (await self.export_tensor_fanout(tensor, (transfer_id,)))[transfer_id] + + async def export_tensor_fanout(self, tensor, transfer_ids): + self.exports.extend(transfer_ids) + return { + transfer_id: NixlTensorRef( + transfer_id=transfer_id, + lease_id=f"lease-{transfer_id}", + shape=tuple(tensor.shape), + dtype="float32", + device="cpu", + rdma_metadata={"opaque": transfer_id}, + ).to_dict() + for transfer_id in transfer_ids + } + + +async def test_server_selects_nixl_for_top_level_tensor_at_runtime() -> None: + class TensorRunner: + contract = StageContract( + id="tensor", inputs={"tensor"}, outputs={"tensor", "metadata"} + ) + + async def run(self, inputs, context): + return {"tensor": inputs["tensor"] * 2, "metadata": {"ok": True}} + + carrier = _TensorCarrier() + response = ( + await RemoteStageServer("tensor", TensorRunner(), carrier) + .generate( + _request( + {"tensor": {"remote": "reference"}}, + carriers={"tensor": "nixl"}, + transfers={"tensor": ("classifier.tensor", "generator.tensor")}, + ) + ) + .__anext__() + ) + envelope = StageResponseEnvelope.from_dict(response) + fanout = NixlTensorFanout.from_dict(envelope.outputs["tensor"]) + + assert set(fanout.transfers) == {"classifier.tensor", "generator.tensor"} + assert envelope.output_carriers == {"tensor": "nixl"} + assert envelope.outputs["metadata"] == {"ok": True} + assert carrier.exports == ["classifier.tensor", "generator.tensor"] + + +async def test_tensor_output_requires_remote_nixl_consumers() -> None: + class TensorRunner: + contract = StageContract(id="tensor", inputs={"value"}, outputs={"tensor"}) + + async def run(self, inputs, context): + return {"tensor": torch.ones(1)} + + with pytest.raises(WorkflowExecutionError, match="no NIXL consumer transfers"): + await RemoteStageServer("tensor", TensorRunner(), _TensorCarrier()).generate( + _request({"value": "ordinary"}) + ).__anext__() + + +async def test_nested_tensor_is_not_implicitly_carried() -> None: + nested = {"tensor": torch.ones(1)} + + class NestedRunner: + contract = StageContract(id="nested", inputs={"value"}, outputs={"result"}) + + async def run(self, inputs, context): + return {"result": nested} + + response = ( + await RemoteStageServer("nested", NestedRunner(), _TensorCarrier()) + .generate(_request({"value": "ordinary"})) + .__anext__() + ) + envelope = StageResponseEnvelope.from_dict(response) + + assert envelope.outputs["result"] is nested + assert envelope.output_carriers == {} + + +TEXT_ENCODER = StageContract(id="text-encoder", inputs={"text"}, outputs={"tokens"}) +TEXT_GENERATOR = StageContract(id="text-generator", inputs={"tokens"}, outputs={"text"}) +RESPONSE = StageContract(id="response", inputs={"text"}, outputs={"chunk"}) + + +class _TextEncoder: + contract = TEXT_ENCODER + + async def run(self, inputs: Mapping[str, Any], context: StageContext): + return {"tokens": inputs["text"].lower().split()} + + +class _TextGenerator: + contract = TEXT_GENERATOR + + async def run(self, inputs: Mapping[str, Any], context: StageContext): + return {"text": " ".join(reversed(inputs["tokens"]))} + + +class _Response: + contract = RESPONSE + + async def run(self, inputs: Mapping[str, Any], context: StageContext): + return {"chunk": {"text": inputs["text"]}} + + +class _LoopbackClient: + def __init__(self, server: RemoteStageServer) -> None: + self._server = server + + async def wait_for_instances(self) -> None: + return None + + async def round_robin(self, request, *, annotated, context=None): + assert annotated is False + return self._server.generate(request, context) + + +class _Endpoint: + def __init__(self, client: Any) -> None: + self._client = client + + async def client(self) -> Any: + return self._client + + +class _Runtime: + def __init__(self, clients: Mapping[str, Any]) -> None: + self._clients = clients + + def endpoint(self, endpoint_id: str) -> _Endpoint: + return _Endpoint(self._clients[endpoint_id]) + + +async def test_ordinary_values_work_on_nixl_capable_remote_bindings() -> None: + workflow = Workflow("remote-text") + text = workflow.input("text") + encoder = workflow.stage("encoder", TEXT_ENCODER, text=text) + generator = workflow.stage("generator", TEXT_GENERATOR, tokens=encoder.tokens) + response = workflow.stage("response", RESPONSE, text=generator.text) + workflow.output("chunk", response.chunk) + endpoints = { + "encoder": "workflows.encoder.generate", + "generator": "workflows.generator.generate", + } + plan = compile_workflow( + workflow, + DeploymentSpec( + { + "encoder": RemoteBinding(endpoints["encoder"], tensor_carrier="nixl"), + "generator": RemoteBinding( + endpoints["generator"], tensor_carrier="nixl" + ), + "response": InlineBinding("response"), + } + ), + ) + clients = { + endpoints["encoder"]: _LoopbackClient( + RemoteStageServer("encoder", _TextEncoder(), _TensorCarrier()) + ), + endpoints["generator"]: _LoopbackClient( + RemoteStageServer("generator", _TextGenerator(), _TensorCarrier()) + ), + } + orchestrator = await WorkflowOrchestrator.bind( + plan, + runtime=_Runtime(clients), + inline_runners={"response": _Response()}, + ) + + assert await orchestrator.run({"text": "Hello Dynamo"}) == { + "chunk": {"text": "dynamo hello"} + } diff --git a/components/src/dynamo/experimental/workflow/tests/test_runtime.py b/components/src/dynamo/experimental/workflow/tests/test_runtime.py new file mode 100644 index 000000000000..a8ae481f5af5 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/tests/test_runtime.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from dataclasses import dataclass +from types import SimpleNamespace + +import pytest + +from dynamo.experimental.workflow import ( + DeploymentSpec, + StageContext, + StageContract, + StageRunner, + Workflow, + WorkflowExecutionError, + WorkflowOrchestrator, + WorkflowValidationError, + compile_workflow, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.gpu_0, + pytest.mark.core, + pytest.mark.timeout(10), +] + + +ENCODER = StageContract( + id="encoder-worker", + inputs={"text"}, + outputs={"embedding"}, +) +CLASSIFIER = StageContract( + id="classifier-worker", + inputs={"embedding"}, + outputs={"scores"}, +) +GENERATOR = StageContract( + id="generator-worker", + inputs={"embedding"}, + outputs={"text"}, +) + + +def _workflow() -> Workflow: + workflow = Workflow("local-execution") + text = workflow.input("text") + encoder = workflow.stage("encoder", _Encoder.contract, text=text) + classifier = workflow.stage( + "classifier", _Classifier.contract, embedding=encoder.embedding + ) + generator = workflow.stage( + "generator", _Generator.contract, embedding=encoder.embedding + ) + workflow.output("scores", classifier.scores) + workflow.output("text", generator.text) + return workflow + + +async def _compile_local( + workflow: Workflow, **runners: StageRunner +) -> WorkflowOrchestrator: + plan = compile_workflow( + workflow, + DeploymentSpec.inline(**{stage_id: stage_id for stage_id in runners}), + ) + return await WorkflowOrchestrator.bind(plan, inline_runners=runners) + + +@dataclass +class _Encoder: + contract = ENCODER + embedding: object + + async def run(self, inputs, context: StageContext): + assert context.stage_id == "encoder" + return {"embedding": self.embedding} + + +@dataclass +class _Classifier: + contract = CLASSIFIER + expected: object + + async def run(self, inputs, context: StageContext): + assert inputs["embedding"] is self.expected + return {"scores": {"class-a": 0.75, "class-b": 0.25}} + + +@dataclass +class _Generator: + contract = GENERATOR + expected: object + + async def run(self, inputs, context: StageContext): + assert inputs["embedding"] is self.expected + return {"text": "generated"} + + +async def test_concise_compile_and_run_preserve_fanout_value_identity(): + embedding = object() + plan = await _compile_local( + _workflow(), + encoder=_Encoder(embedding), + classifier=_Classifier(embedding), + generator=_Generator(embedding), + ) + + result = await plan.run({"text": "hello"}, attempt_id="request-1") + + assert result == { + "scores": {"class-a": 0.75, "class-b": 0.25}, + "text": "generated", + } + + +class _BranchBarrier: + def __init__(self, count: int) -> None: + self._remaining = count + self.open = asyncio.Event() + + async def enter(self) -> None: + self._remaining -= 1 + if self._remaining == 0: + self.open.set() + await self.open.wait() + + +@dataclass +class _BarrierClassifier: + contract = CLASSIFIER + barrier: _BranchBarrier + + async def run(self, inputs, context: StageContext): + await self.barrier.enter() + return {"scores": {"ok": True}} + + +@dataclass +class _BarrierGenerator: + contract = GENERATOR + barrier: _BranchBarrier + + async def run(self, inputs, context: StageContext): + await self.barrier.enter() + return {"text": "joined"} + + +async def test_independent_branches_run_concurrently_before_join(): + embedding = object() + barrier = _BranchBarrier(2) + plan = await _compile_local( + _workflow(), + encoder=_Encoder(embedding), + classifier=_BarrierClassifier(barrier), + generator=_BarrierGenerator(barrier), + ) + + assert await plan.run({"text": "hello"}) == { + "scores": {"ok": True}, + "text": "joined", + } + + +class WorkerFailure(RuntimeError): + pass + + +@dataclass +class _FailingClassifier: + contract = CLASSIFIER + barrier: _BranchBarrier + + async def run(self, inputs, context: StageContext): + await self.barrier.enter() + raise WorkerFailure("classifier failed") + + +@dataclass +class _CancelledGenerator: + contract = GENERATOR + barrier: _BranchBarrier + cancelled: asyncio.Event + + async def run(self, inputs, context: StageContext): + await self.barrier.enter() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled.set() + raise + + +async def test_first_worker_failure_cancels_and_awaits_sibling(): + embedding = object() + barrier = _BranchBarrier(2) + cancelled = asyncio.Event() + plan = await _compile_local( + _workflow(), + encoder=_Encoder(embedding), + classifier=_FailingClassifier(barrier), + generator=_CancelledGenerator(barrier, cancelled), + ) + + with pytest.raises(WorkerFailure, match="classifier failed"): + await plan.run({"text": "hello"}) + + assert cancelled.is_set() + + +@dataclass +class _BlockingClassifier: + contract = CLASSIFIER + started: asyncio.Event + cancelled: asyncio.Event + + async def run(self, inputs, context: StageContext): + self.started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled.set() + raise + + +async def test_timeout_cancels_and_awaits_running_stages(): + embedding = object() + started = asyncio.Event() + cancelled = asyncio.Event() + plan = await _compile_local( + _workflow(), + encoder=_Encoder(embedding), + classifier=_BlockingClassifier(started, cancelled), + generator=_Generator(embedding), + ) + + with pytest.raises(asyncio.TimeoutError): + await plan.run({"text": "hello"}, timeout=0.01) + + assert started.is_set() + assert cancelled.is_set() + + +async def test_caller_cancellation_cleans_up_running_stages(): + embedding = object() + started = asyncio.Event() + cancelled = asyncio.Event() + plan = await _compile_local( + _workflow(), + encoder=_Encoder(embedding), + classifier=_BlockingClassifier(started, cancelled), + generator=_Generator(embedding), + ) + task = asyncio.create_task(plan.run({"text": "hello"})) + await started.wait() + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert cancelled.is_set() + + +async def test_compile_requires_exact_bindings_and_matching_contracts(): + with pytest.raises(WorkflowValidationError, match="missing"): + compile_workflow(_workflow(), DeploymentSpec.inline(encoder="encoder")) + + wrong = SimpleNamespace(contract=CLASSIFIER, run=_Generator(object()).run) + with pytest.raises(WorkflowValidationError, match="does not match"): + await WorkflowOrchestrator.bind( + compile_workflow( + _workflow(), + DeploymentSpec.inline( + encoder="encoder", + classifier="classifier", + generator="generator", + ), + ), + inline_runners={ + "encoder": _Encoder(object()), + "classifier": _Classifier(object()), + "generator": wrong, + }, + ) + + +async def test_runtime_accepts_opaque_values_and_rejects_bad_outputs(): + embedding = object() + plan = await _compile_local( + _workflow(), + encoder=_Encoder(embedding), + classifier=_Classifier(embedding), + generator=_Generator(embedding), + ) + assert await plan.run({"text": object()}) == { + "scores": {"class-a": 0.75, "class-b": 0.25}, + "text": "generated", + } + with pytest.raises(WorkflowExecutionError, match="extra"): + await plan.run({"text": "hello", "extra": "value"}) + + class BadGenerator: + contract = GENERATOR + + async def run(self, inputs, context): + return {"wrong": "value"} + + bad_plan = await _compile_local( + _workflow(), + encoder=_Encoder(embedding), + classifier=_Classifier(embedding), + generator=BadGenerator(), + ) + with pytest.raises(WorkflowExecutionError, match="outputs differ"): + await bad_plan.run({"text": "hello"}) diff --git a/components/src/dynamo/experimental/workflow/vllm/__init__.py b/components/src/dynamo/experimental/workflow/vllm/__init__.py new file mode 100644 index 000000000000..ba155f4b76b2 --- /dev/null +++ b/components/src/dynamo/experimental/workflow/vllm/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Experimental vLLM integrations for Dynamo workflows.""" + +from dynamo.experimental.workflow.vllm.stages import DynamoVllmStage + +__all__ = ["DynamoVllmStage"] diff --git a/components/src/dynamo/experimental/workflow/vllm/stages.py b/components/src/dynamo/experimental/workflow/vllm/stages.py new file mode 100644 index 000000000000..5c0ed005c8de --- /dev/null +++ b/components/src/dynamo/experimental/workflow/vllm/stages.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Contract for a stock Dynamo vLLM workflow stage.""" + +from dynamo.experimental.workflow import StageContract + + +class DynamoVllmStage: + """Contract implemented by a stock aggregated Dynamo vLLM worker.""" + + request_complete_contract = StageContract( + id="dynamo-vllm-request-complete", + inputs={"request"}, + outputs={"completion"}, + ) diff --git a/components/src/dynamo/experimental/workflow/vllm/tests/test_vllm_components.py b/components/src/dynamo/experimental/workflow/vllm/tests/test_vllm_components.py new file mode 100644 index 000000000000..dae2bb5418fe --- /dev/null +++ b/components/src/dynamo/experimental/workflow/vllm/tests/test_vllm_components.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from dynamo.experimental.workflow import GenerateEndpointBinding +from dynamo.experimental.workflow.plan import validate_binding_contract +from dynamo.experimental.workflow.vllm import DynamoVllmStage + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.vllm, + pytest.mark.gpu_0, +] + + +def test_dynamo_vllm_stage_matches_generate_endpoint_contract() -> None: + validate_binding_contract( + GenerateEndpointBinding("workflows.generator.generate"), + DynamoVllmStage.request_complete_contract, + ) + + assert set(DynamoVllmStage.request_complete_contract.inputs) == {"request"} + assert set(DynamoVllmStage.request_complete_contract.outputs) == {"completion"}