Coverage for gco/services/api_routes/queue.py: 87.01%

1"""DynamoDB-backed global job queue endpoints.""" 

2 

3from __future__ import annotations 

4 

5import hashlib 

6import json 

7import logging 

8import os 

9import re 

10import uuid 

11from copy import deepcopy 

12from datetime import UTC, datetime 

13from typing import TYPE_CHECKING, Any 

14 

15from fastapi import APIRouter, Header, HTTPException, Query 

16from fastapi.responses import JSONResponse, Response 

17 

18from gco.services.api_shared import QueuedJobRequest, _check_processor 

19from gco.services.central_queue_worker import process_queued_jobs_once 

20from gco.services.structured_logging import sanitize_log_value 

21from gco.services.template_store import JobSubmissionConflict 

22 

23if TYPE_CHECKING: 

24 from gco.services.template_store import JobStore 

25 

26router = APIRouter(prefix="/api/v1/queue", tags=["Job Queue"]) 

27logger = logging.getLogger(__name__) 

28 

29_AWS_REGION_PATTERN = re.compile(r"^[a-z]{2,4}(?:-[a-z0-9]+)+-[0-9]+$") 

30_DNS_LABEL_PATTERN = re.compile(r"^[a-z0-9](?:[-a-z0-9]*[a-z0-9])?$") 

31_IDEMPOTENCY_KEY_PATTERN = re.compile(r"^[A-Za-z0-9._~:/+=-]{1,128}$") 

32_IDEMPOTENCY_NAMESPACE = uuid.UUID("88284d12-1e04-47d5-8871-607a9e4dac09") 

33 

34 

35def _validated_queue_manifest(request: QueuedJobRequest) -> dict[str, Any]: 

36 """Validate the queue envelope and return an isolated Job manifest.""" 

37 configured_regions = { 

38 value.strip() for value in os.getenv("QUEUE_TARGET_REGIONS", "").split(",") if value.strip() 

39 } 

40 if configured_regions: 40 ↛ 41line 40 didn't jump to line 41 because the condition on line 40 was never true

41 if request.target_region not in configured_regions: 

42 raise HTTPException(status_code=422, detail="target_region is not deployed") 

43 elif not _AWS_REGION_PATTERN.fullmatch(request.target_region): 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true

44 raise HTTPException(status_code=422, detail="target_region is not a valid AWS region") 

45 

46 processor = _check_processor() 

47 if request.namespace not in processor.allowed_namespaces: 47 ↛ 48line 47 didn't jump to line 48 because the condition on line 47 was never true

48 raise HTTPException(status_code=422, detail="namespace is not allowed") 

49 if len(request.namespace) > 63 or not _DNS_LABEL_PATTERN.fullmatch(request.namespace): 49 ↛ 50line 49 didn't jump to line 50 because the condition on line 49 was never true

50 raise HTTPException(status_code=422, detail="namespace is not a valid Kubernetes name") 

51 

52 manifest = deepcopy(request.manifest) 

53 if manifest.get("apiVersion") != "batch/v1" or manifest.get("kind") != "Job": 53 ↛ 54line 53 didn't jump to line 54 because the condition on line 53 was never true

54 raise HTTPException( 

55 status_code=422, 

56 detail="central queue accepts only apiVersion 'batch/v1', kind 'Job'", 

57 ) 

58 metadata = manifest.get("metadata") 

59 if not isinstance(metadata, dict): 59 ↛ 60line 59 didn't jump to line 60 because the condition on line 59 was never true

60 raise HTTPException(status_code=422, detail="manifest.metadata must be an object") 

61 name = metadata.get("name") 

62 if not isinstance(name, str) or len(name) > 63 or not _DNS_LABEL_PATTERN.fullmatch(name): 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true

63 raise HTTPException(status_code=422, detail="manifest.metadata.name is invalid") 

64 declared_namespace = metadata.get("namespace") 

65 if declared_namespace is not None and declared_namespace != request.namespace: 65 ↛ 66line 65 didn't jump to line 66 because the condition on line 65 was never true

66 raise HTTPException( 

67 status_code=422, 

68 detail="manifest namespace must match the queue envelope namespace", 

69 ) 

70 metadata["namespace"] = request.namespace 

71 return manifest 

72 

73 

74def _validated_spot_gate(request: QueuedJobRequest) -> tuple[str, str] | None: 

75 """Validate the optional spot price gate pair on a queue submission.""" 

76 from gco.services.spot_price_gate import validate_spot_gate_fields 

77 

78 error = validate_spot_gate_fields(request.max_spot_price, request.spot_instance_type) 

79 if error: 

80 raise HTTPException(status_code=422, detail=error) 

