flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_internal/runtime/io.py

@@ -5,14 +5,13 @@ It uses the storage module to handle the actual uploading and downloading of fil
 TODO: Convert to use streaming apis
 """
 
-import logging
-
-from flyteidl.core import errors_pb2, execution_pb2
+from flyteidl2.core import errors_pb2, execution_pb2
+from flyteidl2.task import common_pb2
 
 import flyte.storage as storage
-from flyte._protos.workflow import run_definition_pb2
+from flyte._logging import logger
+from flyte.models import PathRewrite
 
-from ..._logging import log
 from .convert import Inputs, Outputs, _clean_error_code
 
 # ------------------------------- CONSTANTS ------------------------------- #
@@ -21,11 +20,11 @@ _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"
+_PKL_EXT = ".pkl.gz"
 
 
-def pkl_path(base_path: str) -> str:
-    return storage.join(base_path, _PKL_FILE_NAME)
+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:
@@ -55,16 +54,25 @@ async def upload_inputs(inputs: Inputs, input_path: str):
     await storage.put_stream(data_iterable=inputs.proto_inputs.SerializeToString(), to_path=input_path)
 
 
-async def upload_outputs(outputs: Outputs, output_path: str):
+async def upload_outputs(outputs: Outputs, output_path: str, max_bytes: int = -1):
     """
     :param outputs: Outputs
     :param output_path: The path to upload the output file.
+    :param max_bytes: Maximum number of bytes to write to the output file. Default is -1, which means no limit.
     """
+    if max_bytes != -1 and outputs.proto_outputs.ByteSize() > max_bytes:
+        import flyte.errors
+
+        raise flyte.errors.InlineIOMaxBytesBreached(
+            f"Output file at {output_path} exceeds max_bytes limit of {max_bytes},"
+            f" size: {outputs.proto_outputs.ByteSize()}"
+        )
     output_uri = outputs_path(output_path)
     await storage.put_stream(data_iterable=outputs.proto_outputs.SerializeToString(), to_path=output_uri)
+    logger.debug(f"Uploaded {output_uri} to {output_path}")
 
 
-async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str):
+async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str) -> str:
     """
     :param err: execution_pb2.ExecutionError
     :param output_prefix: The output prefix of the remote uri.
@@ -76,34 +84,76 @@ async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str):
             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)
+    return 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:
+async def load_inputs(path: str, max_bytes: int = -1, path_rewrite_config: PathRewrite | None = None) -> Inputs:
     """
     :param path: Input file to be downloaded
+    :param max_bytes: Maximum number of bytes to read from the input file. Default is -1, which means no limit.
+    :param path_rewrite_config: If provided, rewrites paths in the input blobs according to the configuration.
     :return: Inputs object
     """
-    lm = run_definition_pb2.Inputs()
-    proto_str = b"".join([c async for c in storage.get_stream(path=path)])
+    lm = common_pb2.Inputs()
+    if max_bytes == -1:
+        proto_str = b"".join([c async for c in storage.get_stream(path=path)])
+    else:
+        proto_bytes = []
+        total_bytes = 0
+        async for chunk in storage.get_stream(path=path):
+            if total_bytes + len(chunk) > max_bytes:
+                import flyte.errors
+
+                raise flyte.errors.InlineIOMaxBytesBreached(
+                    f"Input file at {path} exceeds max_bytes limit of {max_bytes}"
+                )
+            proto_bytes.append(chunk)
+            total_bytes += len(chunk)
+        proto_str = b"".join(proto_bytes)
+
     lm.ParseFromString(proto_str)
+
+    if path_rewrite_config is not None:
+        for inp in lm.literals:
+            if inp.value.HasField("scalar") and inp.value.scalar.HasField("blob"):
+                scalar_blob = inp.value.scalar.blob
+                if scalar_blob.uri.startswith(path_rewrite_config.old_prefix):
+                    scalar_blob.uri = scalar_blob.uri.replace(
+                        path_rewrite_config.old_prefix, path_rewrite_config.new_prefix, 1
+                    )
+
     return Inputs(proto_inputs=lm)
 
 
