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

flyte/errors.py

@@ -10,6 +10,16 @@ from typing import Literal
 ErrorKind = Literal["system", "unknown", "user"]
 
 
+def silence_grpc_polling_error(loop, context):
+    """
+    Suppress specific gRPC polling errors in the event loop.
+    """
+    exc = context.get("exception")
+    if isinstance(exc, BlockingIOError):
+        return  # suppress
+    loop.default_exception_handler(context)
+
+
 class BaseRuntimeError(RuntimeError):
     """
     Base class for all Union runtime errors. These errors are raised when the underlying task execution fails, either
@@ -86,6 +96,9 @@ class TaskTimeoutError(RuntimeUserError):
     This error is raised when the underlying task execution runs for longer than the specified timeout.
     """
 
+    def __init__(self, message: str):
+        super().__init__("TaskTimeoutError", message, "user")
+
 
 class RetriesExhaustedError(RuntimeUserError):
     """
@@ -119,7 +132,9 @@ class CustomError(RuntimeUserError):
         Create a CustomError from an exception. The exception's class name is used as the error code and the exception
         message is used as the error message.
         """
-        return cls(e.__class__.__name__, str(e))
+        new_exc = cls(e.__class__.__name__, str(e))
+        new_exc.__cause__ = e
+        return new_exc
 
 
 class NotInTaskContextError(RuntimeUserError):
@@ -134,10 +149,127 @@ class ActionNotFoundError(RuntimeError):
     """
 
 
+# NOTE: Use RemoteTaskError instead, since "reference tasks" (from v1) are now
+# simply "remote tasks" in v2.
 class ReferenceTaskError(RuntimeUserError):
     """
     This error is raised when the user tries to access a task that does not exist.
     """
 
+    CODE = "ReferenceTaskUsageError"
+
+    def __init__(self, message: str):
+        super().__init__(self.CODE, message, "user")
+
+
+class RemoteTaskError(ReferenceTaskError):
+    """
+    This error is raised when the user tries to access a task that does not exist.
+    """
+
+    CODE = "RemoteTaskUsageError"
+
+
+class LogsNotYetAvailableError(BaseRuntimeError):
+    """
+    This error is raised when the logs are not yet available for a task.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("LogsNotYetAvailable", "system", message, None)
+
+
+class RuntimeDataValidationError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access a resource that does not exist or is invalid.
+    """
+
+    def __init__(self, var: str, e: Exception | str, task_name: str = ""):
+        super().__init__(
+            "DataValidationError", f"In task {task_name} variable {var}, failed to serialize/deserialize because of {e}"
+        )
+
+
+class DeploymentError(RuntimeUserError):
+    """
+    This error is raised when the deployment of a task fails, or some preconditions for deployment are not met.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("DeploymentError", message, "user")
+
+
+class ImageBuildError(RuntimeUserError):
+    """
+    This error is raised when the image build fails.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("ImageBuildError", message, "user")
+
+
+class ModuleLoadError(RuntimeUserError):
+    """
+    This error is raised when the module cannot be loaded, either because it does not exist or because of a
+     syntax error.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("ModuleLoadError", message, "user")
+
+
+class InlineIOMaxBytesBreached(RuntimeUserError):
+    """
+    This error is raised when the inline IO max bytes limit is breached.
+    This can be adjusted per task by setting max_inline_io_bytes in the task definition.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("InlineIOMaxBytesBreached", message, "user")
+
+
+class RunAbortedError(RuntimeUserError):
+    """
+    This error is raised when the run is aborted by the user.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("RunAbortedError", message, "user")
+
+
+class SlowDownError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access a resource that does not exist or is invalid.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("SlowDownError", message, "user")
+
+
+class OnlyAsyncIOSupportedError(RuntimeUserError):
+    """
+    This error is raised when the user tries to use sync IO in an async task.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("OnlyAsyncIOSupportedError", message, "user")
+
+
+class ParameterMaterializationError(RuntimeUserError):
+    """
+    This error is raised when the user tries to use a Parameter in an App, that has delayed Materialization,
+    but the materialization fails.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("ParameterMaterializationError", message, "user")
+
+
+class RestrictedTypeError(RuntimeUserError):
+    """
+    This error is raised when the user uses a restricted type, for example current a Tuple is not supported for one
+     value.
+    """
+
     def __init__(self, message: str):
-        super().__init__("ReferenceTaskUsageError", message, "user")
+        super().__init__("RestrictedTypeUsage", message, "user")

flyte/extend.py