81 if request.max_spot_price is None or request.spot_instance_type is None: 

82 return None 

83 # Serialize the cap as a plain decimal string — DynamoDB items must not 

84 # carry floats, and a stable rendering keeps idempotency hashes exact. 

85 return (f"{request.max_spot_price:.6f}", request.spot_instance_type) 

86 

87 

88def _submission_hash(request: QueuedJobRequest, manifest: dict[str, Any]) -> str: 

89 payload = { 

90 "manifest": manifest, 

91 "target_region": request.target_region, 

92 "namespace": request.namespace, 

93 "priority": request.priority, 

94 "labels": request.labels or {}, 

95 } 

96 # Only price-capped submissions carry the gate fields, so historical 

97 # idempotency keys hash identically to pre-gate deployments. 

98 if request.max_spot_price is not None or request.spot_instance_type is not None: 

99 payload["max_spot_price"] = request.max_spot_price 

100 payload["spot_instance_type"] = request.spot_instance_type 

101 canonical = json.dumps(payload, sort_keys=True, separators=(",", ":")) 

102 return hashlib.sha256(canonical.encode("utf-8")).hexdigest() 

103 

104 

105def _get_job_store() -> JobStore: 

106 from gco.services.manifest_api import job_store 

107 

108 if job_store is None: 

109 raise HTTPException(status_code=503, detail="Job store not initialized") 

110 return job_store 

111 

112 

113@router.post("/jobs") 

114async def submit_job_to_queue( 

115 request: QueuedJobRequest, 

116 idempotency_key: str | None = Header( 

117 None, 

118 alias="Idempotency-Key", 

119 description="Stable key for safely replaying an identical submission", 

120 ), 

121) -> Response: 

122 """Submit one validated ``batch/v1`` Job exactly once.""" 

123 if idempotency_key is not None and not _IDEMPOTENCY_KEY_PATTERN.fullmatch(idempotency_key): 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true

124 raise HTTPException(status_code=422, detail="Idempotency-Key is invalid") 

125 

126 manifest = _validated_queue_manifest(request) 

127 spot_gate = _validated_spot_gate(request) 

128 request_hash = _submission_hash(request, manifest) 

129 job_id = ( 

130 str(uuid.uuid5(_IDEMPOTENCY_NAMESPACE, idempotency_key)) 

131 if idempotency_key 

132 else str(uuid.uuid4()) 

133 ) 

134 store = _get_job_store() 

135 

136 try: 

137 job = store.submit_job( 

138 job_id=job_id, 

139 manifest=manifest, 

140 target_region=request.target_region, 

141 namespace=request.namespace, 

142 priority=request.priority, 

143 labels=request.labels, 

144 idempotency_key=idempotency_key, 

145 request_hash=request_hash, 

146 spot_max_price=spot_gate[0] if spot_gate else None, 

147 spot_instance_type=spot_gate[1] if spot_gate else None, 

148 ) 

149 except JobSubmissionConflict as error: 

150 raise HTTPException(status_code=409, detail=str(error)) from error 

151 except Exception as error: 

152 logger.exception("Failed to queue job") 

153 raise HTTPException(status_code=500, detail="Failed to queue job") from error 

154 

155 replay = bool(job.pop("idempotent_replay", False)) 

156 return JSONResponse( 

157 status_code=200 if replay else 201, 

158 content={ 

159 "timestamp": datetime.now(UTC).isoformat(), 

160 "message": "Idempotent job replay" if replay else "Job queued successfully", 

161 "job": job, 

162 }, 

163 ) 

164 

165 

166@router.get("/jobs") 

167async def list_queued_jobs( 

168 target_region: str | None = Query(None, description="Filter by target region"), 

169 status: str | None = Query(None, description="Filter by status"), 

170 namespace: str | None = Query(None, description="Filter by namespace"), 

171 limit: int = Query(100, description="Maximum results", ge=1, le=1000), 

172 cursor: str | None = Query( 

173 None, 

174 description="Opaque continuation cursor returned by the previous page", 

175 max_length=2048, 

176 ), 

177) -> Response: 

178 """List one bounded page of jobs with optional filters.""" 

179 store = _get_job_store() 

180 try: 

181 jobs, next_cursor, partial = store.list_jobs_page( 

182 target_region=target_region, 

183 status=status, 

184 namespace=namespace, 

185 limit=limit, 

186 cursor=cursor, 

187 ) 

188 return JSONResponse( 

189 status_code=200, 

190 content={ 

191 "timestamp": datetime.now(UTC).isoformat(), 

192 "count": len(jobs), 

193 "jobs": jobs, 

194 "next_cursor": next_cursor, 

195 "partial": partial, 

196 }, 

197 ) 

