From e62bcf96ca024c8e11b8e6b90e95538a31591598 Mon Sep 17 00:00:00 2001
From: Sandro Wenzel <sandro.wenzel@cern.ch>
Date: Tue, 16 Dec 2025 14:42:38 +0100
Subject: [PATCH 1/2] Refactor of TPC clusterization configuration

* beautification of configuration
* coupling to task-finalizer to be able to inject config-keys from outside
* make sensitive to GPU_proc_nn for the ML clusterization keys
---
 MC/bin/o2dpg_sim_workflow.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/MC/bin/o2dpg_sim_workflow.py b/MC/bin/o2dpg_sim_workflow.py
index b6686e425..ca364a011 100755
--- a/MC/bin/o2dpg_sim_workflow.py
+++ b/MC/bin/o2dpg_sim_workflow.py
@@ -1253,7 +1253,22 @@ def getDigiTaskName(det):
        tpcclustertasks.append(taskname)
        tpcclussect = createTask(name=taskname, needs=tpcclusterneed, tf=tf, cwd=timeframeworkdir, lab=["RECO"], cpu='2', mem='8000')
        digitmergerstr = '${O2_ROOT}/bin/o2-tpc-chunkeddigit-merger --tpc-sectors ' + str(s)+'-'+str(s+sectorpertask-1) + ' --tpc-lanes ' + str(NWORKERS_TF) + ' | '
-       tpcclussect['cmd'] = (digitmergerstr,'')[args.no_tpc_digitchunking] + ' ${O2_ROOT}/bin/o2-tpc-reco-workflow ' + getDPL_global_options(bigshm=True) + ' --input-type ' + ('digitizer','digits')[args.no_tpc_digitchunking] + ' --output-type clusters,send-clusters-per-sector --tpc-native-cluster-writer \" --outfile tpc-native-clusters-part'+ str((int)(s/sectorpertask)) + '.root\" --tpc-sectors ' + str(s)+'-'+str(s+sectorpertask-1) + ' ' + putConfigValues(["GPU_global"], {"GPU_proc.ompThreads" : 4}) + ('',' --disable-mc')[args.no_mc_labels]
+       tpcclussect['cmd'] = (digitmergerstr,'')[args.no_tpc_digitchunking] + task_finalizer(
+          [
+            '${O2_ROOT}/bin/o2-tpc-reco-workflow',
+            getDPL_global_options(bigshm=True),
+            '--input-type ' + ('digitizer','digits')[args.no_tpc_digitchunking],
+            '--output-type clusters,send-clusters-per-sector',
+            f'--tpc-native-cluster-writer \" --outfile tpc-native-clusters-part{(int)(s/sectorpertask)}.root\"',
+            f'--tpc-sectors {s}-{s+sectorpertask-1}',
+            putConfigValues(["GPU_global",
+                             "GPU_proc_nn",
+                             "TPCCorrMap",
+                             "TPCGasParam"], {"GPU_proc.ompThreads" : 4}),
+            '--disable-mc' if args.no_mc_labels else None
+          ], configname="tpcclusterizertask"
+       )
+
        tpcclussect['env'] = { "OMP_NUM_THREADS" : "4" , "TBB_NUM_THREADS" : "4" }
        tpcclussect['semaphore'] = "tpctriggers.root"
        tpcclussect['retry_count'] = 2  # the task has a race condition --> makes sense to retry
@@ -1264,6 +1279,7 @@ def getDigiTaskName(det):
      workflow['stages'].append(TPCCLUSMERGEtask)
      tpcreconeeds.append(TPCCLUSMERGEtask['name'])
    else:
+     # TODO: adapt this to the case above and merge code / avoid code duplication
      tpcclus = createTask(name='tpccluster_' + str(tf), needs=tpcclusterneed, tf=tf, cwd=timeframeworkdir, lab=["RECO"], cpu=NWORKERS_TF, mem='2000')
      tpcclus['cmd'] = '${O2_ROOT}/bin/o2-tpc-chunkeddigit-merger --tpc-lanes ' + str(NWORKERS_TF)
      tpcclus['cmd'] += ' | ${O2_ROOT}/bin/o2-tpc-reco-workflow ' + getDPL_global_options() + ' --input-type digitizer --output-type clusters,send-clusters-per-sector ' + putConfigValues(["GPU_global","TPCGasParam","TPCCorrMap"],{"GPU_proc.ompThreads" : 1}) + ('',' --disable-mc')[args.no_mc_labels]

From 73c5c079a1b55e7f3f61fb8130703343837d05ab Mon Sep 17 00:00:00 2001
From: Sandro Wenzel <sandro.wenzel@cern.ch>
Date: Tue, 16 Dec 2025 14:47:43 +0100
Subject: [PATCH 2/2] Example script to run workflow with ML clusterizer

