stackpilot/backend/services/compose_edit_service.py

"""Server-side merging of wizard fragments into a compose YAML string.

PyYAML round-trips lose comments; that is acceptable here because these helpers
are invoked from the editor where the returned YAML replaces the buffer and the
user reviews before saving.
"""
from __future__ import annotations

import yaml

from services import gpu_service, volume_service


class EditError(Exception):
    pass


def _load(yaml_str: str) -> dict:
    try:
        data = yaml.safe_load(yaml_str) or {}
    except yaml.YAMLError as exc:
        raise EditError(f"Invalid YAML: {exc}") from exc
    if not isinstance(data, dict):
        raise EditError("Top-level compose document must be a mapping")
    data.setdefault("services", {})
    if not isinstance(data["services"], dict):
        raise EditError("`services` must be a mapping")
    return data


def _dump(data: dict) -> str:
    # Drop empty top-level keys we may have created.
    if not data.get("services"):
        data.pop("services", None)
    if "volumes" in data and not data["volumes"]:
        data.pop("volumes", None)
    return yaml.safe_dump(data, sort_keys=False, default_flow_style=False)


def _get_service(data: dict, service: str) -> dict:
    svc = data["services"].get(service)
    if svc is None:
        svc = {}
        data["services"][service] = svc
    if not isinstance(svc, dict):
        raise EditError(f"Service '{service}' is not a mapping")
    return svc


# --------------------------------------------------------------------------- #
# Volumes
# --------------------------------------------------------------------------- #


def add_volume(yaml_str: str, service: str, spec: dict) -> str:
    data = _load(yaml_str)
    svc = _get_service(data, service)

    # Top-level volume declaration (named/nfs/smb).
    definition = volume_service.build_volume_definition(spec)
    if definition:
        top = data.setdefault("volumes", {})
        if not isinstance(top, dict):
            raise EditError("`volumes` must be a mapping")
        top.update(definition)

    mount = volume_service.build_service_mount(spec)
    vols = svc.setdefault("volumes", [])
    if not isinstance(vols, list):
        raise EditError(f"Service '{service}' volumes must be a list")
    if mount not in vols:
        vols.append(mount)
    return _dump(data)


# --------------------------------------------------------------------------- #
# GPU
# --------------------------------------------------------------------------- #


def set_gpu(yaml_str: str, service: str, config: dict) -> str:
    """config = { mode: 'none'|'nvidia'|'dri', ... }"""
    data = _load(yaml_str)
    svc = _get_service(data, service)
    mode = config.get("mode", "none")

    gpu_service.remove_gpu(svc)

    if mode == "nvidia":
        gpu_service.inject_nvidia(
            svc,
            device_ids=config.get("device_ids") or None,
            count=config.get("count"),
            capabilities=config.get("capabilities") or ["gpu"],
        )
    elif mode == "dri":
        gpu_service.inject_dri(
            svc,
            vendor=config.get("vendor", "intel"),
            add_render_group=config.get("add_render_group", True),
            add_video_group=config.get("add_video_group", False),
            set_libva=config.get("set_libva", False),
        )
    return _dump(data)


# --------------------------------------------------------------------------- #
# Devices / privileged
# --------------------------------------------------------------------------- #


def add_device(yaml_str: str, service: str, host_path: str, target: str | None = None) -> str:
    data = _load(yaml_str)
    svc = _get_service(data, service)
    entry = f"{host_path}:{target or host_path}"
    devices = svc.setdefault("devices", [])
    if not isinstance(devices, list):
        raise EditError(f"Service '{service}' devices must be a list")
    if entry not in devices:
        devices.append(entry)
    return _dump(data)


def remove_device(yaml_str: str, service: str, host_path: str) -> str:
    data = _load(yaml_str)
    svc = _get_service(data, service)
    devices = svc.get("devices", [])
    if isinstance(devices, list):
        svc["devices"] = [
            d for d in devices if not str(d).startswith(host_path + ":") and d != host_path
        ]
        if not svc["devices"]:
            svc.pop("devices", None)
    return _dump(data)


def set_privileged(yaml_str: str, service: str, value: bool) -> str:
    data = _load(yaml_str)
    svc = _get_service(data, service)
    if value:
        svc["privileged"] = True
    else:
        svc.pop("privileged", None)
    return _dump(data)


def set_resources(
    yaml_str: str,
    service: str,
    *,
    cpus: float | None = None,
    memory: str | None = None,
    cpus_reserve: float | None = None,
    memory_reserve: str | None = None,
) -> str:
    """Set deploy.resources limits/reservations for a service.

    Pass None / empty to clear an individual value.
    """
    data = _load(yaml_str)
    svc = _get_service(data, service)
    deploy = svc.setdefault("deploy", {})
    resources = deploy.setdefault("resources", {})

    limits = resources.get("limits", {})
    if cpus:
        limits["cpus"] = str(cpus)
    else:
        limits.pop("cpus", None)
    if memory:
        limits["memory"] = memory
    else:
        limits.pop("memory", None)
    if limits:
        resources["limits"] = limits
    else:
        resources.pop("limits", None)

    reservations = resources.get("reservations", {})
    if cpus_reserve:
        reservations["cpus"] = str(cpus_reserve)
    else:
        reservations.pop("cpus", None)
    if memory_reserve:
        reservations["memory"] = memory_reserve
    else:
        reservations.pop("memory", None)
    # Don't drop reservations entirely — it may hold GPU devices.
    if reservations:
        resources["reservations"] = reservations
    elif "reservations" in resources and not resources["reservations"]:
        resources.pop("reservations", None)

    if not resources:
        deploy.pop("resources", None)
    if not deploy:
        svc.pop("deploy", None)
    return _dump(data)


def list_services(yaml_str: str) -> list[str]:
    data = _load(yaml_str)
    return list(data["services"].keys())