flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_internal/runtime/task_serde.py

@@ -3,66 +3,68 @@ This module provides functionality to serialize and deserialize tasks to and fro
 It includes a Resolver interface for loading tasks, and functions to load classes and tasks.
 """
 
-import importlib
+import copy
+import typing
 from datetime import timedelta
-from typing import Optional, Type
+from typing import Optional, cast
 
-from flyteidl.core import identifier_pb2, literals_pb2, security_pb2, tasks_pb2
+from flyteidl2.core import identifier_pb2, literals_pb2, security_pb2, tasks_pb2
+from flyteidl2.core.execution_pb2 import TaskLog
+from flyteidl2.task import common_pb2, environment_pb2, task_definition_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._pod import _PRIMARY_CONTAINER_NAME_FIELD, PodTemplate
 from flyte._secret import SecretRequest, secrets_from_request
 from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
+from flyte.models import CodeBundle, SerializationContext
 
+from ... import ReusePolicy
+from ..._context import internal_ctx
 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
 
-
-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)
+_MAX_ENV_NAME_LENGTH = 63  # Maximum length for environment names
+_MAX_TASK_SHORT_NAME_LENGTH = 63  # Maximum length for task short names
 
 
 def translate_task_to_wire(
-    task: TaskTemplate, serialization_context: SerializationContext
+    task: TaskTemplate,
+    serialization_context: SerializationContext,
+    default_inputs: Optional[typing.List[common_pb2.NamedParameter]] = None,
 ) -> task_definition_pb2.TaskSpec:
     """
     Translate a task to a wire format. This is a placeholder function.
 
     :param task: The task to translate.
     :param serialization_context: The serialization context to use for the translation.
