intrinsic.py

import logging
import os
import platform
from collections.abc import Collection, Mapping, Sequence
from decimal import Decimal
from pathlib import Path
from typing import Any

import aiohttp
from aiodocker.docker import Docker, DockerContainer
from aiodocker.exceptions import DockerError
from kubernetes_asyncio import client as K8sClient
from kubernetes_asyncio import config as K8sConfig

from ai.backend.agent import __version__  # pants: no-infer-dep
from ai.backend.agent.alloc_map import AllocationStrategy
from ai.backend.agent.errors import InvalidAllocMapTypeError, InvalidOvercommitFactorError
from ai.backend.agent.resources import (
    AbstractAllocMap,
    AbstractComputeDevice,
    AbstractComputePlugin,
    DeviceSlotInfo,
    DiscretePropertyAllocMap,
)
from ai.backend.agent.stats import (
    ContainerMeasurement,
    NodeMeasurement,
    ProcessMeasurement,
    StatContext,
)
from ai.backend.agent.types import Container, MountInfo
from ai.backend.common.types import (
    AcceleratorMetadata,
    ContainerId,
    DeviceId,
    DeviceModelInfo,
    DeviceName,
    SlotName,
    SlotTypes,
)
from ai.backend.logging import BraceStyleAdapter

from .resources import get_resource_spec_from_container

log = BraceStyleAdapter(logging.getLogger(__spec__.name))


async def fetch_api_stats(container: DockerContainer) -> dict[str, Any] | None:
    short_cid = ContainerId(container.id[:7])
    try:
        # aiodocker may return list[dict] or dict depending on version
        ret: list[dict[str, Any]] | dict[str, Any] = await container.stats(stream=False)
    except RuntimeError as e:
        msg = str(e.args[0]).lower()
        if "event loop is closed" in msg or "session is closed" in msg:
            return None
        raise
    except (DockerError, aiohttp.ClientError) as e:
        log.error(
            "cannot read stats (cid:{}): client error: {!r}.",
            short_cid,
            e,
        )
        return None
    else:
        entry = {"read": "0001-01-01"}
        # aiodocker 0.16 or later returns a list of dict, even when not streaming.
        match ret:
            case list() if ret:
                entry = ret[0]
            case dict() if ret:
                entry = ret
            case _:
                # The API may return an empty result upon container termination.
                log.warning(
                    "cannot read stats (cid:{}): got an empty result: {}",
                    short_cid,
                    ret,
                )
                return None
        if entry["read"].startswith("0001-01-01") or entry["preread"].startswith("0001-01-01"):
            return None
        return entry


# Pseudo-plugins for intrinsic devices (CPU and the main memory)


class CPUDevice(AbstractComputeDevice):
    pass


