feat: matched random control, quieter readout, safe plugin installs

.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@346b90ab5093d520d4c49c10752acce0afa6fb39
+    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@89366a83eb82e41e7bb0dcefe6401210383942de
     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: 346b90ab5093d520d4c49c10752acce0afa6fb39
+      unified_revision: 89366a83eb82e41e7bb0dcefe6401210383942de

brainscore_core/_model/behavior.py

@@ -14,6 +14,7 @@
 `fitting_labels` is an iterable of string labels, one per presentation.
 """
 
+import warnings
 from collections import OrderedDict
 from typing import Iterable
 
@@ -85,7 +86,13 @@ def fit(self, features, labels: Iterable) -> 'ProbabilitiesClassifier':
         self._scaler.fit(X)
         X = self._scaler.transform(X)
         y, self._label_mapping = self._labels_to_indices(list(labels))
-        self._classifier.fit(X, y)
+        # multi_class='multinomial' is deprecated but still honoured below sklearn 1.8,
+        # and dropping it would change binary probabilities; environment_check guards
+        # the version, so the deprecation notice only adds noise to every fit.
+        with warnings.catch_warnings():
+            warnings.filterwarnings('ignore', message=".*'multi_class' was deprecated",
+                                    category=FutureWarning)
+            self._classifier.fit(X, y)
         self._fitted = True
         return self
 

brainscore_core/plugin_management/import_plugin.py

@@ -1,7 +1,9 @@
 import logging
 import os
+import re
 import subprocess
 import sys
+import tempfile
 from pathlib import Path
 
 logger = logging.getLogger(__name__)
@@ -69,32 +71,108 @@ def locate_plugin(self) -> str:
 
     def install_requirements(self):
         """
