flyte 2.0.0b32__py3-none-any.whl

flyte/_internal/runtime/entrypoints.py

@@ -0,0 +1,204 @@
+import importlib
+import os
+import traceback
+from typing import List, Optional, Tuple, Type
+
+import flyte.errors
+from flyte._code_bundle import download_bundle
+from flyte._context import contextual_run
+from flyte._internal import Controller
+from flyte._internal.imagebuild.image_builder import ImageCache
+from flyte._logging import log, logger
+from flyte._task import TaskTemplate
+from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
+
+from ..._utils import adjust_sys_path
+from .convert import Error, Inputs, Outputs
+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,
+    output_path: str,
+    run_base_dir: str,
+    checkpoints: Checkpoints | None = None,
+    code_bundle: CodeBundle | None = None,
+    inputs: Inputs | None = None,
+) -> Tuple[Optional[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 or Inputs.empty(),
+        action=action,
+        raw_data_path=raw_data_path,
+        checkpoints=checkpoints,
+        code_bundle=code_bundle,
+        controller=controller,
+        version=version,
+        output_path=output_path,
+        run_base_dir=run_base_dir,
+    )
+
+
+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()
+    try:
+        return resolver_instance.load_task(resolver_args)
+    except ModuleNotFoundError as e:
+        cwd = os.getcwd()
+        files = []
+        try:
+            for root, dirs, filenames in os.walk(cwd):
+                for name in dirs + filenames:
+                    rel_path = os.path.relpath(os.path.join(root, name), cwd)
+                    files.append(rel_path)
+        except Exception as list_err:
+            files = [f"(Failed to list directory: {list_err})"]
+
+        msg = (
+            "\n\nFull traceback:\n" + "".join(traceback.format_exc()) + f"\n[ImportError Diagnostics]\n"
+            f"Module '{e.name}' not found in either the Python virtual environment or the current working directory.\n"
+            f"Current working directory: {cwd}\n"
+            f"Files found under current directory:\n" + "\n".join(f"  - {f}" for f in files)
+        )
+        raise ModuleNotFoundError(msg) from e
+
+
+def load_pkl_task(code_bundle: CodeBundle) -> TaskTemplate:
+    """
+    Loads a task from a pickled code bundle.
+    :param code_bundle: The code bundle to load the task from.
+    :return: The loaded task template.
+    """
+    logger.debug(f"Loading task from pkl: {code_bundle.downloaded_path}")
+    try:
+        import gzip
+
+        import cloudpickle
+
+        with gzip.open(str(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
+
+
+async def download_code_bundle(code_bundle: CodeBundle) -> CodeBundle:
+    """
+    Downloads the code bundle if it is not already downloaded.
+    :param code_bundle: The code bundle to download.
+    :return: The code bundle with the downloaded path.
+    """
+    adjust_sys_path()
+    logger.debug(f"Downloading {code_bundle}")
+    downloaded_path = await download_bundle(code_bundle)
+    return code_bundle.with_downloaded_path(downloaded_path)
+
+
+async def _download_and_load_task(
+    code_bundle: CodeBundle | None, resolver: str | None = None, resolver_args: List[str] | None = None
+) -> TaskTemplate:
+    if code_bundle and (code_bundle.tgz or code_bundle.pkl):
+        logger.debug(f"Downloading {code_bundle}")
+        code_bundle = await download_code_bundle(code_bundle)
+        if code_bundle.pkl:
+            return load_pkl_task(code_bundle)
+
+        if not resolver or not resolver_args:
+            raise flyte.errors.RuntimeSystemError(
+                "MalformedCommand", "Resolver and resolver args are required. for task"
+            )
+        logger.debug(
+            f"Loading task from tgz: {code_bundle.downloaded_path}, resolver: {resolver}, args: {resolver_args}"
+        )
+        return load_task(resolver, *resolver_args)
+    if not resolver or not resolver_args:
+        raise flyte.errors.RuntimeSystemError("MalformedCommand", "Resolver and resolver args are required. for task")
+    logger.debug(f"No code bundle provided, loading task from resolver: {resolver}, args: {resolver_args}")
+    return load_task(resolver, *resolver_args)
+
+
+@log
+async def load_and_run_task(
+    action: ActionID,
+    raw_data_path: RawDataPath,
+    output_path: str,
+    run_base_dir: str,
+    version: str,
+    controller: Controller,
+    resolver: str,
+    resolver_args: List[str],
+    checkpoints: Checkpoints | None = None,
+    code_bundle: CodeBundle | None = None,
+    input_path: str | None = None,
+    image_cache: ImageCache | None = None,
+    interactive_mode: bool = False,
+):
+    """
+    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 run_base_dir: Base output directory to pass down to child tasks.
+    :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.
+    :param interactive_mode: Whether to run the task in interactive mode.
+    """
+    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,
+        run_base_dir=run_base_dir,
+        checkpoints=checkpoints,
+        code_bundle=code_bundle,
+        input_path=input_path,
+        image_cache=image_cache,
+        interactive_mode=interactive_mode,
+    )

flyte/_internal/runtime/io.py

@@ -0,0 +1,188 @@
+"""
+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
+"""
+
+from flyteidl.core import errors_pb2
+from flyteidl2.core import execution_pb2
+from flyteidl2.task import common_pb2
+
+import flyte.storage as storage
+from flyte.models import PathRewrite
+
+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, 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)
+
+
+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.
+    """
+    # 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)
+    return await storage.put_stream(data_iterable=error_document.SerializeToString(), to_path=error_uri)
+
+
+# ------------------------------- DOWNLOAD Methods ------------------------------- #
+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 = 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, 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 = 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)
+
+
+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,152 @@
+from typing import Dict, List, Optional, Tuple
+
+from flyteidl2.core import tasks_pb2
+
+from flyte._resources import CPUBaseType, DeviceClass, 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-l40s",
+    "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",
+}
+
+_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(
+        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
+
+    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=device_type,
+        partition_size=device.partition if device.partition else None,
+        device_class=device_class,
+    )
+
+
+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/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.ttl is not None:
+        components += f":{reuse_policy.ttl.total_seconds()}"
+    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.ttl.total_seconds() if reuse_policy.ttl else None,
+            "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