flyte 0.2.0b23__py3-none-any.whl → 0.2.0b25__py3-none-any.whl

flyte/_internal/runtime/reuse.py

@@ -0,0 +1,121 @@
+import hashlib
+import typing
+from venv import logger
+
+from flyteidl.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.reuse_salt is None and code_bundle is not None:
+        components += f":{code_bundle.computed_version}"
+    else:
+        components += f":{reuse_policy.reuse_salt}"
+    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
+    )
+
+    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,
+            "replica_count": reuse_policy.max_replicas,
+            "ttl_seconds": reuse_policy.ttl.total_seconds() if reuse_policy.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,165 @@
+import sys
+from typing import Any, 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.models import ActionID, Checkpoints, CodeBundle, 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}")
+    sys.path.insert(0, ".")
+
+    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}")
+    sys.path.insert(0, ".")
+
+    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}")
+    from flyte._initialize import init
+
+    # TODO Currently refrence tasks are not supported in Rusty.
+    await init.aio()
+    controller_kwargs: dict[str, Any] = {"insecure": insecure}
+    if api_key:
+        logger.info("Using api key from environment")
+        controller_kwargs["api_key"] = api_key
+    else:
+        controller_kwargs["endpoint"] = endpoint
+        if "localhost" in endpoint or "docker" in endpoint:
+            controller_kwargs["insecure"] = True
+        logger.debug(f"Using controller endpoint: {endpoint} with kwargs: {controller_kwargs}")
+
+    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,
+):
+    """
+    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.
+    :return: The loaded task template.
+    """
+    logger.info(f"[rusty] Running task {task.name}")
+    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),
+        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,
+    )
+    logger.info(f"[rusty] Finished task {task.name}")
+
+
+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}!")

flyte/_internal/runtime/task_serde.py

