chore: change workflow commons from dict to a pydantic model

Author: AcquaDiGiorgio

src/dirac_cwl/commands/analyze_xml_summary.py

@@ -2,84 +2,61 @@
 
 import os
 
-from DIRAC.TransformationSystem.Client.FileReport import FileReport
-from DIRAC.WorkloadManagementSystem.Client.JobReport import JobReport
-from LHCbDIRAC.Core.Utilities.XMLSummaries import XMLSummary
 from LHCbDIRAC.Workflow.Modules.AnalyseXMLSummary import _areInputsOK, _isXMLSummaryOK
 from LHCbDIRAC.Workflow.Modules.BookkeepingReport import _generate_xml_object
 
 from dirac_cwl.core.exceptions import WorkflowProcessingException
 
 from .core import PostProcessCommand
-from .utils import prepare_lhcb_workflow_commons, save_workflow_commons
+from .workflow_commons import StepStatus, WorkflowCommons
 
 
 class AnalyseXmlSummary(PostProcessCommand):
     """Performs a series of checks on the XMLSummary output to make sure the execution was done correctly."""
 
-    def execute(self, job_path, **kwargs):
+    def execute(self, job_path: os.PathLike, **kwargs):
         """Execute the command.
 
         :param job_path: Path to the job working directory.
         :param kwargs: Additional keyword arguments.
         """
         failed = False
-        workflow_commons = {}
+        workflow_commons = None
         try:
-            workflow_commons_path = kwargs.get("workflow_commons_path", os.path.join(job_path, "workflow_commons.json"))
+            workflow_commons = WorkflowCommons.load(job_path)
 
-            workflow_commons = prepare_lhcb_workflow_commons(
-                workflow_commons_path,
-                extra_mandatory_values=[
-                    "bk_step_id",
-                ],
-                extra_default_values={
-                    "bookkeeping_LFNs": [],
-                    "size": {},
-                    "md5": {},
-                    "guid": {},
-                    "sim_description": "NoSimConditions",
-                },
-            )
-
-            if not workflow_commons["step_status"]["OK"]:
+            if workflow_commons.step_status == StepStatus.Failed:
                 return
 
