intrinsic.py

from __future__ import annotations

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

import aiohttp
import psutil
from aiodocker.docker import Docker, DockerContainer
from aiodocker.exceptions import DockerError

from ai.backend.agent import __version__  # pants: no-infer-dep
from ai.backend.agent.alloc_map import AllocationStrategy
from ai.backend.agent.docker.kernel import DockerKernel
from ai.backend.agent.errors import (
    InvalidResourceConfigError,
)
from ai.backend.agent.exception import InvalidArgumentError
from ai.backend.agent.plugin.network import (
    AbstractNetworkAgentPlugin,
    ContainerNetworkCapability,
    ContainerNetworkInfo,
)
from ai.backend.agent.resources import (
    AbstractAllocMap,
    AbstractComputeDevice,
    AbstractComputePlugin,
    DeviceSlotInfo,
    DiscretePropertyAllocMap,
)
from ai.backend.agent.stats import (
    ContainerMeasurement,
    Measurement,
    MetricTypes,
    NodeMeasurement,
    ProcessMeasurement,
    StatContext,
    StatModes,
)
from ai.backend.agent.types import Container, MountInfo
from ai.backend.agent.utils import read_sysfs
from ai.backend.agent.vendor.linux import libnuma
from ai.backend.common.asyncio import current_loop
from ai.backend.common.json import dump_json
from ai.backend.common.netns import nsenter
from ai.backend.common.types import (
    AcceleratorMetadata,
    ClusterInfo,
    DeviceId,
    DeviceModelInfo,
    DeviceName,
    KernelCreationConfig,
    MetricKey,
    SlotName,
    SlotTypes,
)
from ai.backend.common.utils import nmget
from ai.backend.logging import BraceStyleAdapter

from .resources import get_resource_spec_from_container

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


# The list of pruned fstype when checking the filesystem usage statistics.
# Note that psutil's linux implementation automatically filters out "non-device" filesystems by
# checking /proc/filesystems so we don't have to put all the details virtual filesystems like
# "sockfs", "debugfs", etc.
_CONTAINER_INSPECT_TIMEOUT: float = 2.0
_INVALID_PID: int = 0

# Stats stream reconnect parameters
_STATS_STREAM_INITIAL_BACKOFF: float = 1.0
_STATS_STREAM_MAX_BACKOFF: float = 30.0
_STATS_STREAM_BACKOFF_FACTOR: float = 2.0
_STATS_STREAM_MAX_RETRIES: int = 8

# The list of pruned fstype when checking the filesystem usage statistics.
pruned_disk_types = frozenset([
    "vfat",
    "lxcfs",
    "squashfs",
    "tmpfs",
    "iso9660",  # cdrom
])


@dataclasses.dataclass(frozen=True)
class ContainerNetStat:
    rx_bytes: int
    tx_bytes: int


def _parse_proc_net_dev(content: str) -> ContainerNetStat:
    """Parse /proc/net/dev content and return stats for non-lo interfaces."""
    rx_bytes = 0
    tx_bytes = 0
    for line in content.splitlines():
        if ":" not in line:
            continue
        iface, _, stats_str = line.partition(":")
        iface = iface.strip()
        if iface == "lo":
            continue
        fields = stats_str.split()
        # fields[0] = rx_bytes, fields[8] = tx_bytes
        rx_bytes += int(fields[0])
        tx_bytes += int(fields[8])
    return ContainerNetStat(rx_bytes=rx_bytes, tx_bytes=tx_bytes)


def read_proc_net_dev(container_pid: int) -> ContainerNetStat:
    """Read network stats from /proc/[pid]/net/dev for the given container PID.

    Parses the kernel's net/dev format directly from the container's proc entry,
    avoiding the need for namespace switching (setns) which is unreliable in
    threaded Python processes.
    """
    content = Path(f"/proc/{container_pid}/net/dev").read_text()
    return _parse_proc_net_dev(content)


def read_netns_net_dev(ns_path: Path) -> ContainerNetStat:
    """Read network stats by switching into the given network namespace.

    Uses setns() to enter the namespace, then reads /proc/thread-self/net/dev
    which reflects the calling thread's namespace (not the process-level one).

    This is the fallback for when the container PID is unavailable (PID=0).
    """
    with nsenter(ns_path):
        content = Path("/proc/thread-self/net/dev").read_text()
    return _parse_proc_net_dev(content)


def _validate_stats_entry(entry: dict[str, Any]) -> dict[str, Any] | None:
    if entry["read"].startswith("0001-01-01") or entry["preread"].startswith("0001-01-01"):
        return None
    return entry


