flyte 2.0.0b23__py3-none-any.whl → 2.0.0b25__py3-none-any.whl

flyte/remote/_secret.py

@@ -4,9 +4,9 @@ from dataclasses import dataclass
 from typing import AsyncIterator, Literal, Union
 
 import rich.repr
+from flyteidl2.secret import definition_pb2, payload_pb2
 
-from flyte._initialize import ensure_client, get_client, get_common_config
-from flyte._protos.secret import definition_pb2, payload_pb2
+from flyte._initialize import ensure_client, get_client, get_init_config
 from flyte.remote._common import ToJSONMixin
 from flyte.syncify import syncify
 
@@ -21,12 +21,19 @@ class Secret(ToJSONMixin):
     @classmethod
     async def create(cls, name: str, value: Union[str, bytes], type: SecretTypes = "regular"):
         ensure_client()
-        cfg = get_common_config()
-        secret_type = (
-            definition_pb2.SecretType.SECRET_TYPE_GENERIC
-            if type == "regular"
-            else definition_pb2.SecretType.SECRET_TYPE_IMAGE_PULL_SECRET
-        )
+        cfg = get_init_config()
+        project = cfg.project
+        domain = cfg.domain
+
+        if type == "regular":
+            secret_type = definition_pb2.SecretType.SECRET_TYPE_GENERIC
+
+        else:
+            secret_type = definition_pb2.SecretType.SECRET_TYPE_IMAGE_PULL_SECRET
+            if project or domain:
+                raise ValueError(
+                    f"Project `{project}` or domain `{domain}` should not be set when creating the image pull secret."
+                )
 
         if isinstance(value, str):
             secret = definition_pb2.SecretSpec(
@@ -42,8 +49,8 @@ class Secret(ToJSONMixin):
             request=payload_pb2.CreateSecretRequest(
                 id=definition_pb2.SecretIdentifier(
                     organization=cfg.org,
-                    project=cfg.project,
-                    domain=cfg.domain,
+                    project=project,
+                    domain=domain,
                     name=name,
                 ),
                 secret_spec=secret,
@@ -54,7 +61,7 @@ class Secret(ToJSONMixin):
     @classmethod
     async def get(cls, name: str) -> Secret:
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         resp = await get_client().secrets_service.GetSecret(
             request=payload_pb2.GetSecretRequest(
                 id=definition_pb2.SecretIdentifier(
@@ -71,7 +78,7 @@ class Secret(ToJSONMixin):
     @classmethod
     async def listall(cls, limit: int = 100) -> AsyncIterator[Secret]:
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         token = None
         while True:
             resp = await get_client().secrets_service.ListSecrets(  # type: ignore
@@ -93,7 +100,7 @@ class Secret(ToJSONMixin):
     @classmethod
     async def delete(cls, name):
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         await get_client().secrets_service.DeleteSecret(  # type: ignore
             request=payload_pb2.DeleteSecretRequest(
                 id=definition_pb2.SecretIdentifier(

flyte/remote/_task.py

@@ -6,19 +6,18 @@ from dataclasses import dataclass
 from typing import Any, AsyncIterator, Callable, Coroutine, Dict, Iterator, Literal, Optional, Tuple, Union, cast
 
 import rich.repr
-from flyteidl.core import literals_pb2
-from google.protobuf import timestamp
+from flyteidl2.common import identifier_pb2, list_pb2
+from flyteidl2.core import literals_pb2
+from flyteidl2.task import task_definition_pb2, task_service_pb2
 
 import flyte
 import flyte.errors
 from flyte._cache.cache import CacheBehavior
 from flyte._context import internal_ctx
-from flyte._initialize import ensure_client, get_client, get_common_config
+from flyte._initialize import ensure_client, get_client, get_init_config
 from flyte._internal.runtime.resources_serde import get_proto_resources
 from flyte._internal.runtime.task_serde import get_proto_retry_strategy, get_proto_timeout, get_security_context
 from flyte._logging import logger
-from flyte._protos.common import identifier_pb2, list_pb2
-from flyte._protos.workflow import task_definition_pb2, task_service_pb2
 from flyte.models import NativeInterface
 from flyte.syncify import syncify
 
@@ -35,7 +34,7 @@ def _repr_task_metadata(metadata: task_definition_pb2.TaskMetadata) -> rich.repr
         else:
             yield "deployed_by", f"App: {metadata.deployed_by.application.spec.name}"
     yield "short_name", metadata.short_name
-    yield "deployed_at", timestamp.to_datetime(metadata.deployed_at)
+    yield "deployed_at", metadata.deployed_at.ToDatetime()
     yield "environment_name", metadata.environment_name
 
 
@@ -151,7 +150,7 @@ class TaskDetails(ToJSONMixin):
                     if ctx is None:
                         raise ValueError("auto_version=current can only be used within a task context.")
                     _version = ctx.version
-            cfg = get_common_config()
+            cfg = get_init_config()
             task_id = task_definition_pb2.TaskIdentifier(
                 org=cfg.org,
                 project=project or cfg.project,
@@ -451,7 +450,7 @@ class Task(ToJSONMixin):
         sort_pb2 = list_pb2.Sort(
             key=sort_by[0], direction=list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING
         )
-        cfg = get_common_config()
+        cfg = get_init_config()
         filters = []
         if by_task_name:
             filters.append(

flyte/remote/_trigger.py

@@ -5,12 +5,13 @@ from functools import cached_property
 from typing import AsyncIterator
 
 import grpc.aio
+from flyteidl2.common import identifier_pb2, list_pb2
+from flyteidl2.task import common_pb2, task_definition_pb2
+from flyteidl2.trigger import trigger_definition_pb2, trigger_service_pb2
 
 import flyte
-from flyte._initialize import ensure_client, get_client, get_common_config
+from flyte._initialize import ensure_client, get_client, get_init_config
 from flyte._internal.runtime import trigger_serde
-from flyte._protos.common import identifier_pb2, list_pb2
-from flyte._protos.workflow import common_pb2, task_definition_pb2, trigger_definition_pb2, trigger_service_pb2
 from flyte.syncify import syncify
 
 from ._common import ToJSONMixin
@@ -28,7 +29,7 @@ class TriggerDetails(ToJSONMixin):
         Retrieve detailed information about a specific trigger by its name.
         """
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         resp = await get_client().trigger_service.GetTriggerDetails(
             request=trigger_service_pb2.GetTriggerDetailsRequest(
                 name=identifier_pb2.TriggerName(
@@ -101,7 +102,7 @@ class Trigger(ToJSONMixin):
         :param task_name: Optional name of the task to associate with the trigger.
         """
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
 
         # Fetch the task to ensure it exists and to get its input definitions
         try:
@@ -121,15 +122,12 @@ class Trigger(ToJSONMixin):
 
             resp = await get_client().trigger_service.DeployTrigger(
                 request=trigger_service_pb2.DeployTriggerRequest(
-                    id=identifier_pb2.TriggerIdentifier(
-                        name=identifier_pb2.TriggerName(
-                            name=trigger.name,
-                            task_name=task_name,
-                            org=cfg.org,
-                            project=cfg.project,
-                            domain=cfg.domain,
-                        ),
-                        revision=1,
+                    name=identifier_pb2.TriggerName(
+                        name=trigger.name,
+                        task_name=task_name,
+                        org=cfg.org,
+                        project=cfg.project,
+                        domain=cfg.domain,
                     ),
                     spec=trigger_definition_pb2.TriggerSpec(
                         active=task_trigger.spec.active,
@@ -155,7 +153,7 @@ class Trigger(ToJSONMixin):
         """
         Retrieve a trigger by its name and associated task name.
         """
-        return await TriggerDetails.get(name=name, task_name=task_name)
+        return await TriggerDetails.get.aio(name=name, task_name=task_name)
 
     @syncify
     @classmethod
@@ -166,9 +164,9 @@ class Trigger(ToJSONMixin):
         List all triggers associated with a specific task or all tasks if no task name is provided.
         """
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         token = None
-        # task_name_id = None  TODO: implement listing by task name only
+        task_name_id = None
         project_id = None
         task_id = None
         if task_name and task_version:
@@ -179,13 +177,13 @@ class Trigger(ToJSONMixin):
                 org=cfg.org,
                 version=task_version,
             )
-        # elif task_name:  TODO: implement listing by task name only
-        #     task_name_id = task_definition_pb2.TaskName(
-        #         name=task_name,
-        #         project=cfg.project,
-        #         domain=cfg.domain,
-        #         org=cfg.org,
-        #     )
+        elif task_name:
+            task_name_id = task_definition_pb2.TaskName(
+                name=task_name,
+                project=cfg.project,
+                domain=cfg.domain,
+                org=cfg.org,
+            )
         else:
             project_id = identifier_pb2.ProjectIdentifier(
                 organization=cfg.org,
@@ -198,7 +196,7 @@ class Trigger(ToJSONMixin):
                 request=trigger_service_pb2.ListTriggersRequest(
                     project_id=project_id,
                     task_id=task_id,
-                    # task_name=task_name_id,
+                    task_name=task_name_id,
                     request=list_pb2.ListRequest(
                         limit=limit,
                         token=token,
@@ -218,7 +216,7 @@ class Trigger(ToJSONMixin):
         Pause a trigger by its name and associated task name.
         """
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         await get_client().trigger_service.UpdateTriggers(
             request=trigger_service_pb2.UpdateTriggersRequest(
                 names=[
@@ -241,7 +239,7 @@ class Trigger(ToJSONMixin):
         Delete a trigger by its name.
         """
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         await get_client().trigger_service.DeleteTriggers(
             request=trigger_service_pb2.DeleteTriggersRequest(
                 names=[

flyte/storage/_parallel_reader.py

@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+import io
+import os
+import pathlib
+import sys
+import tempfile
+import typing
+from typing import Any, Hashable, Protocol
+
+import aiofiles
+import aiofiles.os
+import obstore
+
+if typing.TYPE_CHECKING:
+    from obstore import Bytes, ObjectMeta
+    from obstore.store import ObjectStore
+
+CHUNK_SIZE = int(os.getenv("FLYTE_IO_CHUNK_SIZE", str(16 * 1024 * 1024)))
+MAX_CONCURRENCY = int(os.getenv("FLYTE_IO_MAX_CONCURRENCY", str(32)))
+
+
+class DownloadQueueEmpty(RuntimeError):
+    pass
+
+
+class BufferProtocol(Protocol):
+    async def write(self, offset, length, value: Bytes) -> None: ...
+
+    async def read(self) -> memoryview: ...
+
+    @property
+    def complete(self) -> bool: ...
+
+
+@dataclasses.dataclass
+class _MemoryBuffer:
+    arr: bytearray
+    pending: int
+    _closed: bool = False
+
+    async def write(self, offset: int, length: int, value: Bytes) -> None:
+        self.arr[offset : offset + length] = memoryview(value)
+        self.pending -= length
+
+    async def read(self) -> memoryview:
+        return memoryview(self.arr)
+
+    @property
+    def complete(self) -> bool:
+        return self.pending == 0
+
+    @classmethod
+    def new(cls, size):
+        return cls(arr=bytearray(size), pending=size)
+
+
+@dataclasses.dataclass
+class _FileBuffer:
+    path: pathlib.Path
+    pending: int
+    _handle: io.FileIO | None = None
+    _closed: bool = False
+
+    async def write(self, offset: int, length: int, value: Bytes) -> None:
+        async with aiofiles.open(self.path, mode="r+b") as f:
+            await f.seek(offset)
+            await f.write(value)
+        self.pending -= length
+
+    async def read(self) -> memoryview:
+        async with aiofiles.open(self.path, mode="rb") as f:
+            return memoryview(await f.read())
+
+    @property
+    def complete(self) -> bool:
+        return self.pending == 0
+
+    @classmethod
+    def new(cls, path: pathlib.Path, size: int):
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.touch()
+        return cls(path=path, pending=size)
+
+
+@dataclasses.dataclass
+class Chunk:
+    offset: int
+    length: int
+
+
+@dataclasses.dataclass
+class Source:
+    id: Hashable
+    path: pathlib.Path  # Should be str, represents the fully qualified prefix of a file (no bucket)
+    length: int
+    metadata: Any | None = None
+
+
+@dataclasses.dataclass
+class DownloadTask:
+    source: Source
+    chunk: Chunk
+    target: pathlib.Path | None = None
+
+
+class ObstoreParallelReader:
+    def __init__(
+        self,
+        store: ObjectStore,
+        *,
+        chunk_size=CHUNK_SIZE,
+        max_concurrency=MAX_CONCURRENCY,
+    ):
+        self._store = store
+        self._chunk_size = chunk_size
+        self._max_concurrency = max_concurrency
+
+    def _chunks(self, size) -> typing.Iterator[tuple[int, int]]:
+        cs = self._chunk_size
+        for offset in range(0, size, cs):
+            length = min(cs, size - offset)
+            yield offset, length
+
+    async def _as_completed(self, gen: typing.AsyncGenerator[DownloadTask, None], transformer=None):
+        inq: asyncio.Queue = asyncio.Queue(self._max_concurrency * 2)
+        outq: asyncio.Queue = asyncio.Queue()
+        sentinel = object()
+        done = asyncio.Event()
+
+        active: dict[Hashable, _FileBuffer | _MemoryBuffer] = {}
+
+        async def _fill():
+            # Helper function to fill the input queue, this is because the generator is async because it does list/head
+            # calls on the object store which are async.
+            try:
+                counter = 0
+                async for task in gen:
+                    if task.source.id not in active:
+                        active[task.source.id] = (
+                            _FileBuffer.new(task.target, task.source.length)
+                            if task.target is not None
+                            else _MemoryBuffer.new(task.source.length)
+                        )
+                    await inq.put(task)
+                    counter += 1
+                await inq.put(sentinel)
+                if counter == 0:
+                    raise DownloadQueueEmpty
+            except asyncio.CancelledError:
+                # document why we need to swallow this
+                pass
+
+        async def _worker():
+            try:
+                while not done.is_set():
+                    task = await inq.get()
+                    if task is sentinel:
+                        inq.put_nowait(sentinel)
+                        break
+                    chunk_source_offset = task.chunk.offset
+                    buf = active[task.source.id]
+                    data_to_write = await obstore.get_range_async(
+                        self._store,
+                        str(task.source.path),
+                        start=chunk_source_offset,
+                        end=chunk_source_offset + task.chunk.length,
+                    )
+                    await buf.write(
+                        task.chunk.offset,
+                        task.chunk.length,
+                        data_to_write,
+                    )
+                    if not buf.complete:
+                        continue
+                    if transformer is not None:
+                        result = await transformer(buf)
+                    elif task.target is not None:
+                        result = task.target
+                    else:
+                        result = task.source
+                    outq.put_nowait((task.source.id, result))
+                    del active[task.source.id]
+            except asyncio.CancelledError:
+                pass
+            finally:
+                done.set()
+
+        # Yield results as they are completed
+        if sys.version_info >= (3, 11):
+            async with asyncio.TaskGroup() as tg:
+                tg.create_task(_fill())
+                for _ in range(self._max_concurrency):
+                    tg.create_task(_worker())
+                while not done.is_set():
+                    yield await outq.get()
+        else:
+            fill_task = asyncio.create_task(_fill())
+            worker_tasks = [asyncio.create_task(_worker()) for _ in range(self._max_concurrency)]
+            try:
+                while not done.is_set():
+                    yield await outq.get()
+            except Exception as e:
+                if not fill_task.done():
+                    fill_task.cancel()
+                for wt in worker_tasks:
+                    if not wt.done():
+                        wt.cancel()
+                raise e
+            finally:
+                await asyncio.gather(fill_task, *worker_tasks, return_exceptions=True)
+
+        # Drain the output queue
+        try:
+            while True:
+                yield outq.get_nowait()
+        except asyncio.QueueEmpty:
+            pass
+
+    async def download_files(
+        self, src_prefix: pathlib.Path, target_prefix: pathlib.Path, *paths, destination_file_name: str | None = None
+    ) -> None:
+        """
+        src_prefix: Prefix you want to download from in the object store, not including the bucket name, nor file name.
+                    Should be replaced with string
+        target_prefix: Local directory to download to
+        paths: Specific paths (relative to src_prefix) to download. If empty, download everything
+        """
+
+        async def _list_downloadable() -> typing.AsyncGenerator[ObjectMeta, None]:
+            if paths:
+                # For specific file paths, use async head
+                for path_ in paths:
+                    path = src_prefix / path_
+                    x = await obstore.head_async(self._store, str(path))
+                    yield x
+                return
+
+            # Use obstore.list() for recursive listing (all files in all subdirectories)
+            # obstore.list() returns an async iterator that yields batches (lists) of objects
+            async for batch in obstore.list(self._store, prefix=str(src_prefix)):
+                for obj in batch:
+                    yield obj
+
+        async def _gen(tmp_dir: str) -> typing.AsyncGenerator[DownloadTask, None]:
+            async for obj in _list_downloadable():
+                path = pathlib.Path(obj["path"])  # e.g. Path(prefix/file.txt), needs to be changed to str.
+                size = obj["size"]
+                source = Source(id=path, path=path, length=size)
+                # Strip src_prefix from path for destination
+                rel_path = path.relative_to(src_prefix)  # doesn't work on windows
+                for offset, length in self._chunks(size):
+                    yield DownloadTask(
+                        source=source,
+                        target=tmp_dir / rel_path,  # doesn't work on windows
+                        chunk=Chunk(offset, length),
+                    )
+
+        def _transform_decorator(tmp_dir: str):
+            async def _transformer(buf: _FileBuffer) -> None:
+                if len(paths) == 1 and destination_file_name is not None:
+                    target = target_prefix / destination_file_name
+                else:
+                    target = target_prefix / buf.path.relative_to(tmp_dir)
+                await aiofiles.os.makedirs(target.parent, exist_ok=True)
+                return await aiofiles.os.replace(buf.path, target)  # mv buf.path target
+
+            return _transformer
+
+        with tempfile.TemporaryDirectory() as temporary_dir:
+            async for _ in self._as_completed(_gen(temporary_dir), transformer=_transform_decorator(temporary_dir)):
+                pass

flyte/storage/_storage.py

@@ -147,12 +147,76 @@ def _get_anonymous_filesystem(from_path):
     return get_underlying_filesystem(get_protocol(from_path), anonymous=True, asynchronous=True)
 
 
+async def _get_obstore_bypass(from_path: str, to_path: str | pathlib.Path, recursive: bool = False, **kwargs) -> str:
+    from obstore.store import ObjectStore
+
+    from flyte.storage._parallel_reader import ObstoreParallelReader
+
+    fs = get_underlying_filesystem(path=from_path)
+    bucket, prefix = fs._split_path(from_path)  # pylint: disable=W0212
+    store: ObjectStore = fs._construct_store(bucket)
+
+    download_kwargs = {}
+    if "chunk_size" in kwargs:
+        download_kwargs["chunk_size"] = kwargs["chunk_size"]
+    if "max_concurrency" in kwargs:
+        download_kwargs["max_concurrency"] = kwargs["max_concurrency"]
+
+    reader = ObstoreParallelReader(store, **download_kwargs)
+    target_path = pathlib.Path(to_path) if isinstance(to_path, str) else to_path
+
+    # if recursive, just download the prefix to the target path
+    if recursive:
+        logger.debug(f"Downloading recursively {prefix=} to {target_path=}")
+        await reader.download_files(
+            prefix,
+            target_path,
+        )
+        return str(to_path)
+
+    # if not recursive, we need to split out the file name from the prefix
+    else:
+        path_for_reader = pathlib.Path(prefix).name
+        final_prefix = pathlib.Path(prefix).parent
+        logger.debug(f"Downloading single file {final_prefix=}, {path_for_reader=} to {target_path=}")
+        await reader.download_files(
+            final_prefix,
+            target_path.parent,
+            path_for_reader,
+            destination_file_name=target_path.name,
+        )
+        return str(target_path)
+
+
 async def get(from_path: str, to_path: Optional[str | pathlib.Path] = None, recursive: bool = False, **kwargs) -> str:
     if not to_path:
-        name = pathlib.Path(from_path).name
+        name = pathlib.Path(from_path).name  # may need to be adjusted for windows
         to_path = get_random_local_path(file_path_or_file_name=name)
         logger.debug(f"Storing file from {from_path} to {to_path}")
+    else:
+        # Only apply directory logic for single files (not recursive)
+        if not recursive:
+            to_path_str = str(to_path)
+            # Check for trailing separator BEFORE converting to Path (which normalizes and removes it)
+            ends_with_sep = to_path_str.endswith(os.sep)
+            to_path_obj = pathlib.Path(to_path)
+
+            # If path ends with os.sep or is an existing directory, append source filename
+            if ends_with_sep or (to_path_obj.exists() and to_path_obj.is_dir()):
+                source_filename = pathlib.Path(from_path).name  # may need to be adjusted for windows
+                to_path = to_path_obj / source_filename
+        # For recursive=True, keep to_path as-is (it's the destination directory for contents)
+
     file_system = get_underlying_filesystem(path=from_path)
+
+    # Check if we should use obstore bypass
+    if (
+        _is_obstore_supported_protocol(file_system.protocol)
+        and hasattr(file_system, "_split_path")
+        and hasattr(file_system, "_construct_store")
+    ):
+        return await _get_obstore_bypass(from_path, to_path, recursive, **kwargs)
+
     try:
         return await _get_from_filesystem(file_system, from_path, to_path, recursive=recursive, **kwargs)
     except (OSError, GenericError) as oe:
@@ -195,7 +259,7 @@ async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = F
         from flyte._context import internal_ctx
 
         ctx = internal_ctx()
-        name = pathlib.Path(from_path).name if not recursive else None  # don't pass a name for folders
+        name = pathlib.Path(from_path).name
         to_path = ctx.raw_data.get_random_remote_path(file_name=name)
 
     file_system = get_underlying_filesystem(path=to_path)

flyte/types/_interface.py

@@ -1,9 +1,9 @@
 import inspect
 from typing import Any, Dict, Iterable, Tuple, Type, cast
 
-from flyteidl.core import interface_pb2, literals_pb2
+from flyteidl2.core import interface_pb2, literals_pb2
+from flyteidl2.task import common_pb2
 
-from flyte._protos.workflow import common_pb2
 from flyte.models import NativeInterface
 
 

flyte/types/_pickle.py

@@ -6,7 +6,7 @@ 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
 

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():