mplang-nightly 0.1.dev246__tar.gz → 0.1.dev266__tar.gz

{mplang_nightly-0.1.dev246 → mplang_nightly-0.1.dev266}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev246
+Version: 0.1.dev266
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License
@@ -211,7 +211,7 @@ Requires-Dist: cryptography>=43.0.0
 Requires-Dist: duckdb>=1.0.0
 Requires-Dist: fastapi
 Requires-Dist: flax>=0.12.0
-Requires-Dist: httpx>=0.27.0
+Requires-Dist: httpx<1.0.0,>=0.27.0
 Requires-Dist: jax[cpu]==0.8.0
 Requires-Dist: lightphe<0.1.0,>=0.0.15
 Requires-Dist: numpy>=2.0.0
@@ -219,7 +219,7 @@ Requires-Dist: pandas>=2.0.0
 Requires-Dist: protobuf<6.0,>=5.0
 Requires-Dist: pyarrow>=14.0.0
 Requires-Dist: pyyaml>=6.0
-Requires-Dist: spu>=0.10.0.dev20251211
+Requires-Dist: spu>=0.10.0.dev20251208
 Requires-Dist: sqlglot>=23.0.0
 Requires-Dist: tenseal==0.3.16
 Requires-Dist: typing-extensions

{mplang_nightly-0.1.dev246 → mplang_nightly-0.1.dev266}/mplang/v2/__init__.py

@@ -52,15 +52,19 @@ from mplang.v2.edsl import (
     TracedFunction,
     Tracer,
     Value,
+    find_context,
+    find_context_with_state,
+    find_interpreter,
     format_graph,
     get_current_context,
     get_default_context,
-    get_root_context,
+    is_tracing,
     jit,
     pop_context,
     primitive,
     push_context,
     register_default_context_factory,
+    set_root_context,
     trace,
 )
 from mplang.v2.edsl.registry import get_profiler
@@ -89,6 +93,7 @@ from mplang.v2.libs.device import (
     put,
     set_dev_attr,
 )
+from mplang.v2.libs.device import fetch as device_fetch
 from mplang.v2.runtime.interpreter import Interpreter
 
 # =============================================================================
@@ -96,34 +101,6 @@ from mplang.v2.runtime.interpreter import Interpreter
 # =============================================================================
 
 
-def set_root_context(context: Interpreter, force: bool = False) -> None:
-    """Set the global/root execution context.
-
-    This explicitly sets the provided interpreter as the Root Context.
-    All subsequent operations (compile, evaluate, device resolution) will
-    use this context as the default environment.
-
-    Args:
-        context: Interpreter to use as the root context.
-        force: If True, clears the existing context stack before setting.
-               If False (default), pushes onto the stack.
-    """
-    from mplang.v2.edsl.context import _context_stack, get_current_context
-
-    if force:
-        _context_stack.clear()
-        _context_stack.append(context)
-        return
-
-    if get_current_context() is not None:
-        raise RuntimeError(
-            "Cannot set root context: Context stack is not empty. "
-            "Use force=True to overwrite the existing root context."
-        )
-
-    push_context(context)
-
-
 def _get_context(context: Interpreter | None) -> Interpreter:
     """Get context from parameter or context stack."""
     if context is not None:
