Coverage for src / quber / playground / gpu.py: 25%
122 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""The GPU stages of an upload, run on the RunPod worker.
3The hosted playground has no GPU. Two stages of an upload need one: the
4docling parse of the document, and the embedding of its chunks at
5ingestion. Both go to the RunPod serverless worker under the contract the
6inbound-drop Lambda already uses with it: the job payload carries presigned
7S3 URLs for everything the worker reads and writes, the worker holds no
8credentials, and a completion marker written strictly after the outputs says
9the outputs are whole. The files land in the document's own prefix in the
10extracts bucket, so a parse's artifacts and an embedding's vectors sit
11beside everything else the document owns.
13``configured`` says whether this host sends work there at all. On the
14developer host, where nothing names an endpoint, the parse runs in-process
15on the local GPU and the chunks embed locally, as they always have.
16"""
18from __future__ import annotations
20import json
21import time
22from pathlib import Path
23from typing import Any, Callable, Dict, List, Optional
24from urllib.parse import urlparse
26import boto3
27import httpx
28from botocore.config import Config
29from loguru import logger
31from quber.settings import get_settings
33URL_EXPIRY_SECONDS = 4 * 3600
34POLL_SECONDS = 5
35LOG_MIRROR_SECONDS = 10
36#: A job on a worker freshly placed on a new host pays a ten-gigabyte image
37#: pull before it runs: eleven minutes measured once, over twenty on another
38#: night. Either task can be the one that meets a fresh host.
39PARSE_DEADLINE_SECONDS = 35 * 60
40EMBED_DEADLINE_SECONDS = 35 * 60
42#: The worker's artifact names, each with the file the name lands in. The
43#: worker keys its upload URLs by the first and writes the second.
44PARSE_ARTIFACTS = {"document": "docling", "confidence": "confidence", "cells": "cells"}
46Log = Callable[[str], None]
49def configured() -> bool:
50 p = get_settings().playground
51 return bool(p.runpod_endpoint_id and p.runpod_api_key and p.artifacts_uri)
54def _bucket_and_prefix(doc_key: str) -> tuple[str, str]:
55 uri = get_settings().playground.artifacts_uri or ""
56 parsed = urlparse(uri)
57 return parsed.netloc, parsed.path.strip("/") + "/" + doc_key + "/"
60def _s3():
61 # Signature version 4, explicitly: the extracts bucket is KMS-encrypted
62 # and S3 refuses a version 2 presigned URL against it with a bare 400,
63 # which is what a client without a pinned region produces for us-east-1.
64 return boto3.client(
65 "s3",
66 region_name=get_settings().s3_cache.region or "us-east-1",
67 config=Config(signature_version="s3v4"),
68 )
71def _presign(method: str, bucket: str, key: str) -> str:
72 return _s3().generate_presigned_url(
73 method, Params={"Bucket": bucket, "Key": key}, ExpiresIn=URL_EXPIRY_SECONDS
74 )
77def _submit(payload: Dict[str, Any]) -> str:
78 p = get_settings().playground
79 response = httpx.post(
80 f"https://api.runpod.ai/v2/{p.runpod_endpoint_id}/run",
81 json={"input": payload},
82 headers={"Authorization": f"Bearer {p.runpod_api_key}"},
83 timeout=30,
84 )
85 response.raise_for_status()
86 job = response.json()
87 logger.info("runpod job {} ({})", job.get("id"), job.get("status"))
88 return job["id"]
91def _read_if_fresh(bucket: str, key: str, since: float) -> Optional[bytes]:
92 """The object's bytes, or None when it does not exist or predates ``since``.
94 Nothing under a document's prefix is ever deleted, so a previous run's
95 marker and log stay where they are; a run tells its own apart by the
96 object's write time against the moment it submitted the job."""
97 s3 = _s3()
98 try:
99 response = s3.get_object(Bucket=bucket, Key=key)
100 except s3.exceptions.NoSuchKey:
101 return None
102 if response["LastModified"].timestamp() < since:
103 return None
104 return response["Body"].read()
107def _read_json_if_fresh(bucket: str, key: str, since: float) -> Optional[Dict[str, Any]]:
108 body = _read_if_fresh(bucket, key, since)
109 return json.loads(body) if body is not None else None
112def _read_text_if_fresh(bucket: str, key: str, since: float) -> Optional[str]:
113 body = _read_if_fresh(bucket, key, since)
114 return body.decode("utf-8", "replace") if body is not None else None
117def _job_status(job_id: str) -> Dict[str, Any]:
118 p = get_settings().playground
119 response = httpx.get(
120 f"https://api.runpod.ai/v2/{p.runpod_endpoint_id}/status/{job_id}",
121 headers={"Authorization": f"Bearer {p.runpod_api_key}"},
122 timeout=30,
123 )
124 response.raise_for_status()
125 return response.json()
128def _wait_for_marker(
129 bucket: str, marker_key: str, log_key: Optional[str], deadline: float, log: Log, job_id: str, since: float
130) -> Dict[str, Any]:
131 """Poll until the marker exists, mirroring the worker's log into ``log``
132 as it grows. The marker is the truth about the outputs; the job's own
133 status is watched beside it so a job that ends without writing one, or
134 that RunPod fails or cancels, is reported at once rather than at the
135 deadline."""
136 started = time.time()
137 mirrored = 0
138 last_mirror = 0.0
139 ended_at: Optional[float] = None
140 while True:
141 marker = _read_json_if_fresh(bucket, marker_key, since)
142 if marker is not None:
143 if log_key:
144 text = _read_text_if_fresh(bucket, log_key, since) or ""
145 for line in text.splitlines()[mirrored:]:
146 log(line)
147 if marker.get("status") != "complete":
148 raise RuntimeError(
149 f"worker failed at {marker.get('stage', '?')}: {marker.get('error', 'no error recorded')}"
150 )
151 return marker
152 elapsed = time.time() - started
153 if elapsed > deadline:
154 raise TimeoutError(f"no completion marker after {int(elapsed)}s ({marker_key})")
155 status = _job_status(job_id)
156 state = status.get("status")
157 if state in ("FAILED", "CANCELLED", "TIMED_OUT"):
158 raise RuntimeError(
159 f"worker job {job_id} {state}: {json.dumps(status.get('error') or status.get('output'))[:800]}"
160 )
161 if state == "COMPLETED":
162 # The marker is written before the handler returns, so a completed
163 # job with no marker a moment later returned without doing the work.
164 ended_at = ended_at or time.time()
165 if time.time() - ended_at > 3 * POLL_SECONDS:
166 raise RuntimeError(
167 f"worker job {job_id} completed without a completion marker; output: "
168 f"{json.dumps(status.get('output'))[:800]}"
169 )
170 if log_key and time.time() - last_mirror >= LOG_MIRROR_SECONDS:
171 last_mirror = time.time()
172 text = _read_text_if_fresh(bucket, log_key, since)
173 if text:
174 lines = text.splitlines()
175 for line in lines[mirrored:]:
176 log(line)
177 mirrored = len(lines)
178 time.sleep(POLL_SECONDS)
181def parse(doc_key: str, workdir: Path, log: Log) -> None:
182 """Parse the document on the worker and bring the three parse artifacts
183 into ``workdir`` for fusion. The PDF must already be in the document's
184 prefix as ``<doc_key>.pdf``."""
185 bucket, prefix = _bucket_and_prefix(doc_key)
186 marker_key = f"{prefix}{doc_key}.parse.complete.json"
187 log_key = f"{prefix}{doc_key}.parse.log"
188 s3 = _s3()
189 # Clock skew between this host and S3 is absorbed by a minute's grace.
190 since = time.time() - 60
191 payload = {
192 "base": doc_key,
193 "pdf_url": _presign("get_object", bucket, f"{prefix}{doc_key}.pdf"),
194 "artifact_urls": {
195 name: _presign("put_object", bucket, f"{prefix}{doc_key}.{stem}.json")
196 for name, stem in PARSE_ARTIFACTS.items()
197 },
198 "marker_url": _presign("put_object", bucket, marker_key),
199 "log_url": _presign("put_object", bucket, log_key),
200 }
201 job_id = _submit(payload)
202 log(f"parse submitted to the GPU worker as job {job_id}")
203 marker = _wait_for_marker(bucket, marker_key, log_key, PARSE_DEADLINE_SECONDS, log, job_id, since)
204 log(
205 f"parse complete: {marker.get('pages')} pages, {marker.get('tables')} tables, {marker.get('total_seconds')}s"
206 )
207 for stem in PARSE_ARTIFACTS.values():
208 s3.download_file(bucket, f"{prefix}{doc_key}.{stem}.json", str(workdir / f"{doc_key}.{stem}.json"))
211def embed(doc_key: str, texts: List[str], log: Log = logger.info) -> List[List[float]]:
212 """Embed the chunks on the worker. The texts go up as one file, the
213 vectors come back as one file, both kept in the document's prefix."""
214 bucket, prefix = _bucket_and_prefix(doc_key)
215 chunks_key = f"{prefix}{doc_key}.chunks.json"
216 vectors_key = f"{prefix}{doc_key}.vectors.json"
217 marker_key = f"{prefix}{doc_key}.embed.complete.json"
218 log_key = f"{prefix}{doc_key}.embed.log"
219 s3 = _s3()
220 since = time.time() - 60
221 s3.put_object(Bucket=bucket, Key=chunks_key, Body=json.dumps({"texts": texts}).encode("utf-8"))
222 payload = {
223 "task": "embed",
224 "base": doc_key,
225 "chunks_url": _presign("get_object", bucket, chunks_key),
226 "vectors_url": _presign("put_object", bucket, vectors_key),
227 "marker_url": _presign("put_object", bucket, marker_key),
228 "log_url": _presign("put_object", bucket, log_key),
229 }
230 job_id = _submit(payload)
231 log(f"embedding of {len(texts)} chunks submitted to the GPU worker as job {job_id}")
232 marker = _wait_for_marker(bucket, marker_key, log_key, EMBED_DEADLINE_SECONDS, log, job_id, since)
233 vectors = json.loads(s3.get_object(Bucket=bucket, Key=vectors_key)["Body"].read())["vectors"]
234 if len(vectors) != len(texts):
235 raise RuntimeError(f"worker returned {len(vectors)} vectors for {len(texts)} chunks")
236 log(f"embedding complete: {len(vectors)} vectors in {marker.get('total_seconds')}s")
237 return vectors