flyte 0.1.0__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

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_EXT = ".pkl.gz"
+
+
+def pkl_path(base_path: str, pkl_name: str) -> str:
+    return storage.join(base_path, f"{pkl_name}{_PKL_EXT}")
+
+
+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,
+    )

flyte/_internal/runtime/resources_serde.py

@@ -0,0 +1,138 @@
+from typing import List, Optional, Tuple
+
+from flyteidl.core import tasks_pb2
+
+from flyte._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 is None:
+        return None
+    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 | None) -> 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="flyte-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: List[tasks_pb2.Resources.ResourceEntry] = []
+    limit_entries: List[tasks_pb2.Resources.ResourceEntry] = []
+    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()
+        if device is not None:
+            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 | None) -> 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
+    """
+    if resources is None:
+        return None
+    request_entries, limit_entries = _convert_resources_to_resource_entries(resources)
+    if not request_entries and not limit_entries:
+        return None
+
+    return tasks_pb2.Resources(requests=request_entries, limits=limit_entries)

flyte/_internal/runtime/task_serde.py

@@ -0,0 +1,330 @@
+"""
+This module provides functionality to serialize and deserialize tasks to and from the wire format.
+It includes a Resolver interface for loading tasks, and functions to load classes and tasks.
+"""
+
+import copy
+import importlib
+import typing
+from datetime import timedelta
+from typing import Optional, Type, cast
+
+from flyteidl.core import identifier_pb2, literals_pb2, security_pb2, tasks_pb2
+from google.protobuf import duration_pb2, wrappers_pb2
+
+import flyte.errors
+from flyte._cache.cache import VersionParameters, cache_from_request
+from flyte._logging import logger
+from flyte._pod import _PRIMARY_CONTAINER_NAME_FIELD, PodTemplate
+from flyte._protos.workflow import common_pb2, environment_pb2, task_definition_pb2
+from flyte._secret import SecretRequest, secrets_from_request
+from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
+from flyte.models import CodeBundle, SerializationContext
+
+from ..._retry import RetryStrategy
+from ..._timeout import TimeoutType, timeout_from_request
+from .resources_serde import get_proto_extended_resources, get_proto_resources
+from .types_serde import transform_native_to_typed_interface
+
+_MAX_ENV_NAME_LENGTH = 63  # Maximum length for environment names
+_MAX_TASK_SHORT_NAME_LENGTH = 63  # Maximum length for task short names
+
+
+def load_class(qualified_name) -> Type:
+    """
+    Load a class from a qualified name. The qualified name should be in the format 'module.ClassName'.
+    :param qualified_name: The qualified name of the class to load.
+    :return: The class object.
+    """
+    module_name, class_name = qualified_name.rsplit(".", 1)  # Split module and class
+    module = importlib.import_module(module_name)  # Import the module
+    return getattr(module, class_name)  # Retrieve the class
+
+
+def load_task(resolver: str, *resolver_args: str) -> TaskTemplate:
+    """
+    Load a task from a resolver. This is a placeholder function.
+
+    :param resolver: The resolver to use to load the task.
+    :param resolver_args: Arguments to pass to the resolver.
+    :return: The loaded task.
+    """
+    resolver_class = load_class(resolver)
+    resolver_instance = resolver_class()
+    return resolver_instance.load_task(resolver_args)
+
+
+def translate_task_to_wire(
+    task: TaskTemplate,
+    serialization_context: SerializationContext,
+    default_inputs: Optional[typing.List[common_pb2.NamedParameter]] = None,
+) -> task_definition_pb2.TaskSpec:
+    """
+    Translate a task to a wire format. This is a placeholder function.
+
+    :param task: The task to translate.
+    :param serialization_context: The serialization context to use for the translation.
+    :param default_inputs: Optional list of default inputs for the task.
+
+    :return: The translated task.
+    """
+    tt = get_proto_task(task, serialization_context)
+    env: environment_pb2.Environment | None = None
+    if task.parent_env and task.parent_env():
+        _env = task.parent_env()
+        if _env:
+            env = environment_pb2.Environment(name=_env.name[:_MAX_ENV_NAME_LENGTH])
+    return task_definition_pb2.TaskSpec(
+        task_template=tt,
+        default_inputs=default_inputs,
+        short_name=task.friendly_name[:_MAX_TASK_SHORT_NAME_LENGTH],
+        environment=env,
+    )
+
+
+def get_security_context(secrets: Optional[SecretRequest]) -> Optional[security_pb2.SecurityContext]:
+    """
+    Get the security context from a list of secrets. This is a placeholder function.
+
+    :param secrets: The list of secrets to use for the security context.
+
+    :return: The security context.
+    """
+    if secrets is None:
+        return None
+
+    secret_list = secrets_from_request(secrets)
+    return security_pb2.SecurityContext(
+        secrets=[
+            security_pb2.Secret(
+                group=secret.group,
+                key=secret.key,
+                mount_requirement=(
+                    security_pb2.Secret.MountType.ENV_VAR if secret.as_env_var else security_pb2.Secret.MountType.FILE
+                ),
+                env_var=secret.as_env_var,
+            )
+            for secret in secret_list
+        ]
+    )
+
+
+def get_proto_retry_strategy(retries: RetryStrategy | int | None) -> Optional[literals_pb2.RetryStrategy]:
+    if retries is None:
+        return None
+
+    if isinstance(retries, int):
+        raise AssertionError("Retries should be an instance of RetryStrategy, not int")
+
+    return literals_pb2.RetryStrategy(retries=retries.count)
+
+
+def get_proto_timeout(timeout: TimeoutType | None) -> Optional[duration_pb2.Duration]:
+    if timeout is None:
+        return None
+    max_runtime_timeout = timeout_from_request(timeout).max_runtime
+    if isinstance(max_runtime_timeout, int):
+        max_runtime_timeout = timedelta(seconds=max_runtime_timeout)
+    return duration_pb2.Duration(seconds=max_runtime_timeout.seconds)
+
+
+def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext) -> tasks_pb2.TaskTemplate:
+    task_id = identifier_pb2.Identifier(
+        resource_type=identifier_pb2.ResourceType.TASK,
+        project=serialize_context.project,
+        domain=serialize_context.domain,
+        org=serialize_context.org,
+        name=task.name,
+        version=serialize_context.version,
+    )
+    # TODO, there will be tasks that do not have images, handle that case
+    # if task.parent_env is None:
+    # raise ValueError(f"Task {task.name} must have a parent environment")
+
+    # TODO Add support for SQL, extra_config, custom
+    extra_config: typing.Dict[str, str] = {}
+    custom = {}  # type: ignore
+
+    sql = None
+    if task.pod_template and not isinstance(task.pod_template, str):
+        container = None
+        pod = _get_k8s_pod(_get_urun_container(serialize_context, task), task.pod_template)
+        extra_config[_PRIMARY_CONTAINER_NAME_FIELD] = task.pod_template.primary_container_name
+    else:
+        container = _get_urun_container(serialize_context, task)
+        pod = None
+
+    # -------------- CACHE HANDLING ----------------------
+    task_cache = cache_from_request(task.cache)
+    cache_enabled = task_cache.is_enabled()
+    cache_version = None
+
+    if task_cache.is_enabled():
+        logger.debug(f"Cache enabled for task {task.name}")
+        if serialize_context.code_bundle and serialize_context.code_bundle.pkl:
+            logger.debug(f"Detected pkl bundle for task {task.name}, using computed version as cache version")
+            cache_version = serialize_context.code_bundle.computed_version
+        else:
+            version_parameters = None
+            if isinstance(task, AsyncFunctionTaskTemplate):
+                version_parameters = VersionParameters(func=task.func, image=task.image)
+            else:
+                version_parameters = VersionParameters(func=None, image=task.image)
+            cache_version = task_cache.get_version(version_parameters)
+            logger.debug(f"Cache version for task {task.name} is {cache_version}")
+    else:
+        logger.debug(f"Cache disabled for task {task.name}")
+
+    return tasks_pb2.TaskTemplate(
+        id=task_id,
+        type=task.task_type,
+        metadata=tasks_pb2.TaskMetadata(
+            discoverable=cache_enabled,
+            discovery_version=cache_version,
+            cache_serializable=task_cache.serialize,
+            cache_ignore_input_vars=task_cache.get_ignored_inputs() if cache_enabled else None,
+            runtime=tasks_pb2.RuntimeMetadata(
+                version=flyte.version(),
+                type=tasks_pb2.RuntimeMetadata.RuntimeType.FLYTE_SDK,
+                flavor="python",
+            ),
+            retries=get_proto_retry_strategy(task.retries),
+            timeout=get_proto_timeout(task.timeout),
+            pod_template_name=task.pod_template if task.pod_template and isinstance(task.pod_template, str) else None,
+            interruptible=task.interruptable,
+            generates_deck=wrappers_pb2.BoolValue(value=task.report),
+        ),
+        interface=transform_native_to_typed_interface(task.native_interface),
+        custom=custom,
+        container=container,
+        task_type_version=task.task_type_version,
+        security_context=get_security_context(task.secrets),
+        config=extra_config,
+        k8s_pod=pod,
+        sql=sql,
+        extended_resources=get_proto_extended_resources(task.resources),
+    )
+
+
+def _get_urun_container(
+    serialize_context: SerializationContext, task_template: TaskTemplate
+) -> Optional[tasks_pb2.Container]:
+    env = (
+        [literals_pb2.KeyValuePair(key=k, value=v) for k, v in task_template.env.items()] if task_template.env else None
+    )
+    resources = get_proto_resources(task_template.resources)
+    # pr: under what conditions should this return None?
+    if isinstance(task_template.image, str):
+        raise flyte.errors.RuntimeSystemError("BadConfig", "Image is not a valid image")
+    image_id = task_template.image.identifier
+    if not serialize_context.image_cache or image_id not in serialize_context.image_cache.image_lookup:
+        # This computes the image uri, computing hashes as necessary so can fail if done remotely.
+        img_uri = task_template.image.uri
+    else:
+        img_uri = serialize_context.image_cache.image_lookup[image_id]
+
+    return tasks_pb2.Container(
+        image=img_uri,
+        command=[],
+        args=task_template.container_args(serialize_context),
+        resources=resources,
+        env=env,
+        data_config=task_template.data_loading_config(serialize_context),
+        config=task_template.config(serialize_context),
+    )
+
+
+def _sanitize_resource_name(resource: tasks_pb2.Resources.ResourceEntry) -> str:
+    return tasks_pb2.Resources.ResourceName.Name(resource.name).lower().replace("_", "-")
+
+
+def _get_k8s_pod(primary_container: tasks_pb2.Container, pod_template: PodTemplate) -> Optional[tasks_pb2.K8sPod]:
+    """
+    Get the K8sPod representation of the task template.
+    :param task: The task to convert.
+    :return: The K8sPod representation of the task template.
+    """
+    from kubernetes.client import ApiClient, V1PodSpec
+    from kubernetes.client.models import V1EnvVar, V1ResourceRequirements
+
+    pod_template = copy.deepcopy(pod_template)
+    containers = cast(V1PodSpec, pod_template.pod_spec).containers
+    primary_exists = False
+
+    for container in containers:
+        if container.name == pod_template.primary_container_name:
+            primary_exists = True
+            break
+
+    if not primary_exists:
+        raise ValueError(
+            "No primary container defined in the pod spec."
+            f" You must define a primary container with the name '{pod_template.primary_container_name}'."
+        )
+    final_containers = []
+
+    for container in containers:
+        # We overwrite the primary container attributes with the values given to ContainerTask.
+        # The attributes include: image, command, args, resource, and env (env is unioned)
+
+        if container.name == pod_template.primary_container_name:
+            if container.image is None:
+                # Copy the image from primary_container only if the image is not specified in the pod spec.
+                container.image = primary_container.image
+
+            container.command = list(primary_container.command)
+            container.args = list(primary_container.args)
+
+            limits, requests = {}, {}
+            for resource in primary_container.resources.limits:
+                limits[_sanitize_resource_name(resource)] = resource.value
+            for resource in primary_container.resources.requests:
+                requests[_sanitize_resource_name(resource)] = resource.value
+
+            resource_requirements = V1ResourceRequirements(limits=limits, requests=requests)
+            if len(limits) > 0 or len(requests) > 0:
+                # Important! Only copy over resource requirements if they are non-empty.
+                container.resources = resource_requirements
+
+            if primary_container.env is not None:
+                container.env = [V1EnvVar(name=e.key, value=e.value) for e in primary_container.env] + (
+                    container.env or []
+                )
+
+        final_containers.append(container)
+
+    cast(V1PodSpec, pod_template.pod_spec).containers = final_containers
+    pod_spec = ApiClient().sanitize_for_serialization(pod_template.pod_spec)
+
+    metadata = tasks_pb2.K8sObjectMetadata(labels=pod_template.labels, annotations=pod_template.annotations)
+    return tasks_pb2.K8sPod(pod_spec=pod_spec, metadata=metadata)
+
+
+def extract_code_bundle(task_spec: task_definition_pb2.TaskSpec) -> Optional[CodeBundle]:
+    """
+    Extract the code bundle from the task spec.
+    :param task_spec: The task spec to extract the code bundle from.
+    :return: The extracted code bundle or None if not present.
+    """
+    container = task_spec.task_template.container
+    if container and container.args:
+        pkl_path = None
+        tgz_path = None
+        dest_path: str = "."
+        version = ""
+        for i, v in enumerate(container.args):
+            if v == "--pkl":
+                # Extract the code bundle path from the argument
+                pkl_path = container.args[i + 1] if i + 1 < len(container.args) else None
+            elif v == "--tgz":
+                # Extract the code bundle path from the argument
+                tgz_path = container.args[i + 1] if i + 1 < len(container.args) else None
+            elif v == "--dest":
+                # Extract the destination path from the argument
+                dest_path = container.args[i + 1] if i + 1 < len(container.args) else "."
+            elif v == "--version":
+                # Extract the version from the argument
+                version = container.args[i + 1] if i + 1 < len(container.args) else ""
+        if pkl_path or tgz_path:
+            return CodeBundle(destination=dest_path, tgz=tgz_path, pkl=pkl_path, computed_version=version)
+    return None