flyte 0.0.1b3__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/_protos/workflow/task_service_pb2.pyi

@@ -1,8 +1,11 @@
+from flyte._protos.common import identifier_pb2 as _identifier_pb2
+from flyte._protos.common import list_pb2 as _list_pb2
 from flyte._protos.validate.validate import validate_pb2 as _validate_pb2
 from flyte._protos.workflow import task_definition_pb2 as _task_definition_pb2
+from google.protobuf.internal import containers as _containers
 from google.protobuf import descriptor as _descriptor
 from google.protobuf import message as _message
-from typing import ClassVar as _ClassVar, Mapping as _Mapping, Optional as _Optional, Union as _Union
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
 
 DESCRIPTOR: _descriptor.FileDescriptor
 
@@ -29,3 +32,28 @@ class GetTaskDetailsResponse(_message.Message):
     DETAILS_FIELD_NUMBER: _ClassVar[int]
     details: _task_definition_pb2.TaskDetails
     def __init__(self, details: _Optional[_Union[_task_definition_pb2.TaskDetails, _Mapping]] = ...) -> None: ...
+
+class ListTasksRequest(_message.Message):
+    __slots__ = ["request", "org", "project_id", "known_filters"]
+    class KnownFilter(_message.Message):
+        __slots__ = ["deployed_by"]
+        DEPLOYED_BY_FIELD_NUMBER: _ClassVar[int]
+        deployed_by: str
+        def __init__(self, deployed_by: _Optional[str] = ...) -> None: ...
+    REQUEST_FIELD_NUMBER: _ClassVar[int]
+    ORG_FIELD_NUMBER: _ClassVar[int]
+    PROJECT_ID_FIELD_NUMBER: _ClassVar[int]
+    KNOWN_FILTERS_FIELD_NUMBER: _ClassVar[int]
+    request: _list_pb2.ListRequest
+    org: str
+    project_id: _identifier_pb2.ProjectIdentifier
+    known_filters: _containers.RepeatedCompositeFieldContainer[ListTasksRequest.KnownFilter]
+    def __init__(self, request: _Optional[_Union[_list_pb2.ListRequest, _Mapping]] = ..., org: _Optional[str] = ..., project_id: _Optional[_Union[_identifier_pb2.ProjectIdentifier, _Mapping]] = ..., known_filters: _Optional[_Iterable[_Union[ListTasksRequest.KnownFilter, _Mapping]]] = ...) -> None: ...
+
+class ListTasksResponse(_message.Message):
+    __slots__ = ["tasks", "token"]
+    TASKS_FIELD_NUMBER: _ClassVar[int]
+    TOKEN_FIELD_NUMBER: _ClassVar[int]
+    tasks: _containers.RepeatedCompositeFieldContainer[_task_definition_pb2.Task]
+    token: str
+    def __init__(self, tasks: _Optional[_Iterable[_Union[_task_definition_pb2.Task, _Mapping]]] = ..., token: _Optional[str] = ...) -> None: ...

flyte/_protos/workflow/task_service_pb2_grpc.py

@@ -25,6 +25,11 @@ class TaskServiceStub(object):
                 request_serializer=workflow_dot_task__service__pb2.GetTaskDetailsRequest.SerializeToString,
                 response_deserializer=workflow_dot_task__service__pb2.GetTaskDetailsResponse.FromString,
                 )
+        self.ListTasks = channel.unary_unary(
+                '/cloudidl.workflow.TaskService/ListTasks',
+                request_serializer=workflow_dot_task__service__pb2.ListTasksRequest.SerializeToString,
+                response_deserializer=workflow_dot_task__service__pb2.ListTasksResponse.FromString,
+                )
 
 
 class TaskServiceServicer(object):
