flyte 0.0.1b0__py3-none-any.whl

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,210 @@
+"""
+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 importlib
+from datetime import timedelta
+from typing import Optional, Type
+
+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._datastructures import SerializationContext
+from flyte._logging import logger
+from flyte._protos.workflow import task_definition_pb2
+from flyte._secret import SecretRequest, secrets_from_request
+from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
+
+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
+
+
+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
+) -> 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.
+
+    :return: The translated task.
+    """
+    return get_proto_task(task, serialization_context)
+
+
+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) -> task_definition_pb2.TaskSpec:
+    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")
+
+    #
+    # This pod will be incorrect when doing fast serialize
+    #
+    container = _get_urun_container(serialize_context, task)
+
+    # TODO Add support for SQL, Pod, extra_config, custom
+    pod = None
+    sql = None
+    # pod = task.get_k8s_pod(serialize_context)
+    extra_config = {}  # type: ignore
+    custom = {}  # type: ignore
+
+    # -------------- 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}")
+
+    tt = 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.ignored_inputs,
+            runtime=tasks_pb2.RuntimeMetadata(),
+            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=False),  # TODO add support for reports
+        ),
+        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),
+    )
+    return task_definition_pb2.TaskSpec(task_template=tt)
+
+
+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),
+    )

flyte/_internal/runtime/taskrunner.py

