ops_utils.tdr_utils.tdr_job_utils

  1"""Utilities for interacting with TDR jobs."""
  2import json
  3import logging
  4import time
  5from typing import Callable, Optional, Any
  6
  7from ..vars import ARG_DEFAULTS
  8
  9
 10class MonitorTDRJob:
 11    """A class to monitor the status of a TDR job until completion."""
 12
 13    def __init__(self, tdr: Any, job_id: str, check_interval: int, return_json: bool):
 14        """
 15        Initialize the MonitorTDRJob class.
 16
 17        **Args:**
 18        - tdr (`ops_utils.tdr_utils.tdr_api_utils.TDR`): An instance of the TDR class.
 19        - job_id (str): The ID of the job to be monitored.
 20        - check_interval (int): The interval in seconds to wait between status checks.
 21        - return_json (bool): Whether to get and return the result of the job as json.
 22        """
 23        self.tdr = tdr
 24        """@private"""
 25        self.job_id = job_id
 26        """@private"""
 27        self.check_interval = check_interval
 28        """@private"""
 29        self.return_json = return_json
 30        """@private"""
 31
 32    def _raise_for_failed_job(self) -> None:
 33        """
 34        Raise an error with useful information if the job has failed.
 35
 36        Raises:
 37            ValueError: If the job has failed.
 38        """
 39        job_result = self.tdr.get_job_result(self.job_id, expect_failure=True)
 40        raise Exception(
 41            f"Status code {job_result.status_code}: {job_result.text}")
 42
 43    def run(self) -> Optional[dict]:
 44        """
 45        Monitor the job until completion.
 46
 47        **Returns:**
 48        - dict: The job results
 49        """
 50        while True:
 51            ingest_response = self.tdr.get_job_status(self.job_id)
 52            if ingest_response.status_code == 202:
 53                logging.info(f"TDR job {self.job_id} is still running")
 54                # Check every x seconds if ingest is still running
 55                time.sleep(self.check_interval)
 56            elif ingest_response.status_code == 200:
 57                response_json = json.loads(ingest_response.text)
 58                if response_json["job_status"] == "succeeded":
 59                    logging.info(f"TDR job {self.job_id} succeeded")
 60                    if self.return_json:
 61                        request = self.tdr.get_job_result(self.job_id)
 62                        return json.loads(request.text)
 63                    # If not returning json, return None
 64                    return None
 65                else:
 66                    logging.error(f"TDR job {self.job_id} failed")
 67                    self._raise_for_failed_job()
 68            else:
 69                logging.error(f"TDR job {self.job_id} failed")
 70                self._raise_for_failed_job()
 71
 72
 73class SubmitAndMonitorMultipleJobs:
 74    """A class to batch submit and monitor TDR jobs."""
 75
 76    def __init__(
 77            self, tdr: Any,
 78            job_function: Callable,
 79            job_args_list: list[tuple],
 80            batch_size: int = ARG_DEFAULTS["batch_size"],  # type: ignore[assignment]
 81            check_interval: int = ARG_DEFAULTS["waiting_time_to_poll"],  # type: ignore[assignment]
 82            verbose: bool = False
 83    ):
 84        """
 85        Initialize the SubmitAndMonitorMultipleJobs class.
 86
 87        **Args:**
 88        - tdr (`ops_utils.tdr_utils.tdr_api_utils.TDR`): An instance of the TDR class.
 89        - job_function (Callable): The function to submit a job.
 90        - job_args_list (list[tuple]): A list of tuples containing the arguments for each job.
 91        - batch_size (int, optional): The number of jobs to process in each batch. Defaults to `500`.
 92        - check_interval (int, optional): The interval in seconds to wait between status checks. Defaults to `90`.
 93        - verbose (bool, optional): Whether to log detailed information about each job. Defaults to `False`.
 94        """
 95        self.tdr = tdr
 96        """@private"""
 97        self.job_function = job_function
 98        """@private"""
 99        self.job_args_list = job_args_list
100        """@private"""
101        self.batch_size = batch_size
102        """@private"""
103        self.check_interval = check_interval
104        """@private"""
105        self.verbose = verbose
106        """@private"""
107
108    def run(self) -> None:
109        """
110        Run the process to submit and monitor multiple jobs in batches.
111
112        Logs the progress and status of each batch and job.
113
114        Failed jobs are collected and printed out at the end of processing.
115        """
116        failed_jobs = []
117        total_jobs = len(self.job_args_list)
118        logging.info(f"Processing {total_jobs} {self.job_function.__name__} jobs in batches of {self.batch_size}")
119
120        # Process jobs in batches
121        for i in range(0, total_jobs, self.batch_size):
122            job_ids = []
123            current_batch = self.job_args_list[i:i + self.batch_size]
124            logging.info(
125                f"Submitting jobs for batch {i // self.batch_size + 1} with {len(current_batch)} jobs."
126            )
127
128            # Submit jobs for the current batch
129            for job_args in current_batch:
130                # Submit job with arguments and store the job ID
131                job_id = self.job_function(*job_args).json()["id"]
132                if self.verbose:
133                    logging.info(f"Submitted job {job_id} with args {job_args}")
134                job_ids.append(job_id)
135
136            # Monitor jobs for the current batch
137            logging.info(f"Monitoring {len(current_batch)} jobs in batch {i // self.batch_size + 1}")
138            for job_id in job_ids:
139                try:
140                    MonitorTDRJob(
141                        tdr=self.tdr,
142                        job_id=job_id,
143                        check_interval=self.check_interval,
144                        return_json=False
145                    ).run()
146                except Exception as e:
147                    logging.error(f"Job {job_id} failed: {e}")
148                    failed_jobs.append(job_id)
149
150            logging.info(f"Completed batch {i // self.batch_size + 1} with {len(current_batch)} jobs.")
151
152        logging.info(f"Successfully processed {total_jobs - len(failed_jobs)} jobs.")
153
154        if len(failed_jobs) > 0:
155            raise Exception(
156                f"The following job IDs failed: {', '.join(failed_jobs)}")