@@ -4,10 +4,9 @@ It includes a Resolver interface for loading tasks, and functions to load classe
 """
 
 import copy
-import importlib
 import typing
 from datetime import timedelta
-from typing import Optional, Type, cast
+from typing import Optional, cast
 
 from flyteidl.core import identifier_pb2, literals_pb2, security_pb2, tasks_pb2
 from google.protobuf import duration_pb2, wrappers_pb2
@@ -21,39 +20,17 @@ from flyte._secret import SecretRequest, secrets_from_request
 from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
 from flyte.models import CodeBundle, SerializationContext
 
+from ... import ReusePolicy
 from ..._retry import RetryStrategy
 from ..._timeout import TimeoutType, timeout_from_request
 from .resources_serde import get_proto_extended_resources, get_proto_resources
+from .reuse import add_reusable
 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,
@@ -137,23 +114,22 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
         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
+        container = None
     else:
         container = _get_urun_container(serialize_context, task)
         pod = None
 
+    custom = task.custom_config(serialize_context)
+
+    sql = None
+
     # -------------- CACHE HANDLING ----------------------
     task_cache = cache_from_request(task.cache)
     cache_enabled = task_cache.is_enabled()
@@ -175,7 +151,7 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
     else:
         logger.debug(f"Cache disabled for task {task.name}")
 
-    return tasks_pb2.TaskTemplate(
+    task_template = tasks_pb2.TaskTemplate(
         id=task_id,
         type=task.task_type,
         metadata=tasks_pb2.TaskMetadata(
@@ -195,7 +171,7 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
             generates_deck=wrappers_pb2.BoolValue(value=task.report),
         ),
         interface=transform_native_to_typed_interface(task.native_interface),
-        custom=custom,
+        custom=custom if len(custom) > 0 else None,
         container=container,
         task_type_version=task.task_type_version,
         security_context=get_security_context(task.secrets),
@@ -205,6 +181,20 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
         extended_resources=get_proto_extended_resources(task.resources),
     )
 
+    if task.reusable is not None:
+        if not isinstance(task.reusable, ReusePolicy):
+            raise flyte.errors.RuntimeUserError(
+                "BadConfig", f"Expected ReusePolicy, got {type(task.reusable)} for task {task.name}"
+            )
+        env_name = None
+        if task.parent_env is not None:
+            env = task.parent_env()
+            if env is not None:
+                env_name = env.name
+        return add_reusable(task_template, task.reusable, serialize_context.code_bundle, env_name)
+
+    return task_template
+
 
 def _get_urun_container(
     serialize_context: SerializationContext, task_template: TaskTemplate
@@ -220,9 +210,10 @@ def _get_urun_container(
     if not serialize_context.image_cache:
         # This computes the image uri, computing hashes as necessary so can fail if done remotely.
         img_uri = task_template.image.uri
-    elif image_id not in serialize_context.image_cache.image_lookup:
+    elif serialize_context.image_cache and image_id not in serialize_context.image_cache.image_lookup:
         img_uri = task_template.image.uri
         from flyte._version import __version__
+
         logger.warning(
             f"Image {task_template.image} not found in the image cache: {serialize_context.image_cache.image_lookup}.\n"
             f"This typically occurs when the Flyte SDK version (`{__version__}`) used in the task environment "

flyte/_pod.py

@@ -2,7 +2,8 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, Optional
 
 if TYPE_CHECKING:
-    from kubernetes.client import V1PodSpec
+    from flyteidl.core.tasks_pb2 import K8sPod
+    from kubernetes.client import ApiClient, V1PodSpec
 
 
 _PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
@@ -17,3 +18,12 @@ class PodTemplate(object):
     primary_container_name: str = _PRIMARY_CONTAINER_DEFAULT_NAME
     labels: Optional[Dict[str, str]] = None
     annotations: Optional[Dict[str, str]] = None
+
+    def to_k8s_pod(self) -> "K8sPod":
+        from flyteidl.core.tasks_pb2 import K8sObjectMetadata, K8sPod
+
+        return K8sPod(
+            metadata=K8sObjectMetadata(labels=self.labels, annotations=self.annotations),
+            pod_spec=ApiClient().sanitize_for_serialization(self.pod_spec),
+            primary_container_name=self.primary_container_name,
+        )

flyte/_resources.py

@@ -1,8 +1,15 @@
-from dataclasses import dataclass
-from typing import Literal, Optional, Tuple, Union, get_args
+from dataclasses import dataclass, fields
+from typing import TYPE_CHECKING, Literal, Optional, Tuple, Union, get_args
 
 import rich.repr
 
+from flyte._pod import _PRIMARY_CONTAINER_DEFAULT_NAME
+
+if TYPE_CHECKING:
+    from kubernetes.client import V1PodSpec
+
+PRIMARY_CONTAINER_DEFAULT_NAME = "primary"
+
 GPUType = Literal["T4", "A100", "A100 80G", "H100", "L4", "L40s"]
 GPUQuantity = Literal[1, 2, 3, 4, 5, 6, 7, 8]
 A100Parts = Literal["1g.5gb", "2g.10gb", "3g.20gb", "4g.20gb", "7g.40gb"]
@@ -224,3 +231,61 @@ class Resources:
         if self.shm == "auto":
             return ""
         return self.shm
+
+
+def _check_resource_is_singular(resource: Resources):
+    """
+    Raise a value error if the resource has a tuple.
+    """
+    for field in fields(resource):
+        value = getattr(resource, field.name)
+        if isinstance(value, (tuple, list)):
+            raise ValueError(f"{value} can not be a list or tuple")
+    return resource
+
+
+def pod_spec_from_resources(
+    primary_container_name: str = _PRIMARY_CONTAINER_DEFAULT_NAME,
+    requests: Optional[Resources] = None,
+    limits: Optional[Resources] = None,
+    k8s_gpu_resource_key: str = "nvidia.com/gpu",
+) -> "V1PodSpec":
+    from kubernetes.client import V1Container, V1PodSpec, V1ResourceRequirements
+
+    def _construct_k8s_pods_resources(resources: Optional[Resources], k8s_gpu_resource_key: str):
+        if resources is None:
+            return None
+
+        resources_map = {
+            "cpu": "cpu",
+            "memory": "memory",
+            "gpu": k8s_gpu_resource_key,
+            "ephemeral_storage": "ephemeral-storage",
+        }
+
+        k8s_pod_resources = {}
+
+        _check_resource_is_singular(resources)
+        for resource in fields(resources):
+            resource_value = getattr(resources, resource.name)
+            if resource_value is not None:
+                k8s_pod_resources[resources_map[resource.name]] = resource_value
+
+        return k8s_pod_resources
+
+    requests = _construct_k8s_pods_resources(resources=requests, k8s_gpu_resource_key=k8s_gpu_resource_key)
+    limits = _construct_k8s_pods_resources(resources=limits, k8s_gpu_resource_key=k8s_gpu_resource_key)
+    requests = requests or limits
+    limits = limits or requests
+
+    return V1PodSpec(
+        containers=[
+            V1Container(
+                name=primary_container_name,
+                resources=V1ResourceRequirements(
+                    requests=requests,
+                    limits=limits,
+                ),
+            )
+        ]
+    )

flyte/_reusable_environment.py

@@ -2,6 +2,8 @@ from dataclasses import dataclass
 from datetime import timedelta
 from typing import Optional, Tuple, Union
 
+from flyte._logging import logger
+
 
 @dataclass
 class ReusePolicy:
@@ -14,12 +16,65 @@ class ReusePolicy:
     Caution: It is important to note that the environment is shared, so managing memory and resources is important.
 
     :param replicas: Either a single int representing number of replicas or a tuple of two ints representing
-     the min and max
+     the min and max.
     :param idle_ttl: The maximum idle duration for an environment replica, specified as either seconds (int) or a
         timedelta. If not set, the environment's global default will be used.
         When a replica remains idle — meaning no tasks are running — for this duration, it will be automatically
         terminated.
+    :param concurrency: The maximum number of tasks that can run concurrently in one instance of the environment.
+          Concurrency of greater than 1 is only supported only for `async` tasks.
+    :param reuse_salt: Optional string used to control environment reuse.
+        If set, the environment will be reused even if the code bundle changes.
+        To force a new environment, either set this to `None` or change its value.
+
+        Example:
+            reuse_salt = "v1"  # Environment is reused
+            reuse_salt = "v2"  # Forces environment recreation
     """
 
