Set output datasets to VALID in DBS3 before announcing a standard workflow

src/python/WMCore/MicroService/MSOutput/MSOutput.py

@@ -25,6 +25,7 @@
 from WMCore.MicroService.MSOutput.MSOutputTemplate import MSOutputTemplate
 from WMCore.WMException import WMException
 from WMCore.Services.AlertManager.AlertManagerAPI import AlertManagerAPI
+from WMCore.Services.DBS.DBS3Writer import DBS3Writer
 
 
 class MSOutputException(WMException):
@@ -248,6 +249,53 @@ def _executeConsumer(self, summary):
             self.logger.exception(msg)
             self.updateReportDict(summary, "error", msg)
 
+    def setDBSStatus(self, workflow):
+        """
+        The function to set the DBS status of outputs as VALID
+        :param workflow: a MSOutputTemplate object workflow
+        :return: the MSOutputTemplate object itself (with the necessary updates in place)
+        """
+        if not isinstance(workflow, MSOutputTemplate):
+            msg = "Unsupported type object '{}' for workflows! ".format(type(workflow))
+            msg += "It needs to be of type: MSOutputTemplate"
+            raise UnsupportedError(msg)
+
+        dbs3Writer = DBS3Writer(url=self.msConfig["dbsReadUrl"],
+                                writeUrl=self.msConfig["dbsWriteUrl"])
+
+        # if anything fail along the way, set it back to "pending"
+        dbsUpdateStatus = "done"
+        for dMap in workflow['OutputMap']:
+
+            if self.msConfig['enableDbsStatusChange']:
+
+                res = dbs3Writer.setDBSStatus(dataset=dMap["Dataset"],
+                                              status=self.msConfig['dbsStatus']["valid"])
+
+                if res:
+                    dMap["DBSStatus"] = self.msConfig['dbsStatus']["valid"]
+                else:
+                    # There is at least one dataset whose dbs status update is unsuccessful
+                    dbsUpdateStatus = "pending"
+
+            else:
+                msg = "DRY-RUN DBS: DBS status change for DID: {}, ".format(dMap['Dataset'])
+                self.logger.info(msg)
+
+        # Finally, update the MSOutput template document with either partial or
+        # complete dbs statuses
+        self.docKeyUpdate(workflow, OutputMap=workflow['OutputMap'])
+        workflow.updateTime()
+        if dbsUpdateStatus == "done":
+            self.logger.info("All the DBS status updates succeeded for: %s. Marking it as 'done'",
+                             workflow['RequestName'])
+            self.docKeyUpdate(workflow, DBSUpdateStatus='done')
+        else:
+            self.logger.info("DBS status updates partially successful for: %s. Keeping it 'pending'",
+                             workflow['RequestName'])
+
+        return workflow
+
     def makeSubscriptions(self, workflow):
         """
         The common function to make the final subscriptions
@@ -448,28 +496,32 @@ def msOutputConsumer(self):
         #    Done: To build it through a pipe
         #    Done: To write back the updated document to MonogoDB
         msPipelineRelVal = Pipeline(name="MSOutputConsumer PipelineRelVal",
-                                    funcLine=[Functor(self.makeSubscriptions),
+                                    funcLine=[Functor(self.setDBSStatus),
+                                              Functor(self.makeSubscriptions),
                                               Functor(self.makeTapeSubscriptions),
                                               Functor(self.docUploader,
                                                       update=True,
                                                       keys=['LastUpdate',
                                                             'TransferStatus',
+                                                            'DBSUpdateStatus',
                                                             'OutputMap']),
                                               Functor(self.docDump, pipeLine='PipelineRelVal'),
                                               Functor(self.docCleaner)])
         msPipelineNonRelVal = Pipeline(name="MSOutputConsumer PipelineNonRelVal",
-                                       funcLine=[Functor(self.makeSubscriptions),
+                                       funcLine=[Functor(self.setDBSStatus),
+                                                 Functor(self.makeSubscriptions),
                                                  Functor(self.makeTapeSubscriptions),
                                                  Functor(self.docUploader,
                                                          update=True,
                                                          keys=['LastUpdate',
                                                                'TransferStatus',
+                                                               'DBSUpdateStatus',
                                                                'OutputMap']),
                                                  Functor(self.docDump, pipeLine='PipelineNonRelVal'),
                                                  Functor(self.docCleaner)])
 
         wfCounterTotal = 0
-        mQueryDict = {'TransferStatus': 'pending'}
+        mQueryDict = { "$or": [ { "TransferStatus": "pending" }, { "DBSUpdateStatus": "pending" } ] }
         pipeCollections = [(msPipelineRelVal, self.msOutRelValColl),
                            (msPipelineNonRelVal, self.msOutNonRelValColl)]
         for pipeColl in pipeCollections:

src/python/WMCore/MicroService/MSOutput/MSOutputTemplate.py

@@ -102,6 +102,7 @@ def docSchema(self):
                            'Copies': 1,
                            ...}],
                     "TransferStatus": "pending"|"done,
+                    "DBSUpdateStatus": "pending"|"done,
             "RequestType": ""
             }
         :return: a list of tuples
@@ -116,7 +117,8 @@ def docSchema(self):
             ('IsRelVal', False, bool),
             ('OutputDatasets', [], list),
             ('OutputMap', [], list),
-            ('TransferStatus', "pending", (bytes, str))]
+            ('TransferStatus', "pending", (bytes, str)),
+            ('DBSUpdateStatus', "pending", (bytes, str))]
         return docTemplate
 
     def outputMapSchema(self):
@@ -135,7 +137,8 @@ def outputMapSchema(self):
              'DiskDestination': "",
              'TapeDestination': "",
              'DiskRuleID': "",
-             'TapeRuleID': ""}
+             'TapeRuleID': "",
+             'DBSStatus': ""}
         :return: a list of tuples
         """
         outMapTemplate = [
@@ -146,7 +149,8 @@ def outputMapSchema(self):
             ('DiskDestination', "", (bytes, str)),
             ('TapeDestination', "", (bytes, str)),
             ('DiskRuleID', "", (bytes, str)),
-            ('TapeRuleID', "", (bytes, str))]
+            ('TapeRuleID', "", (bytes, str)),
+            ('DBSStatus', "", (bytes, str))]
         return outMapTemplate
 
     def _checkAttr(self, myDoc, update=False, throw=False, **kwargs):

