Recommended way to specialize

[darsnack]

I'm migrating to tyro from Hydra, and after spending a lot of time reading through issues + docs, I still can't figure out the recommended way to specialize.

Specifically, in Hydra's docs, they go over specializing on dataset + model which is a fairly common pattern. Currently I have the following:

@dataclass
class Config:
    """Main configuration for train-bp script."""
    data: ALL_DATA_CFGS = field(default_factory=MNISTConfig)
    optimizer: ALL_OPT_CFGS = field(default_factory=lambda: SGDConfig(learning_rate=0.01))
    logger: ALL_LOGGER_CFGS = field(default_factory=lambda: WandbLoggerConfig())
    model: ALL_MODEL_CFGS = field(default_factory=ConditionedMLPConfig)

If isinstance(cfg.model, CNNConfig) and isinstance(cfg.data, MNISTConfig), then I want to do things like set cfg.optimizer.learning_rate to specific value, etc. Currently, I tried to do this in __post_init__, but this clearly seems like the wrong place. Namely, I might have something like cfg.model.features that is unknown until the dataset is specified, but overwriting this in __post_init__ does not help since tyro flags it as a missing argument that's required.

[darsnack]

It also occurs to me that __post_init__ isn't great when you want to also be able to override the learning rate at the CLI. This is what tyro.extras.overridable_config_clis seems to be solving, but I don't think it is sufficient for the use case served by Hydra's specialization. You would need a separate "experiment" for each dataset and model combination which seems like defining special naming convention that experiments must follow.

[brentyi]

Hi @darsnack!

We don't generally have specific recommendations for things like this. The ideal we're striving for is your config objects are just code so there are many options and you can do whatever's aligned with your personal preferences.

Some options:

• Having some presets that can be selected/overridden seems like it would solve some of your issues: https://brentyi.github.io/tyro/examples/overriding_configs/#choosing-base-configs
  • For example, having an MNSIT + CNN preset with the default learning rate you want.
• I personally prefer to avoid __post_init__ and/or dynamically mutating configuration objects, but lots of projects make it work. If you want to make it explicit you could make a populate_default_fields() method that sets all of the "dependent" fields. You could work around the required argument problem with patterns like features: int = -1 or features: int | None = None as defaults.
• If possible, I'd personally try to structure the config objects so there aren't any dataset-dependent fields in the model config, etc. And if there's a field that's dependent on others I'd see if it can be implemented as a method or @property instead.

Let me know if anything's unclear!

[brentyi]

( I wrote the above before I saw your followup comment. but I think overall it still holds, lmk if not! )

[darsnack]

Hey sorry for the slow response, but I played around with a couple alternatives last week and came up with the following solution. Unfortunately, I don't think that the options suggested above would really meet the feature that I am seeking.

First, I create a "lazy missing" sentinel class. I think this is a fundamental feature in Hydra that would be useful to have in Tyro. It allows the user to mark something as required but defer the error until runtime access of a field. For me, I ended up creating such a class and adding the following rule to Tyro's registry so it didn't complain about the types:

_lazy_registry = tyro_constructors.ConstructorRegistry()

@_lazy_registry.primitive_rule
def _lazy_missing_rule(type_info: tyro_constructors.PrimitiveTypeInfo):
    """Intercept ``Union[T, LazyMissing]`` (produced by :func:`Lazy`) and
    present it to tyro as plain ``T``.

    The returned spec delegates entirely to ``T``'s own primitive spec but
    additionally accepts ``LAZY_MISSING`` as a valid instance so that the
    sentinel can serve as a default value without tripping validation.
    """
    # Only catch Unions with LazyMissing
    if get_origin(type_info.type) is not Union:
        return None
    args = get_args(type_info.type)
    non_lazy = [a for a in args if not isinstance(a, LazyMissing)]
    if len(non_lazy) == len(args):
        return None
    # Reconstruct the inner type, unwrapping a single-element tuple.
    inner_type = non_lazy[0] if len(non_lazy) == 1 else Union[tuple(non_lazy)]
    inner_info = tyro_constructors.PrimitiveTypeInfo.make(inner_type, type_info.markers)
    inner_spec = tyro_constructors.ConstructorRegistry.get_primitive_spec(inner_info)
    if isinstance(inner_spec, tyro_constructors.UnsupportedTypeAnnotationError):
        # Inner type is not a primitive (e.g. a nested dataclass) — let
        # tyro's default struct handling take over.
        return None

    return tyro_constructors.PrimitiveConstructorSpec(
        nargs=inner_spec.nargs,
        metavar=inner_spec.metavar,
        instance_from_str=inner_spec.instance_from_str,
        is_instance=lambda val, _is=inner_spec.is_instance: (
            True if isinstance(val, LazyMissing) else _is(val)
        ),
        str_from_instance=lambda val, _str=inner_spec.str_from_instance: (
            [] if isinstance(val, LazyMissing) else _str(val)
        ),
        choices=inner_spec.choices,
    )

The second step was adding a helper for a two pass CLI. Start by creating a config object (via tyro.cli) without command line overrides such as --model.features 20 10 1, keeping only subcommand args. Then run a user specified specialization function that replaces certain fields in the config. Finally, run tyro.cli again, using the specialized config as the default and including all command line args.

def specialize_cli(config_type: type[T], specializer: Callable[[T], T], *,
                   args: Sequence[str] | None = None, **tyro_kwargs: Any) -> T:
    """Two-pass CLI parser that applies specialization defaults before user
    overrides.
    """
    argv = list(args) if args is not None else sys.argv[1:]
    # Merge the lazy-sentinel registry with any caller-supplied registry.
    # We will use context for _lazy_registry and pass in the user registry.
    _registry = tyro_kwargs.pop("registry", None)
    # Pass 1: keep only subcommand selectors (tokens containing ':').
    # Plain bare tokens like '16' or '32' are positional values belonging
    # to --flags and must NOT be forwarded.
    bare_argv = [a for a in argv if ":" in a and not a.startswith("-")]
    with _lazy_registry:
        skeleton = tyro.cli(config_type,
                            args=bare_argv, console_outputs=False,
                            registry=_registry,
                            **tyro_kwargs)
    # Apply specialization to fill in LAZY_MISSING / computed defaults.
    specialized = specializer(skeleton)
    # Pass 2: full parse with specialized defaults.
    with _lazy_registry:
        return tyro.cli(config_type,
                        default=specialized, args=argv,
                        registry=_registry,
                        **tyro_kwargs)

While I think the second step could be something to prefer users to manually do, I think a tyro.LAZY_MISSING would be helpful to have in the library.

---

An alternative I considered was specialized experiments with a generic type that includes the data / model. I think this would be cleanest except that it doesn't really scale beyond a few type parameters. More complex specialization is just cleaner with a simple if/else block---part of my reason for adopting Tyro was avoiding complex logic implemented in YAML via Hydra.

Even if none of this makes it into Tyro, I think having a doc page with possible patterns is useful for new users. I agree that Tyro doesn't need to take a stance of what's best, but for anyone not thinking about configuration deeply, a suite of options is a good starting point for adapting something to their needs.

[darsnack]

(if you have any suggestions for how the above could be better, let me know too!)

[brentyi]

Thanks, appreciate the update + feedback! Glad you found something that works, it seems reasonable given the goal of approximating specialization in Hydra.

I've thought about adding support for more sentinels. It may... require some patience, but one thing I'm ~waiting for is PEP 661: https://peps.python.org/pep-0661/