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

scripts/tools/server.py

@@ -2,9 +2,14 @@ from pathlib import Path
 
 import click
 import torch
-import yaml
 from click.core import ParameterSource
 
+from astrai.config.cli import (
+    GroupedCommand,
+    OptSpec,
+    apply_specs,
+    merge_yaml_into_kwargs,
+)
 from astrai.inference import run_server
 
 _DTYPES = ["bfloat16", "float16", "float32"]
@@ -26,29 +31,65 @@ def _merge_yaml_into_kwargs(
     explicit_keys: set[str] | None = None,
 ) -> dict:
     """Merge Click defaults, YAML server values, then explicit CLI values."""
-    with open(config_path, encoding="utf-8") as file:
-        config = yaml.safe_load(file) or {}
-    if not isinstance(config, dict):
-        raise click.UsageError(f"Serving config must be a mapping: {config_path}")
-    server = config.get("server") or {}
-    if not isinstance(server, dict):
-        raise click.UsageError("top-level server section must be a mapping")
+    return merge_yaml_into_kwargs(
+        config_path,
+        passed_kwargs,
+        explicit_keys,
+        sections=("server",),
+        allowed_keys=_SERVER_KEYS,
+    )
 
-    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)
-    merged.update({key: server[key] for key in _SERVER_KEYS if key in server})
-    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
+_SPECS = [
+    OptSpec(
+        "config_path",
+        "Server",
+        type=click.Path(exists=True, dir_okay=False),
+        param_decls=("--config", "-c", "config_path"),
+        help="Serving YAML config. CLI flags override YAML values.",
+    ),
+    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:
@@ -87,84 +128,22 @@ def _resolve_server_config(
     return resolved
 
 
-@click.command(name="serve", help="Launch inference server (OpenAI-compatible API).")
-@click.option(
-    "--config",
-    "-c",
-    "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.",
+@click.command(
+    name="serve",
+    cls=GroupedCommand,
+    help="Launch inference server (OpenAI-compatible API).",
 )
+@apply_specs(_SPECS)
 @click.pass_context
-def server_command(
-    ctx,
-    config_path,
-    host,
-    port,
-    reload,
-    param_path,
-    device,
-    dtype,
-    max_batch_size,
-    max_seq_len,
-):
+def server_command(ctx, config_path, **kwargs):
     """Launch inference server (OpenAI-compatible API)."""
     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 = {
             key
-            for key in passed_kwargs
+            for key in kwargs
             if ctx.get_parameter_source(key) is ParameterSource.COMMANDLINE
         }
-        resolved = _resolve_server_config(config_path, passed_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"]
+        kwargs = _resolve_server_config(config_path, kwargs, explicit_keys)
         click.echo(f"Config: {config_path}")
 
     dtype_map = {
@@ -173,19 +152,22 @@ def server_command(
         "float32": torch.float32,
     }
     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"Model: {param_path}  |  Device: {device}  |  Dtype: {dtype}")
+    click.echo(f"Starting server on http://{kwargs['host']}:{kwargs['port']}")
+    click.echo(
+        f"Model: {kwargs['param_path']}  |  "
+        f"Device: {kwargs['device']}  |  Dtype: {kwargs['dtype']}"
+    )
     run_server(
-        host=host,
-        port=port,
-        reload=reload,
-        device=device,
-        dtype=dtype_map[dtype],
-        param_path=Path(param_path),
-        max_batch_size=max_batch_size,
-        max_seq_len=max_seq_len,
+        host=kwargs["host"],
+        port=kwargs["port"],
+        reload=kwargs["reload"],
+        device=kwargs["device"],
+        dtype=dtype_map[kwargs["dtype"]],
+        param_path=Path(kwargs["param_path"]),
+        max_batch_size=kwargs["max_batch_size"],
+        max_seq_len=kwargs["max_seq_len"],
     )