From 8096551bcdca76cf664978df5ab230fd5bd0a096 Mon Sep 17 00:00:00 2001 From: benjcooley Date: Fri, 21 Aug 2026 20:33:47 -0700 Subject: [PATCH] =?UTF-8?q?feat(sdk):=20ops-first=20asset=20interface=20?= =?UTF-8?q?=E2=80=94=20OpsProvider=20seam,=20SDK=5FREFS=20marshaling,=20ra?= =?UTF-8?q?w()=20escape=20hatch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nodes operate on assets (image.invert()) and never receive buffers; compute runs on the trusted plane via the OpsProvider seam. Raw buffer access becomes a permissioned, discouraged escape hatch (raw(); forces dedicated tier under the overlay). The execution seam wraps heavy inputs as refs for SDK_REFS nodes and resolves output refs for downstream legacy nodes. The .pyi contract no longer imports torch. POC stand-ins (interface debt, ledgered in the overlay repo DEBT.md): invert/scale enumerated on OpsProvider; duck-typed wrap_inputs; SDK_REFS class-attr opt-in. --- comfy_api/latest/_sdk.py | 90 +++++++++++++++++-- comfy_api/latest/_sdk_public.pyi | 22 +++-- custom_nodes/comfy_sdk_poc/__init__.py | 25 +++--- execution.py | 13 ++- .../comfy_api_test/test_secure_sdk_seam.py | 22 +++-- 5 files changed, 135 insertions(+), 37 deletions(-) diff --git a/comfy_api/latest/_sdk.py b/comfy_api/latest/_sdk.py index 092d999d2..19704f3d4 100644 --- a/comfy_api/latest/_sdk.py +++ b/comfy_api/latest/_sdk.py @@ -76,17 +76,31 @@ class _TypedRef(Ref): class TensorRef(_TypedRef): KIND = "TENSOR" - async def tensor(self) -> "torch.Tensor": + # --- RAW ESCAPE HATCH (permissioned, discouraged) -------------------- # + # `raw()` returns the underlying buffer object. It is NOT the preferred + # interface: prefer operations on the asset (see ImageRef below). Raw + # access requires the `raw`/`tensor.read` capability and forces a node to + # the dedicated tier under the overlay. The SDK contract does not depend on + # torch; the return is deliberately untyped (Any) here. + async def raw(self) -> Any: return await current_runtime().refs.resolve(self) @classmethod - async def from_tensor(cls, t: "torch.Tensor") -> "TensorRef": - return cls._wrap(await current_runtime().refs.create(cls.KIND, t)) # type: ignore[return-value] + async def _from_raw(cls, obj: Any) -> "TensorRef": + return cls._wrap(await current_runtime().refs.create(cls.KIND, obj)) # type: ignore[return-value] class ImageRef(TensorRef): KIND = "IMAGE" + # --- PREFERRED INTERFACE: operations on the asset. The heavy compute runs + # engine-side (trusted plane); the node never receives a buffer. ---- # + async def invert(self) -> "ImageRef": + return await current_runtime().ops.invert(self) + + async def scale(self, factor: float) -> "ImageRef": + return await current_runtime().ops.scale(self, factor) + class MaskRef(TensorRef): KIND = "MASK" @@ -142,6 +156,7 @@ class AssetRef(_TypedRef): class Runtime: refs: "RefResolver" ctx: "Context" + ops: "OpsProvider" = None # engine-side operations (the preferred interface) _active_runtime: "contextvars.ContextVar[Optional[Runtime]]" = contextvars.ContextVar( @@ -178,8 +193,10 @@ class _RuntimeScope: _active_runtime.reset(self._token) -def bind_runtime(refs: "RefResolver", ctx: "Context") -> _RuntimeScope: - return _RuntimeScope(Runtime(refs=refs, ctx=ctx)) +def bind_runtime( + refs: "RefResolver", ctx: "Context", ops: "OpsProvider" = None +) -> _RuntimeScope: + return _RuntimeScope(Runtime(refs=refs, ctx=ctx, ops=ops)) # --------------------------------------------------------------------------- # @@ -221,6 +238,16 @@ class CtxProvider(Protocol): def build(self, plan: ExecutionPlan) -> "Context": ... +@runtime_checkable +class OpsProvider(Protocol): + """Engine-side operations on assets — the preferred node interface. The + node passes/receives refs; the buffer math happens here, on the trusted + plane (in-process default) or in the engine (overlay).""" + + async def invert(self, image: "ImageRef") -> "ImageRef": ... + async def scale(self, image: "ImageRef", factor: float) -> "ImageRef": ... + + # --------------------------------------------------------------------------- # # ctx — the brokered side-effect surface. Interfaces first; in-process defaults # implement the important ones over real core. Domains not needed by the POC @@ -370,6 +397,20 @@ class InProcessCtxProvider: ) +class InProcessOps: + """Default operations. Runs in the trusted process; uses the real buffers + via the resolver but never hands them to the node. Kept torch-free (tensor + arithmetic works on the resolved objects directly).""" + + async def invert(self, image: "ImageRef") -> "ImageRef": + t = await current_runtime().refs.resolve(image) + return ImageRef._wrap(await current_runtime().refs.create("IMAGE", 1.0 - t)) # type: ignore[return-value] + + async def scale(self, image: "ImageRef", factor: float) -> "ImageRef": + t = await current_runtime().refs.resolve(image) + return ImageRef._wrap(await current_runtime().refs.create("IMAGE", t * factor)) # type: ignore[return-value] + + class InProcessExecutionBackend: async def dispatch( self, @@ -379,6 +420,40 @@ class InProcessExecutionBackend: return await local_call() +# --------------------------------------------------------------------------- # +# Input/output ref marshaling for SDK nodes (those declaring ``SDK_REFS``). +# Heavy inputs become refs before execute(); output refs resolve back to real +# objects for downstream (legacy) nodes. This is what makes execute() see +# assets, not buffers. Under the overlay this happens at the process boundary. +# --------------------------------------------------------------------------- # +def _looks_like_tensor(v: Any) -> bool: + return type(v).__name__ == "Tensor" and hasattr(v, "shape") + + +async def wrap_inputs(resolver: "RefResolver", inputs: dict) -> dict: + out = {} + for k, v in inputs.items(): + if _looks_like_tensor(v): + out[k] = ImageRef._wrap(await resolver.create("IMAGE", v)) + elif isinstance(v, dict) and "samples" in v: + out[k] = LatentRef._wrap(await resolver.create("LATENT", v)) + else: + out[k] = v + return out + + +async def unwrap_outputs(resolver: "RefResolver", node_output: Any) -> Any: + args = getattr(node_output, "result", None) + if not args: + return node_output + resolved = [] + for a in args: + resolved.append(await resolver.resolve(a) if isinstance(a, Ref) else a) + from ._io import NodeOutput + + return NodeOutput(*resolved) + + # --------------------------------------------------------------------------- # # Provider registry — the seam the overlay attaches to. # --------------------------------------------------------------------------- # @@ -386,6 +461,7 @@ class Providers: def __init__(self) -> None: self.execution_backend: ExecutionBackend = InProcessExecutionBackend() self.ctx_provider: CtxProvider = InProcessCtxProvider() + self.ops_provider: OpsProvider = InProcessOps() self.ref_resolver_factory: Callable[[], RefResolver] = InProcessRefResolver self._overlay_name: Optional[str] = None @@ -398,6 +474,10 @@ class Providers: logger.info("SDK: ctx provider -> %s", type(impl).__name__) self.ctx_provider = impl + def register_ops_provider(self, impl: OpsProvider) -> None: + logger.info("SDK: ops provider -> %s", type(impl).__name__) + self.ops_provider = impl + def register_ref_resolver_factory(self, factory: Callable[[], RefResolver]) -> None: logger.info("SDK: ref resolver -> %s", getattr(factory, "__name__", factory)) self.ref_resolver_factory = factory diff --git a/comfy_api/latest/_sdk_public.pyi b/comfy_api/latest/_sdk_public.pyi index 6a022e18a..d46f3d7c2 100644 --- a/comfy_api/latest/_sdk_public.pyi +++ b/comfy_api/latest/_sdk_public.pyi @@ -19,14 +19,17 @@ authoritative. """ from typing import Any, Awaitable, Callable, Optional, Protocol, TypeVar, runtime_checkable -import torch +# NOTE: the contract deliberately does NOT import torch or any backend module. +# Node code has no direct access to torch/CUDA/filesystem; it works through +# refs, operations, and ctx — all brokered. # =========================================================================== # # 1. NODE-AUTHOR SURFACE # =========================================================================== # -# --- Refs: opaque, typed resource handles. A ref never exposes the buffer; in -# process it is zero-copy, under the overlay it is a shm/CUDA-IPC handle. --- # +# --- Refs: opaque, typed asset handles. A ref never exposes the buffer. The +# PREFERRED interface is operations on the asset (below); raw buffer access +# is a permissioned escape hatch. --- # class Ref: kind: str id: str @@ -35,16 +38,19 @@ _T = TypeVar("_T", bound="TensorRef") class TensorRef(Ref): KIND: str - async def tensor(self) -> torch.Tensor: ... - @classmethod - async def from_tensor(cls: type[_T], t: torch.Tensor) -> _T: ... + # RAW ESCAPE HATCH — permissioned (`raw`/`tensor.read`), discouraged; forces + # the dedicated tier under the overlay. Return is untyped by design (the + # contract does not depend on torch). + async def raw(self) -> Any: ... class ImageRef(TensorRef): - """IMAGE — torch.Tensor [B,H,W,C].""" + """IMAGE asset. Preferred interface = operations; the buffer stays engine-side.""" KIND: str + async def invert(self) -> "ImageRef": ... + async def scale(self, factor: float) -> "ImageRef": ... class MaskRef(TensorRef): - """MASK — torch.Tensor.""" + """MASK asset.""" KIND: str _L = TypeVar("_L", bound="LatentRef") diff --git a/custom_nodes/comfy_sdk_poc/__init__.py b/custom_nodes/comfy_sdk_poc/__init__.py index c367fc2dc..7b1a0c7f5 100644 --- a/custom_nodes/comfy_sdk_poc/__init__.py +++ b/custom_nodes/comfy_sdk_poc/__init__.py @@ -1,15 +1,19 @@ """POC custom node authored against the v0_0_3 custom-node SDK. -Demonstrates the SDK in real execution: the node uses ``ctx`` for progress and -round-trips its image through a ref (``ImageRef.from_tensor`` / ``.tensor()``). -In-process (OSS) that ref is zero-copy; under the overlay the same code runs in -an isolated guest with a shm/CUDA-IPC ref — unchanged.""" +The node receives its image as an ``ImageRef`` (an asset handle), NOT a buffer, +and transforms it through an engine-side operation (``image.invert()``). It +never imports torch and never touches a raw tensor. In-process (OSS) the +operation runs on the trusted plane; under the overlay the same code runs in an +isolated guest and the op RPCs to the engine — unchanged.""" from __future__ import annotations from comfy_api.v0_0_3 import ComfyExtension, io, sdk class SandboxInvert(io.ComfyNode): + # Opt in to the SDK asset model: execute() receives refs, not buffers. + SDK_REFS = True + @classmethod def define_schema(cls) -> io.Schema: return io.Schema( @@ -21,19 +25,12 @@ class SandboxInvert(io.ComfyNode): ) @classmethod - async def execute(cls, image) -> io.NodeOutput: + async def execute(cls, image) -> io.NodeOutput: # image: sdk.ImageRef ctx = sdk.ctx() await ctx.progress.update(0.0, 1.0) - - # Ref round-trip: create a ref from the input, materialize it, invert. - # In-process this is the real tensor (zero-copy); isolated it is a - # shm/CUDA-IPC handle mapped into this process. - ref = await sdk.ImageRef.from_tensor(image) - t = await ref.tensor() - out = 1.0 - t - + out = await image.invert() # engine-side op on the asset; no buffer here await ctx.progress.update(1.0, 1.0) - return io.NodeOutput(out) + return io.NodeOutput(out) # returns an ImageRef; the engine resolves it class PocExtension(ComfyExtension): diff --git a/execution.py b/execution.py index e34ea8d03..844f18b10 100644 --- a/execution.py +++ b/execution.py @@ -308,10 +308,16 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f node_id=str(unique_id), node_type=getattr(type_obj, "__name__", "node"), ) + _sdk_refs = _comfy_sdk.providers.ref_resolver_factory() _sdk_runtime = _comfy_sdk.bind_runtime( - _comfy_sdk.providers.ref_resolver_factory(), + _sdk_refs, _comfy_sdk.providers.ctx_provider.build(_sdk_plan), + _comfy_sdk.providers.ops_provider, ) + # SDK nodes see assets (refs), not buffers: wrap heavy inputs. + _sdk_refs_mode = getattr(type_obj, "SDK_REFS", False) + if _sdk_refs_mode: + inputs = await _comfy_sdk.wrap_inputs(_sdk_refs, inputs) # V1 else: f = getattr(obj, func) @@ -337,6 +343,11 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f if _sdk_plan is not None: result = await _comfy_sdk.providers.execution_backend.dispatch(_sdk_plan, local_call) + # Resolve output refs back to real objects for downstream nodes. + if _sdk_refs_mode: + if isinstance(result, asyncio.Task): + result = await result + result = await _comfy_sdk.unwrap_outputs(_sdk_refs, result) else: result = await local_call() results.append(result) diff --git a/tests-unit/comfy_api_test/test_secure_sdk_seam.py b/tests-unit/comfy_api_test/test_secure_sdk_seam.py index b02f2fe8b..512f2a605 100644 --- a/tests-unit/comfy_api_test/test_secure_sdk_seam.py +++ b/tests-unit/comfy_api_test/test_secure_sdk_seam.py @@ -20,6 +20,10 @@ from comfy_api.v0_0_3 import io class _InvertAsync(io.ComfyNode): + # SDK asset node: receives an ImageRef, transforms via an engine-side op, + # never touches a buffer. + SDK_REFS = True + @classmethod def define_schema(cls): return io.Schema( @@ -31,23 +35,23 @@ class _InvertAsync(io.ComfyNode): async def execute(cls, image): ctx = sdk.ctx() await ctx.progress.update(0.0, 1.0) - ref = await sdk.ImageRef.from_tensor(image) - t = await ref.tensor() + out = await image.invert() # operation on the asset await ctx.progress.update(1.0, 1.0) - return io.NodeOutput(1.0 - t) + return io.NodeOutput(out) -class _ScaleSync(io.ComfyNode): +class _ScaleSyncLegacy(io.ComfyNode): + # Legacy (non-SDK) v3 node: sync execute, receives a raw tensor. Confirms + # the sync dispatch branch + that non-SDK nodes are unaffected by the seam. @classmethod def define_schema(cls): return io.Schema( - node_id="_TestScaleSync", category="test", + node_id="_TestScaleSyncLegacy", category="test", inputs=[io.Image.Input("image")], outputs=[io.Image.Output()], ) @classmethod def execute(cls, image): - # sync execute -> exercises the non-coroutine dispatch branch return io.NodeOutput(image * 0.5) @@ -77,9 +81,9 @@ def test_async_sdk_node_inverts_through_real_engine(): assert torch.allclose(got, 1.0 - img) -def test_sync_sdk_node_scales_through_real_engine(): +def test_legacy_sync_node_scales_through_real_engine(): img = torch.rand(1, 8, 8, 3) - got = _output_of(_ScaleSync, img) + got = _output_of(_ScaleSyncLegacy, img) assert torch.allclose(got, img * 0.5) @@ -107,6 +111,6 @@ if __name__ == "__main__": # Runnable without pytest. test_default_backend_is_in_process() test_async_sdk_node_inverts_through_real_engine() - test_sync_sdk_node_scales_through_real_engine() + test_legacy_sync_node_scales_through_real_engine() test_overlay_backend_intercepts_dispatch() print("PASS: all secure-SDK seam checks")