Function define_flags

axlearn/cloud/common/utils.py:489–532

Define flags on `fv` by recursively invoking `cfg.klass.define_flags`. Flags are defined in topological order, i.e., parent flags will be defined prior to defining child flags. Args: cfg: A config instance. It does not need to be a `FlagConfigurable` config. fv: Parsed

(cfg: ConfigBase, fv: flags.FlagValues)

def define_flags(cfg: ConfigBase, fv: flags.FlagValues):
    """Define flags on `fv` by recursively invoking `cfg.klass.define_flags`.

    Flags are defined in topological order, i.e., parent flags will be defined prior to defining
    child flags.

    Args:
        cfg: A config instance. It does not need to be a `FlagConfigurable` config.
        fv: Parsed flag values instance. The same `fv` is used for the entire config hierarchy.

    Raises:
        ValueError: If `cfg` is not a config.
    """
    # pylint: disable=protected-access
    if not isinstance(cfg, ConfigBase):
        raise ValueError(f"Expected {ConfigBase}, got: {type(cfg)}")

    def visit_fn(_, value):
        if not isinstance(value, FlagConfigurable.Config):
            return
        for namespace, child in _get_namespaced_config(value).items():
            sub_fv = flags.FlagValues()
            define_flags(child, sub_fv)
            # Flatten the child flags into `fv`. See `FlagValues.append_flag_values` for ref.
            # The main difference is that we namespace the flags by child name.
            for flag_name, flag in sub_fv._flags().items():
                # absl flattens short names into `fv` during __setattr__.
                # Keep things simple for now by limiting to verbose names.
                if flag.short_name:
                    raise NotImplementedError(
                        f"Short names are currently not supported: {flag.short_name}"
                    )
                fv[f"{namespace}.{flag_name}"] = flag

    def enter_fn(_, value, default_kv):
        if not isinstance(value, Configurable.Config) or not hasattr(value.klass, "define_flags"):
            return default_kv
        klass: FlagConfigurable = value.klass
        klass.define_flags(fv)
        if hasattr(klass, _FLAG_NAMESPACE_ATTRIBUTE):
            return None  # Enter visit_fn.
        return default_kv

    cfg.visit(visit_fn=visit_fn, enter_fn=enter_fn)


_C = TypeVar("_C", bound=ConfigBase)