diff --git a/packages/sie_placement/README.md b/packages/sie_placement/README.md index be45f5483..326cbcbf5 100644 --- a/packages/sie_placement/README.md +++ b/packages/sie_placement/README.md @@ -3,8 +3,9 @@ Versioned documents for a measured-only placement planner: the service level each served identity must meet, the measurements a placement may be chosen from, the prices those placements cost, the demand they serve, and the plan -that results. The package uses only the Python standard library. It is a -workspace package and is not published to PyPI. +that results. The package uses only the Python standard library, except the +KEDA applier, whose `keda` extra adds the SIE Python SDK. It is a workspace +package and is not published to PyPI. A served identity is a model id or a `model:profile` pair, the routing key a request resolves to. A measurement cell is a machine profile together with the @@ -133,6 +134,156 @@ start: read back as `current`, the settings climb one step per tick under steady demand until the lane serves the demand or reaches the top of its envelope. +## KEDA applier + +`sie-placement tune` keeps each plan lane's warm floor inside its envelope on a +`sie-cluster` deployment that autoscales workers with KEDA. It needs the `keda` +extra. + +```bash +sie-placement tune --plan plan.json \ + --prometheus-url http://prometheus:9090 \ + --gateway-url http://sie-gateway:8080 \ + --api-key-file /run/secrets/gateway-client-token +``` + +Every tick, 60 seconds apart by default or once with `--once`, it reads the +four lane series below from Prometheus, reads each lane's floor-carrier pool +from the gateway with `SIEClient.get_pool`, and calls `step` with the carrier's +floor as `current`. It prints one JSON audit row per lane to standard output. Without +`--apply` it runs in shadow mode: it decides and records and writes nothing. +With `--apply` it upserts every lane's carrier on every tick through +`SIEClient.create_pool`, which renews the carrier's lease and sets its floor to +the step's floor. Each write uses its own SDK client and closes it after the +write, so the SDK's background lease renewal never runs: a carrier's lease is +renewed only by a tick that writes it. A lane that has no carrier and whose +floor stays at 0 gets none. + +Under KEDA only the floor moves per lane. The scale-down window and the queue +threshold are chart-wide settings, and there is no buffer to tune, so the +applier passes `Settings(floor=n)` and ignores the envelope's buffer and window +ranges. + +### Floor carriers + +A floor carrier is an API logical pool named +`placement-floor---`, with `queue_pool` and +`bundle` set to the lane's, `gpus` set to `{: 0}`, the +gateway's default lease, no GPU caps, no pinned models, and +`minimum_worker_count` set to the step's floor. The gateway keeps each lane at +the largest floor among the logical pools on it, so the carrier can raise the +lane above its chart floor and never lowers another pool's floor. The applier +writes only its carriers: a pool that has a carrier's name and any other spec is +reported and left unchanged. A carrier carries no pins, because every floor +change is a spec change, which resets the pool's worker assignment; pinned +models stay in the chart values. + +A tick that cannot read a lane's carrier, or finds a pool with its name that is +not its carrier, does not write it. A carrier that no tick writes, because those +reads keep failing or the applier has stopped, lapses 1200 seconds after its +last write by default, the gateway removes it, and the lane returns to its +chart floor. On a lane whose chart `minReplicas` is 0, a floor of 1 or more +activates the lane, which then holds at least the chart's fallback floor. + +### Series + +| Series | Query | Step input | +| --- | --- | --- | +| Queue depth | `max by (pool, machine_profile, bundle) (sie_gateway_lane_queue_depth{} and on(producer_instance, collector_generation, pool, machine_profile, bundle) (abs(time() - sie_gateway_lane_queue_snapshot_timestamp_seconds{}) < 20))` | `running` | +| Pending demand | `max by (pool, machine_profile, bundle) (sie_gateway_pending_demand{} and on(producer_instance, collector_generation) (abs(time() - sie_gateway_capacity_snapshot_timestamp_seconds{}) < 20))` | `backlog` | +| Warm floor | `max by (pool, machine_profile, bundle) (sie_gateway_pool_warm_floor{} and on(producer_instance, collector_generation) (abs(time() - sie_gateway_capacity_snapshot_timestamp_seconds{}) < 20))` | Reported only | +| Running workers | `count by (lane) (max by (lane, producer_instance) (sie_worker_gpu_slots{state="ready"}) > 0)` | `runners` | + +`--series-selector` adds label matchers to every selector, timestamps included, +for example `namespace="sie"` when one Prometheus scrapes several releases. +Lane labels match the plan's lane tokens after trimming and lowering them, as +the chart names lanes. + +The collector keeps a producer's last point for minutes after the producer +stops reporting, so each gateway series keeps only the producers whose snapshot +is at most 20 seconds old, as the chart's KEDA triggers do: the lane's own +queue snapshot for queue depth, and the producer's capacity snapshot for +pending demand and warm floor. A lane whose gateway values are all stale reads +as missing. The worker series has no snapshot timestamp, so a worker that +stopped reporting still counts among the runners until the collector drops it. + +Pending demand is a 0 or 1 flag per lane, not a count. The gateway sets it +while a request that found no ready capacity is recent, while any request's +durable publish to the lane awaits its acknowledgement, and for 120 seconds +after a failed dispatch. The attainment is `unmeasured`, so a lane is under +pressure whenever the flag is set, and idle when neither the flag nor queued or +in-flight work is reported. Because a publish in flight sets the flag, any +traffic can read as pressure even when the lane has capacity to spare, and a +lane that serves steady traffic seldom reads idle, so its floor can climb +towards `floorMax` while traffic lasts and falls one step per idle tick +afterwards. Run the applier in shadow mode first to see how often each lane +reads as under pressure. A lane with no +ready worker has 0 runners. A lane whose queue depth, pending demand or worker +series is missing is held: the step is given a busy lane with nothing waiting, +so it only brings the carrier's floor inside the envelope. + +### Audit rows + +Each row carries `at`, `mode` (`shadow` or `apply`), the lane's `lane`, +`queuePool`, `machineProfile` and `bundle`, the `carrier` pool name, the +`telemetry` read for the lane (`pendingDemand` is the 0 or 1 flag), +`currentFloor`, the step's `action` and `floor`, whether the write was +`applied`, `laneFloorAboveEnvelope` and `error`. +`laneFloorAboveEnvelope` is true when the lane's warm floor exceeds `floorMax`. +The applier does not lower another pool's floor, so a lane held above its +envelope by another pool stays reported until that pool changes. + +### Refusals + +The command exits with status 2 before it calls Prometheus or the gateway when: + +- a lane's `floorMax` exceeds `--max-pool-floor`, which must match the gateway's + `SIE_GATEWAY_POOL_MAX_MINIMUM_WORKER_COUNT` (default 4); +- the plan has more lanes than `--max-pools`, which must match the gateway's + `SIE_GATEWAY_MAX_POOLS` (default 64); +- a lane's `queuePool`, machine profile or bundle is not 1 to 63 letters, + digits, `_` or `-`, two lanes need the same carrier, or a carrier name is + longer than the gateway's 128 characters; +- the `sie-cluster` chart cannot run the plan: a lane on the `default` queue + pool pins an identity, whose pins no static queue pool carries, or an + identity is `bridged`, `snapshot-cold` or `scheduled-warm`, warmths the chart + does not realize; +- a token file or `SIE_ADMIN_TOKEN` holds anything but one run of printable + ASCII characters without spaces after trimming; +- `--prometheus-url` or `--gateway-url` is not a valid absolute http or https + URL with a host, embeds credentials, carries a query or a fragment, or has an + invalid port. The refusal does not repeat the URL. + +A failed query, pool read or pool write fails only its tick, which is recorded, +and the next tick runs as usual. + +With `--once`, a tick whose query, pool read or pool write failed exits with +status 1. + +### Credentials + +| Call | Permission | Source | +| --- | --- | --- | +| Prometheus `/api/v1/query` | Read access to the query API | None sent | +| `GET /v1/pools/{name}` | A gateway client token, when gateway token auth is on | `--api-key-file`, or `SIE_API_KEY` when `SIE_BASE_URL` names the same origin, as the SDK resolves it | +| `POST /v1/pools`, which also renews the lease | The gateway admin token | `--admin-token-file`, or `SIE_ADMIN_TOKEN`, read only with `--apply` | + +The gateway admin token authorizes every gateway admin route, not only pools. +The applier reads it only with `--apply`, sends it only on pool writes to +`--gateway-url`, never takes it as a command-line value and never prints or +logs it. Without it, `--apply` sends its writes with no credential, which a +gateway with token auth refuses. The client token is never sent on writes. A +refused token is named by its source, never by its value. A gateway call that +gets an error reply is reported by its HTTP status alone, and so is a reply +that is neither a success nor an error, such as a redirect; a success reply the +SDK cannot parse is reported as an unexpected response. No text a server +returns reaches an audit row, so an intermediary that quotes the request's +`Authorization` header in its error body cannot put the token there. A failed +connection is reported by the SDK's connection error message, sent before any +header, and any other transport failure by its exception type alone. A message +the SDK raises before any reply is cut to 512 characters, with the token the +client holds replaced by `[redacted]`. + ## Tests ```bash diff --git a/packages/sie_placement/pyproject.toml b/packages/sie_placement/pyproject.toml index 91b0c19d6..60fd9258a 100644 --- a/packages/sie_placement/pyproject.toml +++ b/packages/sie_placement/pyproject.toml @@ -6,6 +6,15 @@ requires-python = ">=3.12,<3.13" license = { text = "Apache-2.0" } dependencies = [] +[project.optional-dependencies] +keda = [ + "httpx>=0.28,<1", + "sie-sdk>=0.11", +] + +[project.scripts] +sie-placement = "sie_placement.cli:main" + [build-system] requires = ["hatchling"] build-backend = "hatchling.build" diff --git a/packages/sie_placement/src/sie_placement/cli.py b/packages/sie_placement/src/sie_placement/cli.py new file mode 100644 index 000000000..ccd1df9b9 --- /dev/null +++ b/packages/sie_placement/src/sie_placement/cli.py @@ -0,0 +1,183 @@ +"""Command-line entry point of the placement tools.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import re +import sys +import time +from collections.abc import Callable, Sequence +from pathlib import Path + +from sie_placement import keda +from sie_placement.schema import Plan + +ADMIN_TOKEN_ENV = "SIE_ADMIN_TOKEN" # noqa: S105 - an environment variable name +SDK_API_KEY_ENV = "SIE_API_KEY" +DEFAULT_INTERVAL_S = 60.0 +EXIT_FAILED = 1 +EXIT_REFUSED = 2 + +_TOKEN = re.compile(r"[\x21-\x7e]+") + + +class _UsageError(ValueError): + pass + + +def _positive_float(text: str) -> float: + value = float(text) + if not math.isfinite(value) or value <= 0: + raise argparse.ArgumentTypeError(f"must be greater than 0, got {text!r}") + return value + + +def _integer_at_least(minimum: int) -> Callable[[str], int]: + def parse(text: str) -> int: + value = int(text) + if value < minimum: + raise argparse.ArgumentTypeError(f"must be at least {minimum}, got {text!r}") + return value + + return parse + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="sie-placement", description="Placement planner tools.") + commands = parser.add_subparsers(dest="command", required=True) + tune = commands.add_parser( + "tune", + help="keep each plan lane's warm floor inside its envelope through gateway floor-carrier pools", + description=( + "Read each plan lane's series from Prometheus, decide one tuner step per lane and, with --apply, " + "upsert one floor-carrier pool per lane through the gateway pool API. Without --apply nothing is written." + ), + ) + tune.add_argument("--plan", required=True, type=Path, help="placement-plan document") + tune.add_argument("--prometheus-url", required=True, help="Prometheus base URL that serves /api/v1/query") + tune.add_argument("--gateway-url", required=True, help="gateway base URL") + tune.add_argument( + "--series-selector", + default="", + help="label matchers added to every series selector, for example 'namespace=\"sie\"'", + ) + tune.add_argument("--apply", action="store_true", help="write the floor-carrier pools") + tune.add_argument("--once", action="store_true", help="run one tick and exit") + tune.add_argument( + "--interval-s", type=_positive_float, default=DEFAULT_INTERVAL_S, help="seconds between ticks (default 60)" + ) + tune.add_argument( + "--max-pool-floor", + type=_integer_at_least(0), + default=keda.DEFAULT_MAX_POOL_FLOOR, + help="the gateway's SIE_GATEWAY_POOL_MAX_MINIMUM_WORKER_COUNT (default 4)", + ) + tune.add_argument( + "--max-pools", + type=_integer_at_least(1), + default=keda.DEFAULT_MAX_POOLS, + help="the gateway's SIE_GATEWAY_MAX_POOLS (default 64)", + ) + tune.add_argument("--api-key-file", type=Path, help="file holding a gateway client token for pool reads") + tune.add_argument( + "--admin-token-file", + type=Path, + help=f"file holding the gateway admin token for pool writes; read only with --apply (default ${ADMIN_TOKEN_ENV})", + ) + return parser + + +def _token(value: str, source: str, label: str) -> str: + if _TOKEN.fullmatch(value) is None: + raise _UsageError(f"the {label} in {source} must be one run of printable ASCII characters without spaces") + return value + + +def _secret(path: Path | None, env_name: str | None, label: str) -> str | None: + if path is not None: + try: + value = path.read_text(encoding="utf-8").strip() + except (OSError, UnicodeError) as error: + reason = error.strerror if isinstance(error, OSError) else "it is not UTF-8 text" + raise _UsageError(f"cannot read the {label} file {str(path)!r}: {reason}") from None + if not value: + raise _UsageError(f"the {label} file {str(path)!r} is empty") + return _token(value, f"the file {str(path)!r}", label) + if env_name is None: + return None + value = os.environ.get(env_name, "").strip() + return _token(value, f"${env_name}", label) if value else None + + +def _load_plan(path: Path) -> Plan: + try: + with path.open(encoding="utf-8") as handle: + document = json.load(handle) + except OSError as error: + raise _UsageError(f"cannot read the plan {str(path)!r}: {error.strerror}") from None + except ValueError as error: + raise _UsageError(f"the plan {str(path)!r} is not JSON: {error}") from None + return Plan.from_dict(document) + + +def _emit(result: keda.Tick) -> None: + for row in result.audit_rows(): + print(json.dumps(row), flush=True) + for error in result.series_errors: + print(f"sie-placement tune: Prometheus series {error}", file=sys.stderr, flush=True) + + +def _tune(args: argparse.Namespace) -> int: + try: + carriers = keda.carriers(_load_plan(args.plan), max_pool_floor=args.max_pool_floor, max_pools=args.max_pools) + keda.series_queries(args.series_selector) + prometheus = keda.PrometheusHttp(args.prometheus_url) + keda.require_base_url("gateway", args.gateway_url) + read_token = _secret(args.api_key_file, None, "gateway client token") + if read_token is None: + _secret(None, SDK_API_KEY_ENV, "gateway client token") + admin_token = _secret(args.admin_token_file, ADMIN_TOKEN_ENV, "gateway admin token") if args.apply else None + except ValueError as error: + print(f"sie-placement tune: {error}", file=sys.stderr) + return EXIT_REFUSED + try: + from sie_placement import pools # noqa: PLC0415 + except ImportError: + print("sie-placement tune needs the sie-placement keda extra, which installs the SIE SDK", file=sys.stderr) + return EXIT_REFUSED + if args.apply and admin_token is None: + print( + f"sie-placement tune: no gateway admin token in --admin-token-file or ${ADMIN_TOKEN_ENV}; " + "pool writes carry no credential", + file=sys.stderr, + ) + try: + reader = pools.connect(args.gateway_url, read_token) + writer = pools.connect(args.gateway_url, admin_token or "") if args.apply else None + except ValueError as error: + print(f"sie-placement tune: {error}", file=sys.stderr) + return EXIT_REFUSED + try: + while True: + started = time.monotonic() + result = keda.tick( + carriers, prometheus=prometheus, reader=reader, writer=writer, selector=args.series_selector + ) + _emit(result) + if args.once: + return EXIT_FAILED if result.failed else 0 + time.sleep(max(0.0, args.interval_s - (time.monotonic() - started))) + except KeyboardInterrupt: + return 0 + finally: + reader.close() + if writer is not None: + writer.close() + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parser().parse_args(argv) + return _tune(args) diff --git a/packages/sie_placement/src/sie_placement/keda.py b/packages/sie_placement/src/sie_placement/keda.py new file mode 100644 index 000000000..2ae451ed3 --- /dev/null +++ b/packages/sie_placement/src/sie_placement/keda.py @@ -0,0 +1,494 @@ +"""Keep each plan lane's warm floor inside its envelope through one gateway floor-carrier pool per lane.""" + +from __future__ import annotations + +import http.client +import json +import math +import re +import urllib.error +import urllib.parse +import urllib.request +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Protocol + +from sie_placement.schema import Plan, PlanLane, Warmth +from sie_placement.tuner import Observation, Settings, Step, step + +CARRIER_PREFIX = "placement-floor" +DEFAULT_QUEUE_POOL = "default" +DEFAULT_MAX_POOL_FLOOR = 4 +DEFAULT_MAX_POOLS = 64 +MAX_POOL_NAME_LENGTH = 128 +PROMETHEUS_TIMEOUT_S = 10.0 +SNAPSHOT_FRESHNESS_S = 20 + +QUEUE_DEPTH = "queueDepth" +PENDING_DEMAND = "pendingDemand" +WARM_FLOOR = "warmFloor" +RUNNERS = "runners" + +UNREALIZED_WARMTHS = { + Warmth.BRIDGED: "the overlay configures no remote bridge to serve it while the lane boots", + Warmth.SNAPSHOT_COLD: "the chart restores no snapshot, so the lane would boot cold", + Warmth.SCHEDULED_WARM: "the chart keeps no warm window, so the lane would stay cold inside it", +} + +_LANE_TOKEN = re.compile(r"[a-z0-9_-]{1,63}") +_SELECTOR_FORBIDDEN = re.compile(r"[{}\r\n]") +_LANE_LABELS = "pool, machine_profile, bundle" +_PRODUCER_LABELS = "producer_instance, collector_generation" +_LANE_TOKEN_COUNT = 3 +_VALUE_PAIR = 2 + + +class PlanRefusedError(ValueError): + def __init__(self, problems: Sequence[str]) -> None: + super().__init__("\n".join(problems)) + self.problems = tuple(problems) + + +class GatewayError(RuntimeError): + pass + + +class PrometheusError(RuntimeError): + pass + + +class CarrierConflictError(ValueError): + pass + + +@dataclass(frozen=True, slots=True) +class LaneKey: + queue_pool: str + machine_profile: str + bundle: str + + @classmethod + def canonical(cls, queue_pool: str, machine_profile: str, bundle: str) -> LaneKey | None: + tokens = [value.strip().lower() for value in (queue_pool, machine_profile, bundle)] + if any(_LANE_TOKEN.fullmatch(token) is None for token in tokens): + return None + return cls(queue_pool=tokens[0], machine_profile=tokens[1], bundle=tokens[2]) + + @property + def otel_lane(self) -> str: + return f"{self.queue_pool}|{self.machine_profile}|{self.bundle}" + + @property + def chart_lane(self) -> str: + return f"{self.queue_pool}/{self.machine_profile}/{self.bundle}" + + +@dataclass(frozen=True, slots=True) +class Carrier: + lane: PlanLane + key: LaneKey + + @property + def name(self) -> str: + return f"{CARRIER_PREFIX}-{self.key.queue_pool}-{self.key.machine_profile}-{self.key.bundle}" + + @property + def gpus(self) -> dict[str, int]: + return {self.key.machine_profile: 0} + + +def carriers( + plan: Plan, *, max_pool_floor: int = DEFAULT_MAX_POOL_FLOOR, max_pools: int = DEFAULT_MAX_POOLS +) -> tuple[Carrier, ...]: + problems: list[str] = [] + if len(plan.lanes) > max_pools: + problems.append( + f"$.lanes: {len(plan.lanes)} lanes need as many floor-carrier pools, " + f"above the gateway's limit of {max_pools} API pools" + ) + result: list[Carrier] = [] + first_index: dict[str, int] = {} + pins: dict[str, dict[str, int]] = {} + for index, entry in enumerate(plan.identities): + if entry.warmth is Warmth.PINNED: + pins.setdefault(entry.lane, {})[entry.identity] = index + for index, lane in enumerate(plan.lanes): + path = f"$.lanes[{index}]" + key = LaneKey.canonical(lane.queue_pool, lane.cell.machine_profile, lane.bundle) + if key is None: + problems.append( + f"{path}: queuePool, cell.machineProfile and bundle must each be 1 to 63 letters, digits, '_' or '-'" + ) + continue + pinned = pins.get(lane.name) + if key.queue_pool == DEFAULT_QUEUE_POOL and pinned: + identity, pin = min(pinned.items(), key=lambda item: item[1]) + problems.append( + f"$.identities[{pin}].warmth: pins {identity!r} on chart lane {key.chart_lane!r}, but the gateway " + f"keeps no static queue pool for the {DEFAULT_QUEUE_POOL!r} queue pool, so its pinnedModels have no " + "carrier; place pinned identities on a named queuePool" + ) + if lane.envelope.floor_max > max_pool_floor: + problems.append( + f"{path}.envelope.floorMax: {lane.envelope.floor_max} exceeds the gateway's " + f"per-pool floor limit of {max_pool_floor}" + ) + carrier = Carrier(lane=lane, key=key) + if len(carrier.name) > MAX_POOL_NAME_LENGTH: + problems.append( + f"{path}: floor-carrier pool name {carrier.name!r} is longer than {MAX_POOL_NAME_LENGTH} characters" + ) + earlier = first_index.setdefault(carrier.name, index) + if earlier != index: + problems.append(f"{path}: needs the same floor-carrier pool {carrier.name!r} as $.lanes[{earlier}]") + result.append(carrier) + for index, entry in enumerate(plan.identities): + reason = UNREALIZED_WARMTHS.get(entry.warmth) + if reason is not None: + problems.append( + f"$.identities[{index}].warmth: is {entry.warmth.value!r} for {entry.identity!r} on lane " + f"{entry.lane!r}, which the sie-cluster chart does not realize: {reason}; plan the identity pinned, " + "lazy-resident or cold-with-bound" + ) + if problems: + raise PlanRefusedError(problems) + return tuple(result) + + +@dataclass(frozen=True, slots=True) +class Sample: + labels: Mapping[str, str] + value: float + + +class Prometheus(Protocol): + def query(self, promql: str) -> Sequence[Sample]: ... + + +def _object(value: object) -> dict[str, object] | None: + if not isinstance(value, Mapping): + return None + fields: dict[str, object] = {} + for key, item in value.items(): + if not isinstance(key, str): + return None + fields[key] = item + return fields + + +def parse_vector(document: object) -> tuple[Sample, ...]: + response = _object(document) + if response is None or response.get("status") != "success": + raise PrometheusError("the query did not succeed") + data = _object(response.get("data")) + if data is None or data.get("resultType") != "vector": + raise PrometheusError("the query did not return an instant vector") + result = data.get("result") + if not isinstance(result, list): + raise PrometheusError("the query result is not a list") + samples: list[Sample] = [] + for item in result: + entry = _object(item) or {} + metric = _object(entry.get("metric")) + value = entry.get("value") + if metric is None or not isinstance(value, list) or len(value) != _VALUE_PAIR: + raise PrometheusError("the query returned a malformed sample") + text = value[1] + try: + number = float(text) if isinstance(text, str) else None + except ValueError: + number = None + if number is None: + raise PrometheusError("the query returned a sample that is not a number") + samples.append(Sample(labels={str(name): str(label) for name, label in metric.items()}, value=number)) + return tuple(samples) + + +def require_base_url(label: str, base_url: str) -> None: + try: + parts = urllib.parse.urlsplit(base_url) + scheme, hostname, username, password = parts.scheme, parts.hostname, parts.username, parts.password + except ValueError: + raise ValueError(f"the {label} URL is not a valid URL") from None + if scheme not in {"http", "https"} or not hostname: + raise ValueError(f"the {label} URL must be an absolute http or https URL") + if username is not None or password is not None: + raise ValueError(f"the {label} URL must not embed credentials") + if parts.query or parts.fragment: + raise ValueError(f"the {label} URL must not carry a query or a fragment") + try: + parts.port # noqa: B018 + except ValueError: + raise ValueError(f"the {label} URL has an invalid port") from None + + +class PrometheusHttp: + def __init__(self, base_url: str, *, timeout_s: float = PROMETHEUS_TIMEOUT_S) -> None: + require_base_url("Prometheus", base_url) + self._endpoint = f"{base_url.rstrip('/')}/api/v1/query" + self._timeout_s = timeout_s + + def query(self, promql: str) -> tuple[Sample, ...]: + url = f"{self._endpoint}?{urllib.parse.urlencode({'query': promql})}" + try: + with urllib.request.urlopen(url, timeout=self._timeout_s) as response: # noqa: S310 + document = json.load(response) + except urllib.error.HTTPError as error: + raise PrometheusError(f"the query failed with HTTP {error.code}") from None + except (OSError, ValueError) as error: + raise PrometheusError(f"the query failed: {error}") from None + except (http.client.HTTPException, RecursionError) as error: + raise PrometheusError(f"the query failed: {type(error).__name__}") from None + return parse_vector(document) + + +def series_queries(selector: str = "") -> dict[str, str]: + if _SELECTOR_FORBIDDEN.search(selector): + raise ValueError("the series selector must be label matchers without braces or line breaks") + + def matchers(*extra: str) -> str: + return ",".join(part for part in (selector.strip(), *extra) if part) + + def fresh(metric: str, timestamp: str, join_labels: str) -> str: + return ( + f"max by ({_LANE_LABELS}) ({metric}{{{matchers()}}} and on({join_labels}) " + f"(abs(time() - {timestamp}{{{matchers()}}}) < {SNAPSHOT_FRESHNESS_S}))" + ) + + capacity = "sie_gateway_capacity_snapshot_timestamp_seconds" + return { + QUEUE_DEPTH: fresh( + "sie_gateway_lane_queue_depth", + "sie_gateway_lane_queue_snapshot_timestamp_seconds", + f"{_PRODUCER_LABELS}, {_LANE_LABELS}", + ), + PENDING_DEMAND: fresh("sie_gateway_pending_demand", capacity, _PRODUCER_LABELS), + WARM_FLOOR: fresh("sie_gateway_pool_warm_floor", capacity, _PRODUCER_LABELS), + RUNNERS: ( + "count by (lane) (max by (lane, producer_instance) " + f"(sie_worker_gpu_slots{{{matchers('state="ready"')}}}) > 0)" + ), + } + + +def _sample_lane(labels: Mapping[str, str], *, otel_lane: bool) -> LaneKey | None: + if otel_lane: + tokens = labels.get("lane", "").split("|") + if len(tokens) != _LANE_TOKEN_COUNT: + return None + return LaneKey.canonical(*tokens) + return LaneKey.canonical(labels.get("pool", ""), labels.get("machine_profile", ""), labels.get("bundle", "")) + + +def _count(value: float) -> int | None: + if not math.isfinite(value) or value < 0: + return None + return round(value) + + +def _lane_values(samples: Sequence[Sample], *, otel_lane: bool) -> dict[LaneKey, int]: + values: dict[LaneKey, int] = {} + for sample in samples: + key = _sample_lane(sample.labels, otel_lane=otel_lane) + count = _count(sample.value) + if key is not None and count is not None: + values[key] = max(count, values.get(key, 0)) + return values + + +@dataclass(frozen=True, slots=True, kw_only=True) +class LaneTelemetry: + queue_depth: int | None = None + pending_demand: int | None = None + warm_floor: int | None = None + runners: int | None = None + + def observation(self, current_floor: int) -> Observation: + if self.queue_depth is None or self.pending_demand is None or self.runners is None: + return Observation(runners=current_floor, backlog=0, running=1) + return Observation(runners=self.runners, backlog=self.pending_demand, running=self.queue_depth) + + def to_dict(self) -> dict[str, object]: + return { + QUEUE_DEPTH: self.queue_depth, + PENDING_DEMAND: self.pending_demand, + WARM_FLOOR: self.warm_floor, + RUNNERS: self.runners, + } + + +@dataclass(frozen=True, slots=True) +class Telemetry: + series: Mapping[str, Mapping[LaneKey, int] | None] + errors: tuple[str, ...] = () + + def _value(self, name: str, key: LaneKey, absent: int | None) -> int | None: + values = self.series.get(name) + if values is None: + return None + return values.get(key, absent) + + def lane(self, key: LaneKey) -> LaneTelemetry: + return LaneTelemetry( + queue_depth=self._value(QUEUE_DEPTH, key, None), + pending_demand=self._value(PENDING_DEMAND, key, None), + warm_floor=self._value(WARM_FLOOR, key, None), + runners=self._value(RUNNERS, key, 0), + ) + + +def read_telemetry(prometheus: Prometheus, selector: str = "") -> Telemetry: + series: dict[str, Mapping[LaneKey, int] | None] = {} + errors: list[str] = [] + for name, promql in series_queries(selector).items(): + try: + samples = prometheus.query(promql) + except PrometheusError as error: + errors.append(f"{name}: {error}") + series[name] = None + continue + series[name] = _lane_values(samples, otel_lane=name == RUNNERS) + return Telemetry(series=series, errors=tuple(errors)) + + +class PoolApi(Protocol): + def get_pool(self, name: str) -> Mapping[str, object] | None: ... + + def create_pool( + self, name: str, gpus: dict[str, int], *, bundle: str, minimum_worker_count: int, queue_pool: str + ) -> None: ... + + +def _matches(value: object, expected: str) -> bool: + return isinstance(value, str) and value.strip().lower() == expected + + +def carrier_floor(carrier: Carrier, pool: Mapping[str, object]) -> int: + spec = _object(pool.get("spec")) + if spec is None: + raise CarrierConflictError("it has no spec") + if not _matches(spec.get("queue_pool", "default"), carrier.key.queue_pool): + raise CarrierConflictError(f"it does not draw from queue pool {carrier.key.queue_pool!r}") + if not _matches(spec.get("bundle"), carrier.key.bundle): + raise CarrierConflictError(f"it is not filtered to bundle {carrier.key.bundle!r}") + gpus = _object(spec.get("gpus")) + if not ( + gpus is not None + and len(gpus) == 1 + and all( + _matches(profile, carrier.key.machine_profile) and count == 0 and not isinstance(count, bool) + for profile, count in gpus.items() + ) + ): + raise CarrierConflictError(f"its gpus are not {carrier.gpus!r}") + for field in ("gpu_caps", "pinned_models", "ttl_seconds"): + if spec.get(field): + raise CarrierConflictError(f"it sets {field}") + floor = spec.get("minimum_worker_count", 0) + if isinstance(floor, bool) or not isinstance(floor, int) or floor < 0: + raise CarrierConflictError("its minimum_worker_count is not a non-negative integer") + return floor + + +@dataclass(frozen=True, slots=True, kw_only=True) +class LaneRecord: + carrier: Carrier + telemetry: LaneTelemetry + current_floor: int | None = None + step: Step | None = None + applied: bool = False + error: str | None = None + + @property + def lane_floor_above_envelope(self) -> bool: + warm_floor = self.telemetry.warm_floor + return warm_floor is not None and warm_floor > self.carrier.lane.envelope.floor_max + + def to_dict(self) -> dict[str, object]: + return { + "lane": self.carrier.lane.name, + "queuePool": self.carrier.key.queue_pool, + "machineProfile": self.carrier.key.machine_profile, + "bundle": self.carrier.key.bundle, + "carrier": self.carrier.name, + "telemetry": self.telemetry.to_dict(), + "currentFloor": self.current_floor, + "action": None if self.step is None else self.step.action.value, + "floor": None if self.step is None else self.step.settings.floor, + "applied": self.applied, + "laneFloorAboveEnvelope": self.lane_floor_above_envelope, + "error": self.error, + } + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Tick: + at: datetime + apply: bool + records: tuple[LaneRecord, ...] + series_errors: tuple[str, ...] = () + + @property + def failed(self) -> bool: + return bool(self.series_errors) or any(record.error is not None for record in self.records) + + def audit_rows(self) -> list[dict[str, object]]: + at = self.at.astimezone(UTC).strftime("%Y-%m-%dT%H:%M:%SZ") + mode = "apply" if self.apply else "shadow" + return [{"at": at, "mode": mode, **record.to_dict()} for record in self.records] + + +def _tune_lane(carrier: Carrier, telemetry: LaneTelemetry, reader: PoolApi, writer: PoolApi | None) -> LaneRecord: + try: + pool = reader.get_pool(carrier.name) + except GatewayError as error: + return LaneRecord(carrier=carrier, telemetry=telemetry, error=f"reading pool {carrier.name!r} failed: {error}") + try: + current = 0 if pool is None else carrier_floor(carrier, pool) + except CarrierConflictError as error: + return LaneRecord( + carrier=carrier, + telemetry=telemetry, + error=f"pool {carrier.name!r} is not this lane's floor carrier ({error}) and is left unchanged", + ) + result = step(carrier.lane.envelope, telemetry.observation(current), Settings(floor=current)) + record = LaneRecord(carrier=carrier, telemetry=telemetry, current_floor=current, step=result) + if writer is None or (pool is None and result.settings.floor == 0): + return record + try: + writer.create_pool( + carrier.name, + carrier.gpus, + bundle=carrier.key.bundle, + minimum_worker_count=result.settings.floor, + queue_pool=carrier.key.queue_pool, + ) + except GatewayError as error: + return LaneRecord( + carrier=carrier, + telemetry=telemetry, + current_floor=current, + step=result, + error=f"writing pool {carrier.name!r} failed: {error}", + ) + return LaneRecord(carrier=carrier, telemetry=telemetry, current_floor=current, step=result, applied=True) + + +def _utc_now() -> datetime: + return datetime.now(UTC) + + +def tick( + lanes: Sequence[Carrier], + *, + prometheus: Prometheus, + reader: PoolApi, + writer: PoolApi | None = None, + selector: str = "", + now: Callable[[], datetime] = _utc_now, +) -> Tick: + at = now() + telemetry = read_telemetry(prometheus, selector) + records = tuple(_tune_lane(carrier, telemetry.lane(carrier.key), reader, writer) for carrier in lanes) + return Tick(at=at, apply=writer is not None, records=records, series_errors=telemetry.errors) diff --git a/packages/sie_placement/src/sie_placement/pools.py b/packages/sie_placement/src/sie_placement/pools.py new file mode 100644 index 000000000..f91556f37 --- /dev/null +++ b/packages/sie_placement/src/sie_placement/pools.py @@ -0,0 +1,102 @@ +"""Gateway pool reads and floor-carrier writes through the SIE Python SDK, installed by the keda extra.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import httpx +from sie_sdk import SIEClient, SIEError + +from sie_placement.keda import GatewayError + +TIMEOUT_S = 10.0 +HTTP_OK_MIN = 200 +HTTP_OK_MAX = 299 +HTTP_ERROR_MIN = 400 +MAX_ERROR_LENGTH = 512 +REDACTED = "[redacted]" +_BEARER = "Bearer " +_UNEXPECTED_REPLY = "the gateway returned an unexpected response" + + +class _Replies: + def __init__(self) -> None: + self.status: int | None = None + + def __call__(self, response: httpx.Response) -> None: + status = response.status_code + self.status = status + if status < HTTP_OK_MIN or HTTP_OK_MAX < status < HTTP_ERROR_MIN: + raise GatewayError(f"the gateway answered HTTP {status}, which is neither a success nor an error") + + +def _client(base_url: str, api_key: str | None) -> tuple[SIEClient, _Replies, str | None]: + replies = _Replies() + try: + http_client = httpx.Client( + base_url=base_url, + timeout=TIMEOUT_S, + follow_redirects=False, + event_hooks={"response": [replies]}, + ) + except httpx.InvalidURL: + raise ValueError("the gateway URL is not a valid URL") from None + try: + client = SIEClient(base_url, api_key=api_key, timeout_s=TIMEOUT_S, http_client=http_client) + except UnicodeError: + http_client.close() + raise ValueError("a gateway token holds a character that an HTTP header cannot carry") from None + except BaseException: + http_client.close() + raise + authorization = http_client.headers.get("Authorization", "") + token = authorization.removeprefix(_BEARER) if authorization.startswith(_BEARER) else "" + return client, replies, token or None + + +def _gateway_error(error: Exception, status: int | None, token: str | None) -> GatewayError: + if isinstance(error, httpx.HTTPError): + return GatewayError(f"the request failed ({type(error).__name__})") + if status is not None and status >= HTTP_ERROR_MIN: + return GatewayError(f"the gateway answered HTTP {status}") + if status is not None or isinstance(error, AttributeError | RecursionError): + return GatewayError(_UNEXPECTED_REPLY) + message = str(error) + if token: + message = message.replace(token, REDACTED) + return GatewayError(message[:MAX_ERROR_LENGTH]) + + +class SdkPools: + def __init__(self, base_url: str, api_key: str | None) -> None: + self._base_url = base_url + self._api_key = api_key + self._client, self._replies, self._token = _client(base_url, api_key) + + def get_pool(self, name: str) -> Mapping[str, object] | None: + self._replies.status = None + try: + pool = self._client.get_pool(name) + except (SIEError, httpx.HTTPError, ValueError, AttributeError, RecursionError) as error: + raise _gateway_error(error, self._replies.status, self._token) from None + return pool + + def create_pool( + self, name: str, gpus: dict[str, int], *, bundle: str, minimum_worker_count: int, queue_pool: str + ) -> None: + replies = None + try: + client, replies, _ = _client(self._base_url, self._api_key) + with client: + client.create_pool( + name, gpus, bundle=bundle, minimum_worker_count=minimum_worker_count, queue_pool=queue_pool + ) + except (SIEError, httpx.HTTPError, ValueError, AttributeError, RecursionError) as error: + raise _gateway_error(error, None if replies is None else replies.status, self._token) from None + + def close(self) -> None: + self._client.close() + + +def connect(base_url: str, api_key: str | None) -> SdkPools: + return SdkPools(base_url, api_key) diff --git a/packages/sie_placement/tests/fixtures/tune-plan.json b/packages/sie_placement/tests/fixtures/tune-plan.json new file mode 100644 index 000000000..b36c3b28f --- /dev/null +++ b/packages/sie_placement/tests/fixtures/tune-plan.json @@ -0,0 +1,68 @@ +{ + "kind": "placement-plan", + "schemaVersion": 1, + "placement": "us-central1", + "lanes": [ + { + "name": "default-l4-default", + "queuePool": "default", + "bundle": "default", + "cell": { + "machineProfile": "l4", + "gpuProduct": "NVIDIA-L4" + }, + "containerShape": { + "gpuCount": 1, + "cpuCores": 7.5, + "memoryMiB": 30720 + }, + "concurrency": 16, + "maxContainers": 6, + "envelope": { + "floorMin": 1, + "floorMax": 3, + "bufferMin": 0, + "bufferMax": 1, + "scaledownWindowMinS": 60, + "scaledownWindowMaxS": 900 + } + }, + { + "name": "default-a100-80gb-generation", + "queuePool": "default", + "bundle": "generation", + "cell": { + "machineProfile": "a100-80gb", + "gpuProduct": "NVIDIA-A100-SXM4-80GB" + }, + "containerShape": { + "gpuCount": 1, + "cpuCores": 12, + "memoryMiB": 87040 + }, + "concurrency": 32, + "maxContainers": 2, + "envelope": { + "floorMin": 0, + "floorMax": 1 + } + } + ], + "identities": [ + { + "identity": "BAAI/bge-m3", + "lane": "default-l4-default", + "warmth": "lazy-resident" + }, + { + "identity": "BAAI/bge-m3:sparse", + "lane": "default-l4-default", + "warmth": "lazy-resident" + }, + { + "identity": "Qwen/Qwen3-8B", + "lane": "default-a100-80gb-generation", + "warmth": "cold-with-bound" + } + ] +} diff --git a/packages/sie_placement/tests/test_cli.py b/packages/sie_placement/tests/test_cli.py new file mode 100644 index 000000000..5a0a1fc16 --- /dev/null +++ b/packages/sie_placement/tests/test_cli.py @@ -0,0 +1,709 @@ +from __future__ import annotations + +import json +import logging +import sys +import threading +import time +from collections.abc import Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from types import SimpleNamespace +from urllib.parse import parse_qs, urlsplit + +import pytest + +pytest.importorskip("sie_sdk") + +from sie_placement import cli, keda, pools +from sie_sdk import SIEError + +FIXTURES = Path(__file__).parent / "fixtures" +READ_TOKEN = "read-token-for-tests" # noqa: S105 - a test value +ADMIN_TOKEN = "admin-token-for-tests" # noqa: S105 - a test value +L4 = {"pool": "default", "machine_profile": "l4", "bundle": "default"} +A100 = {"pool": "default", "machine_profile": "a100-80gb", "bundle": "generation"} + + +def vector(*samples: tuple[dict[str, str], float]) -> dict: + return { + "status": "success", + "data": { + "resultType": "vector", + "result": [{"metric": labels, "value": [0, str(value)]} for labels, value in samples], + }, + } + + +class _Handler(BaseHTTPRequestHandler): + server: _Server + + def _reply(self, status: int, body: object) -> None: + payload = json.dumps(body).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def _body(self) -> object: + length = int(self.headers.get("Content-Length") or 0) + return json.loads(self.rfile.read(length)) if length else None + + def do_GET(self) -> None: + self.server.respond(self, "GET", None) + + def do_POST(self) -> None: + self.server.respond(self, "POST", self._body()) + + def log_message(self, format: str, *args: object) -> None: + return + + +class _Server(ThreadingHTTPServer): + def __init__(self) -> None: + super().__init__(("127.0.0.1", 0), _Handler) + self.requests: list[dict] = [] + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.server_address[1]}" + + def respond(self, handler: _Handler, method: str, body: object) -> None: + raise NotImplementedError + + +class FakeGateway(_Server): + def __init__(self, *, auth: bool = True) -> None: + super().__init__() + self.auth = auth + self.pools: dict[str, dict] = {} + self.fail_reads = False + self.read_reply: object | None = None + self.raw_reply: tuple[int, bytes] | None = None + self.raw_replies: dict[str, tuple[int, bytes]] = {} + + def respond(self, handler: _Handler, method: str, body: object) -> None: + path = urlsplit(handler.path).path + authorization = handler.headers.get("Authorization") + self.requests.append({"method": method, "path": path, "authorization": authorization, "body": body}) + raw = self.raw_replies.get(method, self.raw_reply) + if raw is not None: + status, payload = raw + handler.send_response(status) + handler.send_header("Content-Type", "application/json") + handler.send_header("Content-Length", str(len(payload))) + handler.send_header("Location", "http://127.0.0.1:9/v1/pools") + handler.end_headers() + handler.wfile.write(payload) + return + if method == "GET" and path.startswith("/v1/pools/"): + if self.read_reply is not None: + handler._reply(200, self.read_reply) + return + if self.fail_reads or (self.auth and authorization != f"Bearer {READ_TOKEN}"): + handler._reply(401, {"detail": {"code": "UNAUTHORIZED", "message": "Invalid token"}}) + return + spec = self.pools.get(path.removeprefix("/v1/pools/")) + if spec is None: + handler._reply(404, {"detail": {"code": "POOL_NOT_FOUND", "message": "Pool not found"}}) + return + handler._reply(200, {"spec": spec, "status": {"state": "active", "assigned_workers": []}}) + return + if method == "POST" and path.startswith("/v1/pools"): + if self.auth and authorization != f"Bearer {ADMIN_TOKEN}": + handler._reply(403, {"detail": {"code": "FORBIDDEN", "message": "Admin token required"}}) + return + if path == "/v1/pools" and isinstance(body, dict): + spec = { + "name": body["name"], + "queue_pool": body.get("queue_pool", "default"), + "bundle": body.get("bundle"), + "gpus": body.get("gpus", {}), + "ttl_seconds": body.get("ttl_seconds"), + "minimum_worker_count": body.get("minimum_worker_count", 0), + } + self.pools[body["name"]] = spec + handler._reply(201, {"spec": spec, "status": {"state": "pending", "assigned_workers": []}}) + return + handler._reply(200, {"message": "Pool renewed"}) + return + handler._reply(404, {"detail": {"code": "NOT_FOUND", "message": "not found"}}) + + +class FakePrometheus(_Server): + def __init__(self) -> None: + super().__init__() + self.series: dict[str, dict] = { + "sie_gateway_lane_queue_depth": vector((L4, 3), (A100, 0)), + "sie_gateway_pending_demand": vector((L4, 1), (A100, 0)), + "sie_gateway_pool_warm_floor": vector((L4, 1), (A100, 0)), + "sie_worker_gpu_slots": vector(({"lane": "default|l4|default"}, 1)), + } + self.truncate = False + + def respond(self, handler: _Handler, method: str, body: object) -> None: + url = urlsplit(handler.path) + query = parse_qs(url.query).get("query", [""])[0] + self.requests.append({"path": url.path, "query": query, "authorization": handler.headers.get("Authorization")}) + if self.truncate: + handler.send_response(200) + handler.send_header("Content-Type", "application/json") + handler.send_header("Content-Length", "500") + handler.end_headers() + handler.wfile.write(b'{"status": "succ') + return + for metric, document in self.series.items(): + if metric in query: + handler._reply(200, document) + return + handler._reply(200, vector()) + + +def _serve(server: _Server) -> Iterator[_Server]: + thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + yield server + server.shutdown() + server.server_close() + + +@pytest.fixture(autouse=True) +def clean_environment(monkeypatch): + for name in ( + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + "http_proxy", + "https_proxy", + "all_proxy", + "SIE_API_KEY", + "SIE_BASE_URL", + "SIE_ADMIN_TOKEN", + ): + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def gateway() -> Iterator[FakeGateway]: + yield from _serve(FakeGateway()) + + +@pytest.fixture +def prometheus() -> Iterator[FakePrometheus]: + yield from _serve(FakePrometheus()) + + +@pytest.fixture +def read_token_file(tmp_path) -> Path: + path = tmp_path / "read-token" + path.write_text(READ_TOKEN + "\n", encoding="utf-8") + return path + + +def run(gateway: FakeGateway, prometheus: FakePrometheus, *extra: str, plan: Path = FIXTURES / "tune-plan.json") -> int: + return cli.main( + ["tune", "--plan", str(plan), "--prometheus-url", prometheus.url, "--gateway-url", gateway.url, *extra] + ) + + +def tune( + gateway: FakeGateway, prometheus: FakePrometheus, *extra: str, plan: Path = FIXTURES / "tune-plan.json" +) -> int: + return run(gateway, prometheus, "--once", *extra, plan=plan) + + +def rows(captured: str) -> list[dict]: + return [json.loads(line) for line in captured.splitlines()] + + +def assert_no_secret(capsys_text: str, caplog) -> None: + for secret in (READ_TOKEN, ADMIN_TOKEN): + assert secret not in capsys_text + assert secret not in caplog.text + + +def test_shadow_mode_records_decisions_and_writes_no_pool(gateway, prometheus, read_token_file, capsys, caplog): + caplog.set_level(logging.DEBUG) + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file)) == 0 + captured = capsys.readouterr() + lines = rows(captured.out) + assert [(row["mode"], row["carrier"], row["action"], row["floor"], row["applied"]) for row in lines] == [ + ("shadow", "placement-floor-default-l4-default", "raise_to_floor", 1, False), + ("shadow", "placement-floor-default-a100-80gb-generation", "hold", 0, False), + ] + assert lines[0]["telemetry"] == {"queueDepth": 3, "pendingDemand": 1, "warmFloor": 1, "runners": 1} + assert [request["method"] for request in gateway.requests] == ["GET", "GET"] + assert {request["authorization"] for request in gateway.requests} == {f"Bearer {READ_TOKEN}"} + assert [request["path"] for request in prometheus.requests] == ["/api/v1/query"] * 4 + assert_no_secret(captured.out + captured.err, caplog) + + +def test_shadow_mode_never_reads_the_admin_token(gateway, prometheus, read_token_file, monkeypatch, capsys): + monkeypatch.setenv("SIE_ADMIN_TOKEN", ADMIN_TOKEN) + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file), "--admin-token-file", "/nonexistent") == 0 + assert all(request["method"] == "GET" for request in gateway.requests) + assert all(request["authorization"] != f"Bearer {ADMIN_TOKEN}" for request in gateway.requests) + capsys.readouterr() + + +def test_apply_upserts_carriers_with_the_admin_token_only_on_writes( + gateway, prometheus, read_token_file, monkeypatch, capsys, caplog +): + caplog.set_level(logging.DEBUG) + monkeypatch.setenv("SIE_ADMIN_TOKEN", ADMIN_TOKEN) + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file), "--apply") == 0 + captured = capsys.readouterr() + writes = [request for request in gateway.requests if request["method"] == "POST"] + reads = [request for request in gateway.requests if request["method"] == "GET"] + assert [request["body"] for request in writes] == [ + { + "name": "placement-floor-default-l4-default", + "gpus": {"l4": 0}, + "queue_pool": "default", + "bundle": "default", + "minimum_worker_count": 1, + } + ] + assert {request["authorization"] for request in writes} == {f"Bearer {ADMIN_TOKEN}"} + assert {request["authorization"] for request in reads} == {f"Bearer {READ_TOKEN}"} + assert [(row["mode"], row["applied"]) for row in rows(captured.out)] == [("apply", True), ("apply", False)] + assert_no_secret(captured.out + captured.err, caplog) + + +def test_apply_reads_the_admin_token_from_a_file(gateway, prometheus, read_token_file, tmp_path, capsys): + admin = tmp_path / "admin-token" + admin.write_text(ADMIN_TOKEN, encoding="utf-8") + assert ( + tune(gateway, prometheus, "--api-key-file", str(read_token_file), "--apply", "--admin-token-file", str(admin)) + == 0 + ) + writes = [request for request in gateway.requests if request["method"] == "POST"] + assert {request["authorization"] for request in writes} == {f"Bearer {ADMIN_TOKEN}"} + capsys.readouterr() + + +def test_apply_renews_an_existing_carrier_and_steers_from_its_floor( + gateway, prometheus, read_token_file, monkeypatch, capsys +): + monkeypatch.setenv("SIE_ADMIN_TOKEN", ADMIN_TOKEN) + gateway.pools["placement-floor-default-l4-default"] = { + "name": "placement-floor-default-l4-default", + "queue_pool": "default", + "bundle": "default", + "gpus": {"l4": 0}, + "ttl_seconds": None, + "minimum_worker_count": 1, + } + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file), "--apply") == 0 + l4_row = rows(capsys.readouterr().out)[0] + assert (l4_row["currentFloor"], l4_row["action"], l4_row["floor"]) == (1, "scale_up", 2) + assert gateway.pools["placement-floor-default-l4-default"]["minimum_worker_count"] == 2 + + +def test_the_client_token_from_the_environment_is_never_sent_on_writes(prometheus, monkeypatch, capsys): + unauthenticated = FakeGateway(auth=False) + for server in _serve(unauthenticated): + monkeypatch.setenv("SIE_BASE_URL", server.url) + monkeypatch.setenv("SIE_API_KEY", READ_TOKEN) + assert tune(server, prometheus, "--apply") == 0 + reads = [request for request in server.requests if request["method"] == "GET"] + writes = [request for request in server.requests if request["method"] == "POST"] + assert {request["authorization"] for request in reads} == {f"Bearer {READ_TOKEN}"} + assert writes + assert {request["authorization"] for request in writes} == {None} + assert "pool writes carry no credential" in capsys.readouterr().err + + +def test_a_rejected_write_fails_the_tick_and_names_the_pool(gateway, prometheus, read_token_file, capsys): + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file), "--apply") == 1 + l4_row = rows(capsys.readouterr().out)[0] + assert l4_row["applied"] is False + assert l4_row["error"] == "writing pool 'placement-floor-default-l4-default' failed: the gateway answered HTTP 403" + + +def test_an_envelope_above_the_pool_floor_limit_is_refused_before_any_call( + gateway, prometheus, read_token_file, tmp_path, capsys +): + document = json.loads((FIXTURES / "tune-plan.json").read_text(encoding="utf-8")) + document["lanes"][0]["envelope"]["floorMax"] = 5 + plan = tmp_path / "plan.json" + plan.write_text(json.dumps(document), encoding="utf-8") + assert tune(gateway, prometheus, "--apply", plan=plan) == cli.EXIT_REFUSED + assert "floorMax: 5 exceeds the gateway's per-pool floor limit of 4" in capsys.readouterr().err + assert gateway.requests == [] + assert prometheus.requests == [] + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file), "--max-pool-floor", "5", plan=plan) == 0 + capsys.readouterr() + + +def test_a_plan_the_chart_cannot_run_is_refused_before_any_call(gateway, prometheus, tmp_path, capsys): + document = json.loads((FIXTURES / "tune-plan.json").read_text(encoding="utf-8")) + document["identities"][0]["warmth"] = "pinned" + plan = tmp_path / "plan.json" + plan.write_text(json.dumps(document), encoding="utf-8") + assert tune(gateway, prometheus, plan=plan) == cli.EXIT_REFUSED + assert "pins 'BAAI/bge-m3' on chart lane 'default/l4/default'" in capsys.readouterr().err + assert gateway.requests == [] + assert prometheus.requests == [] + + +def test_more_lanes_than_the_pool_limit_are_refused_before_any_call(gateway, prometheus, capsys): + assert tune(gateway, prometheus, "--apply", "--max-pools", "1") == cli.EXIT_REFUSED + assert "above the gateway's limit of 1 API pools" in capsys.readouterr().err + assert gateway.requests == [] + assert prometheus.requests == [] + + +@pytest.mark.parametrize( + ("extra", "message"), + [ + (("--api-key-file", "/nonexistent/token"), "cannot read the gateway client token file"), + (("--apply", "--admin-token-file", "/nonexistent/token"), "cannot read the gateway admin token file"), + (("--series-selector", 'a="b"}'), "without braces"), + ], +) +def test_bad_inputs_are_refused_before_any_call(gateway, prometheus, extra, message, capsys): + assert tune(gateway, prometheus, *extra) == cli.EXIT_REFUSED + assert message in capsys.readouterr().err + assert gateway.requests == [] + + +def test_an_empty_token_file_is_refused(gateway, prometheus, tmp_path, capsys): + empty = tmp_path / "token" + empty.write_text("\n", encoding="utf-8") + assert tune(gateway, prometheus, "--api-key-file", str(empty)) == cli.EXIT_REFUSED + assert "is empty" in capsys.readouterr().err + + +def test_a_plan_that_is_not_a_plan_is_refused(gateway, prometheus, tmp_path, capsys): + plan = tmp_path / "plan.json" + plan.write_text("{}", encoding="utf-8") + assert tune(gateway, prometheus, plan=plan) == cli.EXIT_REFUSED + assert "$.kind: is required" in capsys.readouterr().err + plan.write_text("{", encoding="utf-8") + assert tune(gateway, prometheus, plan=plan) == cli.EXIT_REFUSED + assert "is not JSON" in capsys.readouterr().err + + +def test_an_unreachable_prometheus_holds_every_lane_and_fails_the_tick(gateway, read_token_file, capsys): + unreachable = FakePrometheus() + unreachable.server_close() + assert tune(gateway, unreachable, "--api-key-file", str(read_token_file)) == 1 + captured = capsys.readouterr() + assert [row["action"] for row in rows(captured.out)] == ["raise_to_floor", "hold"] + assert "Prometheus series queueDepth: the query failed" in captured.err + + +def test_the_core_imports_without_the_sdk_and_tune_names_the_missing_extra(gateway, prometheus, monkeypatch, capsys): + monkeypatch.setitem(sys.modules, "sie_sdk", None) + monkeypatch.delitem(sys.modules, "sie_placement.pools", raising=False) + monkeypatch.delattr(sys.modules["sie_placement"], "pools", raising=False) + assert tune(gateway, prometheus) == cli.EXIT_REFUSED + assert "needs the sie-placement keda extra" in capsys.readouterr().err + assert gateway.requests == [] + + +TWO_LINE_SECRET = "first-line-of-a-secret\nsecond-line-of-a-secret" # noqa: S105 - a test value + + +def assert_secret_absent(secret: str, captured, caplog) -> None: + for part in (secret, *secret.split("\n")): + assert part not in captured.out + assert part not in captured.err + assert part not in caplog.text + + +@pytest.mark.parametrize("secret", [TWO_LINE_SECRET, "secret\x00value"]) +@pytest.mark.parametrize("flag", ["--api-key-file", "--admin-token-file"]) +def test_a_token_file_with_a_control_character_is_refused_without_echoing_it( + gateway, prometheus, tmp_path, flag, secret, capsys, caplog +): + caplog.set_level(logging.DEBUG) + token = tmp_path / "token" + token.write_text(secret + "\n", encoding="utf-8") + assert tune(gateway, prometheus, "--apply", flag, str(token)) == cli.EXIT_REFUSED + captured = capsys.readouterr() + assert "must be one run of printable ASCII characters without spaces" in captured.err + assert_secret_absent(secret, captured, caplog) + assert gateway.requests == [] + assert prometheus.requests == [] + + +@pytest.mark.parametrize("value", [TWO_LINE_SECRET, "secret with spaces", "caf\u00e9-secret"]) +def test_an_admin_token_variable_that_a_header_cannot_carry_is_refused( + gateway, prometheus, read_token_file, monkeypatch, value, capsys, caplog +): + caplog.set_level(logging.DEBUG) + monkeypatch.setenv("SIE_ADMIN_TOKEN", value) + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file), "--apply") == cli.EXIT_REFUSED + captured = capsys.readouterr() + assert "the gateway admin token in $SIE_ADMIN_TOKEN must be one run of printable ASCII" in captured.err + assert_secret_absent(value, captured, caplog) + assert gateway.requests == [] + + +def test_a_token_file_that_is_not_utf8_is_refused_without_echoing_it(gateway, prometheus, tmp_path, capsys): + token = tmp_path / "token" + token.write_bytes(b"secret\xff\xfevalue") + assert tune(gateway, prometheus, "--api-key-file", str(token)) == cli.EXIT_REFUSED + err = capsys.readouterr().err + assert "it is not UTF-8 text" in err + assert "secret" not in err + assert "0xff" not in err + + +@pytest.mark.parametrize("value", [TWO_LINE_SECRET, "caf\u00e9-secret"]) +def test_an_sdk_client_token_that_a_header_cannot_carry_is_refused( + gateway, prometheus, monkeypatch, value, capsys, caplog +): + caplog.set_level(logging.DEBUG) + monkeypatch.setenv("SIE_BASE_URL", gateway.url) + monkeypatch.setenv("SIE_API_KEY", value) + assert tune(gateway, prometheus) == cli.EXIT_REFUSED + captured = capsys.readouterr() + assert "the gateway client token in $SIE_API_KEY must be one run of printable ASCII" in captured.err + assert_secret_absent(value, captured, caplog) + assert gateway.requests == [] + + +def test_a_transport_failure_is_reported_without_the_header_it_carried(gateway): + reader = pools.connect(gateway.url, TWO_LINE_SECRET) + try: + with pytest.raises(keda.GatewayError) as raised: + reader.get_pool("placement-floor-default-l4-default") + finally: + reader.close() + assert str(raised.value) == "the request failed (LocalProtocolError)" + writer = pools.connect(gateway.url, TWO_LINE_SECRET) + try: + with pytest.raises(keda.GatewayError) as raised: + writer.create_pool( + "placement-floor-default-l4-default", + {"l4": 0}, + bundle="default", + minimum_worker_count=1, + queue_pool="default", + ) + finally: + writer.close() + assert str(raised.value) == "the request failed (LocalProtocolError)" + assert gateway.pools == {} + + +PASSWORD = "s3cret-password" # noqa: S105 - a test value + + +@pytest.mark.parametrize( + ("url", "message"), + [ + ("sie-gateway:8080", "must be an absolute http or https URL"), + ("ftp://sie-gateway:8080", "must be an absolute http or https URL"), + ("http:///v1", "must be an absolute http or https URL"), + (f"http://user:{PASSWORD}@sie-gateway:8080", "must not embed credentials"), + (f"http://user:{PASSWORD}@sie-gateway:notaport", "must not embed credentials"), + (f"http://sie-gateway:8080/?token={PASSWORD}", "must not carry a query or a fragment"), + (f"http://sie-gateway:8080/#{PASSWORD}", "must not carry a query or a fragment"), + ("http://sie-gateway:notaport", "has an invalid port"), + (f"http://user:{PASSWORD}@gw\uff0fx:8080", "is not a valid URL"), + (f"http://user:{PASSWORD}\uff20gw:8080", "is not a valid URL"), + ], +) +def test_a_gateway_url_that_is_not_a_plain_base_url_is_refused_without_echoing_it( + gateway, prometheus, read_token_file, monkeypatch, url, message, capsys, caplog +): + caplog.set_level(logging.DEBUG) + monkeypatch.setenv("SIE_ADMIN_TOKEN", ADMIN_TOKEN) + argv = ["tune", "--plan", str(FIXTURES / "tune-plan.json"), "--prometheus-url", prometheus.url] + argv += ["--gateway-url", url, "--api-key-file", str(read_token_file), "--apply", "--once"] + assert cli.main(argv) == cli.EXIT_REFUSED + captured = capsys.readouterr() + assert f"the gateway URL {message}" in captured.err + for text in (captured.out, captured.err, caplog.text): + assert url not in text + assert PASSWORD not in text + assert gateway.requests == [] + assert prometheus.requests == [] + + +def test_a_prometheus_url_that_is_not_a_valid_url_is_refused_without_echoing_it(gateway, prometheus, capsys, caplog): + caplog.set_level(logging.DEBUG) + url = f"http://user:{PASSWORD}@prometheus\uff0fx:9090" + argv = ["tune", "--plan", str(FIXTURES / "tune-plan.json"), "--prometheus-url", url] + assert cli.main([*argv, "--gateway-url", gateway.url, "--once"]) == cli.EXIT_REFUSED + captured = capsys.readouterr() + assert "the Prometheus URL is not a valid URL" in captured.err + for text in (captured.out, captured.err, caplog.text): + assert url not in text + assert PASSWORD not in text + assert gateway.requests == [] + assert prometheus.requests == [] + + +def test_a_gateway_url_httpx_cannot_parse_is_refused_without_echoing_it(): + with pytest.raises(ValueError, match=r"^the gateway URL is not a valid URL$"): + pools.connect(f"http://user:{PASSWORD}@sie-gateway:notaport", READ_TOKEN) + + +def test_a_token_an_http_header_cannot_encode_is_refused_without_echoing_it(gateway): + with pytest.raises(ValueError, match=r"^a gateway token holds a character that an HTTP header cannot carry$"): + pools.connect(gateway.url, "caf\u00e9-secret") + + +CARRIER = "placement-floor-default-l4-default" +NESTED = b"[" * 100_000 + + +def write_carrier(writer: pools.SdkPools) -> None: + writer.create_pool(CARRIER, {"l4": 0}, bundle="default", minimum_worker_count=1, queue_pool="default") + + +@pytest.mark.parametrize( + ("reply", "message"), + [ + ((200, NESTED), "the gateway returned an unexpected response"), + ((307, json.dumps({"spec": {"name": CARRIER}}).encode()), "the gateway answered HTTP 307"), + ((302, b"{}"), "the gateway answered HTTP 302"), + ], +) +def test_a_pool_reply_the_applier_cannot_use_fails_the_call(gateway, reply, message): + gateway.raw_reply = reply + reader = pools.connect(gateway.url, READ_TOKEN) + try: + with pytest.raises(keda.GatewayError) as raised: + reader.get_pool(CARRIER) + finally: + reader.close() + assert str(raised.value).startswith(message) + with pytest.raises(keda.GatewayError) as raised: + write_carrier(pools.connect(gateway.url, ADMIN_TOKEN)) + assert str(raised.value).startswith(message) + assert [request["method"] for request in gateway.requests] == ["GET", "POST"] + + +def test_a_deeply_nested_pool_reply_fails_only_the_tick(gateway, prometheus, read_token_file, capsys): + gateway.raw_reply = (200, NESTED) + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file)) == cli.EXIT_FAILED + errors = [row["error"] for row in rows(capsys.readouterr().out)] + assert errors == [ + f"reading pool {CARRIER!r} failed: the gateway returned an unexpected response", + "reading pool 'placement-floor-default-a100-80gb-generation' failed: " + "the gateway returned an unexpected response", + ] + + +def test_a_redirected_write_is_not_recorded_as_applied(gateway, prometheus, read_token_file, monkeypatch, capsys): + monkeypatch.setenv("SIE_ADMIN_TOKEN", ADMIN_TOKEN) + gateway.auth = False + original = gateway.respond + + def redirect_writes(handler: _Handler, method: str, body: object) -> None: + gateway.raw_reply = (307, json.dumps({"spec": body}).encode()) if method == "POST" else None + original(handler, method, body) + + monkeypatch.setattr(gateway, "respond", redirect_writes) + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file), "--apply") == cli.EXIT_FAILED + written = [row for row in rows(capsys.readouterr().out) if row["floor"]] + assert written + for row in written: + assert row["applied"] is False + assert "the gateway answered HTTP 307" in row["error"] + assert gateway.pools == {} + + +BACKSLASH_READ_TOKEN = "read\\token-for-tests" # noqa: S105 - a test value +BACKSLASH_ADMIN_TOKEN = "admin\\token-for-tests" # noqa: S105 - a test value +ECHOES = { + "403-message": lambda token: ( + 403, + json.dumps({"detail": {"code": "FORBIDDEN", "message": f"token not allowed: Bearer {token}"}}).encode(), + ), + "403-detail": lambda token: ( + 403, + json.dumps({"detail": {"code": "FORBIDDEN", "seen": f"Bearer {token}"}}).encode(), + ), + "502-text": lambda token: (502, f"upstream echoed Authorization: Bearer {token}".encode()), +} + + +@pytest.mark.parametrize("echo", sorted(ECHOES)) +@pytest.mark.parametrize("method", ["GET", "POST"]) +def test_a_reply_that_echoes_the_clients_token_is_reported_by_its_status_alone( + gateway, prometheus, tmp_path, monkeypatch, method, echo, capsys +): + read_token = tmp_path / "read-token" + read_token.write_text(BACKSLASH_READ_TOKEN + "\n", encoding="utf-8") + monkeypatch.setenv("SIE_ADMIN_TOKEN", BACKSLASH_ADMIN_TOKEN) + gateway.auth = False + status, body = ECHOES[echo](BACKSLASH_READ_TOKEN if method == "GET" else BACKSLASH_ADMIN_TOKEN) + gateway.raw_replies[method] = (status, body) + assert tune(gateway, prometheus, "--api-key-file", str(read_token), "--apply") == cli.EXIT_FAILED + captured = capsys.readouterr() + errors = [row["error"] for row in rows(captured.out) if row["error"]] + assert errors + assert all(error.endswith(f" failed: the gateway answered HTTP {status}") for error in errors) + for secret in (BACKSLASH_READ_TOKEN, BACKSLASH_ADMIN_TOKEN): + for form in (secret, secret.replace("\\", "\\\\")): + assert form not in captured.out + assert form not in captured.err + + +def test_an_error_the_sdk_raises_before_any_reply_is_redacted_and_cut(): + message = f"connection error after sending Bearer {BACKSLASH_ADMIN_TOKEN} " + "x" * 5000 + error = str(pools._gateway_error(SIEError(message), None, BACKSLASH_ADMIN_TOKEN)) + assert BACKSLASH_ADMIN_TOKEN not in error + assert error.startswith("connection error after sending Bearer [redacted] x") + assert len(error) == pools.MAX_ERROR_LENGTH + + +def test_a_malformed_pool_reply_fails_only_the_tick(gateway, prometheus, read_token_file, capsys): + gateway.read_reply = ["not", "a", "pool"] + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file)) == cli.EXIT_FAILED + errors = [row["error"] for row in rows(capsys.readouterr().out)] + assert errors == [ + "reading pool 'placement-floor-default-l4-default' failed: the gateway returned an unexpected response", + "reading pool 'placement-floor-default-a100-80gb-generation' failed: " + "the gateway returned an unexpected response", + ] + + +def test_the_loop_survives_failed_ticks_and_never_renews_a_lease_in_the_background( + gateway, prometheus, read_token_file, monkeypatch, capsys +): + monkeypatch.setenv("SIE_ADMIN_TOKEN", ADMIN_TOKEN) + leases: list[list[str]] = [] + closed: list[pools.SdkPools] = [] + close = pools.SdkPools.close + + def record_close(self: pools.SdkPools) -> None: + closed.append(self) + close(self) + + def between_ticks(_seconds: float) -> None: + leases.append([thread.name for thread in threading.enumerate() if thread.name.startswith("pool-lease-")]) + if len(leases) == 1: + prometheus.truncate = True + elif len(leases) == 2: + prometheus.truncate = False + gateway.fail_reads = True + else: + raise KeyboardInterrupt + + monkeypatch.setattr(pools.SdkPools, "close", record_close) + monkeypatch.setattr(cli, "time", SimpleNamespace(monotonic=time.monotonic, sleep=between_ticks)) + assert run(gateway, prometheus, "--api-key-file", str(read_token_file), "--apply", "--interval-s", "1") == 0 + captured = capsys.readouterr() + lines = rows(captured.out) + assert [(row["action"], row["applied"], row["error"] is None) for row in lines] == [ + ("raise_to_floor", True, True), + ("hold", False, True), + ("hold", True, True), + ("hold", False, True), + (None, False, False), + (None, False, False), + ] + assert "Prometheus series queueDepth: the query failed: IncompleteRead" in captured.err + assert leases == [[], [], []] + posts = [request["path"] for request in gateway.requests if request["method"] == "POST"] + assert posts == ["/v1/pools", "/v1/pools"] + assert len(closed) == 2 diff --git a/packages/sie_placement/tests/test_keda.py b/packages/sie_placement/tests/test_keda.py new file mode 100644 index 000000000..3fa82a226 --- /dev/null +++ b/packages/sie_placement/tests/test_keda.py @@ -0,0 +1,702 @@ +from __future__ import annotations + +import copy +import itertools +import json +import socketserver +import threading +from collections.abc import Iterator, Mapping, Sequence +from datetime import UTC, datetime +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from urllib.parse import parse_qs, urlsplit + +import pytest +from sie_placement import keda +from sie_placement.schema import Plan +from sie_placement.tuner import Action + +FIXTURES = Path(__file__).parent / "fixtures" +AT = datetime(2026, 1, 1, 12, 0, tzinfo=UTC) + +QUEUE_METRIC = "sie_gateway_lane_queue_depth" +PENDING_METRIC = "sie_gateway_pending_demand" +FLOOR_METRIC = "sie_gateway_pool_warm_floor" +SLOTS_METRIC = "sie_worker_gpu_slots" + + +def plan_document() -> dict: + return json.loads((FIXTURES / "tune-plan.json").read_text(encoding="utf-8")) + + +def make_plan(*lanes: dict) -> Plan: + document = plan_document() + template = document["lanes"][0] + document["lanes"] = [] + document["identities"] = [] + for index, overrides in enumerate(lanes): + lane = copy.deepcopy(template) + lane["name"] = f"lane-{index}" + envelope = overrides.pop("envelope", None) + machine_profile = overrides.pop("machineProfile", None) + lane.update(overrides) + lane["envelope"] = envelope or {"floorMin": 0, "floorMax": 3} + if machine_profile is not None: + lane["cell"]["machineProfile"] = machine_profile + document["lanes"].append(lane) + warmth = "lazy-resident" if lane["envelope"]["floorMin"] >= 1 else "cold-with-bound" + document["identities"].append({"identity": f"model-{index}", "lane": lane["name"], "warmth": warmth}) + return Plan.from_dict(document) + + +def one_lane(floor_min: int = 0, floor_max: int = 3) -> tuple[keda.Carrier, ...]: + return keda.carriers(make_plan({"envelope": {"floorMin": floor_min, "floorMax": floor_max}})) + + +def lane_labels(queue_pool: str = "default", machine_profile: str = "l4", bundle: str = "default") -> dict[str, str]: + return {"pool": queue_pool, "machine_profile": machine_profile, "bundle": bundle} + + +class FakePrometheus: + def __init__(self, series: Mapping[str, Sequence[tuple[dict[str, str], float]]] | None = None) -> None: + self.series = dict(series or {}) + self.failing: set[str] = set() + self.queries: list[str] = [] + + @classmethod + def lane( + cls, + *, + queue_depth: int = 0, + pending: int = 0, + warm_floor: int = 0, + runners: int = 0, + labels: dict | None = None, + ) -> FakePrometheus: + labels = labels or lane_labels() + otel = f"{labels['pool']}|{labels['machine_profile']}|{labels['bundle']}" + series: dict[str, list[tuple[dict[str, str], float]]] = { + QUEUE_METRIC: [(labels, queue_depth)], + PENDING_METRIC: [(labels, pending)], + FLOOR_METRIC: [(labels, warm_floor)], + SLOTS_METRIC: [({"lane": otel}, runners)] if runners else [], + } + return cls(series) + + def query(self, promql: str) -> Sequence[keda.Sample]: + self.queries.append(promql) + for metric, samples in self.series.items(): + if metric in promql: + if metric in self.failing: + raise keda.PrometheusError("the query failed with HTTP 503") + return tuple(keda.Sample(labels=labels, value=value) for labels, value in samples) + for metric in self.failing: + if metric in promql: + raise keda.PrometheusError("the query failed with HTTP 503") + return () + + +class FakePools: + def __init__(self, pools: Mapping[str, dict] | None = None) -> None: + self.pools: dict[str, dict] = {name: copy.deepcopy(spec) for name, spec in (pools or {}).items()} + self.raw: dict[str, dict] = {} + self.reads: list[str] = [] + self.writes: list[tuple[tuple, dict]] = [] + self.fail_reads = False + self.fail_writes = False + + def get_pool(self, name: str) -> Mapping[str, object] | None: + self.reads.append(name) + if self.fail_reads: + raise keda.GatewayError("connection refused") + if name in self.raw: + return self.raw[name] + spec = self.pools.get(name) + return None if spec is None else {"name": name, "spec": copy.deepcopy(spec), "status": {"state": "active"}} + + def create_pool( + self, name: str, gpus: dict[str, int], *, bundle: str, minimum_worker_count: int, queue_pool: str + ) -> None: + self.writes.append( + ((name, gpus), {"bundle": bundle, "minimum_worker_count": minimum_worker_count, "queue_pool": queue_pool}) + ) + if self.fail_writes: + raise keda.GatewayError("HTTP 403") + self.pools[name] = { + "name": name, + "queue_pool": queue_pool, + "bundle": bundle, + "gpus": dict(gpus), + "ttl_seconds": None, + "minimum_worker_count": minimum_worker_count, + } + + +def carrier_spec(floor: int, *, machine_profile: str = "l4", bundle: str = "default", **extra: object) -> dict: + return { + "name": f"placement-floor-default-{machine_profile}-{bundle}", + "queue_pool": "default", + "bundle": bundle, + "gpus": {machine_profile: 0}, + "ttl_seconds": None, + "minimum_worker_count": floor, + **extra, + } + + +def run_tick( + lanes: Sequence[keda.Carrier], + prometheus: FakePrometheus, + reader: FakePools, + writer: FakePools | None = None, + selector: str = "", +) -> keda.Tick: + return keda.tick(lanes, prometheus=prometheus, reader=reader, writer=writer, selector=selector, now=lambda: AT) + + +def test_each_fixture_lane_gets_one_floor_carrier(): + lanes = keda.carriers(Plan.from_dict(plan_document())) + assert [carrier.name for carrier in lanes] == [ + "placement-floor-default-l4-default", + "placement-floor-default-a100-80gb-generation", + ] + assert [carrier.gpus for carrier in lanes] == [{"l4": 0}, {"a100-80gb": 0}] + assert lanes[0].key.otel_lane == "default|l4|default" + + +def test_lane_tokens_are_canonicalized_the_way_the_chart_names_lanes(): + lanes = keda.carriers(make_plan({"queuePool": "Default", "machineProfile": "L4-Spot", "bundle": "Default"})) + assert lanes[0].key == keda.LaneKey(queue_pool="default", machine_profile="l4-spot", bundle="default") + assert lanes[0].name == "placement-floor-default-l4-spot-default" + + +def test_an_envelope_above_the_gateway_per_pool_floor_limit_is_refused(): + plan = make_plan({"envelope": {"floorMin": 0, "floorMax": 5}, "maxContainers": 6}) + with pytest.raises(keda.PlanRefusedError) as raised: + keda.carriers(plan) + assert raised.value.problems == ("$.lanes[0].envelope.floorMax: 5 exceeds the gateway's per-pool floor limit of 4",) + assert keda.carriers(plan, max_pool_floor=5)[0].lane.envelope.floor_max == 5 + + +def test_more_lanes_than_the_gateway_pool_limit_are_refused(): + plan = make_plan({"bundle": "a"}, {"bundle": "b"}, {"bundle": "c"}) + with pytest.raises(keda.PlanRefusedError, match="3 lanes need as many floor-carrier pools"): + keda.carriers(plan, max_pools=2) + assert len(keda.carriers(plan, max_pools=3)) == 3 + + +def test_two_lanes_that_need_one_carrier_are_refused(): + plan = make_plan({"bundle": "default"}, {"bundle": "DEFAULT"}) + with pytest.raises(keda.PlanRefusedError, match=r"\$\.lanes\[1\]: needs the same floor-carrier pool"): + keda.carriers(plan) + hyphens = make_plan({"queuePool": "a-b", "machineProfile": "c"}, {"queuePool": "a", "machineProfile": "b-c"}) + with pytest.raises(keda.PlanRefusedError, match="placement-floor-a-b-c-default"): + keda.carriers(hyphens) + + +@pytest.mark.parametrize("token", ["l4.x", "l4/x", "x" * 64, 'l4"']) +def test_a_lane_token_the_gateway_would_reject_is_refused(token): + with pytest.raises(keda.PlanRefusedError, match="must each be 1 to 63 letters"): + keda.carriers(make_plan({"machineProfile": token})) + + +def test_a_carrier_name_longer_than_the_gateway_accepts_is_refused(): + plan = make_plan({"queuePool": "q" * 40, "machineProfile": "m" * 40, "bundle": "b" * 40}) + with pytest.raises(keda.PlanRefusedError, match="longer than 128 characters"): + keda.carriers(plan) + + +def test_a_pin_on_the_default_queue_pool_is_refused_as_the_helm_emitter_refuses_it(): + document = plan_document() + document["lanes"][0]["queuePool"] = "Default" + document["identities"][1]["warmth"] = "pinned" + document["identities"][0]["warmth"] = "pinned" + with pytest.raises(keda.PlanRefusedError) as raised: + keda.carriers(Plan.from_dict(document)) + assert raised.value.problems == ( + "$.identities[0].warmth: pins 'BAAI/bge-m3' on chart lane 'default/l4/default', but the gateway keeps no " + "static queue pool for the 'default' queue pool, so its pinnedModels have no carrier; place pinned " + "identities on a named queuePool", + ) + document["lanes"][0]["queuePool"] = "embed" + (pinned, _) = keda.carriers(Plan.from_dict(document)) + assert pinned.name == "placement-floor-embed-l4-default" + + +@pytest.mark.parametrize("warmth", sorted(keda.UNREALIZED_WARMTHS)) +def test_a_warmth_the_chart_does_not_realize_is_refused_as_the_helm_emitter_refuses_it(warmth): + document = plan_document() + document["identities"][2]["warmth"] = warmth.value + with pytest.raises(keda.PlanRefusedError) as raised: + keda.carriers(Plan.from_dict(document)) + assert raised.value.problems == ( + f"$.identities[2].warmth: is {warmth.value!r} for 'Qwen/Qwen3-8B' on lane 'default-a100-80gb-generation', " + f"which the sie-cluster chart does not realize: {keda.UNREALIZED_WARMTHS[warmth]}; plan the identity " + "pinned, lazy-resident or cold-with-bound", + ) + + +def test_every_problem_of_a_refused_plan_is_reported(): + document = plan_document() + document["identities"][0]["warmth"] = "pinned" + document["identities"][2]["warmth"] = "snapshot-cold" + document["lanes"][1]["envelope"]["floorMax"] = 5 + document["lanes"][1]["maxContainers"] = 6 + with pytest.raises(keda.PlanRefusedError) as raised: + keda.carriers(Plan.from_dict(document)) + assert [problem.split(":", 1)[0] for problem in raised.value.problems] == [ + "$.identities[0].warmth", + "$.lanes[1].envelope.floorMax", + "$.identities[2].warmth", + ] + + +def test_the_series_queries_read_each_lane_from_fresh_gateway_producers_and_the_worker_series(): + capacity_fresh = ( + "and on(producer_instance, collector_generation) " + "(abs(time() - sie_gateway_capacity_snapshot_timestamp_seconds{}) < 20))" + ) + assert keda.series_queries() == { + "queueDepth": ( + "max by (pool, machine_profile, bundle) (sie_gateway_lane_queue_depth{} " + "and on(producer_instance, collector_generation, pool, machine_profile, bundle) " + "(abs(time() - sie_gateway_lane_queue_snapshot_timestamp_seconds{}) < 20))" + ), + "pendingDemand": f"max by (pool, machine_profile, bundle) (sie_gateway_pending_demand{{}} {capacity_fresh}", + "warmFloor": f"max by (pool, machine_profile, bundle) (sie_gateway_pool_warm_floor{{}} {capacity_fresh}", + "runners": ('count by (lane) (max by (lane, producer_instance) (sie_worker_gpu_slots{state="ready"}) > 0)'), + } + scoped = keda.series_queries('namespace="sie"') + assert scoped["queueDepth"] == ( + 'max by (pool, machine_profile, bundle) (sie_gateway_lane_queue_depth{namespace="sie"} ' + "and on(producer_instance, collector_generation, pool, machine_profile, bundle) " + '(abs(time() - sie_gateway_lane_queue_snapshot_timestamp_seconds{namespace="sie"}) < 20))' + ) + for name in ("pendingDemand", "warmFloor"): + assert scoped[name].endswith( + '{namespace="sie"} and on(producer_instance, collector_generation) ' + '(abs(time() - sie_gateway_capacity_snapshot_timestamp_seconds{namespace="sie"}) < 20))' + ) + assert scoped["runners"].endswith('(sie_worker_gpu_slots{namespace="sie",state="ready"}) > 0)') + + +def test_a_lane_whose_gateway_producers_are_all_stale_is_held(): + stale = lane_labels(bundle="stale") + prometheus = FakePrometheus.lane(queue_depth=4, pending=1, runners=2) + prometheus.series[QUEUE_METRIC] = [(lane_labels(), 4)] + prometheus.series[PENDING_METRIC] = [(lane_labels(), 1)] + lanes = keda.carriers(make_plan({"bundle": "default"}, {"bundle": "stale"})) + pools = FakePools({carrier.name: carrier_spec(1, bundle=carrier.key.bundle) for carrier in lanes}) + fresh, held = run_tick(lanes, prometheus, pools, pools).records + assert held.carrier.key.bundle == stale["bundle"] + assert (fresh.step.action, fresh.step.settings.floor) == (Action.SCALE_UP, 2) + assert (held.telemetry.queue_depth, held.telemetry.pending_demand) == (None, None) + assert (held.step.action, held.step.settings.floor) == (Action.HOLD, 1) + + +@pytest.mark.parametrize("selector", ['namespace="sie"}', '{namespace="sie"', 'a="b"\nc="d"']) +def test_a_selector_that_could_leave_its_braces_is_rejected(selector): + with pytest.raises(ValueError, match="without braces or line breaks"): + keda.series_queries(selector) + + +def test_parse_vector_reads_an_instant_vector(): + document = { + "status": "success", + "data": {"resultType": "vector", "result": [{"metric": {"pool": "default"}, "value": [1.0, "3"]}]}, + } + assert keda.parse_vector(document) == (keda.Sample(labels={"pool": "default"}, value=3.0),) + + +@pytest.mark.parametrize( + "document", + [ + {"status": "error", "error": "bad query"}, + {"status": "success", "data": {"resultType": "matrix", "result": []}}, + {"status": "success", "data": {"resultType": "vector", "result": {}}}, + {"status": "success", "data": {"resultType": "vector", "result": [{"metric": {}}]}}, + {"status": "success", "data": {"resultType": "vector", "result": [{"metric": {}, "value": [1, "x"]}]}}, + [], + ], +) +def test_parse_vector_rejects_anything_else(document): + with pytest.raises(keda.PrometheusError): + keda.parse_vector(document) + + +class _PrometheusHandler(BaseHTTPRequestHandler): + server: _PrometheusServer + + def do_GET(self) -> None: + url = urlsplit(self.path) + self.server.requests.append((url.path, parse_qs(url.query))) + status, body = self.server.reply + payload = json.dumps(body).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format: str, *args: object) -> None: + return + + +class _PrometheusServer(ThreadingHTTPServer): + def __init__(self) -> None: + super().__init__(("127.0.0.1", 0), _PrometheusHandler) + self.requests: list[tuple[str, dict]] = [] + self.reply: tuple[int, object] = (200, {"status": "success", "data": {"resultType": "vector", "result": []}}) + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.server_address[1]}" + + +@pytest.fixture +def prometheus_server(monkeypatch) -> Iterator[_PrometheusServer]: + for name in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"): + monkeypatch.delenv(name, raising=False) + server = _PrometheusServer() + thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + yield server + server.shutdown() + server.server_close() + + +def test_the_http_reader_sends_the_query_to_the_query_api(prometheus_server): + prometheus_server.reply = ( + 200, + {"status": "success", "data": {"resultType": "vector", "result": [{"metric": {}, "value": [1, "2"]}]}}, + ) + reader = keda.PrometheusHttp(prometheus_server.url + "/") + assert reader.query('max(up{job="x"})') == (keda.Sample(labels={}, value=2.0),) + assert prometheus_server.requests == [("/api/v1/query", {"query": ['max(up{job="x"})']})] + + +def test_the_http_reader_reports_a_failed_query(prometheus_server): + prometheus_server.reply = (400, {"status": "error", "error": "parse error"}) + with pytest.raises(keda.PrometheusError, match="HTTP 400"): + keda.PrometheusHttp(prometheus_server.url).query("up{") + + +def test_the_http_reader_reports_an_unreachable_prometheus(): + with pytest.raises(keda.PrometheusError, match="the query failed"): + keda.PrometheusHttp("http://127.0.0.1:9", timeout_s=2).query("up") + + +class _RawHandler(socketserver.StreamRequestHandler): + server: _RawServer + + def handle(self) -> None: + self.rfile.readline() + self.wfile.write(self.server.reply) + + +class _RawServer(socketserver.ThreadingTCPServer): + daemon_threads = True + + def __init__(self, reply: bytes) -> None: + super().__init__(("127.0.0.1", 0), _RawHandler) + self.reply = reply + + +@pytest.mark.parametrize( + ("reply", "error"), + [ + (b"garbage\r\n\r\n", "BadStatusLine"), + ( + b'HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 500\r\n\r\n{"status": "s', + "IncompleteRead", + ), + ], +) +def test_the_http_reader_reports_a_malformed_reply(monkeypatch, reply, error): + for name in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"): + monkeypatch.delenv(name, raising=False) + server = _RawServer(reply) + thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + try: + with pytest.raises(keda.PrometheusError, match=f"the query failed: {error}"): + keda.PrometheusHttp(f"http://127.0.0.1:{server.server_address[1]}", timeout_s=2).query("up") + finally: + server.shutdown() + server.server_close() + + +@pytest.mark.parametrize( + "url", + [ + "file:///etc/passwd", + "ftp://prometheus", + "prometheus:9090", + "http://user:secret@prometheus", + "http://p/?x=1", + "http://prometheus:notaport", + "http://prometheus:99999", + ], +) +def test_the_http_reader_accepts_only_a_plain_http_base_url(url): + with pytest.raises(ValueError, match="Prometheus URL"): + keda.PrometheusHttp(url) + + +def test_shadow_mode_decides_every_lane_and_writes_nothing(): + lanes = keda.carriers(Plan.from_dict(plan_document())) + prometheus = FakePrometheus.lane(queue_depth=4, pending=1, warm_floor=1, runners=1) + reader = FakePools({"placement-floor-default-l4-default": carrier_spec(1)}) + result = run_tick(lanes, prometheus, reader) + assert result.apply is False + assert reader.writes == [] + assert reader.reads == [carrier.name for carrier in lanes] + first, second = result.records + assert (first.current_floor, first.step.action, first.step.settings.floor, first.applied) == ( + 1, + Action.SCALE_UP, + 2, + False, + ) + assert (second.current_floor, second.step.action, second.step.settings.floor) == (0, Action.HOLD, 0) + assert second.telemetry == keda.LaneTelemetry(queue_depth=None, pending_demand=None, warm_floor=None, runners=0) + + +def test_apply_upserts_one_unpinned_carrier_per_lane_through_the_pool_api(): + lanes = one_lane(1, 3) + pools = FakePools() + result = run_tick(lanes, FakePrometheus.lane(), pools, pools) + assert pools.writes == [ + ( + ("placement-floor-default-l4-default", {"l4": 0}), + {"bundle": "default", "minimum_worker_count": 1, "queue_pool": "default"}, + ) + ] + assert result.records[0].step.action is Action.RAISE_TO_FLOOR + assert result.records[0].applied is True + + +def test_apply_renews_the_carrier_lease_on_every_tick(): + lanes = one_lane(1, 3) + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(2)}) + busy = FakePrometheus.lane(queue_depth=3, runners=2) + for _ in range(3): + result = run_tick(lanes, busy, pools, pools) + assert result.records[0].step.action is Action.HOLD + assert [kwargs["minimum_worker_count"] for _, kwargs in pools.writes] == [2, 2, 2] + + +def test_a_lane_without_a_carrier_that_stays_at_zero_gets_no_carrier(): + pools = FakePools() + result = run_tick(one_lane(0, 3), FakePrometheus.lane(), pools, pools) + assert pools.writes == [] + assert result.records[0].step.action is Action.HOLD + assert result.records[0].applied is False + + +def test_an_existing_carrier_at_zero_is_still_renewed(): + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(0)}) + run_tick(one_lane(0, 3), FakePrometheus.lane(), pools, pools) + assert [kwargs["minimum_worker_count"] for _, kwargs in pools.writes] == [0] + + +@pytest.mark.parametrize(("floor_min", "floor_max"), [(0, 0), (0, 3), (1, 1), (1, 4), (2, 3)]) +def test_the_written_floor_never_leaves_the_envelope(floor_min, floor_max): + lanes = one_lane(floor_min, floor_max) + for current, queue_depth, pending, runners in itertools.product(range(7), (0, 3), (0, 1), (0, 1, 3, 6)): + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(current)}) + prometheus = FakePrometheus.lane(queue_depth=queue_depth, pending=pending, runners=runners) + record = run_tick(lanes, prometheus, pools, pools).records[0] + (_, kwargs) = pools.writes[-1] + assert floor_min <= kwargs["minimum_worker_count"] <= floor_max + assert kwargs["minimum_worker_count"] == record.step.settings.floor + assert record.current_floor == current + + +def test_pending_demand_raises_the_floor_and_an_idle_lane_lowers_it(): + lanes = one_lane(0, 3) + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(1)}) + waiting = FakePrometheus.lane(queue_depth=2, pending=1, runners=1) + serving = FakePrometheus.lane(queue_depth=2, pending=0, runners=1) + idle = FakePrometheus.lane(queue_depth=0, pending=0, runners=1) + assert run_tick(lanes, waiting, pools, pools).records[0].step.action is Action.SCALE_UP + assert pools.pools["placement-floor-default-l4-default"]["minimum_worker_count"] == 2 + assert run_tick(lanes, serving, pools, pools).records[0].step.action is Action.HOLD + assert run_tick(lanes, idle, pools, pools).records[0].step.action is Action.SCALE_DOWN + assert pools.pools["placement-floor-default-l4-default"]["minimum_worker_count"] == 1 + + +def test_a_floor_whose_workers_have_not_started_is_not_raised_again(): + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(2)}) + booting = FakePrometheus.lane(queue_depth=5, pending=1, runners=1) + record = run_tick(one_lane(0, 3), booting, pools, pools).records[0] + assert (record.step.action, record.step.settings.floor) == (Action.HOLD, 2) + + +def test_current_is_read_from_the_gateway_on_every_tick(): + lanes = one_lane(0, 3) + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(1)}) + waiting = FakePrometheus.lane(queue_depth=2, pending=1, runners=3) + assert run_tick(lanes, waiting, pools, pools).records[0].step.settings.floor == 2 + pools.pools["placement-floor-default-l4-default"]["minimum_worker_count"] = 0 + record = run_tick(lanes, waiting, pools, pools).records[0] + assert (record.current_floor, record.step.settings.floor) == (0, 1) + + +def test_a_floor_held_above_the_envelope_by_another_pool_is_reported_and_left_alone(): + lanes = one_lane(1, 3) + other = { + "name": "team-pool", + "queue_pool": "default", + "bundle": "default", + "gpus": {"l4": 1}, + "minimum_worker_count": 4, + } + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(1), "team-pool": other}) + record = run_tick(lanes, FakePrometheus.lane(warm_floor=4, runners=4), pools, pools).records[0] + assert record.lane_floor_above_envelope is True + assert record.to_dict()["laneFloorAboveEnvelope"] is True + assert [args[0] for args, _ in pools.writes] == ["placement-floor-default-l4-default"] + assert pools.pools["team-pool"] == other + assert ( + run_tick(lanes, FakePrometheus.lane(warm_floor=3), pools, pools).records[0].lane_floor_above_envelope is False + ) + + +def test_a_carrier_above_a_lowered_envelope_is_brought_down_to_the_cap(): + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(4)}) + record = run_tick(one_lane(0, 2), FakePrometheus.lane(warm_floor=4, runners=4), pools, pools).records[0] + assert (record.step.action, record.step.settings.floor) == (Action.HOLD_AT_CAP, 2) + assert pools.pools["placement-floor-default-l4-default"]["minimum_worker_count"] == 2 + + +@pytest.mark.parametrize( + ("change", "reason"), + [ + ({"pinned_models": ["BAAI/bge-m3"]}, "it sets pinned_models"), + ({"gpu_caps": {"l4": 2}}, "it sets gpu_caps"), + ({"ttl_seconds": 3600}, "it sets ttl_seconds"), + ({"gpus": {"l4": 1}}, "its gpus are not"), + ({"gpus": {"l4": 0, "a100": 0}}, "its gpus are not"), + ({"bundle": "other"}, "bundle 'default'"), + ({"bundle": None}, "bundle 'default'"), + ({"queue_pool": "dedicated"}, "queue pool 'default'"), + ({"minimum_worker_count": -1}, "minimum_worker_count"), + ({"minimum_worker_count": True}, "minimum_worker_count"), + ], +) +def test_a_pool_that_is_not_the_lane_carrier_is_never_written(change, reason): + spec = carrier_spec(2) | change + pools = FakePools({"placement-floor-default-l4-default": spec}) + record = run_tick(one_lane(0, 3), FakePrometheus.lane(pending=1, runners=2), pools, pools).records[0] + assert pools.writes == [] + assert record.step is None + assert reason in (record.error or "") + assert "is not this lane's floor carrier" in (record.error or "") + + +def test_a_pool_without_a_spec_is_never_written(): + pools = FakePools() + pools.raw["placement-floor-default-l4-default"] = {"name": "placement-floor-default-l4-default"} + record = run_tick(one_lane(0, 3), FakePrometheus.lane(), pools, pools).records[0] + assert record.error is not None + assert pools.writes == [] + + +def test_missing_telemetry_holds_the_lane_and_only_brings_it_inside_the_envelope(): + lanes = one_lane(1, 3) + prometheus = FakePrometheus.lane(pending=1, runners=3) + prometheus.failing.add(QUEUE_METRIC) + for current, expected, action in ((0, 1, Action.RAISE_TO_FLOOR), (2, 2, Action.HOLD), (5, 3, Action.HOLD_AT_CAP)): + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(current)}) + result = run_tick(lanes, prometheus, pools, pools) + record = result.records[0] + assert (record.step.action, record.step.settings.floor) == (action, expected) + assert record.telemetry.queue_depth is None + assert result.series_errors == ("queueDepth: the query failed with HTTP 503",) + assert result.failed is True + + +def test_a_lane_the_gateway_does_not_report_is_held(): + prometheus = FakePrometheus.lane(pending=1, runners=1, labels=lane_labels(bundle="other")) + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(1)}) + result = run_tick(one_lane(0, 3), prometheus, pools, pools) + record = result.records[0] + assert (record.step.action, record.step.settings.floor) == (Action.HOLD, 1) + assert result.failed is False + + +def test_a_lane_with_no_ready_worker_reports_zero_runners(): + prometheus = FakePrometheus.lane(pending=1) + record = run_tick(one_lane(0, 3), prometheus, FakePools(), FakePools()).records[0] + assert record.telemetry.runners == 0 + assert (record.step.action, record.step.settings.floor) == (Action.SCALE_UP, 1) + + +def test_samples_are_matched_to_lanes_by_canonical_labels(): + labels = {"pool": "Default", "machine_profile": "L4", "bundle": "Default"} + prometheus = FakePrometheus( + { + QUEUE_METRIC: [(labels, 2), (lane_labels(), 5), ({"pool": "default"}, 9)], + PENDING_METRIC: [(labels, 0), (lane_labels(), float("nan"))], + FLOOR_METRIC: [(labels, -1)], + SLOTS_METRIC: [({"lane": "Default|L4|Default"}, 2), ({"lane": "default|l4"}, 7)], + } + ) + record = run_tick(one_lane(0, 3), prometheus, FakePools()).records[0] + assert record.telemetry == keda.LaneTelemetry(queue_depth=5, pending_demand=0, warm_floor=None, runners=2) + + +def test_a_gateway_read_failure_skips_only_that_lane(): + lanes = keda.carriers(make_plan({"bundle": "a"}, {"bundle": "b"})) + reader = FakePools() + reader.fail_reads = True + writer = FakePools() + result = run_tick(lanes, FakePrometheus.lane(), reader, writer) + assert writer.writes == [] + assert [record.error for record in result.records] == [ + "reading pool 'placement-floor-default-l4-a' failed: connection refused", + "reading pool 'placement-floor-default-l4-b' failed: connection refused", + ] + assert result.failed is True + + +def test_a_gateway_write_failure_is_recorded(): + pools = FakePools() + pools.fail_writes = True + record = run_tick(one_lane(1, 3), FakePrometheus.lane(), pools, pools).records[0] + assert record.applied is False + assert record.step.settings.floor == 1 + assert record.error == "writing pool 'placement-floor-default-l4-default' failed: HTTP 403" + + +def test_audit_rows_carry_the_tick_time_mode_and_decision(): + pools = FakePools({"placement-floor-default-l4-default": carrier_spec(1)}) + prometheus = FakePrometheus.lane(queue_depth=4, pending=1, warm_floor=1, runners=1) + rows = run_tick(one_lane(1, 3), prometheus, pools, pools).audit_rows() + assert rows == [ + { + "at": "2026-01-01T12:00:00Z", + "mode": "apply", + "lane": "lane-0", + "queuePool": "default", + "machineProfile": "l4", + "bundle": "default", + "carrier": "placement-floor-default-l4-default", + "telemetry": {"queueDepth": 4, "pendingDemand": 1, "warmFloor": 1, "runners": 1}, + "currentFloor": 1, + "action": "scale_up", + "floor": 2, + "applied": True, + "laneFloorAboveEnvelope": False, + "error": None, + } + ] + assert json.loads(json.dumps(rows)) == rows + + +def test_the_selector_scopes_every_series_query(): + prometheus = FakePrometheus.lane() + run_tick(one_lane(), prometheus, FakePools(), selector='namespace="sie"') + assert len(prometheus.queries) == 4 + assert all('namespace="sie"' in query for query in prometheus.queries) diff --git a/uv.lock b/uv.lock index d07c8964a..0fc3486a8 100644 --- a/uv.lock +++ b/uv.lock @@ -4766,12 +4766,22 @@ name = "sie-placement" version = "0.0.0" source = { editable = "packages/sie_placement" } +[package.optional-dependencies] +keda = [ + { name = "httpx" }, + { name = "sie-sdk" }, +] + [package.dev-dependencies] dev = [ { name = "pytest" }, ] [package.metadata] +requires-dist = [ + { name = "httpx", marker = "extra == 'keda'", specifier = ">=0.28,<1" }, + { name = "sie-sdk", marker = "extra == 'keda'", editable = "packages/sie_sdk" }, +] [package.metadata.requires-dev] dev = [{ name = "pytest", specifier = ">=9.0.3" }]