flyte 0.0.1b0__py3-none-any.whl

flyte/_internal/resolvers/_task_module.py

@@ -0,0 +1,54 @@
+import inspect
+import os
+import pathlib
+import sys
+from typing import Tuple
+
+from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
+
+
+def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path | None = None) -> Tuple[str, str]:
+    """
+    Extract the task module from the task template.
+
+    :param task: The task template to extract the module from.
+    :param source_dir: The source directory to use for relative paths.
+    :return: A tuple containing the entity name, module
+    """
+    entity_name = task.name
+    if isinstance(task, AsyncFunctionTaskTemplate):
+        entity_module = inspect.getmodule(task.func)
+        if entity_module is None:
+            raise ValueError(f"Task {entity_name} has no module.")
+
+        fp = entity_module.__file__
+        if fp is None:
+            raise ValueError(f"Task {entity_name} has no module.")
+
+        file_path = pathlib.Path(fp)
+        # Get the relative path to the current directory
+        # Will raise ValueError if the file is not in the source directory
+        relative_path = file_path.relative_to(str(source_dir))
+
+        if relative_path == pathlib.Path("."):
+            entity_module_name = entity_module.__name__
+        else:
+            # Replace file separators with dots and remove the '.py' extension
+            dotted_path = os.path.splitext(str(relative_path))[0].replace(os.sep, ".")
+            entity_module_name = dotted_path
+
+        entity_name = task.func.__name__
+    else:
+        raise NotImplementedError(f"Task module {entity_name} not implemented.")
+
+    if entity_module_name == "__main__":
+        """
+        This case is for the case in which the task is run from the main module.
+        """
+        fp = sys.modules["__main__"].__file__
+        if fp is None:
+            raise ValueError(f"Task {entity_name} has no module.")
+        main_path = pathlib.Path(fp)
+        entity_module_name = main_path.stem
+
+    return entity_name, entity_module_name

flyte/_internal/resolvers/common.py

@@ -0,0 +1,31 @@
+from asyncio import Protocol
+from pathlib import Path
+from typing import List, Optional
+
+from flyte._task import TaskTemplate
+
+
+class Resolver(Protocol):
+    """
+    Resolver interface for loading tasks. This interface should be implemented by Resolvers.
+    """
+
+    @property
+    def import_path(self) -> str:
+        """
+        The import path of the resolver. This should be a valid python import path.
+        """
+        return ""
+
+    def load_task(self, loader_args: List[str]) -> TaskTemplate:
+        """
+        Given the set of identifier keys, should return one TaskTemplate or raise an error if not found
+        """
+        raise NotImplementedError
+
+    def loader_args(self, t: TaskTemplate, root_dir: Optional[Path]) -> List[str]:
+        """
+        Return a list of strings that can help identify the parameter TaskTemplate. Each string should not have
+        spaces or special characters. This is used to identify the task in the resolver.
+        """
+        return []

flyte/_internal/resolvers/default.py

@@ -0,0 +1,28 @@
+import importlib
+from pathlib import Path
+from typing import List, Optional
+
+from flyte._internal.resolvers._task_module import extract_task_module
+from flyte._internal.resolvers.common import Resolver
+from flyte._task import TaskTemplate
+
+
+class DefaultTaskResolver(Resolver):
+    """
+    Please see the notes in the TaskResolverMixin as it describes this default behavior.
+    """
+
+    @property
+    def import_path(self) -> str:
+        return "flyte._internal.resolvers.default.DefaultTaskResolver"
+
+    def load_task(self, loader_args: List[str]) -> TaskTemplate:
+        _, task_module, _, task_name, *_ = loader_args
+
+        task_module = importlib.import_module(name=task_module)  # type: ignore
+        task_def = getattr(task_module, task_name)
+        return task_def
+
+    def loader_args(self, task: TaskTemplate, root_dir: Optional[Path] = None) -> List[str]:  # type:ignore
+        t, m = extract_task_module(task, root_dir)
+        return ["mod", m, "instance", t]

flyte/_internal/runtime/convert.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Tuple, Union
+
+from flyteidl.core import execution_pb2, literals_pb2
+
+import flyte.errors
+import flyte.storage as storage
+from flyte._datastructures import ActionID, NativeInterface, TaskContext
+from flyte._internal.controllers import pbhash
+from flyte._protos.workflow import run_definition_pb2
+from flyte.types import TypeEngine
+
+
+@dataclass(frozen=True)
+class Inputs:
+    proto_inputs: run_definition_pb2.Inputs
+
+    @classmethod
+    def empty(cls) -> "Inputs":
+        return cls(proto_inputs=run_definition_pb2.Inputs())
+
+
+@dataclass(frozen=True)
+class Outputs:
+    proto_outputs: run_definition_pb2.Outputs
+
+
+@dataclass
+class Error:
+    err: execution_pb2.ExecutionError
+
+
+# ------------------------------- CONVERT Methods ------------------------------- #
+
+
+def _clean_error_code(code: str) -> Tuple[str, str | None]:
+    """
+    The error code may have a server injected code and is of the form `RetriesExhausedError|<code>` or `<code>`.
+
+    :param code:
+    :return: "user code", optional server code
+    """
+    if "|" in code:
+        server_code, user_code = code.split("|", 1)
+        return user_code.strip(), server_code.strip()
+    return code.strip(), None
+
+
+async def convert_inputs_to_native(inputs: Inputs, python_interface: NativeInterface) -> Dict[str, Any]:
+    literals = {named_literal.name: named_literal.value for named_literal in inputs.proto_inputs.literals}
+    native_vals = await TypeEngine.literal_map_to_kwargs(
+        literals_pb2.LiteralMap(literals=literals), python_interface.get_input_types()
+    )
+    return native_vals
+
+
+async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwargs) -> Inputs:
+    kwargs = interface.convert_to_kwargs(*args, **kwargs)
+    if len(kwargs) == 0:
+        return Inputs.empty()
+    if len(kwargs) < len(interface.inputs):
+        raise ValueError(
+            f"Received {len(kwargs)} inputs but interface has {len(interface.inputs)}. "
+            f"Please provide all required inputs."
+        )
+    literal_map = await TypeEngine.dict_to_literal_map(kwargs, interface.get_input_types())
+    return Inputs(
+        proto_inputs=run_definition_pb2.Inputs(
+            literals=[run_definition_pb2.NamedLiteral(name=k, value=v) for k, v in literal_map.literals.items()]
+        )
+    )
+
+
+async def convert_from_native_to_outputs(o: Any, interface: NativeInterface) -> Outputs:
+    # Always make it a tuple even if it's just one item to simplify logic below
+    if not isinstance(o, tuple):
+        o = (o,)
+
+    assert len(interface.outputs) == len(interface.outputs), (
+        f"Received {len(o)} outputs but interface has {len(interface.outputs)}"
+    )
+    named = []
+    for (output_name, python_type), v in zip(interface.outputs.items(), o):
+        lit = await TypeEngine.to_literal(v, python_type, TypeEngine.to_literal_type(python_type))
+        named.append(run_definition_pb2.NamedLiteral(name=output_name, value=lit))
+
+    return Outputs(proto_outputs=run_definition_pb2.Outputs(literals=named))
+
+
+async def convert_outputs_to_native(interface: NativeInterface, outputs: Outputs) -> Union[Any, Tuple[Any, ...]]:
+    lm = literals_pb2.LiteralMap(
+        literals={named_literal.name: named_literal.value for named_literal in outputs.proto_outputs.literals}
+    )
+    kwargs = await TypeEngine.literal_map_to_kwargs(lm, interface.outputs)
+    if len(kwargs) == 0:
+        return None
+    elif len(kwargs) == 1:
+        return next(iter(kwargs.values()))
+    else:
+        # Return as tuple if multiple outputs, make sure to order correctly as it seems proto maps can change ordering
+        return tuple(kwargs[k] for k in interface.outputs.keys())
+
+
+def convert_error_to_native(err: execution_pb2.ExecutionError | Exception | Error) -> Exception | None:
+    if not err:
+        return None
+
+    if isinstance(err, Exception):
+        return err
+
+    if isinstance(err, Error):
+        err = err.err
+
+    user_code, server_code = _clean_error_code(err.code)
+    match err.kind:
+        case execution_pb2.ExecutionError.UNKNOWN:
+            return flyte.errors.RuntimeUnknownError(code=user_code, message=err.message, worker=err.worker)
+        case execution_pb2.ExecutionError.USER:
+            if "OOM" in err.code.upper():
+                return flyte.errors.OOMError(code=user_code, message=err.message, worker=err.worker)
+            elif "Interrupted" in err.code:
+                return flyte.errors.TaskInterruptedError(code=user_code, message=err.message, worker=err.worker)
+            elif "PrimaryContainerNotFound" in err.code:
+                return flyte.errors.PrimaryContainerNotFoundError(
+                    code=user_code, message=err.message, worker=err.worker
+                )
+            elif "RetriesExhausted" in err.code:
+                return flyte.errors.RetriesExhaustedError(code=user_code, message=err.message, worker=err.worker)
+            elif "Unknown" in err.code:
+                return flyte.errors.RuntimeUnknownError(code=user_code, message=err.message, worker=err.worker)
+            elif "InvalidImageName" in err.code:
+                return flyte.errors.InvalidImageNameError(code=user_code, message=err.message, worker=err.worker)
+            elif "ImagePullBackOff" in err.code:
+                return flyte.errors.ImagePullBackOffError(code=user_code, message=err.message, worker=err.worker)
+            return flyte.errors.RuntimeUserError(code=user_code, message=err.message, worker=err.worker)
+        case execution_pb2.ExecutionError.SYSTEM:
+            return flyte.errors.RuntimeSystemError(code=user_code, message=err.message, worker=err.worker)
+    return None
+
+
+def convert_from_native_to_error(err: BaseException) -> Error:
+    if isinstance(err, flyte.errors.RuntimeUnknownError):
+        return Error(
+            err=execution_pb2.ExecutionError(
+                kind=execution_pb2.ExecutionError.UNKNOWN,
+                code=err.code,
+                message=str(err),
+                worker=err.worker,
+            )
+        )
+    elif isinstance(err, flyte.errors.RuntimeUserError):
+        return Error(
+            err=execution_pb2.ExecutionError(
+                kind=execution_pb2.ExecutionError.USER,
+                code=err.code,
+                message=str(err),
+                worker=err.worker,
+            )
+        )
+    elif isinstance(err, flyte.errors.RuntimeSystemError):
+        return Error(
+            err=execution_pb2.ExecutionError(
+                kind=execution_pb2.ExecutionError.SYSTEM,
+                code=err.code,
+                message=str(err),
+                worker=err.worker,
+            )
+        )
+    else:
+        return Error(
+            err=execution_pb2.ExecutionError(
+                kind=execution_pb2.ExecutionError.UNKNOWN,
+                code=type(err).__name__,
+                message=str(err),
+                worker="UNKNOWN",
+            )
+        )
+
+
+def generate_sub_action_id_and_output_path(tctx: TaskContext, task_name: str, inputs: Inputs) -> Tuple[ActionID, str]:
+    """
+    Generate a sub-action ID and output path based on the current task context, task name, and inputs.
+    :param tctx:
+    :param task_name:
+    :param inputs:
+    :return:
+    """
+    current_action_id = tctx.action
+    current_output_path = tctx.run_base_dir
+    inputs_hash = pbhash.compute_hash_string(inputs.proto_inputs)
+    sub_action_id = current_action_id.new_sub_action_from(
+        task_name=task_name,
+        input_hash=inputs_hash,
+        group=tctx.group_data.name if tctx.group_data else None,
+    )
+    sub_run_output_path = storage.join(current_output_path, sub_action_id.name)
+    return sub_action_id, sub_run_output_path

flyte/_internal/runtime/entrypoints.py

@@ -0,0 +1,135 @@
+from typing import List, Optional, Tuple
+
+import flyte.errors
+from flyte._code_bundle import download_bundle
+from flyte._context import contextual_run
+from flyte._datastructures import ActionID, Checkpoints, CodeBundle, RawDataPath
+from flyte._internal import Controller
+from flyte._internal.imagebuild.image_builder import ImageCache
+from flyte._logging import log, logger
+from flyte._task import TaskTemplate
+
+from .convert import Error, Inputs, Outputs
+from .task_serde import load_task
+from .taskrunner import (
+    convert_and_run,
+    extract_download_run_upload,
+)
+
+
+async def direct_dispatch(
+    task: TaskTemplate,
+    *,
+    action: ActionID,
+    raw_data_path: RawDataPath,
+    controller: Controller,
+    version: str,
+    output_path: str,
+    run_base_dir: str,
+    checkpoints: Checkpoints | None = None,
+    code_bundle: CodeBundle | None = None,
+    inputs: Inputs | None = None,
+) -> Tuple[Optional[Outputs], Optional[Error]]:
+    """
+    This method is used today by the local_controller and is positioned to be used by a rust core in the future.
+    The caller, loads the task and invokes this method. This method is used to convert the inputs to native types,
+    The reason for this is that the rust entrypoint will not have access to the python context, and
+    will not be able to run the tasks in the context tree.
+    """
+    return await contextual_run(
+        convert_and_run,
+        task=task,
+        inputs=inputs or Inputs.empty(),
+        action=action,
+        raw_data_path=raw_data_path,
+        checkpoints=checkpoints,
+        code_bundle=code_bundle,
+        controller=controller,
+        version=version,
+        output_path=output_path,
+        run_base_dir=run_base_dir,
+    )
+
+
+async def _download_and_load_task(
+    code_bundle: CodeBundle | None, resolver: str | None = None, resolver_args: List[str] | None = None
+) -> TaskTemplate:
+    if code_bundle and (code_bundle.tgz or code_bundle.pkl):
+        logger.debug(f"Downloading {code_bundle}")
+        downloaded_path = await download_bundle(code_bundle)
+        code_bundle = code_bundle.with_downloaded_path(downloaded_path)
+        if code_bundle.pkl:
+            try:
+                logger.debug(f"Loading task from pkl: {code_bundle.downloaded_path}")
+                import gzip
+
+                import cloudpickle
+
+                with gzip.open(str(code_bundle.downloaded_path), "rb") as f:
+                    return cloudpickle.load(f)
+            except Exception as e:
+                logger.exception(f"Failed to load pickled task from {code_bundle.downloaded_path}. Reason: {e!s}")
+                raise
+
+        if not resolver or not resolver_args:
+            raise flyte.errors.RuntimeSystemError(
+                "MalformedCommand", "Resolver and resolver args are required. for task"
+            )
+        logger.debug(
+            f"Loading task from tgz: {code_bundle.downloaded_path}, resolver: {resolver}, args: {resolver_args}"
+        )
+        return load_task(resolver, *resolver_args)
+    if not resolver or not resolver_args:
+        raise flyte.errors.RuntimeSystemError("MalformedCommand", "Resolver and resolver args are required. for task")
+    logger.debug(f"No code bundle provided, loading task from resolver: {resolver}, args: {resolver_args}")
+    return load_task(resolver, *resolver_args)
+
+
+@log
+async def load_and_run_task(
+    action: ActionID,
+    raw_data_path: RawDataPath,
+    output_path: str,
+    run_base_dir: str,
+    version: str,
+    controller: Controller,
+    resolver: str,
+    resolver_args: List[str],
+    checkpoints: Checkpoints | None = None,
+    code_bundle: CodeBundle | None = None,
+    input_path: str | None = None,
+    image_cache: ImageCache | None = None,
+):
+    """
+    This method is invoked from the runtime/CLI and is used to run a task. This creates the context tree,
+    for the tasks to run in. It also handles the loading of the task.
+
+    :param controller: Controller to use for the task.
+    :param resolver: The resolver to use to load the task.
+    :param resolver_args: The arguments to pass to the resolver.
+    :param action: The ActionID to use for the task.
+    :param raw_data_path: The raw data path to use for the task.
+    :param output_path: The output path to use for the task.
+    :param run_base_dir: Base output directory to pass down to child tasks.
+    :param version: The version of the task to run.
+    :param checkpoints: The checkpoints to use for the task.
+    :param code_bundle: The code bundle to use for the task.
+    :param input_path: The input path to use for the task.
+    :param image_cache: Mappings of Image identifiers to image URIs.
+    """
+    task = await _download_and_load_task(code_bundle, resolver, resolver_args)
+
+    await contextual_run(
+        extract_download_run_upload,
+        task,
+        action=action,
+        version=version,
+        controller=controller,
+        raw_data_path=raw_data_path,
+        output_path=output_path,
+        run_base_dir=run_base_dir,
+        checkpoints=checkpoints,
+        code_bundle=code_bundle,
+        input_path=input_path,
+        image_cache=image_cache,
+    )

flyte/_internal/runtime/io.py

@@ -0,0 +1,136 @@
+"""
+This module contains the methods for uploading and downloading inputs and outputs.
+It uses the storage module to handle the actual uploading and downloading of files.
+
+TODO: Convert to use streaming apis
+"""
+
+import logging
+
+from flyteidl.core import errors_pb2, execution_pb2
+
+import flyte.storage as storage
+from flyte._protos.workflow import run_definition_pb2
+
+from ..._logging import log
+from .convert import Inputs, Outputs, _clean_error_code
+
+# ------------------------------- CONSTANTS ------------------------------- #
+_INPUTS_FILE_NAME = "inputs.pb"
+_OUTPUTS_FILE_NAME = "outputs.pb"
+_CHECKPOINT_FILE_NAME = "_flytecheckpoints"
+_ERROR_FILE_NAME = "error.pb"
+_REPORT_FILE_NAME = "report.html"
+_PKL_FILE_NAME = "code_bundle.pkl.gz"
+
+
+def pkl_path(base_path: str) -> str:
+    return storage.join(base_path, _PKL_FILE_NAME)
+
+
+def inputs_path(base_path: str) -> str:
+    return storage.join(base_path, _INPUTS_FILE_NAME)
+
+
+def outputs_path(base_path: str) -> str:
+    return storage.join(base_path, _OUTPUTS_FILE_NAME)
+
+
+def error_path(base_path: str) -> str:
+    return storage.join(base_path, _ERROR_FILE_NAME)
+
+
+def report_path(base_path: str) -> str:
+    return storage.join(base_path, _REPORT_FILE_NAME)
+
+
+# ------------------------------- UPLOAD Methods ------------------------------- #
+
+
+async def upload_inputs(inputs: Inputs, input_path: str):
+    """
+    :param Inputs inputs: Inputs
+    :param str input_path: The path to upload the input file.
+    """
+    await storage.put_stream(data_iterable=inputs.proto_inputs.SerializeToString(), to_path=input_path)
+
+
+async def upload_outputs(outputs: Outputs, output_path: str):
+    """
+    :param outputs: Outputs
+    :param output_path: The path to upload the output file.
+    """
+    output_uri = outputs_path(output_path)
+    await storage.put_stream(data_iterable=outputs.proto_outputs.SerializeToString(), to_path=output_uri)
+
+
+async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str):
+    """
+    :param err: execution_pb2.ExecutionError
+    :param output_prefix: The output prefix of the remote uri.
+    """
+    # TODO - clean this up + conditionally set kind
+    error_document = errors_pb2.ErrorDocument(
+        error=errors_pb2.ContainerError(
+            code=err.code,
+            message=err.message,
+            kind=errors_pb2.ContainerError.RECOVERABLE,
+            origin=err.kind,
+            timestamp=err.timestamp,
+            worker=err.worker,
+        )
+    )
+    error_uri = error_path(output_prefix)
+    await storage.put_stream(data_iterable=error_document.SerializeToString(), to_path=error_uri)
+
+
+# ------------------------------- DOWNLOAD Methods ------------------------------- #
+@log(level=logging.INFO)
+async def load_inputs(path: str) -> Inputs:
+    """
+    :param path: Input file to be downloaded
+    :return: Inputs object
+    """
+    lm = run_definition_pb2.Inputs()
+    proto_str = b"".join([c async for c in storage.get_stream(path=path)])
+    lm.ParseFromString(proto_str)
+    return Inputs(proto_inputs=lm)
+
+
+async def load_outputs(path: str) -> Outputs:
+    """
+    :param path: output file to be loaded
+    :return: Outputs object
+    """
+    lm = run_definition_pb2.Outputs()
+    proto_str = b"".join([c async for c in storage.get_stream(path=path)])
+    lm.ParseFromString(proto_str)
+    return Outputs(proto_outputs=lm)
+
+
+async def load_error(path: str) -> execution_pb2.ExecutionError:
+    """
+    :param path: error file to be downloaded
+    :return: execution_pb2.ExecutionError
+    """
+    err = errors_pb2.ErrorDocument()
+    proto_str = b"".join([c async for c in storage.get_stream(path=path)])
+    err.ParseFromString(proto_str)
+
+    if err.error is not None:
+        user_code, server_code = _clean_error_code(err.error.code)
+        return execution_pb2.ExecutionError(
+            code=user_code,
+            message=err.error.message,
+            kind=err.error.origin,
+            error_uri=path,
+            timestamp=err.error.timestamp,
+            worker=err.error.worker,
+        )
+
+    return execution_pb2.ExecutionError(
+        code="Unknown",
+        message=f"Received unloadable error from path {path}",
+        kind=execution_pb2.ExecutionError.SYSTEM,
+        error_uri=path,
+    )