diff --git a/agent/nix_builder_autoscaler/cli.py b/agent/nix_builder_autoscaler/cli.py index eb3ff70..28d9645 100644 --- a/agent/nix_builder_autoscaler/cli.py +++ b/agent/nix_builder_autoscaler/cli.py @@ -94,131 +94,13 @@ def _print_reservations(data: list[dict[str, Any]]) -> None: _print_table(["reservation_id", "phase", "system", "slot", "instance_id"], rows) -def _print_status_summary(data: dict[str, Any]) -> None: - slots = data.get("slots", {}) - reservations = data.get("reservations", {}) - ec2 = data.get("ec2", {}) - haproxy = data.get("haproxy", {}) - rows = [ - ["slots.total", str(slots.get("total", 0))], - ["slots.ready", str(slots.get("ready", 0))], - ["slots.launching", str(slots.get("launching", 0))], - ["slots.booting", str(slots.get("booting", 0))], - ["slots.binding", str(slots.get("binding", 0))], - ["slots.terminating", str(slots.get("terminating", 0))], - ["slots.empty", str(slots.get("empty", 0))], - ["slots.error", str(slots.get("error", 0))], - ["reservations.pending", str(reservations.get("pending", 0))], - ["reservations.ready", str(reservations.get("ready", 0))], - ["reservations.failed", str(reservations.get("failed", 0))], - ["ec2.api_ok", str(ec2.get("api_ok", False))], - ["haproxy.socket_ok", str(haproxy.get("socket_ok", False))], - ] - _print_table(["metric", "value"], rows) - - -def _bulk_slot_action(socket_path: str, action: str) -> dict[str, Any]: - if action == "drain": - eligible_states = {"ready"} - action_path = "/v1/admin/drain" - elif action == "unquarantine": - eligible_states = {"error"} - action_path = "/v1/admin/unquarantine" - else: - msg = f"unknown bulk action: {action}" - raise ValueError(msg) - - status, data = _uds_request(socket_path, "GET", "/v1/slots") - if status < 200 or status >= 300 or not isinstance(data, list): - msg = "failed to list slots for bulk action" - raise RuntimeError(msg) - - results: list[dict[str, Any]] = [] - summary: dict[str, Any] = { - "action": action, - "matched": 0, - "attempted": 0, - "succeeded": 0, - "failed": 0, - "skipped": 0, - "results": results, - } - - for slot in data: - slot_id = str(slot.get("slot_id", "")) - state = str(slot.get("state", "")) - if not slot_id: - continue - - if state not in eligible_states: - summary["skipped"] += 1 - results.append( - { - "slot_id": slot_id, - "state": state, - "result": "skipped", - "reason": "ineligible_state", - } - ) - continue - - summary["matched"] += 1 - summary["attempted"] += 1 - try: - action_status, action_data = _uds_request( - socket_path, - "POST", - action_path, - body={"slot_id": slot_id}, - ) - except OSError as err: - summary["failed"] += 1 - results.append( - { - "slot_id": slot_id, - "state": state, - "result": "failed", - "error": str(err), - } - ) - continue - - if 200 <= action_status < 300: - summary["succeeded"] += 1 - results.append( - { - "slot_id": slot_id, - "state": state, - "result": "ok", - } - ) - else: - summary["failed"] += 1 - results.append( - { - "slot_id": slot_id, - "state": state, - "result": "failed", - "status": action_status, - "response": action_data, - } - ) - - return summary - - -def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: +def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(prog="autoscalerctl", description="Autoscaler CLI") parser.add_argument( "--socket", default="/run/nix-builder-autoscaler/daemon.sock", help="Daemon Unix socket path", ) - parser.add_argument( - "--json", - action="store_true", - help="Output JSON for status command.", - ) subparsers = parser.add_subparsers(dest="command") subparsers.add_parser("status", help="Show state summary") subparsers.add_parser("slots", help="List slots") @@ -228,14 +110,8 @@ def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser_drain.add_argument("slot_id") parser_unq = subparsers.add_parser("unquarantine", help="Unquarantine a slot") parser_unq.add_argument("slot_id") - subparsers.add_parser("drain-all", help="Drain all eligible slots (state=ready)") - subparsers.add_parser("unquarantine-all", help="Unquarantine all error slots") subparsers.add_parser("reconcile-now", help="Trigger immediate reconcile tick") - args = parser.parse_args(argv) - if not args.command: - parser.print_help() - raise SystemExit(0) - return args + return parser.parse_args() def _print_error(data: object) -> None: @@ -248,19 +124,8 @@ def _print_error(data: object) -> None: def main() -> None: """Entry point for the autoscalerctl CLI.""" args = _parse_args() - - if args.command in {"drain-all", "unquarantine-all"}: - action = "drain" if args.command == "drain-all" else "unquarantine" - try: - summary = _bulk_slot_action(args.socket, action) - except OSError as err: - print(f"Error: cannot connect to daemon at {args.socket}") - raise SystemExit(1) from err - except RuntimeError as err: - print(str(err)) - raise SystemExit(1) from err - print(json.dumps(summary, indent=2)) - raise SystemExit(0 if summary["failed"] == 0 else 1) + if not args.command: + raise SystemExit(1) method = "GET" path = "" @@ -295,15 +160,7 @@ def main() -> None: _print_error(data) raise SystemExit(1) - if args.command == "status": - if not isinstance(data, dict): - _print_error(data) - raise SystemExit(1) - if args.json: - print(json.dumps(data, indent=2)) - else: - _print_status_summary(data) - elif args.command in {"drain", "unquarantine", "reconcile-now"}: + if args.command in {"status", "drain", "unquarantine", "reconcile-now"}: print(json.dumps(data, indent=2)) elif args.command == "slots": if isinstance(data, list): diff --git a/agent/nix_builder_autoscaler/tests/test_cli.py b/agent/nix_builder_autoscaler/tests/test_cli.py deleted file mode 100644 index 2440b7d..0000000 --- a/agent/nix_builder_autoscaler/tests/test_cli.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Unit tests for autoscalerctl CLI argument and display behavior.""" - -from __future__ import annotations - -import pytest - -from nix_builder_autoscaler import cli -from nix_builder_autoscaler.cli import _parse_args, _print_status_summary - - -def test_parse_args_without_command_prints_help_and_exits_zero( - capsys: pytest.CaptureFixture[str], -) -> None: - with pytest.raises(SystemExit) as exc: - _parse_args([]) - assert exc.value.code == 0 - captured = capsys.readouterr() - assert "Autoscaler CLI" in captured.out - assert "status" in captured.out - - -def test_parse_args_json_status() -> None: - args = _parse_args(["--json", "status"]) - assert args.command == "status" - assert args.json is True - - -def test_parse_args_bulk_commands() -> None: - assert _parse_args(["drain-all"]).command == "drain-all" - assert _parse_args(["unquarantine-all"]).command == "unquarantine-all" - - -def test_print_status_summary_renders_metrics_table(capsys: pytest.CaptureFixture[str]) -> None: - _print_status_summary( - { - "slots": { - "total": 4, - "ready": 1, - "launching": 1, - "booting": 1, - "binding": 0, - "terminating": 0, - "empty": 1, - "error": 0, - }, - "reservations": {"pending": 2, "ready": 1, "failed": 0}, - "ec2": {"api_ok": True}, - "haproxy": {"socket_ok": True}, - } - ) - out = capsys.readouterr().out - assert "metric" in out - assert "slots.total" in out - assert "reservations.pending" in out - assert "haproxy.socket_ok" in out - - -def test_bulk_drain_only_targets_ready_slots(monkeypatch: pytest.MonkeyPatch) -> None: - def _fake_request(socket_path: str, method: str, path: str, body=None): # noqa: ANN001 - assert socket_path == "/tmp/sock" - if method == "GET" and path == "/v1/slots": - return 200, [ - {"slot_id": "slot001", "state": "ready"}, - {"slot_id": "slot002", "state": "booting"}, - ] - if method == "POST" and path == "/v1/admin/drain" and body == {"slot_id": "slot001"}: - return 200, {"state": "draining"} - raise AssertionError(f"unexpected request: {method} {path} {body}") - - monkeypatch.setattr(cli, "_uds_request", _fake_request) - summary = cli._bulk_slot_action("/tmp/sock", "drain") - assert summary["matched"] == 1 - assert summary["attempted"] == 1 - assert summary["succeeded"] == 1 - assert summary["failed"] == 0 - assert summary["skipped"] == 1 - - -def test_bulk_unquarantine_only_targets_error_slots(monkeypatch: pytest.MonkeyPatch) -> None: - def _fake_request(socket_path: str, method: str, path: str, body=None): # noqa: ANN001 - assert socket_path == "/tmp/sock" - if method == "GET" and path == "/v1/slots": - return 200, [ - {"slot_id": "slot001", "state": "error"}, - {"slot_id": "slot002", "state": "ready"}, - ] - if method == "POST" and path == "/v1/admin/unquarantine" and body == {"slot_id": "slot001"}: - return 200, {"state": "empty"} - raise AssertionError(f"unexpected request: {method} {path} {body}") - - monkeypatch.setattr(cli, "_uds_request", _fake_request) - summary = cli._bulk_slot_action("/tmp/sock", "unquarantine") - assert summary["matched"] == 1 - assert summary["attempted"] == 1 - assert summary["succeeded"] == 1 - assert summary["failed"] == 0 - assert summary["skipped"] == 1 diff --git a/flake.nix b/flake.nix index 3293071..a5f8570 100644 --- a/flake.nix +++ b/flake.nix @@ -62,8 +62,7 @@ ); mkBuildbotAutoscaleExtPythonPackagesExtension = pkgs: pyFinal: _pyPrev: { buildbot-autoscale-ext = - pyFinal.toPythonModule - (mkBuildbotExtPythonSet pkgs pyFinal.python)."buildbot-autoscale-ext"; + pyFinal.toPythonModule (mkBuildbotExtPythonSet pkgs pyFinal.python)."buildbot-autoscale-ext"; }; in { @@ -95,7 +94,8 @@ pyprojectOverrides ] ); - buildbotExtPythonSet = mkBuildbotExtPythonSet pkgs pkgs.python3; + buildbotExtPythonSet = + mkBuildbotExtPythonSet pkgs pkgs.python3; venv = agentPythonSet.mkVirtualEnv "nix-builder-autoscaler-env" agentWorkspace.deps.default; buildbotExtVenv = buildbotExtPythonSet.mkVirtualEnv "buildbot-autoscale-ext-env" buildbotExtWorkspace.deps.default; in @@ -145,7 +145,8 @@ pyprojectOverrides ] ); - buildbotExtPythonSet = mkBuildbotExtPythonSet pkgs pkgs.python3; + buildbotExtPythonSet = + mkBuildbotExtPythonSet pkgs pkgs.python3; testVenv = agentPythonSet.mkVirtualEnv "nix-builder-autoscaler-test-env" { nix-builder-autoscaler = [ "dev" ]; }; diff --git a/nix/modules/nixos/services/nix-builder-autoscaler.nix b/nix/modules/nixos/services/nix-builder-autoscaler.nix index 7697ac5..803b7a3 100644 --- a/nix/modules/nixos/services/nix-builder-autoscaler.nix +++ b/nix/modules/nixos/services/nix-builder-autoscaler.nix @@ -275,7 +275,7 @@ in backend ${cfg.haproxy.backend} balance leastconn option tcp-check - tcp-check expect rstring SSH-2\\.0-.* + tcp-check expect rstring SSH-2\\.0-OpenSSH.* ${lib.concatMapStrings ( i: "server ${cfg.haproxy.slotPrefix}${lib.fixedWidthNumber 3 i} 127.0.0.2:22 disabled check inter 5s fall 2 rise 2 maxconn 2\n "