@@ -45,6 +50,13 @@ class TaskServiceServicer(object):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def ListTasks(self, request, context):
+        """Lists tasks, one per task name, returning the latest version and who it was deployed by.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
 
 def add_TaskServiceServicer_to_server(servicer, server):
     rpc_method_handlers = {
@@ -58,6 +70,11 @@ def add_TaskServiceServicer_to_server(servicer, server):
                     request_deserializer=workflow_dot_task__service__pb2.GetTaskDetailsRequest.FromString,
                     response_serializer=workflow_dot_task__service__pb2.GetTaskDetailsResponse.SerializeToString,
             ),
+            'ListTasks': grpc.unary_unary_rpc_method_handler(
+                    servicer.ListTasks,
+                    request_deserializer=workflow_dot_task__service__pb2.ListTasksRequest.FromString,
+                    response_serializer=workflow_dot_task__service__pb2.ListTasksResponse.SerializeToString,
+            ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
             'cloudidl.workflow.TaskService', rpc_method_handlers)
@@ -102,3 +119,20 @@ class TaskService(object):
             workflow_dot_task__service__pb2.GetTaskDetailsResponse.FromString,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def ListTasks(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/cloudidl.workflow.TaskService/ListTasks',
+            workflow_dot_task__service__pb2.ListTasksRequest.SerializeToString,
+            workflow_dot_task__service__pb2.ListTasksResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)

flyte/_run.py

@@ -1,20 +1,18 @@
 from __future__ import annotations
 
+import asyncio
 import pathlib
 import uuid
 from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple, Union, cast
 
-import flyte
-import flyte.report
-from flyte import S3
+import flyte.errors
+from flyte.errors import InitializationError
+from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath, SerializationContext, TaskContext
+from flyte.syncify import syncify
 
-from ._api_commons import syncer
 from ._context import contextual_run, internal_ctx
-from ._datastructures import ActionID, Checkpoints, RawDataPath, SerializationContext, TaskContext
 from ._environment import Environment
 from ._initialize import (
-    ABFS,
-    GCS,
     _get_init_config,
     get_client,
     get_common_config,
@@ -22,24 +20,35 @@ from ._initialize import (
     requires_initialization,
     requires_storage,
 )
-from ._internal import create_controller
-from ._internal.runtime.io import _CHECKPOINT_FILE_NAME
-from ._internal.runtime.taskrunner import run_task
 from ._logging import logger
-from ._protos.common import identifier_pb2
 from ._task import P, R, TaskTemplate
 from ._tools import ipython_check
-from .errors import InitializationError
 
 if TYPE_CHECKING:
     from flyte.remote import Run
+    from flyte.remote._task import LazyEntity
 
     from ._code_bundle import CopyFiles
 
 Mode = Literal["local", "remote", "hybrid"]
 
 
-@syncer.wrap
+async def _get_code_bundle_for_run(name: str) -> CodeBundle | None:
+    """
+    Get the code bundle for the run with the given name.
+    This is used to get the code bundle for the run when running in hybrid mode.
+    """
+    from flyte._internal.runtime.task_serde import extract_code_bundle
+    from flyte.remote import Run
+
+    run = await Run.get.aio(name=name)
+    if run:
+        run_details = await run.details()
+        spec = run_details.action_details.pb2.resolved_task_spec
+        return extract_code_bundle(spec)
+    return None
+
+
 class _Runner:
     def __init__(
         self,
@@ -54,13 +63,14 @@ class _Runner:
         raw_data_path: str | None = None,
         metadata_path: str | None = None,
         run_base_dir: str | None = None,
+        overwrite_cache: bool = False,
     ):
         init_config = _get_init_config()
         client = init_config.client if init_config else None
         if not force_mode and client is not None:
             force_mode = "remote"
         force_mode = force_mode or "local"
-        logger.debug(f"Effective run mode: {force_mode}, client configured: {client is not None}")
+        logger.debug(f"Effective run mode: `{force_mode}`, client configured: `{client is not None}`")
         self._mode = force_mode
         self._name = name
         self._service_account = service_account
@@ -72,50 +82,65 @@ class _Runner:
         self._raw_data_path = raw_data_path
         self._metadata_path = metadata_path or "/tmp"
         self._run_base_dir = run_base_dir or "/tmp/base"
+        self._overwrite_cache = overwrite_cache
 
     @requires_initialization
-    async def _run_remote(self, obj: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> Run:
+    async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
+        import grpc
+
         from flyte.remote import Run
+        from flyte.remote._task import LazyEntity
 
         from ._code_bundle import build_code_bundle, build_pkl_bundle
         from ._deploy import build_images, plan_deploy
         from ._internal.runtime.convert import convert_from_native_to_inputs
         from ._internal.runtime.task_serde import translate_task_to_wire
+        from ._protos.common import identifier_pb2
         from ._protos.workflow import run_definition_pb2, run_service_pb2
 
         cfg = get_common_config()
 
-        if obj.parent_env is None:
-            raise ValueError("Task is not attached to an environment. Please attach the task to an environment.")
+        if isinstance(obj, LazyEntity):
+            task = await obj.fetch.aio()
+            task_spec = task.pb2.spec
+            inputs = await convert_from_native_to_inputs(task.interface, *args, **kwargs)
+            version = task.pb2.task_id.version
+            code_bundle = None
+        else:
+            if obj.parent_env is None:
+                raise ValueError("Task is not attached to an environment. Please attach the task to an environment")
 
-        deploy_plan = plan_deploy(cast(Environment, obj.parent_env()))
-        image_cache = await build_images(deploy_plan)
+            deploy_plan = plan_deploy(cast(Environment, obj.parent_env()))
+            image_cache = await build_images(deploy_plan)
 
-        if self._interactive_mode:
-            code_bundle = await build_pkl_bundle(
-                obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
-            )
-        else:
-            if self._copy_files != "none":
-                code_bundle = await build_code_bundle(
-                    from_dir=cfg.root_dir, dryrun=self._dry_run, copy_bundle_to=self._copy_bundle_to
+            if self._interactive_mode:
+                code_bundle = await build_pkl_bundle(
+                    obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
                 )
             else:
-                code_bundle = None
+                if self._copy_files != "none":
+                    code_bundle = await build_code_bundle(
+                        from_dir=cfg.root_dir,
+                        dryrun=self._dry_run,
+                        copy_bundle_to=self._copy_bundle_to,
+                        copy_style=self._copy_files,
+                    )
+                else:
+                    code_bundle = None
 
-        version = self._version or (
-            code_bundle.computed_version if code_bundle and code_bundle.computed_version else None
-        )
-        if not version:
-            raise ValueError("Version is required when running a task")
-        s_ctx = SerializationContext(
-            code_bundle=code_bundle,
-            version=version,
-            image_cache=image_cache,
-            root_dir=cfg.root_dir,
-        )
-        task_spec = translate_task_to_wire(obj, s_ctx)
-        inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
+            version = self._version or (
+                code_bundle.computed_version if code_bundle and code_bundle.computed_version else None
+            )
+            if not version:
+                raise ValueError("Version is required when running a task")
+            s_ctx = SerializationContext(
+                code_bundle=code_bundle,
+                version=version,
+                image_cache=image_cache,
+                root_dir=cfg.root_dir,
+            )
+            task_spec = translate_task_to_wire(obj, s_ctx)
+            inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
 
         if not self._dry_run:
             if get_client() is None:
@@ -152,15 +177,38 @@ class _Runner:
             if task_spec.task_template.id.version == "":
                 task_spec.task_template.id.version = version
 
-            resp = await get_client().run_service.CreateRun(
-                run_service_pb2.CreateRunRequest(
-                    run_id=run_id,
-                    project_id=project_id,
-                    task_spec=task_spec,
-                    inputs=inputs.proto_inputs,
-                ),
-            )
-            return Run(pb2=resp.run)
+            try:
+                resp = await get_client().run_service.CreateRun(
+                    run_service_pb2.CreateRunRequest(
+                        run_id=run_id,
+                        project_id=project_id,
+                        task_spec=task_spec,
+                        inputs=inputs.proto_inputs,
+                        run_spec=run_definition_pb2.RunSpec(
+                            overwrite_cache=self._overwrite_cache,
+                        ),
+                    ),
+                )
+                return Run(pb2=resp.run)
+            except grpc.aio.AioRpcError as e:
+                if e.code() == grpc.StatusCode.UNAVAILABLE:
+                    raise flyte.errors.RuntimeSystemError(
+                        "SystemUnavailableError",
+                        "Flyte system is currently unavailable. check your configuration, or the service status.",
+                    ) from e
+                elif e.code() == grpc.StatusCode.INVALID_ARGUMENT:
+                    raise flyte.errors.RuntimeUserError("InvalidArgumentError", e.details())
+                elif e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                    # TODO maybe this should be a pass and return existing run?
+                    raise flyte.errors.RuntimeUserError(
+                        "RunAlreadyExistsError",
+                        f"A run with the name '{self._name}' already exists. Please choose a different name.",
+                    )
+                else:
+                    raise flyte.errors.RuntimeSystemError(
+                        "RunCreationError",
+                        f"Failed to create run: {e.details()}",
+                    ) from e
 
         class DryRun(Run):
             def __init__(self, _task_spec, _inputs, _code_bundle):
@@ -187,9 +235,14 @@ class _Runner:
         run in the cluster remotely. This is currently only used for testing,
         over the longer term we will productize this.
         """