---
 MC/run/ANCHOR/tests/ML_clusterizer/run.sh | 68 +++++++++++++++++++++++
 1 file changed, 68 insertions(+)
 create mode 100755 MC/run/ANCHOR/tests/ML_clusterizer/run.sh

diff --git a/MC/run/ANCHOR/tests/ML_clusterizer/run.sh b/MC/run/ANCHOR/tests/ML_clusterizer/run.sh
new file mode 100755
index 000000000..9f800d868
--- /dev/null
+++ b/MC/run/ANCHOR/tests/ML_clusterizer/run.sh
@@ -0,0 +1,68 @@
+#!/bin/bash
+
+#JDL_OUTPUT=*.txt@disk=1,AO2D*.root@disk=2,*.log@disk=1,*stat*@disk=1,*.json@disk=1,debug*tgz@disk=2,tf*/coll*.root
+#JDL_ERROROUTPUT=*.txt@disk=1,AO2D*.root@disk=2,*.log@disk=1,*.json@disk=1,debug*tgz@disk=2
+#JDL_PACKAGE=O2PDPSuite::daily-20251215-0000-1
+
+#
+# An **EXAMPLE** showing injection of GPU_proc_nn config_keys into the workflow creation, so
+# that TPC clusterizer runs with ML kernel.
+#
+
+# example anchoring
+export ALIEN_JDL_LPMANCHORPASSNAME=apass4
+export ALIEN_JDL_MCANCHOR=apass4
+export ALIEN_JDL_CPULIMIT=8
+export ALIEN_JDL_LPMRUNNUMBER=544124
+export ALIEN_JDL_LPMPRODUCTIONTYPE=MC
+export ALIEN_JDL_LPMINTERACTIONTYPE=PbPb
+export ALIEN_JDL_LPMPRODUCTIONTAG=MLClusterTest
+export ALIEN_JDL_LPMANCHORRUN=544124
+export ALIEN_JDL_LPMANCHORPRODUCTION=LHC23zz
+export ALIEN_JDL_LPMANCHORYEAR=2023
+
+export NTIMEFRAMES=1
+export PRODSPLIT=${ALIEN_O2DPG_GRIDSUBMIT_PRODSPLIT:-100}
+export SPLITID=${ALIEN_O2DPG_GRIDSUBMIT_SUBJOBID:-50}
+export CYCLE=0
+
+# this file modifies few config key values for GPU_proc_nn for application in TPC clusterization
+LOCAL_CONFIG="customize_ml_clusterizing.json"
+cat > ${LOCAL_CONFIG} <<EOF
+{
+  "ConfigParams": {
+    "GPU_proc_nn" : {
+      "nnInferenceDevice" : "CPU",
+      "nnInferenceAllocateDevMem" : "0",
+      "nnInferenceInputDType" : "FP16",
+      "nnInferenceOutputDType" : "FP16",
+      "nnSigmoidTrafoClassThreshold" : "1",
+      "nnClusterizerAddIndexData" : "1",
+      "nnInferenceIntraOpNumThreads" : "1",
+      "nnInferenceInterOpNumThreads" : "1",
+      "nnInferenceVerbosity" : "3",
+      "nnClusterizerVerbosity" : "2",
+      "nnClusterizerApplyCfDeconvolution" : "0",
+      "nnLoadFromCCDB" : "1",
+      "nnClusterizerBoundaryFillValue" : "0",
+      "nnInferenceOrtProfiling" : "0",
+      "nnClusterizerSetDeconvolutionFlags" : "1",
+      "nnClusterizerSizeInputRow" : "1",
+      "nnClusterizerSizeInputPad" : "4",
+      "nnClusterizerSizeInputTime" : "4",
+      "nnClassThreshold" : "0.05",
+      "nnClusterizerUseCfRegression" : "0",
+      "nnClusterizerBatchedMode" : "262144",
+      "applyNNclusterizer" : "1"
+    }
+  }
+  "Executables": {   
+  }
+}
+EOF
+
+# we use pp event gen just to get faster results
+export ALIEN_JDL_ANCHOR_SIM_OPTIONS="-gen pythia8pp -confKey \"GeometryManagerParam.useParallelWorld=1;GeometryManagerParam.usePwGeoBVH=1;GeometryManagerParam.usePwCaching=1\" ${LOCAL_CONFIG:+--overwrite-config ${LOCAL_CONFIG}}"
+
+export ALIEN_JDL_O2DPGWORKFLOWTARGET="tpcclus"
+${O2DPG_ROOT}/MC/run/ANCHOR/anchorMC.sh