Visibility of parameter options across subcommands with defaults

[nishi951]

Hi Brent,

First off, I'm really liking this framework! I have a use case that kind of combines "base configs as subcommands" with "sequenced subcommands".

Say I have a module that has two submodules, A and B. Furthermore, say each submodule has several possible "typical" configurations, e.g. A1, A2..., B1, B2,...

What I would like to do is simultaneously:

1. Set up base configs for all combinations of the typical configs for both A and B, without having to enumerate all combinations e.g. A1B1, A1B2, etc...
2. View, from -h, the possible options for both submodules A and B.

Is there currently a way of doing this? I've attached 2 examples. The first one sets up all base configs for both, but doesn't list all options with -h (it only lists options for the most recent subcommand). The second one will display all the possible parameter options for both A and B with -h (after one of the subcommands is specified).

I'm not even sure if what I'm trying to do is possible in a "subcommand" sense? I've also tried the AvoidSubcommands type but I can't really make that work either.

Thanks,
Mark

a.py:

from dataclasses import dataclass
from typing import Annotated, Union

import tyro
from tyro.conf import subcommand

@dataclass(frozen=True)
class SubModuleAConfig:
    param1: float
submoda_defaults = {
    'basic': SubModuleAConfig(param1=1.),
    'fancy': SubModuleAConfig(param1=2.2),
}
submoda_descriptions = {
    'basic': 'Basic config',
    'fancy': 'Fancy config'
}
SubModuleADefaultsType = tyro.extras.subcommand_type_from_defaults(
    submoda_defaults, submoda_descriptions
)

@dataclass(frozen=True)
class SubModuleBConfig:
    param2: int
submodb_defaults = {
    'basic': SubModuleBConfig(param2=0),
    'fancy': SubModuleBConfig(param2=-5),
}
submodb_descriptions = {
    'basic': 'Basic config',
    'fancy': 'Fancy config'
}
SubModuleBDefaultsType = tyro.extras.subcommand_type_from_defaults(
    submodb_defaults, submodb_descriptions
)

@dataclass
class FullModuleConfig:
    suba: SubModuleADefaultsType
    subb: SubModuleBDefaultsType

if __name__ == '__main__':
    full_module_config = tyro.cli(FullModuleConfig)
    print(full_module_config)

Output:

$ python a.py suba:basic subb:basic -h
usage: a.py suba:basic subb:basic [-h] [--subb.param2 INT]

Basic config

╭─ arguments ─────────────────────────────────────────────╮
│ -h, --help              show this help message and exit │
╰─────────────────────────────────────────────────────────╯
╭─ subb arguments ────────────────────────────────────────╮
│ --subb.param2 INT       (default: 0)                    │
╰─────────────────────────────────────────────────────────╯

b.py

from dataclasses import dataclass
from itertools import product
from typing import Annotated, Union

import tyro
from tyro.conf import subcommand

@dataclass(frozen=True)
class SubModuleAConfig:
    param1: float
submoda_defaults = {
    'basic': SubModuleAConfig(param1=1.),
    'fancy': SubModuleAConfig(param1=2.2),
}
submoda_descriptions = {
    'basic': 'Basic config',
    'fancy': 'Fancy config'
}

@dataclass(frozen=True)
class SubModuleBConfig:
    param2: int
submodb_defaults = {
    'basic': SubModuleBConfig(param2=0),
    'fancy': SubModuleBConfig(param2=-5),
}
submodb_descriptions = {
    'basic': 'Basic config',
    'fancy': 'Fancy config'
}

@dataclass
class FullModuleConfig:
    suba: SubModuleAConfig
    subb: SubModuleBConfig

all_defaults = {}
all_descriptions = {}
combos = product(submoda_defaults.items(), submodb_defaults.items())
for (suba_name, suba_config), (subb_name, subb_config) in combos:
    name = f'A{suba_name}_B{subb_name}'
    all_defaults[name] = FullModuleConfig(
        suba=suba_config,
        subb=subb_config,
    )
    all_descriptions[name] = f'A: {submoda_descriptions[suba_name]}, ' \
        + f'B: {submodb_descriptions[subb_name]}'

if __name__ == '__main__':
    full_module_config = tyro.cli(
        tyro.extras.subcommand_type_from_defaults(
            all_defaults,
            all_descriptions,
        )
    )
    print(full_module_config)

Output:

$ python b.py Abasic_Bbasic -h
usage: b.py Abasic_Bbasic [-h] [--suba.param1 FLOAT] [--subb.param2 INT]

A: Basic config, B: Basic config

╭─ arguments ─────────────────────────────────────────────╮
│ -h, --help              show this help message and exit │
╰─────────────────────────────────────────────────────────╯
╭─ suba arguments ────────────────────────────────────────╮
│ --suba.param1 FLOAT     (default: 1.0)                  │
╰─────────────────────────────────────────────────────────╯
╭─ subb arguments ────────────────────────────────────────╮
│ --subb.param2 INT       (default: 0)                    │
╰─────────────────────────────────────────────────────────╯

[brentyi]

Hi Mark,

Thanks for giving tyro a try!

My guess is that this is possible with some argparse hacking (tyro.extras.get_parser() might be helpful), but doing so in a general purpose way would likely be fairly involved.

The main reason is that each flag is only applied to the most recent preceding subcommand. This is how most CLI conventions (consider: git commit -a works, but git -a commit doesn't because -a is applied to the commit subcommand) are set up, and how argparse is designed.

A question/suggestion pair below -- please let me know if any of it's unclear! (I'm overseas at the moment though, so responses might be slow)

(1)
In your first example, would moving the help flag get you partially what you want?

So while running

python a.py suba:basic subb:basic -h

would print the help message for subb:basic, you could also run

python a.py suba:basic -h subb:basic

to print the help message and options for suba:basic.

(2)
If you don't care about tab completion script generation (for bash/zsh etc), you can get all of the options in one place by eschewing the standard argparse subcommands and implementing your own "base config" logic with tyro.cli()'s default parameter.

The easiest option would likely be with environment variables: you might pull SUBA={basic,fancy} and SUBB={basic,fancy} from os.environ, construct a FullModelConfig object on the fly with them, and then pass that in via something like:

tyro.cli(
    tyro.conf.AvoidSubcommands[FullModelConfig],
    default=FullModelConfig(
        suba=submoda_defaults[os.environ['SUBA']],
        subb=submodb_defaults[os.environ['SUBB']],
    )
)

To avoid environment variables, you can also do some manual sys.argv parsing before tyro.cli() is called. For the single "subcommand" case, see:
https://github.com/brentyi/tyro/blob/3c3167f4c424ee3ed074fd1122815ee4a8698f8b/examples/06_base_configs.py

[nishi951]

Yes, actually I've been making do with suggestion (1) for now! In some ways, I am starting to prefer this, since it makes the help screens a bit shorter and more specific. However, (2) is a fantastic suggestion that I will probably explore (leaning towards the manual sys.argv parsing), since I haven't been using the tab completion as much.

Thanks for your help!