flyte 0.2.0b35__py3-none-any.whl → 0.2.0b37__py3-none-any.whl

flyte/_protos/workflow/state_service_pb2.pyi

@@ -1,3 +1,4 @@
+from flyte._protos.common import identifier_pb2 as _identifier_pb2
 from flyteidl.core import execution_pb2 as _execution_pb2
 from google.rpc import status_pb2 as _status_pb2
 from flyte._protos.validate.validate import validate_pb2 as _validate_pb2
@@ -13,40 +14,40 @@ class PutRequest(_message.Message):
     ACTION_ID_FIELD_NUMBER: _ClassVar[int]
     PARENT_ACTION_NAME_FIELD_NUMBER: _ClassVar[int]
     STATE_FIELD_NUMBER: _ClassVar[int]
-    action_id: _run_definition_pb2.ActionIdentifier
+    action_id: _identifier_pb2.ActionIdentifier
     parent_action_name: str
     state: str
-    def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ..., parent_action_name: _Optional[str] = ..., state: _Optional[str] = ...) -> None: ...
+    def __init__(self, action_id: _Optional[_Union[_identifier_pb2.ActionIdentifier, _Mapping]] = ..., parent_action_name: _Optional[str] = ..., state: _Optional[str] = ...) -> None: ...
 
 class PutResponse(_message.Message):
     __slots__ = ["action_id", "status"]
     ACTION_ID_FIELD_NUMBER: _ClassVar[int]
     STATUS_FIELD_NUMBER: _ClassVar[int]
-    action_id: _run_definition_pb2.ActionIdentifier
+    action_id: _identifier_pb2.ActionIdentifier
     status: _status_pb2.Status
-    def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ..., status: _Optional[_Union[_status_pb2.Status, _Mapping]] = ...) -> None: ...
+    def __init__(self, action_id: _Optional[_Union[_identifier_pb2.ActionIdentifier, _Mapping]] = ..., status: _Optional[_Union[_status_pb2.Status, _Mapping]] = ...) -> None: ...
 
 class GetRequest(_message.Message):
     __slots__ = ["action_id"]
     ACTION_ID_FIELD_NUMBER: _ClassVar[int]
-    action_id: _run_definition_pb2.ActionIdentifier
-    def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ...) -> None: ...
+    action_id: _identifier_pb2.ActionIdentifier
+    def __init__(self, action_id: _Optional[_Union[_identifier_pb2.ActionIdentifier, _Mapping]] = ...) -> None: ...
 
 class GetResponse(_message.Message):
     __slots__ = ["action_id", "status", "state"]
     ACTION_ID_FIELD_NUMBER: _ClassVar[int]
     STATUS_FIELD_NUMBER: _ClassVar[int]
     STATE_FIELD_NUMBER: _ClassVar[int]
-    action_id: _run_definition_pb2.ActionIdentifier
+    action_id: _identifier_pb2.ActionIdentifier
     status: _status_pb2.Status
     state: str
-    def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ..., status: _Optional[_Union[_status_pb2.Status, _Mapping]] = ..., state: _Optional[str] = ...) -> None: ...
+    def __init__(self, action_id: _Optional[_Union[_identifier_pb2.ActionIdentifier, _Mapping]] = ..., status: _Optional[_Union[_status_pb2.Status, _Mapping]] = ..., state: _Optional[str] = ...) -> None: ...
 
 class WatchRequest(_message.Message):
     __slots__ = ["parent_action_id"]
     PARENT_ACTION_ID_FIELD_NUMBER: _ClassVar[int]
-    parent_action_id: _run_definition_pb2.ActionIdentifier
-    def __init__(self, parent_action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ...) -> None: ...
+    parent_action_id: _identifier_pb2.ActionIdentifier
+    def __init__(self, parent_action_id: _Optional[_Union[_identifier_pb2.ActionIdentifier, _Mapping]] = ...) -> None: ...
 
 class WatchResponse(_message.Message):
     __slots__ = ["action_update", "control_message"]
@@ -68,8 +69,8 @@ class ActionUpdate(_message.Message):
     PHASE_FIELD_NUMBER: _ClassVar[int]
     ERROR_FIELD_NUMBER: _ClassVar[int]
     OUTPUT_URI_FIELD_NUMBER: _ClassVar[int]