-        Install all the requirements of the given plugin directory.
-        - If a `setup.py` file exists, it is run in the current interpreter.
-        - Alternatively, if a `requirements.txt` file exists, this is done via
-          `pip install` in the current interpreter.
-        - If both `setup.py` and `requirements.txt` are present, `setup.py` is installed first.
+        Install the requirements of the given plugin directory into the current interpreter.
+
+        - A `setup.py` installs the directory; a `requirements.txt` is installed with `-r`.
+        - Packages the environment is pinned to (see `PROTECTED_PACKAGES`) are held at their
+          installed versions. A plugin that needs other versions raises `PluginDependencyError`
+          instead of changing them, since that would move every other model's scores.
+        - Lines naming Python's own modules (e.g. `imp`, `collections`) are skipped: they are
+          not packages, and pip would fail on them.
         """
-        setup_file = self.plugins_dir / self.plugin_dirname / 'setup.py'
-        requirements_file = self.plugins_dir / self.plugin_dirname / 'requirements.txt'
-
+        plugin_dir = self.plugins_dir / self.plugin_dirname
+        setup_file, requirements_file = plugin_dir / 'setup.py', plugin_dir / 'requirements.txt'
         if not setup_file.is_file() and not requirements_file.is_file():
-            logger.debug(
-                f"Plugin {self.plugin_dirname} has no requirements file {requirements_file} "
-                f"or setup file {setup_file}"
-            )
-
-        if setup_file.is_file():
-            subprocess.run(
-                [sys.executable, '-m', 'pip', 'install', str(setup_file.parent)],
-                check=True,
-            )
-
-        if requirements_file.is_file():
-            subprocess.run(
-                [sys.executable, '-m', 'pip', 'install', '-r', str(requirements_file)],
-                check=True,
-            )
+            logger.debug(f"Plugin {self.plugin_dirname} has no requirements.txt or setup.py")
+            return
+        with tempfile.TemporaryDirectory(prefix='bs-plugin-') as scratch:
+            scratch = Path(scratch)
+            args = [str(plugin_dir)] if setup_file.is_file() else []
+            if requirements_file.is_file():
+                args += ['-r', str(_without_stdlib_modules(requirements_file, scratch / 'requirements.txt'))]
+            # In a dedicated environment ('newenv'), the plugin may set any versions it needs.
+            if installation_preference() == 'yes':
+                args = ['-c', str(_hold_installed_versions(scratch / 'constraints.txt')), *args]
+            completed = subprocess.run([sys.executable, '-m', 'pip', 'install', *args],
+                                       capture_output=True, text=True)
+        if completed.returncode != 0:
+            raise PluginDependencyError(_explain_install_failure(
+                self.identifier, completed.stdout + completed.stderr, plugin_dir))
+
+
+# Pinned by the shared environment; a plugin may add packages but not move these.
+PROTECTED_PACKAGES = ('numpy', 'pandas', 'scipy', 'scikit-learn', 'xarray', 'pyarrow', 'netCDF4',
+                      'torch', 'torchvision', 'transformers', 'datasets', 'importlib-metadata')
+
+
+class PluginDependencyError(RuntimeError):
+    """A plugin's requirements cannot be installed into this environment."""
+
+
+def _hold_installed_versions(path: Path) -> Path:
+    from importlib.metadata import PackageNotFoundError, version
+    lines = []
+    for name in PROTECTED_PACKAGES:
+        try:
+            lines.append(f'{name}=={version(name)}')
+        except PackageNotFoundError:
+            continue
+    path.write_text('\n'.join(lines) + '\n')
+    return path
+
+
+def _without_stdlib_modules(requirements: Path, path: Path) -> Path:
+    kept = []
+    for line in requirements.read_text().splitlines():
+        name = re.match(r'\s*([A-Za-z0-9_.-]+)\s*$', line.split('#')[0])
+        if name and name.group(1) in sys.stdlib_module_names:
+            logger.warning(f"Skipping '{name.group(1)}' in {requirements}: it is part of Python, not a package.")
+            continue
+        kept.append(line)
+    path.write_text('\n'.join(kept) + '\n')
+    return path
+
+
+def _explain_install_failure(identifier: str, output: str, plugin_dir: Path = None) -> str:
+    """A readable reason, the pip lines behind it, and what to try next."""
+    stderr = output
+    lines = [line for line in output.splitlines() if line.strip()]
+    requirements = (plugin_dir / 'requirements.txt') if plugin_dir is not None else Path('<plugin>/requirements.txt')
+    runner = ("Next step: give it an environment of its own, so this one keeps its versions:\n"
+              "    conda create --name bs-model --clone \"$CONDA_DEFAULT_ENV\" && conda activate bs-model\n"
+              f"    python -m pip install -r \"{requirements}\"\n"
+              f"then load or score '{identifier}' there.")
+    linux = ("Next step: load it on Linux (for example an EC2 instance), where more packages have "
+             "ready-made builds. If it fails there too, please tell a maintainer: "
+             "https://github.com/brain-score/vision/issues")
+    missing = [line for line in lines if 'No matching distribution found for' in line]
+    if 'ResolutionImpossible' in stderr or 'conflict' in stderr.lower():
+        detail = []
+        for line in lines:
+            line = line.strip()
+            if line.startswith('The user requested (constraint) '):
+                detail.append('this environment holds ' + line.split(') ', 1)[1])
+            elif line.startswith('The user requested '):
+                detail.append('the plugin asks for ' + line[len('The user requested '):])
+            elif ' depends on ' in line:
+                detail.append(line)
+        detail = detail[:6]
+        reason = (f"'{identifier}' needs package versions that conflict with this environment. "
+                  "Installing them would change the versions scores are verified against, "
+                  "so nothing was installed.")
+    elif missing:
+        detail = [line.strip() for line in missing][:3]
+        reason = (f"'{identifier}' needs a package with no release for this platform or Python version. "
+                  f"It may install on Linux.")
+        runner = linux
+    elif 'Failed to build' in stderr:
+        detail = [line.strip() for line in lines if 'Failed to build' in line][:3]
+        reason = f"'{identifier}' needs a package that failed to build from source here."
+        runner = linux
+    else:
+        detail = lines[-5:]
+        reason = f"Installing the requirements of '{identifier}' failed."
+    return '\n'.join([reason, *('  ' + line for line in detail), runner])
 
 
 def installation_preference():

