Fix some issues with config registration in tests

Signed-off-by: Fabrice Normandin <[email protected]>

project/algorithms/__init__.py

@@ -14,7 +14,7 @@
 
 # If you add a configuration file under `configs/algorithm`, it will also be available as an option
 # from the command-line, and be validated against the schema.
-
+# todo: It might be nicer if we did this this `configs/algorithms` instead of here, no?
 algorithm_store = store(group="algorithm")
 algorithm_store(ExampleAlgorithm.HParams(), name="example_algo")
 algorithm_store(ManualGradientsExample.HParams(), name="manual_optimization")

project/algorithms/bases/algorithm_test.py

@@ -23,7 +23,7 @@
 from torch.utils.data import DataLoader
 from typing_extensions import ParamSpec
 
-from project.configs.config import Config, cs
+from project.configs import Config, cs
 from project.conftest import setup_hydra_for_tests_and_compose
 from project.datamodules.image_classification import (
     ImageClassificationDataModule,

project/configs/__init__.py

@@ -11,12 +11,11 @@
 )
 from .network import network_store
 
-# todo: look into using this instead:
-# from hydra_zen import store
-
 cs = ConfigStore.instance()
+cs.store(name="base_config", node=Config)
 datamodule_store.add_to_hydra_store()
 network_store.add_to_hydra_store()
+# todo: move the algorithm_store.add_to_hydra_store() here?
 
 __all__ = [
     "Config",

project/configs/config.py

@@ -3,8 +3,6 @@
 from logging import getLogger as get_logger
 from typing import Any, Literal
 
-from hydra.core.config_store import ConfigStore
-
 logger = get_logger(__name__)
 LogLevel = Literal["debug", "info", "warning", "error", "critical"]
 
@@ -39,7 +37,3 @@ class Config:
     debug: bool = False
 
     verbose: bool = False
-
-
-cs = ConfigStore.instance()
-cs.store(name="base_config", node=Config)

project/configs/network/__init__.py

@@ -1,7 +1,26 @@
+import hydra_zen
+import torchvision.models
 from hydra_zen import store
 
-from project.networks import FcNetConfig, ResNet18Config
+from project.networks.fcnet import FcNet
+from project.utils.hydra_utils import interpolate_config_attribute
 
 network_store = store(group="network")
-network_store(FcNetConfig, name="fcnet")
-network_store(ResNet18Config, name="resnet18")
+network_store(
+    hydra_zen.builds(
+        torchvision.models.resnet18,
+        populate_full_signature=True,
+        num_classes=interpolate_config_attribute("datamodule.num_classes"),
+    ),
+    name="resnet18",
+)
+network_store(
+    hydra_zen.builds(
+        FcNet,
+        hydra_convert="object",
+        hydra_recursive=True,
+        populate_full_signature=True,
+        output_dims=interpolate_config_attribute("datamodule.num_classes"),
+    ),
+    name="fcnet",
+)

project/configs/network/jax_cnn.yaml

@@ -0,0 +1,2 @@
+_target_: project.algorithms.jax_algo.CNN
+num_classes: ${instance_attr:datamodule.num_classes}

project/configs/network/jax_fcnet.yaml

@@ -0,0 +1,3 @@
+_target_: project.algorithms.jax_algo.JaxFcNet
+num_classes: ${instance_attr:datamodule.num_classes}
+num_features: 256

project/configs/network/resnet50.yaml

@@ -0,0 +1,3 @@
+_target_: torchvision.models.resnet50
+pretrained: true
+num_classes: "${instance_attr:datamodule.num_classes,datamodule.action_dims:1000}"

project/main_test.py

@@ -4,14 +4,15 @@
 import typing
 from pathlib import Path
 
+import hydra_zen
 import pytest
 
 from project.algorithms import Algorithm, ExampleAlgorithm
 from project.configs.config import Config
 from project.configs.datamodule import CIFAR10DataModuleConfig
 from project.conftest import setup_hydra_for_tests_and_compose, use_overrides
 from project.datamodules.image_classification.cifar10 import CIFAR10DataModule
-from project.networks import FcNetConfig
+from project.networks.fcnet import FcNet
 from project.utils.hydra_utils import resolve_dictconfig
 
 if typing.TYPE_CHECKING:
@@ -77,13 +78,13 @@ def test_setting_algorithm(
 @pytest.mark.parametrize(
     ("overrides", "expected_type"),
     [
-        (["algorithm=example_algo", "network=fcnet"], FcNetConfig),
+        (["algorithm=example_algo", "network=fcnet"], FcNet),
     ],
     ids=_ids,
 )
 def test_setting_network(
     overrides: list[str],
-    expected_type: type[Algorithm.HParams],
+    expected_type: type,
     testing_overrides: list[str],
     tmp_path: Path,
 ) -> None:
@@ -93,7 +94,7 @@ def test_setting_network(
     ) as dictconfig:
         options = resolve_dictconfig(dictconfig)
     assert isinstance(options, Config)
-    assert isinstance(options.network, expected_type)
+    assert hydra_zen.get_target(options.network) is expected_type
 
 
 # TODO: Add some more integration tests:

project/networks/__init__.py

@@ -13,30 +13,7 @@
 # _cs.store(group="network", name="fcnet", node=FcNetConfig)
 # _cs.store(group="network", name="resnet18", node=ResNet18Config)
 # Add your network configs here.
-from dataclasses import field
-
-from hydra_zen import hydrated_dataclass
-from torchvision.models import resnet18
-
-from project.utils.hydra_utils import interpolated_field
 
 from .fcnet import FcNet
 
-
-@hydrated_dataclass(target=FcNet, hydra_convert="object", hydra_recursive=True)
-class FcNetConfig:
-    output_dims: int = interpolated_field(
-        "${instance_attr:datamodule.num_classes,datamodule.action_dims}", default=-1
-    )
-    hparams: FcNet.HParams = field(default_factory=FcNet.HParams)
-
-
-@hydrated_dataclass(target=resnet18)
-class ResNet18Config:
-    pretrained: bool = False
-    num_classes: int = interpolated_field(
-        "${instance_attr:datamodule.num_classes,datamodule.action_dims}", default=1000
-    )
-
-
 __all__ = ["FcNet"]

project/utils/hydra_utils.py

@@ -27,6 +27,86 @@
 T = TypeVar("T")
 
 
+def interpolate_config_attribute(*attributes: str, default: Any | Literal[MISSING] = MISSING):
+    """Use this in a config to to get an attribute from another config after it is instantiated.
+
+    Multiple attributes can be specified, which will lead to trying each of them in order until the
+    attribute is found. If none are found, then an error will be raised.
+
+    For example, if we only know the number of classes in the datamodule after it is instantiated,
+    we can set this in the network config so it is created with the right number of output dims.
+
+    ```yaml
+    _target_: torchvision.models.resnet50
+    num_classes: ${instance_attr:datamodule.num_classes}
+    ```
+
+    This is equivalent to:
+
+    >>> import hydra_zen
+    >>> import torchvision.models
+    >>> resnet50_config = hydra_zen.builds(
+    ...     torchvision.models.resnet50,
+    ...     num_classes=interpolate_config_attribute("datamodule.num_classes"),
+    ...     populate_full_signature=True,
+    ... )
+    >>> print(hydra_zen.to_yaml(resnet50_config))  # doctest: +NORMALIZE_WHITESPACE
+    _target_: torchvision.models.resnet.resnet50
+    weights: null
+    progress: true
+    num_classes: ${instance_attr:datamodule.num_classes}
+    """
+    if default is MISSING:
+        return "${instance_attr:" + ",".join(attributes) + "}"
+    return "${instance_attr:" + ",".join(attributes) + ":" + str(default) + "}"
+
+
+def interpolated_field(
+    interpolation: str,
+    default: T | Literal[MISSING] = MISSING,
+    default_factory: Callable[[], T] | Literal[MISSING] = MISSING,
+    instance_attr: bool = False,
+) -> T:
+    """Field with a default value computed with a OmegaConf-style interpolation when appropriate.
+
+    When the dataclass is created by Hydra / OmegaConf, the interpolation is used.
+    Otherwise, behaves as usual (either using default or calling the default_factory).
+
+    Parameters
+    ----------
+    interpolation: The string interpolation to use to get the default value.
+    default: The default value to use when not in a hydra/OmegaConf context.
+    default_factory: The default value to use when not in a hydra/OmegaConf context.
+    instance_attr: Whether to use the `instance_attr` custom resolver to run the interpolation \
+        with respect to instantiated objects instead of their configs.
+        Passing `interpolation='${instance_attr:some_config.some_attr}'` has the same effect.
+
+    This last parameter is important, since in order to retrieve the instance attribute, we need to
+    instantiate the objects, which could be expensive. These instantiated objects are reused at
+    least, but still, be mindful when using this parameter.
+    """
+    assert "${" in interpolation and "}" in interpolation
+
+    if instance_attr:
+        if not interpolation.startswith("${instance_attr:"):
+            interpolation = interpolation.removeprefix("${")
+            interpolation = "${instance_attr:" + interpolation
+
+    if default is MISSING and default_factory is MISSING:
+        raise RuntimeError(
+            "Interpolated fields currently still require a default value or default factory for "
+            "when they are used outside the Hydra/OmegaConf context."
+        )
+    return field(
+        default_factory=functools.partial(
+            _default_factory,
+            interpolation=interpolation,
+            default=default,
+            default_factory=default_factory,
+        )
+    )
+
+
 # @dataclass(init=False)
 class Partial(functools.partial[T], _Partial[T]):
     def __getattr__(self, name: str):
@@ -262,52 +342,6 @@ def get_instantiated_attr(
     )
 
 
-def interpolated_field(
-    interpolation: str,
-    default: T | Literal[MISSING] = MISSING,
-    default_factory: Callable[[], T] | Literal[MISSING] = MISSING,
-    instance_attr: bool = False,
-) -> T:
-    """Field with a default value computed with a OmegaConf-style interpolation when appropriate.
-
-    When the dataclass is created by Hydra / OmegaConf, the interpolation is used.
-    Otherwise, behaves as usual (either using default or calling the default_factory).
-
-    Parameters
-    ----------
-    interpolation: The string interpolation to use to get the default value.
-    default: The default value to use when not in a hydra/OmegaConf context.
-    default_factory: The default value to use when not in a hydra/OmegaConf context.
-    instance_attr: Whether to use the `instance_attr` custom resolver to run the interpolation \
-        with respect to instantiated objects instead of their configs.
-        Passing `interpolation='${instance_attr:some_config.some_attr}'` has the same effect.
-
-    This last parameter is important, since in order to retrieve the instance attribute, we need to
-    instantiate the objects, which could be expensive. These instantiated objects are reused at
-    least, but still, be mindful when using this parameter.
-    """
-    assert "${" in interpolation and "}" in interpolation
-
-    if instance_attr:
-        if not interpolation.startswith("${instance_attr:"):
-            interpolation = interpolation.removeprefix("${")
-            interpolation = "${instance_attr:" + interpolation
-
-    if default is MISSING and default_factory is MISSING:
-        raise RuntimeError(
-            "Interpolated fields currently still require a default value or default factory for "
-            "when they are used outside the Hydra/OmegaConf context."
-        )
-    return field(
-        default_factory=functools.partial(
-            _default_factory,
-            interpolation=interpolation,
-            default=default,
-            default_factory=default_factory,
-        )
-    )
-
-
 def being_called_in_hydra_context() -> bool:
     import hydra.core.utils
     import omegaconf._utils

project/utils/testutils.py

@@ -22,7 +22,7 @@
 from torch import Tensor, nn
 from torch.optim import Optimizer
 
-from project.configs.config import Config, cs
+from project.configs import Config, cs
 from project.configs.datamodule import DATA_DIR, SLURM_JOB_ID
 from project.datamodules.image_classification import (
     ImageClassificationDataModule,
@@ -125,11 +125,24 @@ def _parametrized_fixture_method(request: pytest.FixtureRequest):
 
 
 def get_all_configs_in_group(group_name: str) -> list[str]:
-    names_yaml = cs.list(group_name)
-    names = [name.rpartition(".")[0] for name in names_yaml]
-    if "base" in names:
-        names.remove("base")
-    return names
+    # note: here we're copying a bit of the internal code from Hydra so that we also get the
+    # configs that are just yaml files, in addition to the configs we added programmatically to the
+    # configstores.
+
+    # names_yaml = cs.list(group_name)
+    # names = [name.rpartition(".")[0] for name in names_yaml]
+    # if "base" in names:
+    #     names.remove("base")
+    # return names
+
+    from hydra._internal.config_loader_impl import ConfigLoaderImpl
+    from hydra._internal.utils import create_automatic_config_search_path
+
+    search_path = create_automatic_config_search_path(
+        calling_file=None, calling_module=None, config_path="pkg://project.configs"
+    )
+    config_loader = ConfigLoaderImpl(config_search_path=search_path)
+    return config_loader.get_group_options(group_name)
 
 
 def get_all_algorithm_names() -> list[str]:

pyproject.toml

@@ -63,6 +63,7 @@ build-backend = "setuptools.build_meta"
 
 [tool.pytest.ini_options]
 testpaths = ["project"]
+addopts = ["--doctest-modules"]
 
 [tool.ruff]
 line-length = 99