Force renew .dbtoken before workflow submission (#224)

* Force renew `.dbtoken` before workflow submission

* Minor change

* Simplification

* Update README, and no pre-commit error

Author: dachengx

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
         language_version: python3
 
 -   repo: https://github.com/pycqa/docformatter
-    rev: v1.7.5
+    rev: master
     hooks:
     -   id: docformatter
         additional_dependencies: [tomli]

README.md

@@ -158,10 +158,9 @@ After installation and setting up the environment, it is time to submit jobs. Th
 
 ```
 [whoami@ap23 ~]$ outsource --help
-usage: Outsource [-h] --context CONTEXT --xedocs_version XEDOCS_VERSION [--image IMAGE]
-                 [--detector {all,tpc,muon_veto,neutron_veto}] [--workflow_id WORKFLOW_ID] [--ignore_processed]
-                 [--debug] [--from NUMBER_FROM] [--to NUMBER_TO] [--run [RUN ...]] [--runlist RUNLIST] [--rucio_upload]
-                 [--rundb_update] [--local_transfer]
+usage: Outsource [-h] --context CONTEXT --xedocs_version XEDOCS_VERSION [--image IMAGE] [--detector {all,tpc,muon_veto,neutron_veto}]
+                 [--workflow_id WORKFLOW_ID] [--ignore_processed] [--debug] [--from NUMBER_FROM] [--to NUMBER_TO] [--run [RUN ...]]
+                 [--runlist RUNLIST] [--rucio_upload] [--rundb_update] [--local_transfer] [--keep_dbtoken]
 
 optional arguments:
   -h, --help            show this help message and exit
@@ -171,20 +170,20 @@ optional arguments:
   --image IMAGE         Singularity image. Accepts either a full path or a single name and assumes a format like this:
                         /cvmfs/singularity.opensciencegrid.org/xenonnt/base-environment:{image}
   --detector {all,tpc,muon_veto,neutron_veto}
-                        Detector to focus on. If 'all' (default) will consider all three detectors. Otherwise pass a
-                        single one of 'tpc', 'neutron_veto', 'muon_veto'. Pairs of detectors not yet supported.
+                        Detector to focus on. If 'all' (default) will consider all three detectors. Otherwise pass a single one of
+                        'tpc', 'neutron_veto', 'muon_veto'. Pairs of detectors not yet supported.
   --workflow_id WORKFLOW_ID
                         Custom workflow_id of workflow. If not passed, inferred from today's date.
   --ignore_processed    Ignore runs that have already been processed
-  --debug               Debug mode. Does not automatically submit the workflow, and jobs do not update RunDB nor upload
-                        to rucio.
+  --debug               Debug mode. Does not automatically submit the workflow, and jobs do not update RunDB nor upload to rucio.
   --from NUMBER_FROM    Run number to start with
   --to NUMBER_TO        Run number to end with
-  --run [RUN ...]       Space separated specific run number(s) to process
+  --run [RUN ...]       Space separated specific run_id(s) to process
   --runlist RUNLIST     Path to a runlist file
   --rucio_upload        Upload data to rucio after processing
   --rundb_update        Update RunDB after processing
   --local_transfer      Transfer data to local after processing
+  --keep_dbtoken        Do not renew .dbtoken
 ```
 
 This script requires at minimum the name of context (which must reside in the cutax version installed in the environment you are in). If no other arguments are passed, this script will try to find all data that can be processed, and process it. Some inputs from the configuration file at environmental variable `XENON_CONFIG` are also used, specifically:

outsource/scripts/submit.py

@@ -1,5 +1,6 @@
+import os
 import argparse
-from utilix import xent_collection, uconfig
+from utilix import xent_collection, uconfig, DB
 from utilix.io import load_runlist
 from utilix.config import setup_logger
 
@@ -77,8 +78,20 @@ def main():
         action="store_true",
         help="Transfer data to local after processing",
     )
+    parser.add_argument(
+        "--keep_dbtoken",
+        dest="keep_dbtoken",
+        action="store_true",
+        help="Do not renew .dbtoken",
+    )
     args = parser.parse_args()
 
+    if not args.keep_dbtoken:
+        os.remove(os.path.join(os.environ["HOME"], ".dbtoken"))
+        # Remove the cached DB instance and reinitialize it
+        DB._instances = dict()
+        DB()
+
     if "development" in args.image:
         raise RuntimeError("Cannot use development images/container for processing!")
 

outsource/submitter.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import json
+import shutil
 import getpass
 from itertools import chain
 from datetime import datetime
@@ -171,6 +172,10 @@ def __init__(
         self.outputs_dir = os.path.join(self.workflow_dir, "outputs")
         self.scratch_dir = os.path.join(self.workflow_dir, "scratch")
 
+    @property
+    def dbtoken(self):
+        return os.path.join(self.generated_dir, ".dbtoken")
+
     @property
     def workflow(self):
         return os.path.join(self.generated_dir, "workflow.yml")
@@ -752,11 +757,11 @@ def _generate_workflow(self):
         xenon_config = File(".xenon_config")
         rc.add_replica("local", ".xenon_config", f"file://{uconfig.config_path}")
 
-        # token needed for DB connection
+        # Token needed for DB connection
         token = File(".dbtoken")
-        rc.add_replica(
-            "local", ".dbtoken", "file://" + os.path.join(os.environ["HOME"], ".dbtoken")
-        )
+        # Avoid its change after the job submission
+        shutil.copy(os.path.join(os.environ["HOME"], ".dbtoken"), self.dbtoken)
+        rc.add_replica("local", ".dbtoken", f"file://{self.dbtoken}")
 
         tarballs, tarball_paths = self.make_tarballs()
         for tarball, tarball_path in zip(tarballs, tarball_paths):