src/python/WMCore/Services/DBS/DBS3Reader.py

@@ -971,3 +971,32 @@ def getParentDatasetTrio(self, childDataset):
                     frozenKey = frozenset(runLumiPair)
                     parentFrozenData[frozenKey] = fileId
         return parentFrozenData
+
+    def getDBSStatus(self, dataset):
+        """
+        The function to get the DBS status of outputs
+        :param dataset: dataset name
+        :return: DBS status of the given dataset
+        """
+
+        allowedDbsStatuses = ["VALID", "INVALID", "PRODUCTION"]
+
+        response = None
+        try:
+            response = self.dbs.listDatasets(dataset=dataset, dataset_access_type='*', detail=True)
+        except Exception as ex:
+            msg = "Exception while getting the status of following dataset on DBS: {} ".format(dataset)
+            msg += "Error: {}".format(str(ex))
+            self.logger.exception(msg)
+
+        if response:
+            dbsStatus = response[0]['dataset_access_type']
+            isAllowedStatus = dbsStatus in allowedDbsStatuses
+
+            if isAllowedStatus:
+                self.logger.info("%s is %s", dataset, dbsStatus)
+                return dbsStatus
+            else:
+                raise Exception("This is not an allowed DBS status: {}".format(str(dbsStatus)))
+        else:
+            return None

src/python/WMCore/Services/DBS/DBS3Writer.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+"""
+_DBSReader_
+
+Read/Write DBS Interface
+
+"""
+from __future__ import print_function, division
+from builtins import str
+import logging
+
+from dbs.apis.dbsClient import DbsApi
+from dbs.exceptions.dbsClientException import dbsClientException
+
+from WMCore.Services.DBS.DBSErrors import DBSWriterError, formatEx3
+from WMCore.Services.DBS.DBS3Reader import DBS3Reader
+
+
+class DBS3Writer(DBS3Reader):
+    """
+    _DBSReader_
+
+    General API for writing data to DBS
+    """
+
+    def __init__(self, url, writeUrl, logger=None, **contact):
+
+        # instantiate dbs api object
+        try:
+            super(DBS3Writer, self).__init__(url=url)
+            self.dbs = DbsApi(writeUrl, **contact)
+            self.logger = logger or logging.getLogger(self.__class__.__name__)
+        except dbsClientException as ex:
+            msg = "Error in DBSWriter with DbsApi\n"
+            msg += "%s\n" % formatEx3(ex)
+            raise DBSWriterError(msg)
+
+    def setDBSStatus(self, dataset, status):
+        """
+        The function to set the DBS status of an output dataset
+        :param dataset: Dataset name
+        :return: True if operation is successful, False otherwise
+        """
+
+        try:
+            self.dbs.updateDatasetType(dataset=dataset,
+                                     dataset_access_type=status)
+        except Exception as ex:
+            msg = "Exception while setting the status of following dataset on DBS: {} ".format(dataset)
+            msg += "Error: {}".format(str(ex))
+            self.logger.exception(msg)
+
+        dbsStatus = super(DBS3Writer, self).getDBSStatus(dataset)
+
+        if dbsStatus == status:
+            return True
+        else:
+            return False

src/python/WMQuality/Emulators/DBSClient/MockedDBSGlobalCalls.py

@@ -84,6 +84,15 @@
           {'block_name': '/Cosmics/Commissioning2015-PromptReco-v1/RECO#004ac3ba-d09e-11e4-afad-001e67ac06a0'}],
          ['listBlockParents',
           {'block_name': '/Cosmics/Commissioning2015-v1/RAW#942d76fe-cf0e-11e4-afad-001e67ac06a0'}],
+         ['listDatasets',
+          {'dataset_access_type': '*',
+           'dataset': '/ggXToJPsiJPsi_JPsiToMuMu_M6p2_JPCZeroMinusPlus_TuneCP5_13TeV-pythia8-JHUGen/RunIIFall17pLHE-93X_mc2017_realistic_v3-v2/LHE',
+           'detail': True}],
+         ['listDatasets',
+          {'dataset_access_type': '*',
+           'dataset': '/DYToEE_M-50_NNPDF31_TuneCP5_14TeV-powheg-pythia8/Run3Summer19DRPremix-BACKFILL_2024Scenario_106X_mcRun3_2024_realistic_v4-v6/GEN-SIM-RAW',
+           'detail': True}],
+
 
          # Exception throwing calls