Move NERSC job/refdata transfers from SFAPI tasks to Perlmutter Slurm jobs

cdmtaskservice/app_state.py

@@ -332,7 +332,7 @@ async def _register_nersc_job_flows(
         kafka_notifier,
         coman,
         cfg.service_group,
-        cfg.service_root_url
+        cfg.service_root_url,
     )
     dest.register("JAWS flow provider", jaws_job_flows.close())
     flowman.register_flow(NERSCJAWSRunner.CLUSTER, jaws_job_flows.get_nersc_job_flow)

cdmtaskservice/config.py

@@ -48,7 +48,6 @@ class CDMTaskServiceConfig:
     nersc_jaws_user: str - the user name of the user associated with the NERSC and JAWS
         credentials.
     jaws_refdata_root_dir: str - the JAWS refdata root directory to use for refdata storage.
-    jaws_staging_dir_dtn: str - the JAWS staging directory for the `kbase` site on a NERSC DTN.
     jaws_staging_dir_prl: str - the JAWS staging directory for the `kbase` site on the NERSC
         Perlmutter system.
     sfapi_cred_path: str - the path to a NERSC Superfacility API credential file. The file is
@@ -153,15 +152,10 @@ def __init__(self, config_file: BinaryIO, version: str):
         self.jaws_refdata_root_dir = _get_string_required(
             config, _SEC_NERSC_JAWS, "refdata_root_dir"
         )
-        # These typically have $PSCRATCH for the kbjaws kbase site user embedded in it,
+        # This typically has $PSCRATCH for the kbjaws kbase site user embedded in it,
         # which I'd prefer not to specify literally in a config. Not sure if there's a better
         # way to deal with it since the service doesn't know anything about the jaws site user.
-        # Worry about it later 
-        # Also not a fan of having to specify both but so far the other solutions I've considered
-        # are even uglier
-        self.jaws_staging_dir_dtn = _get_string_required(
-            config, _SEC_NERSC_JAWS, "jaws_staging_dir_dtn"
-        )
+        # Worry about it later
         self.jaws_staging_dir_prl = _get_string_required(
             config, _SEC_NERSC_JAWS, "jaws_staging_dir_perlmutter"
         )
@@ -262,7 +256,6 @@ def __init__(self, config_file: BinaryIO, version: str):
         self._nersc_paths = NERSCPaths(  # fail early if paths fail validation
             f"{self.nersc_remote_code_dir}/{version}",
             self.jaws_refdata_root_dir,
-            self.jaws_staging_dir_dtn,
             self.jaws_staging_dir_prl
         )
         self._check_path_overlap()
@@ -347,7 +340,6 @@ def print_config(self, output: TextIO):
             f"Authentication refdata service role: {self.refdata_service_role}",
             f"NERSC / JAWS user: {self.nersc_jaws_user}",
             f"NERSC / JAWS refdata root dir: {self.jaws_refdata_root_dir}",
-            f"NERSC / JAWS DTN staging dir: {self.jaws_staging_dir_dtn}",
             f"NERSC / JAWS Perlmutter staging dir: {self.jaws_staging_dir_prl}",
             f"NERSC client credential path: {self.sfapi_cred_path}",
             f"NERSC remote code dir: {self.nersc_remote_code_dir}",

cdmtaskservice/jobflows/jaws_flows_provider.py