class DockerStatsStreamer:
    """
    Maintains one long-lived ``container.stats(stream=True)`` reader per container
    and exposes the most recent decoded sample from an in-memory cache.

    Callers read the cached sample via :meth:`get_latest` instead of issuing a new
    HTTP round-trip every collection cycle.

    Reader lifecycle is driven by the agent's container lifecycle hooks:

    * :meth:`start` is called eagerly from the agent's ``_handle_start_event`` so
      a reader is spawned as soon as the container transitions to RUNNING.
    * :meth:`stop` is called from the agent's ``_handle_clean_event`` so the
      reader task is cancelled promptly when the container goes away.

    :meth:`get_latest` also lazily spawns a reader as a safety net for events
    that were missed.

    On transient transport failures (``ClientConnectionError`` /
    :class:`asyncio.TimeoutError`) the reader reconnects with bounded exponential
    backoff. If reconnection budget is exhausted, it logs an error and exits;
    a subsequent :meth:`get_latest` will start a fresh reader.
    """

    _docker: Docker
    _latest: dict[str, dict[str, Any]]
    _tasks: dict[str, asyncio.Task[None]]
    _closed: bool

    def __init__(self, docker: Docker) -> None:
        self._docker = docker
        self._latest = {}
        self._tasks = {}
        self._closed = False

    def start(self, container_id: str) -> None:
        """Eagerly start the stream reader for ``container_id`` if not already
        running. Idempotent; safe to call from container lifecycle hooks."""
        if self._closed:
            return
        task = self._tasks.get(container_id)
        if task is not None and not task.done():
            return
        self._tasks[container_id] = asyncio.create_task(
            self._read_stream(container_id),
            name=f"docker-stats-stream:{container_id[:7]}",
        )

    def get_latest(self, container_id: str) -> dict[str, Any] | None:
        """Return the most recent cached sample for ``container_id``.

        Lazily spawns a reader if none is running — intended as a safety net
        for cases where the container start event was missed. For newly-created
        containers the primary entry point should be :meth:`start` so the first
        frame lands in the cache by the time the next collection cycle runs.
        """
        if self._closed:
            return None
        task = self._tasks.get(container_id)
        if task is None or task.done():
            self.start(container_id)
        return self._latest.get(container_id)

    async def stop(self, container_id: str) -> None:
        """Cancel and await the reader for ``container_id``, and drop its
        cached sample. Re-raises :class:`asyncio.CancelledError` so the caller's
        cancellation propagates; other exceptions are logged and swallowed."""
        task = self._tasks.pop(container_id, None)
        self._latest.pop(container_id, None)
        if task is None or task.done():
            return
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            # Propagate only if the CURRENT task (the caller) is being cancelled;
            # a CancelledError bubbling out of the awaited task itself is expected.
            current = asyncio.current_task()
            if current is not None and current.cancelling() > 0:
                raise
        except Exception as e:
            log.warning(
                "stats stream stop: reader task for cid:{} raised: {!r}",
                container_id[:7],
                e,
            )

    async def close(self) -> None:
        """Cancel and await every in-flight reader task. Idempotent."""
        self._closed = True
        tasks = list(self._tasks.values())
        self._tasks.clear()
        self._latest.clear()
        for task in tasks:
            if not task.done():
                task.cancel()
        for task in tasks:
            try:
                await task
            except asyncio.CancelledError:
                current = asyncio.current_task()
                if current is not None and current.cancelling() > 0:
                    raise
            except Exception as e:
                log.warning("stats stream close: reader task raised: {!r}", e)

    async def _read_stream(self, container_id: str) -> None:
        """Run the long-lived reader for ``container_id`` with bounded
        exponential-backoff reconnect on transient transport failures."""
        short_cid = container_id[:7]
        backoff = _STATS_STREAM_INITIAL_BACKOFF
        retries = 0
        try:
            while True:
                consumed_any = False
                try:
                    consumed_any = await self._consume_stream(container_id)
                except asyncio.CancelledError:
                    raise
                except RuntimeError as e:
                    msg = str(e.args[0]).lower() if e.args else ""
                    if "event loop is closed" in msg or "session is closed" in msg:
                        return
                    log.warning(
                        "stats stream stopped unexpectedly (cid:{}): {!r}",
                        short_cid,
                        e,
                    )
                    return
                except (aiohttp.ClientConnectionError, TimeoutError) as e:
                    if retries >= _STATS_STREAM_MAX_RETRIES:
                        log.error(
                            "stats stream exhausted retries for cid:{}: {!r}",
                            short_cid,
                            e,
                        )
                        return
                    retries += 1
                    wait = min(backoff, _STATS_STREAM_MAX_BACKOFF)
                    log.warning(
                        "stats stream transient failure (cid:{}) retry {}/{} in {:.1f}s: {!r}",
                        short_cid,
                        retries,
                        _STATS_STREAM_MAX_RETRIES,
                        wait,
                        e,
                    )
                    await asyncio.sleep(wait)
                    backoff = min(
                        backoff * _STATS_STREAM_BACKOFF_FACTOR,
                        _STATS_STREAM_MAX_BACKOFF,
                    )
                    continue
                except DockerError as e:
                    # 404 / container removed / etc. — stop cleanly.
                    log.debug("stats stream ended (cid:{}): {!r}", short_cid, e)
                    return
                # Normal upstream-closed exit (e.g. container removed).
                if consumed_any:
                    # Reset backoff after a successful run; stream may just have
                    # ended because the container was removed.
                    return
                # Stream ended before yielding any frame — treat as transient.
                if retries >= _STATS_STREAM_MAX_RETRIES:
                    return
                retries += 1
                await asyncio.sleep(min(backoff, _STATS_STREAM_MAX_BACKOFF))
                backoff = min(
                    backoff * _STATS_STREAM_BACKOFF_FACTOR,
                    _STATS_STREAM_MAX_BACKOFF,
                )
        finally:
            self._latest.pop(container_id, None)

    async def _consume_stream(self, container_id: str) -> bool:
        """Consume one docker stats stream iterator. Returns True if at least
        one frame was received before the stream ended normally."""
        container = DockerContainer(self._docker, id=container_id)
        consumed_any = False
        async for frame in container.stats(stream=True):
            validated = _validate_stats_entry(frame)
            if validated is not None:
                self._latest[container_id] = validated
                consumed_any = True
        return consumed_any