@@ -0,0 +1,190 @@
+"""
+This module is responsible for running tasks in the V2 runtime. All methods in this file should be
+invoked within a context tree.
+"""
+
+import pathlib
+from typing import Any, Dict, List, Optional, Tuple
+
+import flyte.report
+from flyte._context import internal_ctx
+from flyte._datastructures import ActionID, Checkpoints, CodeBundle, RawDataPath, TaskContext
+from flyte._internal.imagebuild.image_builder import ImageCache
+from flyte._logging import log, logger
+from flyte._task import TaskTemplate
+from flyte.errors import CustomError, RuntimeSystemError, RuntimeUnknownError, RuntimeUserError
+
+from .. import Controller
+from .convert import (
+    Error,
+    Inputs,
+    Outputs,
+    convert_from_native_to_error,
+    convert_from_native_to_outputs,
+    convert_inputs_to_native,
+)
+from .io import load_inputs, upload_error, upload_outputs
+
+
+def replace_task_cli(args: List[str], inputs: Inputs, tmp_path: pathlib.Path, action: ActionID) -> List[str]:
+    """
+    This method can be used to run an task from the cli, if you have cli for the task. It will replace,
+    all the args with the task args.
+
+    The urun cli is of the format
+    ```python
+    ['urun', '--inputs', '{{.Inputs}}', '--outputs-path', '{{.Outputs}}', '--version', '',
+     '--raw-data-path', '{{.rawOutputDataPrefix}}',
+      '--checkpoint-path', '{{.checkpointOutputPrefix}}', '--prev-checkpoint', '{{.prevCheckpointPrefix}}',
+       '--run-name', '{{.runName}}', '--name', '{{.actionName}}',
+        '--tgz', 'some-path', '--dest', '.',
+         '--resolver', 'flyte._internal.resolvers.default.DefaultTaskResolver', '--resolver-args',
+          'mod', 'test_round_trip', 'instance', 'task1']
+    ```
+    We will replace, inputs, outputs, raw_data_path, checkpoint_path, prev_checkpoint, run_name, name
+    with supplied values.
+
+    :param args: urun command
+    :param inputs: converted inputs to the task
+    :param tmp_path: temporary path to use for the task
+    :param action: run id to use for the task
+    :return: modified args
+    """
+    # Iterate over all the args and replace the inputs, outputs, raw_data_path, checkpoint_path, prev_checkpoint,
+    # root_name, run_name with the supplied values
+    # first we will write the inputs to a file called inputs.pb
+    inputs_path = tmp_path / "inputs.pb"
+    with open(inputs_path, "wb") as f:
+        f.write(inputs.proto_inputs.SerializeToString())
+    # now modify the args
+    args = list(args)  # copy first because it's a proto container
+    for i, arg in enumerate(args):
+        match arg:
+            case "--inputs":
+                args[i + 1] = str(inputs_path)
+            case "--outputs-path":
+                args[i + 1] = str(tmp_path)
+            case "--raw-data-path":
+                args[i + 1] = str(tmp_path / "raw_data_path")
+            case "--checkpoint-path":
+                args[i + 1] = str(tmp_path / "checkpoint_path")
+            case "--prev-checkpoint":
+                args[i + 1] = str(tmp_path / "prev_checkpoint")
+            case "--run-name":
+                args[i + 1] = action.run_name or ""
+            case "--name":
+                args[i + 1] = action.name
+    insert_point = args.index("--raw-data-path")
+    args.insert(insert_point, str(tmp_path))
+    args.insert(insert_point, "--run-base-dir")
+    return args
+
+
+@log
+async def run_task(
+    tctx: TaskContext, controller: Controller, task: TaskTemplate, inputs: Dict[str, Any]
+) -> Tuple[Any, Optional[Exception]]:
+    try:
+        logger.info(f"Parent task executing {tctx.action}")
+        outputs = await task.execute(**inputs)
+        logger.info(f"Parent task completed successfully, {tctx.action}")
+        return outputs, None
+    except RuntimeSystemError as e:
+        logger.exception(f"Task failed with error: {e}")
+        return {}, e
+    except RuntimeUnknownError as e:
+        logger.exception(f"Task failed with error: {e}")
+        return {}, e
+    except RuntimeUserError as e:
+        logger.exception(f"Task failed with error: {e}")
+        return {}, e
+    except Exception as e:
+        logger.exception(f"Task failed with error: {e}")
+        return {}, CustomError.from_exception(e)
+    finally:
+        logger.info(f"Parent task finalized {tctx.action}")
+        # reconstruct run id here
+        await controller.finalize_parent_action(tctx.action)
+
+
+async def convert_and_run(
+    *,
+    task: TaskTemplate,
+    inputs: Inputs,
+    action: ActionID,
+    controller: Controller,
+    raw_data_path: RawDataPath,
+    version: str,
+    output_path: str,
+    run_base_dir: str,
+    checkpoints: Checkpoints | None = None,
+    code_bundle: CodeBundle | None = None,
+    image_cache: ImageCache | None = None,
+) -> Tuple[Optional[Outputs], Optional[Error]]:
+    """
+    This method is used to convert the inputs to native types, and run the task. It assumes you are running
+    in a context tree.
+    """
+    ctx = internal_ctx()
+    tctx = TaskContext(
+        action=action,
+        checkpoints=checkpoints,
+        code_bundle=code_bundle,
+        output_path=output_path,
+        run_base_dir=run_base_dir,
+        version=version,
+        raw_data_path=raw_data_path,
+        compiled_image_cache=image_cache,
+        report=flyte.report.Report(name=action.name),
+    )
+    async with ctx.replace_task_context(tctx):
+        inputs_kwargs = await convert_inputs_to_native(inputs, task.native_interface)
+        out, err = await run_task(tctx=tctx, controller=controller, task=task, inputs=inputs_kwargs)
+        if err is not None:
+            return None, convert_from_native_to_error(err)
+        if task.report:
+            await flyte.report.flush.aio()
+        return await convert_from_native_to_outputs(out, task.native_interface), None
+
+
+async def extract_download_run_upload(
+    task: TaskTemplate,
+    *,
+    action: ActionID,
+    controller: Controller,
+    raw_data_path: RawDataPath,
+    output_path: str,
+    run_base_dir: str,
+    version: 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 CLI (urun) and is used to run a task. This assumes that the context tree
+    has already been created, and the task has been loaded. It also handles the loading of the task.
+    """
+    inputs = await load_inputs(input_path) if input_path else None
+    outputs, err = await convert_and_run(
+        task=task,
+        inputs=inputs or Inputs.empty(),
+        action=action,
+        controller=controller,
+        raw_data_path=raw_data_path,
+        output_path=output_path,
+        run_base_dir=run_base_dir,
+        version=version,
+        checkpoints=checkpoints,
+        code_bundle=code_bundle,
+        image_cache=image_cache,
+    )
+    if err is not None:
+        path = await upload_error(err.err, output_path)
+        logger.error(f"Task {task.name} failed with error: {err}. Uploaded error to {path}")
+        return
+    if outputs is None:
+        logger.info(f"Task {task.name} completed successfully, no outputs")
+        return
+    await upload_outputs(outputs, output_path) if output_path else None
+    logger.info(f"Task {task.name} completed successfully, uploaded outputs to {output_path}")

flyte/_internal/runtime/types_serde.py

@@ -0,0 +1,54 @@
+from typing import Dict, Optional, TypeVar
+
+from flyteidl.core import interface_pb2
+
+from flyte._datastructures import NativeInterface
+from flyte.types._type_engine import TypeEngine
+
+T = TypeVar("T")
+
+
+def transform_variable_map(
+    variable_map: Dict[str, type],
+    descriptions: Optional[Dict[str, str]] = None,
+) -> Dict[str, interface_pb2.Variable]:
+    """
+    Given a map of str (names of inputs for instance) to their Python native types, return a map of the name to a
+    Flyte Variable object with that type.
+    """
+    res = {}
+    descriptions = descriptions or {}
+    if variable_map:
+        for k, v in variable_map.items():
+            res[k] = transform_type(v, descriptions.get(k, k))
+    return res
+
+
+def transform_native_to_typed_interface(
+    interface: Optional[NativeInterface],
+) -> Optional[interface_pb2.TypedInterface]:
+    """
+    Transform the given simple python native interface to FlyteIDL's interface
+    """
+    if interface is None:
+        return None
+    input_descriptions: Dict[str, str] = {}
+    output_descriptions: Dict[str, str] = {}
+    if interface.docstring:
+        # Fill in descriptions from docstring in the future
+        input_descriptions = {}
+        output_descriptions = {}
+
+    inputs_map = transform_variable_map(interface.get_input_types(), input_descriptions)
+    outputs_map = transform_variable_map(interface.outputs, output_descriptions)
+    return interface_pb2.TypedInterface(
+        inputs=interface_pb2.VariableMap(variables=inputs_map), outputs=interface_pb2.VariableMap(variables=outputs_map)
+    )
+
+
+def transform_type(x: type, description: Optional[str] = None) -> interface_pb2.Variable:
+    # add artifact handling eventually
+    return interface_pb2.Variable(
+        type=TypeEngine.to_literal_type(x),
+        description=description,
+    )