scripts/set_modal_vllm_autoscaler.py

import argparse
import json


DEFAULT_APP_NAME = "virtual-characters-vllm-gemma"
DEFAULT_FUNCTION_NAME = "serve"


def require_modal():
    try:
        import modal
    except Exception as exc:
        raise SystemExit(
            "Modal Python package is not installed. Install with `python -m pip install -r requirements.txt`, "
            "then run `modal setup` to log in."
        ) from exc
    return modal


def update_autoscaler(args) -> int:
    modal = require_modal()
    fn = modal.Function.from_name(
        args.app_name,
        args.function_name,
        environment_name=args.environment_name,
    )
    fn.update_autoscaler(
        min_containers=args.min_containers,
        buffer_containers=args.buffer_containers,
        scaledown_window=args.scaledown_window,
    )
    print(
        json.dumps(
            {
                "app": args.app_name,
                "function": args.function_name,
                "environment": args.environment_name,
                "min_containers": args.min_containers,
                "buffer_containers": args.buffer_containers,
                "scaledown_window": args.scaledown_window,
                "note": "Autoscaler overrides reset on the next deploy of this app.",
            },
            ensure_ascii=False,
            indent=2,
        )
    )
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="Update the deployed Modal vLLM autoscaler.")
    parser.add_argument(
        "mode",
        choices=["on", "off"],
        help="on keeps one vLLM container warm; off allows the deployment to scale to zero.",
    )
    parser.add_argument("--app-name", default=DEFAULT_APP_NAME)
    parser.add_argument("--function-name", default=DEFAULT_FUNCTION_NAME)
    parser.add_argument("--environment-name")
    parser.add_argument("--min-containers", type=int)
    parser.add_argument("--buffer-containers", type=int)
    parser.add_argument("--scaledown-window", type=int)
    args = parser.parse_args()

    if args.mode == "on":
        args.min_containers = 1 if args.min_containers is None else args.min_containers
        args.buffer_containers = 0 if args.buffer_containers is None else args.buffer_containers
        args.scaledown_window = 20 * 60 if args.scaledown_window is None else args.scaledown_window
    else:
        args.min_containers = 0 if args.min_containers is None else args.min_containers
        args.buffer_containers = 0 if args.buffer_containers is None else args.buffer_containers
        args.scaledown_window = 5 * 60 if args.scaledown_window is None else args.scaledown_window

    return update_autoscaler(args)


if __name__ == "__main__":
    raise SystemExit(main())