brainscore_core/selection.py

@@ -129,6 +129,29 @@ class RandomSelection(UnitSelection):
     axis: Optional[int] = -1
     expected_shape: Optional[Tuple[int, ...]] = None
 
+    @classmethod
+    def matched_to(cls, selection: 'Selection', seed: int = 0) -> 'RandomSelection':
+        """The random control for a resolved selection: the same layer, layout and
+        number of units, drawn from the same population.
+
+        ``RandomSelection.matched_to(units, seed=0)`` is shorthand for passing
+        ``layer``, ``n_units``, ``n_total``, ``population``, ``axis`` and
+        ``expected_shape`` by hand.
+        """
+        metadata = selection.metadata or {}
+        population = metadata.get('unit_population')
+        if population is not None:
+            n_total = len(population)
+        elif 'n_recorded' in metadata:
+            n_total = int(metadata['n_recorded'])
+        else:
+            raise ValueError('matched_to needs a selection that records its unit population '
+                             "(metadata['unit_population'] or metadata['n_recorded']), "
+                             'such as a resolved FunctionalSelection.')
+        return cls(layer=selection.layer, n_units=len(selection.indices), n_total=n_total,
+                   seed=seed, population=None if population is None else list(population),
+                   axis=selection.axis, expected_shape=selection.expected_shape)
+
     def resolve(self, model) -> 'Selection':
         import numpy as np
         k = min(self.n_units, self.n_total)

tests/test_plugin_management/test_import_plugin.py

@@ -85,3 +85,33 @@ def test_stimulus_set(self):
                                 plugin_type='data', identifier='dummy-stimulus_set',
                                 registry_prefix='stimulus_set')
         assert importer.plugin_dirname == 'dummy_stimulus_set'
+
+
+class TestInstallSafeguards:
+    def test_python_modules_are_skipped(self, tmp_path):
+        from brainscore_core.plugin_management.import_plugin import _without_stdlib_modules
+        requirements = tmp_path / 'requirements.txt'
+        requirements.write_text('torch\nimp\ncollections  # stdlib\ntimm>=1.0\n')
+        kept = _without_stdlib_modules(requirements, tmp_path / 'out.txt').read_text().split()
+        assert kept == ['torch', 'timm>=1.0']
+
+    def test_installed_pins_are_held(self, tmp_path):
+        import numpy
+        from brainscore_core.plugin_management.import_plugin import _hold_installed_versions
+        assert f'numpy=={numpy.__version__}' in _hold_installed_versions(tmp_path / 'c.txt').read_text()
+
+    def test_conflict_is_explained_with_a_next_step(self):
+        from brainscore_core.plugin_management.import_plugin import _explain_install_failure
+        stderr = ('ERROR: Cannot install transformers==4.30.2 because these package versions have conflicting dependencies.\n'
+                  'The conflict is caused by:\n    The user requested transformers==4.30.2\n'
+                  '    The user requested (constraint) transformers==4.57.6\nERROR: ResolutionImpossible\n')
+        message = _explain_install_failure('cvt_13', stderr)
+        assert 'conflict with this environment' in message
+        assert 'the plugin asks for transformers==4.30.2' in message and 'this environment holds transformers==4.57.6' in message
+        assert 'conda create --name bs-model --clone' in message
+
+    def test_missing_platform_package_is_explained(self):
+        from brainscore_core.plugin_management.import_plugin import _explain_install_failure
+        message = _explain_install_failure('openstl', 'ERROR: No matching distribution found for decord\n')
+        assert 'no release for this platform' in message and 'decord' in message
+        assert 'Linux' in message

