refactor: generate train and serve CLIs from config-backed option specs

- Add astrai/config/cli.py: OptSpec tables plus apply_specs infer click types and defaults from config fields, covering Optional[X], Union[X, None], PEP 604 X | None, stringified PEP 563 annotations, bool flag pairs, and repeatable list options
- Move GroupedCommand/GroupedOption and the three-layer YAML merge (option defaults < YAML < explicit CLI) into the config package, adding unknown-key warning and mapping validation
- Replace ~420 lines of hand-written @opt decorators in scripts/tools/train.py with a 66-entry spec table; option names, defaults, flag styles, and YAML semantics verified unchanged
- Migrate scripts/tools/server.py to the same mechanism with its section binding, integer coercion, and dtype validation preserved locally
- Add tests/config/test_cli.py covering type inference across annotation styles, default overrides, flag pairs, merge precedence, scientific notation, and help ordering
This commit is contained in:
2026-09-03 20:56:11 +08:00
parent 853aaeefde
commit 76f1c10feb
5 changed files with 953 additions and 598 deletions
+14
View File
@@ -1,3 +1,11 @@
from astrai.config.cli import (
GroupedCommand,
GroupedOption,
OptSpec,
apply_specs,
merge_yaml_into_kwargs,
opt,
)
from astrai.config.model_config import ( from astrai.config.model_config import (
AutoRegressiveLMConfig, AutoRegressiveLMConfig,
BaseModelConfig, BaseModelConfig,
@@ -22,4 +30,10 @@ __all__ = [
"OutputConfig", "OutputConfig",
"PipelineConfig", "PipelineConfig",
"ProcessingConfig", "ProcessingConfig",
"GroupedCommand",
"GroupedOption",
"OptSpec",
"apply_specs",
"merge_yaml_into_kwargs",
"opt",
] ]
+315
View File
@@ -0,0 +1,315 @@
"""Generate grouped click CLIs from pydantic config fields.
A config class alone does not make a CLI: some options need defaults that
differ from the config defaults (e.g. ``num_workers``), choices come from
factory registries or frozenset validators, and values merge across three
layers (option defaults -> YAML -> explicit CLI flags). ``OptSpec`` records
those overrides in a declarative table; types and defaults are inferred
from the backing config field wherever the spec leaves them ``AUTO``.
"""
import dataclasses
import re
import types
import typing as t
from collections import OrderedDict
from collections.abc import Sequence
import click
import yaml
class GroupedOption(click.Option):
"""A ``click.Option`` that carries a ``group`` label for help output."""
def __init__(self, *args, group: str = "Options", **kwargs):
super().__init__(*args, **kwargs)
self.group = group
class GroupedCommand(click.Command):
"""A ``click.Command`` that renders options grouped by their ``group``."""
def format_options(self, ctx, formatter):
groups: OrderedDict[str, list] = OrderedDict()
for param in self.get_params(ctx):
record = param.get_help_record(ctx)
if record is None:
continue
group = getattr(param, "group", "Options")
groups.setdefault(group, []).append(record)
for group_name, records in groups.items():
with formatter.section(group_name):
formatter.write_dl(records)
def opt(*param_decls, group: str, **kwargs):
"""Shorthand for ``click.option`` that tags the option with a group."""
kwargs.setdefault("cls", GroupedOption)
kwargs["group"] = group
return click.option(*param_decls, **kwargs)
class _Auto:
def __repr__(self) -> str:
return "AUTO"
AUTO = _Auto()
_YAML_FLOAT_PATTERN = re.compile(
r"""^(?:[-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+]?[0-9]+)?
|[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
|[-+]?\.(?:inf|Inf|INF)
|\.(?:nan|NaN|NAN))$""",
re.X,
)
DEFAULT_YAML_SECTIONS = ("model", "data", "parallel", "training", "ckpt", "log")
def _enable_yaml12_floats() -> None:
"""PyYAML implements YAML 1.1, where ``2e-5`` parses as a string; switch its
float resolver to the YAML 1.2 core schema so scientific notation works."""
yaml.SafeLoader.add_implicit_resolver(
"tag:yaml.org,2002:float", _YAML_FLOAT_PATTERN, list("-+0123456789.")
)
def merge_yaml_into_kwargs(
config_path: str,
passed_kwargs: dict,
explicit_keys: set[str] | None = None,
sections: Sequence[str] = DEFAULT_YAML_SECTIONS,
allowed_keys: Sequence[str] | None = None,
) -> dict:
"""Merge option defaults, YAML values, then explicit CLI values.
``sections`` selects which top-level YAML mappings feed the flat kwargs
namespace. ``allowed_keys`` optionally restricts the accepted keys across
those sections: unknown keys are warned about once and dropped.
"""
_enable_yaml12_floats()
with open(config_path) as f:
cfg = yaml.safe_load(f) or {}
if not isinstance(cfg, dict):
raise click.UsageError(f"config must be a mapping: {config_path}")
merged = dict(passed_kwargs)
seen: set[str] = set()
for section in sections:
values = cfg.get(section) or {}
if not isinstance(values, dict):
raise click.UsageError(f"top-level {section} section must be a mapping")
if allowed_keys is None:
merged.update(values)
seen.update(values)
else:
merged.update({k: v for k, v in values.items() if k in allowed_keys})
seen.update(values)
if allowed_keys is not None:
unknown = sorted(seen - set(allowed_keys))
if unknown:
click.echo(
f"Warning: ignoring unknown config keys: {', '.join(unknown)}",
err=True,
)
if explicit_keys is None:
explicit_keys = set(passed_kwargs)
for key in explicit_keys:
if key in passed_kwargs:
merged[key] = passed_kwargs[key]
return merged
@dataclasses.dataclass(frozen=True)
class OptSpec:
"""One CLI option, optionally backed by a field of a config class.
``type``/``default`` left as ``AUTO`` are inferred from the backing
config field (bool becomes a ``--x/--no-x`` pair, lists become
repeatable options). Standalone specs for CLI-only options must carry
``type`` or ``default`` explicitly. ``choices`` overrides the inferred
type with ``click.Choice``.
"""
name: str
group: str
type: t.Any = AUTO
default: t.Any = AUTO
help: str | None = None
choices: Sequence[str] | None = None
multiple: bool = False
is_flag: bool = False
required: bool = False
param_decls: tuple[str, ...] = ()
def _resolve_hints(config_cls: type) -> dict[str, t.Any]:
"""Resolve field annotations with ``typing.get_type_hints``.
Raw ``Field.type`` stays a string under PEP 563 (``from __future__
import annotations``) and never mentions ``types.UnionType``; resolved
hints cover old-style ``Optional[X]``/``Union[X, None]``, new-style
``X | None``, and stringified forward references alike.
"""
try:
return t.get_type_hints(config_cls)
except Exception:
return {}
def _unwrap_optional(annotation: t.Any) -> t.Any | None:
"""Return the single non-None member of an optional annotation."""
if annotation is None or isinstance(annotation, str):
return None
origin = t.get_origin(annotation)
if origin is not t.Union and origin is not types.UnionType:
return None
args = [a for a in t.get_args(annotation) if a is not type(None)]
if len(args) == 1:
return args[0]
return None
def _annotation(
spec: OptSpec,
field: dataclasses.Field | None,
hints: dict[str, t.Any] | None = None,
) -> t.Any:
if field is not None:
if hints and spec.name in hints:
return hints[spec.name]
return field.type
if spec.type is not AUTO:
return spec.type
return None
def _click_type(
spec: OptSpec,
field: dataclasses.Field | None,
hints: dict[str, t.Any] | None = None,
annotation: t.Any = None,
) -> t.Any:
if spec.choices is not None:
return click.Choice(list(spec.choices))
if annotation is None:
if spec.type is not AUTO and spec.type is not bool:
return spec.type
annotation = _annotation(spec, field, hints)
inner = _unwrap_optional(annotation)
if inner is not None:
return _click_type(spec, field, hints, annotation=inner)
origin = t.get_origin(annotation)
if annotation is bool or spec.type is bool:
return click.BOOL
if annotation is int:
return click.INT
if annotation is float:
return click.FLOAT
if annotation is str:
return click.STRING
if origin in (list, tuple) or annotation in (list, tuple):
return click.STRING
raise TypeError(
f"cannot infer a click type for option {spec.name!r}; set OptSpec.type"
)
def _resolve_default(spec: OptSpec, field: dataclasses.Field | None) -> t.Any:
if spec.default is not AUTO:
return spec.default
if field is not None:
if field.default is not dataclasses.MISSING:
return field.default
if field.default_factory is not dataclasses.MISSING: # type: ignore[misc]
return field.default_factory() # type: ignore[misc]
return None
def _is_flag_pair(
spec: OptSpec,
field: dataclasses.Field | None,
hints: dict[str, t.Any] | None = None,
) -> bool:
if spec.is_flag:
return False
annotation = _annotation(spec, field, hints)
return annotation is bool
def option_from_spec(
spec: OptSpec,
fields_by_name: dict[str, dataclasses.Field] | None = None,
hints: dict[str, t.Any] | None = None,
) -> t.Callable:
"""Build a grouped ``click.option`` decorator from one spec."""
fields_by_name = fields_by_name or {}
hints = hints or {}
field = fields_by_name.get(spec.name)
kwargs: dict[str, t.Any] = {
"cls": GroupedOption,
"group": spec.group,
}
if spec.help is not None:
kwargs["help"] = spec.help
if spec.required:
kwargs["required"] = True
if spec.is_flag:
decls = spec.param_decls or (f"--{spec.name}",)
kwargs["is_flag"] = True
kwargs["default"] = _resolve_default(spec, field)
elif _is_flag_pair(spec, field, hints):
if spec.param_decls:
raise ValueError(
f"flag pair {spec.name!r} does not support custom param_decls"
)
decls = (f"--{spec.name}/--no-{spec.name}",)
kwargs["type"] = click.BOOL
kwargs["default"] = _resolve_default(spec, field)
else:
decls = spec.param_decls or (f"--{spec.name}",)
kwargs["type"] = _click_type(spec, field, hints)
kwargs["default"] = _resolve_default(spec, field)
multiple = spec.multiple
if not multiple and field is not None:
annotation = _annotation(spec, field, hints)
origin = t.get_origin(annotation)
multiple = origin in (list, tuple) or annotation in (list, tuple)
if multiple:
kwargs["multiple"] = True
if isinstance(kwargs["default"], list):
kwargs["default"] = tuple(kwargs["default"])
return opt(*decls, **kwargs)
def apply_specs(
specs: Sequence[OptSpec],
config_cls: type | None = None,
) -> t.Callable:
"""Apply a table of specs as click options, in table order.
Options render top-to-bottom in the command's help output following the
table order (click reverses decorator application, so specs are applied
reversed). ``config_cls`` supplies type/default inference for specs whose
name matches one of its fields.
"""
fields_by_name: dict[str, dataclasses.Field] = {}
hints: dict[str, t.Any] = {}
if config_cls is not None:
fields_by_name = {f.name: f for f in dataclasses.fields(config_cls)}
hints = _resolve_hints(config_cls)
def decorator(func):
for spec in reversed(specs):
func = option_from_spec(spec, fields_by_name, hints)(func)
return func
return decorator
+85 -103
View File
@@ -2,9 +2,14 @@ from pathlib import Path
import click import click
import torch import torch
import yaml
from click.core import ParameterSource from click.core import ParameterSource
from astrai.config.cli import (
GroupedCommand,
OptSpec,
apply_specs,
merge_yaml_into_kwargs,
)
from astrai.inference import run_server from astrai.inference import run_server
_DTYPES = ["bfloat16", "float16", "float32"] _DTYPES = ["bfloat16", "float16", "float32"]
@@ -26,29 +31,65 @@ def _merge_yaml_into_kwargs(
explicit_keys: set[str] | None = None, explicit_keys: set[str] | None = None,
) -> dict: ) -> dict:
"""Merge Click defaults, YAML server values, then explicit CLI values.""" """Merge Click defaults, YAML server values, then explicit CLI values."""
with open(config_path, encoding="utf-8") as file: return merge_yaml_into_kwargs(
config = yaml.safe_load(file) or {} config_path,
if not isinstance(config, dict): passed_kwargs,
raise click.UsageError(f"Serving config must be a mapping: {config_path}") explicit_keys,
server = config.get("server") or {} sections=("server",),
if not isinstance(server, dict): allowed_keys=_SERVER_KEYS,
raise click.UsageError("top-level server section must be a mapping") )
unknown = sorted(set(server) - set(_SERVER_KEYS))
if unknown:
click.echo(
f"Warning: ignoring unknown server config keys: {', '.join(unknown)}",
err=True,
)
merged = dict(passed_kwargs) _SPECS = [
merged.update({key: server[key] for key in _SERVER_KEYS if key in server}) OptSpec(
if explicit_keys is None: "config_path",
explicit_keys = set(passed_kwargs) "Server",
for key in explicit_keys: type=click.Path(exists=True, dir_okay=False),
if key in passed_kwargs: param_decls=("--config", "-c", "config_path"),
merged[key] = passed_kwargs[key] help="Serving YAML config. CLI flags override YAML values.",
return merged ),
OptSpec("host", "Server", type=str, default="0.0.0.0", help="Host address."),
OptSpec("port", "Server", type=int, default=8000, help="Port number."),
OptSpec(
"reload",
"Server",
is_flag=True,
default=False,
help="Enable auto-reload for development.",
),
OptSpec(
"param_path",
"Model",
type=click.Path(exists=True),
default=None,
help="Path to model parameters.",
),
OptSpec(
"device", "Model", type=str, default="cuda", help="Device to load model on."
),
OptSpec(
"dtype",
"Model",
choices=_DTYPES,
default="bfloat16",
help="Data type for model weights.",
),
OptSpec(
"max_batch_size",
"Performance",
type=int,
default=16,
help="Maximum batch size for continuous batching.",
),
OptSpec(
"max_seq_len",
"Performance",
type=int,
default=None,
help="Maximum sequence length (KV cache size + prompt truncation). "
"Uses model config if not set.",
),
]
def _as_int(value, name: str) -> int | None: def _as_int(value, name: str) -> int | None:
@@ -87,84 +128,22 @@ def _resolve_server_config(
return resolved return resolved
@click.command(name="serve", help="Launch inference server (OpenAI-compatible API).") @click.command(
@click.option( name="serve",
"--config", cls=GroupedCommand,
"-c", help="Launch inference server (OpenAI-compatible API).",
"config_path",
type=click.Path(exists=True, dir_okay=False),
default=None,
help="Serving YAML config. CLI flags override YAML values.",
)
@click.option("--host", default="0.0.0.0", help="Host address.")
@click.option("--port", type=int, default=8000, help="Port number.")
@click.option(
"--reload", is_flag=True, default=False, help="Enable auto-reload for development."
)
@click.option(
"--param_path",
type=click.Path(exists=True),
default=None,
help="Path to model parameters.",
)
@click.option("--device", default="cuda", help="Device to load model on.")
@click.option(
"--dtype",
type=click.Choice(_DTYPES),
default="bfloat16",
help="Data type for model weights.",
)
@click.option(
"--max_batch_size",
type=int,
default=16,
help="Maximum batch size for continuous batching.",
)
@click.option(
"--max_seq_len",
type=int,
default=None,
help="Maximum sequence length (KV cache size + prompt truncation). Uses model config if not set.",
) )
@apply_specs(_SPECS)
@click.pass_context @click.pass_context
def server_command( def server_command(ctx, config_path, **kwargs):
ctx,
config_path,
host,
port,
reload,
param_path,
device,
dtype,
max_batch_size,
max_seq_len,
):
"""Launch inference server (OpenAI-compatible API).""" """Launch inference server (OpenAI-compatible API)."""
if config_path: if config_path:
passed_kwargs = {
"host": host,
"port": port,
"reload": reload,
"param_path": param_path,
"device": device,
"dtype": dtype,
"max_batch_size": max_batch_size,
"max_seq_len": max_seq_len,
}
explicit_keys = { explicit_keys = {
key key
for key in passed_kwargs for key in kwargs
if ctx.get_parameter_source(key) is ParameterSource.COMMANDLINE if ctx.get_parameter_source(key) is ParameterSource.COMMANDLINE
} }
resolved = _resolve_server_config(config_path, passed_kwargs, explicit_keys) kwargs = _resolve_server_config(config_path, kwargs, explicit_keys)
host = resolved["host"]
port = resolved["port"]
reload = resolved["reload"]
param_path = resolved["param_path"]
device = resolved["device"]
dtype = resolved["dtype"]
max_batch_size = resolved["max_batch_size"]
max_seq_len = resolved["max_seq_len"]
click.echo(f"Config: {config_path}") click.echo(f"Config: {config_path}")
dtype_map = { dtype_map = {
@@ -173,19 +152,22 @@ def server_command(
"float32": torch.float32, "float32": torch.float32,
} }
project_root = Path(__file__).parent.parent.parent project_root = Path(__file__).parent.parent.parent
param_path = param_path or str(project_root / "params") kwargs["param_path"] = kwargs.get("param_path") or str(project_root / "params")
click.echo(f"Starting server on http://{host}:{port}") click.echo(f"Starting server on http://{kwargs['host']}:{kwargs['port']}")
click.echo(f"Model: {param_path} | Device: {device} | Dtype: {dtype}") click.echo(
f"Model: {kwargs['param_path']} | "
f"Device: {kwargs['device']} | Dtype: {kwargs['dtype']}"
)
run_server( run_server(
host=host, host=kwargs["host"],
port=port, port=kwargs["port"],
reload=reload, reload=kwargs["reload"],
device=device, device=kwargs["device"],
dtype=dtype_map[dtype], dtype=dtype_map[kwargs["dtype"]],
param_path=Path(param_path), param_path=Path(kwargs["param_path"]),
max_batch_size=max_batch_size, max_batch_size=kwargs["max_batch_size"],
max_seq_len=max_seq_len, max_seq_len=kwargs["max_seq_len"],
) )
+309 -495
View File
@@ -1,16 +1,19 @@
import os import os
import re
from collections import OrderedDict
from collections.abc import Callable from collections.abc import Callable
from functools import partial from functools import partial
import click import click
import torch import torch
import yaml
from click.core import ParameterSource from click.core import ParameterSource
from torch import optim from torch import optim
from astrai.config import AutoRegressiveLMConfig, TrainConfig from astrai.config import AutoRegressiveLMConfig, TrainConfig
from astrai.config.cli import (
GroupedCommand,
OptSpec,
apply_specs,
merge_yaml_into_kwargs,
)
from astrai.config.train_config import ( from astrai.config.train_config import (
BACKENDS, BACKENDS,
PARALLEL_MODES, PARALLEL_MODES,
@@ -24,79 +27,8 @@ from astrai.optim import OptimizerFactory
from astrai.trainer import SchedulerFactory, Trainer from astrai.trainer import SchedulerFactory, Trainer
from astrai.trainer.rollout import BaseRewardModel from astrai.trainer.rollout import BaseRewardModel
# Re-exported under its historical name for tests importing it from here.
class GroupedOption(click.Option): _merge_yaml_into_kwargs = merge_yaml_into_kwargs
"""A ``click.Option`` that carries a ``group`` label for help output."""
def __init__(self, *args, group: str = "Options", **kwargs):
super().__init__(*args, **kwargs)
self.group = group
class GroupedCommand(click.Command):
"""A ``click.Command`` that renders options grouped by their ``group``."""
def format_options(self, ctx, formatter):
groups: OrderedDict[str, list] = OrderedDict()
for param in self.get_params(ctx):
record = param.get_help_record(ctx)
if record is None:
continue
group = getattr(param, "group", "Options")
groups.setdefault(group, []).append(record)
for group_name, records in groups.items():
with formatter.section(group_name):
formatter.write_dl(records)
def opt(*param_decls, group: str, **kwargs):
"""Shorthand for ``click.option`` that tags the option with a group."""
kwargs.setdefault("cls", GroupedOption)
kwargs["group"] = group
return click.option(*param_decls, **kwargs)
_YAML_FLOAT_PATTERN = re.compile(
r"""^(?:[-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+]?[0-9]+)?
|[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
|[-+]?\.(?:inf|Inf|INF)
|\.(?:nan|NaN|NAN))$""",
re.X,
)
def _enable_yaml12_floats() -> None:
"""PyYAML implements YAML 1.1, where ``2e-5`` parses as a string; switch its
float resolver to the YAML 1.2 core schema so scientific notation works."""
yaml.SafeLoader.add_implicit_resolver(
"tag:yaml.org,2002:float", _YAML_FLOAT_PATTERN, list("-+0123456789.")
)
def _merge_yaml_into_kwargs(
config_path: str,
passed_kwargs: dict,
explicit_keys: set[str] | None = None,
) -> dict:
"""Merge Click defaults, YAML values, then explicit CLI values."""
_enable_yaml12_floats()
with open(config_path) as f:
cfg = yaml.safe_load(f) or {}
merged = dict(passed_kwargs)
for section in ("model", "data", "parallel", "training", "ckpt", "log"):
if section in cfg:
merged.update(cfg[section])
if explicit_keys is None:
explicit_keys = set(passed_kwargs)
for key in explicit_keys:
if key in passed_kwargs:
merged[key] = passed_kwargs[key]
return merged
_TRAIN_TYPE = sorted(TRAIN_TYPES) _TRAIN_TYPE = sorted(TRAIN_TYPES)
_PARALLEL = sorted(PARALLEL_MODES) _PARALLEL = sorted(PARALLEL_MODES)
@@ -105,6 +37,306 @@ _OPTIMIZERS = OptimizerFactory.list_registered()
_BACKENDS = sorted(BACKENDS) _BACKENDS = sorted(BACKENDS)
_START_METHODS = sorted(START_METHODS) _START_METHODS = sorted(START_METHODS)
# Option table: types/defaults marked AUTO are inferred from TrainConfig
# fields; everything else (CLI-only options and default overrides) is
# declared inline. Table order is the --help order.
_SPECS = [
OptSpec(
"config_path",
"Paths & Setup",
type=click.Path(exists=True),
param_decls=("--config", "-c", "config_path"),
help="YAML config file. CLI flags override YAML values.",
),
OptSpec("train_type", "Paths & Setup", choices=_TRAIN_TYPE, help="Training type."),
OptSpec(
"data_root_path",
"Paths & Setup",
type=click.Path(exists=True),
help="Root directory of the dataset.",
),
OptSpec(
"param_path",
"Paths & Setup",
type=click.Path(exists=True),
help="Path to model parameters or resume checkpoint.",
),
OptSpec(
"resume",
"Paths & Setup",
is_flag=True,
default=False,
help="Resume from checkpoint.",
),
OptSpec("n_epoch", "Training", help="Number of epochs."),
OptSpec("batch_per_device", "Training", default=1, help="Batch size per GPU."),
OptSpec(
"grad_accum_steps",
"Training",
help="Gradient accumulation steps.",
),
OptSpec("max_grad_norm", "Training", help="Max gradient norm for clipping."),
OptSpec(
"warmup_ratio",
"LR Schedule",
type=float,
default=0.05,
help="Fraction of total steps for LR warmup.",
),
OptSpec(
"max_lr",
"Optimizer",
type=float,
default=3e-4,
help="Max learning rate.",
),
OptSpec(
"optimizer",
"Optimizer",
choices=_OPTIMIZERS,
default="muon_adamw",
help="Built-in optimizer.",
),
OptSpec(
"weight_decay",
"Optimizer",
type=float,
default=0.1,
help="Weight decay for eligible optimizer parameters.",
),
OptSpec(
"nora_lr", "Optimizer", type=float, default=5e-3, help="Nora learning rate."
),
OptSpec(
"nora_beta", "Optimizer", type=float, default=0.95, help="Nora EMA factor."
),
OptSpec(
"nora_momentum",
"Optimizer",
type=float,
default=0.95,
help="Nora update momentum.",
),
OptSpec(
"nora_weight_decay",
"Optimizer",
type=float,
default=0.0,
help="Nora weight decay.",
),
OptSpec(
"muon_momentum",
"Optimizer",
type=float,
default=0.95,
help="Muon momentum factor.",
),
OptSpec(
"muon_nesterov",
"Optimizer",
type=bool,
default=True,
help="Muon Nesterov.",
),
OptSpec(
"muon_ns_steps",
"Optimizer",
type=int,
default=5,
help="Muon Newton-Schulz steps.",
),
OptSpec(
"muon_adjust_lr",
"Optimizer",
choices=["original", "match_rms_adamw"],
default="match_rms_adamw",
help="Muon LR adjustment strategy.",
),
OptSpec(
"mano_momentum",
"Optimizer",
type=float,
default=0.95,
help="Mano momentum factor.",
),
OptSpec(
"mano_nesterov",
"Optimizer",
type=bool,
default=True,
help="Mano Nesterov momentum.",
),
OptSpec("random_seed", "Data Loading", help="Random seed."),
OptSpec("num_workers", "Data Loading", default=4, help="DataLoader workers."),
OptSpec("pin_memory", "Data Loading", default=True, help="Pin memory."),
OptSpec(
"persistent_workers",
"Data Loading",
default=True,
help="Keep DataLoader workers alive between epochs.",
),
OptSpec(
"window_size",
"Data Loading",
type=int,
default=None,
help="Max input sequence length.",
),
OptSpec(
"stride",
"Data Loading",
type=int,
default=None,
help="Step size for sliding window.",
),
OptSpec(
"label_smoothing",
"Data Loading",
type=float,
default=0.0,
help="Label smoothing.",
),
OptSpec("dpo_beta", "Algorithm", type=float, default=0.1, help="DPO beta."),
OptSpec("group_size", "Algorithm", type=int, default=4, help="GRPO group size."),
OptSpec(
"grpo_clip_eps",
"Algorithm",
type=float,
default=0.2,
help="GRPO clip epsilon.",
),
OptSpec(
"grpo_kl_coef",
"Algorithm",
type=float,
default=0.01,
help="GRPO KL penalty coefficient.",
),
OptSpec(
"moe_aux_loss_coef",
"Algorithm",
help="MoE load balancing auxiliary loss coefficient (0=disable).",
),
OptSpec("rollout_interval", "Algorithm", help="Steps between rollouts."),
OptSpec(
"rollout_max_policy_lag",
"Algorithm",
help="Maximum accepted rollout/live policy-version gap.",
),
OptSpec("rollout_temperature", "Algorithm", help="Rollout temperature."),
OptSpec("rollout_top_k", "Algorithm", help="Rollout top-k (0=disable)."),
OptSpec("rollout_top_p", "Algorithm", help="Rollout top-p."),
OptSpec("rollout_max_tokens", "Algorithm", help="Max tokens per rollout response."),
OptSpec("neftune_alpha", "Algorithm", help="NEFTune noise alpha."),
OptSpec("val_split", "Validation", help="Validation split ratio."),
OptSpec("val_step", "Validation", help="Steps between validation runs."),
OptSpec(
"metrics",
"Validation",
default=("loss", "lr", "grad_norm", "grad_snr"),
help="Metrics to log (repeatable).",
),
OptSpec("ckpt_interval", "Checkpoint", help="Steps between checkpoints."),
OptSpec(
"ckpt_dir",
"Checkpoint",
type=click.Path(),
default="checkpoint",
help="Checkpoint directory.",
),
OptSpec("start_epoch", "Checkpoint", help="Start epoch."),
OptSpec("start_samples", "Checkpoint", help="Start samples (per rank)."),
OptSpec("master_addr", "Distributed", help="Master node address."),
OptSpec("master_port", "Distributed", help="Master node port."),
OptSpec("backend", "Distributed", choices=_BACKENDS, help="Distributed backend."),
OptSpec("nprocs", "Distributed", help="Number of GPUs."),
OptSpec(
"parallel_mode",
"Distributed",
choices=_PARALLEL,
default="fsdp",
help="Parallel strategy.",
),
OptSpec("device_type", "Distributed", help="Device type."),
OptSpec(
"start_method",
"Distributed",
choices=_START_METHODS,
help="Multiprocessing start method.",
),
OptSpec(
"tp_size",
"Distributed",
type=int,
default=None,
help="Tensor parallelism (future).",
),
OptSpec(
"gradient_checkpointing",
"Misc",
type=bool,
default=False,
help="Enable activation checkpointing.",
),
OptSpec(
"compile_mode",
"Misc",
choices=["default", "reduce-overhead", "max-autotune"],
param_decls=("--compile", "compile_mode"),
help="torch.compile mode. Omit to disable.",
),
OptSpec(
"dry_run",
"Misc",
is_flag=True,
default=False,
param_decls=("--dry-run",),
help="Validate config and print plan, do not train.",
),
OptSpec(
"schedule_type",
"LR Schedule",
choices=_SCHEDULES,
default="cosine",
help="LR scheduler.",
),
OptSpec(
"min_rate",
"LR Schedule",
type=float,
default=None,
help="Minimum LR as fraction of base LR.",
),
OptSpec(
"cycle_length",
"LR Schedule",
type=int,
default=None,
help="SGDR first cycle length.",
),
OptSpec(
"t_mult",
"LR Schedule",
type=int,
default=2,
help="SGDR cycle length multiplier.",
),
OptSpec(
"stable_steps",
"LR Schedule",
type=int,
default=None,
help="WSD stable plateau steps.",
),
OptSpec(
"decay_steps",
"LR Schedule",
type=int,
default=None,
help="WSD decay steps.",
),
]
@click.command( @click.command(
name="train", name="train",
@@ -112,425 +344,7 @@ _START_METHODS = sorted(START_METHODS)
help="Start model training (pretrain / SFT / DPO / GRPO).", help="Start model training (pretrain / SFT / DPO / GRPO).",
context_settings={"show_default": True}, context_settings={"show_default": True},
) )
@opt( @apply_specs(_SPECS, TrainConfig)
"--config",
"-c",
"config_path",
type=click.Path(exists=True),
group="Paths & Setup",
help="YAML config file. CLI flags override YAML values.",
)
@opt(
"--train_type",
type=click.Choice(_TRAIN_TYPE),
required=False,
group="Paths & Setup",
help="Training type.",
)
@opt(
"--data_root_path",
type=click.Path(exists=True),
group="Paths & Setup",
help="Root directory of the dataset.",
)
@opt(
"--param_path",
type=click.Path(exists=True),
group="Paths & Setup",
help="Path to model parameters or resume checkpoint.",
)
@opt(
"--resume",
is_flag=True,
default=False,
group="Paths & Setup",
help="Resume from checkpoint.",
)
@opt("--n_epoch", type=int, default=1, group="Training", help="Number of epochs.")
@opt(
"--batch_per_device",
type=int,
default=1,
group="Training",
help="Batch size per GPU.",
)
@opt(
"--grad_accum_steps",
type=int,
default=1,
group="Training",
help="Gradient accumulation steps.",
)
@opt(
"--warmup_ratio",
type=float,
default=0.05,
group="LR Schedule",
help="Fraction of total steps for LR warmup.",
)
@opt(
"--max_lr",
type=float,
default=3e-4,
group="Optimizer",
help="Max learning rate.",
)
@opt(
"--optimizer",
type=click.Choice(_OPTIMIZERS),
default="muon_adamw",
group="Optimizer",
help="Built-in optimizer.",
)
@opt(
"--max_grad_norm",
type=float,
default=1.0,
group="Training",
help="Max gradient norm for clipping.",
)
@opt(
"--weight_decay",
type=float,
default=0.1,
group="Optimizer",
help="Weight decay for eligible optimizer parameters.",
)
@opt(
"--nora_lr", type=float, default=5e-3, group="Optimizer", help="Nora learning rate."
)
@opt(
"--nora_beta", type=float, default=0.95, group="Optimizer", help="Nora EMA factor."
)
@opt(
"--nora_momentum",
type=float,
default=0.95,
group="Optimizer",
help="Nora update momentum.",
)
@opt(
"--nora_weight_decay",
type=float,
default=0.0,
group="Optimizer",
help="Nora weight decay.",
)
@opt(
"--muon_momentum",
type=float,
default=0.95,
group="Optimizer",
help="Muon momentum factor.",
)
@opt(
"--muon_nesterov/--no-muon_nesterov",
default=True,
group="Optimizer",
help="Muon Nesterov.",
)
@opt(
"--muon_ns_steps",
type=int,
default=5,
group="Optimizer",
help="Muon Newton-Schulz steps.",
)
@opt(
"--muon_adjust_lr",
type=click.Choice(["original", "match_rms_adamw"]),
default="match_rms_adamw",
group="Optimizer",
help="Muon LR adjustment strategy.",
)
@opt(
"--mano_momentum",
type=float,
default=0.95,
group="Optimizer",
help="Mano momentum factor.",
)
@opt(
"--mano_nesterov/--no-mano_nesterov",
default=True,
group="Optimizer",
help="Mano Nesterov momentum.",
)
@opt(
"--random_seed",
type=int,
default=3407,
group="Data Loading",
help="Random seed.",
)
@opt(
"--num_workers",
type=int,
default=4,
group="Data Loading",
help="DataLoader workers.",
)
@opt(
"--pin_memory/--no-pin_memory",
default=True,
group="Data Loading",
help="Pin memory.",
)
@opt(
"--persistent_workers/--no-persistent_workers",
default=True,
group="Data Loading",
help="Keep DataLoader workers alive between epochs.",
)
@opt(
"--window_size",
type=int,
default=None,
group="Data Loading",
help="Max input sequence length.",
)
@opt(
"--stride",
type=int,
default=None,
group="Data Loading",
help="Step size for sliding window.",
)
@opt("--dpo_beta", type=float, default=0.1, group="Algorithm", help="DPO beta.")
@opt("--group_size", type=int, default=4, group="Algorithm", help="GRPO group size.")
@opt(
"--grpo_clip_eps",
type=float,
default=0.2,
group="Algorithm",
help="GRPO clip epsilon.",
)
@opt(
"--grpo_kl_coef",
type=float,
default=0.01,
group="Algorithm",
help="GRPO KL penalty coefficient.",
)
@opt(
"--label_smoothing",
type=float,
default=0.0,
group="Data Loading",
help="Label smoothing.",
)
@opt(
"--moe_aux_loss_coef",
type=float,
default=0.01,
group="Algorithm",
help="MoE load balancing auxiliary loss coefficient (0=disable).",
)
@opt(
"--rollout_interval",
type=int,
default=512,
group="Algorithm",
help="Steps between rollouts.",
)
@opt(
"--rollout_max_policy_lag",
type=int,
default=None,
group="Algorithm",
help="Maximum accepted rollout/live policy-version gap.",
)
@opt(
"--rollout_temperature",
type=float,
default=0.7,
group="Algorithm",
help="Rollout temperature.",
)
@opt(
"--rollout_top_k",
type=int,
default=0,
group="Algorithm",
help="Rollout top-k (0=disable).",
)
@opt(
"--rollout_top_p",
type=float,
default=0.9,
group="Algorithm",
help="Rollout top-p.",
)
@opt(
"--rollout_max_tokens",
type=int,
default=1024,
group="Algorithm",
help="Max tokens per rollout response.",
)
@opt(
"--gradient_checkpointing/--no-gradient_checkpointing",
default=False,
group="Misc",
help="Enable activation checkpointing.",
)
@opt(
"--compile",
"compile_mode",
type=click.Choice(["default", "reduce-overhead", "max-autotune"]),
default=None,
group="Misc",
help="torch.compile mode. Omit to disable.",
)
@opt(
"--ckpt_interval",
type=int,
default=5000,
group="Checkpoint",
help="Steps between checkpoints.",
)
@opt(
"--ckpt_dir",
type=click.Path(),
default="checkpoint",
group="Checkpoint",
help="Checkpoint directory.",
)
@opt(
"--val_split",
type=float,
default=None,
group="Validation",
help="Validation split ratio.",
)
@opt(
"--val_step",
type=int,
default=1000,
group="Validation",
help="Steps between validation runs.",
)
@opt(
"--metrics",
multiple=True,
default=("loss", "lr", "grad_norm", "grad_snr"),
group="Validation",
help="Metrics to log (repeatable).",
)
@opt("--start_epoch", type=int, default=0, group="Checkpoint", help="Start epoch.")
@opt(
"--start_samples",
type=int,
default=0,
group="Checkpoint",
help="Start samples (per rank).",
)
@opt(
"--master_addr",
type=str,
default="localhost",
group="Distributed",
help="Master node address.",
)
@opt(
"--master_port",
type=str,
default="29500",
group="Distributed",
help="Master node port.",
)
@opt(
"--backend",
type=click.Choice(_BACKENDS),
default="nccl",
group="Distributed",
help="Distributed backend.",
)
@opt("--nprocs", type=int, default=1, group="Distributed", help="Number of GPUs.")
@opt(
"--parallel_mode",
type=click.Choice(_PARALLEL),
default="fsdp",
group="Distributed",
help="Parallel strategy.",
)
@opt(
"--device_type",
type=str,
default="cuda",
group="Distributed",
help="Device type.",
)
@opt(
"--start_method",
type=click.Choice(_START_METHODS),
default="spawn",
group="Distributed",
help="Multiprocessing start method.",
)
@opt(
"--neftune_alpha",
type=float,
default=0.0,
group="Algorithm",
help="NEFTune noise alpha.",
)
@opt(
"--schedule_type",
type=click.Choice(_SCHEDULES),
default="cosine",
group="LR Schedule",
help="LR scheduler.",
)
@opt(
"--min_rate",
type=float,
default=None,
group="LR Schedule",
help="Minimum LR as fraction of base LR.",
)
@opt(
"--cycle_length",
type=int,
default=None,
group="LR Schedule",
help="SGDR first cycle length.",
)
@opt(
"--t_mult",
type=int,
default=2,
group="LR Schedule",
help="SGDR cycle length multiplier.",
)
@opt(
"--stable_steps",
type=int,
default=None,
group="LR Schedule",
help="WSD stable plateau steps.",
)
@opt(
"--decay_steps",
type=int,
default=None,
group="LR Schedule",
help="WSD decay steps.",
)
@opt(
"--tp_size",
type=int,
default=None,
group="Distributed",
help="Tensor parallelism (future).",
)
@opt(
"--dry-run",
is_flag=True,
default=False,
group="Misc",
help="Validate config and print plan, do not train.",
)
@click.pass_context @click.pass_context
def train_command(ctx, config_path, dry_run, metrics, **kwargs): def train_command(ctx, config_path, dry_run, metrics, **kwargs):
"""Start model training (pretrain / SFT / DPO / GRPO).""" """Start model training (pretrain / SFT / DPO / GRPO)."""
+230
View File
@@ -0,0 +1,230 @@
import dataclasses
import typing as t
import click
import pytest
from click.testing import CliRunner
from astrai.config import TrainConfig, merge_yaml_into_kwargs
from astrai.config.cli import (
GroupedCommand,
OptSpec,
apply_specs,
option_from_spec,
)
def test_merge_yaml_overrides_defaults_but_not_explicit_cli(tmp_path):
config_path = tmp_path / "train.yaml"
config_path.write_text(
"training:\n"
" optimizer: nora_nadamw\n"
" max_lr: 2e-4\n"
" nora_lr: 0.004\n"
" batch_per_device: 8\n",
encoding="utf-8",
)
click_values = {
"optimizer": "nora_nadamw",
"max_lr": 3e-4,
"nora_lr": 5e-3,
"batch_per_device": 16,
}
merged = merge_yaml_into_kwargs(
str(config_path), click_values, explicit_keys={"batch_per_device"}
)
assert merged["max_lr"] == 2e-4
assert merged["nora_lr"] == 4e-3
assert merged["batch_per_device"] == 16
def test_merge_yaml_parses_scientific_notation_as_float(tmp_path):
config_path = tmp_path / "train.yaml"
config_path.write_text("training:\n max_lr: 2e-5\n", encoding="utf-8")
merged = merge_yaml_into_kwargs(
str(config_path), {"max_lr": 3e-4}, explicit_keys=set()
)
assert merged["max_lr"] == 2e-5
assert isinstance(merged["max_lr"], float)
def test_merge_yaml_ignores_unknown_sections(tmp_path):
config_path = tmp_path / "train.yaml"
config_path.write_text("unknown:\n max_lr: 1.0\n", encoding="utf-8")
merged = merge_yaml_into_kwargs(str(config_path), {"max_lr": 3e-4})
assert merged["max_lr"] == 3e-4
def test_type_inference_from_config_fields():
@click.command()
@apply_specs(
[
OptSpec("n_epoch", "G"),
OptSpec("rollout_max_policy_lag", "G"),
OptSpec("pin_memory", "G"),
OptSpec("metrics", "G"),
],
TrainConfig,
)
def cmd(**kwargs):
pass
params = {p.name: p for p in cmd.params}
assert params["n_epoch"].default == 1
assert params["n_epoch"].type.name == "integer"
assert params["rollout_max_policy_lag"].default is None
assert params["rollout_max_policy_lag"].type.name == "integer"
assert params["pin_memory"].secondary_opts == ["--no-pin_memory"]
assert params["pin_memory"].default is False # config default, no override
assert params["metrics"].multiple
assert params["metrics"].default == ("loss", "lr", "grad_norm")
def test_spec_overrides_beat_config_defaults():
@click.command()
@apply_specs(
[
OptSpec("num_workers", "G", default=4),
OptSpec("parallel_mode", "G", default="fsdp"),
],
TrainConfig,
)
def cmd(**kwargs):
pass
params = {p.name: p for p in cmd.params}
assert params["num_workers"].default == 4
assert params["parallel_mode"].default == "fsdp"
def test_cli_only_flag_pair_and_one_way_flag():
@click.command()
@apply_specs(
[
OptSpec("muon_nesterov", "G", type=bool, default=True),
OptSpec("dry_run", "G", is_flag=True, default=False),
]
)
def cmd(**kwargs):
pass
params = {p.name: p for p in cmd.params}
assert params["muon_nesterov"].secondary_opts == ["--no-muon_nesterov"]
assert params["muon_nesterov"].default is True
assert params["dry_run"].is_flag
assert not params["dry_run"].secondary_opts
def test_uninferrable_type_raises_without_spec_type():
with pytest.raises(TypeError, match="cannot infer"):
option_from_spec(OptSpec("mystery", "G"), {})
def test_annotation_styles_old_and_new():
"""Optional[X], Union[X, None], X | None, List[str], and list[str] all
infer the same click types."""
@dataclasses.dataclass
class OldStyle:
opt_int: t.Optional[int] = None
opt_float: t.Optional[float] = None
union_int: t.Union[int, None] = None
names: t.List[str] = dataclasses.field(default_factory=lambda: ["loss", "lr"])
flag: bool = True
@dataclasses.dataclass
class NewStyle:
opt_int: int | None = None
names: list[str] = dataclasses.field(default_factory=lambda: ["loss"])
flag: bool = False
@click.command()
@apply_specs(
[
OptSpec("opt_int", "G"),
OptSpec("opt_float", "G"),
OptSpec("union_int", "G"),
OptSpec("names", "G"),
OptSpec("flag", "G"),
],
OldStyle,
)
def old_cmd(**kwargs):
pass
params = {p.name: p for p in old_cmd.params}
assert params["opt_int"].type.name == "integer"
assert params["opt_int"].default is None
assert params["opt_float"].type.name == "float"
assert params["union_int"].type.name == "integer"
assert params["names"].multiple
assert params["names"].type.name == "text"
assert params["flag"].secondary_opts == ["--no-flag"]
assert params["flag"].default is True
@click.command()
@apply_specs(
[
OptSpec("opt_int", "G"),
OptSpec("names", "G"),
OptSpec("flag", "G"),
],
NewStyle,
)
def new_cmd(**kwargs):
pass
params = {p.name: p for p in new_cmd.params}
assert params["opt_int"].type.name == "integer"
assert params["names"].multiple
assert params["names"].default == ("loss",)
assert params["flag"].secondary_opts == ["--no-flag"]
assert params["flag"].default is False
def test_stringified_annotations_resolve_via_get_type_hints():
"""PEP 563 modules (``from __future__ import annotations``) leave
``Field.type`` as a string; hints resolution still infers types."""
@dataclasses.dataclass
class FutureStyle:
opt_int: "t.Optional[int]" = None
count: "int" = 3
@click.command()
@apply_specs(
[
OptSpec("opt_int", "G"),
OptSpec("count", "G"),
],
FutureStyle,
)
def cmd(**kwargs):
pass
params = {p.name: p for p in cmd.params}
assert params["opt_int"].type.name == "integer"
assert params["count"].type.name == "integer"
assert params["count"].default == 3
def test_help_order_follows_spec_table():
@click.command(cls=GroupedCommand)
@apply_specs(
[
OptSpec("first", "G", type=int, default=1),
OptSpec("second", "G", type=int, default=2),
]
)
def cmd(**kwargs):
pass
result = CliRunner().invoke(cmd, ["--help"])
assert result.exit_code == 0
assert result.output.index("--first") < result.output.index("--second")