flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/syncify/_api.py

@@ -0,0 +1,414 @@
+from __future__ import annotations
+
+import asyncio
+import atexit
+import concurrent.futures
+import functools
+import inspect
+import logging
+import threading
+from typing import (
+    Any,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Coroutine,
+    Iterator,
+    ParamSpec,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from flyte._logging import logger
+
+P = ParamSpec("P")
+
+# Track loops that have already had the gRPC error handler installed
+_configured_loops: set[int] = set()
+_configured_loops_lock = threading.Lock()
+
+
+def _ensure_grpc_error_handler_installed() -> None:
+    """
+    Install the gRPC error handler on the current event loop if not already done.
+    Uses a thread-safe set to track which loops have been configured.
+    """
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        return  # No running loop, nothing to do
+
+    loop_id = id(loop)
+    with _configured_loops_lock:
+        if loop_id in _configured_loops:
+            return  # Already configured
+        _configured_loops.add(loop_id)
+
+    import flyte.errors
+
+    loop.set_exception_handler(flyte.errors.silence_grpc_polling_error)
+
+
+R_co = TypeVar("R_co", covariant=True)
+T = TypeVar("T")
+
+
+class SyncFunction(Protocol[P, R_co]):
+    """
+    A protocol that defines the interface for synchronous functions or methods that can be converted from asynchronous
+     ones.
+    """
+
+    def __call__(self, *args: Any, **kwargs: Any) -> R_co: ...
+
+    def aio(self, *args: Any, **kwargs: Any) -> Awaitable[R_co]: ...
+
+
+class SyncGenFunction(Protocol[P, R_co]):
+    """
+    A protocol that defines the interface for synchronous functions or methods that can be converted from asynchronous
+     ones.
+    """
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Iterator[R_co]: ...
+
+    def aio(self, *args: Any, **kwargs: Any) -> AsyncIterator[R_co]: ...
+
+
+class _BackgroundLoop:
+    """
+    A background event loop that runs in a separate thread and used the `Syncify` decorator to run asynchronous
+    functions or methods synchronously.
+    """
+
+    def __init__(self, name: str):
+        self.loop = asyncio.new_event_loop()
+        self.thread = threading.Thread(name=name, target=self._run, daemon=True)
+        self.thread.start()
+        atexit.register(self.stop)
+
+    def _run(self):
+        # Delay import to avoid deadlock during module initialization
+        def exception_handler(loop, context):
+            import flyte.errors
+
+            flyte.errors.silence_grpc_polling_error(loop, context)
+
+        # Set the exception handler to silence specific gRPC polling errors
+        self.loop.set_exception_handler(exception_handler)
+        asyncio.set_event_loop(self.loop)
+        self.loop.run_forever()
+
+    def stop(self):
+        # stop the loop and wait briefly for thread to exit
+        self.loop.call_soon_threadsafe(self.loop.stop)
+        self.thread.join(timeout=1)
+
+    def is_in_loop(self) -> bool:
+        """
+        Check if the current thread is the background loop thread.
+        """
+        # If the current thread is not the background loop thread, return False
+        if threading.current_thread() != self.thread:
+            return False
+
+        if not self.thread.is_alive():
+            # If the thread is not alive, we cannot be in the loop
+            return False
+
+        # Lets get the current event loop and check if it matches the background loop
+        loop = None
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            pass
+
+        return loop == self.loop
+
+    def iterate_in_loop_sync(self, async_gen: AsyncIterator[R_co]) -> Iterator[R_co]:
+        # Create an iterator that pulls items from the async generator
+        assert self.thread.name != threading.current_thread().name, (
+            f"Cannot run coroutine in the same thread {self.thread.name}"
+        )
+        while True:
+            try:
+                # use __anext__() and cast to Coroutine so mypy is happy
+                future: concurrent.futures.Future[R_co] = asyncio.run_coroutine_threadsafe(
+                    cast(Coroutine[Any, Any, R_co], async_gen.__anext__()),
+                    self.loop,
+                )
+                yield future.result()
+            except (StopAsyncIteration, StopIteration):
+                break
+            except Exception as e:
+                if logger.getEffectiveLevel() > logging.DEBUG:
+                    # If the log level is not DEBUG, we will remove the extra stack frames to avoid confusion for the
+                    # user
+                    # This is because the stack trace will include the Syncify wrapper and the background loop thread
+                    tb = e.__traceback__
+                    while tb and tb.tb_next:
+                        if tb.tb_frame.f_code.co_name == "":
+                            break
+                        tb = tb.tb_next
+                    raise e.with_traceback(tb)
+                # If the log level is DEBUG, we will keep the extra stack frames to help with debugging
+                raise e
+
+    def call_in_loop_sync(self, coro: Coroutine[Any, Any, R_co]) -> R_co | Iterator[R_co]:
+        """
+        Run the given coroutine in the background loop and return its result.
+        """
+        future: concurrent.futures.Future[R_co | AsyncIterator[R_co]] = asyncio.run_coroutine_threadsafe(
+            coro, self.loop
+        )
+        result = future.result()
+        if result is not None and hasattr(result, "__aiter__"):
+            # If the result is an async iterator, we need to convert it to a sync iterator
+            return cast(Iterator[R_co], self.iterate_in_loop_sync(cast(AsyncIterator[R_co], result)))
+        # Otherwise, just return the result
+        return result
+
+    async def iterate_in_loop(self, async_gen: AsyncIterator[R_co]) -> AsyncIterator[R_co]:
+        """
+        Run the given async iterator in the background loop and yield its results.
+        """
+        if self.is_in_loop():
+            # If we are already in the background loop, just return the async iterator
+            async for r in async_gen:
+                yield r
+            return
+
+        # Install the gRPC error handler on the caller's event loop as well.
+        # This is needed because gRPC's async polling events may be delivered to
+        # the caller's loop (e.g., FastAPI's event loop) when using .aio().
+        _ensure_grpc_error_handler_installed()
+
+        while True:
+            try:
+                # same replacement here for the async path
+                future: concurrent.futures.Future[R_co] = asyncio.run_coroutine_threadsafe(
+                    cast(Coroutine[Any, Any, R_co], async_gen.__anext__()),
+                    self.loop,
+                )
+                # Wrap the future in an asyncio Future to yield it in an async context
+                aio_future: asyncio.Future[R_co] = asyncio.wrap_future(future)
+                # await for the future to complete and yield its result
+                v = await aio_future
+                yield v
+            except StopAsyncIteration:
+                break
+            except Exception as e:
+                if logger.getEffectiveLevel() > logging.DEBUG:
+                    # If the log level is not DEBUG, we will remove the extra stack frames to avoid confusion for the
+                    # user.
+                    # This is because the stack trace will include the Syncify wrapper and the background loop thread
+                    tb = e.__traceback__
+                    while tb and tb.tb_next:
+                        if tb.tb_frame.f_code.co_name == "":
+                            break
+                        tb = tb.tb_next
+                    raise e.with_traceback(tb)
+                # If the log level is DEBUG, we will keep the extra stack frames to help with debugging
+                raise e
+
+    async def aio(self, coro: Coroutine[Any, Any, R_co]) -> R_co:
+        """
+        Run the given coroutine in the background loop and return its result.
+        """
+        if self.is_in_loop():
+            # If we are already in the background loop, just run the coroutine
+            return await coro
+        try:
+            # Install the gRPC error handler on the caller's event loop as well.
+            # This is needed because gRPC's async polling events may be delivered to
+            # the caller's loop (e.g., FastAPI's event loop) when using .aio().
+            _ensure_grpc_error_handler_installed()
+
+            # Otherwise, run it in the background loop and wait for the result
+            future: concurrent.futures.Future[R_co] = asyncio.run_coroutine_threadsafe(coro, self.loop)
+            # Wrap the future in an asyncio Future to await it in an async context
+            aio_future: asyncio.Future[R_co] = asyncio.wrap_future(future)
+            # await for the future to complete and return its result
+            return await aio_future
+        except Exception as e:
+            if logger.getEffectiveLevel() > logging.DEBUG:
+                # If the log level is not DEBUG, we will remove the extra stack frames to avoid confusion for the user
+                # This is because the stack trace will include the Syncify wrapper and the background loop thread
+                tb = e.__traceback__
+                while tb and tb.tb_next:
+                    if tb.tb_frame.f_code.co_name == "":
+                        break
+                    tb = tb.tb_next
+                raise e.with_traceback(tb)
+            # If the log level is DEBUG, we will keep the extra stack frames to help with debugging
+            raise e
+
+
+class _SyncWrapper:
+    """
+    A wrapper class that the Syncify decorator uses to convert asynchronous functions or methods into synchronous ones.
+    """
+
+    def __init__(
+        self,
+        fn: Any,
+        bg_loop: _BackgroundLoop,
+        underlying_obj: Any = None,
+    ):
+        self.fn = fn
+        self._bg_loop = bg_loop
+        self._underlying_obj = underlying_obj
+
+    def __get__(self, instance: Any, owner: Any) -> Any:
+        """
+        This method is called when the wrapper is accessed as a method of a class instance.
+        :param instance:
+        :param owner:
+        :return:
+        """
+        fn: Any = self.fn
+        if instance is not None:
+            # If we have an instance, we need to bind the method to the instance (for instance methods)
+            fn = self.fn.__get__(instance, owner)
+
+        if instance is None and owner is not None and self._underlying_obj is not None:
+            # If we have an owner, we need to bind the method to the owner (for classmethods or staticmethods)
+            fn = self._underlying_obj.__get__(None, owner)
+
+        wrapper = _SyncWrapper(fn, bg_loop=self._bg_loop, underlying_obj=self._underlying_obj)
+        functools.update_wrapper(wrapper, self.fn)
+        return wrapper
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        if threading.current_thread().name == self._bg_loop.thread.name:
+            # If we are already in the background loop thread, we can call the function directly
+            raise AssertionError(
+                f"Deadlock detected: blocking call used in syncify thread {self._bg_loop.thread.name} "
+                f"when calling function {self.fn}, use .aio() if in an async call."
+            )
+        try:
+            # bind method if needed
+            coro_fn = self.fn
+
+            if inspect.isasyncgenfunction(coro_fn):
+                # Handle async iterator by converting to sync iterator
+                async_gen = coro_fn(*args, **kwargs)
+                return self._bg_loop.iterate_in_loop_sync(async_gen)
+            else:
+                return self._bg_loop.call_in_loop_sync(coro_fn(*args, **kwargs))
+        except Exception as e:
+            if logger.getEffectiveLevel() > logging.DEBUG:
+                # If the log level is not DEBUG, we will remove the extra stack frames to avoid confusion for the user
+                # This is because the stack trace will include the Syncify wrapper and the background loop thread
+                tb = e.__traceback__
+                while tb and tb.tb_next:
+                    if tb.tb_frame.f_code.co_name == self.fn.__name__:
+                        break
+                    tb = tb.tb_next
+                raise e.with_traceback(tb)
+            # If the log level is DEBUG, we will keep the extra stack frames to help with debugging
+            raise e
+
+    def aio(self, *args: Any, **kwargs: Any) -> Any:
+        fn = self.fn
+
+        try:
+            if inspect.isasyncgenfunction(fn):
+                # If the function is an async generator, we need to handle it differently
+                async_iter = fn(*args, **kwargs)
+                return self._bg_loop.iterate_in_loop(async_iter)
+            else:
+                # If we are already in the background loop, just return the coroutine
+                coro = fn(*args, **kwargs)
+                if hasattr(coro, "__aiter__"):
+                    # If the coroutine is an async iterator, we need to handle it differently
+                    return self._bg_loop.iterate_in_loop(coro)
+                return self._bg_loop.aio(coro)
+        except Exception as e:
+            if logger.getEffectiveLevel() > logging.DEBUG:
+                # If the log level is not DEBUG, we will remove the extra stack frames to avoid confusion for the user
+                # This is because the stack trace will include the Syncify wrapper and the background loop thread
+                tb = e.__traceback__
+                while tb and tb.tb_next:
+                    if tb.tb_frame.f_code.co_name == self.fn.__name__:
+                        break
+                    tb = tb.tb_next
+                raise e.with_traceback(tb)
+            # If the log level is DEBUG, we will keep the extra stack frames to help with debugging
+            raise e
+
+
+class Syncify:
+    """
+    A decorator to convert asynchronous functions or methods into synchronous ones.
+
+    This is useful for integrating async code into synchronous contexts.
+
+    Example::
+
+    ```python
+    syncer = Syncify()
+
+    @syncer
+    async def async_function(x: str) -> str:
+        return f"Hello, Async World {x}!"
+
+
+    # now you can call it synchronously
+    result = async_function("Async World")
+    print(result)
+    # Output: Hello, Async World Async World!
+
+    # or call it asynchronously
+    async def main():
+        result = await async_function.aio("World")
+        print(result)
+    ```
+
+    """
+
+    def __init__(self, name: str = "flyte_syncify"):
+        self._bg_loop = _BackgroundLoop(name=name)
+
+    @overload
+    def __call__(self, func: Callable[P, Awaitable[R_co]]) -> SyncFunction[P, R_co]: ...
+
+    @overload
+    def __call__(self, func: Callable[P, Iterator[R_co] | AsyncIterator[R_co]]) -> SyncGenFunction[P, R_co]: ...
+
+    # def __call__(self, func: Callable[[Type[T], *P.args, *P.kwargs], Awaitable[R_co]])
+    # -> SyncFunction[[Type[T], *P.args, *P.kwargs], R_co]: ...
+    @overload
+    def __call__(self, func: classmethod) -> Union[SyncFunction[P, R_co], SyncGenFunction[P, R_co]]: ...
+
+    @overload
+    def __call__(self, func: staticmethod) -> staticmethod: ...
+
+    def __call__(self, obj):
+        if isinstance(obj, classmethod):
+            wrapper = _SyncWrapper(obj.__func__, bg_loop=self._bg_loop, underlying_obj=obj)
+            functools.update_wrapper(wrapper, obj.__func__)
+            return wrapper
+
+        if isinstance(obj, staticmethod):
+            fn = obj.__func__
+            wrapper = _SyncWrapper(fn, bg_loop=self._bg_loop)
+            functools.update_wrapper(wrapper, fn)
+            return staticmethod(wrapper)
+
+        if inspect.isasyncgenfunction(obj):
+            wrapper = _SyncWrapper(obj, bg_loop=self._bg_loop)
+            functools.update_wrapper(wrapper, obj)
+            return cast(Callable[P, Iterator[R_co]], wrapper)
+
+        if inspect.iscoroutinefunction(obj):
+            wrapper = _SyncWrapper(obj, bg_loop=self._bg_loop)
+            functools.update_wrapper(wrapper, obj)
+            return wrapper
+
+        raise TypeError(
+            "Syncify can only be applied to async functions, async generators, async classmethods or staticmethods."
+        )

flyte/types/__init__.py

@@ -1,9 +1,36 @@
+"""
+# Flyte Type System
+
+The Flyte type system provides a way to define, transform, and manipulate types in Flyte workflows.
+Since the data flowing through Flyte has to often cross process, container and langauge boundaries, the type system
+is designed to be serializable to a universal format that can be understood across different environments. This
+universal format is based on Protocol Buffers. The types are called LiteralTypes and the runtime
+representation of data is called Literals.
+
+The type system includes:
+- **TypeEngine**: The core engine that manages type transformations and serialization. This is the main entry point for
+  for all the internal type transformations and serialization logic.
+- **TypeTransformer**: A class that defines how to transform one type to another. This is extensible
+    allowing users to define custom types and transformations.
+- **Renderable**: An interface for types that can be rendered as HTML, that can be outputted to a flyte.report.
+
+It is always possible to bypass the type system and use the `FlytePickle` type to serialize any python object
+ into a pickle format. The pickle format is not human-readable, but can be passed between flyte tasks that are
+ written in python. The Pickled objects cannot be represented in the UI, and may be in-efficient for large datasets.
+"""
+
+from importlib.metadata import entry_points
+
+from flyte._logging import logger
+
 from ._interface import guess_interface
+from ._pickle import FlytePickle
 from ._renderer import Renderable
 from ._string_literals import literal_string_repr
 from ._type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError
 
 __all__ = [
+    "FlytePickle",
     "Renderable",
     "TypeEngine",
     "TypeTransformer",
@@ -11,3 +38,15 @@ __all__ = [
     "guess_interface",
     "literal_string_repr",
 ]
+
+
+def _load_custom_type_transformers():
+    plugins = entry_points(group="flyte.plugins.types")
+    for ep in plugins:
+        try:
+            logger.info(f"Loading type transformer: {ep.name}")
+            loaded = ep.load()
+            if callable(loaded):
+                loaded()
+        except Exception as e:
+            logger.warning(f"Failed to load type transformer {ep.name} with error: {e}")

flyte/types/_interface.py

@@ -1,11 +1,15 @@
-from typing import Any, Dict, Type, cast
+import inspect
+from typing import Any, Dict, Iterable, Tuple, Type, cast
 
-from flyteidl.core import interface_pb2
+from flyteidl2.core import interface_pb2, literals_pb2
+from flyteidl2.task import common_pb2
 
-from flyte._datastructures import NativeInterface
+from flyte.models import NativeInterface
 
 
-def guess_interface(interface: interface_pb2.TypedInterface) -> NativeInterface:
+def guess_interface(
+    interface: interface_pb2.TypedInterface, default_inputs: Iterable[common_pb2.NamedParameter] | None = None
+) -> NativeInterface:
     """
     Returns the interface of the task with guessed types, as types may not be present in current env.
     """
@@ -14,12 +18,23 @@ def guess_interface(interface: interface_pb2.TypedInterface) -> NativeInterface:
     if interface is None:
         return NativeInterface({}, {})
 
-    guessed_inputs: Dict[str, Type[Any]] = {}
+    default_input_literals: Dict[str, literals_pb2.Literal] = {}
+    if default_inputs is not None:
+        for param in default_inputs:
+            if param.parameter.HasField("default"):
+                default_input_literals[param.name] = param.parameter.default
+
+    guessed_inputs: Dict[str, Tuple[Type[Any], Any] | Any] = {}
     if interface.inputs is not None and len(interface.inputs.variables) > 0:
-        guessed_inputs = flyte.types.TypeEngine.guess_python_types(cast(dict, interface.inputs.variables))
+        input_types = flyte.types.TypeEngine.guess_python_types(cast(dict, interface.inputs.variables))
+        for name, t in input_types.items():
+            if name not in default_input_literals:
+                guessed_inputs[name] = (t, inspect.Parameter.empty)
+            else:
+                guessed_inputs[name] = (t, NativeInterface.has_default)
 
     guessed_outputs: Dict[str, Type[Any]] = {}
     if interface.outputs is not None and len(interface.outputs.variables) > 0:
         guessed_outputs = flyte.types.TypeEngine.guess_python_types(cast(dict, interface.outputs.variables))
 
-    return NativeInterface.from_types(guessed_inputs, guessed_outputs)
+    return NativeInterface.from_types(guessed_inputs, guessed_outputs, default_input_literals)

flyte/{io/pickle/transformer.py → types/_pickle.py}

@@ -1,16 +1,19 @@
 import hashlib
 import os
+import sys
 import typing
 from typing import Type
 
 import aiofiles
 import cloudpickle
-from flyteidl.core import literals_pb2, types_pb2
+from flyteidl2.core import literals_pb2, types_pb2
 
 import flyte.storage as storage
-from flyte.types import TypeEngine, TypeTransformer
+
+from ._type_engine import TypeEngine, TypeTransformer
 
 T = typing.TypeVar("T")
+DEFAULT_PICKLE_BYTES_LIMIT = 2**10 * 10  # 10KB
 
 
 class FlytePickle(typing.Generic[T]):
@@ -75,8 +78,13 @@ class FlytePickleTransformer(TypeTransformer[FlytePickle]):
         ...
 
     async def to_python_value(self, lv: literals_pb2.Literal, expected_python_type: Type[T]) -> T:
-        uri = lv.scalar.blob.uri
-        return await FlytePickle.from_pickle(uri)
+        if lv.scalar.blob and lv.scalar.blob.uri:
+            uri = lv.scalar.blob.uri
+            return await FlytePickle.from_pickle(uri)
+        elif lv.scalar.binary and lv.scalar.binary.value:
+            return cloudpickle.loads(lv.scalar.binary.value)
+        else:
+            raise ValueError(f"Cannot convert {lv} to {expected_python_type}")
 
     async def to_literal(
         self,
@@ -91,8 +99,15 @@ class FlytePickleTransformer(TypeTransformer[FlytePickle]):
                 format=self.PYTHON_PICKLE_FORMAT, dimensionality=types_pb2.BlobType.BlobDimensionality.SINGLE
             )
         )
-        remote_path = await FlytePickle.to_pickle(python_val)
-        return literals_pb2.Literal(scalar=literals_pb2.Scalar(blob=literals_pb2.Blob(metadata=meta, uri=remote_path)))
+        if sys.getsizeof(python_val) > DEFAULT_PICKLE_BYTES_LIMIT:
+            remote_path = await FlytePickle.to_pickle(python_val)
+            return literals_pb2.Literal(
+                scalar=literals_pb2.Scalar(blob=literals_pb2.Blob(metadata=meta, uri=remote_path))
+            )
+        else:
+            return literals_pb2.Literal(
+                scalar=literals_pb2.Scalar(binary=literals_pb2.Binary(value=cloudpickle.dumps(python_val)))
+            )
 
     def guess_python_type(self, literal_type: types_pb2.LiteralType) -> typing.Type[FlytePickle[typing.Any]]:
         if (
@@ -101,14 +116,27 @@ class FlytePickleTransformer(TypeTransformer[FlytePickle]):
             and literal_type.blob.format == FlytePickleTransformer.PYTHON_PICKLE_FORMAT
         ):
             return FlytePickle
+        if literal_type.simple == types_pb2.SimpleType.BINARY:
+            return FlytePickle
 
         raise ValueError(f"Transformer {self} cannot reverse {literal_type}")
 
     def get_literal_type(self, t: Type[T]) -> types_pb2.LiteralType:
         lt = types_pb2.LiteralType(
-            blob=types_pb2.BlobType(
-                format=self.PYTHON_PICKLE_FORMAT, dimensionality=types_pb2.BlobType.BlobDimensionality.SINGLE
-            )
+            union_type=types_pb2.UnionType(
+                variants=[
+                    types_pb2.LiteralType(
+                        blob=types_pb2.BlobType(
+                            format=self.PYTHON_PICKLE_FORMAT,
+                            dimensionality=types_pb2.BlobType.BlobDimensionality.SINGLE,
+                        ),
+                        structure=types_pb2.TypeStructure(tag=self.name),
+                    ),
+                    types_pb2.LiteralType(
+                        simple=types_pb2.SimpleType.BINARY, structure=types_pb2.TypeStructure(tag=self.name)
+                    ),
+                ]
+            ),
         )
         lt.metadata = {"python_class_name": str(t)}
         return lt

flyte/types/_string_literals.py

@@ -3,11 +3,10 @@ import json
 from typing import Any, Dict, Union
 
 import msgpack
-from flyteidl.core import literals_pb2
+from flyteidl2.core import literals_pb2
+from flyteidl2.task import common_pb2
 from google.protobuf.json_format import MessageToDict
 
-from flyte._protos.workflow import run_definition_pb2
-
 
 def _primitive_to_string(primitive: literals_pb2.Primitive) -> Any:
     """
@@ -88,9 +87,9 @@ def _dict_literal_repr(lmd: Dict[str, literals_pb2.Literal]) -> Dict[str, Any]:
 def literal_string_repr(
     lm: Union[
         literals_pb2.Literal,
-        run_definition_pb2.NamedLiteral,
-        run_definition_pb2.Inputs,
-        run_definition_pb2.Outputs,
+        common_pb2.NamedLiteral,
+        common_pb2.Inputs,
+        common_pb2.Outputs,
         literals_pb2.LiteralMap,
         Dict[str, literals_pb2.Literal],
     ],
@@ -105,13 +104,13 @@ def literal_string_repr(
             return _literal_string_repr(lm)
         case literals_pb2.LiteralMap():
             return _dict_literal_repr(lm.literals)
-        case run_definition_pb2.NamedLiteral():
+        case common_pb2.NamedLiteral():
             lmd = {lm.name: lm.value}
             return _dict_literal_repr(lmd)
-        case run_definition_pb2.Inputs():
+        case common_pb2.Inputs():
             lmd = {n.name: n.value for n in lm.literals}
             return _dict_literal_repr(lmd)
-        case run_definition_pb2.Outputs():
+        case common_pb2.Outputs():
             lmd = {n.name: n.value for n in lm.literals}
             return _dict_literal_repr(lmd)
         case dict():