198 except ValueError as error: 

199 raise HTTPException(status_code=422, detail=str(error)) from error 

200 except Exception as e: 

201 logger.error(f"Failed to list queued jobs: {e}") 

202 raise HTTPException(status_code=500, detail=f"Failed to list jobs: {e!s}") from e 

203 

204 

205@router.get("/jobs/{job_id}") 

206async def get_queued_job(job_id: str) -> Response: 

207 """Get details of a specific queued job.""" 

208 store = _get_job_store() 

209 try: 

210 job = store.get_job(job_id) 

211 if job is None: 

212 raise HTTPException(status_code=404, detail=f"Job '{job_id}' not found") 

213 return JSONResponse( 

214 status_code=200, 

215 content={"timestamp": datetime.now(UTC).isoformat(), "job": job}, 

216 ) 

217 except HTTPException: 

218 raise 

219 except Exception as e: 

220 logger.error("Failed to get job %s: %s", sanitize_log_value(job_id), e) 

221 raise HTTPException(status_code=500, detail=f"Failed to get job: {e!s}") from e 

222 

223 

224@router.delete("/jobs/{job_id}") 

225async def cancel_queued_job( 

226 job_id: str, reason: str | None = Query(None, description="Cancellation reason") 

227) -> Response: 

228 """Cancel a job only while it remains unclaimed in the queue.""" 

229 store = _get_job_store() 

230 try: 

231 cancelled = store.cancel_job(job_id, reason=reason) 

232 if not cancelled: 

233 raise HTTPException( 

234 status_code=409, 

235 detail=f"Job '{job_id}' cannot be cancelled (already running or completed)", 

236 ) 

237 return JSONResponse( 

238 status_code=200, 

239 content={ 

240 "timestamp": datetime.now(UTC).isoformat(), 

241 "message": f"Job '{job_id}' cancelled successfully", 

242 }, 

243 ) 

244 except HTTPException: 

245 raise 

246 except Exception as e: 

247 logger.error("Failed to cancel job %s: %s", sanitize_log_value(job_id), e) 

248 raise HTTPException(status_code=500, detail=f"Failed to cancel job: {e!s}") from e 

249 

250 

251@router.get("/stats") 

252async def get_queue_stats() -> Response: 

253 """Get job queue statistics by region and status.""" 

254 store = _get_job_store() 

255 try: 

256 counts, records_evaluated, truncated = store.get_job_count_summary() 

257 total_jobs = sum(sum(statuses.values()) for statuses in counts.values()) 

258 total_queued = sum(statuses.get("queued", 0) for statuses in counts.values()) 

259 total_running = sum(statuses.get("running", 0) for statuses in counts.values()) 

260 

261 return JSONResponse( 

262 status_code=200, 

263 content={ 

264 "timestamp": datetime.now(UTC).isoformat(), 

265 "summary": { 

266 "total_jobs": total_jobs, 

267 "total_queued": total_queued, 

268 "total_running": total_running, 

269 "complete": not truncated, 

270 "records_evaluated": records_evaluated, 

271 }, 

272 "by_region": counts, 

273 }, 

274 ) 

275 except Exception as e: 

276 logger.error(f"Failed to get queue stats: {e}") 

277 raise HTTPException(status_code=500, detail=f"Failed to get stats: {e!s}") from e 

278 

279 

280@router.post("/poll") 

281async def poll_and_process_jobs( 

282 limit: int = Query(5, description="Maximum jobs to process", ge=1, le=20), 

283) -> Response: 

284 """Run one immediate queue-worker pass for this region. 

285 

286 The manifest API also runs this same processor continuously when the 

287 deployment enables ``CENTRAL_QUEUE_WORKER_ENABLED``. This endpoint remains 

288 useful for an authenticated operator-triggered pass and diagnostics. 

289 """ 

290 processor = _check_processor() 

291 store = _get_job_store() 

292 

293 try: 

294 jobs_polled, processed_jobs = await process_queued_jobs_once( 

295 processor, 

296 store, 

297 limit=limit, 

298 ) 

299 return JSONResponse( 

300 status_code=200, 

301 content={ 

302 "timestamp": datetime.now(UTC).isoformat(), 

303 "region": processor.region, 

304 "jobs_polled": jobs_polled, 

305 "jobs_processed": len(processed_jobs), 

306 "results": processed_jobs, 

307 }, 

308 ) 

309 except Exception as e: 

310 logger.error(f"Failed to poll jobs: {e}") 

311 raise HTTPException(status_code=500, detail=f"Failed to poll jobs: {e!s}") from e