daggerml.contrib.docker_gc

  1"""Conservative Docker garbage collection and optional user-systemd scheduling."""
  2
  3import json
  4import logging
  5import os
  6import shutil
  7import subprocess
  8import sys
  9from datetime import datetime, timedelta, timezone
 10from pathlib import Path
 11from typing import Any
 12
 13from daggerml.api import DmlRepoError
 14from daggerml.contrib.executors.docker import DockerExecutor, _image_lock
 15
 16_OWNED_REPOSITORIES = {"daggerml-archive", "daggerml-execution"}
 17logger = logging.getLogger(__name__)
 18
 19
 20def _owned_tag(tag: str) -> bool:
 21    tag = tag.removeprefix("docker.io/library/").removeprefix("index.docker.io/library/")
 22    return tag.partition(":")[0] in _OWNED_REPOSITORIES
 23
 24
 25def _timestamp(value: str) -> datetime:
 26    return datetime.fromisoformat(value.replace("Z", "+00:00"))
 27
 28
 29def _run_if_present(*args: str) -> str | None:
 30    """Ignore resources removed by Docker or other clients since enumeration."""
 31    try:
 32        return DockerExecutor._run_docker(*args)
 33    except DmlRepoError as exc:
 34        if any(marker in str(exc).lower() for marker in (f"no such {args[0]}:", "no such object:")):
 35            return None
 36        raise
 37
 38
 39def collect(*, all_resources: bool = False, older_than_hours: float = 24, dry_run: bool = False) -> dict[str, Any]:
 40    """Collect old stopped containers and unreferenced image tags.
 41
 42    Parameters
 43    ----------
 44    all_resources : bool
 45        Include foreign resources, rather than only DML labels and tag prefixes.
 46    older_than_hours : float
 47        Minimum time since container exit/creation or image creation.
 48    dry_run : bool
 49        Report candidates without deleting them.
 50
 51    Returns
 52    -------
 53    dict
 54        Container IDs and image references selected for removal.
 55    """
 56    if not 0 < older_than_hours < float("inf"):
 57        raise ValueError("older-than-hours must be a positive finite number")
 58    cutoff = datetime.now(timezone.utc) - timedelta(hours=older_than_hours)
 59    run = DockerExecutor._run_docker
 60    removed: dict[str, Any] = {"containers": [], "images": [], "dry_run": dry_run}
 61    with _image_lock():
 62        container_ids = run("container", "ls", "-aq", "--no-trunc").split()
 63        containers = []
 64        for ident in container_ids:
 65            output = _run_if_present("container", "inspect", ident)
 66            if output is not None:
 67                containers.append(json.loads(output)[0])
 68        retained = []
 69        for container in containers:
 70            state = container["State"]
 71            owned = (container.get("Config", {}).get("Labels") or {}).get("daggerml.managed") == "true"
 72            finished = state.get("FinishedAt") or container["Created"]
 73            if finished.startswith("0001-"):
 74                finished = container["Created"]
 75            if (state.get("Status") in {"exited", "dead", "created"}
 76                    and (all_resources or owned)
 77                    and max(_timestamp(finished), _timestamp(container["Created"])) < cutoff):
 78                if not dry_run:
 79                    if _run_if_present("container", "rm", container["Id"]) is None:
 80                        continue
 81                removed["containers"].append(container["Id"])
 82            else:
 83                retained.append(container)
 84        referenced = {container["Image"] for container in retained}
 85        image_ids = set(run("image", "ls", "-aq", "--no-trunc").split())
 86        pending = []
 87        for ident in sorted(image_ids):
 88            output = _run_if_present("image", "inspect", ident)
 89            if output is None:
 90                continue
 91            image = json.loads(output)[0]
 92            if image["Id"] in referenced or _timestamp(image["Created"]) >= cutoff:
 93                continue
 94            tags = image.get("RepoTags") or []
 95            refs = tags if all_resources else [tag for tag in tags if _owned_tag(tag)]
 96            if all_resources and not refs:
 97                refs = [image["Id"]]
 98            pending.extend(refs)
 99        while pending:
