feat(contract): unify subjects and preserve legacy benchmarks

.github/workflows/umi-v2.yml

@@ -11,9 +11,9 @@ jobs:
   coordinated:
     name: UMI v2
     # Update the workflow and runner revision together.
-    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@a95ad140a226cf5dd56ca9559e2504214078f2c1
+    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@b8880276242296c6857363777b4bd12a02ab5213
     with:
       candidate: core
       candidate_repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }}
       candidate_revision: ${{ github.event.pull_request.head.sha || github.sha }}
-      unified_revision: a95ad140a226cf5dd56ca9559e2504214078f2c1
+      unified_revision: b8880276242296c6857363777b4bd12a02ab5213

brainscore_core/behavioral.py

@@ -17,11 +17,66 @@ def __init__(self, owner) -> None:
         self.task_context: Optional[TaskContext] = None
         self.readout_classifier: Any = None
         self.use_generation_for_task = False
+        # The TaskContext object the current configuration came from, so a
+        # repeated send of the same task keeps the fitted readout.
+        self.configured_from: Optional[TaskContext] = None
+        self._configured_as: Optional[set] = None  # fingerprints of the fitted configuration
+        self._stale = False
 
     @classmethod
     def requires_behavioral_readout(cls, task_type: str) -> bool:
         return task_type in cls.BEHAVIORAL_TASK_TYPES
 
+    @staticmethod
+    def _fingerprint(task_context: TaskContext) -> tuple:
+        """Everything the fitted readout depends on: task, fitting rows and labels, generation settings."""
+        fitting = task_context.fitting_stimuli
+        rows = labels = None
+        columns = getattr(fitting, 'columns', None)
+        if columns is not None:
+            try:
+                if 'stimulus_id' in columns:
+                    rows = tuple(fitting['stimulus_id'])
+                for column in ('image_label', 'label'):
+                    if column in columns:
+                        labels = tuple(fitting[column])
+                        break
+            except (TypeError, KeyError, AttributeError):
+                rows = labels = None  # not a table; identity alone must do
+        label_set = task_context.label_set
+        return (task_context.task_type, id(fitting), rows, labels, task_context.instruction,
+                tuple(label_set) if label_set is not None else None,
+                task_context.prefer_path or 'auto')
+
+    def is_configured_for(self, task_context: TaskContext) -> bool:
+        """Whether ``task_context`` is the task already configured, unchanged since.
+
+        True for the object the configuration came from and for the context
+        this component holds (the same task after input projection), as long
+        as its task, fitting table, labels and generation settings are the
+        ones fitted. TaskContext is mutable: a caller that swaps the fitting
+        table or relabels it and calls start_task again gets a new fit.
+        """
+        if self._stale or self.task_context is None or not self._configured_as:
+            return False
+        if task_context is not self.configured_from and task_context is not self.task_context:
+            return False
+        return self._fingerprint(task_context) in self._configured_as
+
+    def mark_configured(self, task_context: Optional[TaskContext]) -> None:
+        """Remember the configuration just fitted, as the caller's context and as the one held."""
+        self.configured_from = task_context
+        self._configured_as = None
+        if task_context is not None:
+            self._configured_as = {self._fingerprint(task_context)}
+            if self.task_context is not None and self.task_context is not task_context:
+                self._configured_as.add(self._fingerprint(self.task_context))
+        self._stale = False
+
+    def invalidate(self) -> None:
+        """The model changed (a state change was applied or undone): fit again next time."""
+        self._stale = True
+
     def start_task(self, task_context_or_task, fitting_stimuli=None,
                    **kwargs) -> None:
         del kwargs
@@ -108,6 +163,9 @@ def reset(self) -> None:
         self.task_context = None
         self.readout_classifier = None
         self.use_generation_for_task = False
+        self.configured_from = None
+        self._configured_as = None
+        self._stale = False
 
     def fit_behavioral_readout(self, fitting_stimuli) -> None:
         """Extract readout-layer features and fit a probabilities classifier."""
@@ -145,29 +203,29 @@ def extract_labels(stimuli):
     def extract_behavioral_features(self, stimuli):
         """Run the model once at the behavioral readout layer."""
         owner = self.owner
-        saved_layer = owner._recording_layer
-        saved_layers = owner._recording_layers
-        saved_regions = owner._recording_regions
-        saved_multi = owner._is_multi_region
-        saved_composite = owner._composite_recording
+        saved_layer = owner._recorder.recording_layer
+        saved_layers = owner._recorder.recording_layers
+        saved_regions = owner._recorder.recording_regions
+        saved_multi = owner._recorder.is_multi_region
+        saved_composite = owner._recorder.composite_recording
         saved_classifier = self.readout_classifier
-        owner._recording_layer = owner._behavioral_readout_layer
-        owner._recording_layers = (
+        owner._recorder.recording_layer = owner._behavioral_readout_layer
+        owner._recorder.recording_layers = (
             [owner._behavioral_readout_layer]
             if owner._behavioral_readout_layer else []
         )
-        owner._recording_regions = []
-        owner._is_multi_region = False
-        owner._composite_recording = False
+        owner._recorder.recording_regions = []
+        owner._recorder.is_multi_region = False
+        owner._recorder.composite_recording = False
         self.readout_classifier = None
         try:
             features = owner.process(stimuli)
         finally:
-            owner._recording_layer = saved_layer
-            owner._recording_layers = saved_layers
-            owner._recording_regions = saved_regions
-            owner._is_multi_region = saved_multi
-            owner._composite_recording = saved_composite
+            owner._recorder.recording_layer = saved_layer
+            owner._recorder.recording_layers = saved_layers
+            owner._recorder.recording_regions = saved_regions
+            owner._recorder.is_multi_region = saved_multi
+            owner._recorder.composite_recording = saved_composite
             self.readout_classifier = saved_classifier
 
         if 'presentation' in features.dims and 'neuroid' in features.dims: