diff --git a/.github/workflows/model-smoke.yml b/.github/workflows/model-smoke.yml index 072c98c..a3dfd3b 100644 --- a/.github/workflows/model-smoke.yml +++ b/.github/workflows/model-smoke.yml @@ -18,6 +18,8 @@ on: - worker/ai_bridge/jobs.py # The worker's HTTP handling, which every model test and the MCP workers' health check go through. - worker/ai_bridge/server.py + # Imported by both: the counters every model job passes through. + - worker/ai_bridge/metrics.py - scripts/benchmark.py - scripts/ner_measure.py - tests/fixtures/ner-sample.json diff --git a/CHANGELOG.md b/CHANGELOG.md index 7219f85..84a76db 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ All notable changes to this project are documented in this file. The format foll ## [Unreleased] +### Added + +- `GET /metrics` on the worker, in the Prometheus text format and behind the service token: jobs by task and terminal status, submissions, 429 refusals, queue depth, running jobs, pool size and readiness, the concurrency and capacity settings, and per-task histograms of job and queue durations. Aggregates only, with labels from fixed sets; they live in memory and reset when the worker restarts. + ## [0.2.0] - 2026-10-21 ### Added diff --git a/README.md b/README.md index 9c14a19..8b44f9f 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ Includes a FrankenPHP worker-mode example, deterministic fault tests, and an opt - Redaction: checksum and pattern rules for card numbers, IBANs, emails, IPv4 and IPv6 addresses and phone numbers on every backend, plus named-entity masking with the int8 cased and uncased `bert-base-NER` models on the ONNX profile. Assistive, not a compliance control. - Bounded concurrency, queue/result capacity, retention, request sizes and execution deadlines. - Cancellation, progress polling, crash detection and generic errors that omit internal exception text. +- `GET /metrics` on the worker: aggregate job counts, queue depth, pool readiness and duration histograms in the Prometheus text format, behind the service token and carrying no job content; counters reset on restart ([the protocol](docs/protocol.md#metrics)). - A FrankenPHP example that submits work and returns a job ID immediately. - A stdio MCP server with one tool per task, a second client of the same protocol beside the PHP client, plus an index builder over the protocol and a `bridge_search` tool that searches what it built, returning offsets and snippets rather than vectors, a compose service that puts the pinned ONNX models behind it, and a ten-question evaluation for LLM hosts with a runner for the Claude Code CLI; see [the MCP guide](docs/mcp.md). - CPU inference with `cross-encoder/ms-marco-TinyBERT-L2-v2`, `sentence-transformers/all-MiniLM-L6-v2` and int8 `dslim/bert-base-NER` and `dslim/bert-base-NER-uncased` through ONNX Runtime, using pinned model files and wheel hashes. diff --git a/docs/protocol.md b/docs/protocol.md index 20bf62b..5c6796f 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -5,6 +5,7 @@ Every route requires `Authorization: Bearer `. JSON requests requ | Route | Result | | --- | --- | | `GET /healthz` | 200 with protocol version and configured backend; process liveness only, not model readiness | +| `GET /metrics` | 200 with aggregate counts and timings in the Prometheus text format ([below](#metrics)) | | `POST /v1/jobs` | 202 with a queued job, or with the retained job a repeated `idempotency_key` names | | `GET /v1/jobs/{id}` | 200 with a snapshot, or 404 | | `POST /v1/jobs/{id}/cancel` | Send `{}`; returns the terminal/current snapshot, or 404 | @@ -40,3 +41,24 @@ With `BRIDGE_STATE_DIR` set, the job table (never a job's input) is written thro The PHP client uses a fresh cURL handle per call, refuses redirects, ignores proxy environment variables, caps response bytes, and verifies that GET/cancel responses match the requested ID. Configure its request timeout separately from job execution and local waiting deadlines. There are no transparent submission retries, a 429 included. The MCP server in `mcp_server/` is a second, independent client with the same checks; unlike the PHP client it cancels a job whose local wait expires, because no later request will collect it. Adding a task requires trusted source changes to `worker/ai_bridge/tasks.py`: an explicit name, input validation, bounded output and execution logic. Add a PHP result type where useful, plus successful, malformed-input and lifecycle tests. Task imports, file paths and command strings must not come from request data. + +## Metrics + +`GET /metrics` takes the same bearer token as every other route (401 without it; a `POST` is answered as on any unknown route, 404 for `{}`, and another spelling, a query string included, is 404) and answers `Content-Type: text/plain; version=0.0.4; charset=utf-8`, the Prometheus text exposition format, written by the standard library. Every value is an aggregate since the worker process started, held in memory only: a restart resets every counter and histogram to zero, the job table in `BRIDGE_STATE_DIR` does not feed them, and a job read back from that table counts nothing. Each family has one `# HELP` and one `# TYPE` line, and every series exists from the start, at zero. Label values come only from fixed sets: `task` is a task of the allowlist (`rerank`, `embed`, `redact`, plus the test tasks on a worker started with `BRIDGE_TEST_TASKS=1`), `status` one of the four terminal statuses, `reason` is `capacity`, and `backend` the configured backend; no job ID, idempotency key, input, result, error, model path or token appears. + +| Family | Type | Labels | What it holds | +| --- | --- | --- | --- | +| `bridge_jobs_total` | counter | `task`, `status` | Jobs that reached `succeeded`, `failed`, `timed_out` or `cancelled` | +| `bridge_jobs_submitted_total` | counter | `task` | Accepted submissions that created a job; a repeated `idempotency_key` creates none and is not counted | +| `bridge_submissions_rejected_total` | counter | `reason="capacity"` | Submissions answered 429, for every cause the 429 section names | +| `bridge_queue_depth` | gauge | | Jobs `queued` now | +| `bridge_jobs_running` | gauge | | Jobs `running` now | +| `bridge_pool_size` | gauge | | Child processes in the pool, starting or started | +| `bridge_pool_ready` | gauge | | Of those, the ones that have said they are ready, idle or running a job | +| `bridge_concurrency` | gauge | | The `BRIDGE_CONCURRENCY` setting | +| `bridge_capacity` | gauge | | The `BRIDGE_CAPACITY` setting | +| `bridge_job_duration_seconds` | histogram | `task` | Seconds from submission to terminal status, for every job counted by `bridge_jobs_total` | +| `bridge_job_queue_seconds` | histogram | `task` | Seconds from submission to `running`, for jobs a child started; one that ended while queued adds nothing | +| `bridge_info` | gauge | `protocol="1"`, `backend` | Always 1 | + +Both histograms have the buckets 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60 and 300 seconds and `+Inf`, cumulative as the format requires, with `_sum` and `_count`. A scrape copies the numbers under the store's lock and formats them without it, so it never waits for a job. An invalid submission (400) creates no job and counts nothing. diff --git a/docs/security.md b/docs/security.md index 658c9c2..cbf8afa 100644 --- a/docs/security.md +++ b/docs/security.md @@ -12,6 +12,7 @@ This experimental service is designed for trusted application-to-service traffic - Requests and results may contain sensitive documents. They live in RAM, may be present in process/core dumps, and are not encrypted by this service. With `BRIDGE_STATE_DIR` set, `jobs.sqlite` (and its `-wal` and `-shm` files) in that directory holds each retained job's ID, task, status, error, progress, deadline, idempotency key, submission and finish times, and its result until retention expires; it never holds a job's input, though results such as redacted text, rankings and embedding vectors reveal much of it. SQLite's `secure_delete` overwrites an expired row in the database file, but older copies of its page can stay in the `-wal` file until a checkpoint reuses it, so keep the directory on tmpfs or an encrypted volume readable only by the service user (compose's `worker` uses a 64 MB tmpfs owned by uid 65532 with mode 0700). Minimize retention and use appropriate host/storage policies. Embedding vectors are derived from their input text and can reveal much of it; store and transmit them with the same care as the text. - Redaction is assistive, not a compliance control. The rules find only the structured identifiers they recognise, and the optional named-entity models were trained on English news, so their recall on names from other domains or languages is lower and unmeasured here; review redacted output before treating it as free of personal data. The model files are third-party ONNX conversions of the original weights, pinned by revision and hash but not signed by the original author. - The MCP server is a client of this service for an AI agent over stdio, where there is no authentication: the host that launches it inherits its trust, and it holds the token from its environment and never returns it. Running it over HTTP or SSE would be a new authorisation surface that the bearer token does not cover; it is out of scope. An agent can be prompt-injected and can then fill a worker's capacity, so point the MCP server at a dedicated worker, never at one serving a production application; compose's `mcp` and `mcp-model` each have their own. It reads files only under `BRIDGE_MCP_ROOT` and writes none; snippets it returns are the agent's own file content and enter the agent's context like any file it reads. Index files built for `bridge_search` are derived data: `index.json` holds every chunk's text verbatim and `vectors.f32` holds vectors that can reveal it, so keep them with the same care as their sources. The builder runs with the service token like the PHP client and is the only component that writes, into the directory the operator names. +- `GET /metrics` reveals aggregate numbers only: how many jobs of each task were submitted, refused for capacity and ended in each terminal status, how many are queued and running, the pool's size and readiness, the concurrency and capacity settings, the backend name, and histograms of job and queue durations per task. It never reveals a job ID, idempotency key, input, result, error message, model path or the token; its label values come from fixed sets, never from a request. Counts and timings still describe the traffic (how busy the service is, which tasks run, how long they take), so the endpoint is behind the same service token as every route: scrape it from inside the private network with that token, never through a public route. How often it is scraped is the operator's concern; a scrape copies a few counters under the store's lock and formats them outside it. - Access logs are disabled to avoid accidentally recording identifiers and inputs. Internal exceptions become generic task errors. This is not a full observability or compliance solution. - Cancel/timeout terminates the task process, but cannot roll back previously completed external effects. No task in the prototype modifies an external system. - The default health endpoint proves HTTP service availability, not model integrity, ability to spawn processes or successful inference. An ONNX worker checks before it starts that each model directory holds its tokenizer and model files and refuses to start otherwise; that proves presence, not integrity, which the fetcher's hashes establish. Use a controlled model smoke test for readiness validation. diff --git a/tests/integration.py b/tests/integration.py index c55d007..bac2b66 100644 --- a/tests/integration.py +++ b/tests/integration.py @@ -8,6 +8,7 @@ from urllib.request import Request, urlopen BASE = os.environ.get("FRANKENPHP_URL", "http://frankenphp:8080") +WORKER = os.environ.get("WORKER_URL", "http://worker:8090") def request(path, body=None): @@ -17,6 +18,20 @@ def request(path, body=None): return response.status, json.load(response), int(response.headers["X-Demo-Worker-Requests"]) +def metrics(token=None): + """The worker's counters as {(name, labels): value}; the format itself is checked by the unit tests.""" + headers = {"Authorization": "Bearer " + (token if token is not None else os.environ["BRIDGE_TOKEN"])} + with urlopen(Request(WORKER + "/metrics", headers=headers), timeout=5) as response: + assert response.headers["Content-Type"] == "text/plain; version=0.0.4; charset=utf-8" + text = response.read().decode("utf-8") + samples = {} + for line in text.splitlines(): + if not line.startswith("#"): + series, value = line.rsplit(" ", 1) + samples[series] = float(value) + return text, samples + + def wait_ready(): deadline = time.monotonic() + 30 while time.monotonic() < deadline: @@ -56,6 +71,7 @@ def scenario(index): def main(): wait_ready() + _, counted = metrics() _, _, before = request("/healthz") with concurrent.futures.ThreadPoolExecutor(max_workers=6) as pool: ids = list(pool.map(scenario, range(12))) @@ -132,7 +148,21 @@ def main(): assert error.code == 400, (body, error.code) else: raise AssertionError("Out-of-contract redact input was accepted") - print("PASS: 12 concurrent isolated jobs through a reused FrankenPHP worker; top_k narrowing; embedding; redaction; invalid and missing-job cases") + # The worker counted those jobs, and its metrics hold numbers only, none of the documents. + try: + metrics("not-the-token") + except HTTPError as error: + assert error.code == 401, error.code + else: + raise AssertionError("Metrics were served without the token") + text, moved = metrics() + for task, least in [("rerank", 13), ("embed", 1), ("redact", 1)]: + for name in [f'bridge_jobs_total{{task="{task}",status="succeeded"}}', f'bridge_jobs_submitted_total{{task="{task}"}}', + f'bridge_job_duration_seconds_count{{task="{task}"}}']: + assert moved[name] - counted[name] >= least, (name, counted[name], moved[name]) + for document in ["unique0 token0", "needle token", "x@y.io", "4111 1111 1111 1111", "red apple", "café"]: + assert document not in text, document + print("PASS: 12 concurrent isolated jobs through a reused FrankenPHP worker; top_k narrowing; embedding; redaction; invalid and missing-job cases; worker metrics") if __name__ == "__main__": diff --git a/tests/python/test_jobs.py b/tests/python/test_jobs.py index 6bf7ab2..328dc66 100644 --- a/tests/python/test_jobs.py +++ b/tests/python/test_jobs.py @@ -1616,5 +1616,176 @@ def test_a_running_job_is_not_held_up_by_another_being_cancelled(self): self.assertEqual(self.wait(running["id"])["result"]["value"], "kept") +class MetricsSnapshotTests(unittest.TestCase): + """What the store counts for GET /metrics: plain numbers keyed only by task, status and reason.""" + + def setUp(self): + self.store = JobStore(Settings(concurrency=1, test_tasks=True)) + + def tearDown(self): + self.store.close() + + def wait(self, job_id, states=TERMINAL, timeout=5): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + job = self.store.get(job_id) + if job["status"] in states: + return job + time.sleep(0.01) + self.fail(f"job did not reach {states}") + + def wait_pool(self, size, ready, timeout=10): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + snapshot = self.store.snapshot() + if (snapshot["pool_size"], snapshot["pool_ready"]) == (size, ready): + return snapshot + time.sleep(0.01) + self.fail(f"the pool never reached {size} children, {ready} ready") + + def test_a_fresh_store_reports_every_series_at_zero(self): + snapshot = self.wait_pool(1, 1) + tasks = ("rerank", "embed", "redact", "test.delay", "test.crash", "test.error", "test.inspect") + self.assertEqual(snapshot["tasks"], tasks) + self.assertEqual(snapshot["jobs"], {(task, status): 0 for task in tasks + for status in ("succeeded", "failed", "timed_out", "cancelled")}) + self.assertEqual(snapshot["submitted"], dict.fromkeys(tasks, 0)) + self.assertEqual(snapshot["rejected"], {"capacity": 0}) + self.assertEqual((snapshot["queued"], snapshot["running"]), (0, 0)) + self.assertEqual((snapshot["backend"], snapshot["concurrency"], snapshot["capacity"]), ("lexical", 1, 128)) + for histogram in list(snapshot["duration"].values()) + list(snapshot["queue"].values()): + self.assertEqual(histogram, {"buckets": [0] * 12, "sum": 0.0, "count": 0}) + + def test_the_test_tasks_are_labels_only_where_they_exist(self): + store = JobStore(Settings(concurrency=1)) + try: + snapshot = store.snapshot() + finally: + store.close() + self.assertEqual(snapshot["tasks"], ("rerank", "embed", "redact")) + self.assertEqual(set(snapshot["duration"]), {"rerank", "embed", "redact"}) + + def test_each_terminal_status_is_counted_once(self): + succeeded = self.wait(self.store.submit("test.delay", {"seconds": 0})["id"]) + failed = self.wait(self.store.submit("test.error", {})["id"]) + cancelled = self.store.submit("test.delay", {"seconds": 5}) + self.wait(cancelled["id"], {"running"}) + self.store.cancel(cancelled["id"]) + self.store.cancel(cancelled["id"]) # a repeat changes nothing, so it counts nothing + timed_out = self.wait(self.store.submit("test.delay", {"seconds": 5}, 150)["id"]) + crashed = self.wait(self.store.submit("test.crash", {})["id"]) + self.assertEqual([job["status"] for job in (succeeded, failed, timed_out, crashed)], + ["succeeded", "failed", "timed_out", "failed"]) + snapshot = self.store.snapshot() + self.assertEqual({key: count for key, count in snapshot["jobs"].items() if count}, { + ("test.delay", "succeeded"): 1, ("test.error", "failed"): 1, ("test.delay", "cancelled"): 1, + ("test.delay", "timed_out"): 1, ("test.crash", "failed"): 1}) + self.assertEqual({task: count for task, count in snapshot["submitted"].items() if count}, + {"test.delay": 3, "test.error": 1, "test.crash": 1}) + + def test_histograms_agree_with_the_jobs_run(self): + started = time.monotonic() + self.wait(self.store.submit("test.delay", {"seconds": 0})["id"]) + self.wait(self.store.submit("test.delay", {"seconds": 0.3})["id"]) + # Ends at its deadline while queued behind nothing but a long job: never started, so no queue time. + first = self.store.submit("test.delay", {"seconds": 5}, 150) + never_started = self.store.submit("test.delay", {"seconds": 0}, 100) + self.wait(first["id"]) + self.assertEqual(self.wait(never_started["id"])["status"], "timed_out") + elapsed = time.monotonic() - started + snapshot = self.store.snapshot() + duration, queue = snapshot["duration"]["test.delay"], snapshot["queue"]["test.delay"] + self.assertEqual(duration["count"], 4) + self.assertEqual(sum(duration["buckets"]), 4) + # 0.3 s of sleep and a 0.15 s deadline are lower bounds; the wall clock of the test bounds the sum. + self.assertGreaterEqual(duration["sum"], 0.3 + 0.15 + 0.1) + self.assertLessEqual(duration["sum"], 4 * elapsed) + self.assertLessEqual(duration["buckets"][0] + duration["buckets"][1], 2) # le=0.1 cannot hold the 0.3 s job + self.assertEqual(queue["count"], 3) + self.assertEqual(sum(queue["buckets"]), 3) + self.assertGreaterEqual(queue["sum"], 0.0) + self.assertLessEqual(queue["sum"], duration["sum"]) + self.assertEqual(snapshot["queue"]["rerank"]["count"], 0) + + def test_queue_depth_while_a_job_waits_for_the_single_child(self): + running = self.store.submit("test.delay", {"seconds": 5}) + self.wait(running["id"], {"running"}) + waiting = self.store.submit("test.delay", {"seconds": 0}) + snapshot = self.store.snapshot() + self.assertEqual((snapshot["queued"], snapshot["running"]), (1, 1)) + self.assertEqual((snapshot["pool_size"], snapshot["pool_ready"]), (1, 1)) + self.store.cancel(waiting["id"]) + self.store.cancel(running["id"]) + snapshot = self.store.snapshot() + self.assertEqual((snapshot["queued"], snapshot["running"]), (0, 0)) + + def test_pool_gauges_follow_a_crash_and_the_replacement(self): + self.wait_pool(1, 1) + with patch.object(self.store._context, "Pipe", side_effect=OSError(24, "Too many open files")): + self.wait(self.store.submit("test.crash", {})["id"]) + # The crashed child is gone and no replacement can start yet. + self.wait_pool(0, 0) + self.wait_pool(1, 1) + self.assertEqual(self.wait(self.store.submit("test.delay", {"seconds": 0})["id"])["status"], "succeeded") + + def test_a_capacity_refusal_is_counted_and_a_repeat_or_invalid_input_is_not(self): + self.store.close() + self.store = JobStore(Settings(concurrency=1, capacity=1, test_tasks=True)) + self.store.submit("test.delay", {"seconds": 5}, idempotency_key="only") + self.store.submit("test.delay", {"seconds": 5}, idempotency_key="only") + for _ in range(2): + with self.assertRaises(CapacityError): + self.store.submit("test.delay", {"seconds": 0}) + for bad in ({"seconds": -1}, {"unknown": 1}): + with self.assertRaises(ValueError): + self.store.submit("test.delay", bad) + with self.assertRaises(ValueError): + self.store.submit("rerank", {"query": "", "documents": ["x"]}) + snapshot = self.store.snapshot() + self.assertEqual(snapshot["rejected"], {"capacity": 2}) + self.assertEqual(snapshot["submitted"]["test.delay"], 1) + self.assertEqual(snapshot["submitted"]["rerank"], 0) + + def test_a_submission_the_table_refuses_counts_as_a_refusal(self): + FailingDatabase(self.store, "INSERT") + with self.assertRaises(CapacityError): + self.store.submit("test.delay", {"seconds": 0}) + snapshot = self.store.snapshot() + self.assertEqual(snapshot["rejected"], {"capacity": 1}) + self.assertEqual(snapshot["submitted"]["test.delay"], 0) + + def test_the_snapshot_is_a_copy(self): + self.wait(self.store.submit("test.delay", {"seconds": 0})["id"]) + snapshot = self.store.snapshot() + snapshot["duration"]["test.delay"]["buckets"][0] = 99 + snapshot["jobs"][("test.delay", "succeeded")] = 99 + again = self.store.snapshot() + self.assertNotEqual(again["duration"]["test.delay"]["buckets"][0], 99) + self.assertEqual(again["jobs"][("test.delay", "succeeded")], 1) + + def test_jobs_restored_from_the_table_count_nothing(self): + with tempfile.TemporaryDirectory() as state_dir: + first = JobStore(Settings(concurrency=1, test_tasks=True, state_dir=state_dir)) + try: + job = first.submit("test.delay", {"seconds": 0}) + self.wait_for_store(first, job["id"]) + finally: + first.close() + second = JobStore(Settings(concurrency=1, test_tasks=True, state_dir=state_dir)) + try: + self.assertEqual(second.get(job["id"])["status"], "succeeded") + snapshot = second.snapshot() + finally: + second.close() + self.assertEqual(sum(snapshot["jobs"].values()), 0) + self.assertEqual(snapshot["duration"]["test.delay"]["count"], 0) + + def wait_for_store(self, store, job_id, timeout=5): + deadline = time.monotonic() + timeout + while store.get(job_id)["status"] not in TERMINAL: + self.assertLess(time.monotonic(), deadline) + time.sleep(0.01) + + if __name__ == "__main__": unittest.main() diff --git a/tests/python/test_server.py b/tests/python/test_server.py index 9b3a628..dbda6d5 100644 --- a/tests/python/test_server.py +++ b/tests/python/test_server.py @@ -1,15 +1,17 @@ import http.client import json import os +import re import sqlite3 import sys import tempfile import threading import time import unittest +import uuid from unittest.mock import MagicMock, patch -from ai_bridge import server, tasks +from ai_bridge import metrics, server, tasks from ai_bridge.jobs import Settings from ai_bridge.server import BridgeServer @@ -169,6 +171,22 @@ def test_a_refusal_for_capacity_says_when_to_retry_in_whole_seconds(self): self.assertRegex(retry_after, r"\A[1-9][0-9]*\Z") self.assertIn(int(retry_after), (299, 300)) + def test_a_refusal_for_capacity_is_counted_once_per_429(self): + def rejected(): + connection = http.client.HTTPConnection(*self.server.server_address, timeout=2) + try: + connection.request("GET", "/metrics", headers={"Authorization": "Bearer " + TOKEN}) + return sample(parse_metrics(connection.getresponse().read().decode()), + "bridge_submissions_rejected_total", reason="capacity") + finally: + connection.close() + + self.assertEqual(self.submit()[0], 202) + self.assertEqual(rejected(), 0) + for count in (1, 2): + self.assertEqual(self.submit()[0], 429) + self.assertEqual(rejected(), count) + def test_a_shutting_down_store_says_one_second(self): self.server.jobs.close() status, retry_after, _ = self.submit() @@ -299,3 +317,309 @@ def test_an_out_of_range_job_limit_is_refused_by_the_store(self): for value in ("0", "1000001", "-1"): with self.subTest(value=value), self.assertRaisesRegex(ValueError, "invalid worker limits"): self.main({"BRIDGE_WORKER_MAX_JOBS": value}, real_server=True) + + +SAMPLE = re.compile(r"([a-zA-Z_:][a-zA-Z0-9_:]*)(?:\{(.*)\})? (\S+)") +LABEL = r'([a-zA-Z_][a-zA-Z0-9_]*)="((?:[^"\\\n]|\\[\\"n])*)"' +NUMBER = re.compile(r"[+-]?(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+)(?:[eE][+-]?[0-9]+)?|\+Inf") + + +def parse_metrics(text): + """A strict reading of the Prometheus text format, version 0.0.4: {family: (type, {(name, labels): value})}. + Raises ValueError for anything a scraper could misread.""" + if not text.endswith("\n"): + raise ValueError("the body must end with a newline") + families, seen, current = {}, set(), None + for line in text[:-1].split("\n"): + if line.startswith("# HELP "): + if len(line.split(" ", 3)) != 4: + raise ValueError(f"HELP without text: {line!r}") + continue + if line.startswith("# TYPE "): + parts = line.split(" ") + if len(parts) != 4 or parts[3] not in {"counter", "gauge", "histogram"}: + raise ValueError(f"malformed TYPE: {line!r}") + if parts[2] in families: + raise ValueError(f"a second TYPE for {parts[2]}") + current = parts[2] + families[current] = (parts[3], {}) + continue + match = SAMPLE.fullmatch(line) + if match is None or current is None: + raise ValueError(f"not a sample line under a TYPE: {line!r}") + name, labels_text, value = match.groups() + kind = families[current][0] + allowed = {current + suffix for suffix in ("_bucket", "_sum", "_count")} if kind == "histogram" else {current} + if name not in allowed: + raise ValueError(f"{name} outside its family {current}") + labels = {} + if labels_text is not None: + if not re.fullmatch(f"{LABEL}(?:,{LABEL})*", labels_text): + raise ValueError(f"malformed labels: {line!r}") + for key, label in re.findall(LABEL, labels_text): + if key in labels: + raise ValueError(f"label {key} twice: {line!r}") + labels[key] = re.sub(r"\\(.)", lambda m: "\n" if m[1] == "n" else m[1], label) + if not NUMBER.fullmatch(value): + raise ValueError(f"malformed value: {line!r}") + key = (name, tuple(sorted(labels.items()))) + if key in seen: + raise ValueError(f"duplicate series: {line!r}") + seen.add(key) + families[current][1][key] = float(value) + for family, (kind, samples) in families.items(): + if kind == "histogram": + _check_histogram(family, samples) + return families + + +def _check_histogram(family, samples): + series = {} + for (name, labels), value in samples.items(): + rest = tuple(item for item in labels if item[0] != "le") + entry = series.setdefault(rest, {"buckets": [], "sum": None, "count": None}) + if name.endswith("_bucket"): + bound = dict(labels).get("le") + if bound is None or not NUMBER.fullmatch(bound): + raise ValueError(f"a bucket of {family} without a numeric le") + entry["buckets"].append((float(bound), value)) + else: + entry[name.rsplit("_", 1)[1]] = value + for rest, entry in series.items(): + counts = [count for _, count in sorted(entry["buckets"])] + if not entry["buckets"] or max(entry["buckets"])[0] != float("inf"): + raise ValueError(f"{family}{rest} has no +Inf bucket") + if counts != sorted(counts): + raise ValueError(f"{family}{rest} buckets are not cumulative") + if entry["sum"] is None or entry["count"] is None or counts[-1] != entry["count"]: + raise ValueError(f"{family}{rest} +Inf bucket differs from _count, or a total is missing") + + +def sample(families, name, **labels): + family = next(family for family in families if name == family or name.startswith(family + "_")) + return families[family][1][(name, tuple(sorted(labels.items())))] + + +class MetricsParserTests(unittest.TestCase): + """The parser the endpoint's tests rely on refuses what a scraper would misread.""" + + GOOD = ("# HELP h_seconds Durations.\n# TYPE h_seconds histogram\n" + 'h_seconds_bucket{task="a",le="1"} 1\nh_seconds_bucket{task="a",le="+Inf"} 2\n' + 'h_seconds_sum{task="a"} 3.5\nh_seconds_count{task="a"} 2\n' + '# TYPE c_total counter\nc_total{x="q\\"\\\\\\n"} 4\n') + + def test_a_well_formed_body_parses(self): + families = parse_metrics(self.GOOD) + self.assertEqual(sample(families, "h_seconds_count", task="a"), 2) + self.assertEqual(sample(families, "c_total", x='q"\\\n'), 4) + + def test_malformed_bodies_are_refused(self): + for broken in [ + self.GOOD.rstrip("\n"), + self.GOOD + "# TYPE c_total counter\n", + self.GOOD + 'c_total{x="q\\"\\\\\\n"} 5\n', + self.GOOD.replace('le="+Inf"} 2', 'le="9"} 2'), + self.GOOD.replace('le="1"} 1', 'le="1"} 3'), + self.GOOD.replace('h_seconds_count{task="a"} 2', 'h_seconds_count{task="a"} 5'), + self.GOOD.replace('h_seconds_sum{task="a"} 3.5\n', ""), + self.GOOD.replace("c_total{", "other_total{"), + "x 1\n", self.GOOD + "c_total 1 2\n", self.GOOD + "c_total NaN\n", + self.GOOD + 'c_total{x="a",x="b"} 1\n', self.GOOD + "c_total{x=unquoted} 1\n", + self.GOOD + "# TYPE s summary\n"]: + with self.subTest(broken=broken[-40:]), self.assertRaises(ValueError): + parse_metrics(broken) + + +class MetricsRenderTests(unittest.TestCase): + def test_label_values_are_escaped(self): + self.assertEqual(metrics.escape('a\\b"c\nd'), 'a\\\\b\\"c\\nd') + families = parse_metrics(metrics.gauge("g", "A gauge.", [({"v": 'a\\b"c\nd'}, 1)])) + self.assertEqual(sample(families, "g", v='a\\b"c\nd'), 1) + + def test_values_are_written_exactly(self): + self.assertEqual([metrics.number(value) for value in (0, 7, 0.1, 2.5, 1e-07, float("inf"))], + ["0", "7", "0.1", "2.5", "1e-07", "+Inf"]) + + def test_a_histogram_is_cumulative_with_its_totals(self): + text = metrics.histogram("h_seconds", "Durations.", [({"task": "a"}, { + "buckets": [1, 0, 2] + [0] * 8 + [1], "sum": 400.5, "count": 4})]) + families = parse_metrics(text) + self.assertEqual(sample(families, "h_seconds_bucket", task="a", le="0.05"), 1) + self.assertEqual(sample(families, "h_seconds_bucket", task="a", le="0.25"), 3) + self.assertEqual(sample(families, "h_seconds_bucket", task="a", le="300"), 3) + self.assertEqual(sample(families, "h_seconds_bucket", task="a", le="+Inf"), 4) + self.assertEqual(sample(families, "h_seconds_sum", task="a"), 400.5) + + +class MetricsHttpTests(unittest.TestCase): + """GET /metrics on a store with the test tasks, so failures and timeouts can be produced.""" + + @classmethod + def setUpClass(cls): + cls.server = BridgeServer(("127.0.0.1", 0), TOKEN, Settings(concurrency=1, capacity=16, test_tasks=True)) + cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True) + cls.thread.start() + + @classmethod + def tearDownClass(cls): + cls.server.shutdown() + cls.thread.join() + cls.server.server_close() + + def raw(self, method, path, body=None, headers=None): + connection = http.client.HTTPConnection(*self.server.server_address, timeout=2) + try: + connection.request(method, path, body, headers or {}) + response = connection.getresponse() + return response.status, dict(response.getheaders()), response.read() + finally: + connection.close() + + def auth(self): + return {"Authorization": "Bearer " + TOKEN, "Content-Type": "application/json"} + + def scrape(self): + status, headers, body = self.raw("GET", "/metrics", headers=self.auth()) + self.assertEqual(status, 200) + self.assertEqual(headers["Content-Type"], "text/plain; version=0.0.4; charset=utf-8") + self.assertEqual(headers["Content-Length"], str(len(body))) + self.assertEqual(headers["Cache-Control"], "no-store") + self.assertEqual(headers["Connection"], "close") + return body.decode("utf-8"), parse_metrics(body.decode("utf-8")) + + def submit(self, task, value, timeout_ms=30000, key=None): + body = {"task": task, "input": value, "timeout_ms": timeout_ms} + if key is not None: + body["idempotency_key"] = key + status, _, data = self.raw("POST", "/v1/jobs", json.dumps(body), self.auth()) + self.assertEqual(status, 202, data) + return json.loads(data) + + def wait(self, job_id, states=("succeeded", "failed", "timed_out", "cancelled"), timeout=5): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + job = self.server.jobs.get(job_id) + if job["status"] in states: + return job + time.sleep(0.01) + self.fail(f"job did not reach {states}") + + def test_the_token_is_required_exactly_as_on_every_route(self): + for index, headers in enumerate([{}, {"Authorization": "Bearer wrong"}, + {"Authorization": "Bearer " + TOKEN + "x"}, + {"Authorization": "Basic " + TOKEN}]): + with self.subTest(case=index): + status, _, body = self.raw("GET", "/metrics", headers=headers) + self.assertEqual(status, 401) + self.assertEqual(json.loads(body)["error"]["code"], "unauthorized") + connection = http.client.HTTPConnection(*self.server.server_address, timeout=2) + try: + connection.putrequest("GET", "/metrics") + connection.putheader("Authorization", "Bearer " + TOKEN) + connection.putheader("Authorization", "Bearer " + TOKEN) + connection.endheaders() + self.assertEqual(connection.getresponse().status, 401) + finally: + connection.close() + + def test_other_methods_and_spellings_are_unknown_routes(self): + self.assertEqual(self.raw("POST", "/metrics", "{}", self.auth())[0], 404) + self.assertEqual(self.raw("POST", "/unknown", "{}", self.auth())[0], 404) + for path in ["/metrics?x=1", "/metrics/", "/Metrics", "/v1/metrics"]: + with self.subTest(path=path): + self.assertEqual(self.raw("GET", path, headers=self.auth())[0], 404) + + def test_the_body_parses_and_names_every_family(self): + _, families = self.scrape() + self.assertEqual({family: kind for family, (kind, _) in families.items()}, { + "bridge_jobs_total": "counter", "bridge_jobs_submitted_total": "counter", + "bridge_submissions_rejected_total": "counter", "bridge_queue_depth": "gauge", + "bridge_jobs_running": "gauge", "bridge_pool_size": "gauge", "bridge_pool_ready": "gauge", + "bridge_concurrency": "gauge", "bridge_capacity": "gauge", + "bridge_job_duration_seconds": "histogram", "bridge_job_queue_seconds": "histogram", + "bridge_info": "gauge"}) + self.assertEqual(sample(families, "bridge_info", protocol="1", backend="lexical"), 1) + self.assertEqual(sample(families, "bridge_concurrency"), 1) + self.assertEqual(sample(families, "bridge_capacity"), 16) + self.assertIn(sample(families, "bridge_submissions_rejected_total", reason="capacity"), (0, 1, 2)) + bounds = sorted({dict(labels)["le"] for name, labels in families["bridge_job_duration_seconds"][1] + if name.endswith("_bucket")}, key=float) + self.assertEqual(bounds, ["0.05", "0.1", "0.25", "0.5", "1", "2.5", "5", "10", "30", "60", "300", "+Inf"]) + statuses = {dict(labels)["status"] for _, labels in families["bridge_jobs_total"][1]} + self.assertEqual(statuses, {"succeeded", "failed", "timed_out", "cancelled"}) + tasks = {dict(labels)["task"] for _, labels in families["bridge_jobs_total"][1]} + self.assertEqual(tasks, {"rerank", "embed", "redact", "test.delay", "test.crash", "test.error", "test.inspect"}) + + def test_counters_and_histograms_move_with_the_jobs_and_nothing_of_a_job_is_exposed(self): + _, before = self.scrape() + marker = "metrics-marker-" + uuid.uuid4().hex + key = "metrics-key-" + uuid.uuid4().hex + succeeded = self.submit("test.delay", {"seconds": 0, "value": marker}, key=key) + failed = self.submit("test.error", {"value": marker}) + ids = [succeeded["id"], failed["id"]] + for job_id in ids: + self.wait(job_id) + cancelled = self.submit("test.delay", {"seconds": 5, "value": marker}) + self.wait(cancelled["id"], {"running"}) + self.raw("POST", "/v1/jobs/" + cancelled["id"] + "/cancel", "{}", self.auth()) + # The cancelled job's child is replaced; the next one starts only on the fresh child. + deadline = time.monotonic() + 10 + while self.server.jobs.snapshot()["pool_ready"] != 1: + self.assertLess(time.monotonic(), deadline, "no fresh child") + time.sleep(0.01) + timed_out = self.submit("test.delay", {"seconds": 5, "value": marker}, 150) + ids += [cancelled["id"], timed_out["id"]] + self.assertEqual(self.wait(timed_out["id"])["status"], "timed_out") + text, after = self.scrape() + + def moved(name, **labels): + return sample(after, name, **labels) - sample(before, name, **labels) + + self.assertEqual(moved("bridge_jobs_total", task="test.delay", status="succeeded"), 1) + self.assertEqual(moved("bridge_jobs_total", task="test.error", status="failed"), 1) + self.assertEqual(moved("bridge_jobs_total", task="test.delay", status="cancelled"), 1) + self.assertEqual(moved("bridge_jobs_total", task="test.delay", status="timed_out"), 1) + self.assertEqual(moved("bridge_jobs_submitted_total", task="test.delay"), 3) + self.assertEqual(moved("bridge_jobs_submitted_total", task="test.error"), 1) + self.assertEqual(moved("bridge_job_duration_seconds_count", task="test.delay"), 3) + self.assertEqual(moved("bridge_job_duration_seconds_count", task="test.error"), 1) + self.assertGreaterEqual(moved("bridge_job_duration_seconds_sum", task="test.delay"), 0.15) + self.assertEqual(moved("bridge_job_duration_seconds_bucket", task="test.delay", le="+Inf"), 3) + self.assertEqual(moved("bridge_job_queue_seconds_count", task="test.delay"), 3) + for secret in ids + [marker, key, TOKEN, "This internal exception", "task_failed", "deadline_exceeded", + self.server.jobs.settings.model_dir]: + with self.subTest(secret=secret[:16]): + self.assertNotIn(secret, text) + + def test_a_scrape_while_a_job_runs_answers_at_once(self): + job = self.submit("test.delay", {"seconds": 2}) + self.wait(job["id"], {"running"}) + started = time.monotonic() + _, families = self.scrape() + # Well inside the 2 s request timeout, while the job still runs. + self.assertLess(time.monotonic() - started, 0.5) + self.assertEqual(sample(families, "bridge_jobs_running"), 1) + self.assertEqual(sample(families, "bridge_pool_size"), 1) + self.raw("POST", "/v1/jobs/" + job["id"] + "/cancel", "{}", self.auth()) + self.wait(job["id"]) + + def test_the_body_is_formatted_outside_the_store_lock(self): + held = [] + render = metrics.render + + def checking(snapshot): + # Another thread can take the lock here, so the scrape held it only while it copied the numbers. + def probe(): + taken = self.server.jobs._lock.acquire(blocking=False) + if taken: + self.server.jobs._lock.release() + held.append(not taken) + + thread = threading.Thread(target=probe) + thread.start() + thread.join() + return render(snapshot) + + with patch.object(server, "render", checking): + self.scrape() + self.assertEqual(held, [False]) diff --git a/worker/ai_bridge/jobs.py b/worker/ai_bridge/jobs.py index dbb62f1..7195c90 100644 --- a/worker/ai_bridge/jobs.py +++ b/worker/ai_bridge/jobs.py @@ -20,9 +20,11 @@ import uuid from dataclasses import dataclass -from .tasks import MODEL_NAMES, execute, validate +from . import metrics +from .tasks import MODEL_NAMES, TEST_TASKS, execute, validate -TERMINAL = frozenset({"succeeded", "failed", "timed_out", "cancelled"}) +STATUSES = ("succeeded", "failed", "timed_out", "cancelled") +TERMINAL = frozenset(STATUSES) # How long a retired child gets to exit on its own, after terminate() or after being told to stop, before kill(). GRACE_SECONDS = 0.2 # How long a started child gets to say it is ready before it is given up on; and how long the store waits before @@ -137,6 +139,14 @@ def __init__(self, settings=None): or self.settings.retention_seconds <= 0): raise ValueError("invalid worker settings") self._context = multiprocessing.get_context("spawn") + # Aggregates for GET /metrics, since this process started; every series exists from the start, at zero. + tasks = tuple(MODEL_NAMES[self.settings.backend]) + (TEST_TASKS if self.settings.test_tasks else ()) + self._tasks = tasks + self._ended = {(task, status): 0 for task in tasks for status in STATUSES} + self._submitted = dict.fromkeys(tasks, 0) + self._rejected = {"capacity": 0} + self._durations = {task: metrics.new_histogram() for task in tasks} + self._waits = {task: metrics.new_histogram() for task in tasks} self._jobs = {} self._keys = {} # idempotency key -> job id, for retained jobs only self._active = {} # job id -> the worker it was sent to @@ -226,27 +236,36 @@ def submit(self, task, payload, timeout_ms=30000, idempotency_key=None): raise ValueError("idempotency_key must be 1-128 printable ASCII characters") payload = validate(task, payload, self.settings.test_tasks) with self._lock: - if self._stopping.is_set(): - raise CapacityError("worker is shutting down") - self._expire() - # A repeat adds no job, so it is answered even at capacity. - if idempotency_key in self._keys: - return self._public(self._jobs[self._keys[idempotency_key]]) - if len(self._jobs) >= self.settings.capacity: - raise CapacityError("job capacity reached; wait for retained jobs to expire", self._retry_after()) - now = time.monotonic() - job_id = uuid.uuid4().hex - job = { - "id": job_id, "task": task, "status": "queued", "result": None, - "error": None, "progress": None, "timeout_ms": timeout_ms, - "_input": copy.deepcopy(payload), "_deadline": now + timeout_ms / 1000, - "_finished": None, "_retried": False, "_key": idempotency_key, - } - self._write(job, submitted=True) - self._jobs[job_id] = job - if idempotency_key is not None: - self._keys[idempotency_key] = job_id - return self._public(job) + try: + return self._accept(task, payload, timeout_ms, idempotency_key) + except CapacityError: + # Every 429 the service sends comes from here, so each is counted exactly once. + self._rejected["capacity"] += 1 + raise + + def _accept(self, task, payload, timeout_ms, idempotency_key): + if self._stopping.is_set(): + raise CapacityError("worker is shutting down") + self._expire() + # A repeat adds no job, so it is answered even at capacity. + if idempotency_key in self._keys: + return self._public(self._jobs[self._keys[idempotency_key]]) + if len(self._jobs) >= self.settings.capacity: + raise CapacityError("job capacity reached; wait for retained jobs to expire", self._retry_after()) + now = time.monotonic() + job_id = uuid.uuid4().hex + job = { + "id": job_id, "task": task, "status": "queued", "result": None, + "error": None, "progress": None, "timeout_ms": timeout_ms, + "_input": copy.deepcopy(payload), "_deadline": now + timeout_ms / 1000, + "_finished": None, "_retried": False, "_key": idempotency_key, "_submitted_at": now, + } + self._write(job, submitted=True) + self._jobs[job_id] = job + if idempotency_key is not None: + self._keys[idempotency_key] = job_id + self._submitted[task] += 1 + return self._public(job) def get(self, job_id): with self._lock: @@ -263,6 +282,20 @@ def cancel(self, job_id): self._abort(job, "cancelled") return self._public(job) + def snapshot(self): + """The aggregates for GET /metrics as plain values, copied under the lock so the caller formats them without + holding it.""" + with self._lock: + statuses = [job["status"] for job in self._jobs.values()] + return { + "backend": self.settings.backend, "concurrency": self.settings.concurrency, + "capacity": self.settings.capacity, "tasks": self._tasks, "jobs": dict(self._ended), + "submitted": dict(self._submitted), "rejected": dict(self._rejected), + "queued": statuses.count("queued"), "running": statuses.count("running"), + "pool_size": len(self._pool), "pool_ready": sum(worker.ready for worker in self._pool), + "duration": copy.deepcopy(self._durations), "queue": copy.deepcopy(self._waits), + } + def _retry_after(self): """Seconds until the first retained job expires; a store holding only unfinished jobs gives a short default, since none of them has a known end.""" @@ -294,6 +327,9 @@ def _finish(self, job, status, result=None, code=None): job["_input"] = None job["_finished"] = time.monotonic() self._write(job) + # A job read back from the table finished in an earlier process and never ends here. + self._ended[job["task"], status] += 1 + metrics.observe(self._durations[job["task"]], job["_finished"] - job["_submitted_at"]) def _abort(self, job, status, code=None): """End a job before its child has: the child is retired, since it may be halfway through the input.""" @@ -474,6 +510,7 @@ def _drain_message(self, worker, message): if kind == "started": # Not written through: a restart fails a running job exactly as it fails a queued one. job["status"] = "running" + metrics.observe(self._waits[job["task"]], time.monotonic() - job["_submitted_at"]) return False if kind == "progress": job["progress"] = value diff --git a/worker/ai_bridge/metrics.py b/worker/ai_bridge/metrics.py new file mode 100644 index 0000000..8b5941a --- /dev/null +++ b/worker/ai_bridge/metrics.py @@ -0,0 +1,93 @@ +"""Aggregate counts and timings in the Prometheus text exposition format, version 0.0.4, with the standard library. +Label values come only from fixed sets (the task allowlist, the terminal statuses, the backend name), never from a +job, so nothing a client sent can reach the body and the number of series is bounded.""" + +import bisect +import math + +CONTENT_TYPE = "text/plain; version=0.0.4; charset=utf-8" +# Upper bounds in seconds; a job longer than the last falls only into +Inf. +BUCKETS = (0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60, 300) + + +def new_histogram(): + """Counts per bucket, not cumulative, the last for +Inf only; render() adds them up.""" + return {"buckets": [0] * (len(BUCKETS) + 1), "sum": 0.0, "count": 0} + + +def observe(histogram, seconds): + # bisect_left puts a value equal to a bound in that bound's bucket, as le (less or equal) means. + histogram["buckets"][bisect.bisect_left(BUCKETS, seconds)] += 1 + histogram["sum"] += seconds + histogram["count"] += 1 + + +def escape(value): + return value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n") + + +def number(value): + if isinstance(value, float): + return "+Inf" if value == math.inf else repr(value) + return str(value) + + +def _labels(labels): + if not labels: + return "" + return "{" + ",".join(f'{key}="{escape(value)}"' for key, value in labels.items()) + "}" + + +def _family(name, kind, help_text, lines): + return f"# HELP {name} {help_text}\n# TYPE {name} {kind}\n" + "".join(line + "\n" for line in lines) + + +def _simple(name, kind, help_text, samples): + return _family(name, kind, help_text, [f"{name}{_labels(labels)} {number(value)}" for labels, value in samples]) + + +def counter(name, help_text, samples): + return _simple(name, "counter", help_text, samples) + + +def gauge(name, help_text, samples): + return _simple(name, "gauge", help_text, samples) + + +def histogram(name, help_text, series): + lines = [] + for labels, values in series: + total = 0 + for bound, count in zip(BUCKETS + (math.inf,), values["buckets"]): + total += count + lines.append(f"{name}_bucket{_labels({**labels, 'le': number(bound)})} {total}") + lines.append(f"{name}_sum{_labels(labels)} {number(float(values['sum']))}") + lines.append(f"{name}_count{_labels(labels)} {values['count']}") + return _family(name, "histogram", help_text, lines) + + +def render(snapshot): + """The whole body from JobStore.snapshot(), which is a copy: nothing here holds the store's lock.""" + tasks = snapshot["tasks"] + return "".join([ + counter("bridge_jobs_total", "Jobs that reached a terminal status since the process started.", + [({"task": task, "status": status}, count) for (task, status), count in snapshot["jobs"].items()]), + counter("bridge_jobs_submitted_total", "Jobs accepted since the process started.", + [({"task": task}, snapshot["submitted"][task]) for task in tasks]), + counter("bridge_submissions_rejected_total", "Submissions refused with 429 since the process started.", + [({"reason": reason}, count) for reason, count in snapshot["rejected"].items()]), + gauge("bridge_queue_depth", "Jobs waiting for a child process.", [({}, snapshot["queued"])]), + gauge("bridge_jobs_running", "Jobs a child process has started.", [({}, snapshot["running"])]), + gauge("bridge_pool_size", "Child processes in the pool.", [({}, snapshot["pool_size"])]), + gauge("bridge_pool_ready", "Child processes in the pool that have said they are ready.", + [({}, snapshot["pool_ready"])]), + gauge("bridge_concurrency", "Configured number of child processes.", [({}, snapshot["concurrency"])]), + gauge("bridge_capacity", "Configured limit of queued, running and retained jobs.", + [({}, snapshot["capacity"])]), + histogram("bridge_job_duration_seconds", "Seconds from submission to terminal status.", + [({"task": task}, snapshot["duration"][task]) for task in tasks]), + histogram("bridge_job_queue_seconds", "Seconds from submission to a child starting the job.", + [({"task": task}, snapshot["queue"][task]) for task in tasks]), + gauge("bridge_info", "Protocol version and configured backend.", + [({"protocol": "1", "backend": snapshot["backend"]}, 1)]), + ]) diff --git a/worker/ai_bridge/server.py b/worker/ai_bridge/server.py index e3cc0b8..666edbf 100644 --- a/worker/ai_bridge/server.py +++ b/worker/ai_bridge/server.py @@ -12,6 +12,7 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from .jobs import CapacityError, JobStore, Settings +from .metrics import CONTENT_TYPE, render from .tasks import missing_models MAX_BODY = 262144 @@ -70,11 +71,17 @@ def log_message(self, *_): pass def _reply(self, status, value, headers=()): - body = json.dumps(value, allow_nan=False, separators=(",", ":")).encode() + self._send(status, json.dumps(value, allow_nan=False, separators=(",", ":")).encode(), "application/json", + headers) + + def _text(self, status, text): + self._send(status, text.encode("utf-8"), CONTENT_TYPE) + + def _send(self, status, body, content_type, headers=()): self.send_response(status) for name, header in headers: self.send_header(name, header) - self.send_header("Content-Type", "application/json") + self.send_header("Content-Type", content_type) self.send_header("Content-Length", str(len(body))) self.send_header("Cache-Control", "no-store") self.send_header("Connection", "close") @@ -128,6 +135,10 @@ def do_GET(self): if self.path == "/healthz": self._reply(200, {"status": "ok", "protocol": 1, "backend": self.server.jobs.settings.backend}) return + if self.path == "/metrics": + # The snapshot is taken under the store's lock; the formatting is not. + self._text(200, render(self.server.jobs.snapshot())) + return match = re.fullmatch(r"/v1/jobs/([a-f0-9]{32})", self.path) job = self.server.jobs.get(match[1]) if match else None if job is None: diff --git a/worker/ai_bridge/tasks.py b/worker/ai_bridge/tasks.py index 7f6e1a9..9b0c771 100644 --- a/worker/ai_bridge/tasks.py +++ b/worker/ai_bridge/tasks.py @@ -40,6 +40,10 @@ MODEL_DIRECTORIES = ("rerank", "embed", "redact", "redact-uncased") +# Present only with BRIDGE_TEST_TASKS=1, for the fault tests. +TEST_TASKS = ("test.delay", "test.crash", "test.error", "test.inspect") + + class InvalidInput(ValueError): pass @@ -65,7 +69,7 @@ def _strings(values: object, plural: str, singular: str, limit: int, extra: int def validate(task: str, payload: object, test_tasks: bool = False) -> dict: if not isinstance(payload, dict): raise InvalidInput("input must be an object") - if test_tasks and task in {"test.delay", "test.crash", "test.error", "test.inspect"}: + if test_tasks and task in TEST_TASKS: if set(payload) - {"seconds", "value"}: raise InvalidInput("unknown input field") seconds = payload.get("seconds", 0.05)