From c25676543297824e22dfe13c2243febbefdf3c25 Mon Sep 17 00:00:00 2001
From: swenzel <sandro.wenzel@cern.ch>
Date: Mon, 28 Apr 2025 14:15:53 +0200
Subject: [PATCH] Fixes for 2tag

* config transfer needs to happen in new environment
* ability to specify from_stage for async tag env
---
 MC/bin/o2dpg_workflow_utils.py | 10 +++++++++-
 MC/run/ANCHOR/anchorMC.sh      | 29 +++++++++++++++--------------
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/MC/bin/o2dpg_workflow_utils.py b/MC/bin/o2dpg_workflow_utils.py
index d4350e199..c15d9633b 100755
--- a/MC/bin/o2dpg_workflow_utils.py
+++ b/MC/bin/o2dpg_workflow_utils.py
@@ -274,6 +274,14 @@ def adjust_RECO_environment(workflowspec, package = ""):
     if len(package) == 0:
        return
 
+    # we try to extract the stage from the path (can be given via '@' separation)
+    # example O2sim::daily-xxx@DIGI ---> apply this environment from the DIGI phase
+    # example O2sim::daily-xxx@RECO ---> apply this environment from the RECO phase
+    # example O2sim::daily-xxx ---> apply this environment from the RECO phase == default case
+    from_stage = "RECO"
+    if package.count('@') == 1:
+        package, from_stage = package.split('@')
+
     # We essentially need to go through the graph and apply the mapping
     # so take the workflow spec and see if the task itself or any child
     # is labeled RECO ---> typical graph traversal with caching
@@ -305,7 +313,7 @@ def matches_or_inherits_label(taskid, label, cache):
 
     # fills the matches_label dictionary
     for taskid in range(len(workflowspec['stages'])):
-        if (matches_or_inherits_label(taskid, "RECO", matches_label)):
+        if (matches_or_inherits_label(taskid, from_stage, matches_label)):
            # now we do the final adjust (as annotation) in the workflow itself
            if workflowspec['stages'][taskid].get("alternative_alienv_package") == None:
               workflowspec['stages'][taskid]["alternative_alienv_package"] = package
diff --git a/MC/run/ANCHOR/anchorMC.sh b/MC/run/ANCHOR/anchorMC.sh
index 71ed0afe3..39a8d326a 100755
--- a/MC/run/ANCHOR/anchorMC.sh
+++ b/MC/run/ANCHOR/anchorMC.sh
@@ -187,6 +187,7 @@ fi
 #<----- START OF part that should run under a clean alternative software environment if this was given ------
 if [ "${ALIEN_JDL_O2DPG_ASYNC_RECO_TAG}" ]; then
   if [ "${LOADEDMODULES}" ]; then
+    export > env_before_stashing.env
     echo "Stashing initial modules"
     module save initial_modules.list # we stash the current modules environment
     module list --no-pager
@@ -253,19 +254,6 @@ fi
 ALIEN_JDL_LPMPRODUCTIONTAG=$ALIEN_JDL_LPMPRODUCTIONTAG_KEEP
 echo_info "Setting back ALIEN_JDL_LPMPRODUCTIONTAG to $ALIEN_JDL_LPMPRODUCTIONTAG"
 
-# now create the local MC config file --> config-json.json
-# we create the new config output with blacklist functionality
-ASYNC_CONFIG_BLACKLIST=${ASYNC_CONFIG_BLACKLIST:-${O2DPG_ROOT}/MC/run/ANCHOR/anchor-dpl-options-blacklist.json}
-${O2DPG_ROOT}/MC/bin/o2dpg_dpl_config_tools.py workflowconfig.log ${ASYNC_CONFIG_BLACKLIST} config-json.json
-ASYNC_WF_RC=${?}
-
-# check if config reasonably created
-if [[ "${ASYNC_WF_RC}" != "0" || `grep "ConfigParams" config-json.json 2> /dev/null | wc -l` == "0" ]]; then
-  echo_error "Problem in anchor config creation. Exiting."
-  exit 1
-fi
-
-
 # get rid of the temporary software environment
 if [ "${ALIEN_JDL_O2DPG_ASYNC_RECO_TAG}" ]; then
   module purge --no-pager
@@ -280,6 +268,18 @@ if [ "${ALIEN_JDL_O2DPG_ASYNC_RECO_TAG}" ]; then
 fi
 #<----- END OF part that should run under a clean alternative software environment if this was given ------
 
+# now create the local MC config file --> config-json.json
+# we create the new config output with blacklist functionality
+ASYNC_CONFIG_BLACKLIST=${ASYNC_CONFIG_BLACKLIST:-${O2DPG_ROOT}/MC/run/ANCHOR/anchor-dpl-options-blacklist.json}
+${O2DPG_ROOT}/MC/bin/o2dpg_dpl_config_tools.py workflowconfig.log ${ASYNC_CONFIG_BLACKLIST} config-json.json
+ASYNC_WF_RC=${?}
+
+# check if config reasonably created
+if [[ "${ASYNC_WF_RC}" != "0" || `grep "ConfigParams" config-json.json 2> /dev/null | wc -l` == "0" ]]; then
+  echo_error "Problem in anchor config creation. Exiting."
+  exit 1
+fi
+
 # -- CREATE THE MC JOB DESCRIPTION ANCHORED TO RUN --
 
 MODULES="--skipModules ZDC"
@@ -300,7 +300,8 @@ remainingargs="${remainingargs} -productionTag ${ALIEN_JDL_LPMPRODUCTIONTAG:-ali
 remainingargs="${ALIEN_JDL_ANCHOR_SIM_OPTIONS} ${remainingargs} --anchor-config config-json.json"
 # apply software tagging choice
 # remainingargs="${remainingargs} ${ALIEN_JDL_O2DPG_ASYNC_RECO_TAG:+--alternative-reco-software ${ALIEN_JDL_O2DPG_ASYNC_RECO_TAG}}"
-remainingargs="${remainingargs} ${ALIEN_JDL_O2DPG_ASYNC_RECO_TAG:+--alternative-reco-software ${PWD}/env_async.env}"
+ALIEN_JDL_O2DPG_ASYNC_RECO_FROMSTAGE=${ALIEN_JDL_O2DPG_ASYNC_RECO_FROMSTAGE:-RECO}
+remainingargs="${remainingargs} ${ALIEN_JDL_O2DPG_ASYNC_RECO_TAG:+--alternative-reco-software ${PWD}/env_async.env@${ALIEN_JDL_O2DPG_ASYNC_RECO_FROMSTAGE}}"
 # potentially add CCDB timemachine timestamp
 remainingargs="${remainingargs} ${ALIEN_JDL_CCDB_CONDITION_NOT_AFTER:+--condition-not-after ${ALIEN_JDL_CCDB_CONDITION_NOT_AFTER}}"