-            if "xml_summary_path" in workflow_commons:
-                xf_o = XMLSummary(workflow_commons["xml_summary_path"])
-            else:
-                xf_o = _generate_xml_object(
-                    workflow_commons["cleaned_application_name"],
-                    workflow_commons["production_id"],
-                    workflow_commons["prod_job_id"],
-                    workflow_commons["command_number"],
-                    workflow_commons["command_id"],
+            if not workflow_commons.xf_o:
+                workflow_commons.xf_o = _generate_xml_object(
+                    workflow_commons.cleaned_application_name,
+                    workflow_commons.production_id,
+                    workflow_commons.prod_job_id,
+                    workflow_commons.step_number,
+                    workflow_commons.step_id,
                 )
 
-            file_report = FileReport()
-            job_report = JobReport(workflow_commons["job_id"])
-
-            file_report.statusDict = workflow_commons["file_report_files_dict"]
-
-            jobOk = _isXMLSummaryOK(xf_o)
+            jobOk = _isXMLSummaryOK(workflow_commons.xf_o)
 
             if jobOk:
                 jobOk = _areInputsOK(
-                    xf_o,
-                    workflow_commons["inputs"],
-                    workflow_commons["number_of_events"],
-                    workflow_commons["production_id"],
-                    file_report,
+                    workflow_commons.xf_o,
+                    workflow_commons.inputs,
+                    workflow_commons.number_of_events,
+                    workflow_commons.production_id,
+                    workflow_commons.file_report,
                 )
             if not jobOk:
-                job_report.setApplicationStatus("XMLSummary reports error")
+                workflow_commons.job_report.setApplicationStatus("XMLSummary reports error")
                 raise WorkflowProcessingException("XMLSummary reports error")
 
-            job_report.setApplicationStatus(f"{workflow_commons['application_name']} Step OK")
+            workflow_commons.job_report.setApplicationStatus(f"{workflow_commons.application_name} Step OK")
 
         except Exception as e:
             failed = True
             raise WorkflowProcessingException(e) from e
 
         finally:
-            save_workflow_commons(workflow_commons, workflow_commons_path, failed=failed)
+            if workflow_commons:
+                workflow_commons.save(job_path, failed=failed)

src/dirac_cwl/commands/bookkeeping_report.py

@@ -3,9 +3,7 @@
 import os
 
 from DIRAC.Workflow.Utilities.Utils import getStepCPUTimes
-from LHCbDIRAC.BookkeepingSystem.Client.BookkeepingClient import BookkeepingClient
 from LHCbDIRAC.Core.Utilities.ProductionData import constructProductionLFNs
-from LHCbDIRAC.Core.Utilities.XMLSummaries import XMLSummary
 from LHCbDIRAC.Workflow.Modules.BookkeepingReport import (
     _generate_xml_object,
     _generateInputFiles,
@@ -18,105 +16,85 @@
 from dirac_cwl.core.exceptions import WorkflowProcessingException
 
 from .core import PostProcessCommand
-from .utils import prepare_lhcb_workflow_commons, save_workflow_commons
+from .workflow_commons import StepStatus, WorkflowCommons
 
 
 class BookkeepingReport(PostProcessCommand):
     """Generates a bookkeeping report file based on the XMLSummary and the pool XML catalog."""
 
-    def execute(self, job_path, **kwargs):
+    def execute(self, job_path: os.PathLike, **kwargs):
         """Execute the command.
 
         :param job_path: Path to the job working directory.
         :param kwargs: Additional keyword arguments.
         """
         failed = False
-        workflow_commons = {}
+        workflow_commons = None
         try:
             # Obtain Workflow Commons
-            workflow_commons_path = kwargs.get("workflow_commons_path", os.path.join(job_path, "workflow_commons.json"))
-
-            workflow_commons = prepare_lhcb_workflow_commons(
-                workflow_commons_path,
-                extra_mandatory_values=[
-                    "bk_step_id",
-                ],
-                extra_default_values={
-                    "bookkeeping_LFNs": [],
-                    "size": {},
-                    "md5": {},
-                    "guid": {},
-                    "sim_description": "NoSimConditions",
-                },
-            )
+            workflow_commons = WorkflowCommons.load(job_path)
 
-            if not workflow_commons["step_status"]["OK"]:
+            if workflow_commons.step_status == StepStatus.Failed:
                 return
 
             # Setup variables
-            start_time = workflow_commons.get("start_time", None)
-
             cpu_times = {}
-            if start_time:
-                cpu_times["StartTime"] = start_time
-            if "start_stats" in workflow_commons:
-                cpu_times["StartStats"] = workflow_commons["start_stats"]
+            if workflow_commons.start_time:
+                cpu_times["StartTime"] = workflow_commons.start_time
+            if workflow_commons.start_stats:
+                cpu_times["StartStats"] = workflow_commons.start_stats
 
             exectime, cputime = getStepCPUTimes(cpu_times)
 
             number_of_processors = getNumberOfProcessorsToUse(
-                workflow_commons["job_id"], workflow_commons["max_number_of_processors"]
+                workflow_commons.job_id, workflow_commons.max_number_of_processors
             )
 
-            bk_client = BookkeepingClient()
-
             parameters = {
-                "PRODUCTION_ID": workflow_commons["production_id"],
-                "JOB_ID": workflow_commons["prod_job_id"],
-                "configVersion": workflow_commons["config_version"],
-                "outputList": workflow_commons["outputs"],
-                "configName": workflow_commons["config_name"],
-                "outputDataFileMask": workflow_commons["output_data_file_mask"],
+                "PRODUCTION_ID": workflow_commons.production_id,
+                "JOB_ID": workflow_commons.prod_job_id,
+                "configVersion": workflow_commons.config_version,
+                "outputList": workflow_commons.outputs,
+                "configName": workflow_commons.config_name,
+                "outputDataFileMask": workflow_commons.output_data_file_mask,
             }
 
-            if "bookkeeping_LFNs" in workflow_commons and "production_output_data" in workflow_commons:
-                bk_lfns = workflow_commons["bookkeeping_LFNs"]
+            if workflow_commons.bookkeeping_lfns and workflow_commons.production_output_data:
+                bk_lfns = workflow_commons.bookkeeping_lfns
 
                 if not isinstance(bk_lfns, list):
                     bk_lfns = [i.strip() for i in bk_lfns.split(";")]
 
             else:
-                result = constructProductionLFNs(parameters, bk_client)
+                result = constructProductionLFNs(parameters, workflow_commons.bk_client)
                 if not result["OK"]:
                     raise WorkflowProcessingException("Could not create production LFNs")
 
                 bk_lfns = result["Value"]["BookkeepingLFNs"]
 
-            ldate, ltime, ldatestart, ltimestart = _process_time(start_time)
+            ldate, ltime, ldatestart, ltimestart = _process_time(workflow_commons.start_time)
 
             # Obtain XMLSummary
-            if "xml_summary_path" in workflow_commons:
-                xf_o = XMLSummary(workflow_commons["xml_summary_path"])
-            else:
-                xf_o = _generate_xml_object(
-                    workflow_commons["cleaned_application_name"],
-                    workflow_commons["production_id"],
-                    workflow_commons["prod_job_id"],
-                    workflow_commons["command_number"],
-                    workflow_commons["command_id"],
+            if not workflow_commons.xf_o:
+                workflow_commons.xf_o = _generate_xml_object(
+                    workflow_commons.cleaned_application_name,
+                    workflow_commons.production_id,
+                    workflow_commons.prod_job_id,
+                    workflow_commons.step_number,
+                    workflow_commons.step_id,
                 )
 
             info_dict = {
                 "exectime": exectime,
                 "cputime": cputime,
                 "numberOfProcessors": number_of_processors,
-                "production_id": workflow_commons["production_id"],
-                "jobID": workflow_commons["job_id"],
-                "siteName": workflow_commons["site_name"],
-                "jobType": workflow_commons["job_type"],
-                "applicationName": workflow_commons["application_name"],
-                "applicationVersion": workflow_commons["application_version"],
-                "numberOfEvents": workflow_commons["number_of_events"],
+                "production_id": workflow_commons.production_id,
+                "jobID": workflow_commons.job_id,
+                "siteName": workflow_commons.site_name,
+                "jobType": workflow_commons.job_type,
+                "applicationName": workflow_commons.application_name,
+                "applicationVersion": workflow_commons.application_version,
+                "numberOfEvents": workflow_commons.number_of_events,
             }
 
             # Generate job_info object
@@ -126,38 +104,38 @@ def execute(self, job_path, **kwargs):
                 ltimestart,
                 ldate,
                 ltime,
-                xf_o,
-                workflow_commons["inputs"],
-                workflow_commons["command_id"],
-                workflow_commons["bk_step_id"],
-                bk_client,
-                workflow_commons["config_name"],
-                workflow_commons["config_version"],
+                workflow_commons.xf_o,
+                workflow_commons.inputs,
+                workflow_commons.step_id,
+                workflow_commons.bk_step_id,
+                workflow_commons.bk_client,
+                workflow_commons.config_name,
+                workflow_commons.config_version,
             )
 
             # Add input files to job_info
-            _generateInputFiles(job_info, bk_lfns, workflow_commons["inputs"])
+            _generateInputFiles(job_info, bk_lfns, workflow_commons.inputs)
 
             # Add output files to job_info
             _generateOutputFiles(
                 job_info,
                 bk_lfns,
-                workflow_commons["event_type"],
-                workflow_commons["application_name"],
-                xf_o,
-                workflow_commons["outputs"],
-                workflow_commons["inputs"],
+                workflow_commons.event_type,
+                workflow_commons.application_name,
+                workflow_commons.xf_o,
+                workflow_commons.outputs,
+                workflow_commons.inputs,
             )
 
             # Generate SimulationConditions
-            if workflow_commons["application_name"] == "Gauss":
-                job_info.simulation_condition = workflow_commons["sim_description"]
+            if workflow_commons.application_name == "Gauss":
+                job_info.simulation_condition = workflow_commons.sim_description
 
             # Convert job_info object to XML
             doc = job_info.to_xml()
 
             # Write to file
-            bfilename = f"bookkeeping_{workflow_commons['command_id']}.xml"
+            bfilename = f"bookkeeping_{workflow_commons.step_id}.xml"
             with open(bfilename, "wb") as bfile:
                 bfile.write(doc)
 
@@ -166,4 +144,5 @@ def execute(self, job_path, **kwargs):
             raise WorkflowProcessingException(e) from e
 
         finally:
-            save_workflow_commons(workflow_commons, workflow_commons_path, failed=failed)
+            if workflow_commons:
+                workflow_commons.save(job_path, failed=failed)