app_serde.py

"""
Serialization module for AppEnvironment to AppIDL conversion.

This module provides functionality to serialize an AppEnvironment object into
the AppIDL protobuf format, using SerializationContext for configuration.
"""

from __future__ import annotations

from copy import deepcopy
from dataclasses import replace
from typing import Dict, List, Optional, Union

from flyteidl2.app import app_definition_pb2
from flyteidl2.common import runtime_version_pb2
from flyteidl2.core import literals_pb2, tasks_pb2
from google.protobuf.duration_pb2 import Duration

import flyte
import flyte.io
from flyte._internal.runtime.resources_serde import get_proto_extended_resources, get_proto_resources
from flyte._internal.runtime.task_serde import get_security_context, lookup_image_in_cache
from flyte._logging import logger
from flyte.app import AppEnvironment, Input, Scaling
from flyte.app._input import _DelayedValue
from flyte.models import SerializationContext
from flyte.syncify import syncify


def get_proto_container(
    app_env: AppEnvironment,
    serialization_context: SerializationContext,
    input_overrides: list[Input] | None = None,
) -> tasks_pb2.Container:
    """
    Construct the container specification.

    Args:
        app_env: The app environment
        serialization_context: Serialization context
        input_overrides: Input overrides to apply to the app environment.
    Returns:
        Container protobuf message
    """
    from flyte import Image

    env = [literals_pb2.KeyValuePair(key=k, value=v) for k, v in app_env.env_vars.items()] if app_env.env_vars else None
    resources = get_proto_resources(app_env.resources)

    if app_env.image == "auto":
        img = Image.from_debian_base()
    elif isinstance(app_env.image, str):
        img = Image.from_base(app_env.image)
    else:
        img = app_env.image

    env_name = app_env.name
    img_uri = lookup_image_in_cache(serialization_context, env_name, img)

    p = app_env.get_port()
    container_ports = [tasks_pb2.ContainerPort(container_port=p.port, name=p.name)]

    return tasks_pb2.Container(
        image=img_uri,
        command=app_env.container_cmd(serialization_context, input_overrides),
        args=app_env.container_args(serialization_context),
        resources=resources,
        ports=container_ports,
        env=env,
    )


def _sanitize_resource_name(resource: tasks_pb2.Resources.ResourceEntry) -> str:
    """
    Sanitize resource name for Kubernetes compatibility.

    Args:
        resource: Resource entry

    Returns:
        Sanitized resource name
    """
    return tasks_pb2.Resources.ResourceName.Name(resource.name).lower().replace("_", "-")


def _get_mig_resources_from_extended_resources(
    extended_resources: Optional[tasks_pb2.ExtendedResources],
    device_quantity: Optional[int] = None,
    mig_resource_prefix: Optional[str] = None,
) -> Dict[str, str]:
    """
    Generate MIG-specific resources from GPUAccelerator partition info.

    When a GPU has a partition_size specified, generate a custom resource name
    for that partition instead of using the generic GPU resource.

    :param extended_resources: The extended resources containing GPU accelerator info
    :param mig_resource_prefix: Custom MIG resource prefix (defaults to "nvidia.com/mig")
    :param device_quantity: The quantity of GPUs/partitions requested
    :return: Dict mapping MIG resource name to quantity (e.g., {"nvidia.com/mig-1g. 5gb": "1"})
    """
    mig_resources: Dict[str, str] = {}

    if extended_resources is None or not extended_resources.gpu_accelerator:
        return mig_resources

    gpu_accel = extended_resources.gpu_accelerator
    partition = gpu_accel.partition_size

    if not partition:
        return mig_resources

    # Default to "nvidia.com/mig" if not specified
    prefix = mig_resource_prefix if mig_resource_prefix is not None else "nvidia.com/mig"
    resource_name = f"{prefix}-{partition}"

    quantity = device_quantity if device_quantity is not None else 1
    mig_resources[resource_name] = str(quantity)

    return mig_resources