@@ -79,9 +79,9 @@ async def create(  # lot of arguments, but I'm not seeing a way to simplify that
     ):
         """
         WARNING: this class is not thread safe.
-        
+
         Create the JAWS based job flows provider.
-        
+
         superfacility_api_credential_path - a path to an SFAPI credential file.
             The first line of the file must be the client ID, and the rest the client
             secret in PEM format.
@@ -111,7 +111,7 @@ async def create(  # lot of arguments, but I'm not seeing a way to simplify that
         jfp._coman = _not_falsy(coman, "coman")
         jfp._service_group = _require_string(service_group, "service_group")
         jfp._service_root_url = _require_string(service_root_url, "service_root_url")
-        
+
         # setup other variables
         jfp._logr = logging.getLogger(__name__)
         jfp._nersc_status_cli = NERSCStatus()

cdmtaskservice/jobflows/nersc_jaws.py

@@ -265,14 +265,14 @@ async def start_job(self, job: models.Job, objmeta: list[S3ObjectMeta]):
             callback_url = get_download_complete_callback(self._callback_root, job.id)
             # TODO PERF config / set concurrency
             # TODO DISKSPACE will need to clean up job downloads @ NERSC
-            task_id = await self._nman.download_s3_files(
+            job_id = await self._nman.download_s3_files(
                 job.id, objmeta, presigned, callback_url, insecure_ssl=self._s3insecure
             )
             # Hmm. really this should go through job state but that seems pointless right now.
             # May need to refactor this and the mongo method later to be more generic to
             # remote cluster and have job_state handle choosing the correct mongo method & params
             # to run
-            await self._updates.update_job_state(job.id, submitted_nersc_download(task_id))
+            await self._updates.update_job_state(job.id, submitted_nersc_download(job_id))
         except Exception as e:
             await self._updates.handle_exception(e, job.id, "starting file download for")
 
@@ -296,7 +296,7 @@ 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)
-            # See notes above about adding the NERSC task id to the job
+            # See notes above about adding the NERSC job id to the job
             await self._updates.update_job_state(job.id, submitted_jaws_job(jaws_job_id))
         except Exception as e:
             if jaws_job_id:
@@ -385,14 +385,14 @@ async def presign(output_files: list[Path]) -> list[PresignedPost]:
 
         try:
             # TODO PERF config / set concurrency
-            task_id = await self._nman.upload_JAWS_log_files_on_error(
+            job_id = await self._nman.upload_JAWS_log_files_on_error(
                 job,
                 jaws_info["output_dir"],
                 presign,
                 get_error_log_upload_complete_callback(self._callback_root, job.id),
                 insecure_ssl=self._s3insecure,
             )
-            await self._updates.update_job_state(job.id, submitted_nersc_error_processing(task_id))
+            await self._updates.update_job_state(job.id, submitted_nersc_error_processing(job_id))
         except Exception as e:
             await self._updates.handle_exception(e, job.id, "starting error processing for")
 
@@ -407,15 +407,15 @@ async def presign(output_files: list[Path], crc64nvmes: list[str]) -> list[Presi
 
         try:
             # TODO PERF config / set concurrency
-            task_id = await self._nman.upload_JAWS_job_files(
+            job_id = await self._nman.upload_JAWS_job_files(
                 job,
                 jaws_info["output_dir"],
                 presign,
                 get_upload_complete_callback(self._callback_root, job.id),
                 insecure_ssl=self._s3insecure,
             )
-            # See notes above about adding the NERSC task id to the job
-            await self._updates.update_job_state(job.id, submitted_nersc_upload(task_id))
+            # See notes above about adding the NERSC job id to the job
+            await self._updates.update_job_state(job.id, submitted_nersc_upload(job_id))
         except Exception as e:
             await self._updates.handle_exception(e, job.id, "starting file upload for")
 
@@ -588,7 +588,7 @@ async def stage_refdata(self, refdata: models.ReferenceData, objmeta: S3ObjectMe
             )
             # TODO DISKSPACE clean up no longer used refdata @ NERSC
             #                keep the refdata mongo record so it can be restaged if necessary
-            task_id = await self._nman.download_s3_files(
+            job_id = await self._nman.download_s3_files(
                 refdata.id,
                 [objmeta],
                 presigned,
@@ -598,7 +598,7 @@ async def stage_refdata(self, refdata: models.ReferenceData, objmeta: S3ObjectMe
                 unpack=refdata.unpack,
             )
             await self._updates.update_refdata_state(
-                refdata.id, submitted_nersc_refdata_download(task_id)
+                refdata.id, submitted_nersc_refdata_download(job_id)
             )
         except Exception as e:
             await self._updates.handle_exception(

cdmtaskservice/logfields.py

@@ -17,7 +17,7 @@
 TRANS_ID = "trans_id"
 REFDATA_ID = "refdata_id"
 NERSC_STATUS = "nersc_status"
-NERSC_TASK_ID = "nersc_task"
+NERSC_JOB_ID = "nersc_job"
 JAWS_RUN_ID = "jaws_id"
 NEXT_ACTION_SEC = "next_action_sec"
 REMOTE_ERROR = "remote_error"

cdmtaskservice/models.py

@@ -49,9 +49,12 @@
 FLD_JOB_STATE_TRANSITION_NOTIFICATION_SENT = "notif_sent"
 FLD_JOB_ADMIN_META = "admin_meta"
 FLD_JOB_NERSC_DETAILS = "nersc_details"
-FLD_NERSC_DETAILS_DL_TASK_ID = "download_task_id"
-FLD_NERSC_DETAILS_UL_TASK_ID = "upload_task_id"
-FLD_NERSC_DETAILS_LOG_UL_TASK_ID = "log_upload_task_id"
+FLD_NERSC_DETAILS_DL_TASK_ID = "download_task_id"  # deprecated, see FLD_NERSC_DETAILS_DL_JOB_ID
+FLD_NERSC_DETAILS_DL_JOB_ID = "download_job_id"
+FLD_NERSC_DETAILS_UL_TASK_ID = "upload_task_id"  # deprecated, see FLD_NERSC_DETAILS_UL_JOB_ID
+FLD_NERSC_DETAILS_UL_JOB_ID = "upload_job_id"
+FLD_NERSC_DETAILS_LOG_UL_TASK_ID = "log_upload_task_id"  # deprecated, see *_LOG_UL_JOB_ID
+FLD_NERSC_DETAILS_LOG_UL_JOB_ID = "log_upload_job_id"
 FLD_JOB_JAWS_DETAILS = "jaws_details"
 FLD_JAWS_DETAILS_RUN_ID = "run_id"
 FLD_JOB_HTC_CLUSTER_ID = "cluster_id"
@@ -68,7 +71,8 @@
 FLD_REFDATA_FILE = "file"
 FLD_REFDATA_STATUSES = "statuses"
 FLD_REFDATA_CLUSTER = "cluster"
-FLD_REFDATA_NERSC_DL_TASK_ID = "nersc_download_task_id"
+FLD_REFDATA_NERSC_DL_TASK_ID = "nersc_download_task_id"  # deprecated, see FLD_REFDATA_NERSC_DL_JOB_ID
+FLD_REFDATA_NERSC_DL_JOB_ID = "nersc_download_job_id"
 # Fields that are shared between multiple models for consistency
 # Currently refdata, jobs, and subjobs
 FLD_COMMON_ID = "id"
@@ -1169,24 +1173,43 @@ class NERSCDetails(BaseModel):
     Details about a job run at NERSC.
     """
     # Output only model, no validation
