Coverage for cli / jobs.py: 100.00%
538 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""
2Job management for GCO CLI.
4Provides functionality to submit, query, and manage jobs across GCO clusters.
5"""
7from collections.abc import Callable, Mapping
8from dataclasses import dataclass, field
9from datetime import UTC, datetime
10from pathlib import Path
11from typing import Any
12from urllib.parse import quote
14import requests
15import yaml
17from gco.job_envelope import build_job_message
18from gco.services.manifest_processor import TRAINJOB_API_VERSION, safe_load_all_yaml
20from .aws_client import get_aws_client
21from .config import GCOConfig, get_config
23# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
24# Generated at (UTC): 2026-09-20T23:20:00Z
25# Generated from Git commit: 42cc3f2094423b3f8056fbc5c20652e222b0650f
26# Flowchart(s) generated from this file:
27# * ``JobManager.submit_job`` -> ``diagrams/code_diagrams/cli/jobs.JobManager_submit_job.html``
28# (PNG: ``diagrams/code_diagrams/cli/jobs.JobManager_submit_job.png``)
29# * ``JobManager.submit_job_sqs`` -> ``diagrams/code_diagrams/cli/jobs.JobManager_submit_job_sqs.html``
30# (PNG: ``diagrams/code_diagrams/cli/jobs.JobManager_submit_job_sqs.png``)
31# * ``JobManager.get_job_logs`` -> ``diagrams/code_diagrams/cli/jobs.JobManager_get_job_logs.html``
32# (PNG: ``diagrams/code_diagrams/cli/jobs.JobManager_get_job_logs.png``)
33# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
34# <pyflowchart-code-diagram> END
37logger = __import__("logging").getLogger(__name__)
40def _format_duration(seconds: int) -> str:
41 """Format seconds into a human-readable duration string."""
42 if seconds < 60:
43 return f"{seconds}s"
44 minutes, secs = divmod(seconds, 60)
45 if minutes < 60:
46 return f"{minutes}m{secs:02d}s"
47 hours, mins = divmod(minutes, 60)
48 return f"{hours}h{mins:02d}m{secs:02d}s"
51def resolve_submission_identity(
52 result: Any,
53 *,
54 fallback_name: str | None = None,
55 fallback_namespace: str | None = None,
56) -> tuple[str | None, str | None]:
57 """Resolve the submitted Job name and namespace from supported responses.
59 API submissions return resource-status dictionaries, while direct kubectl
60 submissions return a top-level ``job_name`` and a ``resources`` list of
61 human-readable strings. Only mapping-shaped resources are inspected, so
62 direct response lines can never be mistaken for response envelopes.
63 """
64 if not isinstance(result, Mapping):
65 return fallback_name, fallback_namespace
67 raw_resources = result.get("resources") or []
68 if isinstance(raw_resources, Mapping):
69 raw_resources = [raw_resources]
70 resources = [resource for resource in raw_resources if isinstance(resource, Mapping)]
71 job_resources = [
72 resource for resource in resources if str(resource.get("kind", "")).lower() == "job"
73 ]
75 explicit_job_name = result.get("job_name")
76 resource_with_name = next(
77 (resource for resource in job_resources if resource.get("name")), None
78 )
79 job_name = (
80 str(explicit_job_name)
81 if explicit_job_name
82 else str(resource_with_name.get("name"))
83 if resource_with_name is not None
84 else fallback_name
85 )
87 matching_resource = next(
88 (resource for resource in job_resources if resource.get("name") == job_name),
89 resource_with_name,
90 )
91 resource_namespace = matching_resource.get("namespace") if matching_resource else None
92 envelope_namespace = result.get("namespace")
93 namespace = (
94 str(resource_namespace)
95 if resource_namespace
96 else str(envelope_namespace)
97 if envelope_namespace
98 else fallback_namespace
99 )
100 return job_name, namespace
103def _extract_image_refs(spec: dict[str, Any]) -> list[str]:
104 """Extract container image refs from a parsed Job spec.
106 The API surface for a Job carries ``spec.template.spec.containers`` and
107 ``spec.template.spec.initContainers`` lists, each entry of which has
108 a ``name`` and an ``image`` URI. Returns an alphabetically-sorted,
109 deduplicated list so the output is stable across calls — orphan-image
110 cross-references rely on set equality.
111 """
112 refs: set[str] = set()
113 template = spec.get("template") if isinstance(spec, dict) else None
114 pod_spec = template.get("spec") if isinstance(template, dict) else None
115 if not isinstance(pod_spec, dict):
116 return []
117 for key in ("containers", "initContainers"):
118 items = pod_spec.get(key, [])
119 if not isinstance(items, list):
120 continue
121 for entry in items:
122 if not isinstance(entry, dict):
123 continue
124 image = entry.get("image")
125 if isinstance(image, str) and image:
126 refs.add(image)
127 return sorted(refs)
130def _extract_scheduling(payload: dict[str, Any]) -> dict[str, Any]:
131 """Pull the node-placement block out of a job or pods API response.
133 Absent on responses that predate the field (an older regional bridge) and
134 on the list endpoint, which deliberately does not pay for a Node read per
135 job. Both cases yield empty placement rather than an error: an absent
136 instance type is honest, a guessed one is not.
137 """
138 scheduling = payload.get("scheduling")
139 if not isinstance(scheduling, dict):
140 return {}
142 def _text(key: str) -> str | None:
143 value = scheduling.get(key)
144 return value if isinstance(value, str) else None
146 def _count(key: str) -> int | None:
147 # A bridge that predates the field reports nothing, which must read as
148 # "unknown" (None) rather than as a confident zero.
149 value = scheduling.get(key)
150 return value if isinstance(value, int) and not isinstance(value, bool) else None
152 nodes = scheduling.get("nodes")
153 labels = scheduling.get("node_labels")
154 phases = scheduling.get("pod_phases")
155 unscheduled = scheduling.get("unscheduled")
156 return {
157 "node_name": _text("node_name"),
158 "node_instance_type": _text("node_instance_type"),
159 "node_capacity_type": _text("node_capacity_type"),
160 "node_labels": dict(labels) if isinstance(labels, dict) else {},
161 "nodes": list(nodes) if isinstance(nodes, list) else [],
162 "pod_phase": _text("pod_phase"),
163 "pod_phases": dict(phases) if isinstance(phases, dict) else {},
164 "scheduled_pods": _count("scheduled_pods"),
165 "unscheduled_pods": _count("unscheduled_pods"),
166 "unscheduled": [dict(pod) for pod in unscheduled if isinstance(pod, dict)]
167 if isinstance(unscheduled, list)
168 else [],
169 "node_lookup_error": _text("node_lookup_error"),
170 }
173@dataclass
174class JobInfo:
175 """Information about a Kubernetes job."""
177 name: str
178 namespace: str
179 region: str
180 status: str # "pending", "running", "succeeded", "failed"
181 created_time: datetime | None = None
182 start_time: datetime | None = None
183 completion_time: datetime | None = None
184 active_pods: int = 0
185 succeeded_pods: int = 0
186 failed_pods: int = 0
187 parallelism: int = 1
188 completions: int = 1
189 labels: dict[str, str] = field(default_factory=dict)
190 image_refs: list[str] = field(default_factory=list)
192 # Where the job's pods actually landed. A job constrained to a *set* of
193 # interchangeable instance types is placed by Karpenter within that set,
194 # so the manifest records only what the run was authorized to use. These
195 # record what it used, which is what cost reconciliation and
196 # "did this fail on the smaller box?" both need.
197 #
198 # ``node_*`` describe the earliest-created scheduled pod; ``nodes`` lists
199 # every node involved (a retried job can move between instance types) with
200 # the pods on each. All are unset when nothing is scheduled yet, when the
201 # pods have been garbage-collected, or on the list endpoint, which does not
202 # pay for a Node read per job.
203 node_name: str | None = None
204 node_instance_type: str | None = None
205 node_capacity_type: str | None = None
206 node_labels: dict[str, str] = field(default_factory=dict)
207 nodes: list[dict[str, Any]] = field(default_factory=list)
209 # What the pods themselves are doing. ``status`` above is the *Job's*
210 # verdict, and a Job whose only pod cannot be scheduled still reports
211 # ``running``; these fields are how a caller tells "the container is
212 # training" from "a Job object exists and nothing is placed".
213 #
214 # ``pod_phase`` is the phase every pod shares (None the moment they
215 # disagree — it never summarizes a Pending pod away), ``pod_phases`` counts
216 # pods by phase, ``scheduled_pods`` / ``unscheduled_pods`` split them by
217 # whether a node was assigned, and ``unscheduled`` names each unplaced pod
218 # with the ``PodScheduled`` condition's ``reason`` / ``message`` / ``since``.
219 # The counts are None, not 0, when the regional bridge predates them:
220 # absent stays absent.
221 pod_phase: str | None = None
222 pod_phases: dict[str, int] = field(default_factory=dict)
223 scheduled_pods: int | None = None
224 unscheduled_pods: int | None = None
225 unscheduled: list[dict[str, Any]] = field(default_factory=list)
226 # Why node hardware is missing when it is (RBAC refusal, node reclaimed,
227 # pod listing failed). Evidence for the caller, never folded into ``status``.
228 node_lookup_error: str | None = None
230 @property
231 def is_complete(self) -> bool:
232 return self.status in ("succeeded", "failed")
234 @property
235 def duration_seconds(self) -> int | None:
236 if self.start_time and self.completion_time:
237 return int((self.completion_time - self.start_time).total_seconds())
238 if self.start_time:
239 return int((datetime.now(UTC) - self.start_time).total_seconds())
240 return None
243class JobManager:
244 """
245 Manages jobs across GCO clusters.
247 Provides:
248 - Job submission with region targeting
249 - Job status queries across regions
250 - Job logs retrieval
251 - Job deletion
252 """
254 def __init__(self, config: GCOConfig | None = None):
255 self.config = config or get_config()
256 self._aws_client = get_aws_client(self.config)
258 def load_manifests(self, path: str) -> list[dict[str, Any]]:
259 """
260 Load Kubernetes manifests from a file or directory.
262 Args:
263 path: Path to YAML file or directory containing YAML files
265 Returns:
266 List of manifest dictionaries
267 """
268 manifests = []
269 path_obj = Path(path)
271 if path_obj.is_file():
272 manifests.extend(self._load_yaml_file(path_obj))
273 elif path_obj.is_dir():
274 for yaml_file in sorted(path_obj.glob("*.yaml")):
275 manifests.extend(self._load_yaml_file(yaml_file))
276 for yaml_file in sorted(path_obj.glob("*.yml")):
277 manifests.extend(self._load_yaml_file(yaml_file))
278 else:
279 raise FileNotFoundError(f"Path not found: {path}")
281 return manifests
283 def _load_yaml_file(self, path: Path) -> list[dict[str, Any]]:
284 """Load manifests from a single YAML file."""
285 with open(path, encoding="utf-8") as f:
286 return safe_load_all_yaml(f, allow_aliases=False)
288 def submit_job(
289 self,
290 manifests: str | list[dict[str, Any]],
291 namespace: str | None = None,
292 target_region: str | None = None,
293 dry_run: bool = False,
294 labels: dict[str, str] | None = None,
295 ) -> dict[str, Any]:
296 """
297 Submit a job to GCO.
299 Args:
300 manifests: Path to manifest file/directory or list of manifest dicts
301 namespace: Fallback namespace for manifests that don't declare
302 their own. When set, each manifest's ``metadata.namespace`` is
303 filled in only if missing — existing values are preserved so
304 users who've declared a target namespace in the manifest can
305 rely on it reaching the server untouched. Server-side
306 validation enforces the allowlist.
307 target_region: Force job to specific region
308 dry_run: Validate without applying
309 labels: Additional labels to add to manifests
311 Returns:
312 Submission result dictionary
313 """
314 # Load manifests if path provided
315 manifest_list = self.load_manifests(manifests) if isinstance(manifests, str) else manifests
317 # Apply the explicit or configured namespace as a fallback only —
318 # preserve any namespace declared by an individual manifest.
319 effective_namespace = namespace or self.config.default_namespace
320 for manifest in manifest_list:
321 if "metadata" not in manifest:
322 manifest["metadata"] = {}
323 manifest["metadata"].setdefault("namespace", effective_namespace)
325 # Apply additional labels
326 if labels:
327 for manifest in manifest_list:
328 if "labels" not in manifest["metadata"]:
329 manifest["metadata"]["labels"] = {}
330 manifest["metadata"]["labels"].update(labels)
332 # Submit via API
333 return self._aws_client.submit_manifests(
334 manifests=manifest_list,
335 namespace=effective_namespace,
336 target_region=target_region,
337 dry_run=dry_run,
338 )
340 def submit_job_direct(
341 self,
342 manifests: str | list[dict[str, Any]],
343 region: str,
344 namespace: str | None = None,
345 dry_run: bool = False,
346 labels: dict[str, str] | None = None,
347 ) -> dict[str, Any]:
348 """
349 Submit a job directly to a regional cluster using kubectl.
351 This bypasses the API Gateway and submits directly to the EKS cluster
352 using kubectl. Requires:
353 - kubectl installed and in PATH
354 - EKS access entry configured for your IAM principal
355 - AWS credentials with eks:DescribeCluster permission
357 Args:
358 manifests: Path to manifest file/directory or list of manifest dicts
359 region: Target region for direct submission (required)
360 namespace: Fallback namespace for manifests that don't declare
361 their own. When set, each manifest's ``metadata.namespace`` is
362 filled in only if missing — existing values are preserved so
363 users who've declared a target namespace in the manifest can
364 rely on it reaching ``kubectl apply`` untouched.
365 dry_run: Validate without applying
366 labels: Additional labels to add to manifests
368 Returns:
369 Submission result dictionary
370 """
371 import subprocess
372 import tempfile
373 import uuid
375 # Load manifests if path provided
376 manifest_list = self.load_manifests(manifests) if isinstance(manifests, str) else manifests
378 # Apply the explicit or configured namespace as a fallback only —
379 # preserve any namespace declared by an individual manifest.
380 effective_namespace = namespace or self.config.default_namespace
381 for manifest in manifest_list:
382 if "metadata" not in manifest:
383 manifest["metadata"] = {}
384 manifest["metadata"].setdefault("namespace", effective_namespace)
386 # Apply additional labels
387 if labels:
388 for manifest in manifest_list:
389 if "labels" not in manifest["metadata"]:
390 manifest["metadata"]["labels"] = {}
391 manifest["metadata"]["labels"].update(labels)
393 # Get cluster name from stack
394 stack = self._aws_client.get_regional_stack(region)
395 if not stack:
396 raise ValueError(f"No GCO stack found in region {region}")
398 cluster_name = stack.cluster_name
400 # Update kubeconfig for the cluster
401 from .kubectl_helpers import update_kubeconfig
403 update_kubeconfig(cluster_name, region)
405 # Handle existing Job resources before applying
406 warnings: list[str] = []
407 if not dry_run:
408 for manifest in manifest_list:
409 if manifest.get("kind") != "Job":
410 continue
411 job_name = manifest.get("metadata", {}).get("name")
412 job_ns = manifest.get("metadata", {}).get("namespace", effective_namespace)
413 if not job_name:
414 continue
416 existing_status = self._get_kubectl_job_status(job_name, job_ns)
417 if existing_status is None:
418 # No existing job — nothing to do
419 continue
421 if existing_status in ("complete", "failed"):
422 # Finished job — safe to delete and replace
423 subprocess.run(
424 ["kubectl", "delete", "job", job_name, "-n", job_ns],
425 capture_output=True,
426 text=True,
427 check=False,
428 )
429 else:
430 # Job is still active — auto-rename to avoid collision
431 suffix = uuid.uuid4().hex[:5]
432 new_name = f"{job_name}-{suffix}"
433 original_name = job_name
434 manifest["metadata"]["name"] = new_name
435 warnings.append(
436 f"Job '{original_name}' is still running in namespace "
437 f"'{job_ns}'. Renamed new submission to '{new_name}'."
438 )
439 logger.warning(
440 "Job %s is active in %s, renamed to %s",
441 original_name,
442 job_ns,
443 new_name,
444 )
446 # Write manifests to temp file
447 with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
448 yaml.dump_all(manifest_list, f)
449 f.flush() # Ensure content is written before using f.name
450 temp_path = f.name # nosemgrep: tempfile-without-flush
452 try:
453 # Build kubectl command
454 kubectl_cmd = ["kubectl", "apply", "-f", temp_path]
456 if dry_run:
457 kubectl_cmd.extend(["--dry-run=client"])
459 # Run kubectl apply
460 result = subprocess.run(
461 kubectl_cmd,
462 capture_output=True,
463 text=True,
464 check=False,
465 ) # nosemgrep: dangerous-subprocess-use-audit - kubectl_cmd is a list ["kubectl","apply","-f",temp_path]; temp_path is a secure tempfile, not user input
467 if result.returncode != 0:
468 raise RuntimeError(f"kubectl apply failed: {result.stderr}")
470 # Parse output to get job name
471 output_lines = result.stdout.strip().split("\n")
472 created_resources = []
473 for line in output_lines:
474 if line:
475 created_resources.append(line)
477 # Get the actual Job identity from the submitted manifest. The
478 # name may have been changed above to avoid an active-job collision,
479 # and a manifest-declared namespace takes precedence over the CLI
480 # fallback.
481 job_name = None
482 job_namespace = effective_namespace
483 for manifest in manifest_list:
484 if manifest.get("kind") == "Job":
485 metadata = manifest.get("metadata", {})
486 job_name = metadata.get("name")
487 job_namespace = metadata.get("namespace") or job_namespace
488 break
490 response: dict[str, Any] = {
491 "status": "success",
492 "method": "kubectl",
493 "cluster": cluster_name,
494 "region": region,
495 "namespace": job_namespace,
496 "job_name": job_name,
497 "dry_run": dry_run,
498 "resources": created_resources,
499 "output": result.stdout,
500 }
501 if warnings:
502 response["warnings"] = warnings
503 return response
505 finally:
506 # Clean up temp file
507 import os
509 os.unlink(temp_path)
511 def _get_kubectl_job_status(self, job_name: str, namespace: str) -> str | None:
512 """Check the status of an existing Job via kubectl.
514 Returns:
515 "complete", "failed", "active", or None if the job doesn't exist.
516 """
517 import json
518 import subprocess
520 result = subprocess.run(
521 [
522 "kubectl",
523 "get",
524 "job",
525 job_name,
526 "-n",
527 namespace,
528 "-o",
529 "json",
530 ],
531 capture_output=True,
532 text=True,
533 check=False,
534 )
535 if result.returncode != 0:
536 return None # Job doesn't exist
538 try:
539 job_data = json.loads(result.stdout)
540 except json.JSONDecodeError, KeyError:
541 return None
543 conditions = job_data.get("status", {}).get("conditions") or []
544 for condition in conditions:
545 cond_type = condition.get("type", "")
546 cond_status = condition.get("status", "")
547 if cond_type == "Complete" and cond_status == "True":
548 return "complete"
549 if cond_type == "Failed" and cond_status == "True":
550 return "failed"
551 return "active"
553 def list_jobs(
554 self,
555 region: str | None = None,
556 namespace: str | None = None,
557 status: str | None = None,
558 all_regions: bool = False,
559 ) -> list[JobInfo]:
560 """
561 List jobs across GCO clusters.
563 Args:
564 region: Specific region to query
565 namespace: Filter by namespace
566 status: Filter by status
567 all_regions: Query all discovered regions
569 Returns:
570 List of JobInfo objects
571 """
572 jobs = []
574 if all_regions:
575 # Query all discovered regional stacks
576 stacks = self._aws_client.discover_regional_stacks()
577 for stack_region in stacks:
578 try:
579 region_jobs = self._query_jobs_in_region(stack_region, namespace, status)
580 jobs.extend(region_jobs)
581 except Exception as e:
582 logger.warning("Failed to query jobs in %s: %s", stack_region, e)
583 continue
584 elif region:
585 jobs = self._query_jobs_in_region(region, namespace, status)
586 else:
587 # Use default region
588 jobs = self._query_jobs_in_region(self.config.default_region, namespace, status)
590 return jobs
592 def _query_jobs_in_region(
593 self, region: str, namespace: str | None, status: str | None
594 ) -> list[JobInfo]:
595 """Query jobs in a specific region."""
596 try:
597 response = self._aws_client.get_jobs(region=region, namespace=namespace, status=status)
598 return [self._parse_job_info(job_data, region) for job_data in response.get("jobs", [])]
599 except Exception as exc:
600 logger.warning("Failed to query jobs in %s: %s", region, exc)
601 return []
603 def _parse_job_info(self, job_data: dict[str, Any], region: str) -> JobInfo:
604 """Parse job data into JobInfo object."""
605 metadata = job_data.get("metadata", {})
606 status_data = job_data.get("status", {})
607 spec = job_data.get("spec", {})
609 # Determine job status
610 conditions = status_data.get("conditions", [])
611 job_status = "pending"
612 for condition in conditions:
613 if condition.get("type") == "Complete" and condition.get("status") == "True":
614 job_status = "succeeded"
615 break
616 if condition.get("type") == "Failed" and condition.get("status") == "True":
617 job_status = "failed"
618 break
620 if job_status == "pending" and status_data.get("active", 0) > 0:
621 job_status = "running"
623 # Parse timestamps
624 created_time = None
625 if metadata.get("creationTimestamp"):
626 created_time = datetime.fromisoformat(
627 metadata["creationTimestamp"].replace("Z", "+00:00")
628 )
630 start_time = None
631 if status_data.get("startTime"):
632 start_time = datetime.fromisoformat(status_data["startTime"].replace("Z", "+00:00"))
634 completion_time = None
635 if status_data.get("completionTime"):
636 completion_time = datetime.fromisoformat(
637 status_data["completionTime"].replace("Z", "+00:00")
638 )
640 return JobInfo(
641 name=metadata.get("name", ""),
642 namespace=metadata.get("namespace", "default"),
643 region=region,
644 status=job_status,
645 created_time=created_time,
646 start_time=start_time,
647 completion_time=completion_time,
648 active_pods=status_data.get("active", 0),
649 succeeded_pods=status_data.get("succeeded", 0),
650 failed_pods=status_data.get("failed", 0),
651 parallelism=spec.get("parallelism", 1),
652 completions=spec.get("completions", 1),
653 labels=metadata.get("labels", {}),
654 image_refs=_extract_image_refs(spec),
655 **_extract_scheduling(job_data),
656 )
658 def get_job(self, job_name: str, namespace: str, region: str | None = None) -> JobInfo | None:
659 """
660 Get detailed information about a specific job.
662 A name the batch Job endpoint does not know (404) is retried as a
663 Kubeflow TrainJob through the generic manifests endpoint, so
664 ``gco jobs get`` / ``gco jobs submit --wait`` work unchanged for
665 TrainJobs.
667 ``None`` means the API confirmed the job does not exist (HTTP 404 on
668 both lookups). Any other failure — an unreachable regional API bridge,
669 an auth error, a 5xx — propagates to the caller so it is never
670 misreported as "not found" (see issue #258).
672 Args:
673 job_name: Name of the job
674 namespace: Namespace of the job
675 region: Region where the job is running
677 Returns:
678 JobInfo, or None if the API confirmed the job does not exist
680 Raises:
681 RuntimeError: If the regional API bridge is unreachable or the
682 request fails for reasons other than the job being absent
683 requests.exceptions.HTTPError: For non-404 HTTP responses
684 """
685 target_region = region or self.config.default_region
686 try:
687 response = self._aws_client.get_job_details(
688 job_name=job_name, namespace=namespace, region=target_region
689 )
690 return self._parse_job_info(response, target_region)
691 except requests.exceptions.HTTPError as e:
692 if e.response is not None and e.response.status_code == 404:
693 # The batch Job endpoint confirmed the name is unknown; the
694 # job may still exist as a Kubeflow TrainJob.
695 trainjob = self._get_trainjob_info(job_name, namespace, target_region)
696 if trainjob is not None:
697 return trainjob
698 logger.debug("Job %s not found in %s/%s", job_name, target_region, namespace)
699 return None
700 # Non-404 responses (401/403/5xx) do not mean the job is absent.
701 # Surface the real failure instead of collapsing it into None,
702 # which the CLI would render as "not found".
703 raise
705 def _get_trainjob_info(self, job_name: str, namespace: str, region: str) -> JobInfo | None:
706 """Fetch a TrainJob through the generic manifests endpoint, or None."""
707 try:
708 response = self._aws_client.call_api(
709 "GET",
710 f"/api/v1/manifests/{quote(namespace, safe='')}/{quote(job_name, safe='')}",
711 region=region,
712 params={"api_version": TRAINJOB_API_VERSION, "kind": "TrainJob"},
713 )
714 except RuntimeError as e:
715 # 404 (no such TrainJob either) and transport errors both mean
716 # "not resolvable as a TrainJob"; the caller reports not-found.
717 logger.debug("TrainJob fallback lookup failed for %s: %s", job_name, e)
718 return None
719 resource = response.get("resource") or {}
720 if not resource.get("exists", False):
721 return None
722 info = self._parse_trainjob_info(resource, region)
723 self._attach_trainjob_scheduling(info, job_name, namespace, region)
724 return info
726 def _attach_trainjob_scheduling(
727 self, info: JobInfo, job_name: str, namespace: str, region: str
728 ) -> None:
729 """Fill in node placement for a TrainJob from its child Job's pods.
731 The manifests endpoint that resolves a TrainJob knows nothing about
732 pods, so placement comes from the child Job ``<name>-node-0`` whose
733 indexed pods are the node ranks (same convention the log fetch uses).
734 Best effort: a TrainJob with no pods yet, or a bridge too old to report
735 placement, simply leaves the fields unset.
736 """
737 try:
738 response = self._aws_client.get_job_pods(f"{job_name}-node-0", namespace, region)
739 except Exception as exc:
740 logger.debug("No node placement available for TrainJob %s: %s", job_name, exc)
741 return
743 scheduling = _extract_scheduling(response if isinstance(response, dict) else {})
744 for attr, value in scheduling.items():
745 setattr(info, attr, value)
747 def _parse_trainjob_info(self, resource: dict[str, Any], region: str) -> JobInfo:
748 """Map a TrainJob resource payload onto JobInfo.
750 TrainJob status carries ``conditions`` (terminal types ``Complete`` /
751 ``Failed``, same strings as batch Jobs) plus per-child-Job counts in
752 ``jobsStatus``; there is no ``startTime``/``completionTime``, so the
753 terminal condition's ``lastTransitionTime`` stands in for completion.
754 """
755 metadata = resource.get("metadata", {}) or {}
756 status_data = resource.get("status", {}) or {}
757 spec = resource.get("spec", {}) or {}
759 job_status = "pending"
760 completion_time = None
761 for condition in status_data.get("conditions", []) or []:
762 if condition.get("status") != "True":
763 continue
764 if condition.get("type") == "Complete":
765 job_status = "succeeded"
766 elif condition.get("type") == "Failed":
767 job_status = "failed"
768 else:
769 continue
770 if condition.get("lastTransitionTime"):
771 completion_time = datetime.fromisoformat(
772 condition["lastTransitionTime"].replace("Z", "+00:00")
773 )
774 break
776 active = succeeded = failed = 0
777 for child in status_data.get("jobsStatus", []) or []:
778 active += child.get("active", 0) or 0
779 succeeded += child.get("succeeded", 0) or 0
780 failed += child.get("failed", 0) or 0
781 if job_status == "pending" and active > 0:
782 job_status = "running"
784 created_time = None
785 if metadata.get("creationTimestamp"):
786 created_time = datetime.fromisoformat(
787 metadata["creationTimestamp"].replace("Z", "+00:00")
788 )
790 trainer = spec.get("trainer", {}) or {}
791 try:
792 num_nodes = max(1, int(trainer.get("numNodes") or 1))
793 except TypeError, ValueError:
794 num_nodes = 1
796 return JobInfo(
797 name=metadata.get("name", ""),
798 namespace=metadata.get("namespace", "default"),
799 region=region,
800 status=job_status,
801 created_time=created_time,
802 # TrainJob status reports no startTime; leave it unset rather
803 # than fabricating one from creation time.
804 completion_time=completion_time,
805 active_pods=active,
806 succeeded_pods=succeeded,
807 failed_pods=failed,
808 parallelism=num_nodes,
809 completions=num_nodes,
810 labels=metadata.get("labels", {}) or {},
811 image_refs=sorted({trainer["image"]} if trainer.get("image") else set()),
812 )
814 def get_job_logs(
815 self,
816 job_name: str,
817 namespace: str,
818 region: str | None = None,
819 tail_lines: int = 100,
820 follow: bool = False,
821 since_hours: int = 24,
822 node: int = 0,
823 ) -> str:
824 """
825 Get logs from a job.
827 Tries the Kubernetes API first (via the GCO API). A name the batch
828 Job endpoint does not know is retried as a Kubeflow TrainJob: the
829 torch runtime runs the whole job as one child Job named
830 ``<name>-node-0`` whose indexed pods are the node ranks, so the
831 rank-``node`` pod's logs are fetched (rank 0 by default). If the pod
832 is no longer available (completed/deleted), falls back to CloudWatch
833 Logs where Container Insights stores application logs.
835 Args:
836 job_name: Name of the job
837 namespace: Namespace of the job
838 region: Region where the job is running
839 tail_lines: Number of lines to return
840 follow: Stream logs (not implemented yet)
841 since_hours: Hours to look back in CloudWatch (default 24)
842 node: Node rank to fetch for a distributed TrainJob (default 0)
844 Returns:
845 Log content as string
846 """
847 if follow:
848 raise NotImplementedError("Log streaming not yet implemented")
850 target_region = region or self.config.default_region
852 try:
853 return self._aws_client.get_job_logs(
854 job_name=job_name,
855 namespace=namespace,
856 region=target_region,
857 tail_lines=tail_lines,
858 )
859 except RuntimeError as e:
860 error_msg = str(e)
861 lowered = error_msg.lower()
862 # An unknown batch Job name may be a TrainJob; its pods run
863 # under the JobSet child Job instead.
864 if "not found" in lowered:
865 try:
866 return self._get_trainjob_node_logs(
867 job_name=job_name,
868 namespace=namespace,
869 region=target_region,
870 node=node,
871 tail_lines=tail_lines,
872 )
873 except Exception as tj_err:
874 logger.debug("TrainJob logs fallback failed: %s", tj_err)
875 # If the pod is gone or pending, try CloudWatch
876 if any(
877 hint in lowered
878 for hint in ["not found", "pending", "no pods", "terminated", "completed"]
879 ):
880 logger.info("Pod not available, falling back to CloudWatch Logs")
881 try:
882 return self._get_cloudwatch_logs(
883 job_name=job_name,
884 region=target_region,
885 tail_lines=tail_lines,
886 since_hours=since_hours,
887 )
888 except Exception as cw_err:
889 logger.debug("CloudWatch fallback failed: %s", cw_err)
890 raise RuntimeError(
891 f"{error_msg}\n\n"
892 f"CloudWatch Logs fallback also failed: {cw_err}\n"
893 f"Tip: Container logs appear in CloudWatch within a few minutes. "
894 f"If the job just finished, try again shortly."
895 ) from e
896 raise
898 def _get_trainjob_node_logs(
899 self,
900 job_name: str,
901 namespace: str,
902 region: str,
903 node: int,
904 tail_lines: int,
905 ) -> str:
906 """Fetch logs from the rank-``node`` pod of a TrainJob.
908 The shipped torch-distributed runtime materializes a TrainJob as a
909 JobSet with a single replicated Job ``<name>-node-0`` running
910 ``numNodes`` indexed pods; the completion index is the node rank
911 (``PET_NODE_RANK``). Raises if no rank-``node`` pod exists.
912 """
913 child_job = f"{job_name}-node-0"
914 pods_response = self._aws_client.get_job_pods(child_job, namespace, region)
915 pods = pods_response.get("pods", []) or []
916 if not pods:
917 raise RuntimeError(f"no pods found for TrainJob child job '{child_job}'")
919 target_pod = None
920 for pod in pods:
921 metadata = pod.get("metadata", {}) or {}
922 labels = metadata.get("labels", {}) or {}
923 if labels.get("batch.kubernetes.io/job-completion-index") == str(node):
924 target_pod = metadata.get("name")
925 break
926 if target_pod is None:
927 # Older control planes may omit the index label; fall back to the
928 # deterministic indexed-pod name prefix '<child>-<index>-'.
929 prefix = f"{child_job}-{node}-"
930 for pod in pods:
931 pod_name = (pod.get("metadata", {}) or {}).get("name", "")
932 if pod_name.startswith(prefix):
933 target_pod = pod_name
934 break
935 if target_pod is None:
936 raise RuntimeError(
937 f"no rank-{node} pod found for TrainJob '{job_name}' "
938 f"({len(pods)} pod(s) under child job '{child_job}'); "
939 f"list them with: gco jobs pods {child_job} -n {namespace} -r {region}"
940 )
942 logger.info(
943 "Job '%s' resolved as TrainJob; fetching logs from rank-%d pod '%s'",
944 job_name,
945 node,
946 target_pod,
947 )
948 response = self._aws_client.get_pod_logs(
949 job_name=child_job,
950 pod_name=target_pod,
951 namespace=namespace,
952 region=region,
953 tail_lines=tail_lines,
954 )
955 return str(response.get("logs", ""))
957 def _get_cloudwatch_logs(
958 self,
959 job_name: str,
960 region: str,
961 tail_lines: int = 100,
962 since_hours: int = 24,
963 ) -> str:
964 """
965 Fetch job logs from CloudWatch Logs (Container Insights).
967 The CloudWatch Observability addon ships container stdout/stderr to:
968 /aws/containerinsights/{cluster_name}/application
970 Args:
971 job_name: Name of the job (used to filter log streams)
972 region: AWS region
973 tail_lines: Number of log lines to return
974 since_hours: Hours to look back (default 24)
976 Returns:
977 Log content as string
978 """
979 cluster_name = f"{self.config.project_name}-{region}"
980 log_group = f"/aws/containerinsights/{cluster_name}/application"
982 logs_client = self._aws_client._session.client("logs", region_name=region)
984 import time
986 now = int(time.time())
987 start_time = now - (since_hours * 3600)
989 query = (
990 f"fields @timestamp, @message "
991 f'| filter @logStream like "{job_name}" '
992 f"| sort @timestamp asc "
993 f"| limit {tail_lines}"
994 )
996 start_query = logs_client.start_query(
997 logGroupName=log_group,
998 startTime=start_time,
999 endTime=now,
1000 queryString=query,
1001 )
1002 query_id = start_query["queryId"]
1004 # Poll for results (CloudWatch Insights is async)
1005 result = None
1006 for _ in range(30): # up to 30 seconds
1007 time.sleep(1)
1008 result = logs_client.get_query_results(queryId=query_id)
1009 if result["status"] in ("Complete", "Failed", "Cancelled"):
1010 break
1012 if result is None or result["status"] != "Complete":
1013 status = result["status"] if result else "unknown"
1014 raise RuntimeError(
1015 f"CloudWatch Logs query did not complete (status: {status}). Try again in a moment."
1016 )
1018 if not result["results"]:
1019 raise RuntimeError(
1020 f"No logs found in CloudWatch for job '{job_name}' "
1021 f"in the last {since_hours} hours (log group: {log_group}). "
1022 f"Logs may take 1-2 minutes to appear after a pod runs. "
1023 f"Use --since to search further back, or check the job name "
1024 f"with: gco jobs list -r {region}"
1025 )
1027 # Extract log messages from results.
1028 # CloudWatch Container Insights wraps logs in a JSON envelope:
1029 # {"time":"...","stream":"stdout","log":"actual message","kubernetes":{...}}
1030 # We parse out the "log" field for clean output, falling back to the
1031 # raw message if it's not JSON.
1032 import json as _json
1034 lines = []
1035 for row in result["results"]:
1036 for entry in row:
1037 if entry["field"] == "@message":
1038 raw = entry["value"].rstrip()
1039 try:
1040 parsed = _json.loads(raw)
1041 message = parsed.get("log", raw) if isinstance(parsed, dict) else raw
1042 lines.append(message.rstrip() if isinstance(message, str) else raw)
1043 except ValueError, TypeError:
1044 lines.append(raw)
1045 break
1047 header = f"[CloudWatch Logs — {log_group}]\n"
1048 return header + "\n".join(lines)
1050 def delete_job(
1051 self,
1052 job_name: str,
1053 namespace: str,
1054 region: str | None = None,
1055 expected_uid: str | None = None,
1056 ) -> dict[str, Any]:
1057 """
1058 Delete a job.
1060 A name the batch Job endpoint does not know (404) is retried as a
1061 Kubeflow TrainJob through the generic manifests endpoint. Deleting
1062 the TrainJob cascades to its JobSet, child Jobs, and pods via owner
1063 references. ``expected_uid`` (a batch-Job concurrency guard) is not
1064 supported on the TrainJob path.
1066 Args:
1067 job_name: Name of the job
1068 namespace: Namespace of the job
1069 region: Region where the job is running
1071 Returns:
1072 Deletion result
1073 """
1074 target_region = region or self.config.default_region
1075 try:
1076 return self._aws_client.delete_job(
1077 job_name=job_name,
1078 namespace=namespace,
1079 region=target_region,
1080 expected_uid=expected_uid,
1081 )
1082 except requests.exceptions.HTTPError as e:
1083 if e.response is not None and e.response.status_code == 404:
1084 trainjob_result = self._delete_trainjob(job_name, namespace, target_region)
1085 if trainjob_result is not None:
1086 return trainjob_result
1087 raise
1089 def _delete_trainjob(self, job_name: str, namespace: str, region: str) -> dict[str, Any] | None:
1090 """Delete a TrainJob via the manifests endpoint, or None if that fails.
1092 Returning None (rather than raising) lets the caller surface the
1093 original batch-Job 404, which is the right error for a typo'd name.
1094 """
1095 try:
1096 return self._aws_client.call_api(
1097 "DELETE",
1098 f"/api/v1/manifests/{quote(namespace, safe='')}/{quote(job_name, safe='')}",
1099 region=region,
1100 params={"api_version": TRAINJOB_API_VERSION, "kind": "TrainJob"},
1101 )
1102 except RuntimeError as e:
1103 logger.debug("TrainJob delete fallback failed for %s: %s", job_name, e)
1104 return None
1106 def wait_for_job(
1107 self,
1108 job_name: str,
1109 namespace: str,
1110 region: str | None = None,
1111 timeout_seconds: int = 3600,
1112 poll_interval: int = 10,
1113 progress_callback: Callable[[JobInfo, int], None] | None = None,
1114 ) -> JobInfo:
1115 """
1116 Wait for a job to complete with progress reporting.
1118 Args:
1119 job_name: Name of the job
1120 namespace: Namespace of the job
1121 region: Region where the job is running
1122 timeout_seconds: Maximum time to wait
1123 poll_interval: Seconds between status checks
1124 progress_callback: Optional callable(JobInfo, elapsed_seconds) for progress updates.
1125 If None, a default stderr progress line is printed.
1127 Returns:
1128 Final JobInfo
1130 Raises:
1131 TimeoutError: If job doesn't complete within timeout
1132 """
1133 import sys
1134 import time
1136 start_time = time.time()
1138 while True:
1139 job = self.get_job(job_name, namespace, region)
1141 if job is None:
1142 raise ValueError(f"Job {job_name} not found in namespace {namespace}")
1144 elapsed = time.time() - start_time
1145 elapsed_str = _format_duration(int(elapsed))
1147 if job.is_complete:
1148 # Clear the progress line and return
1149 sys.stderr.write("\r\033[K")
1150 sys.stderr.flush()
1151 return job
1153 # Build progress message
1154 pods_info = (
1155 f"{job.active_pods} active, {job.succeeded_pods}/{job.completions} succeeded"
1156 )
1157 if job.failed_pods:
1158 pods_info += f", {job.failed_pods} failed"
1160 status_line = f" ⏳ {job.status.capitalize()} — {pods_info} — {elapsed_str} elapsed"
1162 if progress_callback:
1163 progress_callback(job, int(elapsed))
1164 else:
1165 # Overwrite the same line on stderr
1166 sys.stderr.write(f"\r\033[K{status_line}")
1167 sys.stderr.flush()
1169 if elapsed >= timeout_seconds:
1170 sys.stderr.write("\r\033[K")
1171 sys.stderr.flush()
1172 raise TimeoutError(
1173 f"Job {job_name} did not complete within {timeout_seconds} seconds "
1174 f"(last status: {job.status}, pods: {pods_info})"
1175 )
1177 time.sleep(poll_interval) # nosemgrep: arbitrary-sleep - intentional polling delay
1179 def submit_job_sqs(
1180 self,
1181 manifests: str | list[dict[str, Any]],
1182 region: str,
1183 namespace: str | None = None,
1184 labels: dict[str, str] | None = None,
1185 priority: int = 0,
1186 ) -> dict[str, Any]:
1187 """
1188 Submit a job to a regional SQS queue for processing.
1190 This is the recommended way to submit jobs as it:
1191 - Decouples submission from processing
1192 - Enables KEDA-based autoscaling
1193 - Provides better fault tolerance
1195 Args:
1196 manifests: Path to manifest file/directory or list of manifest dicts
1197 region: Target region for job submission (required)
1198 namespace: Fallback namespace for manifests that don't declare
1199 their own. When set, each manifest's ``metadata.namespace`` is
1200 filled in only if missing — existing values are preserved so
1201 users who've declared a target namespace in the manifest can
1202 rely on it reaching the queue processor untouched. Server-side
1203 validation enforces the allowlist.
1204 labels: Additional labels to add to manifests
1205 priority: Job priority (higher = more important)
1207 Returns:
1208 Submission result dictionary with message_id and queue info
1209 """
1210 import boto3
1212 # Load manifests if path provided
1213 manifest_list = self.load_manifests(manifests) if isinstance(manifests, str) else manifests
1215 # Apply namespace as a fallback only — preserve any namespace the
1216 # manifest declared itself.
1217 if namespace:
1218 for manifest in manifest_list:
1219 if "metadata" not in manifest:
1220 manifest["metadata"] = {}
1221 manifest["metadata"].setdefault("namespace", namespace)
1223 # Apply additional labels
1224 if labels:
1225 for manifest in manifest_list:
1226 if "metadata" not in manifest:
1227 manifest["metadata"] = {}
1228 if "labels" not in manifest["metadata"]:
1229 manifest["metadata"]["labels"] = {}
1230 manifest["metadata"]["labels"].update(labels)
1232 # Get queue URL from stack
1233 stack = self._aws_client.get_regional_stack(region)
1234 if not stack:
1235 raise ValueError(f"No GCO stack found in region {region}")
1237 # Get queue URL from CloudFormation outputs
1238 cfn = boto3.client("cloudformation", region_name=region)
1239 response = cfn.describe_stacks(StackName=stack.stack_name)
1240 outputs = {
1241 o["OutputKey"]: o["OutputValue"] for o in response["Stacks"][0].get("Outputs", [])
1242 }
1243 queue_url = outputs.get("JobQueueUrl")
1245 if not queue_url:
1246 raise ValueError(f"Job queue not found in stack {stack.stack_name}")
1248 # Build the envelope through the shared protocol module, which the
1249 # in-cluster consumer and the CI end-to-end test also import, so the
1250 # producer cannot drift from the reader.
1251 message = build_job_message(manifest_list, namespace=namespace, priority=priority)
1253 # Send to SQS
1254 sqs = boto3.client("sqs", region_name=region)
1255 response = sqs.send_message(
1256 QueueUrl=queue_url,
1257 MessageBody=message.json(),
1258 MessageAttributes=message.attributes,
1259 )
1261 # Get job name from first manifest
1262 job_name = None
1263 for manifest in manifest_list:
1264 if manifest.get("kind") == "Job":
1265 job_name = manifest.get("metadata", {}).get("name")
1266 break
1268 return {
1269 "status": "queued",
1270 "method": "sqs",
1271 "message_id": response["MessageId"],
1272 "job_id": message.job_id,
1273 "job_name": job_name,
1274 "queue_url": queue_url,
1275 "region": region,
1276 "namespace": message.namespace,
1277 "priority": priority,
1278 }
1280 def get_queue_status(self, region: str) -> dict[str, Any]:
1281 """
1282 Get the status of the job queue in a region.
1284 Args:
1285 region: AWS region
1287 Returns:
1288 Queue status including message counts
1289 """
1290 import boto3
1292 stack = self._aws_client.get_regional_stack(region)
1293 if not stack:
1294 raise ValueError(f"No GCO stack found in region {region}")
1296 # Get queue URLs from CloudFormation outputs
1297 cfn = boto3.client("cloudformation", region_name=region)
1298 response = cfn.describe_stacks(StackName=stack.stack_name)
1299 outputs = {
1300 o["OutputKey"]: o["OutputValue"] for o in response["Stacks"][0].get("Outputs", [])
1301 }
1303 queue_url = outputs.get("JobQueueUrl")
1304 dlq_url = outputs.get("JobDlqUrl")
1306 if not queue_url:
1307 raise ValueError(f"Job queue not found in stack {stack.stack_name}")
1309 sqs = boto3.client("sqs", region_name=region)
1311 # Get main queue attributes
1312 queue_attrs = sqs.get_queue_attributes(
1313 QueueUrl=queue_url,
1314 AttributeNames=[
1315 "ApproximateNumberOfMessages",
1316 "ApproximateNumberOfMessagesNotVisible",
1317 "ApproximateNumberOfMessagesDelayed",
1318 ],
1319 )["Attributes"]
1321 result = {
1322 "region": region,
1323 "queue_url": queue_url,
1324 "messages_available": int(queue_attrs.get("ApproximateNumberOfMessages", 0)),
1325 "messages_in_flight": int(queue_attrs.get("ApproximateNumberOfMessagesNotVisible", 0)),
1326 "messages_delayed": int(queue_attrs.get("ApproximateNumberOfMessagesDelayed", 0)),
1327 }
1329 # Get DLQ attributes if available
1330 if dlq_url:
1331 dlq_attrs = sqs.get_queue_attributes(
1332 QueueUrl=dlq_url,
1333 AttributeNames=["ApproximateNumberOfMessages"],
1334 )["Attributes"]
1335 result["dlq_url"] = dlq_url
1336 result["dlq_messages"] = int(dlq_attrs.get("ApproximateNumberOfMessages", 0))
1338 return result
1340 def list_jobs_global(
1341 self,
1342 namespace: str | None = None,
1343 status: str | None = None,
1344 limit: int = 50,
1345 ) -> dict[str, Any]:
1346 """
1347 List jobs across all regions via the global API endpoint.
1349 This uses the cross-region aggregator Lambda to query all regional
1350 clusters in parallel and return a unified view.
1352 Args:
1353 namespace: Filter by namespace
1354 status: Filter by status
1355 limit: Maximum jobs to return
1357 Returns:
1358 Aggregated job list with region information
1359 """
1360 return self._aws_client.get_global_jobs(
1361 namespace=namespace,
1362 status=status,
1363 limit=limit,
1364 )
1366 def get_global_health(self) -> dict[str, Any]:
1367 """
1368 Get health status across all regions.
1370 Returns:
1371 Aggregated health status from all regional clusters
1372 """
1373 return self._aws_client.get_global_health()
1375 def get_global_status(self) -> dict[str, Any]:
1376 """
1377 Get cluster status across all regions.
1379 Returns:
1380 Aggregated status from all regional clusters
1381 """
1382 return self._aws_client.get_global_status()
1384 def bulk_delete_global(
1385 self,
1386 namespace: str | None = None,
1387 status: str | None = None,
1388 older_than_days: int | None = None,
1389 label_selector: str | None = None,
1390 dry_run: bool = True,
1391 ) -> dict[str, Any]:
1392 """
1393 Bulk delete jobs across all regions.
1395 Args:
1396 namespace: Filter by namespace
1397 status: Filter by status
1398 older_than_days: Delete jobs older than N days
1399 label_selector: Kubernetes label selector
1400 dry_run: If True, only return what would be deleted
1402 Returns:
1403 Deletion results from all regions
1404 """
1405 return self._aws_client.bulk_delete_global(
1406 namespace=namespace,
1407 status=status,
1408 older_than_days=older_than_days,
1409 label_selector=label_selector,
1410 dry_run=dry_run,
1411 )
1413 def get_job_events(
1414 self,
1415 job_name: str,
1416 namespace: str,
1417 region: str | None = None,
1418 ) -> dict[str, Any]:
1419 """
1420 Get Kubernetes events for a job.
1422 Args:
1423 job_name: Name of the job
1424 namespace: Namespace of the job
1425 region: Region where the job is running
1427 Returns:
1428 Events related to the job
1429 """
1430 return self._aws_client.get_job_events(
1431 job_name=job_name,
1432 namespace=namespace,
1433 region=region or self.config.default_region,
1434 )
1436 def get_job_pods(
1437 self,
1438 job_name: str,
1439 namespace: str,
1440 region: str | None = None,
1441 ) -> dict[str, Any]:
1442 """
1443 Get pods for a job.
1445 Args:
1446 job_name: Name of the job
1447 namespace: Namespace of the job
1448 region: Region where the job is running
1450 Returns:
1451 Pod details for the job
1452 """
1453 return self._aws_client.get_job_pods(
1454 job_name=job_name,
1455 namespace=namespace,
1456 region=region or self.config.default_region,
1457 )
1459 def get_pod_logs(
1460 self,
1461 job_name: str,
1462 pod_name: str,
1463 namespace: str,
1464 region: str | None = None,
1465 tail_lines: int = 100,
1466 container: str | None = None,
1467 ) -> dict[str, Any]:
1468 """
1469 Get logs from a specific pod of a job.
1471 Args:
1472 job_name: Name of the job
1473 pod_name: Name of the pod
1474 namespace: Namespace of the job
1475 region: Region where the job is running
1476 tail_lines: Number of lines to return
1477 container: Container name (for multi-container pods)
1479 Returns:
1480 Pod logs response
1481 """
1482 return self._aws_client.get_pod_logs(
1483 job_name=job_name,
1484 pod_name=pod_name,
1485 namespace=namespace,
1486 region=region or self.config.default_region,
1487 tail_lines=tail_lines,
1488 container=container,
1489 )
1491 def get_job_metrics(
1492 self,
1493 job_name: str,
1494 namespace: str,
1495 region: str | None = None,
1496 ) -> dict[str, Any]:
1497 """
1498 Get resource metrics for a job.
1500 Args:
1501 job_name: Name of the job
1502 namespace: Namespace of the job
1503 region: Region where the job is running
1505 Returns:
1506 Resource usage metrics for the job's pods
1507 """
1508 return self._aws_client.get_job_metrics(
1509 job_name=job_name,
1510 namespace=namespace,
1511 region=region or self.config.default_region,
1512 )
1514 def retry_job(
1515 self,
1516 job_name: str,
1517 namespace: str,
1518 region: str | None = None,
1519 ) -> dict[str, Any]:
1520 """
1521 Retry a failed job.
1523 Creates a new job from the failed job's spec with a new name.
1525 Args:
1526 job_name: Name of the failed job
1527 namespace: Namespace of the job
1528 region: Region where the job is running
1530 Returns:
1531 Result with new job name
1532 """
1533 return self._aws_client.retry_job(
1534 job_name=job_name,
1535 namespace=namespace,
1536 region=region or self.config.default_region,
1537 )
1539 def bulk_delete_jobs(
1540 self,
1541 namespace: str | None = None,
1542 status: str | None = None,
1543 older_than_days: int | None = None,
1544 label_selector: str | None = None,
1545 region: str | None = None,
1546 dry_run: bool = True,
1547 ) -> dict[str, Any]:
1548 """
1549 Bulk delete jobs in a region.
1551 Args:
1552 namespace: Filter by namespace
1553 status: Filter by status
1554 older_than_days: Delete jobs older than N days
1555 label_selector: Kubernetes label selector
1556 region: Target region
1557 dry_run: If True, only return what would be deleted
1559 Returns:
1560 Deletion results
1561 """
1562 return self._aws_client.bulk_delete_jobs(
1563 namespace=namespace,
1564 status=status,
1565 older_than_days=older_than_days,
1566 label_selector=label_selector,
1567 region=region or self.config.default_region,
1568 dry_run=dry_run,
1569 )
1572def get_job_manager(config: GCOConfig | None = None) -> JobManager:
1573 """Get a configured job manager instance."""
1574 return JobManager(config)