Method test_flag_utils

axlearn/cloud/common/utils_test.py:440–526

Tests define_flags and from_flags.

(self)

        self.assertEqual(fv.shared_override, "child-default")

    def test_flag_utils(self):
        """Tests define_flags and from_flags."""

        class Inner(utils.FlagConfigurable):
            """An inner config."""

            @config_class
            class Config(utils.FlagConfigurable.Config):
                common_value: Required[str] = REQUIRED
                inner_value: Required[str] = REQUIRED

            @classmethod
            def define_flags(cls, fv):
                super().define_flags(fv)
                common_kwargs = dict(flag_values=fv, allow_override=True)
                flags.DEFINE_string("common_value", None, "", **common_kwargs)
                flags.DEFINE_string("inner_value", None, "", **common_kwargs)

            @classmethod
            def set_defaults(cls, fv):
                super().set_defaults(fv)
                fv.set_default("common_value", "child-default")

        # Test that it traverses non FlagConfigurables.
        class RegularConfigurable(Configurable):
            """A dummy container config to test traversal."""

            @config_class
            class Config(Configurable.Config):
                # Test that it traverses other non-config containers.
                inner: list[Inner.Config] = [Inner.default_config()]

        class Outer(utils.FlagConfigurable):
            """An outer config."""

            @config_class
            class Config(utils.FlagConfigurable.Config):
                common_value: Required[str] = REQUIRED
                outer_value: Required[str] = REQUIRED
                inner_enabled: Optional[bool] = None
                inner: Optional[Configurable.Config] = RegularConfigurable.default_config()

            @classmethod
            def define_flags(cls, fv):
                super().define_flags(fv)
                common_kwargs = dict(flag_values=fv, allow_override=True)
                flags.DEFINE_string("common_value", None, "", **common_kwargs)
                flags.DEFINE_string("outer_value", None, "", **common_kwargs)
                flags.DEFINE_bool("inner_enabled", None, "", **common_kwargs)

            @classmethod
            def set_defaults(cls, fv):
                super().set_defaults(fv)
                fv.set_default("common_value", "parent-default")

            @classmethod
            def from_flags(cls, fv, **kwargs):
                cfg = super().from_flags(fv, **kwargs)