flyte 0.0.1b0__py3-none-any.whl

union/_internal/resolvers/common.py

@@ -0,0 +1,24 @@
+from asyncio import Protocol
+from typing import List
+
+from union._task import TaskTemplate
+
+
+class Resolver(Protocol):
+    """
+    Resolver interface for loading tasks. This interface should be implemented by Resolvers.
+    """
+
+    @property
+    def import_path(self) -> str: ...
+
+    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
+        """
+
+    def loader_args(self, t: TaskTemplate) -> 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.
+        """

union/_internal/resolvers/default.py

@@ -0,0 +1,27 @@
+import importlib
+from typing import List
+
+from union._internal.resolvers._task_module import extract_task_module
+from union._internal.resolvers.common import Resolver
+from union._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 "union._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) -> List[str]:  # type:ignore
+        t, m = extract_task_module(task)
+        return ["mod", m, "instance", t]

union/_internal/runtime/convert.py

@@ -0,0 +1,163 @@
+from dataclasses import dataclass
+from typing import Any, Dict, Tuple, Union
+
+from flyteidl.core import execution_pb2, literals_pb2
+
+import union.errors
+from union._datastructures import NativeInterface
+from union._protos.workflow import run_definition_pb2
+from union.types import TypeEngine
+
+
+@dataclass(frozen=True)
+class Inputs:
+    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}
+    inputs = await TypeEngine.literal_map_to_kwargs(
+        literals_pb2.LiteralMap(literals=literals), python_interface.get_input_types()
+    )
+    return inputs
+
+
+async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwargs) -> Inputs:
+    kwargs = interface.convert_to_kwargs(*args, **kwargs)
+    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) -> BaseException | 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 union.errors.RuntimeUnknownError(code=user_code, message=err.message, worker=err.worker)
+        case execution_pb2.ExecutionError.USER:
+            if "OOM" in err.code.upper():
+                return union.errors.OOMError(code=user_code, message=err.message, worker=err.worker)
+            elif "Interrupted" in err.code:
+                return union.errors.TaskInterruptedError(code=user_code, message=err.message, worker=err.worker)
+            elif "PrimaryContainerNotFound" in err.code:
+                return union.errors.PrimaryContainerNotFoundError(
+                    code=user_code, message=err.message, worker=err.worker
+                )
+            elif "RetriesExhausted" in err.code:
+                return union.errors.RetriesExhaustedError(code=user_code, message=err.message, worker=err.worker)
+            elif "Unknown" in err.code:
+                return union.errors.RuntimeUnknownError(code=user_code, message=err.message, worker=err.worker)
+            elif "InvalidImageName" in err.code:
+                return union.errors.InvalidImageNameError(code=user_code, message=err.message, worker=err.worker)
+            elif "ImagePullBackOff" in err.code:
+                return union.errors.ImagePullBackOffError(code=user_code, message=err.message, worker=err.worker)
+            return union.errors.RuntimeUserError(code=user_code, message=err.message, worker=err.worker)
+        case execution_pb2.ExecutionError.SYSTEM:
+            return union.errors.RuntimeSystemError(code=user_code, message=err.message, worker=err.worker)
+
+
+def convert_from_native_to_error(err: BaseException) -> Error:
+    if isinstance(err, union.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, union.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, union.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",
+            )
+        )

union/_internal/runtime/entrypoints.py

@@ -0,0 +1,121 @@
+from typing import List, Optional, Tuple
+
+from union._code_bundle import download_bundle
+from union._context import contextual_run
+from union._datastructures import ActionID, Checkpoints, CodeBundle, RawDataPath
+from union._internal import Controller
+from union._internal.imagebuild.image_builder import ImageCache
+from union._logging import logger
+from union._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,
+    checkpoints: Checkpoints | None = None,
+    code_bundle: CodeBundle | None = None,
+    inputs: Inputs | None = None,
+    output_path: str | None = None,
+) -> Tuple[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,
+        action=action,
+        raw_data_path=raw_data_path,
+        checkpoints=checkpoints,
+        code_bundle=code_bundle,
+        controller=controller,
+        version=version,
+        output_path=output_path,
+    )
+
+
+async def _download_and_load_task(
+    code_bundle: CodeBundle | None, resolver: str, resolver_args: List[str]
+) -> 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(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
+        logger.debug(
+            f"Loading task from tgz: {code_bundle.downloaded_path}, resolver: {resolver}, args: {resolver_args}"
+        )
+        return load_task(resolver, *resolver_args)
+    logger.debug(f"No code bundle provided, loading task from resolver: {resolver}, args: {resolver_args}")
+    return load_task(resolver, *resolver_args)
+
+
+async def load_and_run_task(
+    action: ActionID,
+    raw_data_path: RawDataPath,
+    output_path: str,
+    version: str,
+    controller: Controller,
+    resolver: str | None = None,
+    resolver_args: List[str] | None = None,
+    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 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,
+        checkpoints=checkpoints,
+        code_bundle=code_bundle,
+        input_path=input_path,
+        image_cache=image_cache,
+    )

union/_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 union.storage as storage
+from union._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,
+    )

union/_internal/runtime/resources_serde.py