-    action_id: _run_definition_pb2.ActionIdentifier
+    action_id: _identifier_pb2.ActionIdentifier
     phase: _run_definition_pb2.Phase
     error: _execution_pb2.ExecutionError
     output_uri: str
-    def __init__(self, action_id: _Optional[_Union[_run_definition_pb2.ActionIdentifier, _Mapping]] = ..., phase: _Optional[_Union[_run_definition_pb2.Phase, str]] = ..., error: _Optional[_Union[_execution_pb2.ExecutionError, _Mapping]] = ..., output_uri: _Optional[str] = ...) -> None: ...
+    def __init__(self, action_id: _Optional[_Union[_identifier_pb2.ActionIdentifier, _Mapping]] = ..., phase: _Optional[_Union[_run_definition_pb2.Phase, str]] = ..., error: _Optional[_Union[_execution_pb2.ExecutionError, _Mapping]] = ..., output_uri: _Optional[str] = ...) -> None: ...

flyte/_run.py

@@ -17,6 +17,7 @@ from flyte._initialize import (
     requires_storage,
 )
 from flyte._logging import logger
+from flyte._protos.common import identifier_pb2
 from flyte._task import P, R, TaskTemplate
 from flyte._tools import ipython_check
 from flyte.errors import InitializationError
@@ -185,7 +186,7 @@ class _Runner:
             run_id = None
             project_id = None
             if self._name:
