ViperEkura
4de42d83c2
- Replace argparse with click in all scripts (train, server, generate, preprocess, benchmark) - Add --config YAML support to train.py with CLI flag override - Add --dry-run mode to validate config before training - Add type annotations throughout benchmark.py - Unify docstring format across all commands - Remove redundant deps httpx, requests, pyyaml, rich from pyproject.toml - Net -346 lines while adding YAML config support
59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
from pathlib import Path
|
|
|
|
import click
|
|
import torch
|
|
|
|
from astrai.inference import run_server
|
|
|
|
_DTYPES = ["bfloat16", "float16", "float32"]
|
|
|
|
|
|
@click.command(name="serve", help="Launch inference server (OpenAI-compatible API).")
|
|
@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, 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.",
|
|
)
|
|
def server_command(host, port, reload, param_path, device, dtype, max_batch_size):
|
|
"""Launch inference server (OpenAI-compatible API)."""
|
|
dtype_map = {
|
|
"bfloat16": torch.bfloat16,
|
|
"float16": torch.float16,
|
|
"float32": torch.float32,
|
|
}
|
|
project_root = Path(__file__).parent.parent.parent
|
|
param_path = 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}")
|
|
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,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
server_command()
|