-async def load_outputs(path: str) -> Outputs:
+async def load_outputs(path: str, max_bytes: int = -1) -> Outputs:
     """
     :param path: output file to be loaded
+    :param max_bytes: Maximum number of bytes to read from the output file.
+                      If -1, reads the entire file.
     :return: Outputs object
     """
-    lm = run_definition_pb2.Outputs()
-    proto_str = b"".join([c async for c in storage.get_stream(path=path)])
+    lm = common_pb2.Outputs()
+
+    if max_bytes == -1:
+        proto_str = b"".join([c async for c in storage.get_stream(path=path)])
+    else:
+        proto_bytes = []
+        total_bytes = 0
+        async for chunk in storage.get_stream(path=path):
+            if total_bytes + len(chunk) > max_bytes:
+                import flyte.errors
+
+                raise flyte.errors.InlineIOMaxBytesBreached(
+                    f"Output file at {path} exceeds max_bytes limit of {max_bytes}"
+                )
+            proto_bytes.append(chunk)
+            total_bytes += len(chunk)
+        proto_str = b"".join(proto_bytes)
+
     lm.ParseFromString(proto_str)
     return Outputs(proto_outputs=lm)
 
@@ -118,14 +168,12 @@ async def load_error(path: str) -> execution_pb2.ExecutionError:
     err.ParseFromString(proto_str)
 
     if err.error is not None:
