feat: Implement dataset upgrade functionality

packages/enclave-model-api-example/docker/inference_server.py

@@ -9,7 +9,11 @@
 from attestation_server import app
 from syft_enclaves.settings import EnclaveSettings
 
-from enclave_model_api.paths import default_syftbox_folder, private_dataset_dir
+from enclave_model_api.paths import (
+    default_syftbox_folder,
+    private_dataset_dir,
+    resolve_weights_dir,
+)
 from enclave_model_api.server import build_router
 from enclave_model_api.service import InferenceService
 from enclave_model_api.settings import InferenceSettings
@@ -32,9 +36,13 @@
 service = InferenceService(
     backend=backend,
     model_size=inference.model_size,
-    weights_dir=private_dataset_dir(
+    # Resolved on each poll: the model owner writes the layout its own release
+    # decided, and this process starts before the weights arrive.
+    weights_dir=lambda: resolve_weights_dir(
         syftbox_folder, inference.model_owner, inference.model_dataset
     ),
+    # The enclave owns its logs dataset and writes it, so the current layout is
+    # the right answer while it is absent.
     logs_dir=private_dataset_dir(
         syftbox_folder, settings.email, inference.logs_dataset
     ),

packages/enclave-model-api-example/scripts/local_server.py

@@ -13,7 +13,11 @@
 import uvicorn
 from syft_enclaves.settings import EnclaveSettings
 
-from enclave_model_api.paths import default_syftbox_folder, private_dataset_dir
+from enclave_model_api.paths import (
+    default_syftbox_folder,
+    private_dataset_dir,
+    resolve_weights_dir,
+)
 from enclave_model_api.server import create_app
 from enclave_model_api.service import InferenceService
 from enclave_model_api.settings import InferenceSettings
@@ -33,7 +37,11 @@
 service = InferenceService(
     backend=backend,
     model_size=inf.model_size,
-    weights_dir=private_dataset_dir(folder, inf.model_owner, inf.model_dataset),
+    # Resolved on each poll: the model owner writes the layout its own release
+    # decided, and this process starts before the weights arrive.
+    weights_dir=lambda: resolve_weights_dir(folder, inf.model_owner, inf.model_dataset),
+    # The enclave owns its logs dataset and writes it, so the current layout is
+    # the right answer while it is absent.
     logs_dir=private_dataset_dir(folder, settings.email, inf.logs_dataset),
 )
 service.start_polling()

packages/enclave-model-api-example/src/enclave_model_api/paths.py

@@ -18,21 +18,56 @@ def default_syftbox_folder(email: str) -> Path:
     return get_jupyter_default_syftbox_folder(email)
 
 
-def resolve_private_dataset_dir(storage: DatasetStorage, owner: str, name: str) -> Path:
-    """Private dir at the dataset's actual on-disk protocol layout.
+def candidate_private_dataset_dirs(
+    storage: DatasetStorage, owner: str, name: str
+) -> list[Path]:
+    """The private dirs a dataset could occupy, newest layout first.
 
-    A dataset may live at protocol 0 (flat) or under a ``v<n>`` segment; the
-    written layout depends on what the audience can read, not the current
-    default. Datasets not yet on disk (e.g. weights still syncing) fall back to
-    the widest-compatible protocol — where a peer running any current release
-    writes them for us.
+    One entry, the layout on disk, once the dataset is there. While it is
+    absent, one entry for each protocol layout this client reads: the owner
+    writes the layout its own release and audience decided, and a dataset that
+    has not arrived cannot tell us which that is. A reader that waits must
+    therefore watch them all.
     """
     try:
         ref = storage.find_dataset_ref(owner, name)
     except DatasetNotFoundError:
-        (widest,) = storage.target_protocol_versions_for_peers(None)
-        ref = DatasetRef(owner=owner, name=name, protocol_version=widest)
-    return storage.private_dataset_dir(ref)
+        return [
+            storage.private_dataset_dir(
+                DatasetRef(owner=owner, name=name, protocol_version=protocol_version)
+            )
+            for protocol_version in storage.supported_protocol_versions
+        ]
+    return [storage.private_dataset_dir(ref)]
+
+
+def resolve_private_dataset_dir(storage: DatasetStorage, owner: str, name: str) -> Path:
+    """Private dir at the dataset's on-disk protocol layout.
+
+    The current layout while the dataset is absent, which is the layout this
+    client writes for a dataset of its own. A reader waiting for a dataset that
+    another datasite writes must use ``candidate_private_dataset_dirs``, because
+    that owner may write an older layout.
+    """
+    return candidate_private_dataset_dirs(storage, owner, name)[0]
+
+
+def resolve_weights_dir(
+    syftbox_folder: Path | str, datasite: str, dataset_name: str
+) -> Path:
+    """The layout that holds the synced weights, or the newest candidate so far.
+
+    Re-resolved on each poll, not fixed at startup: the weights arrive in the
+    layout their owner writes, and an owner on an earlier release writes the
+    flat one.
+    """
+    config = SyftBoxConfig(syftbox_folder=Path(syftbox_folder), email=datasite)
+    storage = DatasetStorage(config=config)
+    candidates = candidate_private_dataset_dirs(storage, datasite, dataset_name)
+    for path in candidates:
+        if weights_ready(path):
+            return path
+    return candidates[0]
 
 
 def private_dataset_dir(

packages/enclave-model-api-example/src/enclave_model_api/service.py

@@ -9,6 +9,7 @@
 import threading
 import time
 from pathlib import Path
+from typing import Callable
 
 from enclave_model_api.log_writer import append_log_record, build_log_record
 from enclave_model_api.paths import weights_ready
@@ -21,16 +22,25 @@ def __init__(
         self,
         backend,
         model_size: str,
-        weights_dir: Path | str,
+        weights_dir: Path | str | Callable[[], Path],
         logs_dir: Path | str,
     ):
         self.backend = backend
         self.model_size = model_size
-        self.weights_dir = Path(weights_dir)
+        # A callable is re-resolved on each read. The weights arrive in the
+        # layout their owner writes, so a path fixed at startup can watch a
+        # layout the owner never writes, and the poll then never ends.
+        self._weights_dir = (
+            weights_dir if callable(weights_dir) else (lambda: Path(weights_dir))
+        )
         self.logs_dir = Path(logs_dir)
         self._loaded = None
         self._lock = threading.Lock()
 
+    @property
+    def weights_dir(self) -> Path:
+        return Path(self._weights_dir())
+
     @property
     def loaded(self) -> bool:
         return self._loaded is not None

packages/enclave-model-api-example/tests/test_inference_server.py

@@ -1,6 +1,8 @@
 """Unit tests for the inference service pieces: paths, log writer, FastAPI app."""
 
+import ast
 import json
+from pathlib import Path
 
 from fastapi.testclient import TestClient
 
@@ -9,9 +11,17 @@
     append_log_record,
     build_log_record,
 )
-from enclave_model_api.paths import private_dataset_dir, weights_ready
+from enclave_model_api.paths import (
+    candidate_private_dataset_dirs,
+    private_dataset_dir,
+    resolve_weights_dir,
+    weights_ready,
+)
 from enclave_model_api.server import create_app
 from enclave_model_api.service import InferenceService
+from syft_datasets.config import SyftBoxConfig
+from syft_datasets.dataset_storage import DatasetStorage
+from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION
 
 from inference_stub import STUB_COMPLETION_PREFIX, StubBackend, make_stub_weights
 
@@ -44,16 +54,46 @@ def test_weights_ready(tmp_path):
 
 
 def test_private_dataset_dir_layout(tmp_path):
+    # A dataset not yet on disk falls back to the current protocol, which is the
+    # layout this client writes for a dataset of its own (the enclave's logs).
+    # A dataset another datasite writes may arrive in an older layout, and
+    # candidate_private_dataset_dirs covers that case.
     path = private_dataset_dir(tmp_path, "enclave@openmined.org", "inference_logs")
     assert path == (
         tmp_path
         / "enclave@openmined.org"
         / "private"
         / "syft_datasets"
+        / f"v{DATASET_PROTOCOL_VERSION}"
         / "inference_logs"
     )
 
 
+def test_private_dataset_dir_follows_the_layout_on_disk(tmp_path):
+    # A dataset already on disk decides its own layout, whatever the fallback is.
+    flat = (
+        tmp_path
+        / "enclave@openmined.org"
+        / "private"
+        / "syft_datasets"
+        / "inference_logs"
+    )
+    flat.mkdir(parents=True)
+    (flat / "private_metadata.yaml").write_text("uid: x\n")
+    public = (
+        tmp_path
+        / "enclave@openmined.org"
+        / "public"
+        / "syft_datasets"
+        / "inference_logs"
+    )
+    public.mkdir(parents=True)
+    (public / "dataset.yaml").write_text("name: inference_logs\n")
+
+    path = private_dataset_dir(tmp_path, "enclave@openmined.org", "inference_logs")
+    assert path == flat
+
+
 def test_inference_server_full_lifecycle(tmp_path):
     """503 before weights → load after weights sync → /infer logs each request."""
     weights_dir = tmp_path / "weights"
@@ -96,3 +136,129 @@ def test_inference_server_full_lifecycle(tmp_path):
     assert len(records) == 1
     assert records[0]["prompt"] == "What is the capital of NL?"
     assert records[0]["completion"] == body["completion"]
+
+
+def _storage(tmp_path, email: str) -> DatasetStorage:
+    return DatasetStorage(config=SyftBoxConfig(syftbox_folder=tmp_path, email=email))
+
+
+def _write_weights(private_dir):
+    private_dir.mkdir(parents=True, exist_ok=True)
+    make_stub_weights(private_dir)
+
+
+def test_candidate_dirs_cover_every_layout_while_the_dataset_is_absent(tmp_path):
+    # The owner writes the layout its own release decided, so a reader that is
+    # still waiting cannot know which one arrives.
+    storage = _storage(tmp_path, "owner@test.org")
+    candidates = candidate_private_dataset_dirs(storage, "owner@test.org", "weights")
+
+    assert [c.name for c in candidates] == ["weights"] * len(candidates)
+    segments = [c.parent.name for c in candidates]
+    # Newest layout first, down to the floor.
+    assert segments[0] == f"v{DATASET_PROTOCOL_VERSION}"
+    assert "syft_datasets" in segments
+
+
+def test_candidate_dirs_follow_the_layout_on_disk(tmp_path):
+    storage = _storage(tmp_path, "owner@test.org")
+    flat_private = tmp_path / "owner@test.org" / "private" / "syft_datasets" / "weights"
+    flat_private.mkdir(parents=True)
+    (flat_private / "private_metadata.yaml").write_text("uid: x\n")
+    flat_public = tmp_path / "owner@test.org" / "public" / "syft_datasets" / "weights"
+    flat_public.mkdir(parents=True)
+    (flat_public / "dataset.yaml").write_text("name: weights\n")
+
+    assert candidate_private_dataset_dirs(storage, "owner@test.org", "weights") == [
+        flat_private
+    ]
+
+
+def test_resolve_weights_dir_finds_the_flat_layout_of_an_earlier_release(tmp_path):
+    # Every data owner in the fleet today writes protocol 0, so the weights land
+    # flat. A guess fixed at the current layout never sees them.
+    owner = "owner@test.org"
+    flat = tmp_path / owner / "private" / "syft_datasets" / "weights"
+    _write_weights(flat)
+
+    assert resolve_weights_dir(tmp_path, owner, "weights") == flat
+    assert weights_ready(resolve_weights_dir(tmp_path, owner, "weights"))
+
+
+def test_resolve_weights_dir_returns_the_newest_candidate_while_absent(tmp_path):
+    owner = "owner@test.org"
+    resolved = resolve_weights_dir(tmp_path, owner, "weights")
+    assert resolved.parent.name == f"v{DATASET_PROTOCOL_VERSION}"
+    assert not weights_ready(resolved)
+
+
+def test_the_service_sees_weights_that_land_in_an_older_layout(tmp_path):
+    # The poll exists so the enclave need not restart. It re-resolves the
+    # layout, so weights arriving flat after startup are picked up.
+    owner = "owner@test.org"
+    service = InferenceService(
+        backend=StubBackend(),
+        model_size="270m",
+        weights_dir=lambda: resolve_weights_dir(tmp_path, owner, "weights"),
+        logs_dir=tmp_path / "logs",
+    )
+    assert not service.weights_present
+
+    _write_weights(tmp_path / owner / "private" / "syft_datasets" / "weights")
+
+    assert service.weights_present
+    assert service.try_load()
+
+
+# --- entrypoint wiring ----------------------------------------------------
+#
+# Neither entrypoint can be imported by a test: each starts a server at import
+# time. The invariant they must hold is checked on their syntax tree instead.
+# A path fixed at startup defeats the poll, and both entrypoints have to pass a
+# callable that re-resolves the layout.
+
+
+def _entrypoints():
+    root = Path(__file__).resolve().parents[1]
+    return sorted((root / "scripts").glob("*.py")) + sorted(
+        (root / "docker").glob("inference_server.py")
+    )
+
+
+def _weights_argument(path: Path):
+    """The ``weights_dir`` argument of the InferenceService call in a module."""
+    tree = ast.parse(path.read_text())
+    for node in ast.walk(tree):
+        if (
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Name)
+            and node.func.id == "InferenceService"
+        ):
+            for keyword in node.keywords:
+                if keyword.arg == "weights_dir":
+                    return keyword.value
+    return None
+
+
+def test_every_entrypoint_re_resolves_the_weights_layout():
+    checked = []
+    for path in _entrypoints():
+        argument = _weights_argument(path)
+        if argument is None:
+            continue
+        checked.append(path.name)
+        assert isinstance(argument, ast.Lambda), (
+            f"{path.name} fixes the weights path at startup; the owner may write "
+            "an older layout, and the poll would never see it"
+        )
+        called = {
+            n.func.id
+            for n in ast.walk(argument)
+            if isinstance(n, ast.Call) and isinstance(n.func, ast.Name)
+        }
+        assert "resolve_weights_dir" in called, (
+            f"{path.name} must resolve the weights layout through "
+            "resolve_weights_dir, which probes every layout it supports"
+        )
+    # Both entrypoints, or the glob stopped matching.
+    assert sorted(checked) == ["inference_server.py", "local_server.py"]