-    download_task_id: Annotated[list[str], Field(
-        description="IDs for tasks run via the NERSC SFAPI to download files from an S3 "
-            + "instance to NERSC. Note that task details only persist for ~10 minutes past "
-            + "completion in the SFAPI. Multiple tasks indicate job retries after failures."
+    download_task_id: Annotated[list[str] | None, Field(
+        default=None,
+        deprecated="Replaced by download_job_id. Only present on records created before NERSC "
+            + "downloads were migrated from SFAPI async tasks to Slurm jobs.",
+        description="IDs for SFAPI download tasks. Deprecated - see download_job_id."
     )]
-    upload_task_id: Annotated[list[str], Field(
-        default_factory=list,
-        description="IDs for tasks run via the NERSC SFAPI to upload files to an S3 "
-            + "instance from NERSC. Note that task details only persist for ~10 minutes past "
-            + "completion in the SFAPI. Multiple tasks indicate job retries after failures."
-            + "Empty if an upload task has not yet been submitted to NERSC."
+    download_job_id: Annotated[list[str] | None, Field(
+        default=None,
+        description="IDs for NERSC Slurm jobs run via the NERSC SFAPI to download files from "
+            + "an S3 instance to NERSC. Multiple job IDs indicate job retries after failures. "
+            + "Missing if the record predates the migration to Slurm jobs, see download_task_id."
     )]
-    log_upload_task_id: Annotated[list[str], Field(
-        default_factory=list,
-        description="IDs for tasks run via the NERSC SFAPI to upload log files to an S3 "
-            + "instance from NERSC. Note that task details only persist for ~10 minutes past "
-            + "completion in the SFAPI. Multiple tasks indicate job retries after failures."
-            + "Empty if a log upload task has not yet been submitted to NERSC."
+    upload_task_id: Annotated[list[str] | None, Field(
+        default=None,
+        deprecated="Replaced by upload_job_id. Only present on records created before NERSC "
+            + "uploads were migrated from SFAPI async tasks to Slurm jobs.",
+        description="IDs for SFAPI upload tasks. Deprecated - see upload_job_id."
+    )]
+    upload_job_id: Annotated[list[str] | None, Field(
+        default=None,
+        description="IDs for NERSC Slurm jobs run via the NERSC SFAPI to upload files to an S3 "
+            + "instance from NERSC. Multiple job IDs indicate job retries after failures. "
+            + "Missing if an upload job has not yet been submitted to NERSC, or if the record "
+            + "predates the migration to Slurm jobs, see upload_task_id."
+    )]
+    log_upload_task_id: Annotated[list[str] | None, Field(
+        default=None,
+        deprecated="Replaced by log_upload_job_id. Only present on records created before NERSC "
+            + "log uploads were migrated from SFAPI async tasks to Slurm jobs.",
+        description="IDs for SFAPI log upload tasks. Deprecated - see log_upload_job_id."
+    )]
+    log_upload_job_id: Annotated[list[str] | None, Field(
+        default=None,
+        description="IDs for NERSC Slurm jobs run via the NERSC SFAPI to upload log files to "
+            + "an S3 instance from NERSC. Multiple job IDs indicate job retries after failures. "
+            + "Missing if a log upload job has not yet been submitted to NERSC, or if the record "
+            + "predates the migration to Slurm jobs, see log_upload_task_id."
     )]
 
 
@@ -1406,11 +1429,17 @@ class AdminReferenceDataStatus(ReferenceDataStatus):
             + "download manifests and results, etc."
     )] = False
     nersc_download_task_id: Annotated[list[str] | None, Field(
-        default_factory=list,
-        description="IDs for tasks run via the NERSC SFAPI to download files from an S3 "
-            + "instance to NERSC. Note that task details only persist for ~10 minutes past "
-            + "completion in the SFAPI. Multiple tasks indicate job retries after failures. "
-            + "Only present if the refdata is being downloaded to NERSC."
+        default=None,
+        deprecated="Replaced by nersc_download_job_id. Only present on records created before "
+            + "NERSC refdata downloads were migrated from SFAPI async tasks to Slurm jobs.",
+        description="IDs for SFAPI download tasks. Deprecated - see nersc_download_job_id."
+    )]
+    nersc_download_job_id: Annotated[list[str] | None, Field(
+        default=None,
+        description="IDs for NERSC Slurm jobs run via the NERSC SFAPI to download files from "
+            + "an S3 instance to NERSC. Multiple job IDs indicate job retries after failures. "
+            + "Only present if the refdata is being downloaded to NERSC via a Slurm job, i.e. "
+            + "the record postdates the migration to Slurm jobs, see nersc_download_task_id."
     )]
     admin_error: Annotated[str | None, Field(
         examples=["The back fell off"],

cdmtaskservice/mongo.py

@@ -639,7 +639,7 @@ async def _update_job_state(
             pre_doc, job_id, update, time, subjob_id, recovery_cooldown, last_update_time
         )
 
-    _FLD_NERSC_DL_TASK = f"{models.FLD_JOB_NERSC_DETAILS}.{models.FLD_NERSC_DETAILS_DL_TASK_ID}"
+    _FLD_NERSC_DL_JOB = f"{models.FLD_JOB_NERSC_DETAILS}.{models.FLD_NERSC_DETAILS_DL_JOB_ID}"
     _FLD_JAWS_RUN_ID = f"{models.FLD_JOB_JAWS_DETAILS}.{models.FLD_JAWS_DETAILS_RUN_ID}"
     _FLD_HTC_CLUSTER_ID = f"{models.FLD_COMMON_HTC_DETAILS}.{models.FLD_JOB_HTC_CLUSTER_ID}"
     _FLD_HTC_CPU_HOURS = f"{models.FLD_COMMON_HTC_DETAILS}.{models.FLD_COMMON_HTC_CPU_HOURS}"
@@ -648,13 +648,13 @@ async def _update_job_state(
     _FLD_HTC_STATS_INCOMPLETE = (
         f"{models.FLD_COMMON_HTC_DETAILS}.{models.FLD_JOB_HTC_STATS_INCOMPLETE}"
     )
-    _FLD_NERSC_UL_TASK = f"{models.FLD_JOB_NERSC_DETAILS}.{models.FLD_NERSC_DETAILS_UL_TASK_ID}"
-    _FLD_NERSC_LOG_UL_TASK = (
-        f"{models.FLD_JOB_NERSC_DETAILS}.{models.FLD_NERSC_DETAILS_LOG_UL_TASK_ID}"
+    _FLD_NERSC_UL_JOB = f"{models.FLD_JOB_NERSC_DETAILS}.{models.FLD_NERSC_DETAILS_UL_JOB_ID}"
+    _FLD_NERSC_LOG_UL_JOB = (
+        f"{models.FLD_JOB_NERSC_DETAILS}.{models.FLD_NERSC_DETAILS_LOG_UL_JOB_ID}"
     )
     def _setup_field_mappings(self):
         self._FIELD_TO_KEY_AND_PUSH = {
-            UpdateField.NERSC_DOWNLOAD_TASK_ID: (self._FLD_NERSC_DL_TASK, True),
+            UpdateField.NERSC_DOWNLOAD_JOB_ID: (self._FLD_NERSC_DL_JOB, True),
             UpdateField.JAWS_RUN_ID: (self._FLD_JAWS_RUN_ID, True),
             UpdateField.HTCONDOR_CLUSTER_ID: (self._FLD_HTC_CLUSTER_ID, True),
             UpdateField.HTCONDOR_CPU_HOURS: (self._FLD_HTC_CPU_HOURS, False),
@@ -664,10 +664,10 @@ def _setup_field_mappings(self):
             UpdateField.CPU_HOURS: (models.FLD_COMMON_CPU_HOURS, False),
             UpdateField.CPU_FACTOR: (models.FLD_JOB_CPU_FACTOR, False),
             UpdateField.MAX_MEMORY: (models.FLD_COMMON_MAX_MEM, False),
-            UpdateField.NERSC_UPLOAD_TASK_ID: (self._FLD_NERSC_UL_TASK, True),
+            UpdateField.NERSC_UPLOAD_JOB_ID: (self._FLD_NERSC_UL_JOB, True),
             UpdateField.OUTPUT_FILE_PATHS: (models.FLD_COMMON_OUTPUTS, False),
             UpdateField.OUTPUT_FILE_COUNT: (models.FLD_JOB_OUTPUT_FILE_COUNT, False),
-            UpdateField.NERSC_LOG_UPLOAD_TASK_ID: (self._FLD_NERSC_LOG_UL_TASK, True),
+            UpdateField.NERSC_LOG_UPLOAD_JOB_ID: (self._FLD_NERSC_LOG_UL_JOB, True),
             UpdateField.EXIT_CODE: (models.FLD_SUBJOB_EXIT_CODE, False),
             UpdateField.RUNTIME: (models.FLD_SUBJOB_RUNTIME, False),
             UpdateField.USER_ERROR: (models.FLD_COMMON_ERROR, False),
@@ -676,7 +676,7 @@ def _setup_field_mappings(self):
             UpdateField.LOG_PATH: (models.FLD_JOB_LOGPATH, False),
         }
         self._REFDATA_FIELD_TO_KEY_AND_PUSH = {
-            UpdateField.NERSC_DOWNLOAD_TASK_ID: (models.FLD_REFDATA_NERSC_DL_TASK_ID, True),
+            UpdateField.NERSC_DOWNLOAD_JOB_ID: (models.FLD_REFDATA_NERSC_DL_JOB_ID, True),
             UpdateField.USER_ERROR: (models.FLD_COMMON_ERROR, False),
             UpdateField.ADMIN_ERROR: (models.FLD_COMMON_ADMIN_ERROR, False),
             UpdateField.TRACEBACK: (models.FLD_COMMON_TRACEBACK, False),
@@ -1371,7 +1371,7 @@ async def save_refdata(self, refdata: models.AdminReferenceData):
         # to ensure unique IDs
 
         # TDOO REFDATA add a force option to allow for file overwrites if needed
-        r = refdata.model_dump()
+        r = refdata.model_dump(exclude_none=True)
         # Could add a check in the refdata model that rds have > 0 statuses,
         # statuses have > 0 transitions and the last
         # transition == the redfdata cluster state... probably not necessary.
@@ -1395,7 +1395,7 @@ async def add_refdata_site(self, refdata_id: str, rds: models.AdminReferenceData
         refdata_id - the ID of the refdata to modify.
         rds - the information for the new site.
         """
-        s = _not_falsy(rds, "rds").model_dump()
+        s = _not_falsy(rds, "rds").model_dump(exclude_none=True)
         s[_FLD_UPDATE_TIME] = rds.transition_times[-1].time
         result = await self._col_refdata.update_one(
             {