apps/library/2.1.49/resources.py

import re
from typing import Any, TYPE_CHECKING

if TYPE_CHECKING:
    from render import Render

try:
    from .error import RenderError
except ImportError:
    from error import RenderError

DEFAULT_CPUS = 2.0
DEFAULT_MEMORY = 4096


class Resources:
    def __init__(self, render_instance: "Render"):
        self._render_instance = render_instance
        self._limits: dict = {}
        self._reservations: dict = {}
        self._nvidia_ids: set[str] = set()
        self._auto_add_cpu_from_values()
        self._auto_add_memory_from_values()
        self._auto_add_gpus_from_values()

    def _set_cpu(self, cpus: Any):
        c = str(cpus)
        if not re.match(r"^[1-9][0-9]*(\.[0-9]+)?$", c):
            raise RenderError(f"Expected cpus to be a number or a float (minimum 1.0), got [{cpus}]")
        self._limits.update({"cpus": c})

    def _set_memory(self, memory: Any):
        m = str(memory)
        if not re.match(r"^[1-9][0-9]*$", m):
            raise RenderError(f"Expected memory to be a number, got [{memory}]")
        self._limits.update({"memory": f"{m}M"})

    def _auto_add_cpu_from_values(self):
        resources = self._render_instance.values.get("resources", {})
        self._set_cpu(resources.get("limits", {}).get("cpus", DEFAULT_CPUS))

    def _auto_add_memory_from_values(self):
        resources = self._render_instance.values.get("resources", {})
        self._set_memory(resources.get("limits", {}).get("memory", DEFAULT_MEMORY))

    def _auto_add_gpus_from_values(self):
        resources = self._render_instance.values.get("resources", {})
        gpus = resources.get("gpus", {}).get("nvidia_gpu_selection", {})
        if not gpus:
            return

        for pci, gpu in gpus.items():
            if gpu.get("use_gpu", False):
                if not gpu.get("uuid"):
                    raise RenderError(f"Expected [uuid] to be set for GPU in slot [{pci}] in [nvidia_gpu_selection]")
                self._nvidia_ids.add(gpu["uuid"])

        if self._nvidia_ids:
            if not self._reservations:
                self._reservations["devices"] = []
            self._reservations["devices"].append(
                {
                    "capabilities": ["gpu"],
                    "driver": "nvidia",
                    "device_ids": sorted(self._nvidia_ids),
                }
            )

    # This is only used on ix-app that we allow
    # disabling cpus and memory. GPUs are only added
    # if the user has requested them.
    def remove_cpus_and_memory(self):
        self._limits.pop("cpus", None)
        self._limits.pop("memory", None)

    # Mainly will be used from dependencies
    # There is no reason to pass devices to
    # redis or postgres for example
    def remove_devices(self):
        self._reservations.pop("devices", None)

    def set_profile(self, profile: str):
        cpu, memory = profile_mapping(profile)
        self._set_cpu(cpu)
        self._set_memory(memory)

    def has_resources(self):
        return len(self._limits) > 0 or len(self._reservations) > 0

    def has_gpus(self):
        gpu_devices = [d for d in self._reservations.get("devices", []) if "gpu" in d["capabilities"]]
        return len(gpu_devices) > 0

    def render(self):
        result = {}
        if self._limits:
            result["limits"] = self._limits
        if self._reservations:
            result["reservations"] = self._reservations

        return result


def profile_mapping(profile: str):
    profiles = {
        "low": (1, 512),
        "medium": (2, 1024),
    }

    if profile not in profiles:
        raise RenderError(
            f"Resource profile [{profile}] is not valid. Valid options are: [{', '.join(profiles.keys())}]"
        )

    return profiles[profile]