From a05b5ba49bbd9e1f7cb8e8208fcfc9d8d97fdb76 Mon Sep 17 00:00:00 2001 From: krisztian-gajdar Date: Thu, 5 Nov 2026 03:15:54 +0100 Subject: [PATCH 01/10] feat(placement): add the sie-placement tune KEDA applier Add a `sie-placement tune` command that keeps each plan lane's warm floor inside its envelope on a KEDA-scaled sie-cluster deployment: - every tick it reads four lane series from Prometheus (lane queue depth, pending demand, warm floor and ready workers), reads the lane's floor-carrier pool with SIEClient.get_pool, and calls the tuner step with that pool's floor as current - shadow mode is the default and writes nothing; with --apply it upserts one floor-carrier logical pool per lane through SIEClient.create_pool on every tick, which also renews the pool's lease - a carrier is named placement-floor--- and has gpus {: 0}, the lane's queue pool and bundle, no pins, no caps and the gateway's default lease; a pool with that name and any other spec is reported and never written - the applier writes only its carriers, so it never lowers another pool's floor, and each audit row reports a lane whose warm floor exceeds floorMax - a lane missing a series is held and only brought inside its envelope - the plan is refused before any call when an envelope's floorMax is above the gateway's per-pool floor limit, when it has more lanes than the gateway's API pool limit, or when its lane tokens cannot name a carrier Pool reads use a gateway client token and pool writes the gateway admin token, each read from a file or an environment variable. The admin token is read only with --apply, is sent only on pool writes, and is never printed or logged. The README documents the permission each call needs. The SIE SDK and httpx come from a new keda extra; the schema, tuner and applier core keep to the standard library. Tests drive the applier with a fake Prometheus and fake pools, and drive the command through the real SDK against fake HTTP Prometheus and gateway servers. --- packages/sie_placement/README.md | 111 +++- packages/sie_placement/pyproject.toml | 9 + .../sie_placement/src/sie_placement/cli.py | 168 +++++ .../sie_placement/src/sie_placement/keda.py | 435 +++++++++++++ .../sie_placement/src/sie_placement/pools.py | 41 ++ packages/sie_placement/tests/test_cli.py | 370 +++++++++++ packages/sie_placement/tests/test_keda.py | 577 ++++++++++++++++++ uv.lock | 10 + 8 files changed, 1719 insertions(+), 2 deletions(-) create mode 100644 packages/sie_placement/src/sie_placement/cli.py create mode 100644 packages/sie_placement/src/sie_placement/keda.py create mode 100644 packages/sie_placement/src/sie_placement/pools.py create mode 100644 packages/sie_placement/tests/test_cli.py create mode 100644 packages/sie_placement/tests/test_keda.py diff --git a/packages/sie_placement/README.md b/packages/sie_placement/README.md index be45f5483..873dd8c19 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,112 @@ 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. 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. + +When the applier stops renewing, each carrier's lease lapses, 1200 seconds 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{})` | `running` | +| Pending demand | `max by (pool, machine_profile, bundle) (sie_gateway_pending_demand{})` | `backlog` | +| Warm floor | `max by (pool, machine_profile, bundle) (sie_gateway_pool_warm_floor{})` | 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, 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 attainment is `unmeasured`, so a lane is under pressure while +the gateway reports pending demand for it, which it does when a request found no +ready capacity, and idle when it reports neither pending demand nor queued or +in-flight work. 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, `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. + +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` and its lease renewal | 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. + ## 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..8f653dc5b --- /dev/null +++ b/packages/sie_placement/src/sie_placement/cli.py @@ -0,0 +1,168 @@ +"""Command-line entry point of the placement tools.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +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 +DEFAULT_INTERVAL_S = 60.0 +EXIT_FAILED = 1 +EXIT_REFUSED = 2 + + +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 _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 as error: + raise _UsageError(f"cannot read the {label} file {str(path)!r}: {error.strerror}") from None + if not value: + raise _UsageError(f"the {label} file {str(path)!r} is empty") + return value + if env_name is None: + return None + return os.environ.get(env_name, "").strip() or 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) + read_token = _secret(args.api_key_file, None, "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..3c297b70e --- /dev/null +++ b/packages/sie_placement/src/sie_placement/keda.py @@ -0,0 +1,435 @@ +"""Keep each plan lane's warm floor inside its envelope through one gateway floor-carrier pool per lane.""" + +from __future__ import annotations + +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 +from sie_placement.tuner import Observation, Settings, Step, step + +CARRIER_PREFIX = "placement-floor" +DEFAULT_MAX_POOL_FLOOR = 4 +DEFAULT_MAX_POOLS = 64 +MAX_POOL_NAME_LENGTH = 128 +PROMETHEUS_TIMEOUT_S = 10.0 + +QUEUE_DEPTH = "queueDepth" +PENDING_DEMAND = "pendingDemand" +WARM_FLOOR = "warmFloor" +RUNNERS = "runners" + +_LANE_TOKEN = re.compile(r"[a-z0-9_-]{1,63}") +_SELECTOR_FORBIDDEN = re.compile(r"[{}\r\n]") +_LANE_LABELS = "pool, machine_profile, bundle" +_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}" + + +@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] = {} + 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 + 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) + 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) + + +class PrometheusHttp: + def __init__(self, base_url: str, *, timeout_s: float = PROMETHEUS_TIMEOUT_S) -> None: + parts = urllib.parse.urlsplit(base_url) + if parts.scheme not in {"http", "https"} or not parts.hostname: + raise ValueError("the Prometheus URL must be an absolute http or https URL") + if parts.username is not None or parts.password is not None: + raise ValueError("the Prometheus URL must not embed credentials") + if parts.query or parts.fragment: + raise ValueError("the Prometheus URL must not carry a query or a fragment") + 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 + 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) + + return { + QUEUE_DEPTH: f"max by ({_LANE_LABELS}) (sie_gateway_lane_queue_depth{{{matchers()}}})", + PENDING_DEMAND: f"max by ({_LANE_LABELS}) (sie_gateway_pending_demand{{{matchers()}}})", + WARM_FLOOR: f"max by ({_LANE_LABELS}) (sie_gateway_pool_warm_floor{{{matchers()}}})", + 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..860d6f8c1 --- /dev/null +++ b/packages/sie_placement/src/sie_placement/pools.py @@ -0,0 +1,41 @@ +"""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 + + +class SdkPools: + def __init__(self, client: SIEClient) -> None: + self._client = client + + def get_pool(self, name: str) -> Mapping[str, object] | None: + try: + pool = self._client.get_pool(name) + except (SIEError, httpx.HTTPError, ValueError) as error: + raise GatewayError(str(error)) from None + return pool + + def create_pool( + self, name: str, gpus: dict[str, int], *, bundle: str, minimum_worker_count: int, queue_pool: str + ) -> None: + try: + self._client.create_pool( + name, gpus, bundle=bundle, minimum_worker_count=minimum_worker_count, queue_pool=queue_pool + ) + except (SIEError, httpx.HTTPError, ValueError) as error: + raise GatewayError(str(error)) from None + + def close(self) -> None: + self._client.close() + + +def connect(base_url: str, api_key: str | None) -> SdkPools: + return SdkPools(SIEClient(base_url, api_key=api_key, timeout_s=TIMEOUT_S)) diff --git a/packages/sie_placement/tests/test_cli.py b/packages/sie_placement/tests/test_cli.py new file mode 100644 index 000000000..54be76b9f --- /dev/null +++ b/packages/sie_placement/tests/test_cli.py @@ -0,0 +1,370 @@ +from __future__ import annotations + +import json +import logging +import sys +import threading +from collections.abc import Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from urllib.parse import parse_qs, urlsplit + +import pytest + +pytest.importorskip("sie_sdk") + +from sie_placement import cli + +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] = {} + + 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}) + if method == "GET" and path.startswith("/v1/pools/"): + if 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)), + } + + 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")}) + 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 tune(gateway: FakeGateway, prometheus: FakePrometheus, *extra: str, plan: Path = FIXTURES / "plan.json") -> int: + return cli.main( + [ + "tune", + "--plan", + str(plan), + "--prometheus-url", + prometheus.url, + "--gateway-url", + gateway.url, + "--once", + *extra, + ] + ) + + +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"].startswith("writing pool 'placement-floor-default-l4-default' failed") + assert "Admin token required" in l4_row["error"] + + +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 / "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_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 == [] diff --git a/packages/sie_placement/tests/test_keda.py b/packages/sie_placement/tests/test_keda.py new file mode 100644 index 000000000..5e80308b1 --- /dev/null +++ b/packages/sie_placement/tests/test_keda.py @@ -0,0 +1,577 @@ +from __future__ import annotations + +import copy +import itertools +import json +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 / "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 = "pinned" 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_the_series_queries_read_each_lane_from_the_gateway_and_worker_series(): + assert keda.series_queries() == { + "queueDepth": "max by (pool, machine_profile, bundle) (sie_gateway_lane_queue_depth{})", + "pendingDemand": "max by (pool, machine_profile, bundle) (sie_gateway_pending_demand{})", + "warmFloor": "max by (pool, machine_profile, bundle) (sie_gateway_pool_warm_floor{})", + "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"].endswith('(sie_gateway_lane_queue_depth{namespace="sie"})') + assert scoped["runners"].endswith('(sie_worker_gpu_slots{namespace="sie",state="ready"}) > 0)') + + +@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") + + +@pytest.mark.parametrize( + "url", + ["file:///etc/passwd", "ftp://prometheus", "prometheus:9090", "http://user:secret@prometheus", "http://p/?x=1"], +) +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" }] From d99d236cb35cf098a50ace102e559ff9625049cd Mon Sep 17 00:00:00 2001 From: krisztian-gajdar Date: Thu, 5 Nov 2026 03:46:49 +0100 Subject: [PATCH 02/10] fix(placement): harden the KEDA applier's tokens, series and leases - refuse a token file, SIE_ADMIN_TOKEN or SIE_API_KEY that is not one run of printable ASCII without spaces, naming its source and never its value, and report a gateway transport failure by its exception type only, so no error text carries an Authorization header value - map malformed Prometheus replies (a bad status line, a truncated body) to a failed series, so the tick holds its lanes and the loop keeps running, and refuse a Prometheus URL with an invalid port at startup - join each gateway series to producers whose snapshot is at most 20 seconds old, as the chart's KEDA triggers do: the lane's queue snapshot for queue depth, the capacity snapshot for pending demand and warm floor, so a stale producer's last point reads as missing - write each carrier through its own SDK client and close it after the write, so the SDK's background lease renewal never runs and a carrier whose reads fail or whose name another pool holds lapses after its last write - report a malformed pool reply from the gateway as a failed lane instead of ending the loop - document that pending demand is a 0 or 1 flag that a publish in flight also sets, so steady traffic can read as pressure, and that the worker series has no freshness timestamp --- packages/sie_placement/README.md | 79 +++++--- .../sie_placement/src/sie_placement/cli.py | 22 ++- .../sie_placement/src/sie_placement/keda.py | 26 ++- .../sie_placement/src/sie_placement/pools.py | 39 +++- packages/sie_placement/tests/test_cli.py | 182 ++++++++++++++++-- packages/sie_placement/tests/test_keda.py | 92 ++++++++- 6 files changed, 379 insertions(+), 61 deletions(-) diff --git a/packages/sie_placement/README.md b/packages/sie_placement/README.md index 873dd8c19..3ff3dfab8 100644 --- a/packages/sie_placement/README.md +++ b/packages/sie_placement/README.md @@ -154,8 +154,10 @@ floor as `current`. It prints one JSON audit row per lane to standard output. Wi `--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. A lane that has no carrier and whose floor stays at 0 gets -none. +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 @@ -176,37 +178,57 @@ 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. -When the applier stops renewing, each carrier's lease lapses, 1200 seconds 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. +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{})` | `running` | -| Pending demand | `max by (pool, machine_profile, bundle) (sie_gateway_pending_demand{})` | `backlog` | -| Warm floor | `max by (pool, machine_profile, bundle) (sie_gateway_pool_warm_floor{})` | Reported only | +| 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, 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 attainment is `unmeasured`, so a lane is under pressure while -the gateway reports pending demand for it, which it does when a request found no -ready capacity, and idle when it reports neither pending demand nor queued or -in-flight work. 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. +`--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, `currentFloor`, the step's `action` and `floor`, -whether the write was `applied`, `laneFloorAboveEnvelope` and `error`. +`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. @@ -221,7 +243,13 @@ The command exits with status 2 before it calls Prometheus or the gateway when: `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. + longer than the gateway's 128 characters; +- a token file or `SIE_ADMIN_TOKEN` holds anything but one run of printable + ASCII characters without spaces after trimming, or the Prometheus URL has an + invalid port. + +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. @@ -232,13 +260,16 @@ status 1. | --- | --- | --- | | 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` and its lease renewal | The gateway admin token | `--admin-token-file`, or `SIE_ADMIN_TOKEN`, read only with `--apply` | +| `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. +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, and a failed gateway +call is reported by the server's error message or, for a transport failure, by +the exception type alone, so no reported error carries a header value. ## Tests diff --git a/packages/sie_placement/src/sie_placement/cli.py b/packages/sie_placement/src/sie_placement/cli.py index 8f653dc5b..03eb2f80a 100644 --- a/packages/sie_placement/src/sie_placement/cli.py +++ b/packages/sie_placement/src/sie_placement/cli.py @@ -6,6 +6,7 @@ import json import math import os +import re import sys import time from collections.abc import Callable, Sequence @@ -15,10 +16,13 @@ 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 @@ -86,18 +90,26 @@ def _parser() -> argparse.ArgumentParser: 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 as error: - raise _UsageError(f"cannot read the {label} file {str(path)!r}: {error.strerror}") from None + 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 value + return _token(value, f"the file {str(path)!r}", label) if env_name is None: return None - return os.environ.get(env_name, "").strip() or 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: @@ -124,6 +136,8 @@ def _tune(args: argparse.Namespace) -> int: keda.series_queries(args.series_selector) prometheus = keda.PrometheusHttp(args.prometheus_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) diff --git a/packages/sie_placement/src/sie_placement/keda.py b/packages/sie_placement/src/sie_placement/keda.py index 3c297b70e..1282e2474 100644 --- a/packages/sie_placement/src/sie_placement/keda.py +++ b/packages/sie_placement/src/sie_placement/keda.py @@ -2,6 +2,7 @@ from __future__ import annotations +import http.client import json import math import re @@ -21,6 +22,7 @@ DEFAULT_MAX_POOLS = 64 MAX_POOL_NAME_LENGTH = 128 PROMETHEUS_TIMEOUT_S = 10.0 +SNAPSHOT_FRESHNESS_S = 20 QUEUE_DEPTH = "queueDepth" PENDING_DEMAND = "pendingDemand" @@ -30,6 +32,7 @@ _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 @@ -180,6 +183,10 @@ def __init__(self, base_url: str, *, timeout_s: float = PROMETHEUS_TIMEOUT_S) -> raise ValueError("the Prometheus URL must not embed credentials") if parts.query or parts.fragment: raise ValueError("the Prometheus URL must not carry a query or a fragment") + try: + parts.port # noqa: B018 + except ValueError: + raise ValueError("the Prometheus URL has an invalid port") from None self._endpoint = f"{base_url.rstrip('/')}/api/v1/query" self._timeout_s = timeout_s @@ -192,6 +199,8 @@ def query(self, promql: str) -> tuple[Sample, ...]: 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) @@ -202,10 +211,21 @@ def series_queries(selector: str = "") -> dict[str, str]: 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: f"max by ({_LANE_LABELS}) (sie_gateway_lane_queue_depth{{{matchers()}}})", - PENDING_DEMAND: f"max by ({_LANE_LABELS}) (sie_gateway_pending_demand{{{matchers()}}})", - WARM_FLOOR: f"max by ({_LANE_LABELS}) (sie_gateway_pool_warm_floor{{{matchers()}}})", + 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)" diff --git a/packages/sie_placement/src/sie_placement/pools.py b/packages/sie_placement/src/sie_placement/pools.py index 860d6f8c1..a6d49ddd0 100644 --- a/packages/sie_placement/src/sie_placement/pools.py +++ b/packages/sie_placement/src/sie_placement/pools.py @@ -10,32 +10,51 @@ from sie_placement.keda import GatewayError TIMEOUT_S = 10.0 +_UNEXPECTED_REPLY = "the gateway returned an unexpected response" + + +def _client(base_url: str, api_key: str | None) -> SIEClient: + try: + return SIEClient(base_url, api_key=api_key, timeout_s=TIMEOUT_S) + except UnicodeError: + raise ValueError("a gateway token holds a character that an HTTP header cannot carry") from None + + +def _gateway_error(error: Exception) -> GatewayError: + if isinstance(error, httpx.HTTPError): + return GatewayError(f"the request failed ({type(error).__name__})") + if isinstance(error, AttributeError): + return GatewayError(_UNEXPECTED_REPLY) + return GatewayError(str(error)) class SdkPools: - def __init__(self, client: SIEClient) -> None: - self._client = client + def __init__(self, base_url: str, api_key: str | None) -> None: + self._base_url = base_url + self._api_key = api_key + self._client = _client(base_url, api_key) def get_pool(self, name: str) -> Mapping[str, object] | None: try: pool = self._client.get_pool(name) - except (SIEError, httpx.HTTPError, ValueError) as error: - raise GatewayError(str(error)) from None + except (SIEError, httpx.HTTPError, ValueError, AttributeError) as error: + raise _gateway_error(error) from None return pool def create_pool( self, name: str, gpus: dict[str, int], *, bundle: str, minimum_worker_count: int, queue_pool: str ) -> None: try: - self._client.create_pool( - name, gpus, bundle=bundle, minimum_worker_count=minimum_worker_count, queue_pool=queue_pool - ) - except (SIEError, httpx.HTTPError, ValueError) as error: - raise GatewayError(str(error)) from None + with _client(self._base_url, self._api_key) as client: + client.create_pool( + name, gpus, bundle=bundle, minimum_worker_count=minimum_worker_count, queue_pool=queue_pool + ) + except (SIEError, httpx.HTTPError, ValueError, AttributeError) as error: + raise _gateway_error(error) from None def close(self) -> None: self._client.close() def connect(base_url: str, api_key: str | None) -> SdkPools: - return SdkPools(SIEClient(base_url, api_key=api_key, timeout_s=TIMEOUT_S)) + return SdkPools(base_url, api_key) diff --git a/packages/sie_placement/tests/test_cli.py b/packages/sie_placement/tests/test_cli.py index 54be76b9f..81eb8aacf 100644 --- a/packages/sie_placement/tests/test_cli.py +++ b/packages/sie_placement/tests/test_cli.py @@ -4,16 +4,18 @@ 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 +from sie_placement import cli, keda, pools FIXTURES = Path(__file__).parent / "fixtures" READ_TOKEN = "read-token-for-tests" # noqa: S105 - a test value @@ -75,13 +77,18 @@ 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 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}) if method == "GET" and path.startswith("/v1/pools/"): - if self.auth and authorization != f"Bearer {READ_TOKEN}": + 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/")) @@ -120,11 +127,19 @@ def __init__(self) -> None: "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) @@ -173,22 +188,16 @@ def read_token_file(tmp_path) -> Path: return path -def tune(gateway: FakeGateway, prometheus: FakePrometheus, *extra: str, plan: Path = FIXTURES / "plan.json") -> int: +def run(gateway: FakeGateway, prometheus: FakePrometheus, *extra: str, plan: Path = FIXTURES / "plan.json") -> int: return cli.main( - [ - "tune", - "--plan", - str(plan), - "--prometheus-url", - prometheus.url, - "--gateway-url", - gateway.url, - "--once", - *extra, - ] + ["tune", "--plan", str(plan), "--prometheus-url", prometheus.url, "--gateway-url", gateway.url, *extra] ) +def tune(gateway: FakeGateway, prometheus: FakePrometheus, *extra: str, plan: Path = FIXTURES / "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()] @@ -368,3 +377,148 @@ def test_the_core_imports_without_the_sdk_and_tune_names_the_missing_extra(gatew 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 == {} + + +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") + + +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 index 5e80308b1..c867cd32a 100644 --- a/packages/sie_placement/tests/test_keda.py +++ b/packages/sie_placement/tests/test_keda.py @@ -3,6 +3,7 @@ import copy import itertools import json +import socketserver import threading from collections.abc import Iterator, Mapping, Sequence from datetime import UTC, datetime @@ -205,18 +206,49 @@ def test_a_carrier_name_longer_than_the_gateway_accepts_is_refused(): keda.carriers(plan) -def test_the_series_queries_read_each_lane_from_the_gateway_and_worker_series(): +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{})", - "pendingDemand": "max by (pool, machine_profile, bundle) (sie_gateway_pending_demand{})", - "warmFloor": "max by (pool, machine_profile, bundle) (sie_gateway_pool_warm_floor{})", + "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"].endswith('(sie_gateway_lane_queue_depth{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"): @@ -309,9 +341,57 @@ def test_the_http_reader_reports_an_unreachable_prometheus(): 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"], + [ + "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"): From 818f848b6995fc050bc4b4d3c6240e0eaffa4103 Mon Sep 17 00:00:00 2001 From: krisztian-gajdar Date: Thu, 5 Nov 2026 19:37:21 +0100 Subject: [PATCH 03/10] fix(placement): refuse plans the chart cannot run before tuning their floors sie-placement tune now refuses, before it calls Prometheus or the gateway, a plan that pins an identity on a lane of the default queue pool or gives an identity a warmth the sie-cluster chart does not realize (bridged, snapshot-cold or scheduled-warm), in the wording the Helm values emitter uses for the same plans, so it never moves floors for a plan the chart cannot run. The tune tests read a plan the chart can run. --- packages/sie_placement/README.md | 4 ++ .../sie_placement/src/sie_placement/keda.py | 33 ++++++++- .../tests/fixtures/tune-plan.json | 68 +++++++++++++++++++ packages/sie_placement/tests/test_cli.py | 19 +++++- packages/sie_placement/tests/test_keda.py | 49 ++++++++++++- 5 files changed, 167 insertions(+), 6 deletions(-) create mode 100644 packages/sie_placement/tests/fixtures/tune-plan.json diff --git a/packages/sie_placement/README.md b/packages/sie_placement/README.md index 3ff3dfab8..7cf3d5ce6 100644 --- a/packages/sie_placement/README.md +++ b/packages/sie_placement/README.md @@ -244,6 +244,10 @@ The command exits with status 2 before it calls Prometheus or the gateway when: - 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, or the Prometheus URL has an invalid port. diff --git a/packages/sie_placement/src/sie_placement/keda.py b/packages/sie_placement/src/sie_placement/keda.py index 1282e2474..70a43e153 100644 --- a/packages/sie_placement/src/sie_placement/keda.py +++ b/packages/sie_placement/src/sie_placement/keda.py @@ -14,10 +14,11 @@ from datetime import UTC, datetime from typing import Protocol -from sie_placement.schema import Plan, PlanLane +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 @@ -29,6 +30,12 @@ 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" @@ -72,6 +79,10 @@ def canonical(cls, queue_pool: str, machine_profile: str, bundle: str) -> LaneKe 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: @@ -98,6 +109,10 @@ def carriers( ) 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) @@ -106,6 +121,14 @@ def carriers( 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 " @@ -120,6 +143,14 @@ def carriers( 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) 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 index 81eb8aacf..7adb7ef84 100644 --- a/packages/sie_placement/tests/test_cli.py +++ b/packages/sie_placement/tests/test_cli.py @@ -188,13 +188,15 @@ def read_token_file(tmp_path) -> Path: return path -def run(gateway: FakeGateway, prometheus: FakePrometheus, *extra: str, plan: Path = FIXTURES / "plan.json") -> int: +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 / "plan.json") -> int: +def tune( + gateway: FakeGateway, prometheus: FakePrometheus, *extra: str, plan: Path = FIXTURES / "tune-plan.json" +) -> int: return run(gateway, prometheus, "--once", *extra, plan=plan) @@ -311,7 +313,7 @@ def test_a_rejected_write_fails_the_tick_and_names_the_pool(gateway, prometheus, 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 / "plan.json").read_text(encoding="utf-8")) + 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") @@ -323,6 +325,17 @@ def test_an_envelope_above_the_pool_floor_limit_is_refused_before_any_call( 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 diff --git a/packages/sie_placement/tests/test_keda.py b/packages/sie_placement/tests/test_keda.py index c867cd32a..3fa82a226 100644 --- a/packages/sie_placement/tests/test_keda.py +++ b/packages/sie_placement/tests/test_keda.py @@ -26,7 +26,7 @@ def plan_document() -> dict: - return json.loads((FIXTURES / "plan.json").read_text(encoding="utf-8")) + return json.loads((FIXTURES / "tune-plan.json").read_text(encoding="utf-8")) def make_plan(*lanes: dict) -> Plan: @@ -44,7 +44,7 @@ def make_plan(*lanes: dict) -> Plan: if machine_profile is not None: lane["cell"]["machineProfile"] = machine_profile document["lanes"].append(lane) - warmth = "pinned" if lane["envelope"]["floorMin"] >= 1 else "cold-with-bound" + 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) @@ -206,6 +206,51 @@ def test_a_carrier_name_longer_than_the_gateway_accepts_is_refused(): 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) " From baef2579ccb63f8b75c8b084572265795eb57af1 Mon Sep 17 00:00:00 2001 From: krisztian-gajdar Date: Fri, 6 Nov 2026 11:08:48 +0100 Subject: [PATCH 04/10] fix(placement): fail only the tick on a pool reply too deep to parse - A gateway reply nested deeper than the JSON decoder's recursion limit raised RecursionError, which escaped the tick and ended the applier. Pool reads and writes now report it as an unexpected response, as the Prometheus reader already does for its replies. --- .../sie_placement/src/sie_placement/pools.py | 6 +-- packages/sie_placement/tests/test_cli.py | 50 +++++++++++++++++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/packages/sie_placement/src/sie_placement/pools.py b/packages/sie_placement/src/sie_placement/pools.py index a6d49ddd0..886cfa29c 100644 --- a/packages/sie_placement/src/sie_placement/pools.py +++ b/packages/sie_placement/src/sie_placement/pools.py @@ -23,7 +23,7 @@ def _client(base_url: str, api_key: str | None) -> SIEClient: def _gateway_error(error: Exception) -> GatewayError: if isinstance(error, httpx.HTTPError): return GatewayError(f"the request failed ({type(error).__name__})") - if isinstance(error, AttributeError): + if isinstance(error, AttributeError | RecursionError): return GatewayError(_UNEXPECTED_REPLY) return GatewayError(str(error)) @@ -37,7 +37,7 @@ def __init__(self, base_url: str, api_key: str | None) -> None: def get_pool(self, name: str) -> Mapping[str, object] | None: try: pool = self._client.get_pool(name) - except (SIEError, httpx.HTTPError, ValueError, AttributeError) as error: + except (SIEError, httpx.HTTPError, ValueError, AttributeError, RecursionError) as error: raise _gateway_error(error) from None return pool @@ -49,7 +49,7 @@ def create_pool( client.create_pool( name, gpus, bundle=bundle, minimum_worker_count=minimum_worker_count, queue_pool=queue_pool ) - except (SIEError, httpx.HTTPError, ValueError, AttributeError) as error: + except (SIEError, httpx.HTTPError, ValueError, AttributeError, RecursionError) as error: raise _gateway_error(error) from None def close(self) -> None: diff --git a/packages/sie_placement/tests/test_cli.py b/packages/sie_placement/tests/test_cli.py index 7adb7ef84..b9a7f63d7 100644 --- a/packages/sie_placement/tests/test_cli.py +++ b/packages/sie_placement/tests/test_cli.py @@ -79,11 +79,21 @@ def __init__(self, *, auth: bool = True) -> None: self.pools: dict[str, dict] = {} self.fail_reads = False self.read_reply: object | None = None + self.raw_reply: tuple[int, bytes] | None = None 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}) + if self.raw_reply is not None: + status, payload = self.raw_reply + 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) @@ -484,6 +494,46 @@ def test_a_token_an_http_header_cannot_encode_is_refused_without_echoing_it(gate 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"), + ], +) +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_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 From 66a01bc79e47ea0434fbce5a325c04b4b8a30c92 Mon Sep 17 00:00:00 2001 From: krisztian-gajdar Date: Fri, 6 Nov 2026 11:09:09 +0100 Subject: [PATCH 05/10] fix(placement): check the gateway URL as the Prometheus URL is checked - tune refuses with exit 2, before it builds any client, a --gateway-url that is not an absolute http or https URL with a host, embeds credentials, carries a query or a fragment, or has an invalid port, and never repeats the URL. Credentials in the URL made httpx send Basic auth in place of the bearer token, and an unparsable URL's error quoted its password. - Both URLs share one check, require_base_url. --- packages/sie_placement/README.md | 6 ++-- .../sie_placement/src/sie_placement/cli.py | 1 + .../sie_placement/src/sie_placement/keda.py | 26 ++++++++------- packages/sie_placement/tests/test_cli.py | 33 +++++++++++++++++++ 4 files changed, 53 insertions(+), 13 deletions(-) diff --git a/packages/sie_placement/README.md b/packages/sie_placement/README.md index 7cf3d5ce6..2daaf00ce 100644 --- a/packages/sie_placement/README.md +++ b/packages/sie_placement/README.md @@ -249,8 +249,10 @@ The command exits with status 2 before it calls Prometheus or the gateway when: 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, or the Prometheus URL has an - invalid port. + ASCII characters without spaces after trimming; +- `--prometheus-url` or `--gateway-url` is not an 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. diff --git a/packages/sie_placement/src/sie_placement/cli.py b/packages/sie_placement/src/sie_placement/cli.py index 03eb2f80a..ccd1df9b9 100644 --- a/packages/sie_placement/src/sie_placement/cli.py +++ b/packages/sie_placement/src/sie_placement/cli.py @@ -135,6 +135,7 @@ def _tune(args: argparse.Namespace) -> int: 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") diff --git a/packages/sie_placement/src/sie_placement/keda.py b/packages/sie_placement/src/sie_placement/keda.py index 70a43e153..167ef842e 100644 --- a/packages/sie_placement/src/sie_placement/keda.py +++ b/packages/sie_placement/src/sie_placement/keda.py @@ -205,19 +205,23 @@ def parse_vector(document: object) -> tuple[Sample, ...]: return tuple(samples) +def require_base_url(label: str, base_url: str) -> None: + parts = urllib.parse.urlsplit(base_url) + if parts.scheme not in {"http", "https"} or not parts.hostname: + raise ValueError(f"the {label} URL must be an absolute http or https URL") + if parts.username is not None or parts.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: - parts = urllib.parse.urlsplit(base_url) - if parts.scheme not in {"http", "https"} or not parts.hostname: - raise ValueError("the Prometheus URL must be an absolute http or https URL") - if parts.username is not None or parts.password is not None: - raise ValueError("the Prometheus URL must not embed credentials") - if parts.query or parts.fragment: - raise ValueError("the Prometheus URL must not carry a query or a fragment") - try: - parts.port # noqa: B018 - except ValueError: - raise ValueError("the Prometheus URL has an invalid port") from None + require_base_url("Prometheus", base_url) self._endpoint = f"{base_url.rstrip('/')}/api/v1/query" self._timeout_s = timeout_s diff --git a/packages/sie_placement/tests/test_cli.py b/packages/sie_placement/tests/test_cli.py index b9a7f63d7..20c2f6427 100644 --- a/packages/sie_placement/tests/test_cli.py +++ b/packages/sie_placement/tests/test_cli.py @@ -489,6 +489,39 @@ def test_a_transport_failure_is_reported_without_the_header_it_carried(gateway): 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"), + ], +) +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_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") From b77337adaffce3ad501e3758895033eac7c400b1 Mon Sep 17 00:00:00 2001 From: krisztian-gajdar Date: Fri, 6 Nov 2026 11:09:30 +0100 Subject: [PATCH 06/10] fix(placement): fail a pool call whose reply is neither a success nor an error - The SDK treats every status below 400 as success, so a proxy answering a pool write with a redirect and a JSON body made the audit row report the write as applied. The applier now gives the SDK its own httpx client, through SIEClient's http_client parameter, with redirects off and a response hook that fails any reply outside 2xx below 400. Error replies still reach the SDK, so get_pool keeps reading a 404 as no pool. - A gateway URL that httpx cannot parse is refused without quoting it. --- .../sie_placement/src/sie_placement/pools.py | 24 ++++++++++++++++- packages/sie_placement/tests/test_cli.py | 26 +++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/packages/sie_placement/src/sie_placement/pools.py b/packages/sie_placement/src/sie_placement/pools.py index 886cfa29c..dcc50f1c9 100644 --- a/packages/sie_placement/src/sie_placement/pools.py +++ b/packages/sie_placement/src/sie_placement/pools.py @@ -10,14 +10,36 @@ from sie_placement.keda import GatewayError TIMEOUT_S = 10.0 +HTTP_OK_MIN = 200 +HTTP_OK_MAX = 299 +HTTP_ERROR_MIN = 400 _UNEXPECTED_REPLY = "the gateway returned an unexpected response" +def _require_success_or_error(response: httpx.Response) -> None: + status = response.status_code + 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) -> SIEClient: try: - return SIEClient(base_url, api_key=api_key, timeout_s=TIMEOUT_S) + http_client = httpx.Client( + base_url=base_url, + timeout=TIMEOUT_S, + follow_redirects=False, + event_hooks={"response": [_require_success_or_error]}, + ) + except httpx.InvalidURL: + raise ValueError("the gateway URL is not a valid URL") from None + try: + return 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 def _gateway_error(error: Exception) -> GatewayError: diff --git a/packages/sie_placement/tests/test_cli.py b/packages/sie_placement/tests/test_cli.py index 20c2f6427..b4cfef5ec 100644 --- a/packages/sie_placement/tests/test_cli.py +++ b/packages/sie_placement/tests/test_cli.py @@ -522,6 +522,11 @@ def test_a_gateway_url_that_is_not_a_plain_base_url_is_refused_without_echoing_i 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") @@ -539,6 +544,8 @@ def write_carrier(writer: pools.SdkPools) -> None: ("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): @@ -567,6 +574,25 @@ def test_a_deeply_nested_pool_reply_fails_only_the_tick(gateway, prometheus, rea ] +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 == {} + + 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 From cdfeef9f209c4a0fa52263e214c25d4a3489eb31 Mon Sep 17 00:00:00 2001 From: krisztian-gajdar Date: Fri, 6 Nov 2026 11:09:36 +0100 Subject: [PATCH 07/10] docs(placement): say how each kind of gateway call failure is reported - A failed connection is reported with the SDK's connection error message, which carries no header value because no header has been sent, and any other transport failure by its exception type alone. - A reply outside 2xx below 400 fails the call with its status code, and a reply the SDK cannot parse fails it as an unexpected response. --- packages/sie_placement/README.md | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/packages/sie_placement/README.md b/packages/sie_placement/README.md index 2daaf00ce..0498d48d1 100644 --- a/packages/sie_placement/README.md +++ b/packages/sie_placement/README.md @@ -273,9 +273,13 @@ 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, and a failed gateway -call is reported by the server's error message or, for a transport failure, by -the exception type alone, so no reported error carries a header value. +refused token is named by its source, never by its value. A failed gateway call +is reported by the server's error message, a failed connection by the SDK's +connection error message, and any other transport failure by its exception type +alone, so no reported error carries a header value: a connection fails before +any header is sent. A reply that is neither a success nor an error, such as a +redirect, fails the call with its status code, and a reply the SDK cannot parse +fails it as an unexpected response. ## Tests From 56ed3d78b84da969c441ed94a1b8d79e26643fa3 Mon Sep 17 00:00:00 2001 From: krisztian-gajdar Date: Fri, 6 Nov 2026 11:32:34 +0100 Subject: [PATCH 08/10] fix(placement): refuse a URL urllib cannot split without repeating it - urlsplit raises a ValueError that quotes the whole netloc, credentials included, when the netloc holds a character that NFKC normalization turns into a separator, such as a fullwidth slash or at sign. require_base_url now reports such a --gateway-url or --prometheus-url as not a valid URL. --- packages/sie_placement/README.md | 4 ++-- packages/sie_placement/src/sie_placement/keda.py | 10 +++++++--- packages/sie_placement/tests/test_cli.py | 16 ++++++++++++++++ 3 files changed, 25 insertions(+), 5 deletions(-) diff --git a/packages/sie_placement/README.md b/packages/sie_placement/README.md index 0498d48d1..f4761707b 100644 --- a/packages/sie_placement/README.md +++ b/packages/sie_placement/README.md @@ -250,8 +250,8 @@ The command exits with status 2 before it calls Prometheus or the gateway when: 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 an absolute http or https URL - with a host, embeds credentials, carries a query or a fragment, or has an +- `--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, diff --git a/packages/sie_placement/src/sie_placement/keda.py b/packages/sie_placement/src/sie_placement/keda.py index 167ef842e..2ae451ed3 100644 --- a/packages/sie_placement/src/sie_placement/keda.py +++ b/packages/sie_placement/src/sie_placement/keda.py @@ -206,10 +206,14 @@ def parse_vector(document: object) -> tuple[Sample, ...]: def require_base_url(label: str, base_url: str) -> None: - parts = urllib.parse.urlsplit(base_url) - if parts.scheme not in {"http", "https"} or not parts.hostname: + 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 parts.username is not None or parts.password is not None: + 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") diff --git a/packages/sie_placement/tests/test_cli.py b/packages/sie_placement/tests/test_cli.py index b4cfef5ec..64c9d7148 100644 --- a/packages/sie_placement/tests/test_cli.py +++ b/packages/sie_placement/tests/test_cli.py @@ -503,6 +503,8 @@ def test_a_transport_failure_is_reported_without_the_header_it_carried(gateway): (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( @@ -522,6 +524,20 @@ def test_a_gateway_url_that_is_not_a_plain_base_url_is_refused_without_echoing_i 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) From e7952b2c2a94491315dba63bd928ac0bedb78d5f Mon Sep 17 00:00:00 2001 From: krisztian-gajdar Date: Fri, 6 Nov 2026 11:32:49 +0100 Subject: [PATCH 09/10] fix(placement): redact the client's own token from reported gateway errors - A gateway call's error text is copied from the reply, so an intermediary that quotes the request's Authorization header in its error body put the token in an audit row. SdkPools now replaces the bearer token its own client sent, the client token on reads and the admin token on writes, with [redacted], and cuts the reported text to 512 characters, which also bounds an audit row against a large error body. --- packages/sie_placement/README.md | 11 +++-- .../sie_placement/src/sie_placement/pools.py | 26 +++++++---- packages/sie_placement/tests/test_cli.py | 43 ++++++++++++++++++- 3 files changed, 66 insertions(+), 14 deletions(-) diff --git a/packages/sie_placement/README.md b/packages/sie_placement/README.md index f4761707b..bb444a26f 100644 --- a/packages/sie_placement/README.md +++ b/packages/sie_placement/README.md @@ -276,10 +276,13 @@ 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 failed gateway call is reported by the server's error message, a failed connection by the SDK's connection error message, and any other transport failure by its exception type -alone, so no reported error carries a header value: a connection fails before -any header is sent. A reply that is neither a success nor an error, such as a -redirect, fails the call with its status code, and a reply the SDK cannot parse -fails it as an unexpected response. +alone; a connection fails before any header is sent. A reply that is neither a +success nor an error, such as a redirect, fails the call with its status code, +and a reply the SDK cannot parse fails it as an unexpected response. The +applier cuts a reported error to 512 characters and replaces the token its +client sent with `[redacted]`, so an intermediary that quotes the request's +`Authorization` header in its error body does not put the token in an audit +row. Other text a server returns is reported as it is. ## Tests diff --git a/packages/sie_placement/src/sie_placement/pools.py b/packages/sie_placement/src/sie_placement/pools.py index dcc50f1c9..46750f7a3 100644 --- a/packages/sie_placement/src/sie_placement/pools.py +++ b/packages/sie_placement/src/sie_placement/pools.py @@ -13,6 +13,9 @@ 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" @@ -22,7 +25,7 @@ def _require_success_or_error(response: httpx.Response) -> None: 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) -> SIEClient: +def _client(base_url: str, api_key: str | None) -> tuple[SIEClient, str | None]: try: http_client = httpx.Client( base_url=base_url, @@ -33,46 +36,53 @@ def _client(base_url: str, api_key: str | None) -> SIEClient: except httpx.InvalidURL: raise ValueError("the gateway URL is not a valid URL") from None try: - return SIEClient(base_url, api_key=api_key, timeout_s=TIMEOUT_S, http_client=http_client) + 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, token or None -def _gateway_error(error: Exception) -> GatewayError: +def _gateway_error(error: Exception, token: str | None) -> GatewayError: if isinstance(error, httpx.HTTPError): return GatewayError(f"the request failed ({type(error).__name__})") if isinstance(error, AttributeError | RecursionError): return GatewayError(_UNEXPECTED_REPLY) - return GatewayError(str(error)) + 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 = _client(base_url, api_key) + self._client, self._token = _client(base_url, api_key) def get_pool(self, name: str) -> Mapping[str, object] | None: try: pool = self._client.get_pool(name) except (SIEError, httpx.HTTPError, ValueError, AttributeError, RecursionError) as error: - raise _gateway_error(error) from None + raise _gateway_error(error, 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: try: - with _client(self._base_url, self._api_key) as client: + client, _ = _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) from None + raise _gateway_error(error, self._token) from None def close(self) -> None: self._client.close() diff --git a/packages/sie_placement/tests/test_cli.py b/packages/sie_placement/tests/test_cli.py index 64c9d7148..95ed302e5 100644 --- a/packages/sie_placement/tests/test_cli.py +++ b/packages/sie_placement/tests/test_cli.py @@ -80,13 +80,15 @@ def __init__(self, *, auth: bool = True) -> None: 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}) - if self.raw_reply is not None: - status, payload = self.raw_reply + 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))) @@ -609,6 +611,43 @@ def redirect_writes(handler: _Handler, method: str, body: object) -> None: assert gateway.pools == {} +ECHOES = { + "403-json": lambda token: ( + 403, + json.dumps({"detail": {"code": "FORBIDDEN", "message": f"token not allowed: 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", "token"), [("GET", READ_TOKEN), ("POST", ADMIN_TOKEN)]) +def test_a_reply_that_echoes_the_clients_token_is_redacted_in_the_audit_row( + gateway, prometheus, read_token_file, monkeypatch, method, token, echo, capsys +): + monkeypatch.setenv("SIE_ADMIN_TOKEN", ADMIN_TOKEN) + gateway.raw_replies[method] = ECHOES[echo](token) + assert tune(gateway, prometheus, "--api-key-file", str(read_token_file), "--apply") == cli.EXIT_FAILED + captured = capsys.readouterr() + errors = [row["error"] for row in rows(captured.out) if row["error"]] + assert errors + assert all("Bearer [redacted]" in error for error in errors) + for secret in (READ_TOKEN, ADMIN_TOKEN): + assert secret not in captured.out + assert secret not in captured.err + + +def test_a_long_gateway_error_is_cut_to_512_characters(gateway, prometheus, read_token_file, capsys): + gateway.raw_replies["GET"] = (502, b"x" * 5000) + 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 + for error in errors: + prefix, message = error.split(" failed: ", 1) + assert prefix.startswith("reading pool ") + assert len(message) == 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 From 440321d761f67ede811cdd7f23a3697646c79f91 Mon Sep 17 00:00:00 2001 From: krisztian-gajdar Date: Fri, 6 Nov 2026 11:58:22 +0100 Subject: [PATCH 10/10] fix(placement): report a gateway error reply by its status alone - Redacting the client's token from the reply text missed any encoding of it: the SDK reports a detail object without a message as its repr, which doubles each backslash, so a token holding one reached the audit row. A pool call that gets an error reply is now reported only by its HTTP status, recorded by the response hook of the applier's own httpx client, and a success reply the SDK cannot parse as an unexpected response. - A message the SDK raises before any reply, such as a connection error, keeps its text, cut to 512 characters with the client's token redacted. --- packages/sie_placement/README.md | 20 +++---- .../sie_placement/src/sie_placement/pools.py | 36 +++++++----- packages/sie_placement/tests/test_cli.py | 56 +++++++++++-------- 3 files changed, 65 insertions(+), 47 deletions(-) diff --git a/packages/sie_placement/README.md b/packages/sie_placement/README.md index bb444a26f..326cbcbf5 100644 --- a/packages/sie_placement/README.md +++ b/packages/sie_placement/README.md @@ -273,16 +273,16 @@ 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 failed gateway call -is reported by the server's error message, a failed connection by the SDK's -connection error message, and any other transport failure by its exception type -alone; a connection fails before any header is sent. A reply that is neither a -success nor an error, such as a redirect, fails the call with its status code, -and a reply the SDK cannot parse fails it as an unexpected response. The -applier cuts a reported error to 512 characters and replaces the token its -client sent with `[redacted]`, so an intermediary that quotes the request's -`Authorization` header in its error body does not put the token in an audit -row. Other text a server returns is reported as it is. +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 diff --git a/packages/sie_placement/src/sie_placement/pools.py b/packages/sie_placement/src/sie_placement/pools.py index 46750f7a3..f91556f37 100644 --- a/packages/sie_placement/src/sie_placement/pools.py +++ b/packages/sie_placement/src/sie_placement/pools.py @@ -19,19 +19,25 @@ _UNEXPECTED_REPLY = "the gateway returned an unexpected response" -def _require_success_or_error(response: httpx.Response) -> None: - status = response.status_code - 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") +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, str | None]: + +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": [_require_success_or_error]}, + event_hooks={"response": [replies]}, ) except httpx.InvalidURL: raise ValueError("the gateway URL is not a valid URL") from None @@ -45,13 +51,15 @@ def _client(base_url: str, api_key: str | None) -> tuple[SIEClient, str | None]: raise authorization = http_client.headers.get("Authorization", "") token = authorization.removeprefix(_BEARER) if authorization.startswith(_BEARER) else "" - return client, token or None + return client, replies, token or None -def _gateway_error(error: Exception, token: str | None) -> GatewayError: +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 isinstance(error, AttributeError | RecursionError): + 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: @@ -63,26 +71,28 @@ 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._token = _client(base_url, 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._token) from None + 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, _ = _client(self._base_url, self._api_key) + 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, self._token) from None + raise _gateway_error(error, None if replies is None else replies.status, self._token) from None def close(self) -> None: self._client.close() diff --git a/packages/sie_placement/tests/test_cli.py b/packages/sie_placement/tests/test_cli.py index 95ed302e5..5a0a1fc16 100644 --- a/packages/sie_placement/tests/test_cli.py +++ b/packages/sie_placement/tests/test_cli.py @@ -16,6 +16,7 @@ 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 @@ -318,8 +319,7 @@ def test_a_rejected_write_fails_the_tick_and_names_the_pool(gateway, prometheus, 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"].startswith("writing pool 'placement-floor-default-l4-default' failed") - assert "Admin token required" in l4_row["error"] + 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( @@ -611,41 +611,49 @@ def redirect_writes(handler: _Handler, method: str, body: object) -> None: 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-json": lambda token: ( + "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", "token"), [("GET", READ_TOKEN), ("POST", ADMIN_TOKEN)]) -def test_a_reply_that_echoes_the_clients_token_is_redacted_in_the_audit_row( - gateway, prometheus, read_token_file, monkeypatch, method, token, echo, capsys +@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 ): - monkeypatch.setenv("SIE_ADMIN_TOKEN", ADMIN_TOKEN) - gateway.raw_replies[method] = ECHOES[echo](token) - assert tune(gateway, prometheus, "--api-key-file", str(read_token_file), "--apply") == cli.EXIT_FAILED + 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("Bearer [redacted]" in error for error in errors) - for secret in (READ_TOKEN, ADMIN_TOKEN): - assert secret not in captured.out - assert secret not in captured.err - - -def test_a_long_gateway_error_is_cut_to_512_characters(gateway, prometheus, read_token_file, capsys): - gateway.raw_replies["GET"] = (502, b"x" * 5000) - 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 - for error in errors: - prefix, message = error.split(" failed: ", 1) - assert prefix.startswith("reading pool ") - assert len(message) == pools.MAX_ERROR_LENGTH + 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):