""" pyserve scale - Scale services """ import asyncio from pathlib import Path from typing import Any import click @click.command("scale") @click.argument("scales", nargs=-1, required=True) @click.option( "--timeout", "timeout", default=60, type=int, help="Timeout in seconds for scaling operation", ) @click.option( "--no-wait", is_flag=True, help="Don't wait for services to be ready", ) @click.pass_obj def scale_cmd(ctx: Any, scales: tuple[str, ...], timeout: int, no_wait: bool) -> None: """ Scale services to specified number of workers. Use SERVICE=NUM format to specify scaling. \b Examples: pyserve scale api=4 # Scale api to 4 workers pyserve scale api=4 admin=2 # Scale multiple services """ from ...config import Config from .._runner import ServiceRunner from ..output import console, print_error, print_info, print_success from ..state import StateManager scale_map = {} for scale in scales: try: service, num = scale.split("=") scale_map[service] = int(num) except ValueError: print_error(f"Invalid scale format: {scale}. Use SERVICE=NUM") raise click.Abort() config_path = Path(ctx.config_file) if not config_path.exists(): print_error(f"Configuration file not found: {config_path}") raise click.Abort() config = Config.from_yaml(str(config_path)) state_manager = StateManager(Path(".pyserve"), ctx.project) all_services = state_manager.get_all_services() for service in scale_map: if service not in all_services: print_error(f"Service '{service}' not found") raise click.Abort() runner = ServiceRunner(config, state_manager) console.print("[bold]Scaling services...[/bold]") async def do_scale() -> None: for service, workers in scale_map.items(): current = all_services[service].workers or 1 print_info(f"Scaling {service}: {current} → {workers} workers") try: success = await runner.scale_service(service, workers, timeout=timeout, wait=not no_wait) if success: print_success(f"Scaled {service} to {workers} workers") else: print_error(f"Failed to scale {service}") except Exception as e: print_error(f"Error scaling {service}: {e}") try: asyncio.run(do_scale()) except Exception as e: print_error(f"Scaling failed: {e}") raise click.Abort()