pynixd-nixkube: conform to pynixd stores dict API

Author: Lillecarl

kubenix/pynixd.nix

@@ -34,23 +34,24 @@ let
       config file mounted in the pynixd pod. Corresponds to the PynixdSettings
       pydantic model (see pynixd.config).
 
-      Common keys include stores (list of StoreSpec), ranking weights,
-      GC intervals, etc. When stores include SSH stores, their client
-      keys are auto-discovered from HOME/.ssh/ if client_keys is omitted.
+      Common keys include stores (dict of StoreSpec keyed by store ID),
+      ranking weights, GC intervals, etc. When stores include SSH stores,
+      their client keys are auto-discovered from HOME/.ssh/ if client_keys
+      is omitted.
     '';
     type = jsonFormat.type;
     default = { };
     example = lib.literalExpression ''
       {
-        stores = [
-          {
+        stores = {
+          builder1 = {
             type = "ssh-subprocess";
             host = "builder.example.com";
             port = 22;
             username = "nix";
             systems = [ "x86_64-linux" ];
-          }
-        ];
+          };
+        };
       }
     '';
   };

pkgs/pynixd-nixkube/pynixd_nixkube/__main__.py

@@ -4,9 +4,10 @@
 from pathlib import Path
 
 import structlog
-from pynixd.config import PynixdSettings
+from pynixd.config import LocalSocketStoreSpec, PynixdSettings
 from pynixd.instance import Server
 from pynixd.store import LocalSocketStore
+from pynixd.types.ids import StoreId
 
 from .setup import configure_logging, install_nss
 from .ssh_keys import watch_authorized_keys
@@ -20,18 +21,19 @@ async def _main():
     install_nss()
 
     local_store = LocalSocketStore(
-        store_id="local",
-        store_path=Path("/"),
-        use_db=False,
-        monitor=False,
+        LocalSocketStoreSpec(
+            store_id=StoreId("local"),
+            store_path=Path("/"),
+            use_db=False,
+            monitor=False,
+        )
     )
 
     settings = PynixdSettings()
 
-    server = Server(local_store=local_store, settings=settings)
+    server = Server(stores={StoreId("local"): local_store}, settings=settings)
 
     async with server:
-        await server.add_store(local_store)
         keys_watch = asyncio.create_task(watch_authorized_keys(server))
         server.background_tasks.append(keys_watch)
         log.info("pynixd_nixkube_running")

pkgs/pynixd-nixkube/pynixd_nixkube/builder_main.py

@@ -4,9 +4,10 @@
 from pathlib import Path
 
 import structlog
-from pynixd.config import PynixdSettings
+from pynixd.config import LocalSocketStoreSpec, PynixdSettings
 from pynixd.instance import Server
 from pynixd.store import LocalSocketStore
+from pynixd.types.ids import StoreId
 
 from .setup import configure_logging, install_nss
 
@@ -19,15 +20,18 @@ async def _main() -> None:
     install_nss()
 
     local_store = LocalSocketStore(
-        store_id="local",
-        store_path=Path("/"),
-        use_db=False,
-        monitor=False,
+        LocalSocketStoreSpec(
+            store_id=StoreId("local"),
+            store_path=Path("/"),
+            use_db=False,
+            monitor=False,
+            extra_args=["--option", "build-dir", "/nix/var/nix/builds"],
+        )
     )
 
     settings = PynixdSettings()
 
-    server = Server(local_store=local_store, settings=settings)
+    server = Server(stores={StoreId("local"): local_store}, settings=settings)
 
     async with server:
         log.info("pynixd_nixkube_builder_running")

pkgs/pynixd-nixkube/pynixd_nixkube/builder_manager.py

@@ -9,7 +9,9 @@
 import kr8s.asyncio as k8s
 import structlog
 from kr8s.asyncio.objects import Job, PodTemplate
+from pynixd.config import SSHSubprocessStoreSpec
 from pynixd.store import SSHSubprocessStore
+from pynixd.types.ids import StoreId
 
 if TYPE_CHECKING:
     from pynixd.instance import Server
@@ -371,14 +373,16 @@ async def _register_builder(
         self, store_id: str, pod_ip: str, job_name: str = "", probe: bool = False
     ) -> None:
         store = SSHSubprocessStore(
-            host=pod_ip,
-            store_id=store_id,
-            port=22,
-            username="nix",
-            client_keys=["/etc/ssh-key/id_ed25519"],
-            nix_bin="/nix/var/result/bin/nix",
-            monitor=False,
-            no_schedule=probe,
+            SSHSubprocessStoreSpec(
+                store_id=StoreId(store_id),
+                host=pod_ip,
+                port=22,
+                username="nix",
+                client_keys=["/etc/ssh-key/id_ed25519"],
+                nix_bin="/nix/var/result/bin/nix",
+                monitor=False,
+                no_schedule=probe,
+            )
         )
         try:
             await self.server.add_store(store, dynamic=True)

pkgs/pynixd-nixkube/pynixd_nixkube/central_main.py

@@ -4,9 +4,10 @@
 from pathlib import Path
 
 import structlog
-from pynixd.config import PynixdSettings
+from pynixd.config import LocalSocketStoreSpec, PynixdSettings
 from pynixd.instance import Server
-from pynixd.store import LocalSocketStore
+from pynixd.store import LocalSocketStore, Store
+from pynixd.types.ids import StoreId
 
 from .builder_manager import BuilderManager
 from .config import NixkubeCentralSettings
@@ -22,30 +23,32 @@ async def _main() -> None:
     install_nss()
 
     local_store = LocalSocketStore(
-        store_id="local",
-        store_path=Path("/"),
-        use_db=False,
-        monitor=False,
+        LocalSocketStoreSpec(
+            store_id=StoreId("local"),
+            store_path=Path("/data"),
+            use_db=False,
+            monitor=False,
+        )
     )
 
     pynixd_settings = PynixdSettings()
     settings = NixkubeCentralSettings()
 
-    server = Server(local_store=local_store, settings=pynixd_settings)
-
+    stores: dict[StoreId, Store] = {StoreId("local"): local_store}
     if pynixd_settings.stores:
-        _, remote_stores = pynixd_settings.to_stores()
-        server.ctx._stores.update(remote_stores)
+        config_stores = pynixd_settings.to_stores()
+        for store_id, store in config_stores.items():
+            if store_id != StoreId("local"):
+                stores[store_id] = store
+
+    server = Server(stores=stores, settings=pynixd_settings)
 
     builder_manager: BuilderManager | None = None
 
     async with server:
         keys_watch = asyncio.create_task(watch_authorized_keys(server))
         server.background_tasks.append(keys_watch)
 
-        if pynixd_settings.schedule_mode != "scheduler":
-            await server.add_store(local_store)
-
         if settings.kube_namespace:
             builder_manager = BuilderManager(
                 server=server,