class CPUPlugin(AbstractComputePlugin):
    """
    Represents the CPU.
    """

    config_watch_enabled = False

    key = DeviceName("cpu")
    slot_types = [
        (SlotName("cpu"), SlotTypes.COUNT),
    ]

    async def init(self, context: Any | None = None) -> None:
        pass

    async def cleanup(self) -> None:
        pass

    async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None:
        pass

    async def list_devices(self) -> Collection[CPUDevice]:
        await K8sConfig.load_kube_config()
        core_api = K8sClient.CoreV1Api()

        nodes = (await core_api.list_node()).to_dict()["items"]
        overcommit_factor = int(os.environ.get("BACKEND_CPU_OVERCOMMIT_FACTOR", "1"))
        if not (1 <= overcommit_factor <= 10):
            raise InvalidOvercommitFactorError(
                f"CPU overcommit factor must be between 1 and 10, got {overcommit_factor}."
            )

        return [
            CPUDevice(
                device_id=DeviceId(node["metadata"]["uid"]),
                hw_location="root",
                numa_node=None,
                memory_size=0,
                processing_units=int(node["status"]["capacity"]["cpu"]) * overcommit_factor,
            )
            for i, node in enumerate(nodes)
            # if 'node-role.kubernetes.io/master' not in node['metadata']['labels'].keys()
        ]

    async def available_slots(self) -> Mapping[SlotName, Decimal]:
        devices = await self.list_devices()
        log.debug("available_slots: {}", devices)
        return {
            SlotName("cpu"): Decimal(sum(dev.processing_units for dev in devices)),
        }

    def get_version(self) -> str:
        return __version__

    async def extra_info(self) -> Mapping[str, str]:
        return {
            "agent_version": __version__,
            "machine": platform.machine(),
            "os_type": platform.system(),
        }

    async def gather_node_measures(self, ctx: StatContext) -> Sequence[NodeMeasurement]:
        # TODO: Create our own k8s metric collector

        return []

    async def gather_container_measures(
        self,
        ctx: StatContext,
        container_ids: Sequence[str],
    ) -> Sequence[ContainerMeasurement]:
        # TODO: Implement Kubernetes-specific container metric collection

        return []

    async def gather_process_measures(
        self, ctx: StatContext, pid_map: Mapping[int, str]
    ) -> Sequence[ProcessMeasurement]:
        return []

    async def create_alloc_map(self) -> AbstractAllocMap:
        devices = await self.list_devices()
        return DiscretePropertyAllocMap(
            device_slots={
                dev.device_id: DeviceSlotInfo(
                    SlotTypes.COUNT, SlotName("cpu"), Decimal(dev.processing_units)
                )
                for dev in devices
            },
        )

    async def get_hooks(self, distro: str, arch: str) -> Sequence[Path]:
        # TODO: move the sysconf hook in libbaihook.so here
        return []

    async def generate_docker_args(
        self,
        docker: Docker,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> Mapping[str, Any]:
        # This function might be needed later to apply fine-grained tuning for
        # K8s resource allocation. NUMA memory pinning is not mirrored from the
        # Docker backend because this backend currently does not emit
        # per-container CPU/memory requests/limits; node-local placement on
        # Kubernetes would require Guaranteed-QoS pod specs plus cluster-level
        # Topology Manager configuration, which is out of scope here.
        return {}

    async def restore_from_container(
        self,
        container: Container,
        alloc_map: AbstractAllocMap,
    ) -> None:
        if not isinstance(alloc_map, DiscretePropertyAllocMap):
            raise InvalidAllocMapTypeError(
                f"Expected DiscretePropertyAllocMap, got {type(alloc_map).__name__}."
            )
        # Docker does not return the original cpuset.... :(
        # We need to read our own records.
        resource_spec = await get_resource_spec_from_container(container.backend_obj)
        if resource_spec is None:
            return
        alloc_map.apply_allocation({
            SlotName("cpu"): resource_spec.allocations[DeviceName("cpu")][SlotName("cpu")],
        })

    async def get_attached_devices(
        self,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> Sequence[DeviceModelInfo]:
        device_ids = [*device_alloc[SlotName("cpu")].keys()]
        available_devices = await self.list_devices()
        attached_devices: list[DeviceModelInfo] = []
        for device in available_devices:
            if device.device_id in device_ids:
                attached_devices.append({
                    "device_id": device.device_id,
                    "model_name": "",
                    "data": {"cores": len(device_ids)},
                })
        return attached_devices

    async def generate_mounts(
        self, source_path: Path, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]]
    ) -> list[MountInfo]:
        return []

    async def get_docker_networks(
        self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]]
    ) -> list[str]:
        return []

    def get_metadata(self) -> AcceleratorMetadata:
        return {
            "slot_name": "cpu",
            "description": "CPU",
            "human_readable_name": "CPU",
            "display_unit": "Core",
            "number_format": {"binary": False, "round_length": 0},
            "display_icon": "cpu",
        }


class MemoryDevice(AbstractComputeDevice):
    pass


