cache: resolve modules= against the tree instead of caching nothing

docs/usage/cache.md

@@ -63,6 +63,17 @@ with model.trace("The Eiffel Tower is in") as tracer:
     ])
 ```
 
+A path string is the **whole** path, spelled the way the cache's keys are
+(`"model.lm_head"`, not `"lm_head"`), and is resolved against the envoy tree when
+the cache is created: one that names no module raises there rather than
+returning an empty cache. Only envoys and paths are accepted — a glob
+(`"model.transformer.h.*"`), a regex, or a predicate matches nothing and raises
+too. To select modules by pattern, filter them yourself and pass the envoys:
+
+```python
+mlps = [envoy for envoy in model.modules() if envoy.path.endswith(".mlp")]
+```
+
 ### Include inputs
 
 ```python

src/nnsight/intervention/cache.py

@@ -20,6 +20,8 @@
 
 from __future__ import annotations
 
+import inspect
+from collections.abc import Iterable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
@@ -66,7 +68,9 @@ class Cache:
 
     Attributes:
         model: The root envoy, used to resolve paths / aliases for [`CacheView`][nnsight.intervention.cache.CacheView].
-        targets: The module paths to keep, or ``None`` to keep every module.
+        targets: The module paths to keep, resolved against the envoy tree (see
+            [`_resolve_targets`][nnsight.intervention.cache.Cache._resolve_targets]),
+            or ``None`` to keep every module.
         entries: Recorded values, ``{module_path: [Entry, ...]}``.
     """
 
@@ -90,15 +94,88 @@ def __init__(
         self.include_inputs = include_inputs
         # None => every module; else the exact set of paths to keep.
         self.targets: set[str] | None = (
-            None
-            if modules is None
-            else {m if isinstance(m, str) else m.path for m in modules}
+            None if modules is None else self._resolve_targets(modules)
         )
         self.entries: dict[str, list[Entry]] = {}
         # Which slots of each path's newest Entry have been written this visit.
         # Recording-time bookkeeping only; see `_record`.
         self._open: dict[str, set[str]] = {}
 
+    def _resolve_targets(self, modules: Any) -> set[str]:
+        """The module paths ``modules=`` names, raising for anything that names none.
+
+        A target is matched against the paths the run reports, so a string that
+        resolves to no module -- a typo, a path missing the model's own name, a
+        glob, a regex -- subscribed to a location nothing ever provides: an empty
+        cache and no error. Resolving each one against the envoy tree here gives
+        the same error reading that module by attribute does.
+        """
+        if isinstance(modules, str) or not isinstance(modules, Iterable):
+            # A lone target rather than a list of them: a string would otherwise
+            # be taken a character at a time, and a callable or a regex is not
+            # iterable at all -- `_resolve` says what `modules=` accepts.
+            modules = [modules]
+
+        return {self._resolve(module) for module in modules}
+
+    def _resolve(self, module: Any) -> str:
+        """One ``modules=`` entry as a module path."""
+        from .envoy import Envoy
+
+        if isinstance(module, Envoy):
+            return module.path
+        if not isinstance(module, str):
+            raise TypeError(
+                f"cache modules= takes envoys (model.transformer.h[0]) or their "
+                f"paths ('model.transformer.h.0'), not {type(module).__name__}. "
+                f"Globs, regexes and predicates are not matched against the tree: "
+                f"select the modules yourself and pass them."
+            )
+
+        root = self.model.path
+        if module == root:
+            return root
+        if not module.startswith(f"{root}."):
+            # A target is spelled the way the cache's own keys are, from the
+            # model's own name down; the common miss is dropping that name, so
+            # say so when adding it back does resolve.
+            try:
+                suggestion = self._resolve(f"{root}.{module}")
+            except (AttributeError, TypeError):
+                suggestion = None
+            raise AttributeError(
+                f"{module!r} is not a module path"
+                + (
+                    f" -- did you mean {suggestion!r}?"
+                    if suggestion is not None
+                    else f": no module of {root!r} is named that."
+                )
+            )
+
+        envoy: Any = self.model
+        for name in module[len(root) + 1 :].split("."):
+            # `.output` / `.input` are values the run serves, not modules; reading
+            # one here would request it from the run instead of resolving a path.
+            if inspect.isdatadescriptor(getattr(type(envoy), name, None)):
+                raise AttributeError(
+                    f"{module!r} names {name!r}, which is a value nnsight serves "
+                    f"on every module, not a module. Cache targets are modules."
+                )
+            try:
+                envoy = getattr(envoy, name)
+            except AttributeError as error:
+                # The envoy's own miss, named against the target it came from —
+                # otherwise a list of paths doesn't say which one is wrong.
+                raise AttributeError(
+                    f"{module!r} is not a module path: {error}"
+                ) from error
+            if not isinstance(envoy, Envoy):
+                raise AttributeError(
+                    f"{module!r} is not a module: {name!r} is a "
+                    f"{type(envoy).__name__}."
+                )
+        return envoy.path
+
     def __getstate__(self) -> dict:
         # `model` is a live Envoy used only for CacheView navigation, and it drags
         # in the interleaver's unpicklable instrumented forwards. Drop it so the

src/nnsight/intervention/tracer.py

@@ -200,7 +200,11 @@ def cache(
         per step (``len(cache[path])`` is the step count).
 
         Args:
-            modules: Envoys or path strings to capture; ``None`` captures every module.
+            modules: Envoys or their paths (``"model.transformer.h.0"``, the
+                whole path, as the cache's keys are) to capture; ``None``
+                captures every module. A path is resolved against the envoy
+                tree, so one that names no module raises rather than caching
+                nothing; globs, regexes and predicates are not matched.
             device: Device to move captured tensors to (default CPU); ``None`` leaves them.
             dtype: Optional dtype to cast captured tensors to.
             detach: Detach captured tensors from the autograd graph.

tests/test_interleaving.py

@@ -1,4 +1,6 @@
 
+import re
+
 import pytest
 import torch
 import nnsight
@@ -753,6 +755,52 @@ def test_subset_by_reference(self, envoy, x):
         assert cache.keys() == ["model.l2"]
         assert "model.l1" not in cache
 
+    # -- modules= is resolved against the tree -----------------------------
+
+    def test_subset_by_path_string(self, envoy, x):
+        with envoy.trace(x) as tracer:
+            cache = tracer.cache(modules=["model.l2"])
+        assert cache.keys() == ["model.l2"]
+
+    def test_a_lone_path_string_is_one_target(self, envoy, x):
+        # Not a list: iterating the string would take it a character at a time.
+        with envoy.trace(x) as tracer:
+            cache = tracer.cache(modules="model.l2")
+        assert cache.keys() == ["model.l2"]
+
+    def test_alias_target_resolves_to_the_real_path(self, model, x):
+        renamed = Envoy(model, rename={"l1": "first"})
+        with renamed.trace(x) as tracer:
+            cache = tracer.cache(modules=["model.first"])
+        assert cache.keys() == ["model.l1"]
+
+    @pytest.mark.parametrize(
+        "path", ["model.ln", "model.l1.lin", "model.l*", r".*\.l1$", "model.l1.weight"]
+    )
+    def test_a_path_naming_no_module_raises(self, envoy, x, path):
+        # Targets used to be taken as given, so a typo (or a glob, or a regex)
+        # subscribed to a location nothing provides: an empty cache, no error.
+        with pytest.raises(AttributeError):
+            with envoy.trace(x) as tracer:
+                tracer.cache(modules=[path])
+
+    def test_a_path_missing_the_model_name_suggests_the_whole_one(self, envoy, x):
+        with pytest.raises(AttributeError, match="did you mean 'model.l1'"):
+            with envoy.trace(x) as tracer:
+                tracer.cache(modules=["l1"])
+
+    def test_a_served_value_is_not_a_module(self, envoy, x):
+        with pytest.raises(AttributeError, match="not a module"):
+            with envoy.trace(x) as tracer:
+                tracer.cache(modules=["model.l1.output"])
+
+    @pytest.mark.parametrize("target", [re.compile("l1"), lambda envoy: True])
+    def test_a_target_that_is_neither_envoy_nor_path_says_so(self, envoy, x, target):
+        # `modules=<callable>` raised "'function' object is not iterable".
+        with pytest.raises(TypeError, match="cache modules="):
+            with envoy.trace(x) as tracer:
+                tracer.cache(modules=target)
+
     def test_navigation_matches_path(self, envoy, x):
         with envoy.trace(x) as tracer:
             cache = tracer.cache()