diff --git a/cdmtaskservice/app_state.py b/cdmtaskservice/app_state.py index 619bc84..e9c5723 100644 --- a/cdmtaskservice/app_state.py +++ b/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) diff --git a/cdmtaskservice/config.py b/cdmtaskservice/config.py index a9c50b6..4b6dd7a 100644 --- a/cdmtaskservice/config.py +++ b/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}", diff --git a/cdmtaskservice/jobflows/jaws_flows_provider.py b/cdmtaskservice/jobflows/jaws_flows_provider.py index b146bde..f3b7b63 100644 --- a/cdmtaskservice/jobflows/jaws_flows_provider.py +++ b/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() diff --git a/cdmtaskservice/jobflows/nersc_jaws.py b/cdmtaskservice/jobflows/nersc_jaws.py index b9f162c..101607b 100644 --- a/cdmtaskservice/jobflows/nersc_jaws.py +++ b/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( diff --git a/cdmtaskservice/logfields.py b/cdmtaskservice/logfields.py index d36a368..3f9bceb 100644 --- a/cdmtaskservice/logfields.py +++ b/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" diff --git a/cdmtaskservice/models.py b/cdmtaskservice/models.py index 155e059..1aa60b0 100644 --- a/cdmtaskservice/models.py +++ b/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"], diff --git a/cdmtaskservice/mongo.py b/cdmtaskservice/mongo.py index 7eb2383..6a1878a 100644 --- a/cdmtaskservice/mongo.py +++ b/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( { diff --git a/cdmtaskservice/nersc/manager.py b/cdmtaskservice/nersc/manager.py index 280b87c..1783c41 100644 --- a/cdmtaskservice/nersc/manager.py +++ b/cdmtaskservice/nersc/manager.py @@ -5,7 +5,6 @@ import asyncio from collections.abc import Callable from enum import Enum -from httpx import HTTPStatusError import io import inspect import json @@ -50,15 +49,32 @@ # servers on a different version. Note in admin docs that old unused installs # can be deleted. Server code is tiny anyway -_DT_TARGET = Machine.dtns -# TODO NERSCUPDATE delete the following line when DTN downloads work normally. -# See https://nersc.servicenowservices.com/sp?sys_id=ad33e85f1b5a5610ac81a820f54bcba0&view=sp&id=ticket&table=incident -_DT_WORKAROUND = "source /etc/bashrc" - -_COMMAND_PATH = "utilities/command" - _MIN_TIMEOUT_SEC = 300 _SEC_PER_GB = 2 * 60 # may want to make this configurable +# Large file transfers are submitted as Slurm jobs on the `xfer` QOS (NERSC's data transfer +# queue) rather than SFAPI async tasks, since tasks are capped at 10 minutes of execution time. +# The `xfer` QOS runs on Perlmutter login nodes, requires `--licenses=SCRATCH`, forbids `-N`/ +# `--nodes`, and has a 48 hour wall time ceiling. See https://docs.nersc.gov/jobs/policy/ +_XFER_QOS = "xfer" +_MAX_SBATCH_TIME_SEC = 48 * 60 * 60 +# Buffer applied to the size-based transfer time estimate before submitting the Slurm job, so +# a hung transfer still hits Slurm's own wall time limit rather than running indefinitely and +# tying up one of the `xfer` QOS' limited (15) concurrent job slots per user. +# +# Kept small (1.5x) rather than large because the estimate it's applied to is already very +# conservative on its own: +# * _SEC_PER_GB (2 min/GB, ~67 Mbps) is a slow rate for a transfer expected to run over ESnet +# between two DOE-facility endpoints. +# * The estimate sums bytes across every file in the batch with no credit for concurrency, i.e. +# it assumes the whole batch transfers serially even though multiple files transfer at once. +# * Actual hangs on individual files/connections are caught independently and much sooner by +# the per-file timeout in s3/remote.py's _timeout(), so this buffer isn't the primary hang +# detector - it mainly needs to absorb variance in the rate estimate, not find hangs itself. +_TIME_BUFFER_MULTIPLIER = 1.5 +_STAGING_DIR_NAME = "staged" +# `ssh dtn` resolves to a NERSC Data Transfer Node from a Perlmutter login node without further +# configuration, so no configurable hostname is needed here. +_DTN_HOST = "dtn" _JOB_MANIFESTS = Path("manifests") _MANIFEST_FILE_PREFIX = "manifest-" @@ -84,29 +100,25 @@ # TODO PERF add start and end time to task output and log / record in db / put in result file) -# TODO NERSCFEATURE if NERSC puts python on the dtns revert to regular load -# Pinned rather than floating on `latest`: pip dependencies (see _install_pip_dependencies) -# are installed once, at server startup, into the site-packages of whatever Python version -# is resolved at that moment. Every later job run resolves `module load python` again from -# scratch. If `latest` moves to a different default Python version in between (NERSC ships -# a new dated PE release roughly monthly), the two resolutions disagree and previously -# installed packages become invisible to the newer interpreter, e.g. NoModuleFoundError: -# awscrt at job runtime despite a successful install at startup. -# To check the current released versions / find a newer one to pin to, on a NERSC -# DTN node run `ls -l /global/common/software/nersc/pe/modulefiles/` (the `latest` entry -# there is a symlink to the currently recommended release). After updating this constant, -# restart the CTS server so _install_pip_dependencies reinstalls under the newly pinned -# version. -_NERSC_PE_MODULEFILES_VERSION = "26.8.1" -_PYTHON_LOAD_HACK = f"module use /global/common/software/nersc/pe/modulefiles/{_NERSC_PE_MODULEFILES_VERSION}" +# Pinned to a specific NERSC `python` module rather than a bare `module load python` (which +# floats to whatever NERSC currently defaults to): pip dependencies (see +# _install_pip_dependencies) are installed once, at server startup, into the site-packages of +# whatever Python version is resolved at that moment. Every later job run resolves the same +# module load again from scratch. If the default moved to a different Python version in +# between, the two resolutions disagree and previously installed packages become invisible to +# the newer interpreter, e.g. NoModuleFoundError: awscrt at job runtime despite a successful +# install at startup. +# To check for newer versions to pin to, run `module avail python` on a NERSC Perlmutter login +# node. After updating this constant, restart the CTS server so _install_pip_dependencies +# reinstalls under the newly pinned version. +_PYTHON_MODULE = "python/3.13-26.8.0" _RUN_CTS_REMOTE_CODE_FILENAME = "run_cts_remote_code.sh" # Might want to make a shared constants module for all these env var names and update this # file and remote.py _RUN_CTS_REMOTE_CODE = f""" #!/usr/bin/env bash -{_PYTHON_LOAD_HACK} -module load python +module load {_PYTHON_MODULE} export PYTHONPATH=$CTS_CODE_LOCATION export CTS_MODE=$CTS_MODE @@ -115,6 +127,11 @@ export CTS_CONTAINER_LOGS_LOCATION=$CTS_CONTAINER_LOGS_LOCATION export CTS_JAWS_OUTPUT_DIR=$CTS_JAWS_OUTPUT_DIR export CTS_CHECKSUM_FILE_LOCATION=$CTS_CHECKSUM_FILE_LOCATION +export CTS_STAGING_DIR=$CTS_STAGING_DIR +export CTS_DTN_HOST=$CTS_DTN_HOST +export CTS_REFDATA_DEST_DIR=$CTS_REFDATA_DEST_DIR +export CTS_COMPLETION_FILE_LOCATION=$CTS_COMPLETION_FILE_LOCATION +export CTS_COMPLETION_FILE_CONTENTS=$CTS_COMPLETION_FILE_CONTENTS export CTS_RESULT_FILE_LOCATION=$CTS_RESULT_FILE_LOCATION export CTS_LOG_FILE_LOCATION=$CTS_LOG_FILE_LOCATION export CTS_CALLBACK_URL=$CTS_CALLBACK_URL @@ -132,6 +149,11 @@ echo "CTS_CONTAINER_LOGS_LOCATION=[$CTS_CONTAINER_LOGS_LOCATION]" echo "CTS_JAWS_OUTPUT_DIR=[$CTS_JAWS_OUTPUT_DIR]" echo "CTS_CHECKSUM_FILE_LOCATION=[$CTS_CHECKSUM_FILE_LOCATION]" +echo "CTS_STAGING_DIR=[$CTS_STAGING_DIR]" +echo "CTS_DTN_HOST=[$CTS_DTN_HOST]" +echo "CTS_REFDATA_DEST_DIR=[$CTS_REFDATA_DEST_DIR]" +echo "CTS_COMPLETION_FILE_LOCATION=[$CTS_COMPLETION_FILE_LOCATION]" +echo "CTS_COMPLETION_FILE_CONTENTS=[$CTS_COMPLETION_FILE_CONTENTS]" echo "CTS_RESULT_FILE_LOCATION=[$CTS_RESULT_FILE_LOCATION]" echo "CTS_LOG_FILE_LOCATION=[$CTS_LOG_FILE_LOCATION]" echo "CTS_CALLBACK_URL=[$CTS_CALLBACK_URL]" @@ -142,14 +164,23 @@ echo "python exited with code $?" """ +# Submitted to the `xfer` QOS on Perlmutter. `-N`/`--nodes` is forbidden on this QOS since it +# runs on shared login nodes rather than dedicated compute nodes. +_SBATCH_SCRIPT_TEMPLATE = f"""#!/usr/bin/env bash +#SBATCH --qos={_XFER_QOS} +#SBATCH --licenses=SCRATCH +#SBATCH --time={{time}} + +{{body}} +""" + # Note there's a race condition that theoretically could happen here if the path is removed # after the existence check but before the rm. Since this is just for clean up not an issue. # Also note I wasted way too much time trying to make this fail cleanly if there was a write # protected file in the path tree, which should never happen. # It'll still fail if it really can't delete the path. -_REMOVE_DTN_PATH_TEMPLATE = f""" -{_DT_WORKAROUND} +_REMOVE_PATH_TEMPLATE = """ if [ ! -e "{{path}}" ]; then exit 0 fi @@ -181,6 +212,22 @@ def _get_dependencies(mod: ModuleType, cts_dep: set[ModuleType], pip_dep: set[Mo _get_dependencies(remote, _CTS_DEPENDENCIES, _PIP_DEPENDENCIES) +def _compute_sbatch_time_sec(total_bytes: int) -> int: + """ + Compute the wall time, in seconds, to request for a `xfer` QOS Slurm job transferring + total_bytes of data, capped at the QOS' 48 hour maximum. + """ + estimate = max(_MIN_TIMEOUT_SEC, _SEC_PER_GB * total_bytes / 1_000_000_000) + return int(min(_MAX_SBATCH_TIME_SEC, estimate * _TIME_BUFFER_MULTIPLIER)) + + +def _seconds_to_slurm_time(seconds: int) -> str: + """ Convert a number of seconds to a Slurm `--time` argument in HH:MM:SS format. """ + hours, remainder = divmod(seconds, 3600) + minutes, secs = divmod(remainder, 60) + return f"{hours:02d}:{minutes:02d}:{secs:02d}" + + class TransferState(Enum): """ The state of the transfer. """ @@ -222,7 +269,7 @@ async def create( ) -> Self: """ Create the NERSC manager. - + client_provider - a function that provides a valid SFAPI client. It is assumed that the user associated with the client does not change. nersc_paths - the set of paths for NERSC manager use. @@ -237,7 +284,7 @@ async def create( nm = NERSCManager(client_provider, nersc_paths, jaws_config.user, service_group) await nm._setup_remote_code(nersc_paths, jaws_config.token, jaws_config.group) return nm - + def __init__( self, client_provider: Callable[[], str], @@ -282,19 +329,12 @@ async def _setup_remote_code(self, nersc_paths: NERSCPaths, jaws_token: str, jaw ), chmod = "600" )) - pm_scratch = tg.create_task(self._set_up_perlmutter_scratch(cli)) - dtn_scratch = tg.create_task(self._set_up_dtn_scratch(cli)) + pm_scratch = tg.create_task(self._set_up_perlmutter_scratch()) if _PIP_DEPENDENCIES: - tg.create_task(self._install_pip_dependencies(cli)) - self._dtn_scratch = dtn_scratch.result() + tg.create_task(self._install_pip_dependencies(perlmutter)) self._perlmutter_scratch = pm_scratch.result() - self._nersc_dtn_file_cache_path = self._make_cache_path(nersc_paths.jaws_staging_dir_dtns) self._nersc_perlmutter_file_cache_path = self._make_cache_path( nersc_paths.jaws_staging_dir_perlmutter) - logr.info( - "NERSC DTN JAWS staging cache path", - extra={logfields.FILE: self._nersc_dtn_file_cache_path} - ) logr.info( "NERSC perlmutter JAWS staging cache path", extra={logfields.FILE: self._nersc_perlmutter_file_cache_path} @@ -312,54 +352,48 @@ def _make_cache_path(self, jaws_staging_dir: Path): / "cache" ) - async def _install_pip_dependencies(self, client: AsyncClient): + async def _install_pip_dependencies(self, compute: AsyncCompute): logr = logging.getLogger(__name__) deps = " ".join( # may need to do something else if module doesn't have __version__ [f"{mod.__name__}=={mod.__version__}" for mod in _PIP_DEPENDENCIES]) logr.info(f"Installing pip modules at NERSC: {deps}") command = ( - f"{_DT_WORKAROUND}; " - + f"{_PYTHON_LOAD_HACK}; " - + f"module load python; " + f"module load {_PYTHON_MODULE}; " # Unlikely, but this could cause problems if multiple versions - # of the server are running at once. Don't worry about it for now + # of the server are running at once. Don't worry about it for now + f"pip install {deps}" # adding notapackage causes a failure ) - dt = await client.compute(_DT_TARGET) - await dt.run(command) + await compute.run(command) logr.info(f"Installed pip modules at NERSC") - - - async def _set_up_perlmutter_scratch(self, client: AsyncClient) -> Path: + + + async def _set_up_perlmutter_scratch(self) -> Path: logr = logging.getLogger(__name__) logr.info("Getting Perlmutter scratch path from NERSC") - perlmutter = await client.compute(Machine.perlmutter) - scratch = (await perlmutter.run("echo $SCRATCH")).strip() - logr.info("NERSC perlmutter scratch path", extra={logfields.FILE: scratch}) - return Path(scratch) - - async def _set_up_dtn_scratch(self, client: AsyncClient) -> Path: - logr = logging.getLogger(__name__) - logr.info("Getting DTN scratch path from NERSC") - dt = await client.compute(_DT_TARGET) - scratch = await dt.run(f"{_DT_WORKAROUND}; echo $SCRATCH") - scratch = scratch.strip() + cli = self._client_provider() + compute = await cli.compute(Machine.perlmutter) + scratch = (await compute.run("echo $SCRATCH")).strip() if not scratch: # have had issues here previously - raise ValueError("Unable to determine $SCRATCH variable for NERSC dtns") - logr.info("NERSC DTN scratch path", extra={logfields.FILE: scratch}) + raise ValueError("Unable to determine $SCRATCH variable for NERSC perlmutter") + logr.info("NERSC perlmutter scratch path", extra={logfields.FILE: scratch}) return Path(scratch) - - def _get_job_scratch(self, job_id, perlmutter=False) -> Path: - sc = self._perlmutter_scratch if perlmutter else self._dtn_scratch - return sc / self._work_loc / _JOBS_DIR / job_id - + + def _get_job_scratch(self, job_id) -> Path: + return self._perlmutter_scratch / self._work_loc / _JOBS_DIR / job_id + def _get_refdata_scratch(self, refdata_id) -> Path: - return self._dtn_scratch / self._work_loc / _REFDATA_DIR / refdata_id - + return self._perlmutter_scratch / self._work_loc / _REFDATA_DIR / refdata_id + + def _get_refdata_staging_loc(self, refdata_id) -> Path: + # Perlmutter login nodes (where `xfer` QOS jobs run) cannot write to the DTN-only + # refdata root, so refdata is staged here first and copied to its final location by a + # DTN-side step at the end of the job. + return self._get_refdata_scratch(refdata_id) / _STAGING_DIR_NAME + def _get_refdata_loc(self, refdata_id) -> Path: return self._refdata_root / self._get_relative_refdata_loc(refdata_id) - + def _get_relative_refdata_loc(self, refdata_id) -> Path: return self._work_loc / refdata_id @@ -381,10 +415,6 @@ def _get_refdata_file_complete_path(self, refdata_id, site: sites.Cluster) -> Pa ) ) - async def _run_command(self, client: AsyncClient, machine: Machine, exe: str): - # TODO ERRORHANDlING deal with errors - return (await client.post(f"{_COMMAND_PATH}/{machine}", data={"executable": exe})).json() - async def _upload_file_to_nersc( self, compute: AsyncCompute, @@ -392,12 +422,11 @@ async def _upload_file_to_nersc( file: Path = None, bio: io.BytesIO = None, chmod: str = None, - ): + ): logr = logging.getLogger(__name__) logr.info("Uploading file to NERSC.", extra={logfields.FILE: target}) - dtw = f"{_DT_WORKAROUND}; " if compute.name == Machine.dtns else "" if target.parent != Path("."): - cmd = f"{dtw}mkdir -p {target.parent}" + cmd = f"mkdir -p {target.parent}" await compute.run(cmd) asrp = self._get_async_path(compute, target) # TODO ERRORHANDLING throw custom errors @@ -408,19 +437,19 @@ async def _upload_file_to_nersc( await asrp.upload(bio) logr.info("Upload of file to NERSC complete.", extra={logfields.FILE: target}) if chmod: - cmd = f"{dtw}chmod {chmod} {target}" + cmd = f"chmod {chmod} {target}" await compute.run(cmd) logr.info("chmod of uploaded file complete.", extra={logfields.FILE: target}) - async def _delete_dtn_paths(self, paths: list[Path]): + async def _delete_paths(self, paths: list[Path]): logr = logging.getLogger(__name__) cli = self._client_provider() - dtns = await cli.compute(Machine.dtns) + perlmutter = await cli.compute(Machine.perlmutter) async with asyncio.TaskGroup() as tg: for p in paths: logr.info("Deleting path at NERSC", extra={logfields.FILE: p}) # May need to catch SFAPI client errors and wrap. YAGNI for now - tg.create_task(dtns.run(_REMOVE_DTN_PATH_TEMPLATE.format(path=p))) + tg.create_task(perlmutter.run(_REMOVE_PATH_TEMPLATE.format(path=p))) def _get_async_path(self, compute: AsyncCompute, target: Path) -> AsyncRemotePath: # skip some API calls vs. the upload example in the NERSC docs @@ -429,36 +458,6 @@ def _get_async_path(self, compute: AsyncCompute, target: Path) -> AsyncRemotePat asrp.perms = "-" # hack to prevent an unnecessary network call return asrp - async def is_task_complete(self, task_id: str, max_retries=10) -> bool: - """ - Returns true if the the task is complete, signified by: - * The task data denoting it as complete, or - * The task not being available in the NERSC SFAPI (e.g. HTTP 404), as tasks are removed - 10 minutes after completion. - The task may be in a successful or errored state. Polls the task 1/s up to `max_retries`. - - Note that if a task ID never existed, this method will return that it is - complete, as there is no way to tell the difference. - """ - logr = logging.getLogger(__name__) - cli = self._client_provider() - retries = 1 - while True: - try: - task = await cli.get(f"tasks/{task_id}") - except HTTPStatusError as e: - if e.response.status_code == 404: - return True - raise - if task.json()["status"] == "completed": - return True - elif retries > max_retries: - return False - else: - logr.info(f"Polling state of task {task_id} in 1s, retry attempt {retries}") - retries += 1 - await asyncio.sleep(1) - async def download_s3_files( self, download_id: str, @@ -485,49 +484,98 @@ async def download_s3_files( refdata - whether this is a refdata download and files should be stored in the NERSC refdata location. unpack - whether to unpack *.gz, *.tar.gz, or *.tgz files. - - Returns the NERSC task ID for the download. + + Returns the NERSC Slurm job ID for the download. """ maniio = self._create_download_manifest( download_id, objects, presigned_urls, concurrency, insecure_ssl, refdata, unpack) + total_bytes = sum(o.size for o in objects) return await self._process_manifest( maniio, download_id, callback_url, "download_manifest.json", "download", - refdata=refdata, + total_bytes, + mode="refdata_manifest" if refdata else "manifest", ) - - async def upload_presigned_files( + + async def _upload_presigned_files( self, job_id: str, remote_files: list[Path], presigned_urls: list[PresignedPost], callback_url: str, + total_bytes: int, concurrency: int = 10, - insecure_ssl: bool = False + insecure_ssl: bool = False, ) -> str: """ Upload a set of files to presigned URLs from NERSC. - - job_id - the ID of the job for which the files are being transferred. + + job_id - the ID of the job for which the files are being transferred. This must be a unique ID, and no other transfers should be occurring for the job. remote_files - the files to upload. presigned_urls - the presigned upload URLs for each file, in the same order as the file. callback_url - the URL to GET as a callback for when the upload is complete. + total_bytes - the total size of the files to upload, used to compute the Slurm job's + wall time limit. concurrency - the number of simultaneous uploads to process. insecure_ssl - whether to skip the cert check for the S3 URL. - - Returns the NERSC task ID for the upload. + + Returns the NERSC Slurm job ID for the upload. """ maniio = self._create_upload_manifest( remote_files, presigned_urls, concurrency, insecure_ssl) return await self._process_manifest( - maniio, job_id, callback_url, "upload_manifest.json", "upload" + maniio, + job_id, + callback_url, + "upload_manifest.json", + "upload", + total_bytes, ) - + + def _build_process_manifest_command( + self, + entity_id: str, + callback_url: str, + mode: str, + manifestpath: Path, + task_base_path: Path, + error_json_file_location: str = None, + container_logs_location: str = None, + ) -> list[str]: + refdata = mode == "refdata_manifest" + command = [ + f"export CTS_MODE={mode}", + f"export CTS_CODE_LOCATION={self._nersc_code_path}", + f"export CTS_MANIFEST_LOCATION={manifestpath}", + f"export CTS_RESULT_FILE_LOCATION={task_base_path}_result.json", + f"export CTS_LOG_FILE_LOCATION={task_base_path}_log.txt", + f"export CTS_CALLBACK_URL={callback_url}", + f"export SCRATCH=$SCRATCH", + ] + if error_json_file_location: + command.append(f"export CTS_ERRORS_JSON_LOCATION={error_json_file_location}") + command.append(f"export CTS_CONTAINER_LOGS_LOCATION={container_logs_location}") + if refdata: + command.append(f"export CTS_STAGING_DIR={self._get_refdata_staging_loc(entity_id)}") + command.append(f"export CTS_DTN_HOST={_DTN_HOST}") + command.append(f"export CTS_REFDATA_DEST_DIR={self._get_refdata_loc(entity_id)}") + # TODO CLEANUP need to delete this and the JAWS written completion file + # see https://jaws-docs.jgi.doe.gov/en/latest/jaws/jaws_refdata.html#adding-data-to-refdata-directory + command.append( + "export CTS_COMPLETION_FILE_LOCATION=" + f"{self._get_refdata_file_change_path(entity_id)}" + ) + command.append( + f"export CTS_COMPLETION_FILE_CONTENTS={self._get_refdata_loc(entity_id)}" + ) + command.append(f'"$CTS_CODE_LOCATION"/{_RUN_CTS_REMOTE_CODE_FILENAME}') + return command + async def _process_manifest( self, manifest: io.BytesIO, @@ -535,45 +583,48 @@ async def _process_manifest( callback_url: str, filename: str, task_type: str, - mode="manifest", + total_bytes: int, + mode: str = "manifest", error_json_file_location: str = None, container_logs_location: str = None, # this is expected to be present if the above is - refdata: bool = False ): + refdata = mode == "refdata_manifest" if refdata: rootpath = self._get_refdata_scratch(entity_id) else: rootpath = self._get_job_scratch(entity_id) manifestpath = rootpath / filename cli = self._client_provider() - dt = await cli.compute(_DT_TARGET) + perl = await cli.compute(Machine.perlmutter) # TODO CLEANUP manifests after some period of time - await self._upload_file_to_nersc(dt, manifestpath, bio=manifest) - command = [ - f"{_DT_WORKAROUND}; ", - f"export CTS_MODE={mode}; ", - f"export CTS_CODE_LOCATION={self._nersc_code_path}; ", - f"export CTS_MANIFEST_LOCATION={manifestpath}; ", - f"export CTS_RESULT_FILE_LOCATION={rootpath / task_type}_result.json; ", - f"export CTS_LOG_FILE_LOCATION={rootpath / task_type}_log.txt; ", - f"export CTS_CALLBACK_URL={callback_url}; ", - f"export SCRATCH=$SCRATCH; ", - f'"$CTS_CODE_LOCATION"/{_RUN_CTS_REMOTE_CODE_FILENAME}', - ] - if error_json_file_location: - command.insert(2, f"export CTS_ERRORS_JSON_LOCATION={error_json_file_location}; ") - command.insert(3, f"export CTS_CONTAINER_LOGS_LOCATION={container_logs_location}; ") - command = "".join(command) - task_id = (await self._run_command(cli, _DT_TARGET, command))["task_id"] + await self._upload_file_to_nersc(perl, manifestpath, bio=manifest) + command = self._build_process_manifest_command( + entity_id, + callback_url, + mode, + manifestpath, + rootpath / task_type, + error_json_file_location=error_json_file_location, + container_logs_location=container_logs_location, + ) + script = _SBATCH_SCRIPT_TEMPLATE.format( + time=_seconds_to_slurm_time(_compute_sbatch_time_sec(total_bytes)), + body="\n".join(command), + ) + # upload script to make debugging easier + scriptpath = rootpath / f"{task_type}_submit.sh" + await self._upload_file_to_nersc(perl, scriptpath, bio=io.BytesIO(script.encode())) + job = await perl.submit_job(str(scriptpath)) + job_id = str(job.jobid) logging.getLogger(__name__).info( - f"Created {task_type} task for {'refdata' if refdata else 'job'}", + f"Submitted {task_type} Slurm job for {'refdata' if refdata else 'job'}", extra={ - logfields.NERSC_TASK_ID: task_id, + logfields.NERSC_JOB_ID: job_id, logfields.REFDATA_ID if refdata else logfields.JOB_ID: entity_id } ) - return task_id - + return job_id + def _create_download_manifest( self, download_id: str, @@ -593,13 +644,13 @@ def _create_download_manifest( raise ValueError("All the S3 objects must have a CRC64/NVME checksum") manifest = self._base_manifest("download", concurrency, insecure_ssl) if refdata: - sc = self._get_refdata_loc(download_id) - # TODO CLEANUP need to delete this and the JAWS written completion file - # see https://jaws-docs.jgi.doe.gov/en/latest/jaws/jaws_refdata.html#adding-data-to-refdata-directory - manifest["completion-file"] = str(self._get_refdata_file_change_path(download_id)) - manifest["completion-file-contents"] = str(sc) + # The final refdata location and completion file live under the DTN-only-writable + # refdata root, so they can't be written directly by this manifest (it runs on a + # Perlmutter login node). Files are staged here and moved into place, and the + # completion file written, by a DTN-side step at the end of the Slurm job. + sc = self._get_refdata_staging_loc(download_id) else: - manifest["cache-dir"] = str(self._nersc_dtn_file_cache_path) + manifest["cache-dir"] = str(self._nersc_perlmutter_file_cache_path) manifest["files"] = [] for url, meta in zip(presigned_urls, objects): fileman = { @@ -700,7 +751,7 @@ async def _get_transfer_result( sc = self._get_refdata_scratch(entity_id) if refdata else self._get_job_scratch(entity_id) path = sc / f"{op}_result.json" res = await self._download_json_file_from_NERSC( - Machine.dtns, path, no_exception_on_missing_file=True + Machine.perlmutter, path, no_exception_on_missing_file=True ) if not res: return TransferResult(state=TransferState.INCOMPLETE), None @@ -725,7 +776,7 @@ async def run_JAWS(self, job: models.Job, file_download_concurrency: int = 10) - cli = self._client_provider() await self._generate_and_load_job_files_to_nersc(cli, job, file_download_concurrency) perl = await cli.compute(Machine.perlmutter) - pre = self._get_job_scratch(job.id, perlmutter=True) + pre = self._get_job_scratch(job.id) try: res = await perl.run(_JAWS_COMMAND_TEMPLATE.format( job_id=job.id, @@ -775,7 +826,7 @@ async def _generate_and_load_job_files_to_nersc( downloads = {pre / fp: f for fp, f in zip(manifest_file_paths, manifest_files)} downloads[pre / _JAWS_INPUT_WDL] = wdljson.wdl downloads[pre / _JAWS_INPUT_JSON] = json.dumps(wdljson.input_json, indent=4) - dt = await cli.compute(_DT_TARGET) + perl = await cli.compute(Machine.perlmutter) semaphore = asyncio.Semaphore(concurrency) async def sem_coro(coro): async with semaphore: @@ -785,7 +836,7 @@ async def sem_coro(coro): async with asyncio.TaskGroup() as tg: for path, file in downloads.items(): coros.append(self._upload_file_to_nersc( - dt, path, bio=io.BytesIO(file.encode()) + perl, path, bio=io.BytesIO(file.encode()) )) tg.create_task(sem_coro(coros[-1])) except ExceptionGroup as eg: @@ -822,8 +873,8 @@ async def upload_JAWS_job_files( callback_url - the URL to GET as a callback for when the upload is complete. concurrency - the number of simultaneous uploads to process. insecure_ssl - whether to skip the cert check for the S3 URL. - - Returns the NERSC task ID for the upload. + + Returns the NERSC Slurm job ID for the upload. """ _not_falsy(job, "job") _not_falsy(files_to_urls, "files_to_urls") @@ -831,11 +882,10 @@ async def upload_JAWS_job_files( cburl = _require_string(callback_url, "callback_url") _check_num(concurrency, "concurrency") cli = self._client_provider() - dtns = await cli.compute(Machine.dtns) + perl = await cli.compute(Machine.perlmutter) rootpath = self._get_job_scratch(job.id) checksum_file = _CRC64NVME_CHECKSUMS_JSON_FILE_NAME command = [ # similar to the command in _process_manifest - f"{_DT_WORKAROUND}; ", f"export CTS_MODE=checksum; ", f"export CTS_CODE_LOCATION={self._nersc_code_path}; ", f"export CTS_RESULT_FILE_LOCATION={rootpath / 'upload_checksums_result.json'}; ", @@ -848,23 +898,26 @@ async def upload_JAWS_job_files( command = "".join(command) # May want to make this non-blocking if calculating checksums takes too long # Would require another set of job states and another callback URL so try to avoid - await dtns.run(command) - + await perl.run(command) + checksumpath = rootpath / checksum_file - checksums = await self._download_json_file_from_NERSC(Machine.dtns, checksumpath) + checksums = await self._download_json_file_from_NERSC(Machine.perlmutter, checksumpath) s3_paths = [] crc64nvmes = [] nersc_rel_paths = [] + total_bytes = 0 for c in checksums["files"]: s3_paths.append(Path(c["s3path"])) crc64nvmes.append(c["crc64nvme"]) nersc_rel_paths.append(c["respath"]) + total_bytes += c["size"] presigns = await files_to_urls(s3_paths, crc64nvmes) - return await self.upload_presigned_files( + return await self._upload_presigned_files( job.id, [os.path.join(jaws_output_dir, nrp) for nrp in nersc_rel_paths], presigns, cburl, + total_bytes, concurrency, insecure_ssl, ) @@ -882,7 +935,7 @@ async def get_uploaded_JAWS_files(self, job: models.Job) -> dict[str, str]: path = self._get_job_scratch(job.id) / _CRC64NVME_CHECKSUMS_JSON_FILE_NAME # This uploads the same file from NERSC again. We could put the results in a temporary DB # collection if it turns out to be too expensive. YAGNI - checksums = await self._download_json_file_from_NERSC(Machine.dtns, path) + checksums = await self._download_json_file_from_NERSC(Machine.perlmutter, path) return {c["s3path"]: c["crc64nvme"] for c in checksums["files"]} async def upload_JAWS_log_files_on_error( @@ -909,7 +962,7 @@ async def upload_JAWS_log_files_on_error( concurrency - the number of simultaneous uploads to process. insecure_ssl - whether to skip the cert check for the S3 URL. - Returns the NERSC task ID for the upload. + Returns the NERSC Slurm job ID for the upload. """ _not_falsy(job, "job") _not_falsy(files_to_urls, "files_to_urls") @@ -924,7 +977,14 @@ async def upload_JAWS_log_files_on_error( rootpath = self._get_job_scratch(job.id) remotelogs = [rootpath / _JOB_LOGS / f for f in logs] - + # The extracted log files don't exist yet - they're written by the same Slurm job this + # method submits below, so they can't be stat'd ahead of time to size that job's wall + # time. Use the size of errors.json itself instead: it already exists (JAWS wrote it), + # a stat is instant regardless of its size, and its on-disk size is a safe upper bound + # for the stdout/stderr content that will be extracted from it (JSON string-escaping + # only inflates size relative to the raw decoded content). + total_bytes = await self._get_remote_file_sizes([errfilepath]) + manifest = self._create_upload_manifest(remotelogs, presigns, concurrency, insecure_ssl) return await self._process_manifest( manifest, @@ -932,11 +992,29 @@ async def upload_JAWS_log_files_on_error( cburl, "error_log_upload_manifest.json", "error_log", + total_bytes, mode="errorsjson", error_json_file_location=errfilepath, container_logs_location=str(rootpath / _JOB_LOGS) ) + async def _get_remote_file_sizes(self, paths: list[Path]) -> int: + """ + Get the total size in bytes of a list of files on NERSC. Raises an error if any file + does not exist. + """ + cli = self._client_provider() + perl = await cli.compute(Machine.perlmutter) + total = 0 + # TODO PERF parallelize the ls calls if this ever needs to handle more than a handful + # of paths + for p in paths: + entries = await perl.ls(str(p)) + if not entries: + raise ValueError(f"File does not exist on NERSC: {p}") + total += int(entries[0].size) + return total + async def setup_refdata_transfer_callback( self, refdata: models.ReferenceData, site: sites.Cluster, callback_url: str ): @@ -960,9 +1038,14 @@ async def setup_refdata_transfer_callback( sfcli = AsyncClient(api_base_url="https://api.nersc.gov/api/beta", access_token=token) payload = { "path_condition": { - "path": str(self._get_refdata_file_complete_path(refdata.id, site)), - "machine": Machine.dtns.value - }, + "path": str(self._get_refdata_file_complete_path(refdata.id, site)), + # The completion file lives in the DTN-only-writable refdata area, but that + # area is readable (just not writable) from Perlmutter login nodes, so the + # watch itself can run from either machine. + # TODO VERIFY confirm against a live NERSC callback (beta endpoint, no + # automated test coverage) before relying on this in production. + "machine": Machine.perlmutter.value + }, "url": cb_url, # seconds. Assume that refdata transfers take less than a day. Make configurable? "timeout": 24 * 60 * 60, @@ -986,7 +1069,7 @@ async def get_refdata_transfer_result( # similar to _get_transfer_result, but not similar enough to warrant DRYing things up path = self._get_refdata_file_complete_path(refdata.id, site) res = await self._download_json_file_from_NERSC( - Machine.dtns, path, no_exception_on_missing_file=True + Machine.perlmutter, path, no_exception_on_missing_file=True ) if not res: return TransferResult(state=TransferState.INCOMPLETE) @@ -1019,7 +1102,7 @@ async def clean_job(self, job: models.Job, jaws_output_dirs: list[Path]): # extra JAWS call for nothing... to_delete = list(jaws_output_dirs or []) # if jaws_output_dirs is None or empty to_delete.append(self._get_job_scratch(_not_falsy(job, "job").id)) - await self._delete_dtn_paths(to_delete) + await self._delete_paths(to_delete) async def clean_refdata(self, refdata: models.ReferenceData): """ @@ -1031,4 +1114,4 @@ async def clean_refdata(self, refdata: models.ReferenceData): Note that running this method on reference data where staging is not in a terminal state may result in undefined behavior. """ - await self._delete_dtn_paths([self._get_refdata_scratch(refdata.id)]) + await self._delete_paths([self._get_refdata_scratch(refdata.id)]) diff --git a/cdmtaskservice/nersc/paths.py b/cdmtaskservice/nersc/paths.py index 70147a5..2a5108c 100644 --- a/cdmtaskservice/nersc/paths.py +++ b/cdmtaskservice/nersc/paths.py @@ -28,9 +28,6 @@ class NERSCPaths: jaws_refdata_root_dir: Path = field(init=False) """The root directory for reference data configured for the JAWS 'kbase` site.""" - jaws_staging_dir_dtns: Path = field(init=False) - """The JAWS staging directory for the `kbase` site on a NERSC DTN.""" - jaws_staging_dir_perlmutter: Path = field(init=False) """The JAWS staging directory for the `kbase` site on the NERSC Perlmutter system.""" @@ -38,13 +35,11 @@ def __init__( self, code_path: str, jaws_refdata_root_dir: str, - jaws_staging_dir_dtns: str, jaws_staging_dir_perlmutter: str ): fields = { 'code_path': code_path, 'jaws_refdata_root_dir': jaws_refdata_root_dir, - 'jaws_staging_dir_dtns': jaws_staging_dir_dtns, 'jaws_staging_dir_perlmutter': jaws_staging_dir_perlmutter, } for name, value in fields.items(): diff --git a/cdmtaskservice/nersc/remote.py b/cdmtaskservice/nersc/remote.py index 9d72567..4f89326 100644 --- a/cdmtaskservice/nersc/remote.py +++ b/cdmtaskservice/nersc/remote.py @@ -12,6 +12,7 @@ import os from pathlib import Path import requests +import subprocess import sys import traceback from typing import Callable @@ -38,9 +39,16 @@ def calculate_checksums(jaws_output_dir: str, checksum_output_file: str): outs = parse_outputs_json(f) res = {"files": [], "stdouts": [], "stderrs": []} # TODO PERF may want to parallelize this with a max process limit + # TODO PERF could also chunk large files and combine CRCs via awscrt.checksums.combine_crc64nvme for s3_path, result_path in outs.output_files.items(): - crc = crc64nvme_b64(jdir / result_path) - res["files"].append({"crc64nvme": crc, "s3path": s3_path, "respath": result_path}) + fpath = jdir / result_path + crc = crc64nvme_b64(fpath) + res["files"].append({ + "crc64nvme": crc, + "s3path": s3_path, + "respath": result_path, + "size": fpath.stat().st_size, + }) for so in outs.stdout: crc = crc64nvme_b64(jdir / so) res["stdouts"].append({"crc64nvme": crc, "respath": so}) @@ -69,6 +77,86 @@ def process_data_transfer_manifest(manifest_file_path: str): return None +def sync_refdata_to_dtn( + staging_dir: str, + dtn_host: str, + dest_dir: str, + completion_file: str, + completion_file_contents: str, +): + """ + Copy reference data downloaded and unpacked into a local staging directory to its final + location on a NERSC DTN-mounted filesystem, and write the completion file JAWS watches for. + + This is necessary because reference data lives on a filesystem (/global/dna) that is only + writable from NERSC DTNs, but `xfer` QOS Slurm jobs run on Perlmutter login nodes. This + relies on the NERSC sshproxy SSH key already present in $HOME being valid and shared between + login and DTN nodes, requiring no further authentication. + + staging_dir - the local directory the reference data was downloaded and unpacked into. + dtn_host - the DTN hostname to copy the data to. + dest_dir - the final, DTN-only-writable destination directory for the reference data. + completion_file - the path, on the DTN host, of the completion file JAWS watches for. + completion_file_contents - the contents to write to the completion file. + """ + # --protect-args: dest_dir / completion_file are always built from a server-generated UUID + # plus fixed config strings, never user input, so this is defense in depth rather than a + # fix for a reachable bug. + # --timeout / ConnectTimeout: without these a stalled connection hangs until the enclosing + # Slurm job's wall-clock limit kills it (up to 48h under the `xfer` QOS). + # --partial: keep partially-transferred files on interruption so a retry isn't guaranteed to + # re-copy everything from scratch. + rsync_base_args = [ + "rsync", "-a", "--mkpath", "--protect-args", "--partial", "--timeout=300", + "-e", "ssh -o BatchMode=yes -o ConnectTimeout=30", + ] + subprocess.run( + [ + *rsync_base_args, + f"{str(staging_dir).rstrip('/')}/", + f"{dtn_host}:{dest_dir}/", + ], + check=True, + ) + # Write the completion file locally, next to the (by now already synced) staging dir, and + # let rsync push it over rather than shelling out to write it remotely. + local_completion_file = Path(staging_dir).parent / Path(completion_file).name + local_completion_file.write_text(f"{completion_file_contents}\n") + subprocess.run( + [ + *rsync_base_args, + str(local_completion_file), + f"{dtn_host}:{completion_file}", + ], + check=True, + ) + + +def process_refdata_download_manifest( + manifest_file_path: str, + staging_dir: str, + dtn_host: str, + dest_dir: str, + completion_file: str, + completion_file_contents: str, +): + """ + Downloads reference data files per a transfer manifest into a local staging directory, then + copies them to their final DTN-only-writable location and writes the JAWS completion file. + + manifest_file_path - the path to the transfer manifest file. Its file entries are expected + to point into staging_dir. + staging_dir - the local directory the reference data will be downloaded and unpacked into. + dtn_host - the DTN hostname to copy the data to. + dest_dir - the final, DTN-only-writable destination directory for the reference data. + completion_file - the path, on the DTN host, of the completion file JAWS watches for. + completion_file_contents - the contents to write to the completion file. + """ + process_data_transfer_manifest(manifest_file_path) + sync_refdata_to_dtn(staging_dir, dtn_host, dest_dir, completion_file, completion_file_contents) + return None + + def process_errorsjson( errorsjson_file_path: str, logfiles_directory: str, @@ -153,6 +241,22 @@ def main(): resfile, callback_url ) + elif mode == "refdata_manifest": + _error_wrapper( + process_refdata_download_manifest, + [ + # TODO CODE staging dir is redundant with the manifest, the file in + # the manifest should be prefixed with the staging dir + os.environ["CTS_MANIFEST_LOCATION"], + os.environ["CTS_STAGING_DIR"], + os.environ["CTS_DTN_HOST"], + os.environ["CTS_REFDATA_DEST_DIR"], + os.environ["CTS_COMPLETION_FILE_LOCATION"], + os.environ["CTS_COMPLETION_FILE_CONTENTS"], + ], + resfile, + callback_url + ) elif mode == "errorsjson": _error_wrapper( process_errorsjson, diff --git a/cdmtaskservice/nersc/status.py b/cdmtaskservice/nersc/status.py index 9af3c8c..d7c6d1e 100644 --- a/cdmtaskservice/nersc/status.py +++ b/cdmtaskservice/nersc/status.py @@ -13,7 +13,13 @@ class Status(NamedTuple): """ The status of the NERSC compute systems. """ ok: bool - """ True if all systems are available. """ + """ + True if all systems are available. Note that job execution no longer uses the DTNs + (SFAPI compute now targets Perlmutter exclusively), so a DTN outage only actually + blocks refdata staging, not job submission - `ok` is not that precise, though, and + still requires both systems to be up. Making job submission proceed independently of + DTN status is more work than is warranted right now. + """ perlmutter_up: bool """ True if perlmutter is available. """ diff --git a/cdmtaskservice/s3/remote.py b/cdmtaskservice/s3/remote.py index c42a966..ce5139e 100644 --- a/cdmtaskservice/s3/remote.py +++ b/cdmtaskservice/s3/remote.py @@ -381,9 +381,12 @@ async def process_data_transfer_manifest(manifest: dict[str, Any]): # stress error checking too much. # Potential performance improvements: # * aiofiles - # * Add multiprocessing; not clear if helpful given low CPU load expected - # * See if multipart uploads are possible with presigned urls - # * Presumably only helpful if disk reads are the bottleneck + # * Add multiprocessing; not clear if helpful given low CPU load expected + # * TODO PERF support multipart transfers via multiple range-based presigned URLs per file, + # fetched/pushed concurrently. A single HTTP stream tends to plateau well below what a + # high-bandwidth WAN path (e.g. NERSC <-> ANL over ESnet) can sustain; splitting large + # files across several concurrent range requests is the standard fix and doesn't require + # exposing real S3 credentials to NERSC the way handing off to an S3-aware CLI would. # TODO TEST add tests for this and its dependency functions. _not_falsy(manifest, "manifest") operation = manifest["op"] @@ -410,11 +413,6 @@ async def process_data_transfer_manifest(manifest: dict[str, Any]): ) else: raise ValueError(f"unknown operation: {operation}") - if "completion-file" in manifest: - with open(manifest["completion-file"], "w") as f: - # just raise a keyerror if it's not there - f.write(manifest["completion-file-contents"] + "\n") - _logr.info(f"Wrote completion file {manifest['completion-file']}") _logr.info(f"{operation} manifest processing complete") diff --git a/cdmtaskservice/update_state.py b/cdmtaskservice/update_state.py index ceceeee..c50af24 100644 --- a/cdmtaskservice/update_state.py +++ b/cdmtaskservice/update_state.py @@ -20,7 +20,7 @@ class UpdateField(StrEnum): Fields which may be present in an update. """ - NERSC_DOWNLOAD_TASK_ID = auto() + NERSC_DOWNLOAD_JOB_ID = auto() """ The NERSC Superfacility ID for an download task. """ JAWS_RUN_ID = auto() @@ -38,10 +38,10 @@ class UpdateField(StrEnum): CPU_FACTOR = auto() """ The ratio of cpu time actually used to cpu time requested. """ - NERSC_UPLOAD_TASK_ID = auto() + NERSC_UPLOAD_JOB_ID = auto() """ The NERSC Superfacility ID for an upload task. """ - NERSC_LOG_UPLOAD_TASK_ID = auto() + NERSC_LOG_UPLOAD_JOB_ID = auto() """ The NERSC Superfacility ID for a log upload task. """ OUTPUT_FILE_PATHS = auto() @@ -187,16 +187,16 @@ def submitted_download() -> JobUpdate: ) -def submitted_nersc_download(task_id: str) -> JobUpdate: +def submitted_nersc_download(job_id: str) -> JobUpdate: """ Update a job's state from created to download submitted and add a NERSC - superfacility API download task ID. + superfacility API download job ID. """ return JobUpdate( )._set_current_state(models.JobState.CREATED )._set_new_state(models.JobState.DOWNLOAD_SUBMITTED )._set_fields( - {UpdateField.NERSC_DOWNLOAD_TASK_ID: _require_string(task_id, "task_id")} + {UpdateField.NERSC_DOWNLOAD_JOB_ID: _require_string(job_id, "job_id")} ) @@ -285,15 +285,15 @@ def submitted_upload() -> JobUpdate: ) -def submitted_nersc_upload(task_id: str) -> JobUpdate: +def submitted_nersc_upload(job_id: str) -> JobUpdate: """ Update a job's state from upload submitting to upload submitted and add a NERSC - superfacility API upload task ID. + superfacility API upload job ID. """ return JobUpdate( )._set_current_state(models.JobState.UPLOAD_SUBMITTING )._set_new_state(models.JobState.UPLOAD_SUBMITTED - )._set_fields({UpdateField.NERSC_UPLOAD_TASK_ID: _require_string(task_id, "task_id")} + )._set_fields({UpdateField.NERSC_UPLOAD_JOB_ID: _require_string(job_id, "job_id")} ) @@ -512,16 +512,16 @@ def submitted_error_processing() -> JobUpdate: ) -def submitted_nersc_error_processing(task_id: str) -> JobUpdate: +def submitted_nersc_error_processing(job_id: str) -> JobUpdate: """ Update a job's state from error processing submitting to error processing submitted and add - a NERSC superfacility API upload task ID. + a NERSC superfacility API upload job ID. """ return JobUpdate( )._set_current_state(models.JobState.ERROR_PROCESSING_SUBMITTING )._set_new_state(models.JobState.ERROR_PROCESSING_SUBMITTED )._set_fields( - {UpdateField.NERSC_LOG_UPLOAD_TASK_ID: _require_string(task_id, "task_id")} + {UpdateField.NERSC_LOG_UPLOAD_JOB_ID: _require_string(job_id, "job_id")} ) @@ -584,16 +584,16 @@ def submitted_refdata_download() -> RefdataUpdate: ) -def submitted_nersc_refdata_download(task_id: str) -> RefdataUpdate: +def submitted_nersc_refdata_download(job_id: str) -> RefdataUpdate: """ Update a refdata staging process's state from created to download submitted and add a NERSC - superfacility API download task ID. + superfacility API download job ID. """ return RefdataUpdate( )._set_current_state(models.ReferenceDataState.CREATED )._set_new_state(models.ReferenceDataState.DOWNLOAD_SUBMITTED )._set_fields( - {UpdateField.NERSC_DOWNLOAD_TASK_ID: _require_string(task_id, "task_id")} + {UpdateField.NERSC_DOWNLOAD_JOB_ID: _require_string(job_id, "job_id")} ) diff --git a/cdmtaskservice_config.toml.jinja b/cdmtaskservice_config.toml.jinja index 090ec60..cd34336 100644 --- a/cdmtaskservice_config.toml.jinja +++ b/cdmtaskservice_config.toml.jinja @@ -40,16 +40,7 @@ user = "{{ KBCTS_NERSC_JAWS_USER or "" }}" # The JAWS installation is expected to be configured to use this directory to look up refdata. refdata_root_dir = "{{ KBCTS_NERSC_JAWS_REFDATA_DIR or "" }}" -# The JAWS site staging directory at NERSC on a Data Transfer Node. -# Note that currently the CTS is hard coded to use the `kbase` site at NERSC Perlmutter. -# This is typically something like -# `/global/pscratch/sd///-prod/` -jaws_staging_dir_dtn = "{{ - KBCTS_NERSC_JAWS_DTN_STAGING_DIR or "/global/pscratch/sd/k/kbjaws/kbase-prod/" - }}" - # The JAWS site staging directory at NERSC on the Perlmutter system. -# This is the equivalent directory to the DTN staging directory. # Note that currently the CTS is hard coded to use the `kbase` site at NERSC Perlmutter. # This is typically something like # `/pscratch/sd///-prod/` diff --git a/docker-compose.yaml b/docker-compose.yaml index ad0fed5..d461d83 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -36,7 +36,6 @@ services: - KBCTS_HAS_NERSC_ACCOUNT_ROLE=HAS_NERSC_ACCOUNT - KBCTS_NERSC_JAWS_USER=cdm_ts - KBCTS_NERSC_JAWS_REFDATA_DIR=/global/dna/kbase/reference/jaws - - KBCTS_NERSC_JAWS_DTN_STAGING_DIR=/global/pscratch/sd/k/kbjaws/kbase-prod/ - KBCTS_NERSC_JAWS_PERMUTTER_STAGING_DIR=/pscratch/sd/k/kbjaws/kbase-prod/ - KBCTS_SFAPI_CRED_PATH=/creds/sfapi_creds - KBCTS_NERSC_REMOTE_CODE_DIR=/global/cfs/cdirs/kbase/cdm_task_service diff --git a/test/mongo_test.py b/test/mongo_test.py index effa99a..abc5dcf 100644 --- a/test/mongo_test.py +++ b/test/mongo_test.py @@ -1502,7 +1502,7 @@ async def test_refdata_redundant_update_time(mondb): state=models.ReferenceDataState.DOWNLOAD_SUBMITTED, time=dt, )) - rd.statuses[0].nersc_download_task_id = ["ntid"] + rd.statuses[0].nersc_download_job_id = ["ntid"] assert got == rd # check that the update time is set correctly