+    :param default_inputs: Optional list of default inputs for the task.
 
     :return: The translated task.
     """
-    return get_proto_task(task, serialization_context)
+    tt = get_proto_task(task, serialization_context)
+    env: environment_pb2.Environment | None = None
+
+    if task.parent_env and task.parent_env():
+        _env = task.parent_env()
+        if _env:
+            env = environment_pb2.Environment(name=_env.name[:_MAX_ENV_NAME_LENGTH])
+    return task_definition_pb2.TaskSpec(
+        task_template=tt,
+        default_inputs=default_inputs,
+        short_name=task.short_name[:_MAX_TASK_SHORT_NAME_LENGTH],
+        environment=env,
+    )
 
 
-def get_security_context(secrets: Optional[SecretRequest]) -> Optional[security_pb2.SecurityContext]:
+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.
 
@@ -89,7 +91,9 @@ def get_security_context(secrets: Optional[SecretRequest]) -> Optional[security_
     )
 
 
-def get_proto_retry_strategy(retries: RetryStrategy | int | None) -> Optional[literals_pb2.RetryStrategy]:
+def get_proto_retry_strategy(
+    retries: RetryStrategy | int | None,
+) -> Optional[literals_pb2.RetryStrategy]:
     if retries is None:
         return None
 
@@ -108,7 +112,7 @@ def get_proto_timeout(timeout: TimeoutType | None) -> Optional[duration_pb2.Dura
     return duration_pb2.Duration(seconds=max_runtime_timeout.seconds)
 
 
-def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext) -> task_definition_pb2.TaskSpec:
+def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext) -> tasks_pb2.TaskTemplate:
     task_id = identifier_pb2.Identifier(
         resource_type=identifier_pb2.ResourceType.TASK,
         project=serialize_context.project,
@@ -117,21 +121,38 @@ 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")
-
-    #
-    # 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
+
+    extra_config: typing.Dict[str, str] = {}
+
+    if task.pod_template and not isinstance(task.pod_template, str):
+        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
+
+    ctx = internal_ctx()
+    task_ctx = ctx.data.task_context
+    log_links = []
+    if task.links and task_ctx:
+        action = task_ctx.action
+        for link in task.links:
+            uri = link.get_link(
+                run_name=action.run_name if action.run_name else "",
+                project=action.project if action.project else "",
+                domain=action.domain if action.domain else "",
+                context=task_ctx.custom_context if task_ctx.custom_context else {},
+                parent_action_name=action.name if action.name else "",
+                action_name="{{.actionName}}",
+                pod_name="{{.podName}}",
+            )
+            task_log = TaskLog(name=link.name, uri=uri, icon_uri=link.icon_uri)
+            log_links.append(task_log)
+
+    custom = task.custom_config(serialize_context)
+
+    sql = task.sql(serialize_context)
 
     # -------------- CACHE HANDLING ----------------------
     task_cache = cache_from_request(task.cache)
@@ -144,7 +165,6 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
             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:
@@ -154,23 +174,29 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
     else:
         logger.debug(f"Cache disabled for task {task.name}")
 
-    tt = tasks_pb2.TaskTemplate(
+    task_template = 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(),
+            cache_ignore_input_vars=(task_cache.get_ignored_inputs() if cache_enabled else None),
+            runtime=tasks_pb2.RuntimeMetadata(
+                version=flyte.version(),
+                type=tasks_pb2.RuntimeMetadata.RuntimeType.FLYTE_SDK,
+                flavor="python",
+            ),
             retries=get_proto_retry_strategy(task.retries),
             timeout=get_proto_timeout(task.timeout),
-            pod_template_name=task.pod_template if task.pod_template and isinstance(task.pod_template, str) else None,
-            interruptible=task.interruptable,
-            generates_deck=wrappers_pb2.BoolValue(value=False),  # TODO add support for reports
+            pod_template_name=(task.pod_template if task.pod_template and isinstance(task.pod_template, str) else None),
+            interruptible=task.interruptible,
+            generates_deck=wrappers_pb2.BoolValue(value=task.report),
+            debuggable=task.debuggable,
+            log_links=log_links,
         ),
         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),
@@ -179,25 +205,69 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
         sql=sql,
         extended_resources=get_proto_extended_resources(task.resources),
     )
-    return task_definition_pb2.TaskSpec(task_template=tt)
+
+    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 lookup_image_in_cache(serialize_context: SerializationContext, env_name: str, image: flyte.Image) -> str:
+    if not serialize_context.image_cache or len(image._layers) == 0:
+        # This computes the image uri, computing hashes as necessary so can fail if done remotely.
+        return image.uri
+    elif serialize_context.image_cache and env_name not in serialize_context.image_cache.image_lookup:
+        raise flyte.errors.RuntimeUserError(
+            "MissingEnvironment",
+            f"Environment '{env_name}' not found in image cache.\n\n"
+            "💡 To fix this:\n"
+            "  1. If your parent environment calls a task in another environment,"
+            " declare that dependency using 'depends_on=[...]'.\n"
+            "     Example:\n"
+            "         env1 = flyte.TaskEnvironment(\n"
+            "             name='outer',\n"
+            "             image=flyte.Image.from_debian_base().with_pip_packages('requests'),\n"
+            "             depends_on=[env2, env3],\n"
+            "         )\n"
+            "  2. If you're using os.getenv() to set the environment name,"
+            " make sure the runtime environment has the same environment variable defined.\n"
+            "     Example:\n"
+            "         env = flyte.TaskEnvironment(\n"
+            '             name=os.getenv("my-name"),\n'
+            '             env_vars={"my-name": os.getenv("my-name")},\n'
+            "         )\n",
+        )
+    return serialize_context.image_cache.image_lookup[env_name]
 
 
 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
+        [literals_pb2.KeyValuePair(key=k, value=v) for k, v in task_template.env_vars.items()]
+        if task_template.env_vars
+        else None
     )
     resources = get_proto_resources(task_template.resources)
-    # pr: under what conditions should this return None?
-    if isinstance(task_template.image, str):
+
+    img = task_template.image
+    if isinstance(img, 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]
+
+    env_name = task_template.parent_env_name
+    if env_name is None:
+        raise flyte.errors.RuntimeSystemError("BadConfig", f"Task {task_template.name} has no parent environment name")
+
+    img_uri = lookup_image_in_cache(serialize_context, env_name, img)
 
     return tasks_pb2.Container(
         image=img_uri,
@@ -208,3 +278,106 @@ def _get_urun_container(
         data_config=task_template.data_loading_config(serialize_context),
         config=task_template.config(serialize_context),
     )
+
+
+def _sanitize_resource_name(resource: tasks_pb2.Resources.ResourceEntry) -> str:
+    return tasks_pb2.Resources.ResourceName.Name(resource.name).lower().replace("_", "-")
+
+
+def _get_k8s_pod(primary_container: tasks_pb2.Container, pod_template: PodTemplate) -> Optional[tasks_pb2.K8sPod]:
+    """
+    Get the K8sPod representation of the task template.
+    :param task: The task to convert.
+    :return: The K8sPod representation of the task template.
+    """
+    from kubernetes.client import ApiClient, V1PodSpec
+    from kubernetes.client.models import V1EnvVar, V1ResourceRequirements
+
+    pod_template = copy.deepcopy(pod_template)
+    containers = cast(V1PodSpec, pod_template.pod_spec).containers
+    primary_exists = False
+
+    for container in containers:
+        if container.name == pod_template.primary_container_name:
+            primary_exists = True
+            break
+
+    if not primary_exists:
+        raise ValueError(
+            "No primary container defined in the pod spec."
+            f" You must define a primary container with the name '{pod_template.primary_container_name}'."
+        )
+    final_containers = []
+
+    for container in containers:
+        # We overwrite the primary container attributes with the values given to ContainerTask.
+        # The attributes include: image, command, args, resource, and env (env is unioned)
+
+        if container.name == pod_template.primary_container_name:
+            if container.image is None:
+                # Copy the image from primary_container only if the image is not specified in the pod spec.
+                container.image = primary_container.image
+
+            container.command = list(primary_container.command)
+            container.args = list(primary_container.args)
+
+            limits, requests = {}, {}
+            for resource in primary_container.resources.limits:
+                limits[_sanitize_resource_name(resource)] = resource.value
+            for resource in primary_container.resources.requests:
+                requests[_sanitize_resource_name(resource)] = resource.value
+
+            resource_requirements = V1ResourceRequirements(limits=limits, requests=requests)
+            if len(limits) > 0 or len(requests) > 0:
+                # Important! Only copy over resource requirements if they are non-empty.
+                container.resources = resource_requirements
+
+            if primary_container.env is not None:
+                container.env = [V1EnvVar(name=e.key, value=e.value) for e in primary_container.env] + (
+                    container.env or []
+                )
+
+        final_containers.append(container)
+
+    cast(V1PodSpec, pod_template.pod_spec).containers = final_containers
+    pod_spec = ApiClient().sanitize_for_serialization(pod_template.pod_spec)
+
+    metadata = tasks_pb2.K8sObjectMetadata(labels=pod_template.labels, annotations=pod_template.annotations)
+    return tasks_pb2.K8sPod(pod_spec=pod_spec, metadata=metadata)
+
+
+def extract_code_bundle(
+    task_spec: task_definition_pb2.TaskSpec,
+) -> Optional[CodeBundle]:
+    """
+    Extract the code bundle from the task spec.
+    :param task_spec: The task spec to extract the code bundle from.
+    :return: The extracted code bundle or None if not present.
+    """
+    container = task_spec.task_template.container
+    if container and container.args:
+        pkl_path = None
+        tgz_path = None
+        dest_path: str = "."
+        version = ""
+        for i, v in enumerate(container.args):
+            if v == "--pkl":
+                # Extract the code bundle path from the argument
+                pkl_path = container.args[i + 1] if i + 1 < len(container.args) else None
+            elif v == "--tgz":
+                # Extract the code bundle path from the argument
+                tgz_path = container.args[i + 1] if i + 1 < len(container.args) else None
+            elif v == "--dest":
+                # Extract the destination path from the argument
+                dest_path = container.args[i + 1] if i + 1 < len(container.args) else "."
+            elif v == "--version":
+                # Extract the version from the argument
+                version = container.args[i + 1] if i + 1 < len(container.args) else ""
+        if pkl_path or tgz_path:
+            return CodeBundle(
+                destination=dest_path,
+                tgz=tgz_path,
+                pkl=pkl_path,
+                computed_version=version,
+            )
+    return None

flyte/_internal/runtime/taskrunner.py

@@ -4,15 +4,17 @@ invoked within a context tree.
 """
 
 import pathlib
