Merge pull request #150 from kbase/dev-service

Save JAWS run ID in mongo after submit

Author: MrCreosote

cdmtaskservice/jobflows/nersc_jaws.py

@@ -98,11 +98,11 @@ async def start_job(self, job: models.Job, objmeta: list[S3ObjectMeta]):
             # remote cluster and have job_state handle choosing the correct mongo method & params
             # to run
             await self._mongo.add_NERSC_download_task_id(
-                # TODO TEST will need a way to mock out timestamps
                 job.id,
                 task_id,
                 models.JobState.CREATED,
                 models.JobState.DOWNLOAD_SUBMITTED,
+                # TODO TEST will need a way to mock out timestamps
                 timestamp.utcdatetime(),
             )
         except Exception as e:
@@ -137,8 +137,16 @@ async def _submit_jaws_job(self, job: models.AdminJobDetails):
         try:
             # TODO PERF configure file download concurrency
             jaws_job_id = await self._nman.run_JAWS(job)
-            # TODO JAWS record job ID in DB
-            logr.info(f"JAWS job id: {jaws_job_id}")
+            # See notes above about adding the NERSC task id to the job
+            await self._mongo.add_JAWS_run_id(
+                job.id,
+                jaws_job_id,
+                models.JobState.JOB_SUBMITTING,
+                models.JobState.JOB_SUBMITTED,
+                # TODO TEST will need a way to mock out timestamps
+                timestamp.utcdatetime(),
+            )
+            # TDOO JOBRUN start polling JAWS to wait for job completion
         except Exception as e:
             # TODO LOGGING figure out how logging it going to work etc.
             logr.exception(f"Error starting JAWS job for job {job.id}")

cdmtaskservice/models.py

@@ -42,6 +42,8 @@
 FLD_JOB_STATE = "state"
 FLD_JOB_NERSC_DETAILS = "nersc_details"
 FLD_NERSC_DETAILS_DL_TASK_ID = "download_task_id"
+FLD_JOB_JAWS_DETAILS = "jaws_details"
+FLD_JAWS_DETAILS_RUN_ID = "run_id"
 
 
 # https://en.wikipedia.org/wiki/Filename#Comparison_of_filename_limitations
@@ -699,9 +701,18 @@ class NERSCDetails(BaseModel):
             + "completion in the SFAPI. Multiple tasks indicate job retries after failures.")]
 
 
+class JAWSDetails(BaseModel):
+    """
+    Details about a JAWS job run.
+    """
+    run_id: Annotated[list[str], Field(
+        description="Run IDs for a JGI JAWS job. Multiple run IDs indicate job retries after "
+            + "failures.")]
+
 class AdminJobDetails(Job):
     """
     Information about a job with added details of interest to service administrators.
     """
     nersc_details: NERSCDetails | None = None
     # TODO NERSC more details, logs, etc.
+    jaws_details: JAWSDetails | None = None

cdmtaskservice/mongo.py

@@ -195,19 +195,39 @@ async def add_NERSC_download_task_id(
         """
         Add a download task_id to the NERSC section of a job and update the state.
         
-        job_id - the job ID.
+        Arguments are as update_job_state except for the addition of:
+        
         task_id - the NERSC task ID.
-        current_state - the expected current state of the job. If the job is not in this state
-            an error is thrown.
-        state - the new state for the job.
-        time - the time at which the job transitioned to the new state.
         """
         # may need to make this more generic where the cluster is passed in and mapped to
         # a job structure location or something if we support more than NERSC
         await self._update_job_state(job_id, current_state, state, time, push={
             self._FLD_NERSC_DL_TASK: _require_string(task_id, "task_id")
         })
 
+    _FLD_JAWS_RUN_ID = f"{models.FLD_JOB_JAWS_DETAILS}.{models.FLD_JAWS_DETAILS_RUN_ID}"
+
+    async def add_JAWS_run_id(
+        self,
+        job_id: str,
+        run_id: str,
+        current_state: models.JobState,
+        state: models.JobState,
+        time: datetime.datetime
+    ):
+        """
+        Add a run ID to the JAWS section of a job and update the state.
+        
+        Arguments are as update_job_state except for the addition of:
+        
+        run_id - the JAWS run ID.
+        """
+        # may need to make this more generic where the cluster is passed in and mapped to
+        # a job structure location or something if we support more than NERSC
+        await self._update_job_state(job_id, current_state, state, time, push={
+            self._FLD_JAWS_RUN_ID: _require_string(run_id, "run_id")
+        })
+
 
 class NoSuchImageError(Exception):
     """ The image does not exist in the system. """

cdmtaskservice/nersc/manager.py

@@ -408,6 +408,10 @@ async def run_JAWS(self, job: models.Job, file_download_concurrency: int = 10) -
         perl = await cli.compute(Machine.perlmutter)
         pre = self._perlmutter_scratch / _CTS_SCRATCH_ROOT_DIR / job.id
         try:
+            # TODO PERF this copies all the files to the jaws staging area, and so could take
+            #           a long time. Hardlink them in first to avoid the copy. Also look into
+            #           caching between jobs, so multiple jobs on the same file don't DL it over
+            #           and over
             res = await perl.run(_JAWS_COMMAND_TEMPLATE.format(
                 job_id=job.id,
                 wdlpath=pre / _JAWS_INPUT_WDL,
@@ -431,7 +435,7 @@ async def run_JAWS(self, job: models.Job, file_download_concurrency: int = 10) -
             logging.getLogger(__name__).info(
                 f"Submitted JAWS job with run id {run_id} for job {job.id}"
             )
-            return run_id
+            return str(run_id)
         except json.JSONDecodeError as e:
             raise ValueError(f"JAWS returned invalid JSON: {e}\n{res}") from e