# 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),
    ]

    _docker: Docker
    _stats_streamer: DockerStatsStreamer

    async def init(self, context: Any | None = None) -> None:
        self._docker = Docker()

    async def cleanup(self) -> None:
        await self._docker.close()

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

    def attach_stats_streamer(self, streamer: DockerStatsStreamer) -> None:
        """Attach the agent-owned :class:`DockerStatsStreamer` used for reading
        per-container stats. Called once by :class:`DockerAgent` after plugin
        init so the streamer is shared across intrinsic plugins."""
        self._stats_streamer = streamer

    async def list_devices(self) -> Collection[CPUDevice]:
        cores = await libnuma.get_available_cores()
        overcommit_factor = int(os.environ.get("BACKEND_CPU_OVERCOMMIT_FACTOR", "1"))
        if not (1 <= overcommit_factor <= 10):
            raise InvalidResourceConfigError(
                f"BACKEND_CPU_OVERCOMMIT_FACTOR must be between 1 and 10, got {overcommit_factor}"
            )
        return [
            CPUDevice(
                device_id=DeviceId(str(core_idx)),
                hw_location="root",
                numa_node=libnuma.node_of_cpu(core_idx),
                memory_size=0,
                processing_units=1 * overcommit_factor,
            )
            for core_idx in sorted(cores)
        ]

    async def available_slots(self) -> Mapping[SlotName, Decimal]:
        devices = await self.list_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]:
        _cstat = psutil.cpu_times(True)
        q = Decimal("0.000")
        total_cpu_used = cast(
            Decimal, sum((Decimal(c.user + c.system) * 1000).quantize(q) for c in _cstat)
        )
        now, raw_interval = ctx.update_timestamp("cpu-node")
        interval = Decimal(raw_interval * 1000).quantize(q)

        return [
            NodeMeasurement(
                MetricKey("cpu_util"),
                MetricTypes.UTILIZATION,
                unit_hint="msec",
                current_hook=lambda metric: metric.stats.diff,
                per_node=Measurement(total_cpu_used, interval),
                per_device={
                    DeviceId(str(idx)): Measurement(
                        (Decimal(c.user + c.system) * 1000).quantize(q),
                        interval,
                    )
                    for idx, c in enumerate(_cstat)
                },
            ),
        ]

    async def gather_container_measures(
        self,
        ctx: StatContext,
        container_ids: Sequence[str],
    ) -> Sequence[ContainerMeasurement]:
        if not container_ids:
            return []

        async def sysfs_impl(container_id: str) -> float | None:
            cpu_path = ctx.agent.get_cgroup_path("cpuacct", container_id)
            version = ctx.agent.docker_info["CgroupVersion"]  # type: ignore[attr-defined]
            try:
                match version:
                    case "1":
                        cpu_used = read_sysfs(cpu_path / "cpuacct.usage", int) / 1e6
                    case "2":
                        cpu_stats = {
                            k: v
                            for k, v in map(
                                lambda line: line.split(" "),
                                (cpu_path / "cpu.stat").read_text().splitlines(),
                            )
                        }
                        cpu_used = int(cpu_stats["usage_usec"]) / 1e3
                    case _:
                        return None
            except OSError as e:
                log.warning(
                    "CPUPlugin: cannot read stats: sysfs unreadable for container {0}\n{1!r}",
                    container_id[:7],
                    e,
                )
                return None
            return cpu_used

        # TODO(#11223): After sysfs-first lands, CPU is sourced from sysfs and this stream becomes redundant here; migrate to a network/IO consumer.
        async def api_impl(container_id: str) -> float | None:
            ret = self._stats_streamer.get_latest(container_id)
            if ret is None:
                return None
            cpu_usage = cast(float, nmget(ret, "cpu_stats.cpu_usage.total_usage", 0))
            return cpu_usage / 1e6

        if ctx.mode == StatModes.CGROUP:
            impl = sysfs_impl
        elif ctx.mode == StatModes.DOCKER:
            impl = api_impl
        else:
            raise RuntimeError("should not reach here")

        tasks = []
        for cid in container_ids:
            tasks.append(asyncio.create_task(impl(cid)))
        results = await asyncio.gather(*tasks)

        q = Decimal("0.000")
        per_container_cpu_used = {}
        per_container_cpu_util = {}
        for cid, cpu_used in zip(container_ids, results, strict=True):
            if cpu_used is None:
                continue
            per_container_cpu_used[cid] = Measurement(Decimal(cpu_used).quantize(q))
            per_container_cpu_util[cid] = Measurement(
                Decimal(cpu_used).quantize(q),
                capacity=Decimal(1000),
            )
        return [
            ContainerMeasurement(
                MetricKey("cpu_util"),
                MetricTypes.UTILIZATION,
                unit_hint="percent",
                current_hook=lambda metric: metric.stats.rate,
                stats_filter=frozenset({"avg", "max"}),
                per_container=per_container_cpu_util,
            ),
            ContainerMeasurement(
                MetricKey("cpu_used"),
                MetricTypes.ACCUMULATION,
                unit_hint="msec",
                per_container=per_container_cpu_used,
            ),
        ]

    async def gather_process_measures(
        self, ctx: StatContext, pid_map: Mapping[int, str]
    ) -> Sequence[ProcessMeasurement]:
        async def psutil_impl(pid: int, cid: str) -> Decimal | None:
            try:
                p = psutil.Process(pid)
                cpu_times = p.cpu_times()
            except psutil.NoSuchProcess:
                log.debug("Process not found for CPU stats (pid:{0}, container id:{1})", pid, cid)
            else:
                return Decimal(cpu_times.user + cpu_times.system) * 1000
            return None

        async def api_impl(_cid: str, _pids: list[int]) -> list[Decimal | None]:
            return []

        per_process_cpu_util = {}
        per_process_cpu_used = {}
        results: list[Decimal | None] = []
        q = Decimal("0.000")
        pid_map_list = list(pid_map.items())
        match self.local_config["agent"]["docker-mode"]:
            case "linuxkit":
                api_tasks: list[asyncio.Task[list[Decimal | None]]] = []
                # group by container ID
                cid_pids_map: dict[str, list[int]] = {}
                for pid, cid in pid_map_list:
                    if cid_pids_map.get(cid) is None:
                        cid_pids_map[cid] = []
                    cid_pids_map[cid].append(pid)
                for cid, pids in cid_pids_map.items():
                    api_tasks.append(asyncio.create_task(api_impl(cid, pids)))
                chunked_results = await asyncio.gather(*api_tasks)
                for chunk in chunked_results:
                    results.extend(chunk)
            case _:
                psutil_tasks = []
                for pid, cid in pid_map_list:
                    psutil_tasks.append(asyncio.create_task(psutil_impl(pid, cid)))
                results = await asyncio.gather(*psutil_tasks)

        for (pid, cid), cpu_used in zip(pid_map_list, results, strict=True):
            if cpu_used is None:
                continue
            per_process_cpu_util[pid] = Measurement(
                Decimal(cpu_used).quantize(q), capacity=Decimal(1000)
            )
            per_process_cpu_used[pid] = Measurement(Decimal(cpu_used).quantize(q))
        return [
            ProcessMeasurement(
                MetricKey("cpu_util"),
                MetricTypes.UTILIZATION,
                unit_hint="percent",
                current_hook=lambda metric: metric.stats.rate,
                stats_filter=frozenset({"avg", "max"}),
                per_process=per_process_cpu_util,
            ),
            ProcessMeasurement(
                MetricKey("cpu_used"),
                MetricTypes.ACCUMULATION,
                unit_hint="msec",
                per_process=per_process_cpu_used,
            ),
        ]

    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]:
        cores = [*map(int, device_alloc[SlotName("cpu")].keys())]
        sorted_core_ids = [*map(str, sorted(cores))]
        return {
            "HostConfig": {
                "Cpus": len(cores),
                "CpusetCpus": ",".join(sorted_core_ids),
                # 'CpusetMems': f'{resource_spec.numa_node}',
            },
        }

    async def restore_from_container(
        self,
        container: Container,
        alloc_map: AbstractAllocMap,
    ) -> None:
        if not isinstance(alloc_map, DiscretePropertyAllocMap):
            raise InvalidArgumentError(
                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 get_docker_networks(
        self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]]
    ) -> list[str]:
        return []

    async def generate_mounts(
        self, source_path: Path, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]]
    ) -> list[MountInfo]:
        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),
    ]

    _docker: Docker
    _stats_streamer: DockerStatsStreamer

    async def init(self, context: Any | None = None) -> None:
        self._docker = Docker()

    async def cleanup(self) -> None:
        await self._docker.close()

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

    def attach_stats_streamer(self, streamer: DockerStatsStreamer) -> None:
        """Attach the agent-owned :class:`DockerStatsStreamer` used for reading
        per-container stats. Called once by :class:`DockerAgent` after plugin
        init so the streamer is shared across intrinsic plugins."""
        self._stats_streamer = streamer

    async def list_devices(self) -> Collection[MemoryDevice]:
        memory_size = psutil.virtual_memory().total
        overcommit_factor = int(os.environ.get("BACKEND_MEM_OVERCOMMIT_FACTOR", "1"))
        return [
            MemoryDevice(
                device_id=DeviceId("root"),
                device_name=self.key,
                hw_location="root",
                numa_node=0,  # the kernel setting will do the job.
                memory_size=overcommit_factor * memory_size,
                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]:
        _mstat = psutil.virtual_memory()
        total_mem_used_bytes = Decimal(_mstat.total - _mstat.available)
        total_mem_capacity_bytes = Decimal(_mstat.total)
        _nstat = psutil.net_io_counters()
        net_rx_bytes = _nstat.bytes_recv
        net_tx_bytes = _nstat.bytes_sent

        def get_disk_stat() -> tuple[Decimal, Decimal, dict[DeviceId, Measurement]]:
            total_disk_usage = Decimal(0)
            total_disk_capacity = Decimal(0)
            per_disk_stat: dict[DeviceId, Measurement] = {}
            for disk_info in psutil.disk_partitions():
                # Skip additional filesystem types not filtered by psutil, like squashfs.
                if disk_info.fstype in pruned_disk_types:
                    continue
                # Skip transient filesystems created/destroyed by Docker.
                if disk_info.mountpoint.startswith("/proc/docker/runtime-runc/moby/"):
                    continue
                # Skip btrfs subvolumes used by Docker if configured.
                if disk_info.mountpoint == "/var/lib/docker/btrfs":
                    continue
                dstat = os.statvfs(disk_info.mountpoint)
                disk_usage = Decimal(dstat.f_frsize * (dstat.f_blocks - dstat.f_bavail))
                disk_capacity = Decimal(dstat.f_frsize * dstat.f_blocks)
                per_disk_stat[DeviceId(disk_info.device)] = Measurement(disk_usage, disk_capacity)
                total_disk_usage += disk_usage
                total_disk_capacity += disk_capacity
            return total_disk_usage, total_disk_capacity, per_disk_stat

        loop = current_loop()
        total_disk_usage, total_disk_capacity, per_disk_stat = await loop.run_in_executor(
            None, get_disk_stat
        )
        return [
            NodeMeasurement(
                MetricKey("mem"),
                MetricTypes.GAUGE,
                unit_hint="bytes",
                stats_filter=frozenset({"max"}),
                per_node=Measurement(total_mem_used_bytes, total_mem_capacity_bytes),
                per_device={
                    DeviceId("root"): Measurement(total_mem_used_bytes, total_mem_capacity_bytes)
                },
            ),
            NodeMeasurement(
                MetricKey("disk"),
                MetricTypes.GAUGE,
                unit_hint="bytes",
                per_node=Measurement(total_disk_usage, total_disk_capacity),
                per_device=per_disk_stat,
            ),
            NodeMeasurement(
                MetricKey("net_rx"),
                MetricTypes.RATE,
                unit_hint="bps",
                current_hook=lambda metric: metric.stats.rate,
                per_node=Measurement(Decimal(net_rx_bytes)),
                per_device={DeviceId("node"): Measurement(Decimal(net_rx_bytes))},
            ),
            NodeMeasurement(
                MetricKey("net_tx"),
                MetricTypes.RATE,
                unit_hint="bps",
                current_hook=lambda metric: metric.stats.rate,
                per_node=Measurement(Decimal(net_tx_bytes)),
                per_device={DeviceId("node"): Measurement(Decimal(net_tx_bytes))},
            ),
        ]

    async def gather_container_measures(
        self, ctx: StatContext, container_ids: Sequence[str]
    ) -> Sequence[ContainerMeasurement]:
        if not container_ids:
            return []

        def get_scratch_size(_container_id: str) -> int:
            # Temporarily disabled as this function incurs too much delay with
            # a large number of files in scratch dirs, causing indefinite accumulation of
            # stat collector tasks and slowing down everything.
            return 0
            # for kernel_id, info in ctx.agent.kernel_registry.items():
            #     if info['container_id'] == container_id:
            #         break
            # else:
            #     return 0
            # work_dir = ctx.agent.local_config['container']['scratch-root'] / str(kernel_id) / 'work'
            # total_size = 0
            # for path in work_dir.rglob('*'):
            #     if path.is_symlink():
            #         total_size += path.lstat().st_size
            #     elif path.is_file():
            #         total_size += path.stat().st_size
            # return total_size

        async def sysfs_impl(
            container_id: str,
        ) -> tuple[int, int, int, int, int, int, int] | None:
            mem_path = ctx.agent.get_cgroup_path("memory", container_id)
            io_path = ctx.agent.get_cgroup_path("blkio", container_id)
            version = ctx.agent.get_cgroup_version()

            try:
                io_read_bytes = 0
                io_write_bytes = 0
                match version:
                    case "1":
                        mem_cur_bytes = read_sysfs(mem_path / "memory.usage_in_bytes", int)
                        mem_max_bytes = read_sysfs(mem_path / "memory.limit_in_bytes", int)

                        for line in (mem_path / "memory.stat").read_text().splitlines():
                            key, _, value = line.partition(" ")
                            if key == "total_inactive_file":
                                try:
                                    mem_cur_bytes -= int(value)
                                except ValueError:
                                    log.warning(
                                        "MemoryPlugin: cannot parse inactive stat. container: {0}",
                                        container_id[:7],
                                    )
                                break

                        # example data:
                        #   8:0 Read 13918208
                        #   8:0 Write 0
                        #   8:0 Sync 0
                        #   8:0 Async 13918208
                        #   8:0 Total 13918208
                        #   Total 13918208
                        for line in (
                            (io_path / "blkio.throttle.io_service_bytes").read_text().splitlines()
                        ):
                            if line.startswith("Total "):
                                continue
                            dev, op, nbytes = line.strip().split()
                            if op == "Read":
                                io_read_bytes += int(nbytes)
                            elif op == "Write":
                                io_write_bytes += int(nbytes)
                    case "2":
                        mem_cur_bytes = read_sysfs(mem_path / "memory.current", int)
                        mem_max_bytes = read_sysfs(mem_path / "memory.max", int)

                        for line in (mem_path / "memory.stat").read_text().splitlines():
                            key, _, value = line.partition(" ")
                            if key == "inactive_file":
                                try:
                                    mem_cur_bytes -= int(value)
                                except ValueError:
                                    log.warning(
                                        "MemoryPlugin: cannot parse inactive stat. container: {0}",
                                        container_id[:7],
                                    )
                                break

                        # example data:
                        # 8:16 rbytes=1459200 wbytes=314773504 rios=192 wios=353 dbytes=0 dios=0
                        # 8:0 rbytes=3387392 wbytes=176128 rios=103 wios=32 dbytes=0 dios=0
                        # 253:0 8:0 rbytes=3387392 wbytes=176128 rios=103 wios=32 dbytes=0 dios=0
                        for line in (io_path / "io.stat").read_text().splitlines():
                            for io_stat in line.split():
                                stat, _, value = io_stat.partition("=")
                                if stat == "rbytes":
                                    io_read_bytes += int(value)
                                if stat == "wbytes":
                                    io_write_bytes += int(value)
                    case _:
                        return None
            except OSError as e:
                log.warning(
                    "MemoryPlugin: cannot read stats: sysfs unreadable for container {0}\n{1!r}",
                    container_id[:7],
                    e,
                )
                return None
            container = DockerContainer(self._docker, id=container_id)
            try:
                async with asyncio.timeout(_CONTAINER_INSPECT_TIMEOUT):
                    data = await container.show()
                    container_pid: int = data.get("State", {}).get("Pid", _INVALID_PID)
            except TimeoutError:
                log.warning(
                    "MemoryPlugin: timeout reading container info for container {0}",
                    container_id[:7],
                )
                return None
            net_stat = ContainerNetStat(rx_bytes=0, tx_bytes=0)
            loop = current_loop()
            if container_pid > 0:
                try:
                    net_stat = await loop.run_in_executor(None, read_proc_net_dev, container_pid)
                except OSError as e:
                    log.warning(
                        "MemoryPlugin: cannot read net stats for container {0} (pid={1}): {2!r}",
                        container_id[:7],
                        container_pid,
                        e,
                    )
            else:
                sandbox_key = data.get("NetworkSettings", {}).get("SandboxKey", "")
                ns_path = Path(sandbox_key) if sandbox_key else None
                if ns_path and ns_path.exists():
                    try:
                        net_stat = await loop.run_in_executor(None, read_netns_net_dev, ns_path)
                    except OSError as e:
                        log.warning(
                            "MemoryPlugin: cannot read net stats via netns for"
                            " container {0} (sandbox_key={1!r}): {2!r}",
                            container_id[:7],
                            sandbox_key,
                            e,
                        )
                else:
                    log.warning(
                        "MemoryPlugin: container {0} has no PID and no valid SandboxKey,"
                        " skipping net stat collection",
                        container_id[:7],
                    )
            loop = current_loop()
            scratch_sz = await loop.run_in_executor(None, get_scratch_size, container_id)
            return (
                mem_cur_bytes,
                mem_max_bytes,
                io_read_bytes,
                io_write_bytes,
                net_stat.rx_bytes,
                net_stat.tx_bytes,
                scratch_sz,
            )

        # TODO(#11223): After sysfs-first lands, memory is sourced from sysfs and this stream becomes redundant here; migrate to a network/IO consumer.
        async def api_impl(
            container_id: str,
        ) -> tuple[int, int, int, int, int, int, int] | None:
            ret = self._stats_streamer.get_latest(container_id)
            if ret is None:
                return None
            mem_cur_bytes = nmget(ret, "memory_stats.usage", 0)
            mem_total_bytes = nmget(ret, "memory_stats.limit", 0)
            io_read_bytes = 0
            io_write_bytes = 0
            for item in nmget(ret, "blkio_stats.io_service_bytes_recursive", []):
                if item["op"] == "Read":
                    io_read_bytes += item["value"]
                elif item["op"] == "Write":
                    io_write_bytes += item["value"]
            net_rx_bytes = 0
            net_tx_bytes = 0
            for name, stat in ret["networks"].items():
                net_rx_bytes += stat["rx_bytes"]
                net_tx_bytes += stat["tx_bytes"]
            loop = current_loop()
            scratch_sz = await loop.run_in_executor(None, get_scratch_size, container_id)
            return (
                mem_cur_bytes,
                mem_total_bytes,
                io_read_bytes,
                io_write_bytes,
                net_rx_bytes,
                net_tx_bytes,
                scratch_sz,
            )

        if ctx.mode == StatModes.CGROUP:
            impl = sysfs_impl
        elif ctx.mode == StatModes.DOCKER:
            impl = api_impl
        else:
            raise RuntimeError("should not reach here")

        per_container_mem_used_bytes = {}
        per_container_io_read_bytes = {}
        per_container_io_write_bytes = {}
        per_container_net_rx_bytes = {}
        per_container_net_tx_bytes = {}
        per_container_io_scratch_size = {}
        tasks = []
        for cid in container_ids:
            tasks.append(asyncio.create_task(impl(cid)))
        results = await asyncio.gather(*tasks, return_exceptions=True)
        for cid, result in zip(container_ids, results, strict=True):
            if result is None:
                continue
            if isinstance(result, Exception):
                log.warning(
                    "gather_container_measures: error collecting stats for {}: {}", cid, result
                )
                continue
            if isinstance(result, BaseException):
                raise result
            per_container_mem_used_bytes[cid] = Measurement(
                Decimal(result[0]), capacity=Decimal(result[1])
            )
            per_container_io_read_bytes[cid] = Measurement(Decimal(result[2]))
            per_container_io_write_bytes[cid] = Measurement(Decimal(result[3]))
            per_container_net_rx_bytes[cid] = Measurement(Decimal(result[4]))
            per_container_net_tx_bytes[cid] = Measurement(Decimal(result[5]))
            per_container_io_scratch_size[cid] = Measurement(Decimal(result[6]))
        return [
            ContainerMeasurement(
                MetricKey("mem"),
                MetricTypes.GAUGE,
                unit_hint="bytes",
                stats_filter=frozenset({"max"}),
                per_container=per_container_mem_used_bytes,
            ),
            ContainerMeasurement(
                MetricKey("io_read"),
                MetricTypes.GAUGE,
                unit_hint="bytes",
                stats_filter=frozenset({"rate"}),
                per_container=per_container_io_read_bytes,
            ),
            ContainerMeasurement(
                MetricKey("io_write"),
                MetricTypes.GAUGE,
                unit_hint="bytes",
                stats_filter=frozenset({"rate"}),
                per_container=per_container_io_write_bytes,
            ),
            ContainerMeasurement(
                MetricKey("net_rx"),
                MetricTypes.RATE,
                unit_hint="bps",
                current_hook=lambda metric: metric.stats.rate,
                per_container=per_container_net_rx_bytes,
            ),
            ContainerMeasurement(
                MetricKey("net_tx"),
                MetricTypes.RATE,
                unit_hint="bps",
                current_hook=lambda metric: metric.stats.rate,
                per_container=per_container_net_tx_bytes,
            ),
            ContainerMeasurement(
                MetricKey("io_scratch_size"),
                MetricTypes.GAUGE,
                unit_hint="bytes",
                stats_filter=frozenset({"max"}),
                per_container=per_container_io_scratch_size,
            ),
        ]

    async def gather_process_measures(
        self, ctx: StatContext, pid_map: Mapping[int, str]
    ) -> Sequence[ProcessMeasurement]:
        async def psutil_impl(pid: int, cid: str) -> tuple[int | None, int | None, int | None]:
            try:
                p = psutil.Process(pid)
                stats = p.as_dict(attrs=["memory_info", "io_counters"])
            except psutil.NoSuchProcess:
                log.debug(
                    "Process not found for memory stats (pid:{0}, container id:{1})",
                    pid,
                    cid,
                )
            else:
                mem_cur_bytes = io_read_bytes = io_write_bytes = None
                if stats["memory_info"] is not None:
                    mem_cur_bytes = stats["memory_info"].rss
                if stats["io_counters"] is not None:
                    io_read_bytes = stats["io_counters"].read_bytes
                    io_write_bytes = stats["io_counters"].write_bytes
                return mem_cur_bytes, io_read_bytes, io_write_bytes
            return None, None, None

        async def api_impl(
            _cid: str, _pids: list[int]
        ) -> list[tuple[int | None, int | None, int | None]]:
            return []

        per_process_mem_used_bytes = {}
        per_process_io_read_bytes = {}
        per_process_io_write_bytes = {}
        results: list[tuple[int | None, int | None, int | None]]
        pid_map_list = list(pid_map.items())
        match self.local_config["agent"]["docker-mode"]:
            case "linuxkit":
                api_tasks = []
                # group by container ID
                cid_pids_map: dict[str, list[int]] = {}
                for pid, cid in pid_map_list:
                    if cid_pids_map.get(cid) is None:
                        cid_pids_map[cid] = []
                    cid_pids_map[cid].append(pid)
                for cid, pids in cid_pids_map.items():
                    api_tasks.append(asyncio.create_task(api_impl(cid, pids)))
                chunked_results = await asyncio.gather(*api_tasks)
                results = []
                for chunk in chunked_results:
                    results.extend(chunk)
            case _:
                psutil_tasks = []
                for pid, cid in pid_map_list:
                    psutil_tasks.append(asyncio.create_task(psutil_impl(pid, cid)))
                results = await asyncio.gather(*psutil_tasks)

        for (pid, _), result in zip(pid_map_list, results, strict=True):
            mem, io_read, io_write = result
            if mem is not None:
                per_process_mem_used_bytes[pid] = Measurement(Decimal(mem))
            if io_read is not None:
                per_process_io_read_bytes[pid] = Measurement(Decimal(io_read))
            if io_write is not None:
                per_process_io_write_bytes[pid] = Measurement(Decimal(io_write))
        return [
            ProcessMeasurement(
                MetricKey("mem"),
                MetricTypes.GAUGE,
                unit_hint="bytes",
                stats_filter=frozenset({"max"}),
                per_process=per_process_mem_used_bytes,
            ),
            ProcessMeasurement(
                MetricKey("io_read"),
                MetricTypes.GAUGE,
                unit_hint="bytes",
                stats_filter=frozenset({"rate"}),
                per_process=per_process_io_read_bytes,
            ),
            ProcessMeasurement(
                MetricKey("io_write"),
                MetricTypes.GAUGE,
                unit_hint="bytes",
                stats_filter=frozenset({"rate"}),
                per_process=per_process_io_write_bytes,
            ),
        ]

    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]:
        memory = sum(device_alloc[SlotName("mem")].values())
        return {
            "HostConfig": {
                "MemorySwap": int(memory),  # prevent using swap!
                "Memory": int(memory),
            },
        }

    async def restore_from_container(
        self,
        container: Container,
        alloc_map: AbstractAllocMap,
    ) -> None:
        if not isinstance(alloc_map, DiscretePropertyAllocMap):
            raise InvalidArgumentError(
                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 get_docker_networks(
        self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]]
    ) -> list[str]:
        return []

    async def generate_mounts(
        self, source_path: Path, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]]
    ) -> list[MountInfo]:
        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": "ram",
        }


