Save job outputs to Mongo

cdmtaskservice/jobflows/nersc_jaws.py

@@ -256,10 +256,9 @@
                     )
                 outfiles.append(models.S3FileOutput(file=o.path, etag=o.e_tag))
             # TODO DISKSPACE will need to clean up job results @ NERSC
-            # TODO NEXT save output files to mongo and remove log line
-            logr.info(f"outfiles:\n{outfiles}")
-            await self._mongo.update_job_state(
+            await self._mongo.add_output_files_to_job(
                 job.id,
+                outfiles,
                 models.JobState.UPLOAD_SUBMITTED,
                 models.JobState.COMPLETE,
                 # TODO TEST will need a way to mock out timestamps

cdmtaskservice/models.py

@@ -45,6 +45,7 @@
 FLD_NERSC_DETAILS_UL_TASK_ID = "upload_task_id"
 FLD_JOB_JAWS_DETAILS = "jaws_details"
 FLD_JAWS_DETAILS_RUN_ID = "run_id"
+FLD_JOB_OUTPUTS = "outputs"
 
 
 # https://en.wikipedia.org/wiki/Filename#Comparison_of_filename_limitations

cdmtaskservice/mongo.py

@@ -144,6 +144,7 @@
         state: models.JobState,
         time: datetime.datetime,
         push: dict[str, Any] | None = None,
+        set_: dict[str, Any] | None = None,
     ):
         res = await self._col_jobs.update_one(
             {
@@ -156,7 +157,7 @@
                           (_not_falsy(state, "state").value, _not_falsy(time, "time")
                     )
                 },
-                "$set": {models.FLD_JOB_STATE: state.value}
+                "$set": (set_ if set_ else {}) | {models.FLD_JOB_STATE: state.value}
             },
         )
         if not res.matched_count:
@@ -250,6 +251,27 @@
         await self._update_job_state(job_id, current_state, state, time, push={
             self._FLD_NERSC_UL_TASK: _require_string(task_id, "task_id")
         })
+
+    async def add_output_files_to_job(
+        self,
+        job_id: str,
+        output: list[models.S3FileOutput],
+        current_state: models.JobState,
+        state: models.JobState,
+        time: datetime.datetime
+    ):
+        """
+        Add output files to a job and update the state.
+
+        Arguments are as update_job_state except for the addition of:
+
+        output - the output files.
+        """
+        out = [o.model_dump() for o in _not_falsy(output, "output")]
+        await self._update_job_state(
+            job_id, current_state, state, time, set_={models.FLD_JOB_OUTPUTS: out}
+        )
+
 
 class NoSuchImageError(Exception):
     """ The image does not exist in the system. """

cdmtaskservice/nersc/manager.py

@@ -532,7 +532,6 @@ async def upload_JAWS_job_files(
         outputs = jaws_output.parse_outputs_json(outputs_json)
         container_files = list(outputs.output_files.keys())
         presigns = await files_to_urls(container_files)
-        # TODO CORRECTNESS IMPORTANT log etags in remote code, upload and return
         # TODO LOGGING upload job stdout and stderr logs
         return await self.upload_s3_files(
             job.id,