tests/test_plugin_management/test_plugin_command_paths.py

@@ -201,12 +201,22 @@ def test_plugin_install_uses_current_python_and_literal_paths(tmp_path, monkeypa
     installer.plugins_dir = tmp_path
     installer.plugin_dirname = plugin.name
     calls = []
-    monkeypatch.setattr(subprocess, 'run', lambda *args, **kwargs: calls.append((args, kwargs)))
+
+    def run(*args, **kwargs):
+        calls.append((args, kwargs))
+        return subprocess.CompletedProcess(args[0], 0, '', '')
+
+    monkeypatch.setenv('BS_INSTALL_DEPENDENCIES', 'yes')
+    monkeypatch.setattr(subprocess, 'run', run)
     installer.install_requirements()
     command = calls[0][0][0]
     assert command[:4] == [sys.executable, '-m', 'pip', 'install']
-    assert command[4:] == ([str(plugin)] if filename == 'setup.py' else ['-r', str(plugin / filename)])
-    assert calls[0][1] == {'check': True}
+    assert command[4] == '-c'                    # pinned packages held at their installed versions
+    if filename == 'setup.py':
+        assert command[6:] == [str(plugin)]
+    else:                                        # a copy without Python's own module names
+        assert command[6] == '-r' and Path(command[7]).name == 'requirements.txt'
+    assert calls[0][1] == {'capture_output': True, 'text': True}
 
 
 def test_failed_plugin_install_stops_before_import(tmp_path, monkeypatch):
@@ -217,9 +227,12 @@ def test_failed_plugin_install_stops_before_import(tmp_path, monkeypatch):
     installer.plugins_dir = tmp_path.parent
     installer.plugin_dirname = tmp_path.name
 
+    installer.identifier = 'dummy'
+    from brainscore_core.plugin_management.import_plugin import PluginDependencyError
+
     def fail(command, **kwargs):
-        raise subprocess.CalledProcessError(1, command)
+        return subprocess.CompletedProcess(command, 1, '', 'ERROR: something went wrong')
 
     monkeypatch.setattr(subprocess, 'run', fail)
-    with pytest.raises(subprocess.CalledProcessError):
+    with pytest.raises(PluginDependencyError, match='dummy'):
         installer.install_requirements()

tests/test_unit_selection.py

@@ -283,3 +283,24 @@ def test_subset_without_original_addresses_is_rejected():
     model.region_layer_selectors = {'VWFA': CompositeSelector((('enc.5', (4, 1)),))}
     with pytest.raises(ValueError, match='original unit_index'):
         FunctionalSelection('VWFA', 'stim', (['word'], ['nonword']), n_units=1).resolve(model)
+
+
+def test_matched_to_copies_the_functional_selection_settings():
+    from brainscore_core.model_interface import Selection
+    from brainscore_core.selection import RandomSelection
+    population = list(range(10, 60))
+    units = Selection(layer='features.12', indices=[11, 12, 13], axis=1, expected_shape=(512, 13, 13),
+                      metadata={'selector': 'functional', 'unit_population': population})
+    by_hand = RandomSelection('features.12', n_units=3, n_total=50, population=population,
+                              seed=4, axis=1, expected_shape=(512, 13, 13)).resolve(None)
+    matched = RandomSelection.matched_to(units, seed=4).resolve(None)
+    assert matched.indices == by_hand.indices
+    assert (matched.layer, matched.axis, matched.expected_shape) == ('features.12', 1, (512, 13, 13))
+
+
+def test_matched_to_needs_a_population():
+    import pytest
+    from brainscore_core.model_interface import Selection
+    from brainscore_core.selection import RandomSelection
+    with pytest.raises(ValueError, match='unit population'):
+        RandomSelection.matched_to(Selection(layer='x', indices=[0], metadata={}))