def _serialized_pod_spec(
    app_env: AppEnvironment,
    pod_template: flyte.PodTemplate,
    serialization_context: SerializationContext,
) -> dict:
    """
    Convert pod spec into a dict for serialization.

    Args:
        app_env: The app environment
        pod_template: Pod template specification
        serialization_context: Serialization context

    Returns:
        Dictionary representation of the pod spec
    """
    from kubernetes.client import ApiClient
    from kubernetes.client.models import V1Container, V1ContainerPort, V1EnvVar, V1ResourceRequirements

    pod_template = deepcopy(pod_template)

    if pod_template.pod_spec is None:
        return {}

    if pod_template.primary_container_name != "app":
        msg = "Primary container name must be 'app'"
        raise ValueError(msg)

    containers: list[V1Container] = pod_template.pod_spec.containers
    primary_exists = any(container.name == pod_template.primary_container_name for container in containers)

    if not primary_exists:
        msg = "Primary container does not exist with name 'app'"
        raise ValueError(msg)

    final_containers = []

    # Process containers
    for container in containers:
        img = container.image
        if isinstance(img, flyte.Image):
            img = lookup_image_in_cache(serialization_context, container.name, img)
        container.image = img

        if container.name == pod_template.primary_container_name:
            container.args = app_env.container_args(serialization_context)
            container.command = app_env.container_cmd(serialization_context)

            limits, requests = {}, {}
            resources = get_proto_resources(app_env.resources)
            extended_resources = get_proto_extended_resources(app_env.resources)
            if resources:
                for resource in resources.limits:
                    limits[_sanitize_resource_name(resource)] = resource.value
                for resource in resources.requests:
                    requests[_sanitize_resource_name(resource)] = resource.value

                # Add MIG resources if GPU partitions are specified
                mig_prefix = app_env.resources.gpu_partition_resource_prefix if app_env.resources else None
                # Get device quantity from resources
                device = app_env.resources.get_device() if app_env.resources else None
                device_quantity = device.quantity if device else None
                mig_resources = _get_mig_resources_from_extended_resources(
                    extended_resources, device_quantity, mig_prefix
                )
                requests.update(mig_resources)
                limits.update(mig_resources)

                resource_requirements = V1ResourceRequirements(limits=limits, requests=requests)

                if limits or requests or mig_resources:
                    container.resources = resource_requirements

            if app_env.env_vars:
                container.env = [V1EnvVar(name=k, value=v) for k, v in app_env.env_vars.items()] + (container.env or [])

            _port = app_env.get_port()
            container.ports = [V1ContainerPort(container_port=_port.port, name=_port.name)] + (container.ports or [])

        final_containers.append(container)

    pod_template.pod_spec.containers = final_containers
    return ApiClient().sanitize_for_serialization(pod_template.pod_spec)


def _get_k8s_pod(
    app_env: AppEnvironment,
    pod_template: flyte.PodTemplate,
    serialization_context: SerializationContext,
) -> tasks_pb2.K8sPod:
    """
    Convert pod_template into a K8sPod IDL.

    Args:
        app_env: The app environment
        pod_template: Pod template specification
        serialization_context: Serialization context

    Returns:
        K8sPod protobuf message
    """
    import json

    from google.protobuf.json_format import Parse
    from google.protobuf.struct_pb2 import Struct

    pod_spec_dict = _serialized_pod_spec(app_env, pod_template, serialization_context)
    pod_spec_idl = Parse(json.dumps(pod_spec_dict), Struct())

    metadata = tasks_pb2.K8sObjectMetadata(
        labels=pod_template.labels,
        annotations=pod_template.annotations,
    )
    return tasks_pb2.K8sPod(pod_spec=pod_spec_idl, metadata=metadata)


def _get_scaling_metric(
    metric: Optional[Union[Scaling.Concurrency, Scaling.RequestRate]],
) -> Optional[app_definition_pb2.ScalingMetric]:
    """
    Convert scaling metric to protobuf format.

    Args:
        metric: Scaling metric (Concurrency or RequestRate)

    Returns:
        ScalingMetric protobuf message or None
    """

    if metric is None:
        return None

    if isinstance(metric, Scaling.Concurrency):
        return app_definition_pb2.ScalingMetric(concurrency=app_definition_pb2.Concurrency(val=metric.val))
    elif isinstance(metric, Scaling.RequestRate):
        return app_definition_pb2.ScalingMetric(request_rate=app_definition_pb2.RequestRate(val=metric.val))

    return None


async def _materialize_inputs_with_delayed_values(inputs: List[Input]) -> List[Input]:
    """
    Materialize the inputs that contain delayed values. This is important for both
    serializing the input for the container command and for the app idl inputs collection.

    Args:
        inputs: The inputs to materialize.

    Returns:
        The materialized inputs.
    """
    _inputs = []
    for input in inputs:
        if isinstance(input.value, _DelayedValue):
            logger.info(f"Materializing {input.name} with delayed values of type {input.value.type}")
            value = await input.value.get()
            assert isinstance(value, (str, flyte.io.File, flyte.io.Dir)), (
                f"Materialized value must be a string, file or directory, found {type(value)}"
            )
            _inputs.append(replace(input, value=await input.value.get()))
        else:
            _inputs.append(input)
    return _inputs


async def translate_inputs(inputs: List[Input]) -> app_definition_pb2.InputList:
    """
    Placeholder for translating inputs to protobuf format.

    Returns:
        InputList protobuf message
    """
    if not inputs:
        return app_definition_pb2.InputList()

    inputs_list = []
    for input in inputs:
        if isinstance(input.value, str):
            inputs_list.append(app_definition_pb2.Input(name=input.name, string_value=input.value))
        elif isinstance(input.value, flyte.io.File):
            inputs_list.append(app_definition_pb2.Input(name=input.name, string_value=str(input.value.path)))
        elif isinstance(input.value, flyte.io.Dir):
            inputs_list.append(app_definition_pb2.Input(name=input.name, string_value=str(input.value.path)))
        else:
            raise ValueError(f"Unsupported input value type: {type(input.value)}")
    return app_definition_pb2.InputList(items=inputs_list)


@syncify
async def translate_app_env_to_idl(
    app_env: AppEnvironment,
    serialization_context: SerializationContext,
    input_overrides: list[Input] | None = None,
    desired_state: app_definition_pb2.Spec.DesiredState = app_definition_pb2.Spec.DesiredState.DESIRED_STATE_ACTIVE,
) -> app_definition_pb2.App:
    """
    Translate an AppEnvironment to AppIDL protobuf format.

    This is the main entry point for serializing an AppEnvironment object into
    the AppIDL protobuf format.

    Args:
        app_env: The app environment to serialize
        serialization_context: Serialization context containing org, project, domain, version, etc.
        input_overrides: Input overrides to apply to the app environment.
        desired_state: Desired state of the app (ACTIVE, INACTIVE, etc.)

    Returns:
        AppIDL protobuf message
    """
    # Build security context
    task_sec_ctx = get_security_context(app_env.secrets)
    allow_anonymous = False
    if not app_env.requires_auth:
        allow_anonymous = True

    security_context = None
    if task_sec_ctx or allow_anonymous:
        security_context = app_definition_pb2.SecurityContext(
            run_as=task_sec_ctx.run_as if task_sec_ctx else None,
            secrets=task_sec_ctx.secrets if task_sec_ctx else [],
            allow_anonymous=allow_anonymous,
        )

    # Build autoscaling config
    scaling_metric = _get_scaling_metric(app_env.scaling.metric)

    dur = None
    if app_env.scaling.scaledown_after:
        dur = Duration()
        dur.FromTimedelta(app_env.scaling.scaledown_after)

    min_replicas, max_replicas = app_env.scaling.get_replicas()
    autoscaling = app_definition_pb2.AutoscalingConfig(
        replicas=app_definition_pb2.Replicas(min=min_replicas, max=max_replicas),
        scaledown_period=dur,
        scaling_metric=scaling_metric,
    )

    # Build spec based on image type
    inputs = await _materialize_inputs_with_delayed_values(input_overrides or app_env.inputs)
    container = None
    pod = None
    if app_env.pod_template:
        if isinstance(app_env.pod_template, str):
            raise NotImplementedError("PodTemplate as str is not supported yet")
        pod = _get_k8s_pod(
            app_env,
            app_env.pod_template,
            serialization_context,
        )
    elif app_env.image:
        container = get_proto_container(
            app_env,
            serialization_context,
            input_overrides=inputs,
        )
    else:
        msg = "image must be a str, Image, or PodTemplate"
        raise ValueError(msg)

    ingress = app_definition_pb2.IngressConfig(
        private=False,
        subdomain=app_env.domain.subdomain if app_env.domain else None,
        cname=app_env.domain.custom_domain if app_env.domain else None,
    )

    # Build links
    links = None
    if app_env.links:
        links = [
            app_definition_pb2.Link(path=link.path, title=link.title, is_relative=link.is_relative)
            for link in app_env.links
        ]

    # Build profile
    profile = app_definition_pb2.Profile(
        type=app_env.type,
        short_description=app_env.description,
    )

    # Build the full App IDL
    return app_definition_pb2.App(
        metadata=app_definition_pb2.Meta(
            id=app_definition_pb2.Identifier(
                org=serialization_context.org,
                project=serialization_context.project,
                domain=serialization_context.domain,
                name=app_env.name,
            ),
        ),
        spec=app_definition_pb2.Spec(
            desired_state=desired_state,
            ingress=ingress,
            autoscaling=autoscaling,
            security_context=security_context,
            cluster_pool=app_env.cluster_pool,
            extended_resources=get_proto_extended_resources(app_env.resources),
            runtime_metadata=runtime_version_pb2.RuntimeMetadata(
                type=runtime_version_pb2.RuntimeMetadata.RuntimeType.FLYTE_SDK,
                version=flyte.version(),
                flavor="python",
            ),
            profile=profile,
            links=links,
            container=container,
            pod=pod,
            inputs=await translate_inputs(inputs),
        ),
    )