From 7c418c6d73b358f884a0bf9c4ca40f3b8d7155f9 Mon Sep 17 00:00:00 2001
From: swenzel <sandro.wenzel@cern.ch>
Date: Mon, 16 Dec 2024 14:48:55 +0100
Subject: [PATCH] EventPools: Timeout option and stat file

This commit introduces the possibility to create event pools under
a given TTL. So even, if the event search times out, events already found
will be saved to the pool. This feature allows to make best use of a TTL GRID job
and there is no danger of overestimating an event count in the script.

The actual number of events found in an event pool job is now
also reported in a special 0_0_0_N.stat file, just like in ordinary AO2D productions.
This will help for (MonaLisa) accounting.
---
 MC/bin/o2dpg_sim_workflow.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/MC/bin/o2dpg_sim_workflow.py b/MC/bin/o2dpg_sim_workflow.py
index b09df8222..51cfa86f7 100755
--- a/MC/bin/o2dpg_sim_workflow.py
+++ b/MC/bin/o2dpg_sim_workflow.py
@@ -791,8 +791,13 @@ def getDPL_global_options(bigshm=False, ccdbbackend=True):
        cmd = 'export HEPMCEVENTSKIP=$(${O2DPG_ROOT}/UTILS/ReadHepMCEventSkip.sh ../HepMCEventSkip.json ' + str(tf) + ');'
      SGNGENtask['cmd'] = cmd
 
-
-   SGNGENtask['cmd'] +='${O2_ROOT}/bin/o2-sim --noGeant -j 1 --field ccdb --vertexMode ' + vtxmode    \
+   generationtimeout = -1 # possible timeout for event pool generation
+   if args.make_evtpool:
+     JOBTTL=environ.get('JOBTTL', None)
+     if JOBTTL != None:
+       generationtimeout = 0.95*int(JOBTTL) # for GRID jobs, determine timeout automatically
+   SGNGENtask['cmd'] +=('','timeout ' + str(generationtimeout) + ' ')[args.make_evtpool and generationtimeout>0] \
+                     + '${O2_ROOT}/bin/o2-sim --noGeant -j 1 --field ccdb --vertexMode ' + vtxmode    \
                      + ' --run ' + str(args.run) + ' ' + str(CONFKEY) + str(TRIGGER)                  \
                      + ' -g ' + str(GENERATOR) + ' ' + str(INIFILE) + ' -o genevents ' + embeddinto   \
                      + ('', ' --timestamp ' + str(args.timestamp))[args.timestamp!=-1]                \
@@ -804,6 +809,10 @@ def getDPL_global_options(bigshm=False, ccdbbackend=True):
       workflow['stages'].append(SGNGENtask)
       signalneeds = signalneeds + [SGNGENtask['name']]
    if args.make_evtpool:
+      if generationtimeout > 0:
+         # final adjustment of command for event pools and timeout --> we need to analyse the return code
+         # if we have a timeout then we finish what we can and are also happy with return code 124
+         SGNGENtask['cmd'] += ' ; RC=$? ; [[ ${RC} == 0 || ${RC} == 124 ]]'
       continue
 
    # GeneratorFromO2Kine parameters are needed only before the transport
@@ -1594,6 +1603,8 @@ def addQCPerTF(taskName, needs, readerCommand, configFilePath, objectsFile=''):
    tfpool=['tf' + str(tf) + '/genevents_Kine.root' for tf in range(1, NTIMEFRAMES + 1)]
    POOL_merge_task = createTask(name='poolmerge', needs=wfneeds, lab=["POOL"], mem='2000', cpu='1')
    POOL_merge_task['cmd'] = '${O2DPG_ROOT}/UTILS/root_merger.py -o evtpool.root -i ' + ','.join(tfpool)
+   # also create the stat file with the event count
+   POOL_merge_task['cmd'] += '; RC=$?; root -l -q -b -e "auto f=TFile::Open(\\\"evtpool.root\\\"); auto t=(TTree*)f->Get(\\\"o2sim\\\"); int n=t->GetEntries(); std::ofstream((\\\"0_0_0_\\\"+std::to_string(n)+\\\".stat\\\").c_str()).close();" ; [[ ${RC} == 0 ]]'
    workflow['stages'].append(POOL_merge_task)
 
 # adjust for alternate (RECO) software environments