@@ -186,92 +163,165 @@ def evaluate(
             inputs = fn.prepare_inputs(*args, **kwargs)
             inputs = [unwrap_if_interp(v) for v in inputs]
             raw_result = interp.evaluate_graph(fn.graph, inputs)
-            return fn.reconstruct_outputs(raw_result)
+            wrapped = [
+                InterpObject(v, fn.graph.outputs[i].type, interp)
+                for i, v in enumerate(raw_result)
+            ]
+            return fn.reconstruct_outputs(wrapped)
 
         return fn(*args, **kwargs)
 
 
 def fetch(
     result: Any,
-    party: int | str | None = None,
     *,
+    follow_device: bool = True,
     context: Interpreter | None = None,
 ) -> Any:
-    """Fetch the result, optionally for a specific party.
+    """Fetch results from interpreter context to Python.
+
+    This is a meta-function that operates at execution boundaries, not a traced
+    dialect operation. It brings data from the distributed/MPC runtime back to
+    the Python host.
+
+    Behavior in different contexts:
+    - **Tracing (compile)**: Returns the input unchanged (identity). The graph
+      outputs are determined by the function's return statement, not fetch calls.
+    - **Execution (evaluate)**: Actually fetches data from workers/parties.
+
+    Design Note (A vs B tradeoff):
+        Two designs were considered for fetch behavior during tracing:
+
+        - **Design A (chosen)**: fetch = identity during tracing. Graph outputs
+          are determined solely by the return statement. This is simpler and
+          avoids ambiguity when fetch and return reference different values.
+
+        - **Design B (alternative)**: fetch marks output points in the graph.
+          This would allow fetch(a), fetch(b), return b to output both a and b.
+          However, it complicates the semantics and requires tracking fetch
+          points separately from return values.
+
+        Design A was chosen for simplicity. If a value needs to be an output,
+        it should be returned. fetch's role is purely for execution-time I/O.
 
     Args:
-        result: The result to fetch (DriverVar or other value).
-        party: Optional party index or device name (e.g., "P0").
-        context: Optional interpreter context. If None, uses current context.
+        result: Object(s) to fetch. Can be a single InterpObject, DriverVar,
+            or nested structure containing them.
+        follow_device: If True and object has device attribute, dispatch to
+            device.fetch which fetches from the correct rank based on device.
+            If False, fetch from all parties.
+        context: Interpreter context. If None, uses current context.
 
     Returns:
-        The fetched data (unwrapped from Value wrappers).
-
-    Example:
-        >>> with mp.make_simulator(3) as sim:
-        ...     data = mp.fetch(result, "P0")  # uses sim from context
+        Fetched Python values. For device objects with follow_device=True,
+        returns single value from the device's rank(s). Otherwise returns
+        list of values (one per party) or single value for world_size=1.
+        During tracing, returns the input unchanged.
     """
-    from typing import cast
+    from jax.tree_util import tree_map
 
-    from mplang.v2.backends.simp_driver.state import SimpDriver
     from mplang.v2.backends.simp_driver.values import DriverVar
-    from mplang.v2.backends.table_impl import TableValue
-    from mplang.v2.backends.tensor_impl import TensorValue
+    from mplang.v2.edsl.context import is_tracing
     from mplang.v2.runtime.interpreter import InterpObject
+    from mplang.v2.runtime.value import WrapValue
 
-    interp = _get_context(context)
+    # Check if we are in tracing context - if so, return identity
+    if is_tracing():
+        # Design A: fetch = identity during tracing
+        # Graph outputs are determined by return statement, not fetch calls
+        return result
 
-    def _unwrap_value(val: Any) -> Any:
-        """Unwrap Value types to get the underlying data."""
-        if isinstance(val, TensorValue):
-            return val.data
-        elif isinstance(val, TableValue):
-            return val.data
-        return val
+    # Execution context - actually fetch data
+    interp = _get_context(context)
 
-    # Unwrap InterpObject to get the runtime value
-    if isinstance(result, InterpObject):
-        result = result.runtime_obj
-
-    # Get simp state for fetching
-    simp_state = cast(SimpDriver | None, interp.get_dialect_state("simp"))
-    cluster_spec = getattr(interp, "_cluster_spec", None)
-
-    # Fetch from DriverVar
-    if isinstance(result, DriverVar):
-        resolved_values = []
-        for rank, val in enumerate(result.values):
-            if isinstance(val, str) and "://" in val:
-                if simp_state is not None:
-                    fut = simp_state.fetch(rank, val)
-                    resolved_values.append(fut.result())
-                else:
-                    resolved_values.append(val)
-            else:
-                resolved_values.append(val)
-
-        # Select party if needed
-        if party is not None:
-            if isinstance(party, str) and cluster_spec is not None:
-                device_info = cluster_spec.devices.get(party)
-                if device_info and device_info.members:
-                    party = device_info.members[0].rank
+    def _fetch_single(var: Any) -> Any:
+        """Fetch a single value from InterpObject."""
+        # InterpObject (from mp.evaluate) - extract runtime_obj
+        if isinstance(var, InterpObject):
+            if follow_device and is_device_obj(var):
+                return device_fetch(var)
+            var = var.runtime_obj  # extract and continue processing
+
+        # DriverVar (simp dialect) - remote fetch from workers
+        if isinstance(var, DriverVar):
+            from mplang.v2.backends.simp_driver.state import SimpDriver
+
+            simp_state = interp.get_dialect_state("simp")
+            assert isinstance(simp_state, SimpDriver), "DriverVar requires simp state"
+
+            resolved: list[Any] = []
+            for rank, uri in enumerate(var.values):
+                if uri is None:
+                    resolved.append(None)
                 else:
-                    raise ValueError(f"Unknown party: {party}")
-            else:
-                # Default logic for int
-                pass
+                    fetched = simp_state.fetch(rank, uri).result()
+                    if isinstance(fetched, WrapValue):
+                        fetched = fetched.data
+                    resolved.append(fetched)
 
-            p_idx = cast(int, party)
-            return _unwrap_value(resolved_values[p_idx])  # type: ignore[no-any-return]
-        return [_unwrap_value(v) for v in resolved_values]
+            return resolved[0] if len(resolved) == 1 else resolved
 
-    # Unwrap Value types to get the underlying data
-    return _unwrap_value(result)
+        # WrapValue (TensorValue, TableValue, etc.) - unwrap
+        if isinstance(var, WrapValue):
+            return var.data
+
+        # Plain values pass through
+        return var
+
+    with interp:
+        return tree_map(_fetch_single, result)
 
 
 # Alias for compatibility
-function = jit  # @mp.function -> @mp2.function (JIT compilation)
+def function(fn: Callable[..., Any] | None = None) -> Callable[..., Any]:
+    """Decorator defining a Multi-Party Function (MP Program).
+
+    This decorator "lifts" a local function into a distributed program by
+    automatically wrapping it in a `simp.pcall_static` that targets ALL available
+    parties in the current context.
+
+    Semantics: f(args) -> pcall(ALL, f, args)
+
+    Args:
+        fn: The function to decorate.
+
+    Returns:
+        A wrapper function that, when called, executes the original function
+        on all workers.
+    """
+    import functools
+
+    from mplang.v2.dialects import simp
+
+    if fn is None:
+        return function
+
+    def has_simp_state(ctx: Any) -> bool:
+        if hasattr(ctx, "get_dialect_state"):
+            state = ctx.get_dialect_state("simp")
+            return state is not None and hasattr(state, "world_size")
+        return False
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        from mplang.v2.edsl.context import find_context
+
+        # Find context with simp dialect state
+        ctx = find_context(has_simp_state)
+        if ctx is None:
+            raise RuntimeError(
+                "mp.function requires a context with world_size information "
+                "(e.g. SimpSimulator or Driver initialized)."
+            )
+
+        # ctx found by predicate so we know it has get_dialect_state
+        simp_state = ctx.get_dialect_state("simp")  # type: ignore[attr-defined]
+        world_size = simp_state.world_size  # type: ignore
+
+        all_parties = tuple(range(world_size))
+        return simp.pcall_static(all_parties, fn, *args, **kwargs)
+
+    return wrapper
 
 
 def compile(
@@ -334,10 +384,14 @@ __all__ = [  # noqa: RUF022
     "compile",
     "evaluate",
     "fetch",
+    "find_context",
+    "find_context_with_state",
+    "find_interpreter",
     "format_graph",
     "function",
     "get_current_context",
     "get_default_context",
+    "is_tracing",
     "jit",
     "mplang",
     "pop_context",
@@ -363,7 +417,6 @@ __all__ = [  # noqa: RUF022
     # Dialects
     "dialects",
     "register_default_context_factory",
-    "get_root_context",
     "get_profiler",
 ]
 

{mplang_nightly-0.1.dev246 → mplang_nightly-0.1.dev266}/mplang/v2/backends/crypto_impl.py

@@ -271,6 +271,27 @@ def bytes_to_point_impl(
 
 # --- Sym / Hash Impl ---
 
+# Supported symmetric encryption algorithms
+_SUPPORTED_ALGOS = {"aes-gcm"}
+
+
+def _validate_algo(algo: str, operation: str) -> None:
+    """Validate that the algorithm is supported.
+
+    Args:
+        algo: Algorithm name to validate
+        operation: Operation name for error message (e.g., "encryption", "decryption")
+
+    Raises:
+        ValueError: If algo is not supported
+    """
+    if algo not in _SUPPORTED_ALGOS:
+        supported = ", ".join(sorted(_SUPPORTED_ALGOS))
+        raise ValueError(
+            f"Unsupported {operation} algorithm: {algo!r}. "
+            f"Supported algorithms: {supported}"
+        )
+
 
 @crypto.hash_p.def_impl
 def hash_impl(interpreter: Interpreter, op: Operation, data: Value) -> Value:
@@ -336,6 +357,10 @@ def sym_encrypt_impl(
     numpy arrays, scalars, etc.). This supports both high-level API usage
     (with TensorValue) and elementwise operations (with raw scalars).
     """
+    # Read and validate algo parameter (must be provided by frontend)
+    algo = op.attrs["algo"]
+    _validate_algo(algo, "encryption")
+
     # Get raw key bytes - strict type checking
     if isinstance(key, SymmetricKeyValue):
         k = key.key_bytes
@@ -376,6 +401,10 @@ def sym_decrypt_impl(
     on what was encrypted - could be a Value subclass (TensorValue, BytesValue),
     a numpy array, or a scalar (int, float, etc.) when used in elementwise ops.
     """
+    # Read and validate algo parameter (must be provided by frontend)
+    algo = op.attrs["algo"]
+    _validate_algo(algo, "decryption")
+
     # Get raw key bytes - strict type checking
     if isinstance(key, SymmetricKeyValue):
         k = key.key_bytes
@@ -585,6 +614,101 @@ def kem_derive_impl(
         return SymmetricKeyValue(suite=suite, key_bytes=secret)
 
 
+@crypto.hkdf_p.def_impl
+def hkdf_impl(
+    interpreter: Interpreter,
+    op: Operation,
+    secret: SymmetricKeyValue | TensorValue,
+) -> SymmetricKeyValue:
+    """HKDF key derivation implementation using SHA-256.
+
+    Implements RFC 5869 HKDF with HMAC-SHA256. This is the NIST SP 800-56C
+    compliant way to derive symmetric keys from ECDH shared secrets.
+
+    Current implementation supports only SHA-256. Future versions will add
+    SHA-512, SHA3-256, and BLAKE2b support.
+
+    Security Notes:
+        - Uses salt=None (defaults to 32-byte all-zero salt per RFC 5869)
+        - ONLY SAFE for high-entropy IKM (e.g., 256-bit ECDH shared secrets)
+        - NOT suitable for: passwords, low-entropy secrets, or repeated key derivations
+        - For session keys with same ECDH pair: use unique 'info' per session
+
+        Per NIST SP 800-56C Rev. 2:
+        "If the IKM is already cryptographically strong (e.g., from ECDH),
+         a salt may not be necessary, but using one does not hurt."
+
+    Args:
+        interpreter: Runtime interpreter context
+        op: Operation node containing attributes (info, hash_algo)
+        secret: Input key material (IKM) as SymmetricKeyValue or TensorValue
+                Must be high-entropy (≥256 bits) for security with salt=None
+
+    Returns:
+        SymmetricKeyValue with suite="hkdf-{hash_algo}" and 32-byte key_bytes
+
+    Raises:
+        TypeError: If secret is not SymmetricKeyValue or TensorValue
+        ValueError: If info parameter is empty (required for domain separation)
+        NotImplementedError: If hash_algo is not "sha256"
+    """
+    from cryptography.hazmat.primitives import hashes
+    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+
+    # Extract operation attributes
+    info_str = op.attrs.get("info", "")
+    hash_algo = (
+        op.attrs.get("hash_algo", "sha256").lower().replace("-", "").replace("_", "")
+    )
+
+    # Validate info parameter (REQUIRED for domain separation per NIST)
+    if not info_str:
+        raise ValueError(
+            "HKDF requires non-empty 'info' parameter for domain separation. "
+            "The info string binds the derived key to a specific protocol/context. "
+            "Recommended format: 'namespace/component/purpose/version'"
+        )
+
+    info_bytes = info_str.encode("utf-8")
+
+    # Extract input key material (IKM) bytes
+    if isinstance(secret, SymmetricKeyValue):
+        ikm = secret.key_bytes
+    elif isinstance(secret, TensorValue):
+        ikm = secret.unwrap().tobytes()
+    else:
+        raise TypeError(
+            f"hkdf secret must be SymmetricKeyValue or TensorValue, "
+            f"got {type(secret).__name__}"
+        )
+
+    # Validate hash algorithm (currently only SHA-256 implemented)
+    if hash_algo != "sha256":
+        raise NotImplementedError(
+            f"HKDF with hash algorithm '{hash_algo}' is not yet implemented. "
+            f"Currently only 'sha256' is supported. "
+            f"Planned future support: sha512, sha3256, blake2b"
+        )
+
+    # Perform HKDF derivation using cryptography library
+    # Note: salt=None uses 32-byte all-zero salt (not random salt!)
+    # This is secure ONLY because ECDH outputs are already high-entropy (256-bit uniform)
+    # For low-entropy inputs or repeated derivations, a random salt would be required
+    hkdf = HKDF(
+        algorithm=hashes.SHA256(),
+        length=32,  # Output length in bytes (AES-256 key = 32 bytes)
+        salt=None,  # 32-byte zero salt (secure for high-entropy ECDH shared secrets)
+        info=info_bytes,  # Context-specific binding for domain separation
+    )
+
+    derived_key = hkdf.derive(ikm)
+
+    # Return SymmetricKeyValue with composite suite name
+    # Format: "hkdf-{hash_algo}" to indicate derivation method and hash function
+    suite = f"hkdf-{hash_algo}"
+    return SymmetricKeyValue(suite=suite, key_bytes=derived_key)
+
+
 @crypto.random_bytes_p.def_impl
 def random_bytes_impl(interpreter: Interpreter, op: Operation) -> TensorValue:
     """Generate random bytes using os.urandom."""

{mplang_nightly-0.1.dev246 → mplang_nightly-0.1.dev266}/mplang/v2/backends/func_impl.py

@@ -102,4 +102,6 @@ def _call_impl(
         raise TypeError(f"func.call expects FunctionValue, got {type(fn_obj)}")
 
     call_args = list(args)
-    return interpreter.evaluate_graph(fn_obj.graph, call_args)
+    result = interpreter.evaluate_graph(fn_obj.graph, call_args)
+    # Return single value or list based on graph outputs
+    return result[0] if len(fn_obj.graph.outputs) == 1 else result

{mplang_nightly-0.1.dev246 → mplang_nightly-0.1.dev266}/mplang/v2/backends/simp_driver/http.py

@@ -26,6 +26,7 @@ import httpx
 from mplang.v2.backends.simp_driver.state import SimpDriver
 from mplang.v2.edsl import serde
 from mplang.v2.runtime.interpreter import Interpreter
+from mplang.v2.runtime.object_store import ObjectStore
 
 if TYPE_CHECKING:
     from concurrent.futures import Future
@@ -159,6 +160,7 @@ def make_driver(endpoints: list[str], *, cluster_spec: Any = None) -> Interprete
         name="DriverInterpreter",
         root_dir=state.driver_root,
         handlers=handlers,
+        store=ObjectStore(fs_root=str(state.driver_root)),
     )
     interp.set_dialect_state("simp", state)
     interp._cluster_spec = cluster_spec  # type: ignore[attr-defined]

{mplang_nightly-0.1.dev246 → mplang_nightly-0.1.dev266}/mplang/v2/backends/simp_driver/mem.py

@@ -154,6 +154,10 @@ class SimpMemDriver(SimpDriver):
         self._workers = workers
         self._mesh = mesh
 
+    def shutdown(self) -> None:
+        """Shutdown the local memory driver and its mesh."""
+        self._mesh.shutdown()
+
     @property
     def world_size(self) -> int:
         return self._world_size
@@ -201,27 +205,20 @@ class SimpMemDriver(SimpDriver):
         worker_interp = self._workers[rank]
         worker_ctx = cast(SimpWorker, worker_interp.get_dialect_state("simp"))
 
-        # Resolve URI inputs
-        resolved_inputs = []
-        for inp in inputs:
-            if isinstance(inp, str) and "://" in inp:
-                resolved_inputs.append(worker_ctx.store.get(inp))
-            else:
-                resolved_inputs.append(inp)
+        # Resolve URI inputs (None means rank has no data)
+        resolved_inputs = [
+            worker_ctx.store.get(inp) if inp is not None else None for inp in inputs
+        ]
 
         # Execute
         results = worker_interp.evaluate_graph(graph, resolved_inputs, job_id)
 
-        # Store results
+        # Store results (results is always a list)
         if not graph.outputs:
             return None
-        if len(graph.outputs) == 1:
-            val = results
-            if val is None:
-                return None
-            return worker_ctx.store.put(val)
-        else:
-            return [worker_ctx.store.put(res) for res in results]
+        return [
+            worker_ctx.store.put(res) if res is not None else None for res in results
+        ]
 
 
 def make_simulator(
@@ -268,6 +265,7 @@ def make_simulator(
         root_dir=cluster.host_root,
         handlers=handlers,
         tracer=cluster.tracer,
+        store=ObjectStore(fs_root=str(cluster.host_root)),
     )
     interp.set_dialect_state("simp", state)
 

{mplang_nightly-0.1.dev246 → mplang_nightly-0.1.dev266}/mplang/v2/backends/simp_driver/ops.py

@@ -68,7 +68,8 @@ def _collect_to_hostvars(results: list[Any], num_outputs: int, world_size: int)
     """Collect worker results into DriverVar(s).
 
     Args:
-        results: List of results from each worker (length = world_size)
+        results: List of results from each worker (length = world_size).
+                 Each result is a list of URIs (one per output).
         num_outputs: Number of outputs per worker
         world_size: Total number of workers
 
@@ -78,15 +79,16 @@ def _collect_to_hostvars(results: list[Any], num_outputs: int, world_size: int)
     if num_outputs == 0:
         return None
 
-    if num_outputs == 1:
-        return DriverVar(results)
-
-    # Multiple outputs: transpose [worker][output] -> [output][worker]
+    # Transpose [worker][output] -> [output][worker]
+    # results[worker_idx] is a list of URIs for that worker's outputs
     transposed = []
-    for i in range(num_outputs):
+    for out_idx in range(num_outputs):
         transposed.append(
-            DriverVar([res[i] if res is not None else None for res in results])
+            DriverVar([res[out_idx] if res is not None else None for res in results])
         )
+
+    if num_outputs == 1:
+        return transposed[0]
     return transposed
 
 

{mplang_nightly-0.1.dev246 → mplang_nightly-0.1.dev266}/mplang/v2/backends/simp_worker/http.py

@@ -49,9 +49,9 @@ from pydantic import BaseModel
 
 from mplang.v2.backends import spu_impl as _spu_impl  # noqa: F401
 from mplang.v2.backends import tensor_impl as _tensor_impl  # noqa: F401
-from mplang.v2.backends.simp_worker import SimpWorker
 
 # Register operation implementations (side-effect imports)
+from mplang.v2.backends.simp_worker import SimpWorker
 from mplang.v2.backends.simp_worker import ops as _simp_worker_ops  # noqa: F401
 from mplang.v2.edsl import serde
 from mplang.v2.edsl.graph import Graph
@@ -232,7 +232,9 @@ def create_worker_app(
     # func_impl is already imported at module level for side-effects
     handlers: dict[str, Callable[..., Any]] = {**WORKER_HANDLERS}  # type: ignore[dict-item]
 
-    worker = Interpreter(tracer=tracer, root_dir=root_dir, handlers=handlers)
+    worker = Interpreter(
+        tracer=tracer, root_dir=root_dir, handlers=handlers, store=store
+    )
     # Register SimpWorker context as 'simp' dialect state
     worker.set_dialect_state("simp", ctx)
 
@@ -242,11 +244,18 @@ def create_worker_app(
 
     def _do_execute(graph: Graph, inputs: list[Any], job_id: str | None = None) -> Any:
         """Execute graph in worker thread."""
-        # Note: we ignore job_id for now as generic Interpreter doesn't take it.
-        # But we could set it in thread local if needed.
-        result = worker.evaluate_graph(graph, inputs)
+        # Resolve URI inputs (None means rank has no data)
+        resolved_inputs = [
+            store.get(inp) if inp is not None else None for inp in inputs
+        ]
+
+        result = worker.evaluate_graph(graph, resolved_inputs)
         comm.wait_pending_sends()
-        return result
+
+        # Store results and return URIs (result is always a list)
+        if not graph.outputs:
+            return None
+        return [store.put(res) if res is not None else None for res in result]
 
     @app.post("/exec")
     async def execute(req: ExecRequest) -> dict[str, str]: