flyte 0.1.0__py3-none-any.whl → 0.2.0a0__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,342 @@
+from __future__ import annotations
+
+import asyncio
+import base64
+import hashlib
+import inspect
+from dataclasses import dataclass
+from types import NoneType
+from typing import Any, Dict, List, Tuple, Union, get_args
+
+from flyteidl.core import execution_pb2, interface_pb2, literals_pb2
+
+import flyte.errors
+import flyte.storage as storage
+from flyte._protos.workflow import common_pb2, run_definition_pb2, task_definition_pb2
+from flyte.models import ActionID, NativeInterface, TaskContext
+from flyte.types import TypeEngine, TypeTransformerFailedError
+
+
+@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_upload_default_inputs(interface: NativeInterface) -> List[common_pb2.NamedParameter]:
+    """
+    Converts the default inputs of a NativeInterface to a list of NamedParameters for upload.
+    This is used to upload default inputs to the Flyte backend.
+    """
+    if not interface.inputs:
+        return []
+
+    vars = []
+    literal_coros = []
+    for input_name, (input_type, default_value) in interface.inputs.items():
+        if default_value is not inspect.Parameter.empty:
+            lt = TypeEngine.to_literal_type(input_type)
+            literal_coros.append(TypeEngine.to_literal(default_value, input_type, lt))
+            vars.append((input_name, lt))
+
+    literals: List[literals_pb2.Literal] = await asyncio.gather(*literal_coros)
+    named_params = []
+    for (name, lt), literal in zip(vars, literals):
+        param = interface_pb2.Parameter(
+            var=interface_pb2.Variable(
+                type=lt,
+            ),
+            default=literal,
+        )
+        named_params.append(
+            common_pb2.NamedParameter(
+                name=name,
+                parameter=param,
+            ),
+        )
+    return named_params
+
+
+def is_optional_type(tp) -> bool:
+    """
+    True if the *annotation* `tp` is equivalent to Optional[…].
+    Works for Optional[T], Union[T, None], and T | None.
+    """
+    return NoneType in get_args(tp)  # fastest check
+
+
+async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwargs) -> Inputs:
+    kwargs = interface.convert_to_kwargs(*args, **kwargs)
+
+    if len(kwargs) < interface.num_required_inputs():
+        raise ValueError(
+            f"Received {len(kwargs)} inputs but interface has {interface.num_required_inputs()} required inputs. "
+            f"Please provide all required inputs. Inputs received: {kwargs}, interface: {interface}"
+        )
+
+    if len(interface.inputs) == 0:
+        return Inputs.empty()
+
+    # fill in defaults if missing
+    type_hints: Dict[str, type] = {}
+    already_converted_kwargs: Dict[str, literals_pb2.Literal] = {}
+    for input_name, (input_type, default_value) in interface.inputs.items():
+        if input_name in kwargs:
+            type_hints[input_name] = input_type
+        elif (default_value is not None and default_value is not inspect.Signature.empty) or (
+            default_value is None and is_optional_type(input_type)
+        ):
+            if default_value == NativeInterface.has_default:
+                if interface._remote_defaults is None or input_name not in interface._remote_defaults:
+                    raise ValueError(f"Input '{input_name}' has a default value but it is not set in the interface.")
+                already_converted_kwargs[input_name] = interface._remote_defaults[input_name]
+            else:
+                kwargs[input_name] = default_value
+                type_hints[input_name] = input_type
+
+    literal_map = await TypeEngine.dict_to_literal_map(kwargs, type_hints)
+    if len(already_converted_kwargs) > 0:
+        copied_literals: Dict[str, literals_pb2.Literal] = {}
+        for k, v in literal_map.literals.items():
+            copied_literals[k] = v
+        # Add the already converted kwargs to the literal map
+        for k, v in already_converted_kwargs.items():
+            copied_literals[k] = v
+        literal_map = literals_pb2.LiteralMap(literals=copied_literals)
+    # Make sure we the interface, not literal_map or kwargs, because those may have a different order
+    return Inputs(
+        proto_inputs=run_definition_pb2.Inputs(
+            literals=[
+                run_definition_pb2.NamedLiteral(name=k, value=literal_map.literals[k]) for k in interface.inputs.keys()
+            ]
+        )
+    )
+
+
+async def convert_from_native_to_outputs(o: Any, interface: NativeInterface, task_name: str = "") -> 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):
+        try:
+            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))
+        except TypeTransformerFailedError as e:
+            raise flyte.errors.RuntimeDataValidationError(output_name, e, task_name)
+
+    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 hash_data(data: Union[str, bytes]) -> str:
+    """
+    Generate a hash for the given data. If the data is a string, it will be encoded to bytes before hashing.
+    :param data: The data to hash, can be a string or bytes.
+    :return: A hexadecimal string representation of the hash.
+    """
+    if isinstance(data, str):
+        data = data.encode("utf-8")
+    digest = hashlib.sha256(data).digest()
+    return base64.b64encode(digest).decode("utf-8")
+
+
+def generate_inputs_hash(serialized_inputs: str | bytes) -> str:
+    """
+    Generate a hash for the inputs. This is used to uniquely identify the inputs for a task.
+    :return: A hexadecimal string representation of the hash.
+    """
+    return hash_data(serialized_inputs)
+
+
+def generate_cache_key_hash(
+    task_name: str,
+    inputs_hash: str,
+    task_interface: interface_pb2.TypedInterface,
+    cache_version: str,
+    ignored_input_vars: List[str],
+    proto_inputs: run_definition_pb2.Inputs,
+) -> str:
+    """
+    Generate a cache key hash based on the inputs hash, task name, task interface, and cache version.
+    This is used to uniquely identify the cache key for a task.
+
+    :param task_name: The name of the task.
+    :param inputs_hash: The hash of the inputs.
+    :param task_interface: The interface of the task.
+    :param cache_version: The version of the cache.
+    :param ignored_input_vars: A list of input variable names to ignore when generating the cache key.
+    :param proto_inputs: The proto inputs for the task, only used if there are ignored inputs.
+    :return: A hexadecimal string representation of the cache key hash.
+    """
+    if ignored_input_vars:
+        filtered = [named_lit for named_lit in proto_inputs.literals if named_lit.name not in ignored_input_vars]
+        final = run_definition_pb2.Inputs(literals=filtered)
+        final_inputs = final.SerializeToString(deterministic=True)
+    else:
+        final_inputs = inputs_hash
+    data = f"{final_inputs}{task_name}{task_interface.SerializeToString(deterministic=True)}{cache_version}"
+    return hash_data(data)
+
+
+def generate_sub_action_id_and_output_path(
+    tctx: TaskContext,
+    task_spec_or_name: task_definition_pb2.TaskSpec | str,
+    inputs_hash: str,
+    invoke_seq: int,
+) -> Tuple[ActionID, str]:
+    """
+    Generate a sub-action ID and output path based on the current task context, task name, and inputs.
+
+    action name = current action name + task name + input hash + group name (if available)
+    :param tctx:
+    :param task_spec_or_name: task specification or task name. Task name is only used in case of trace actions.
+    :param inputs_hash: Consistent hash string of the inputs
+    :param invoke_seq: The sequence number of the invocation, used to differentiate between multiple invocations.
+    :return:
+    """
+    current_action_id = tctx.action
+    current_output_path = tctx.run_base_dir
+    if isinstance(task_spec_or_name, task_definition_pb2.TaskSpec):
+        task_spec_or_name.task_template.interface
+        task_hash = hash_data(task_spec_or_name.SerializeToString(deterministic=True))
+    else:
+        task_hash = task_spec_or_name
+    sub_action_id = current_action_id.new_sub_action_from(
+        task_hash=task_hash,
+        input_hash=inputs_hash,
+        group=tctx.group_data.name if tctx.group_data else None,
+        task_call_seq=invoke_seq,
+    )
+    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._internal import Controller
+from flyte._internal.imagebuild.image_builder import ImageCache
+from flyte._logging import log, logger
+from flyte._task import TaskTemplate
+from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
+
+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,
+    )