+import time
 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._metrics import Stopwatch
 from flyte._task import TaskTemplate
 from flyte.errors import CustomError, RuntimeSystemError, RuntimeUnknownError, RuntimeUserError
+from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath, TaskContext
 
 from .. import Controller
 from .convert import (
@@ -110,41 +112,75 @@ async def run_task(
 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,
+    inputs: Inputs = Inputs.empty(),
+    input_path: str | None = None,
     checkpoints: Checkpoints | None = None,
     code_bundle: CodeBundle | None = None,
     image_cache: ImageCache | None = None,
+    interactive_mode: bool = False,
 ) -> 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()
+
+    # Load inputs first to get context
+    if input_path:
+        sw = Stopwatch("load_inputs")
+        sw.start()
+        inputs = await load_inputs(input_path, path_rewrite_config=raw_data_path.path_rewrite)
+        sw.stop()
+
+    # Extract context from inputs
+    custom_context = inputs.context if inputs else {}
+
     tctx = TaskContext(
         action=action,
         checkpoints=checkpoints,
         code_bundle=code_bundle,
+        input_path=input_path,
         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),
+        mode="remote" if not ctx.data.task_context else ctx.data.task_context.mode,
+        interactive_mode=interactive_mode,
+        custom_context=custom_context,
     )
-    async with ctx.replace_task_context(tctx):
+
+    with ctx.replace_task_context(tctx):
+        sw = Stopwatch("convert_inputs_to_native")
+        sw.start()
         inputs_kwargs = await convert_inputs_to_native(inputs, task.native_interface)
+        sw.stop()
+
+        sw = Stopwatch("run_task")
+        sw.start()
         out, err = await run_task(tctx=tctx, controller=controller, task=task, inputs=inputs_kwargs)
+        sw.stop()
+
         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
+            # Check if report has content before flushing to avoid overwriting
+            # worker reports (from Elastic/distributed tasks) with empty main process report
+            if ctx.get_report():
+                await flyte.report.flush.aio()
+
+        sw = Stopwatch("convert_outputs_from_native")
+        sw.start()
+        result = await convert_from_native_to_outputs(out, task.native_interface, task.name), None
+        sw.stop()
+        return result
 
 
 async def extract_download_run_upload(
@@ -160,15 +196,17 @@ async def extract_download_run_upload(
     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 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
+    t = time.time()
+    logger.info(f"Task {action.name} started at {t}")
     outputs, err = await convert_and_run(
         task=task,
-        inputs=inputs or Inputs.empty(),
+        input_path=input_path,
         action=action,
         controller=controller,
         raw_data_path=raw_data_path,
@@ -178,7 +216,9 @@ async def extract_download_run_upload(
         checkpoints=checkpoints,
         code_bundle=code_bundle,
         image_cache=image_cache,
+        interactive_mode=interactive_mode,
     )
+    logger.debug(f"Task {action.name} completed at {t}, with outputs: {outputs}")
     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}")
@@ -187,4 +227,4 @@ async def extract_download_run_upload(
         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}")
+    logger.info(f"Task {task.name} completed successfully, uploaded outputs to {output_path} in {time.time() - t}s")