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

union/_internal/resolvers/common.py

@@ -1,24 +0,0 @@
-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

@@ -1,27 +0,0 @@
-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

@@ -1,163 +0,0 @@
-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

@@ -1,121 +0,0 @@
-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

@@ -1,136 +0,0 @@
-"""
-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

@@ -1,134 +0,0 @@
-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)