-                run_id = run_definition_pb2.RunIdentifier(
+                run_id = identifier_pb2.RunIdentifier(
                     project=project,
                     domain=domain,
                     org=cfg.org,
@@ -260,9 +261,9 @@ class _Runner:
                 super().__init__(
                     pb2=run_definition_pb2.Run(
                         action=run_definition_pb2.Action(
-                            id=run_definition_pb2.ActionIdentifier(
+                            id=identifier_pb2.ActionIdentifier(
                                 name="a0",
-                                run=run_definition_pb2.RunIdentifier(name="dry-run"),
+                                run=identifier_pb2.RunIdentifier(name="dry-run"),
                             )
                         )
                     )
@@ -422,9 +423,9 @@ class _Runner:
                 super().__init__(
                     pb2=run_definition_pb2.Run(
                         action=run_definition_pb2.Action(
-                            id=run_definition_pb2.ActionIdentifier(
+                            id=identifier_pb2.ActionIdentifier(
                                 name="a0",
-                                run=run_definition_pb2.RunIdentifier(name="dry-run"),
+                                run=identifier_pb2.RunIdentifier(name="dry-run"),
                             )
                         )
                     )
@@ -434,7 +435,11 @@ class _Runner:
             def url(self) -> str:
                 return "local-run"
 
-            def wait(self, quiet: bool = False, wait_for: Literal["terminal", "running"] = "terminal"):
+            def wait(
+                self,
+                quiet: bool = False,
+                wait_for: Literal["terminal", "running"] = "terminal",
+            ):
                 pass
 
             def outputs(self) -> R:

flyte/_trace.py

@@ -1,7 +1,6 @@
 import functools
 import inspect
 import time
-from datetime import timedelta
 from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, TypeGuard, TypeVar, Union, cast
 
 from flyte.models import NativeInterface
@@ -43,14 +42,12 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
                 # Cast to Awaitable to satisfy mypy
                 coroutine_result = cast(Awaitable[Any], func(*args, **kwargs))
                 results = await coroutine_result
-                duration = time.time() - start_time
-                info.add_outputs(results, timedelta(seconds=duration))
+                info.add_outputs(results, start_time=start_time, end_time=time.time())
                 await controller.record_trace(info)
                 return results
             except Exception as e:
                 # If there is an error, we need to record it
-                duration = time.time() - start_time
-                info.add_error(e, timedelta(seconds=duration))
+                info.add_error(e, start_time=start_time, end_time=time.time())
                 await controller.record_trace(info)
                 raise e
         else:
@@ -93,14 +90,11 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
                     async for item in async_iter:
                         items.append(item)
                         yield item
-                duration = time.time() - start_time
-                info.add_outputs(items, timedelta(seconds=duration))
+                info.add_outputs(items, start_time=start_time, end_time=time.time())
                 await controller.record_trace(info)
                 return
             except Exception as e:
-                end_time = time.time()
-                duration = end_time - start_time
-                info.add_error(e, timedelta(seconds=duration))
+                info.add_error(e, start_time=start_time, end_time=time.time())
                 await controller.record_trace(info)
                 raise e
         else:

flyte/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0b35'
-__version_tuple__ = version_tuple = (0, 2, 0, 'b35')
+__version__ = version = '0.2.0b37'
+__version_tuple__ = version_tuple = (0, 2, 0, 'b37')

flyte/migrate/__init__.py

@@ -0,0 +1 @@
+from flyte.migrate import dynamic, task, workflow  # noqa: F401

flyte/migrate/dynamic.py

@@ -0,0 +1,13 @@
+from typing import Callable, Union
+
+import flytekit
+
+import flyte.migrate
+from flyte._task import AsyncFunctionTaskTemplate, P, R
+
+
+def dynamic_shim(**kwargs) -> Union[AsyncFunctionTaskTemplate, Callable[P, R]]:
+    return flyte.migrate.task.task_shim(**kwargs)
+
+
+flytekit.dynamic = dynamic_shim

flyte/migrate/task.py

@@ -0,0 +1,99 @@
+import datetime
+from typing import Callable, Dict, Iterable, List, Literal, Optional, Tuple, Union
+
+import flytekit
+from flytekit.core import launch_plan, workflow
+from flytekit.core.base_task import T, TaskResolverMixin
+from flytekit.core.python_function_task import PythonFunctionTask
+from flytekit.core.task import FuncOut
+from flytekit.deck import DeckField
+from flytekit.extras.accelerators import BaseAccelerator
+
+import flyte
+from flyte import Image, Resources, TaskEnvironment
+from flyte._doc import Documentation
+from flyte._task import AsyncFunctionTaskTemplate, P, R
+
+
+def task_shim(
+    _task_function: Optional[Callable[P, FuncOut]] = None,
+    task_config: Optional[T] = None,
+    cache: Union[bool, flytekit.Cache] = False,
+    retries: int = 0,
+    interruptible: Optional[bool] = None,
+    deprecated: str = "",
+    timeout: Union[datetime.timedelta, int] = 0,
+    container_image: Optional[Union[str, flytekit.ImageSpec]] = None,
+    environment: Optional[Dict[str, str]] = None,
+    requests: Optional[flytekit.Resources] = None,
+    limits: Optional[flytekit.Resources] = None,
+    secret_requests: Optional[List[flytekit.Secret]] = None,
+    execution_mode: PythonFunctionTask.ExecutionBehavior = PythonFunctionTask.ExecutionBehavior.DEFAULT,
+    node_dependency_hints: Optional[
+        Iterable[
+            Union[
+                flytekit.PythonFunctionTask,
+                launch_plan.LaunchPlan,
+                workflow.WorkflowBase,
+            ]
+        ]
+    ] = None,
+    task_resolver: Optional[TaskResolverMixin] = None,
+    docs: Optional[flytekit.Documentation] = None,
+    disable_deck: Optional[bool] = None,
+    enable_deck: Optional[bool] = None,
+    deck_fields: Optional[Tuple[DeckField, ...]] = (
+        DeckField.SOURCE_CODE,
+        DeckField.DEPENDENCIES,
+        DeckField.TIMELINE,
+        DeckField.INPUT,
+        DeckField.OUTPUT,
+    ),
+    pod_template: Optional[flytekit.PodTemplate] = None,
+    pod_template_name: Optional[str] = None,
+    accelerator: Optional[BaseAccelerator] = None,
+    pickle_untyped: bool = False,
+    shared_memory: Optional[Union[Literal[True], str]] = None,
+    resources: Optional[Resources] = None,
+    labels: Optional[dict[str, str]] = None,
+    annotations: Optional[dict[str, str]] = None,
+    **kwargs,
+) -> Union[AsyncFunctionTaskTemplate, Callable[P, R]]:
+    plugin_config = task_config
+    pod_template = (
+        flyte.PodTemplate(
+            pod_spec=pod_template.pod_spec,
+            primary_container_name=pod_template.primary_container_name,
+            labels=pod_template.labels,
+            annotations=pod_template.annotations,
+        )
+        if pod_template
+        else None
+    )
+
+    if isinstance(container_image, flytekit.ImageSpec):
+        image = Image.from_debian_base()
+        if container_image.apt_packages:
+            image = image.with_apt_packages(*container_image.apt_packages)
+        pip_packages = ["flytekit"]
+        if container_image.packages:
+            pip_packages.extend(container_image.packages)
+        image = image.with_pip_packages(*pip_packages)
+    elif isinstance(container_image, str):
+        image = Image.from_base(container_image).with_pip_packages("flyte")
+    else:
+        image = Image.from_debian_base().with_pip_packages("flytekit")
+
+    docs = Documentation(description=docs.short_description) if docs else None
+
+    env = TaskEnvironment(
+        name="flytekit",
+        resources=Resources(cpu=0.8, memory="800Mi"),
+        image=image,
+        cache="enabled" if cache else "disable",
+        plugin_config=plugin_config,
+    )
+    return env.task(retries=retries, pod_template=pod_template_name or pod_template, docs=docs)
+
+
+flytekit.task = task_shim

flyte/migrate/workflow.py

@@ -0,0 +1,13 @@
+import flytekit
+
+from flyte import Image, Resources, TaskEnvironment
+
+env = TaskEnvironment(
+    name="flytekit",
+    resources=Resources(cpu=0.8, memory="800Mi"),
+    image=Image.from_debian_base().with_apt_packages("vim").with_pip_packages("flytekit", "pandas"),
+)
+
+# TODO: Build subtask's image
+
+flytekit.workflow = env.task

flyte/remote/_action.py

@@ -4,18 +4,29 @@ import asyncio
 from collections import UserDict
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from typing import Any, AsyncGenerator, AsyncIterator, Dict, Iterator, List, Literal, Tuple, Union, cast
+from typing import (
+    Any,
+    AsyncGenerator,
+    AsyncIterator,
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Tuple,
+    Union,
+    cast,
+)
 
 import grpc
 import rich.pretty
 import rich.repr
-from google.protobuf import timestamp
+from google.protobuf import timestamp_pb2
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 
 from flyte import types
 from flyte._initialize import ensure_client, get_client, get_common_config
-from flyte._protos.common import list_pb2
+from flyte._protos.common import identifier_pb2, list_pb2
 from flyte._protos.workflow import run_definition_pb2, run_service_pb2
 from flyte._protos.workflow.run_service_pb2 import WatchActionDetailsResponse
 from flyte.remote._logs import Logs
@@ -24,11 +35,13 @@ from flyte.syncify import syncify
 WaitFor = Literal["terminal", "running", "logs-ready"]
 
 
-def _action_time_phase(action: run_definition_pb2.Action | run_definition_pb2.ActionDetails) -> rich.repr.Result:
+def _action_time_phase(
+    action: run_definition_pb2.Action | run_definition_pb2.ActionDetails,
+) -> rich.repr.Result:
     """
     Rich representation of the action time and phase.
     """
-    start_time = timestamp.to_datetime(action.status.start_time, timezone.utc)
+    start_time = action.status.start_time.ToDatetime().replace(tzinfo=timezone.utc)
     yield "start_time", start_time.isoformat()
     if action.status.phase in [
         run_definition_pb2.PHASE_FAILED,
@@ -36,7 +49,7 @@ def _action_time_phase(action: run_definition_pb2.Action | run_definition_pb2.Ac
         run_definition_pb2.PHASE_ABORTED,
         run_definition_pb2.PHASE_TIMED_OUT,
     ]:
-        end_time = timestamp.to_datetime(action.status.end_time, timezone.utc)
+        end_time = action.status.end_time.ToDatetime().replace(tzinfo=timezone.utc)
         yield "end_time", end_time.isoformat()
         yield "run_time", f"{(end_time - start_time).seconds} secs"
     else:
@@ -46,7 +59,7 @@ def _action_time_phase(action: run_definition_pb2.Action | run_definition_pb2.Ac
     if isinstance(action, run_definition_pb2.ActionDetails):
         yield (
             "error",
-            f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA",
+            (f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA"),
         )
 
 
@@ -57,7 +70,10 @@ def _action_rich_repr(action: run_definition_pb2.Action) -> rich.repr.Result:
     yield "run", action.id.run.name
     if action.metadata.HasField("task"):
         yield "task", action.metadata.task.id.name
-        yield "type", "task"
+        yield "type", action.metadata.task.task_type
+    elif action.metadata.HasField("trace"):
+        yield "trace", action.metadata.trace.name
+        yield "type", "trace"
     yield "name", action.id.name
     yield from _action_time_phase(action)
     yield "group", action.metadata.group
@@ -65,14 +81,18 @@ def _action_rich_repr(action: run_definition_pb2.Action) -> rich.repr.Result:
     yield "attempts", action.status.attempts
 
 
-def _attempt_rich_repr(action: List[run_definition_pb2.ActionAttempt]) -> rich.repr.Result:
+def _attempt_rich_repr(
+    action: List[run_definition_pb2.ActionAttempt],
+) -> rich.repr.Result:
     for attempt in action:
         yield "attempt", attempt.attempt
         yield "phase", run_definition_pb2.Phase.Name(attempt.phase)
         yield "logs_available", attempt.logs_available
 
 
-def _action_details_rich_repr(action: run_definition_pb2.ActionDetails) -> rich.repr.Result:
+def _action_details_rich_repr(
+    action: run_definition_pb2.ActionDetails,
+) -> rich.repr.Result:
     """
     Rich representation of the action details.
     """
@@ -82,7 +102,7 @@ def _action_details_rich_repr(action: run_definition_pb2.ActionDetails) -> rich.
     yield "task_type", action.resolved_task_spec.task_template.type
     yield "task_version", action.resolved_task_spec.task_template.id.version
     yield "attempts", action.attempts
-    yield "error", f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA"
+    yield "error", (f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA")
     yield "phase", run_definition_pb2.Phase.Name(action.status.phase)
     yield "group", action.metadata.group
     yield "parent", action.metadata.parent
@@ -129,7 +149,8 @@ class Action:
         token = None
         sort_by = sort_by or ("created_at", "asc")
         sort_pb2 = list_pb2.Sort(
-            key=sort_by[0], direction=list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING
+            key=sort_by[0],
+            direction=(list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING),
         )
         cfg = get_common_config()
         while True:
@@ -141,7 +162,7 @@ class Action:
             resp = await get_client().run_service.ListActions(
                 run_service_pb2.ListActionsRequest(
                     request=req,
-                    run_id=run_definition_pb2.RunIdentifier(
+                    run_id=identifier_pb2.RunIdentifier(
                         org=cfg.org,
                         project=cfg.project,
                         domain=cfg.domain,
@@ -157,7 +178,13 @@ class Action:
 
     @syncify
     @classmethod
-    async def get(cls, uri: str | None = None, /, run_name: str | None = None, name: str | None = None) -> Action:
+    async def get(
+        cls,
+        uri: str | None = None,
+        /,
+        run_name: str | None = None,
+        name: str | None = None,
+    ) -> Action:
         """
         Get a run by its ID or name. If both are provided, the ID will take precedence.
 
@@ -168,8 +195,8 @@ class Action:
         ensure_client()
         cfg = get_common_config()
         details: ActionDetails = await ActionDetails.get_details.aio(
-            run_definition_pb2.ActionIdentifier(
-                run=run_definition_pb2.RunIdentifier(
+            identifier_pb2.ActionIdentifier(
+                run=identifier_pb2.RunIdentifier(
                     org=cfg.org,
                     project=cfg.project,
                     domain=cfg.domain,
@@ -225,7 +252,7 @@ class Action:
         return None
 
     @property
-    def action_id(self) -> run_definition_pb2.ActionIdentifier:
+    def action_id(self) -> identifier_pb2.ActionIdentifier:
         """
         Get the action ID.
         """
@@ -396,7 +423,7 @@ class ActionDetails:
 
     @syncify
     @classmethod
-    async def get_details(cls, action_id: run_definition_pb2.ActionIdentifier) -> ActionDetails:
+    async def get_details(cls, action_id: identifier_pb2.ActionIdentifier) -> ActionDetails:
         """
         Get the details of the action. This is a placeholder for getting the action details.
         """
@@ -411,7 +438,11 @@ class ActionDetails:
     @syncify
     @classmethod
     async def get(
-        cls, uri: str | None = None, /, run_name: str | None = None, name: str | None = None
+        cls,
+        uri: str | None = None,
+        /,
+        run_name: str | None = None,
+        name: str | None = None,
     ) -> ActionDetails:
         """
         Get a run by its ID or name. If both are provided, the ID will take precedence.
@@ -425,8 +456,8 @@ class ActionDetails:
             assert name is not None and run_name is not None, "Either uri or name and run_name must be provided"
         cfg = get_common_config()
         return await cls.get_details.aio(
-            run_definition_pb2.ActionIdentifier(
-                run=run_definition_pb2.RunIdentifier(
+            identifier_pb2.ActionIdentifier(
+                run=identifier_pb2.RunIdentifier(
                     org=cfg.org,
                     project=cfg.project,
                     domain=cfg.domain,
@@ -438,7 +469,7 @@ class ActionDetails:
 
     @syncify
     @classmethod
-    async def watch(cls, action_id: run_definition_pb2.ActionIdentifier) -> AsyncIterator[ActionDetails]:
+    async def watch(cls, action_id: identifier_pb2.ActionIdentifier) -> AsyncIterator[ActionDetails]:
         """
         Watch the action for updates. This is a placeholder for watching the action.
         """
@@ -521,7 +552,7 @@ class ActionDetails:
         return None
 
     @property
-    def action_id(self) -> run_definition_pb2.ActionIdentifier:
+    def action_id(self) -> identifier_pb2.ActionIdentifier:
         """
         Get the action ID.
         """
@@ -552,9 +583,9 @@ class ActionDetails:
         """
         Get the runtime of the action.
         """
-        start_time = timestamp.to_datetime(self.pb2.status.start_time, timezone.utc)
+        start_time = self.pb2.status.start_time.ToDatetime().replace(tzinfo=timezone.utc)
         if self.pb2.status.HasField("end_time"):
-            end_time = timestamp.to_datetime(self.pb2.status.end_time, timezone.utc)
+            end_time = self.pb2.status.end_time.ToDatetime().replace(tzinfo=timezone.utc)
             return end_time - start_time
         return datetime.now(timezone.utc) - start_time
 

flyte/remote/_logs.py

@@ -11,8 +11,9 @@ from rich.text import Text
 
 from flyte._initialize import ensure_client, get_client
 from flyte._logging import logger
+from flyte._protos.common import identifier_pb2
 from flyte._protos.logs.dataplane import payload_pb2
-from flyte._protos.workflow import run_definition_pb2, run_logs_service_pb2
+from flyte._protos.workflow import run_logs_service_pb2
 from flyte._tools import ipython_check, ipywidgets_check
 from flyte.errors import LogsNotYetAvailableError
 from flyte.syncify import syncify
@@ -98,7 +99,7 @@ class Logs:
     @classmethod
     async def tail(
         cls,
-        action_id: run_definition_pb2.ActionIdentifier,
+        action_id: identifier_pb2.ActionIdentifier,
         attempt: int = 1,
         retry: int = 3,
     ) -> AsyncGenerator[payload_pb2.LogLine, None]:
@@ -139,7 +140,7 @@ class Logs:
     @classmethod
     async def create_viewer(
         cls,
-        action_id: run_definition_pb2.ActionIdentifier,
+        action_id: identifier_pb2.ActionIdentifier,
         attempt: int = 1,
         max_lines: int = 30,
         show_ts: bool = False,

flyte/remote/_run.py

@@ -53,7 +53,8 @@ class Run:
         token = None
         sort_by = sort_by or ("created_at", "asc")
         sort_pb2 = list_pb2.Sort(
-            key=sort_by[0], direction=list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING
+            key=sort_by[0],
+            direction=(list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING),
         )
         cfg = get_common_config()
         while True:
@@ -229,7 +230,7 @@ class RunDetails:
 
     @syncify
     @classmethod
-    async def get_details(cls, run_id: run_definition_pb2.RunIdentifier) -> RunDetails:
+    async def get_details(cls, run_id: identifier_pb2.RunIdentifier) -> RunDetails:
         """
         Get the details of the run. This is a placeholder for getting the run details.
         """
@@ -253,7 +254,7 @@ class RunDetails:
         ensure_client()
         cfg = get_common_config()
         return await RunDetails.get_details.aio(
-            run_id=run_definition_pb2.RunIdentifier(
+            run_id=identifier_pb2.RunIdentifier(
                 org=cfg.org,
                 project=cfg.project,
                 domain=cfg.domain,
@@ -276,7 +277,7 @@ class RunDetails:
         return self.action_details.task_name
 
     @property
-    def action_id(self) -> run_definition_pb2.ActionIdentifier:
+    def action_id(self) -> identifier_pb2.ActionIdentifier:
         """
         Get the action ID.
         """