100            deferred = []
101            for ref in pending:
102                try:
103                    # Never force deletion. Docker remains the final reference check.
104                    if not dry_run:
105                        if _run_if_present("image", "rm", "--no-prune", ref) is None:
106                            continue
107                except DmlRepoError as exc:
108                    if "image has dependent child images" not in str(exc).lower():
109                        raise
110                    deferred.append(ref)
111                    continue
112                removed["images"].append(ref)
113            if len(deferred) == len(pending):
114                logger.warning("Retaining images with dependent children: %s", ", ".join(deferred))
115                break
116            pending = deferred
117    return removed
118
119
120def _systemd_quote(value: str, *, exec_argument: bool = True) -> str:
121    if "\n" in value or "\r" in value or "\x00" in value:
122        raise ValueError("systemd values must be single-line strings")
123    value = value.replace("\\", "\\\\").replace('"', '\\"').replace("%", "%%")
124    if exec_argument:
125        value = value.replace("$", "$$")
126    return '"' + value + '"'
127
128
129def install(*, all_resources: bool = False, older_than_hours: float = 24) -> Path:
130    """Install and enable an hourly systemd user timer for Docker GC.
131
132    Parameters
133    ----------
134    all_resources : bool
135        Persist host-wide collection in the installed service.
136    older_than_hours : float
137        Minimum resource age passed to each collection.
138
139    Returns
140    -------
141    Path
142        Directory containing the installed service and timer.
143    """
144    if not sys.platform.startswith("linux") or not shutil.which("systemctl"):
145        raise DmlRepoError(
146            "Automatic installation requires Linux and systemd; schedule 'dml contrib docker-gc' manually"
147        )
148    if not 0 < older_than_hours < float("inf"):
149        raise ValueError("older-than-hours must be a positive finite number")
150    command = [sys.executable, "-m", "daggerml.contrib.docker_gc", "--older-than-hours", str(older_than_hours)]
151    if all_resources:
152        command.append("--all")
153    # User services need the same Docker executable, cache lock, and daemon
154    # selection as the interactive installer. No shell interpolation is used.
155    environment = {"PATH": os.environ.get("PATH", os.defpath)}
156    unset_environment = []
157    for name in ("XDG_CACHE_HOME", "DOCKER_HOST", "DOCKER_CONTEXT", "DOCKER_CONFIG", "DOCKER_TLS_VERIFY",
158                 "DOCKER_CERT_PATH"):
159        if name in os.environ:
160            environment[name] = os.environ[name]
161        else:
162            # User services inherit the manager's environment, which can contain
163            # stale Docker or cache settings absent from this shell.
164            unset_environment.append(name)
165    service = ("[Unit]\nDescription=DaggerML Docker garbage collection\n\n[Service]\nType=oneshot\n"
166               + "".join(f"Environment={_systemd_quote(f'{key}={value}', exec_argument=False)}\n"
167                          for key, value in environment.items())
168               + ("UnsetEnvironment=" + " ".join(unset_environment) + "\n" if unset_environment else "")
169               + "ExecStart=" + " ".join(_systemd_quote(arg) for arg in command) + "\n")
170    timer = ("[Unit]\nDescription=Hourly DaggerML Docker garbage collection\n\n[Timer]\n"
171             "OnCalendar=hourly\nPersistent=true\n\n[Install]\nWantedBy=timers.target\n")
172    directory = Path(os.getenv("XDG_CONFIG_HOME") or Path.home() / ".config") / "systemd" / "user"
173    directory.mkdir(parents=True, exist_ok=True)
174    (directory / "daggerml-docker-gc.service").write_text(service, encoding="utf-8")
175    (directory / "daggerml-docker-gc.timer").write_text(timer, encoding="utf-8")
176    subprocess.run(["systemctl", "--user", "daemon-reload"], check=True)
177    subprocess.run(["systemctl", "--user", "enable", "--now", "daggerml-docker-gc.timer"], check=True)
178    return directory
179
180
181def main() -> None:
182    """Run the standalone collector used by the installed timer."""
183    from daggerml.contrib.status import cli
184
185    raise SystemExit(cli(["docker-gc", *sys.argv[1:]]))
186
187
188if __name__ == "__main__":
189    main()