Add verify to check cached models against their recorded digests

Author: kevinzakka

python/DOC.md

@@ -93,7 +93,7 @@ A warm cache can be read-only. Concurrent cold starts on one model are safe: a m
 **Command line**
 
 ```bash
-uvx mujoco-menagerie names | info NAME | path NAME [ENTRY] | view NAME [ENTRY] | prefetch [NAME...] | prune
+uvx mujoco-menagerie names | info NAME | path NAME [ENTRY] | view NAME [ENTRY] | prefetch [NAME...] | verify [NAME...] | prune
 ```
 
 **Speed**, laptop, local mirror

python/src/mujoco_menagerie/__main__.py

@@ -39,6 +39,10 @@ def main(argv: list[str] | None = None) -> None:
   sub.add_parser(
     'prune', help='delete cached models this version no longer references'
   )
+  sub.add_parser(
+    'verify',
+    help='check cached models against the digests recorded at download',
+  ).add_argument('names', nargs='*')
   a = p.parse_args(argv)
   try:
     if a.cmd == 'names':
@@ -61,6 +65,16 @@ def main(argv: list[str] | None = None) -> None:
       mm.prefetch(a.names)
     elif a.cmd == 'prune':
       print(*mm.prune(), sep='\n')
+    elif a.cmd == 'verify':
+      cache = mm.Cache()
+      bad = {
+        r.name: cache.verify(r)
+        for r in map(mm.get, a.names or mm.names())
+        if cache.is_cached(r)
+      }
+      for name, files in bad.items():
+        print(name, 'ok' if not files else 'changed: ' + ' '.join(files))
+      sys.exit(any(bad.values()))
   except mm.MenagerieError as e:
     sys.exit(f'error: {e}')
 

python/src/mujoco_menagerie/_cache.py

@@ -139,7 +139,14 @@ def _publish(
       tree = stage / robot.name
       if not tree.is_dir():
         raise DownloadError(f'{robot.asset} has no top-level {robot.name}/')
-      (tree / SENTINEL).write_text(json.dumps({'oid': robot.oid}))
+      files = {
+        p.relative_to(tree).as_posix(): sha256_file(p)
+        for p in sorted(tree.rglob('*'))
+        if p.is_file()
+      }
+      (tree / SENTINEL).write_text(
+        json.dumps({'oid': robot.oid, 'files': files})
+      )
       _fsync_tree(tree)
       target.parent.mkdir(parents=True, exist_ok=True)
       with contextlib.suppress(OSError):  # lost a race to another process
@@ -150,6 +157,16 @@ def _publish(
     finally:
       shutil.rmtree(stage, ignore_errors=True)
 
+  def verify(self, robot: Robot) -> list[str]:
+    """Files missing or changed since the model was published."""
+    target = self.model_path(robot)
+    files = json.loads((target / SENTINEL).read_text())['files']
+    return [
+      f
+      for f, h in files.items()
+      if not (target / f).is_file() or sha256_file(target / f) != h
+    ]
+
   def prune(self, keep: Iterable[Robot]) -> list[pathlib.Path]:
     # Like uninstalling a package: not safe for models other processes are using.
     keep = {self.model_path(r).name for r in keep}

python/tests/test_cache.py

@@ -133,6 +133,15 @@ def test_environment_supplies_defaults(monkeypatch, tmp_path):
   assert Cache(base_url='file:///x').base_url == 'file:///x'
 
 
+def test_verify_detects_changed_files(robot, base_url, cache_dir):
+  cache = Cache(cache_dir, base_url)
+  path = cache.resolve(robot)
+  assert cache.verify(robot) == []
+  (path / 'fake.xml').write_text('<mujoco/>')
+  (path / 'assets/cube.obj').unlink()
+  assert cache.verify(robot) == ['assets/cube.obj', 'fake.xml']
+
+
 def test_versions_coexist_and_prune(fake_model, archives, base_url, cache_dir):
   v1 = make_robot(fake_model, archives, oid='1' * 40)
   (fake_model / 'README.md').write_text('# v2\n')