-    replicas: Union[int, Tuple[int, int]] = 1
+    replicas: Union[int, Tuple[int, int]] = 2
     idle_ttl: Optional[Union[int, timedelta]] = None
+    reuse_salt: str | None = None
+    concurrency: int = 1
+
+    def __post_init__(self):
+        if self.replicas is None:
+            raise ValueError("replicas cannot be None")
+        if isinstance(self.replicas, int):
+            self.replicas = (self.replicas, self.replicas)
+        elif not isinstance(self.replicas, tuple):
+            raise ValueError("replicas must be an int or a tuple of two ints")
+        elif len(self.replicas) != 2:
+            raise ValueError("replicas must be an int or a tuple of two ints")
+
+        if self.idle_ttl:
+            if isinstance(self.idle_ttl, int):
+                self.idle_ttl = timedelta(seconds=int(self.idle_ttl))
+            elif not isinstance(self.idle_ttl, timedelta):
+                raise ValueError("idle_ttl must be an int (seconds) or a timedelta")
+
+        if self.replicas[1] == 1 and self.concurrency == 1:
+            logger.warning(
+                "It is recommended to use a minimum of 2 replicas, to avoid starvation. "
+                "Starvation can occur if a task is running and no other replicas are available to handle new tasks."
+                "Options, increase concurrency, increase replicas or turn-off reuse for the parent task, "
+                "that runs child tasks."
+            )
+
+    @property
+    def ttl(self) -> timedelta | None:
+        """
+        Returns the idle TTL as a timedelta. If idle_ttl is not set, returns the global default.
+        """
+        if self.idle_ttl is None:
+            return None
+        if isinstance(self.idle_ttl, timedelta):
+            return self.idle_ttl
+        return timedelta(seconds=self.idle_ttl)
+
+    @property
+    def max_replicas(self) -> int:
+        """
+        Returns the maximum number of replicas.
+        """
+        return self.replicas[1] if isinstance(self.replicas, tuple) else self.replicas