class MemoryPlugin(AbstractComputePlugin):
    """
    Represents the main memory.

    When collecting statistics, it also measures network and I/O usage
    in addition to the memory usage.
    """

    config_watch_enabled = False

    key = DeviceName("mem")
    slot_types = [
        (SlotName("mem"), SlotTypes.BYTES),
    ]

    async def init(self, context: Any | None = None) -> None:
        pass

    async def cleanup(self) -> None:
        pass

    async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None:
        pass

    async def list_devices(self) -> Collection[MemoryDevice]:
        await K8sConfig.load_kube_config()
        core_api = K8sClient.CoreV1Api()

        nodes = (await core_api.list_node()).to_dict()["items"]
        overcommit_factor = int(os.environ.get("BACKEND_MEM_OVERCOMMIT_FACTOR", "1"))
        if not (1 <= overcommit_factor <= 10):
            raise InvalidOvercommitFactorError(
                f"Memory overcommit factor must be between 1 and 10, got {overcommit_factor}."
            )
        mem = 0
        for node in nodes:
            # if 'node-role.kubernetes.io/master' in node['metadata']['labels'].keys():
            #     continue
            mem += int(node["status"]["capacity"]["memory"][:-2]) * 1024
        return [
            MemoryDevice(
                device_id=DeviceId("root"),
                hw_location="root",
                numa_node=0,
                memory_size=mem * overcommit_factor,
                processing_units=0,
            ),
        ]

    async def available_slots(self) -> Mapping[SlotName, Decimal]:
        devices = await self.list_devices()
        return {
            SlotName("mem"): Decimal(sum(dev.memory_size for dev in devices)),
        }

    def get_version(self) -> str:
        return __version__

    async def extra_info(self) -> Mapping[str, str]:
        return {}

    async def gather_node_measures(self, ctx: StatContext) -> Sequence[NodeMeasurement]:
        # TODO: Create our own k8s metric collector
        return []

    async def gather_container_measures(
        self, ctx: StatContext, container_ids: Sequence[str]
    ) -> Sequence[ContainerMeasurement]:
        # TODO: Implement Kubernetes-specific container metric collection
        return []

    async def gather_process_measures(
        self, ctx: StatContext, pid_map: Mapping[int, str]
    ) -> Sequence[ProcessMeasurement]:
        return []

    async def create_alloc_map(self) -> AbstractAllocMap:
        devices = await self.list_devices()
        return DiscretePropertyAllocMap(
            allocation_strategy=AllocationStrategy.FILL,
            device_slots={
                dev.device_id: DeviceSlotInfo(
                    SlotTypes.BYTES, SlotName("mem"), Decimal(dev.memory_size)
                )
                for dev in devices
            },
        )

    async def get_hooks(self, distro: str, arch: str) -> Sequence[Path]:
        return []

    async def generate_docker_args(
        self,
        docker: Docker,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> Mapping[str, Any]:
        # This function might be needed later to apply fine-grained tuning for
        # K8s resource allocation
        return {}

    async def restore_from_container(
        self,
        container: Container,
        alloc_map: AbstractAllocMap,
    ) -> None:
        if not isinstance(alloc_map, DiscretePropertyAllocMap):
            raise InvalidAllocMapTypeError(
                f"Expected DiscretePropertyAllocMap, got {type(alloc_map).__name__}."
            )
        memory_limit = container.backend_obj["HostConfig"]["Memory"]
        alloc_map.apply_allocation({
            SlotName("mem"): {DeviceId("root"): memory_limit},
        })

    async def get_attached_devices(
        self,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> Sequence[DeviceModelInfo]:
        device_ids = [*device_alloc[SlotName("mem")].keys()]
        available_devices = await self.list_devices()
        attached_devices: list[DeviceModelInfo] = []
        for device in available_devices:
            if device.device_id in device_ids:
                attached_devices.append({
                    "device_id": device.device_id,
                    "model_name": "",
                    "data": {},
                })
        return attached_devices

    async def generate_mounts(
        self, source_path: Path, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]]
    ) -> list[MountInfo]:
        return []

    async def get_docker_networks(
        self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]]
    ) -> list[str]:
        return []

    def get_metadata(self) -> AcceleratorMetadata:
        return {
            "slot_name": "ram",
            "description": "Memory",
            "human_readable_name": "RAM",
            "display_unit": "GiB",
            "number_format": {"binary": True, "round_length": 0},
            "display_icon": "cpu",
        }