@@ -0,0 +1,134 @@
+from typing import List, Optional, Tuple
+
+from flyteidl.core import tasks_pb2
+
+from union._resources import CPUBaseType, Resources
+
+ACCELERATOR_DEVICE_MAP = {
+    "A100": "nvidia-tesla-a100",
+    "A100 80G": "nvidia-a100-80gb",
+    "A10": "nvidia-a10",
+    "A10G": "nvidia-a10g",
+    "A100G": "nvidia-a100g",
+    "L4": "nvidia-l4",
+    "L40s": "nvidia-l40",
+    "L4_VWS": "nvidia-l4-vws",
+    "K80": "nvidia-tesla-k80",
+    "M60": "nvidia-tesla-m60",
+    "P4": "nvidia-tesla-p4",
+    "P100": "nvidia-tesla-p100",
+    "T4": "nvidia-tesla-t4",
+    "V100": "nvidia-tesla-v100",
+    "V5E": "tpu-v5-lite-podslice",
+    "V5P": "tpu-v5p-slice",
+    "V6E": "tpu-v6e-slice",
+}
+
+
+def _get_cpu_resource_entry(cpu: CPUBaseType) -> tasks_pb2.Resources.ResourceEntry:
+    return tasks_pb2.Resources.ResourceEntry(
+        name=tasks_pb2.Resources.ResourceName.CPU,
+        value=str(cpu),
+    )
+
+
+def _get_memory_resource_entry(memory: str) -> tasks_pb2.Resources.ResourceEntry:
+    return tasks_pb2.Resources.ResourceEntry(
+        name=tasks_pb2.Resources.ResourceName.MEMORY,
+        value=memory,
+    )
+
+
+def _get_gpu_resource_entry(gpu: int) -> tasks_pb2.Resources.ResourceEntry:
+    return tasks_pb2.Resources.ResourceEntry(
+        name=tasks_pb2.Resources.ResourceName.GPU,
+        value=str(gpu),
+    )
+
+
+def _get_gpu_extended_resource_entry(resources: Resources) -> Optional[tasks_pb2.GPUAccelerator]:
+    if resources is None:
+        return None
+    if resources.gpu is None or isinstance(resources.gpu, int):
+        return None
+    device = resources.get_device()
+    if device.device not in ACCELERATOR_DEVICE_MAP:
+        raise ValueError(f"GPU of type {device.device} unknown, cannot map to device name")
+    return tasks_pb2.GPUAccelerator(
+        device=ACCELERATOR_DEVICE_MAP[device.device],
+        partition_size=device.partition if device.partition else None,
+    )
+
+
+def _get_disk_resource_entry(disk: str) -> tasks_pb2.Resources.ResourceEntry:
+    return tasks_pb2.Resources.ResourceEntry(
+        name=tasks_pb2.Resources.ResourceName.EPHEMERAL_STORAGE,
+        value=disk,
+    )
+
+
+def get_proto_extended_resources(resources: Resources) -> Optional[tasks_pb2.ExtendedResources]:
+    """
+    TODO Implement partitioning logic string handling for GPU
+    :param resources:
+    """
+    if resources is None:
+        return None
+    acc = _get_gpu_extended_resource_entry(resources)
+    shm = resources.get_shared_memory()
+    if acc is None and shm is None:
+        return None
+    proto_shm = None
+    if shm is not None:
+        proto_shm = tasks_pb2.SharedMemory(
+            mount_path="/dev/shm",
+            mount_name="union-shm",
+            size_limit=shm,
+        )
+    return tasks_pb2.ExtendedResources(gpu_accelerator=acc, shared_memory=proto_shm)
+
+
+def _convert_resources_to_resource_entries(
+    resources: Resources | None,
+) -> Tuple[List[tasks_pb2.Resources.ResourceEntry], List[tasks_pb2.Resources.ResourceEntry]]:
+    request_entries, limit_entries = [], []
+    if resources is None:
+        return request_entries, limit_entries
+    if resources.cpu is not None:
+        if isinstance(resources.cpu, tuple):
+            request_entries.append(_get_cpu_resource_entry(resources.cpu[0]))
+            limit_entries.append(_get_cpu_resource_entry(resources.cpu[1]))
+        else:
+            request_entries.append(_get_cpu_resource_entry(resources.cpu))
+
+    if resources.memory is not None:
+        if isinstance(resources.memory, tuple):
+            request_entries.append(_get_memory_resource_entry(resources.memory[0]))
+            limit_entries.append(_get_memory_resource_entry(resources.memory[1]))
+        else:
+            request_entries.append(_get_memory_resource_entry(resources.memory))
+
+    if resources.gpu is not None:
+        device = resources.get_device()
+        request_entries.append(_get_gpu_resource_entry(device.quantity))
+
+    if resources.disk is not None:
+        request_entries.append(_get_disk_resource_entry(resources.disk))
+
+    return request_entries, limit_entries
+
+
+def get_proto_resources(resources: Resources) -> Optional[tasks_pb2.Resources]:
+    """
+    Get main resources IDL representation from the resources object
+
+    :param resources: User facing Resources object containing potentially both requests and limits
+    :return: The given resources as requests and limits
+    """
+    request_entries, limit_entries = _convert_resources_to_resource_entries(resources)
+    request_entries = request_entries if request_entries else None
+    limit_entries = limit_entries if limit_entries else None
+    if not request_entries and not limit_entries:
+        return None
+
+    return tasks_pb2.Resources(requests=request_entries, limits=limit_entries)