Coverage for gco/services/api_shared.py: 98.27%
147 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""
2Shared state, models, and helpers for the Manifest API routers.
4This module holds the global state (manifest processor, DynamoDB stores),
5Pydantic request/response models, and helper functions used across all API
6route modules. Centralizing them here avoids circular imports between
7manifest_api.py and the routers.
8"""
10from __future__ import annotations
12import logging
13from enum import StrEnum
14from typing import Any
16from fastapi import HTTPException
17from kubernetes.client.models import CoreV1Event, V1Job, V1Pod
18from pydantic import BaseModel, Field
20from gco.services.manifest_processor import ManifestProcessor
21from gco.services.metrics_publisher import ManifestProcessorMetrics
22from gco.services.template_store import (
23 JobStore,
24 TemplateStore,
25 WebhookStore,
26)
28logger = logging.getLogger(__name__)
31# ---------------------------------------------------------------------------
32# Shared enums and Pydantic models
33# ---------------------------------------------------------------------------
36class SortOrder(StrEnum):
37 ASC = "asc"
38 DESC = "desc"
41class JobStatus(StrEnum):
42 PENDING = "pending"
43 RUNNING = "running"
44 COMPLETED = "completed"
45 SUCCEEDED = "succeeded"
46 FAILED = "failed"
49class WebhookEvent(StrEnum):
50 JOB_COMPLETED = "job.completed"
51 JOB_FAILED = "job.failed"
52 JOB_STARTED = "job.started"
55class ManifestSubmissionAPIRequest(BaseModel):
56 """API model for manifest submission requests."""
58 manifests: list[dict[str, Any]] = Field(
59 ..., description="List of Kubernetes manifests to apply"
60 )
61 namespace: str | None = Field(
62 None, description="Default namespace for resources without namespace specified"
63 )
64 dry_run: bool = Field(False, description="If true, validate manifests without applying them")
65 validate_manifests: bool = Field(
66 True, description="If true, perform validation checks on manifests", alias="validate"
67 )
69 model_config = {
70 "json_schema_extra": {
71 "example": {
72 "manifests": [
73 {"apiVersion": "batch/v1", "kind": "Job", "metadata": {"name": "example"}}
74 ],
75 "namespace": "gco-jobs",
76 "dry_run": False,
77 }
78 }
79 }
82class ResourceIdentifier(BaseModel):
83 api_version: str = Field(..., description="Kubernetes API version (e.g., 'apps/v1')")
84 kind: str = Field(..., description="Kubernetes resource kind (e.g., 'Deployment')")
85 name: str = Field(..., description="Resource name")
86 namespace: str = Field(..., description="Resource namespace")
89class BulkDeleteRequest(BaseModel):
90 namespace: str | None = Field(None, description="Filter by namespace")
91 status: JobStatus | None = Field(None, description="Filter by status")
92 older_than_days: int | None = Field(
93 None, description="Delete jobs older than N days", ge=1, le=365
94 )
95 label_selector: str | None = Field(
96 None,
97 description="Comma-separated exact-match label filters (key=value only)",
98 max_length=1024,
99 )
100 dry_run: bool = Field(False, description="If true, only return what would be deleted")
102 model_config = {
103 "json_schema_extra": {
104 "example": {
105 "namespace": "gco-jobs",
106 "status": "completed",
107 "older_than_days": 7,
108 "dry_run": False,
109 }
110 }
111 }
114class JobTemplateRequest(BaseModel):
115 name: str = Field(..., description="Template name", min_length=1, max_length=63)
116 description: str | None = Field(None, description="Template description")
117 manifest: dict[str, Any] = Field(..., description="Job manifest template")
118 parameters: dict[str, Any] | None = Field(None, description="Default parameter values")
120 model_config = {
121 "json_schema_extra": {
122 "example": {
123 "name": "gpu-training-template",
124 "description": "Template for GPU training jobs",
125 "manifest": {
126 "apiVersion": "batch/v1",
127 "kind": "Job",
128 "metadata": {"name": "{{name}}"},
129 },
130 "parameters": {"image": "pytorch/pytorch:latest"},
131 }
132 }
133 }
136class JobFromTemplateRequest(BaseModel):
137 name: str = Field(..., description="Job name", min_length=1, max_length=63)
138 namespace: str = Field("gco-jobs", description="Target namespace")
139 parameters: dict[str, Any] | None = Field(None, description="Parameter overrides")
141 model_config = {
142 "json_schema_extra": {
143 "example": {
144 "name": "my-training-job",
145 "namespace": "gco-jobs",
146 "parameters": {"image": "my-custom-image:v1"},
147 }
148 }
149 }
152class WebhookRequest(BaseModel):
153 url: str = Field(..., description="Webhook URL to call")
154 events: list[WebhookEvent] = Field(..., description="Events to subscribe to")
155 namespace: str | None = Field(None, description="Filter by namespace (optional)")
156 secret: str | None = Field(None, description="Secret for HMAC signature (optional)")
158 model_config = {
159 "json_schema_extra": {
160 "example": {
161 "url": "https://example.com/webhook",
162 "events": ["job.completed", "job.failed"],
163 "namespace": "gco-jobs",
164 }
165 }
166 }
169class QueuedJobRequest(BaseModel):
170 manifest: dict[str, Any] = Field(..., description="Kubernetes job manifest")
171 target_region: str = Field(..., description="Target region for job execution")
172 namespace: str = Field("gco-jobs", description="Kubernetes namespace")
173 priority: int = Field(0, description="Job priority (higher = more important)", ge=0, le=100)
174 labels: dict[str, str] | None = Field(None, description="Optional labels for filtering")
175 max_spot_price: float | None = Field(
176 None,
177 gt=0,
178 description=(
179 "Optional spot price cap in USD/hour. The job is not dispatched "
180 "until the current spot price of spot_instance_type in the target "
181 "region drops to or below this value. Requires spot_instance_type."
182 ),
183 )
184 spot_instance_type: str | None = Field(
185 None,
186 description=(
187 "EC2 instance type whose spot price gates dispatch (e.g. "
188 "g5.xlarge). Requires max_spot_price."
189 ),
190 )
192 model_config = {
193 "json_schema_extra": {
194 "example": {
195 "manifest": {
196 "apiVersion": "batch/v1",
197 "kind": "Job",
198 "metadata": {"name": "my-training-job"},
199 },
200 "target_region": "us-east-1",
201 "namespace": "gco-jobs",
202 "priority": 10,
203 "max_spot_price": 0.5,
204 "spot_instance_type": "g5.xlarge",
205 }
206 }
207 }
210class PaginatedResponse(BaseModel):
211 total: int = Field(..., description="Total number of items")
212 limit: int = Field(..., description="Items per page")
213 offset: int = Field(..., description="Current offset")
214 has_more: bool = Field(..., description="Whether more items exist")
217class ErrorResponse(BaseModel):
218 error: str = Field(..., description="Error type")
219 detail: str = Field(..., description="Error details")
220 timestamp: str = Field(..., description="Error timestamp")
223# ---------------------------------------------------------------------------
224# Global state — populated by the lifespan handler in manifest_api.py
225# ---------------------------------------------------------------------------
226manifest_processor: ManifestProcessor | None = None
227manifest_metrics: ManifestProcessorMetrics | None = None
228template_store: TemplateStore | None = None
229webhook_store: WebhookStore | None = None
230job_store: JobStore | None = None
233# ---------------------------------------------------------------------------
234# Helper functions
235# ---------------------------------------------------------------------------
238def _check_processor() -> ManifestProcessor:
239 """Check if manifest processor is initialized and return it."""
240 # Import at call-time to read the global that lifespan populates on
241 # the manifest_api module (tests also patch it there).
242 from gco.services import manifest_api as _api
244 if _api.manifest_processor is None:
245 raise HTTPException(status_code=503, detail="Manifest processor not initialized")
246 return _api.manifest_processor
249def _check_namespace(namespace: str, processor: ManifestProcessor) -> None:
250 """Check if namespace is allowed."""
251 if namespace not in processor.allowed_namespaces:
252 raise HTTPException(
253 status_code=403,
254 detail=f"Namespace '{namespace}' not allowed. Allowed: {list(processor.allowed_namespaces)}",
255 )
258def _parse_job_to_dict(job: V1Job) -> dict[str, Any]:
259 """Parse a Kubernetes Job object to a dictionary."""
260 metadata = job.metadata
261 status = job.status
262 spec = job.spec
264 conditions = status.conditions or []
265 computed_status = "pending"
266 for condition in conditions:
267 if condition.type == "Complete" and condition.status == "True":
268 computed_status = "succeeded"
269 break
270 if condition.type == "Failed" and condition.status == "True": 270 ↛ 266line 270 didn't jump to line 266 because the condition on line 270 was always true
271 computed_status = "failed"
272 break
274 if computed_status == "pending" and (status.active or 0) > 0:
275 computed_status = "running"
277 # Pull container image refs from the pod template so callers (e.g.
278 # the orphan-image cross-reference) can identify which ECR images
279 # are still in use without a second round-trip per job.
280 template = getattr(spec, "template", None)
281 pod_spec = getattr(template, "spec", None) if template is not None else None
282 containers = getattr(pod_spec, "containers", None) or []
283 init_containers = getattr(pod_spec, "init_containers", None) or []
284 container_specs = [
285 {"name": getattr(c, "name", ""), "image": getattr(c, "image", "")} for c in containers
286 ]
287 init_container_specs = [
288 {"name": getattr(c, "name", ""), "image": getattr(c, "image", "")} for c in init_containers
289 ]
291 return {
292 "metadata": {
293 "name": metadata.name,
294 "namespace": metadata.namespace,
295 "creationTimestamp": (
296 metadata.creation_timestamp.isoformat() if metadata.creation_timestamp else None
297 ),
298 "labels": metadata.labels or {},
299 "annotations": metadata.annotations or {},
300 "uid": metadata.uid,
301 },
302 "spec": {
303 "parallelism": spec.parallelism,
304 "completions": spec.completions,
305 "backoffLimit": spec.backoff_limit,
306 "template": {
307 "spec": {
308 "containers": container_specs,
309 "initContainers": init_container_specs,
310 },
311 },
312 },
313 "status": {
314 "active": status.active or 0,
315 "succeeded": status.succeeded or 0,
316 "failed": status.failed or 0,
317 "startTime": status.start_time.isoformat() if status.start_time else None,
318 "completionTime": (
319 status.completion_time.isoformat() if status.completion_time else None
320 ),
321 "conditions": [
322 {
323 "type": c.type,
324 "status": c.status,
325 "reason": c.reason,
326 "message": c.message,
327 "lastTransitionTime": (
328 c.last_transition_time.isoformat() if c.last_transition_time else None
329 ),
330 }
331 for c in conditions
332 ],
333 },
334 "computed_status": computed_status,
335 }
338def _parse_pod_to_dict(pod: V1Pod) -> dict[str, Any]:
339 """Parse a Kubernetes Pod object to a dictionary."""
340 metadata = pod.metadata
341 status = pod.status
342 spec = pod.spec
344 container_statuses = []
345 for cs in status.container_statuses or []:
346 container_status: dict[str, Any] = {
347 "name": cs.name,
348 "ready": cs.ready,
349 "restartCount": cs.restart_count,
350 "image": cs.image,
351 }
352 if cs.state: 352 ↛ 365line 352 didn't jump to line 365 because the condition on line 352 was always true
353 if cs.state.running:
354 container_status["state"] = "running"
355 container_status["startedAt"] = (
356 cs.state.running.started_at.isoformat() if cs.state.running.started_at else None
357 )
358 elif cs.state.waiting:
359 container_status["state"] = "waiting"
360 container_status["reason"] = cs.state.waiting.reason
361 elif cs.state.terminated: 361 ↛ 365line 361 didn't jump to line 365 because the condition on line 361 was always true
362 container_status["state"] = "terminated"
363 container_status["exitCode"] = cs.state.terminated.exit_code
364 container_status["reason"] = cs.state.terminated.reason
365 container_statuses.append(container_status)
367 init_container_statuses = []
368 for cs in status.init_container_statuses or []:
369 init_status = {
370 "name": cs.name,
371 "ready": cs.ready,
372 "restartCount": cs.restart_count,
373 }
374 init_container_statuses.append(init_status)
376 return {
377 "metadata": {
378 "name": metadata.name,
379 "namespace": metadata.namespace,
380 "creationTimestamp": (
381 metadata.creation_timestamp.isoformat() if metadata.creation_timestamp else None
382 ),
383 "labels": metadata.labels or {},
384 "uid": metadata.uid,
385 },
386 "spec": {
387 "nodeName": spec.node_name,
388 "containers": [{"name": c.name, "image": c.image} for c in spec.containers],
389 "initContainers": [
390 {"name": c.name, "image": c.image} for c in (spec.init_containers or [])
391 ],
392 },
393 "status": {
394 "phase": status.phase,
395 "hostIP": status.host_ip,
396 "podIP": status.pod_ip,
397 "startTime": status.start_time.isoformat() if status.start_time else None,
398 "containerStatuses": container_statuses,
399 "initContainerStatuses": init_container_statuses,
400 },
401 }
404def _parse_event_to_dict(event: CoreV1Event) -> dict[str, Any]:
405 """Parse a Kubernetes Event object to a dictionary."""
406 return {
407 "type": event.type,
408 "reason": event.reason,
409 "message": event.message,
410 "count": event.count or 1,
411 "firstTimestamp": (event.first_timestamp.isoformat() if event.first_timestamp else None),
412 "lastTimestamp": (event.last_timestamp.isoformat() if event.last_timestamp else None),
413 "source": {
414 "component": event.source.component if event.source else None,
415 "host": event.source.host if event.source else None,
416 },
417 "involvedObject": {
418 "kind": event.involved_object.kind if event.involved_object else None,
419 "name": event.involved_object.name if event.involved_object else None,
420 "namespace": event.involved_object.namespace if event.involved_object else None,
421 },
422 }
425def _apply_template_parameters(
426 manifest: dict[str, Any], parameters: dict[str, Any]
427) -> dict[str, Any]:
428 """Apply parameter substitutions to a manifest template."""
429 import json
430 import re
432 manifest_str = json.dumps(manifest)
433 for key, value in parameters.items():
434 pattern = r"\{\{\s*" + re.escape(key) + r"\s*\}\}"
435 manifest_str = re.sub(pattern, str(value), manifest_str)
436 result: dict[str, Any] = json.loads(manifest_str)
437 return result