@@ -0,0 +1,24 @@
+from flyte._image import Architecture
+
+from ._initialize import is_initialized
+from ._internal.imagebuild.image_builder import ImageBuildEngine, ImageBuilder, ImageChecker
+from ._internal.runtime.entrypoints import download_code_bundle
+from ._internal.runtime.resources_serde import get_proto_resources
+from ._resources import PRIMARY_CONTAINER_DEFAULT_NAME, pod_spec_from_resources
+from ._task import AsyncFunctionTaskTemplate, TaskTemplate
+from ._task_plugins import TaskPluginRegistry
+
+__all__ = [
+    "PRIMARY_CONTAINER_DEFAULT_NAME",
+    "Architecture",
+    "AsyncFunctionTaskTemplate",
+    "ImageBuildEngine",
+    "ImageBuilder",
+    "ImageChecker",
+    "TaskPluginRegistry",
+    "TaskTemplate",
+    "download_code_bundle",
+    "get_proto_resources",
+    "is_initialized",
+    "pod_spec_from_resources",
+]

flyte/extras/_container.py

@@ -2,27 +2,23 @@ import os
 import pathlib
 from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
 
-from flyteidl.core import tasks_pb2
+from flyteidl2.core import tasks_pb2
 
 from flyte import Image, storage
-from flyte._datastructures import NativeInterface, SerializationContext
 from flyte._logging import logger
 from flyte._task import TaskTemplate
+from flyte.io import Dir, File
+from flyte.models import NativeInterface, SerializationContext
 
-_PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
 
-
-def _extract_command_key(cmd: str, **kwargs) -> Any:
+def _extract_command_key(cmd: str, **kwargs) -> List[Any] | None:
     """
     Extract the key from the command using regex.
     """
     import re
 
-    input_regex = r"^\{\{\s*\.inputs\.(.*?)\s*\}\}$"
-    match = re.match(input_regex, cmd)
-    if match:
-        return match.group(1)
-    return None
+    input_regex = r"\{\{\.inputs\.([a-zA-Z0-9_]+)\}\}"
+    return re.findall(input_regex, cmd)
 
 
 def _extract_path_command_key(cmd: str, input_data_dir: Optional[str]) -> Optional[str]:
@@ -32,8 +28,9 @@ def _extract_path_command_key(cmd: str, input_data_dir: Optional[str]) -> Option
     import re
 
     input_data_dir = input_data_dir or ""
-    input_regex = rf"{re.escape(input_data_dir)}/(.+)$"
-    match = re.match(input_regex, cmd)
+    input_regex = rf"{re.escape(input_data_dir)}/([\w\-.]+)"  # captures file or dir names
+
+    match = re.search(input_regex, cmd)
     if match:
         return match.group(1)
     return None
@@ -70,7 +67,7 @@ class ContainerTask(TaskTemplate):
         input_data_dir: str | pathlib.Path = "/var/inputs",
         output_data_dir: str | pathlib.Path = "/var/outputs",
         metadata_format: MetadataFormat = "JSON",
