Move get_named_trainer_config from experiments to common

GitOrigin-RevId: 3b4449a5a642b42aa32ffea2d4daa85506091b0c

Author: Yi Wang

axlearn/common/config.py

@@ -60,7 +60,6 @@ class Config(ConfigBase):
 config fields with mutable values, including config values.
 """
 
-# Note: config.py should not depend on jax, torch, or tf.
 import copy
 import dataclasses
 import enum
@@ -83,6 +82,9 @@ class Config(ConfigBase):
 # to apply validation on field names and values.
 import attr
 
+# Note: config.py should not depend on jax, torch, or tf.
+from absl import logging
+
 
 def is_named_tuple(x: Any):
     """Returns whether an object is an instance of a collections.namedtuple.
@@ -210,6 +212,32 @@ def validate_config_field_name(name: str) -> None:
         raise InvalidConfigNameError(f'Invalid config field name "{name}"')
 
 
+def validate_config_field_value(value: Any) -> None:
+    """Validates a config field value.
+
+    Validation is handled by validators registered via `register_validator`. `match_fn`s will be
+    invoked in order of registration, and all matched `validate_fn`s will be invoked.
+
+    Args:
+        value: The value to be validated.
+
+    Raises:
+        InvalidConfigValueError: If no validator matched the given value.
+    """
+    matched = False
+    for match_fn, validate_fn in _config_field_validators.items():
+        if match_fn(value):
+            matched = True
+            validate_fn(value)
+
+    # No validators matched.
+    if not matched:
+        raise InvalidConfigValueError(
+            f'Invalid config value type {type(value)} for value "{value}". '
+            f"Consider registering a custom validator with `{register_validator.__name__}`."
+        )
+
+
 # Validate basic types.
 register_validator(
     match_fn=lambda v: (
@@ -280,32 +308,6 @@ def _maybe_register_optional_type(module: str, attribute: str):
 _maybe_register_optional_type("jax.sharding", "PartitionSpec")
 
 
-def validate_config_field_value(value: Any) -> None:
-    """Validates a config field value.
-
-    Validation is handled by validators registered via `register_validator`. `match_fn`s will be
-    invoked in order of registration, and all matched `validate_fn`s will be invoked.
-
-    Args:
-        value: The value to be validated.
-
-    Raises:
-        InvalidConfigValueError: If no validator matched the given value.
-    """
-    matched = False
-    for match_fn, validate_fn in _config_field_validators.items():
-        if match_fn(value):
-            matched = True
-            validate_fn(value)
-
-    # No validators matched.
-    if not matched:
-        raise InvalidConfigValueError(
-            f'Invalid config value type {type(value)} for value "{value}". '
-            f"Consider registering a custom validator with `{register_validator.__name__}`."
-        )
-
-
 def _validate_and_transform_field(instance, attribute, value):
     """Validates an attribute as a config field.
 
@@ -1055,3 +1057,52 @@ class ConfigModifier(Configurable):
     def __call__(self, cfg: InstantiableConfig[T]) -> InstantiableConfig[T]:
         """A function that modifies the input config, should be defined by subclasses."""
         return cfg
+
+
+def _load_trainer_configs(
+    config_module: str, *, optional: bool = False
+) -> dict[str, TrainerConfigFn]:
+    try:
+        module = importlib.import_module(config_module)
+        return module.named_trainer_configs()
+    except (ImportError, AttributeError):
+        if not optional:
+            raise
+        logging.warning(
+            "Missing dependencies for %s but it's marked optional -- skipping.", config_module
+        )
+    return {}
+
+
+def get_named_trainer_config(config_name: str, *, config_module: str) -> TrainerConfigFn:
+    """Looks up TrainerConfigFn by config name.
+
+    Args:
+        config_name: Candidate config name.
+        config_module: Config module name.
+
+    Returns:
+        A TrainerConfigFn corresponding to the config name.
+
+    Raises:
+        KeyError: Error containing the message to show to the user.
+    """
+    config_map = _load_trainer_configs(config_module)
+    if callable(config_map):
+        return config_map(config_name)
+
+    try:
+        return config_map[config_name]
+    except KeyError as e:
+        similar = similar_names(config_name, set(config_map.keys()))
+        if similar:
+            message = f"Unrecognized config {config_name}; did you mean [{', '.join(similar)}]"
+        else:
+            message = (
+                f"Unrecognized config {config_name} under {config_module}; "
+                f"Please make sure that the following conditions are met:\n"
+                f"    1. {config_module} can be imported; "
+                f"    2. {config_module} defines `named_trainer_configs()`; "
+                f"    3. `named_trainer_configs()` returns a dict with '{config_name}' as a key."
+            )
+        raise KeyError(message) from e

axlearn/common/launch_trainer.py

@@ -11,10 +11,9 @@
 
 from axlearn.common import file_system as fs
 from axlearn.common import measurement
-from axlearn.common.config import TrainerConfigFn
+from axlearn.common.config import TrainerConfigFn, get_named_trainer_config
 from axlearn.common.trainer import SpmdTrainer, select_mesh_config
 from axlearn.common.utils import MeshShape, get_data_dir, infer_mesh_shape
-from axlearn.experiments import get_named_trainer_config
 
 # Trainer-specific flags.
 flags.DEFINE_string(

axlearn/experiments/__init__.py

@@ -1,58 +1,3 @@
 # Copyright © 2023 Apple Inc.
 
 """AXLearn experiments."""
-
-from importlib import import_module
-
-from absl import logging
-
-from axlearn.common.config import TrainerConfigFn, similar_names
-
-
-def _load_trainer_configs(
-    config_module: str, *, optional: bool = False
-) -> dict[str, TrainerConfigFn]:
-    try:
-        module = import_module(config_module)
-        return module.named_trainer_configs()
-    except (ImportError, AttributeError):
-        if not optional:
-            raise
-        logging.warning(
-            "Missing dependencies for %s but it's marked optional -- skipping.", config_module
-        )
-    return {}
-
-
-def get_named_trainer_config(config_name: str, *, config_module: str) -> TrainerConfigFn:
-    """Looks up TrainerConfigFn by config name.
-
-    Args:
-        config_name: Candidate config name.
-        config_module: Config module name.
-
-    Returns:
-        A TrainerConfigFn corresponding to the config name.
-
-    Raises:
-        KeyError: Error containing the message to show to the user.
-    """
-    config_map = _load_trainer_configs(config_module)
-    if callable(config_map):
-        return config_map(config_name)
-
-    try:
-        return config_map[config_name]
-    except KeyError as e:
-        similar = similar_names(config_name, set(config_map.keys()))
-        if similar:
-            message = f"Unrecognized config {config_name}; did you mean [{', '.join(similar)}]"
-        else:
-            message = (
-                f"Unrecognized config {config_name} under {config_module}; "
-                f"Please make sure that the following conditions are met:\n"
-                f"    1. {config_module} can be imported; "
-                f"    2. {config_module} defines `named_trainer_configs()`; "
-                f"    3. `named_trainer_configs()` returns a dict with '{config_name}' as a key."
-            )
-        raise KeyError(message) from e

axlearn/experiments/golden_ckpt_test.py

@@ -13,11 +13,11 @@
 from absl import flags
 from absl.testing import absltest
 
+from axlearn.common.config import get_named_trainer_config
 from axlearn.common.inference import InferenceRunner
 from axlearn.common.summary_writer import NoOpWriter
 from axlearn.common.test_utils import TestCase
 from axlearn.common.trainer import SpmdTrainer
-from axlearn.experiments import get_named_trainer_config
 
 flags.DEFINE_boolean("update_golden_checkpoints", False, "If true, update golden config files.")
 

axlearn/experiments/run_aot_compilation.py

@@ -40,11 +40,10 @@
 from jax.experimental.serialize_executable import serialize
 
 from axlearn.common import aot_compilation, compiler_options
-from axlearn.common.config import TrainerConfigFn
+from axlearn.common.config import TrainerConfigFn, get_named_trainer_config
 from axlearn.common.trainer import SpmdTrainer, aot_model_analysis, select_mesh_config
 from axlearn.common.utils import set_data_dir
 from axlearn.common.utils_spmd import setup
-from axlearn.experiments import get_named_trainer_config
 
 flags.DEFINE_string("module", None, "The trainer config module.", required=True)
 flags.DEFINE_string("config", None, "The trainer config name.", required=True)