SF-12: manual coordinator sign-off before a model reaches labs

controller/TASK_GUIDE.md

@@ -713,7 +713,7 @@ For tests and the workbench, write a synthetic store with `python -m starfish.co
 - Training: `local_epochs`, `batch_size` default 1, `grad_accum`, `lr`, `weight_decay`, `amp` default true and used on CUDA only, `grad_checkpointing`, `clip_norm`, `device` default `auto` for CUDA then MPS then CPU, `seed`.
 - `curriculum`: loss weights by round, because FL rounds replace the paper's epochs. The last stage whose `from_round` has been reached applies.
 - Each round's log records train loss, local validation relative l2 and SSIM, seconds per epoch and peak memory. The coordinator refuses any delta trained from a different global model, for another round, or with another architecture, and any non-finite delta.
-- Model registry, SF-12: after a round whose candidate the gate accepted, the coordinator publishes the new global model to the router's registry as the next version for its bucket, such as `250k-v0003`, with the previous latest version as its parent and the round's eval report. Models from rounds with the gate disabled or rejected are never published. Any site can list them with `transfer.list_models(bucket_hz)`, get the newest with `transfer.latest_model(bucket_hz)`, and download one with its checksum verified with `transfer.download_model(entry, folder)`. On the router: `GET registry/?bucket_hz=`, `GET registry/latest/?bucket_hz=`, `GET registry/file/?version=`.
+- Model registry, SF-12: after a round whose candidate the gate accepted, the coordinator publishes the new global model to the router's registry as the next version for its bucket, such as `250k-v0003`, with the newest registered version as its parent and the round's eval report. Models from rounds with the gate disabled or rejected are never published. A published version is pending until the coordinator tests it by hand and signs it off with `POST registry/approve/` or `POST registry/reject/`, body `{"version": ..., "note": ...}`; the router records who decided and when, and a decision is final. Only the superuser and, with site tokens, the site that coordinates the project can sign off or see unapproved versions. Any site can list approved versions with `transfer.list_models(bucket_hz)`, get the newest with `transfer.latest_model(bucket_hz)`, and download one with its checksum verified with `transfer.download_model(entry, folder)`; reviewers pass `status='pending'`, `'rejected'` or `'any'`. On the router: `GET registry/?bucket_hz=&status=`, `GET registry/latest/?bucket_hz=&status=`, `GET registry/file/?version=`.
 - `compression`, SF-07: how each site's delta travels. `{"method": "none"}` by default; `fp16` halves it, `int8` quarters it, and `{"method": "topk", "k": 0.01}` sends only the largest 1% of each tensor's entries, with error feedback kept locally in the controller folder. Measure a method against `none` in E1 before using it for real runs. The global model is always sent in full.
 - `aggregation`, SF-11: `{"screen_factor": 5, "clip_norm": 10.0}`, both off by default. A delta with non-finite values is always left out. `screen_factor` leaves out a delta whose L2 norm is more than that many times the median, from three sites up. `clip_norm` scales longer deltas down to that norm. Left-out deltas are logged by run and listed under `excluded` in the global model's metadata; the round continues with the rest and fails only if none is left.
 - `fedprox_mu`, SF-11: weight of the FedProx proximal term, which pulls local weights towards the global model they started from. 0 by default, which is plain FedAvg.

controller/starfish/controller/file/transfer.py

@@ -166,7 +166,7 @@ def download_all(run_id, file_type, folder, task_seq=None, round_seq=None, all_r
 
 def publish_model(run_id, task_seq, round_seq, bucket_hz, source_version, parent=None,
                   eval_report=None):
-    """Coordinator: register a round's aggregated artifact as an approved model version."""
+    """Coordinator: register a round's gate-approved artifact; it waits for manual sign-off."""
     router, auth = _router()
     body = {'run': run_id, 'task_seq': task_seq, 'round_seq': round_seq, 'bucket_hz': bucket_hz,
             'source_version': source_version, 'parent': parent, 'eval_report': eval_report or {}}
@@ -180,26 +180,35 @@ def attempt():
     return _with_retries('publish of {}'.format(source_version), attempt, _retries())
 
 
-def list_models(bucket_hz):
-    """Approved model versions for a bucket, newest first."""
+def _registry_params(bucket_hz, status):
+    params = {'bucket_hz': bucket_hz}
+    if status:
+        params['status'] = status
+    return params
+
+
+def list_models(bucket_hz, status=None):
+    """Model versions for a bucket, newest first; approved only unless a reviewer asks for ``status``."""
     router, auth = _router()
 
     def attempt():
-        response = requests.get('{}/registry/'.format(router), params={'bucket_hz': bucket_hz},
+        response = requests.get('{}/registry/'.format(router),
+                                params=_registry_params(bucket_hz, status),
                                 auth=auth, timeout=TIMEOUT)
         _check(response, 'registry listing')
         return response.json()
 
     return _with_retries('registry listing', attempt, _retries())
 
 
-def latest_model(bucket_hz):
-    """The newest approved model version for a bucket, or None."""
+def latest_model(bucket_hz, status=None):
+    """The newest approved model version for a bucket, or None; ``status`` as in list_models."""
     router, auth = _router()
 
     def attempt():
         response = requests.get('{}/registry/latest/'.format(router),
-                                params={'bucket_hz': bucket_hz}, auth=auth, timeout=TIMEOUT)
+                                params=_registry_params(bucket_hz, status), auth=auth,
+                                timeout=TIMEOUT)
         if response.status_code == 404:
             return None
         _check(response, 'registry lookup')

controller/starfish/controller/tasks/babel_brain_fno/task.py

@@ -628,13 +628,18 @@ def _aggregate(self):
         return True
 
     def _publish(self, version):
-        """Register a gate-approved global model, SF-12. Failure is logged, not fatal."""
+        """Register a gate-approved global model, SF-12. Failure is logged, not fatal.
+
+        The version waits in the registry until the coordinator signs it off.
+        Its parent is the newest registered version whatever its status, since
+        training continued from it.
+        """
         try:
-            parent = transfer.latest_model(self.bucket_hz())
+            parent = transfer.latest_model(self.bucket_hz(), status='any')
             entry = transfer.publish_model(
                 self.run_id, self.cur_seq, self._round(), self.bucket_hz(), version,
                 parent=parent['version'] if parent else None, eval_report=self._eval_report)
-            self.logger.info('Published approved model {} for {} Hz, parent {}'.format(
+            self.logger.info('Published model {} for {} Hz, parent {}, pending sign-off'.format(
                 entry['version'], self.bucket_hz(), entry.get('parent')))
         except transfer.TransferFailed as e:
             self.logger.error(

controller/starfish/controller/tasks/babel_brain_fno/tests/test_gate.py

@@ -299,6 +299,8 @@ def test_accepted_model_is_published_with_its_parent_and_report(self):
                          'standin-250k-p7-b1-t1-r1'))
         self.assertEqual(kwargs['parent'], '250k-v0004')
         self.assertTrue(kwargs['eval_report']['accepted'])
+        # The parent is the newest registered version, pending or not
+        self.latest.assert_called_once_with(250000, status='any')
 
     def test_rejected_model_is_not_published(self):
         site = self.coordinator()

e2e/test_babelbrain_workflow.py

@@ -188,7 +188,8 @@ def logs_uploaded():
     assert {"current", "candidate", "accepted"} <= set(reports[0])
     # The eval store carries region labels, D6: one sample per region class
     breakdown = reports[0]["region_breakdown"]
-    assert set(breakdown) == {"P7", "P8", "PO7", "TP7", "TP8", "other"}, breakdown
+    assert set(breakdown) == {"P7", "P8", "PO7",
+                              "TP7", "TP8", "other"}, breakdown
     assert all(r["candidate"]["n"] == 1 for r in breakdown.values()), breakdown
     assert "/babelbrain-store" not in combined
     assert "[Agent]" not in combined
@@ -208,16 +209,27 @@ def test_router_holds_one_global_model_per_round(started_run):
 
 
 def test_registry_serves_the_latest_approved_model(started_run):
-    """SF-12: list the approved 250 kHz models, download the latest, verify its checksum."""
+    """SF-12: the coordinator signs off this run's newest model, then a lab lists the
+    approved 250 kHz models, downloads the latest and verifies its checksum."""
     import hashlib
+    project_id, _ = started_run
     # The run's rounds are finished once the earlier tests pass
+    pending = api("GET", "/registry/",
+                  params={"bucket_hz": 250000, "status": "pending"})
+    assert pending.ok, pending.text
+    mine = [v for v in pending.json() if v["project"] == project_id]
+    assert mine, "no pending model: the gate rejected every round"
+    approved = api("POST", "/registry/approve/",
+                   json={"version": mine[0]["version"], "note": "e2e"})
+    assert approved.ok, approved.text
+    assert approved.json()["status"] == "approved"
     listing = api("GET", "/registry/", params={"bucket_hz": 250000})
     assert listing.ok, listing.text
     versions = listing.json()
-    assert versions, "no approved model: the gate rejected every round"
+    assert all(v["status"] == "approved" for v in versions)
     latest = api("GET", "/registry/latest/",
                  params={"bucket_hz": 250000}).json()
-    assert latest["version"] == versions[0]["version"]
+    assert latest["version"] == versions[0]["version"] == mine[0]["version"]
     response = requests.get(ROUTER + "/registry/file/", params={"version": latest["version"]},
                             auth=AUTH, timeout=60)
     assert response.ok

router/starfish/migrations/0013_modelversion_sign_off.py

@@ -0,0 +1,40 @@
+from django.db import migrations, models
+from django.db.models import F
+
+
+def approve_existing(apps, schema_editor):
+    # Versions published before manual sign-off were already released to labs
+    ModelVersion = apps.get_model('starfish', 'ModelVersion')
+    ModelVersion.objects.update(status='approved', reviewed_by='before manual sign-off',
+                                reviewed_at=F('created_at'))
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('starfish', '0012_site_tokens'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='modelversion',
+            name='status',
+            field=models.CharField(choices=[('pending', 'Pending'), ('approved', 'Approved'), ('rejected', 'Rejected')], default='pending', max_length=10),
+        ),
+        migrations.AddField(
+            model_name='modelversion',
+            name='reviewed_by',
+            field=models.CharField(blank=True, default='', max_length=200),
+        ),
+        migrations.AddField(
+            model_name='modelversion',
+            name='reviewed_at',
+            field=models.DateTimeField(blank=True, null=True),
+        ),
+        migrations.AddField(
+            model_name='modelversion',
+            name='review_note',
+            field=models.CharField(blank=True, default='', max_length=500),
+        ),
+        migrations.RunPython(approve_existing, migrations.RunPython.noop),
+    ]

router/starfish/router/models.py

@@ -258,11 +258,18 @@ def __str__(self):
 
 class ModelVersion(models.Model):
     """
-    An approved global model in the registry, SF-12.
+    A gate-approved global model in the registry, SF-12.
 
-    Only a model that passed the coordinator's evaluation gate is published.
-    The registry keeps its own copy of the file, so it outlives the run files.
+    Only a model that passed the coordinator's evaluation gate is published,
+    and it starts as pending. Labs see it only after the coordinator tests it
+    and approves it by hand. The registry keeps its own copy of the file, so
+    it outlives the run files.
     """
+    class Status(models.TextChoices):
+        PENDING = 'pending', 'Pending'
+        APPROVED = 'approved', 'Approved'
+        REJECTED = 'rejected', 'Rejected'
+
     version = models.CharField(max_length=40, unique=True)
     bucket_hz = models.IntegerField()
     sequence = models.IntegerField()
@@ -277,6 +284,11 @@ class ModelVersion(models.Model):
     round_seq = models.IntegerField()
     source_version = models.CharField(max_length=200)
     eval_report = models.JSONField(default=dict, blank=True)
+    status = models.CharField(
+        max_length=10, choices=Status.choices, default=Status.PENDING)
+    reviewed_by = models.CharField(max_length=200, blank=True, default='')
+    reviewed_at = models.DateTimeField(null=True, blank=True)
+    review_note = models.CharField(max_length=500, blank=True, default='')
     created_at = models.DateTimeField(auto_now_add=True)
 
     class Meta: