88 lines
2.5 KiB
Python
88 lines
2.5 KiB
Python
"""
|
|
pyserve scale - Scale services
|
|
"""
|
|
|
|
import asyncio
|
|
from pathlib import Path
|
|
|
|
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, scales: tuple, timeout: int, no_wait: bool):
|
|
"""
|
|
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():
|
|
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()
|