-        local_logs: bool = False,
+        local_logs: bool = True,
         **kwargs,
     ):
         super().__init__(
@@ -83,9 +80,14 @@ class ContainerTask(TaskTemplate):
         self._image = image
         if isinstance(image, str):
             if image == "auto":
-                self._image = Image.auto()
+                self._image = Image.from_debian_base()
             else:
-                self._image = Image.from_prebuilt(image)
+                self._image = Image.from_base(image)
+
+        if command and any(not isinstance(c, str) for c in command):
+            raise ValueError("All elements in the command list must be strings.")
+        if arguments and any(not isinstance(a, str) for a in arguments):
+            raise ValueError("All elements in the arguments list must be strings.")
         self._cmd = command
         self._args = arguments
         self._input_data_dir = input_data_dir
@@ -106,32 +108,34 @@ class ContainerTask(TaskTemplate):
         For FlyteFile and FlyteDirectory commands, e.g., "/var/inputs/inputs", we extract the key from strings that
          begin with the specified `input_data_dir`.
         """
-        # from flytekit.types.directory import FlyteDirectory
-        # from flytekit.types.file import FlyteFile
+        from flyte.io import Dir, File
 
         volume_binding: Dict[str, Dict[str, str]] = {}
         path_k = _extract_path_command_key(cmd, str(self._input_data_dir))
-        k = path_k if path_k else _extract_command_key(cmd)
-
-        if k:
-            input_val = kwargs.get(k)
-            # TODO: Add support file and directory transformer first
-            # if type(input_val) in [FlyteFile, FlyteDirectory]:
-            #     if not path_k:
-            #         raise AssertionError(
-            #             "FlyteFile and FlyteDirectory commands should not use the template syntax like this:
-            #             {{.inputs.infile}}\n"
-            #             "Please use a path-like syntax, such as: /var/inputs/infile.\n"
-            #             "This requirement is due to how Flyte Propeller processes template syntax inputs."
-            #         )
-            #     local_flyte_file_or_dir_path = str(input_val)
-            #     remote_flyte_file_or_dir_path = os.path.join(self._input_data_dir, k)  # type: ignore
-            #     volume_binding[local_flyte_file_or_dir_path] = {
-            #         "bind": remote_flyte_file_or_dir_path,
-            #         "mode": "rw",
-            #     }
-            #     command = remote_flyte_file_or_dir_path
-            command = str(input_val)
+        keys = [path_k] if path_k else _extract_command_key(cmd)
+
+        command = cmd
+
+        if keys:
+            for k in keys:
+                input_val = kwargs.get(k)
+                # TODO: Add support file and directory transformer first
+                if input_val and type(input_val) in [File, Dir]:
+                    if not path_k:
+                        raise AssertionError(
+                            "File and Directory commands should not use the template syntax "
+                            "like this: {{.inputs.infile}}\n"
+                            "Please use a path-like syntax, such as: /var/inputs/infile.\n"
+                            "This requirement is due to how Flyte Propeller processes template syntax inputs."
+                        )
+                    local_flyte_file_or_dir_path = input_val.path
+                    remote_flyte_file_or_dir_path = os.path.join(self._input_data_dir, k)  # type: ignore
+                    volume_binding[local_flyte_file_or_dir_path] = {
+                        "bind": remote_flyte_file_or_dir_path,
+                        "mode": "rw",
+                    }
+                else:
+                    command = command.replace(f"{{{{.inputs.{k}}}}}", str(input_val))
         else:
             command = cmd
 
@@ -193,7 +197,9 @@ class ContainerTask(TaskTemplate):
             microseconds=microseconds,
         )
 
-    def _convert_output_val_to_correct_type(self, output_val: Any, output_type: Type) -> Any:
+    async def _convert_output_val_to_correct_type(
+        self, output_path: pathlib.Path, output_val: Any, output_type: Type
+    ) -> Any:
         import datetime
 
         if issubclass(output_type, bool):
@@ -202,20 +208,31 @@ class ContainerTask(TaskTemplate):
             return datetime.datetime.fromisoformat(output_val)
         elif issubclass(output_type, datetime.timedelta):
             return self._string_to_timedelta(output_val)
+        elif issubclass(output_type, File):
+            return await File.from_local(output_path)
+        elif issubclass(output_type, Dir):
+            return await Dir.from_local(output_path)
         else:
             return output_type(output_val)
 
-    def _get_output_dict(self, output_directory: pathlib.Path) -> Dict[str, Any]:
-        output_dict = {}
+    async def _get_output(self, output_directory: pathlib.Path) -> Tuple[Any]:
+        output_items = []
         if self._outputs:
             for k, output_type in self._outputs.items():
                 output_path = output_directory / k
-                with output_path.open("r") as f:
-                    output_val = f.read()
-                output_dict[k] = self._convert_output_val_to_correct_type(output_val, output_type)
-        return output_dict
-
-    def execute(self, **kwargs) -> Any:
+                if os.path.isfile(output_path):
+                    with output_path.open("r") as f:
+                        output_val = f.read()
+                else:
+                    output_val = None
+                parsed = await self._convert_output_val_to_correct_type(output_path, output_val, output_type)
+                output_items.append(parsed)
+        # return a tuple so that each element is treated as a separate output.
+        # this allows flyte to map the user-defined output types (dict) to individual values.
+        # if we returned a list instead, it would be treated as a single output.
+        return tuple(output_items)
+
+    async def execute(self, **kwargs) -> Any:
         try:
             import docker
         except ImportError:
@@ -235,6 +252,7 @@ class ContainerTask(TaskTemplate):
             raise AssertionError(f"Only Image objects are supported, not strings. Got {self._image} instead.")
         uri = self._image.uri
         self._pull_image_if_not_exists(client, uri)
+        print(f"Command: {commands!r}")
 
         container = client.containers.run(uri, command=commands, remove=True, volumes=volume_bindings, detach=True)
 
@@ -247,8 +265,8 @@ class ContainerTask(TaskTemplate):
 
         container.wait()
 
-        output_dict = self._get_output_dict(output_directory)
-        return output_dict
+        output = await self._get_output(output_directory)
+        return output
 
     def data_loading_config(self, sctx: SerializationContext) -> tasks_pb2.DataLoadingConfig:
         literal_to_protobuf = {
@@ -258,16 +276,11 @@ class ContainerTask(TaskTemplate):
         }
 
         return tasks_pb2.DataLoadingConfig(
-            input_path=self._input_data_dir,
-            output_path=self._output_data_dir,
+            input_path=str(self._input_data_dir) if self._input_data_dir else None,
+            output_path=str(self._output_data_dir) if self._output_data_dir else None,
             enabled=True,
             format=literal_to_protobuf.get(self._metadata_format, "JSON"),
         )
 
     def container_args(self, sctx: SerializationContext) -> List[str]:
         return self._cmd + (self._args if self._args else [])
-
-    def config(self, sctx: SerializationContext) -> Dict[str, str]:
-        if self.pod_template is None:
-            return {}
-        return {_PRIMARY_CONTAINER_NAME_FIELD: self.primary_container_name}

flyte/git/__init__.py

@@ -0,0 +1,3 @@
+from ._config import GitStatus, config_from_root
+
+__all__ = ["GitStatus", "config_from_root"]