Load docker-sbx template images into sbx

docs/deploy.md

@@ -198,6 +198,9 @@ docker save drukbox/sbx-sandbox:latest -o /tmp/sbx-sandbox.tar
 sbx template load /tmp/sbx-sandbox.tar
 ```
 
+drukbox does these steps for each template that `POST /templates`
+builds, and then removes the Docker image.
+
 A sandbox creation takes approximately 20 seconds with a warm template
 cache, and more than 30 seconds at the first pull. Thus a warm pool
 (`POOL_SIZES`) is useful. Each sandbox gets the explicit

src/providers/docker/api.py

@@ -1,4 +1,5 @@
 import io
+from pathlib import Path
 
 import aiodocker
 import aiohttp
@@ -108,6 +109,20 @@ async def remove_image(self, image: str) -> None:
         except aiohttp.ClientError as exc:
             raise DockerTransportError(str(exc)) from exc
 
+    async def save_image(self, image: str, archive: Path) -> None:
+        # An image can hold gigabytes, so it streams to disk and not to memory.
+        try:
+            async with self._get_client().images.export_image(image) as export:
+                with archive.open("wb") as file:
+                    while chunk := await export.read(1 << 20):
+                        file.write(chunk)
+        except aiodocker.DockerError as exc:
+            if exc.status == 404:
+                raise DockerImageNotFoundError(str(exc)) from exc
+            raise DockerTransportError(_detail(exc)) from exc
+        except aiohttp.ClientError as exc:
+            raise DockerTransportError(str(exc)) from exc
+
     async def push_image(
         self,
         image: str,

src/providers/docker/tests/test_api.py

@@ -137,6 +137,18 @@ async def test_missing_image_maps_to_not_found() -> None:
         await _api(fake).remove_image("drukbox-template:missing")
 
 
+async def test_save_image_streams_the_archive_to_the_file(tmp_path) -> None:
+    fake = _fake_docker()
+    fake.images.export_image = MagicMock()
+    export = fake.images.export_image.return_value.__aenter__.return_value
+    export.read = AsyncMock(side_effect=[b"layer-", b"data", b""])
+
+    await _api(fake).save_image("drukbox-template:123456789abc", tmp_path / "image.tar")
+
+    fake.images.export_image.assert_called_once_with("drukbox-template:123456789abc")
+    assert (tmp_path / "image.tar").read_bytes() == b"layer-data"
+
+
 async def test_push_image_sends_per_call_credentials() -> None:
     fake = _fake_docker()
 

src/providers/docker_sbx/api.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import re
+from pathlib import Path
 
 from .exceptions import DockerSbxNotFoundError, DockerSbxTransportError
 
@@ -9,7 +10,7 @@
 
 # "credentials not found" must stay a transport error, or delete_vm takes a
 # live sandbox for a removed one.
-_SANDBOX_NOT_FOUND_RE = re.compile(r"sandbox '[^']*' not found")
+_NOT_FOUND_RE = re.compile(r"sandbox '[^']*' not found|no image \"[^\"]*\"")
 
 
 class SbxCLI:
@@ -119,6 +120,12 @@ async def remove_secret(self, service: str, *, sandbox: str) -> None:
     async def remove_custom_secret(self, *, sandbox: str, placeholder: str) -> None:
         await self._run("secret", "rm", "-f", "--placeholder", placeholder, "--sandbox", sandbox)
 
+    async def load_template(self, archive: Path) -> None:
+        await self._run("template", "load", str(archive))
+
+    async def remove_template(self, template: str) -> None:
+        await self._run("template", "rm", "--force", template)
+
     async def sandbox_count(self) -> int:
         output = await self._run("ls", "--json")
         try:
@@ -154,7 +161,7 @@ async def _run(self, *args: str, stdin: str | None = None) -> str:
             ) from error
         if process.returncode != 0:
             detail = stderr.decode().strip() or f"sbx {args[0]} exited {process.returncode}"
-            if _SANDBOX_NOT_FOUND_RE.search(detail):
+            if _NOT_FOUND_RE.search(detail):
                 raise DockerSbxNotFoundError(detail)
             raise DockerSbxTransportError(detail)
         return stdout.decode()

src/providers/docker_sbx/exceptions.py

@@ -3,7 +3,7 @@ class DockerSbxProviderError(RuntimeError):
 
 
 class DockerSbxNotFoundError(DockerSbxProviderError):
-    """The sandbox was not found."""
+    """The sandbox or the template was not found."""
 
 
 class DockerSbxTransportError(DockerSbxProviderError):

src/providers/docker_sbx/provider.py

@@ -1,6 +1,7 @@
 import contextlib
 import shlex
 import shutil
+import tempfile
 from pathlib import Path
 from typing import ClassVar, Self
 
@@ -10,7 +11,8 @@
 from providers.base import VMCreateResult, VMProvider
 from providers.capabilities import TemplateCapability
 from providers.docker.api import DockerAPI
-from providers.docker.images import build_derived_image, remove_derived_image
+from providers.docker.exceptions import DockerProviderError
+from providers.docker.images import build_derived_image
 from providers.exceptions import (
     ProviderCommandError,
     ProviderNotFoundError,
@@ -25,20 +27,6 @@
 from .settings import DockerSbxSettings
 
 
-def _bootstrap_script(*, public_key: str, env: dict[str, str], ssh_username: str) -> str:
-    home = "/root" if ssh_username == "root" else f"/home/{ssh_username}"
-    owner = shlex.quote(ssh_username)
-    lines = [
-        "set -euo pipefail",
-        f"install -d -m 700 -o {owner} -g {owner} {home}/.ssh",
-        f"printf '%s\\n' {shlex.quote(public_key)} > {home}/.ssh/authorized_keys",
-        f"chmod 600 {home}/.ssh/authorized_keys",
-        f"chown {owner}:{owner} {home}/.ssh/authorized_keys",
-    ]
-    # The runtime takes no environment at create time. pam_env reads this file.
-    return "\n".join([*lines, *environment.get_persist(env)]) + "\n"
-
-
 class DockerSbxProvider(VMProvider, TemplateCapability):
     name: ClassVar[str] = "docker-sbx"
     diagnose_hint: ClassVar[str] = "check_sandboxd_is_running_and_logged_in"
@@ -122,12 +110,22 @@ async def create_vm(
             self._remove_sandbox_files(name)
             raise ProviderTransportError(str(exc)) from exc
 
+        username = self.settings.ssh_username
+        home = "/root" if username == "root" else f"/home/{username}"
+        owner = shlex.quote(username)
+        script = "\n".join(
+            [
+                "set -euo pipefail",
+                f"install -d -m 700 -o {owner} -g {owner} {home}/.ssh",
+                f"printf '%s\\n' {shlex.quote(public_key)} > {home}/.ssh/authorized_keys",
+                f"chmod 600 {home}/.ssh/authorized_keys",
+                f"chown {owner}:{owner} {home}/.ssh/authorized_keys",
+                # The runtime takes no environment at create time. pam_env reads this file.
+                *environment.get_persist(caller_env),
+                "",
+            ]
+        )
         try:
-            script = _bootstrap_script(
-                public_key=public_key,
-                env=caller_env,
-                ssh_username=self.settings.ssh_username,
-            )
             await self.api.run_bootstrap(name, script)
         except DockerSbxProviderError as exc:
             with contextlib.suppress(DockerSbxProviderError):
@@ -139,7 +137,7 @@ async def create_vm(
         return VMCreateResult(
             provider_id=name,
             name=name,
-            ssh_username=self.settings.ssh_username,
+            ssh_username=username,
             private_key=private_key,
             public_key=public_key,
         )
@@ -180,14 +178,29 @@ async def build_template_image(
         setup_script: str,
         label: str,
     ) -> str:
-        return await build_derived_image(
+        image = await build_derived_image(
             self.docker,
             base_image=base_image,
             setup_script=setup_script,
         )
+        # sbx keeps its own image store. The Docker image only carries the build.
+        try:
+            with tempfile.TemporaryDirectory() as directory:
+                archive = Path(directory) / "template.tar"
+                await self.docker.save_image(image, archive)
+                await self.api.load_template(archive)
+            await self.docker.remove_image(image)
+        except (OSError, DockerProviderError, DockerSbxProviderError) as exc:
+            raise ProviderTransportError(str(exc)) from exc
+        return image
 
     async def delete_template_image(self, image: str) -> None:
-        await remove_derived_image(self.docker, image)
+        try:
+            await self.api.remove_template(image)
+        except DockerSbxNotFoundError as exc:
+            raise ProviderNotFoundError(f"sbx template '{image}' was not found") from exc
+        except DockerSbxProviderError as exc:
+            raise ProviderTransportError(str(exc)) from exc
 
     async def diagnose(self) -> str:
         return f"sandboxd reachable, {await self.api.sandbox_count()} sandbox(es)"

src/providers/docker_sbx/tests/test_api.py

@@ -1,4 +1,5 @@
 import asyncio
+from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import AsyncMock
 
@@ -214,9 +215,44 @@ async def test_missing_sandbox_maps_to_not_found(monkeypatch):
         await SbxCLI().remove_sandbox("sb-test")
 
 
+@pytest.mark.asyncio
+async def test_template_load_and_removal_address_the_sbx_image_store(monkeypatch):
+    calls: list = []
+
+    async def fake_exec(*args, **kwargs):
+        calls.append(args)
+        return _process()
+
+    monkeypatch.setattr("providers.docker_sbx.api.asyncio.create_subprocess_exec", fake_exec)
+
+    await SbxCLI().load_template(Path("/tmp/template.tar"))
+    await SbxCLI().remove_template("drukbox-template:123456789abc")
+
+    assert calls == [
+        ("sbx", "template", "load", "/tmp/template.tar"),
+        # Without --force, the CLI asks for confirmation.
+        ("sbx", "template", "rm", "--force", "drukbox-template:123456789abc"),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_missing_template_maps_to_not_found(monkeypatch):
+    create = AsyncMock(
+        return_value=_process(
+            returncode=1,
+            stderr=b"error: delete template: image not found: request failed: 404 Not Found: "
+            b'no image "drukbox-template:missing"',
+        )
+    )
+    monkeypatch.setattr("providers.docker_sbx.api.asyncio.create_subprocess_exec", create)
+
+    with pytest.raises(DockerSbxNotFoundError):
+        await SbxCLI().remove_template("drukbox-template:missing")
+
+
 @pytest.mark.asyncio
 async def test_stderr_merely_containing_not_found_stays_a_transport_error(monkeypatch):
-    # Only the CLI message for a missing sandbox can map to not-found. If an
+    # Only the CLI messages for a missing sandbox or template map to not-found. If an
     # auth error maps to not-found, delete_vm removes the record and the
     # workspace of a live sandbox.
     create = AsyncMock(

src/providers/docker_sbx/tests/test_provider.py

@@ -4,7 +4,6 @@
 
 import pytest
 
-from providers.docker.exceptions import DockerImageNotFoundError
 from providers.docker_sbx.exceptions import (
     DockerSbxNotFoundError,
     DockerSbxTransportError,
@@ -29,13 +28,16 @@ def _api_mock() -> MagicMock:
     api.run_bootstrap = AsyncMock()
     api.remove_sandbox = AsyncMock()
     api.sandbox_count = AsyncMock(return_value=2)
+    api.load_template = AsyncMock()
+    api.remove_template = AsyncMock()
     return api
 
 
 def _docker_mock() -> MagicMock:
     docker = MagicMock()
     docker.build_image = AsyncMock()
     docker.remove_image = AsyncMock()
+    docker.save_image = AsyncMock()
     return docker
 
 
@@ -224,7 +226,7 @@ async def test_delete_vm_keeps_the_files_when_teardown_fails(tmp_path):
 
 
 @pytest.mark.asyncio
-async def test_build_template_image_builds_and_returns_the_local_image_tag(tmp_path):
+async def test_build_template_image_moves_the_built_image_into_sbx(tmp_path):
     api = _api_mock()
     docker = _docker_mock()
     provider = _provider(api, _settings(tmp_path), docker=docker)
@@ -237,25 +239,42 @@ async def test_build_template_image_builds_and_returns_the_local_image_tag(tmp_p
 
     assert image.startswith("drukbox-template:")
     assert docker.build_image.await_args.args[0] == image
+    saved, archive = docker.save_image.await_args.args
+    assert saved == image
+    api.load_template.assert_awaited_once_with(archive)
+    assert not archive.exists()
+    docker.remove_image.assert_awaited_once_with(image)
 
 
 @pytest.mark.asyncio
-async def test_delete_template_image_removes_the_local_image(tmp_path):
+async def test_build_template_image_translates_a_failed_load(tmp_path):
     api = _api_mock()
-    docker = _docker_mock()
-    provider = _provider(api, _settings(tmp_path), docker=docker)
+    api.load_template.side_effect = DockerSbxTransportError("sandboxd unavailable")
+    provider = _provider(api, _settings(tmp_path))
+
+    with pytest.raises(ProviderTransportError, match="sandboxd unavailable"):
+        await provider.build_template_image(
+            base_image="sandbox:base",
+            setup_script="apt-get update",
+            label="Node tools",
+        )
+
+
+@pytest.mark.asyncio
+async def test_delete_template_image_removes_the_sbx_template(tmp_path):
+    api = _api_mock()
+    provider = _provider(api, _settings(tmp_path))
 
     await provider.delete_template_image("drukbox-template:123456789abc")
 
-    docker.remove_image.assert_awaited_once_with("drukbox-template:123456789abc")
+    api.remove_template.assert_awaited_once_with("drukbox-template:123456789abc")
 
 
 @pytest.mark.asyncio
-async def test_delete_template_image_translates_a_missing_image(tmp_path):
+async def test_delete_template_image_translates_a_missing_template(tmp_path):
     api = _api_mock()
-    docker = _docker_mock()
-    docker.remove_image.side_effect = DockerImageNotFoundError("No such image")
-    provider = _provider(api, _settings(tmp_path), docker=docker)
+    api.remove_template.side_effect = DockerSbxNotFoundError('no image "drukbox-template:missing"')
+    provider = _provider(api, _settings(tmp_path))
 
     with pytest.raises(ProviderNotFoundError, match="was not found"):
         await provider.delete_template_image("drukbox-template:missing")