-        user_code, server_code = _clean_error_code(err.error.code)
+        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(

flyte/_internal/runtime/resources_serde.py

@@ -1,8 +1,8 @@
-from typing import List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
-from flyteidl.core import tasks_pb2
+from flyteidl2.core import tasks_pb2
 
-from flyte._resources import CPUBaseType, Resources
+from flyte._resources import CPUBaseType, DeviceClass, Resources
 
 ACCELERATOR_DEVICE_MAP = {
     "A100": "nvidia-tesla-a100",
@@ -11,7 +11,7 @@ ACCELERATOR_DEVICE_MAP = {
     "A10G": "nvidia-a10g",
     "A100G": "nvidia-a100g",
     "L4": "nvidia-l4",
-    "L40s": "nvidia-l40",
+    "L40s": "nvidia-l40s",
     "L4_VWS": "nvidia-l4-vws",
     "K80": "nvidia-tesla-k80",
     "M60": "nvidia-tesla-m60",
@@ -24,6 +24,14 @@ ACCELERATOR_DEVICE_MAP = {
     "V6E": "tpu-v6e-slice",
 }
 
+_DeviceClassToProto: Dict[DeviceClass, "tasks_pb2.GPUAccelerator.DeviceClass"] = {
+    "GPU": tasks_pb2.GPUAccelerator.NVIDIA_GPU,
+    "TPU": tasks_pb2.GPUAccelerator.GOOGLE_TPU,
+    "NEURON": tasks_pb2.GPUAccelerator.AMAZON_NEURON,
+    "AMD_GPU": tasks_pb2.GPUAccelerator.AMD_GPU,
+    "HABANA_GAUDI": tasks_pb2.GPUAccelerator.HABANA_GAUDI,
+}
+
 
 def _get_cpu_resource_entry(cpu: CPUBaseType) -> tasks_pb2.Resources.ResourceEntry:
     return tasks_pb2.Resources.ResourceEntry(
@@ -54,11 +62,17 @@ def _get_gpu_extended_resource_entry(resources: Resources) -> Optional[tasks_pb2
     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")
+
+    device_class = _DeviceClassToProto.get(device.device_class, tasks_pb2.GPUAccelerator.NVIDIA_GPU)
+    if device.device is None:
+        raise RuntimeError("Device type must be specified for GPU string.")
+    else:
+        device_type = device.device
+    device_type = ACCELERATOR_DEVICE_MAP.get(device_type, device_type)
     return tasks_pb2.GPUAccelerator(
-        device=ACCELERATOR_DEVICE_MAP[device.device],
+        device=device_type,
         partition_size=device.partition if device.partition else None,
+        device_class=device_class,
     )
 
 

flyte/_internal/runtime/reuse.py

@@ -0,0 +1,125 @@
+import hashlib
+import typing
+from venv import logger
+
+from flyteidl2.core import tasks_pb2
+
+import flyte.errors
+from flyte import ReusePolicy
+from flyte._pod import _PRIMARY_CONTAINER_DEFAULT_NAME, _PRIMARY_CONTAINER_NAME_FIELD
+from flyte.models import CodeBundle
+
+
+def extract_unique_id_and_image(
+    env_name: str,
+    code_bundle: CodeBundle | None,
+    task: tasks_pb2.TaskTemplate,
+    reuse_policy: ReusePolicy,
+) -> typing.Tuple[str, str]:
+    """
+    Compute a unique ID for the task based on its name, version, image URI, and code bundle.
+    :param env_name: Name of the reusable environment.
+    :param reuse_policy: The reuse policy for the task.
+    :param task: The task template.
+    :param code_bundle: The code bundle associated with the task.
+    :return: A unique ID string and the image URI.
+    """
+    image = ""
+    container_ser = ""
+    if task.HasField("container"):
+        copied_container = tasks_pb2.Container()
+        copied_container.CopyFrom(task.container)
+        copied_container.args.clear()  # Clear args to ensure deterministic serialization
+        container_ser = copied_container.SerializeToString(deterministic=True)
+        image = copied_container.image
+
+    if task.HasField("k8s_pod"):
+        # Clear args to ensure deterministic serialization
+        copied_k8s_pod = tasks_pb2.K8sPod()
+        copied_k8s_pod.CopyFrom(task.k8s_pod)
+        if task.config is not None:
+            primary_container_name = task.config[_PRIMARY_CONTAINER_NAME_FIELD]
+        else:
+            primary_container_name = _PRIMARY_CONTAINER_DEFAULT_NAME
+        for container in copied_k8s_pod.pod_spec["containers"]:
+            if "name" in container and container["name"] == primary_container_name:
+                image = container["image"]
+                del container["args"]
+        container_ser = copied_k8s_pod.SerializeToString(deterministic=True)
+
+    components = f"{env_name}:{container_ser}"
+    if isinstance(reuse_policy.replicas, tuple):
+        components += f":{reuse_policy.replicas[0]}:{reuse_policy.replicas[1]}"
+    else:
+        components += f":{reuse_policy.replicas}"
+    if reuse_policy.idle_ttl:
+        components += f":{reuse_policy.idle_ttl.total_seconds()}"  # type: ignore [union-attr]
+    if reuse_policy.get_scaledown_ttl() is not None:
+        components += f":{reuse_policy.get_scaledown_ttl()}"
+    if code_bundle is not None:
+        components += f":{code_bundle.computed_version}"
+    if task.security_context is not None:
+        security_ctx_str = task.security_context.SerializeToString(deterministic=True)
+        components += f":{security_ctx_str}"
+    if task.metadata.interruptible is not None:
+        components += f":{task.metadata.interruptible}"
+    if task.metadata.pod_template_name is not None:
+        components += f":{task.metadata.pod_template_name}"
+    sha256 = hashlib.sha256()
+    sha256.update(components.encode("utf-8"))
+    return sha256.hexdigest(), image
+
+
+def add_reusable(
+    task: tasks_pb2.TaskTemplate,
+    reuse_policy: ReusePolicy,
+    code_bundle: CodeBundle | None,
+    parent_env_name: str | None = None,
+) -> tasks_pb2.TaskTemplate:
+    """
+    Convert a ReusePolicy to a custom configuration dictionary.
+
+    :param task: The task to which the reusable policy will be added.
+    :param reuse_policy: The reuse policy to apply.
+    :param code_bundle: The code bundle associated with the task.
+    :param parent_env_name: The name of the parent environment, if any.
+    :return: The modified task with the reusable policy added.
+    """
+    if reuse_policy is None:
+        return task
+
+    if task.HasField("custom"):
+        raise flyte.errors.RuntimeUserError(
+            "BadConfiguration", "Plugins do not support reusable policy. Only container tasks and pods."
+        )
+
+    logger.debug(f"Adding reusable policy for task: {task.id.name}")
+    name = parent_env_name if parent_env_name else ""
+    if parent_env_name is None:
+        name = task.id.name.split(".")[0]
+
+    version, image_uri = extract_unique_id_and_image(
+        env_name=name, code_bundle=code_bundle, task=task, reuse_policy=reuse_policy
+    )
+
+    scaledown_ttl = reuse_policy.get_scaledown_ttl()
+
+    task.custom = {
+        "name": name,
+        "version": version[:15],  # Use only the first 15 characters for the version
+        "type": "actor",
+        "spec": {
+            "container_image": image_uri,
+            "backlog_length": None,
+            "parallelism": reuse_policy.concurrency,
+            "min_replica_count": reuse_policy.min_replicas,
+            "replica_count": reuse_policy.max_replicas,
+            "ttl_seconds": reuse_policy.idle_ttl.total_seconds() if reuse_policy.idle_ttl else None,  # type: ignore [union-attr]
+            "scaledown_ttl_seconds": scaledown_ttl.total_seconds() if scaledown_ttl else None,
+        },
+    }
+
+    task.type = "actor"
+    logger.info(f"Reusable task {task.id.name} with config {task.custom}")
+
+    return task

flyte/_internal/runtime/rusty.py

@@ -0,0 +1,196 @@
+import asyncio
+import time
+from typing import List, Tuple
+
+from flyte._context import contextual_run
+from flyte._internal.controllers import Controller
+from flyte._internal.controllers import create_controller as _create_controller
+from flyte._internal.imagebuild.image_builder import ImageCache
+from flyte._internal.runtime.entrypoints import download_code_bundle, load_pkl_task, load_task
+from flyte._internal.runtime.taskrunner import extract_download_run_upload
+from flyte._logging import logger
+from flyte._task import TaskTemplate
+from flyte._utils import adjust_sys_path
+from flyte.models import ActionID, Checkpoints, CodeBundle, PathRewrite, RawDataPath
+
+
+async def download_tgz(destination: str, version: str, tgz: str) -> CodeBundle:
+    """
+    Downloads and loads the task from the code bundle or resolver.
+    :param tgz: The path to the task template in a tar.gz format.
+    :param destination: The path to save the downloaded task template.
+    :param version: The version of the task to load.
+    :return: The CodeBundle object.
+    """
+    logger.info(f"[rusty] Downloading tgz code bundle from {tgz} to {destination} with version {version}")
+    adjust_sys_path()
+
+    code_bundle = CodeBundle(
+        tgz=tgz,
+        destination=destination,
+        computed_version=version,
+    )
+    return await download_code_bundle(code_bundle)
+
+
+async def download_load_pkl(destination: str, version: str, pkl: str) -> Tuple[CodeBundle, TaskTemplate]:
+    """
+    Downloads and loads the task from the code bundle or resolver.
+    :param pkl: The path to the task template in a pickle format.
+    :param destination: The path to save the downloaded task template.
+    :param version: The version of the task to load.
+    :return: The CodeBundle object.
+    """
+    logger.info(f"[rusty] Downloading pkl code bundle from {pkl} to {destination} with version {version}")
+    adjust_sys_path()
+
+    code_bundle = CodeBundle(
+        pkl=pkl,
+        destination=destination,
+        computed_version=version,
+    )
+    code_bundle = await download_code_bundle(code_bundle)
+    return code_bundle, load_pkl_task(code_bundle)
+
+
+def load_task_from_code_bundle(resolver: str, resolver_args: List[str]) -> TaskTemplate:
+    """
+    Loads the task from the code bundle or resolver.
+    :param resolver: The resolver to use to load the task.
+    :param resolver_args: The arguments to pass to the resolver.
+    :return: The loaded task template.
+    """
+    logger.debug(f"[rusty] Loading task from code bundle {resolver} with args: {resolver_args}")
+    return load_task(resolver, *resolver_args)
+
+
+async def create_controller(
+    endpoint: str = "host.docker.internal:8090",
+    insecure: bool = False,
+    api_key: str | None = None,
+) -> Controller:
+    """
+    Creates a controller instance for remote operations.
+    :param endpoint:
+    :param insecure:
+    :param api_key:
+    :return:
+    """
+    logger.info(f"[rusty] Creating controller with endpoint {endpoint}")
+    import flyte.errors
+    from flyte._initialize import init_in_cluster
+
+    loop = asyncio.get_event_loop()
+    loop.set_exception_handler(flyte.errors.silence_grpc_polling_error)
+
+    # TODO Currently remote tasks are not supported in Rusty.
+    controller_kwargs = await init_in_cluster.aio(api_key=api_key, endpoint=endpoint, insecure=insecure)
+    return _create_controller(ct="remote", **controller_kwargs)
+
+
+async def run_task(
+    task: TaskTemplate,
+    controller: Controller,
+    org: str,
+    project: str,
+    domain: str,
+    run_name: str,
+    name: str,
+    raw_data_path: str,
+    output_path: str,
+    run_base_dir: str,
+    version: str,
+    image_cache: str | None = None,
+    checkpoint_path: str | None = None,
+    prev_checkpoint: str | None = None,
+    code_bundle: CodeBundle | None = None,
+    input_path: str | None = None,
+    path_rewrite_cfg: str | None = None,
+):
+    """
+    Runs the task with the provided parameters.
+    :param prev_checkpoint: Previous checkpoint path to resume from.
+    :param checkpoint_path: Checkpoint path to save the current state.
+    :param image_cache: Image cache to use for the task.
+    :param name: Action name to run.
+    :param run_name: Parent run name to use for the task.
+    :param domain: domain to run the task in.
+    :param project: project to run the task in.
+    :param org: organization to run the task in.
+    :param task: The task template to run.
+    :param raw_data_path: The path to the raw data.
+    :param output_path: The path to save the output.
+    :param run_base_dir: The base directory for the run.
+    :param version: The version of the task to run.
+    :param controller: The controller to use for the task.
+    :param code_bundle: Optional code bundle for the task.
+    :param input_path: Optional input path for the task.
+    :param path_rewrite_cfg: Optional path rewrite configuration.
+    :return: The loaded task template.
+    """
+    start_time = time.time()
+    action_id = f"{org}/{project}/{domain}/{run_name}/{name}"
+
+    logger.info(
+        f"[rusty] Running task '{task.name}' (action: {action_id})"
+        f" at {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))}"
+    )
+
+    path_rewrite = PathRewrite.from_str(path_rewrite_cfg) if path_rewrite_cfg else None
+    if path_rewrite:
+        import flyte.storage as storage
+
+        if not await storage.exists(path_rewrite.new_prefix):
+            logger.error(
+                f"[rusty] Path rewrite failed for path {path_rewrite.new_prefix}, "
+                f"not found, reverting to original path {path_rewrite.old_prefix}"
+            )
+            path_rewrite = None
+        else:
+            logger.info(f"[rusty] Using path rewrite: {path_rewrite}")
+
+    try:
+        await contextual_run(
+            extract_download_run_upload,
+            task,
+            action=ActionID(name=name, org=org, project=project, domain=domain, run_name=run_name),
+            version=version,
+            controller=controller,
+            raw_data_path=RawDataPath(path=raw_data_path, path_rewrite=path_rewrite),
+            output_path=output_path,
+            run_base_dir=run_base_dir,
+            checkpoints=Checkpoints(prev_checkpoint_path=prev_checkpoint, checkpoint_path=checkpoint_path),
+            code_bundle=code_bundle,
+            input_path=input_path,
+            image_cache=ImageCache.from_transport(image_cache) if image_cache else None,
+        )
+    except asyncio.CancelledError as e:
+        logger.error(f"[rusty] Task cancellation received: {e!s}")
+        raise
+    except Exception as e:
+        logger.error(f"[rusty] Task failed: {e!s}")
+        raise
+    finally:
+        end_time = time.time()
+        duration = end_time - start_time
+        logger.info(
+            f"[rusty] TASK_EXECUTION_END: Task '{task.name}' (action: {action_id})"
+            f" done after {duration:.2f}s at {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))}"
+        )
+
+
+async def ping(name: str) -> str:
+    """
+    A simple hello world function to test the Rusty entrypoint.
+    """
+    print(f"Received ping request from {name} in Rusty!")
+    return f"pong from Rusty to {name}!"
+
+
+async def hello(name: str):
+    """
+    A simple hello world function to test the Rusty entrypoint.
+    :param name: The name to greet.
+    :return: A greeting message.
+    """
+    print(f"Received hello request in Rusty with name: {name}!")