+        import flyte.report
         from flyte._code_bundle import build_code_bundle, build_pkl_bundle
-        from flyte._datastructures import RawDataPath
         from flyte._deploy import build_images, plan_deploy
+        from flyte.models import RawDataPath
+        from flyte.storage import ABFS, GCS, S3
+
+        from ._internal import create_controller
+        from ._internal.runtime.taskrunner import run_task
 
         cfg = get_common_config()
 
@@ -199,17 +252,26 @@ class _Runner:
         deploy_plan = plan_deploy(cast(Environment, obj.parent_env()))
         image_cache = await build_images(deploy_plan)
 
-        if self._interactive_mode:
-            code_bundle = await build_pkl_bundle(
-                obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
-            )
-        else:
-            if self._copy_files != "none":
-                code_bundle = await build_code_bundle(
-                    from_dir=cfg.root_dir, dryrun=self._dry_run, copy_bundle_to=self._copy_bundle_to
+        code_bundle = None
+        if self._name is not None:
+            # Check if remote run service has this run name already and if exists, then extract the code bundle from it.
+            code_bundle = await _get_code_bundle_for_run(name=self._name)
+
+        if not code_bundle:
+            if self._interactive_mode:
+                code_bundle = await build_pkl_bundle(
+                    obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
                 )
             else:
-                code_bundle = None
+                if self._copy_files != "none":
+                    code_bundle = await build_code_bundle(
+                        from_dir=cfg.root_dir,
+                        dryrun=self._dry_run,
+                        copy_bundle_to=self._copy_bundle_to,
+                        copy_style=self._copy_files,
+                    )
+                else:
+                    code_bundle = None
 
         version = self._version or (
             code_bundle.computed_version if code_bundle and code_bundle.computed_version else None
@@ -217,14 +279,14 @@ class _Runner:
         if not version:
             raise ValueError("Version is required when running a task")
 
-        project = cfg.project or "testproject"
-        domain = cfg.domain or "development"
-        org = cfg.org or "testorg"
+        project = cfg.project
+        domain = cfg.domain
+        org = cfg.org
         action_name = "a0"
         run_name = self._name
         random_id = str(uuid.uuid4())[:6]
 
-        controller = create_controller(ct="remote", endpoint="localhost:8090", insecure=True)
+        controller = create_controller("remote", endpoint="localhost:8090", insecure=True)
         action = ActionID(name=action_name, run_name=run_name, project=project, domain=domain, org=org)
 
         inputs = obj.native_interface.convert_to_kwargs(*args, **kwargs)
@@ -242,7 +304,7 @@ class _Runner:
         output_path = self._run_base_dir
         raw_data_path = f"{output_path}/rd/{random_id}"
         raw_data_path_obj = RawDataPath(path=raw_data_path)
-        checkpoint_path = f"{raw_data_path}/{_CHECKPOINT_FILE_NAME}"
+        checkpoint_path = f"{raw_data_path}/checkpoint"
         prev_checkpoint = f"{raw_data_path}/prev_checkpoint"
         checkpoints = Checkpoints(checkpoint_path, prev_checkpoint)
 
@@ -253,7 +315,7 @@ class _Runner:
                 checkpoints=checkpoints,
                 code_bundle=code_bundle,
                 output_path=output_path,
-                version=version,
+                version=version if version else "na",
                 raw_data_path=raw_data_path_obj,
                 compiled_image_cache=image_cache,
                 run_base_dir=self._run_base_dir,
@@ -269,43 +331,44 @@ class _Runner:
 
     async def _run_local(self, obj: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
         from flyte._internal.controllers import create_controller
-        from flyte._internal.runtime.convert import (
-            convert_error_to_native,
-            convert_from_native_to_inputs,
-            convert_outputs_to_native,
-        )
-        from flyte._internal.runtime.entrypoints import direct_dispatch
+        from flyte._internal.controllers._local_controller import LocalController
+        from flyte.report import Report
 
-        controller = create_controller(ct="local")
+        controller = cast(LocalController, create_controller("local"))
 
-        inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
         if self._name is None:
             action = ActionID.create_random()
         else:
             action = ActionID(name=self._name)
-        out, err = await direct_dispatch(
-            obj,
+
+        ctx = internal_ctx()
+        tctx = TaskContext(
             action=action,
-            raw_data_path=internal_ctx().raw_data,
-            version="na",
-            controller=controller,
-            inputs=inputs,
-            output_path=self._metadata_path,
-            run_base_dir=self._metadata_path,
             checkpoints=Checkpoints(
                 prev_checkpoint_path=internal_ctx().raw_data.path, checkpoint_path=internal_ctx().raw_data.path
             ),
-        )  # type: ignore
-        if err:
-            native_err = convert_error_to_native(err)
-            if native_err:
-                raise native_err
-        if obj.native_interface.outputs and len(obj.native_interface.outputs) > 0:
-            if out is not None:
-                return cast(R, await convert_outputs_to_native(obj.native_interface, out))
-        return cast(R, None)
-
-    async def run(self, task: TaskTemplate[P, Union[R, Run]], *args: P.args, **kwargs: P.kwargs) -> Union[R, Run]:
+            code_bundle=None,
+            output_path=self._metadata_path,
+            run_base_dir=self._metadata_path,
+            version="na",
+            raw_data_path=internal_ctx().raw_data,
+            compiled_image_cache=None,
+            report=Report(name=action.name),
+            mode="local",
+        )
+        with ctx.replace_task_context(tctx):
+            # make the local version always runs on a different thread, returns a wrapped future.
+            if obj._call_as_synchronous:
+                fut = controller.submit_sync(obj, *args, **kwargs)
+                awaitable = asyncio.wrap_future(fut)
+                return await awaitable
+            else:
+                return await controller.submit(obj, *args, **kwargs)
+
+    @syncify
+    async def run(
+        self, task: TaskTemplate[P, Union[R, Run]] | LazyEntity, *args: P.args, **kwargs: P.kwargs
+    ) -> Union[R, Run]:
         """
         Run an async `@env.task` or `TaskTemplate` instance. The existing async context will be used.
 
@@ -327,13 +390,18 @@ class _Runner:
         :param kwargs: Keyword arguments to pass to the Task
         :return: Run instance or the result of the task
         """
+        from flyte.remote._task import LazyEntity
+
+        if isinstance(task, LazyEntity) and self._mode != "remote":
+            raise ValueError("Remote task can only be run in remote mode.")
         if self._mode == "remote":
             return await self._run_remote(task, *args, **kwargs)
+        task = cast(TaskTemplate, task)
         if self._mode == "hybrid":
             return await self._run_hybrid(task, *args, **kwargs)
 
         # TODO We could use this for remote as well and users could simply pass flyte:// or s3:// or file://
-        async with internal_ctx().new_raw_data_path(
+        with internal_ctx().new_raw_data_path(
             raw_data_path=RawDataPath.from_local_folder(local_folder=self._raw_data_path)
         ):
             return await self._run_local(task, *args, **kwargs)
@@ -351,6 +419,7 @@ def with_runcontext(
     interactive_mode: bool | None = None,
     raw_data_path: str | None = None,
     run_base_dir: str | None = None,
+    overwrite_cache: bool = False,
 ) -> _Runner:
     """
     Launch a new run with the given parameters as the context.
@@ -380,6 +449,8 @@ def with_runcontext(
          interactive mode, while scripts are not. This is used to determine how the code bundle is created.
     :param raw_data_path: Use this path to store the raw data for the run. Currently only supported for local runs,
       and can be used to store raw data in specific locations. TODO coming soon for remote runs as well.
+    :param run_base_dir: Optional The base directory to use for the run. This is used to store the metadata for the run,
+     that is passed between tasks.
     :return: runner
     """
     if mode == "hybrid" and not name and not run_base_dir:
@@ -398,7 +469,7 @@ def with_runcontext(
     )
 
 
-@syncer.wrap
+@syncify
 async def run(task: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> Union[R, Run]:
     """
     Run a task with the given parameters