class OverlayNetworkPlugin(AbstractNetworkAgentPlugin[DockerKernel]):
    async def init(self, context: Any = None) -> None:
        pass

    async def cleanup(self) -> None:
        pass

    async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None:
        return await super().update_plugin_config(plugin_config)

    async def get_capabilities(self) -> set[ContainerNetworkCapability]:
        return set()

    async def join_network(
        self,
        kernel_config: KernelCreationConfig,
        cluster_info: ClusterInfo,
        **kwargs: Any,
    ) -> dict[str, Any]:
        network_name: str = kwargs["network_name"]
        return {
            "HostConfig": {
                "NetworkMode": network_name,
            },
            "NetworkingConfig": {
                "EndpointsConfig": {
                    network_name: {
                        "Aliases": [kernel_config["cluster_hostname"]],
                        "DriverOpts": {"com.docker.network.endpoint.ifname": "baimulti0"},
                    },
                },
            },
        }

    async def leave_network(self, kernel: DockerKernel) -> None:
        pass


class HostNetworkPlugin(AbstractNetworkAgentPlugin[DockerKernel]):
    async def init(self, context: Any = None) -> None:
        pass

    async def cleanup(self) -> None:
        pass

    async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None:
        return await super().update_plugin_config(plugin_config)

    async def get_capabilities(self) -> set[ContainerNetworkCapability]:
        return {ContainerNetworkCapability.GLOBAL}

    async def join_network(
        self, kernel_config: KernelCreationConfig, cluster_info: ClusterInfo, **kwargs: Any
    ) -> dict[str, Any]:
        if _cluster_ssh_port_mapping := cluster_info.get("cluster_ssh_port_mapping"):
            return {
                "HostConfig": {
                    "ExtraHosts": [
                        f"{hostname}:{host_port[0]}"
                        for hostname, host_port in _cluster_ssh_port_mapping.items()
                    ],
                    "NetworkMode": "host",
                }
            }
        return {
            "HostConfig": {
                "NetworkMode": "host",
            },
        }

    async def leave_network(self, kernel: DockerKernel) -> None:
        pass

    async def prepare_port_forward(
        self, kernel: DockerKernel, bind_host: str, ports: Iterable[tuple[int, int]], **kwargs: Any
    ) -> None:
        host_ports = [p[0] for p in ports]
        scratch_dir = (
            self.local_config["container"]["scratch-root"] / str(kernel.kernel_id)
        ).resolve()
        config_dir: Path = scratch_dir / "config"

        intrinsic_ports = {
            "replin": host_ports[0],
            "replout": host_ports[1],
        }
        for index, port_info in enumerate(kernel.service_ports):
            port_name = port_info["name"]
            if port_name in ("sshd", "ttyd"):
                intrinsic_ports[port_name] = host_ports[index + 2]

        await current_loop().run_in_executor(
            None,
            lambda: (config_dir / "intrinsic-ports.json").write_bytes(dump_json(intrinsic_ports)),
        )

    async def expose_ports(
        self, kernel: DockerKernel, bind_host: str, ports: Iterable[tuple[int, int]], **kwargs: Any
    ) -> ContainerNetworkInfo:
        host_ports = [p[0] for p in ports]

        intrinsic_ports = {
            "replin": host_ports[0],
            "replout": host_ports[1],
        }
        for index, port_info in enumerate(kernel.service_ports):
            port_name = port_info["name"]
            if port_name in ("sshd", "ttyd"):
                intrinsic_ports[port_name] = host_ports[index + 2]

        return ContainerNetworkInfo(
            bind_host,
            {service_name: {port: port} for service_name, port in intrinsic_ports.items()},
        )