From 1bb82c6fb2d29b8c68d9ce1820654230bb9e5797 Mon Sep 17 00:00:00 2001
From: Marcel R <github.riga@icloud.com>
Date: Thu, 15 Dec 2022 18:08:29 +0100
Subject: [PATCH 1/3] Move tf session to global cache using SessionCache.

---
 .../Phase2L1ParticleFlow/interface/TauNNId.h  |  9 ++----
 .../plugins/L1NNTauProducer.cc                | 28 ++++++++-----------
 L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc |  4 +--
 3 files changed, 15 insertions(+), 26 deletions(-)

diff --git a/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h b/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h
index c26d4abf3b07a..8613c023648a5 100644
--- a/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h
+++ b/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h
@@ -5,14 +5,9 @@
 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
 #include "DataFormats/L1TParticleFlow/interface/PFCandidate.h"
 
-struct TauNNTFCache {
-  TauNNTFCache() : graphDef(nullptr) {}
-  std::atomic<tensorflow::GraphDef *> graphDef;
-};
-
 class TauNNId {
 public:
-  TauNNId(const std::string &iInput, const TauNNTFCache *cache, const std::string &iWeightFile, int iNParticles);
+  TauNNId(const std::string &iInput, const tensorflow::Session *session, const std::string &iWeightFile, int iNParticles);
   ~TauNNId();
 
   void setNNVectorVar();
@@ -20,7 +15,7 @@ class TauNNId {
   float compute(const l1t::PFCandidate &iSeed, l1t::PFCandidateCollection &iParts);
 
 private:
-  tensorflow::Session *session_;
+  const tensorflow::Session *session_;
   std::vector<float> NNvectorVar_;
   std::string fInput_;
   int fNParticles_;
diff --git a/L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc b/L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc
index 474986195d7fc..ef04f91444e46 100644
--- a/L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc
+++ b/L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc
@@ -14,14 +14,14 @@
 
 using namespace l1t;
 
-class L1NNTauProducer : public edm::stream::EDProducer<edm::GlobalCache<TauNNTFCache>> {
+class L1NNTauProducer : public edm::stream::EDProducer<edm::GlobalCache<tensorflow::SessionCache>> {
 public:
-  explicit L1NNTauProducer(const edm::ParameterSet&, const TauNNTFCache*);
+  explicit L1NNTauProducer(const edm::ParameterSet&, const tensorflow::SessionCache*);
   ~L1NNTauProducer() override;
 
   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
-  static std::unique_ptr<TauNNTFCache> initializeGlobalCache(const edm::ParameterSet&);
-  static void globalEndJob(const TauNNTFCache*);
+  static std::unique_ptr<tensorflow::SessionCache> initializeGlobalCache(const edm::ParameterSet&);
+  static void globalEndJob(const tensorflow::SessionCache*){};
 
 private:
   std::unique_ptr<TauNNId> fTauNNId_;
@@ -40,7 +40,7 @@ class L1NNTauProducer : public edm::stream::EDProducer<edm::GlobalCache<TauNNTFC
 
 static constexpr float track_trigger_eta_max = 2.5;
 
-L1NNTauProducer::L1NNTauProducer(const edm::ParameterSet& cfg, const TauNNTFCache* cache)
+L1NNTauProducer::L1NNTauProducer(const edm::ParameterSet& cfg, const tensorflow::SessionCache* cache)
     : fSeedPt_(cfg.getParameter<double>("seedpt")),
       fConeSize_(cfg.getParameter<double>("conesize")),
       fTauSize_(cfg.getParameter<double>("tausize")),
@@ -49,22 +49,16 @@ L1NNTauProducer::L1NNTauProducer(const edm::ParameterSet& cfg, const TauNNTFCach
       fL1PFToken_(consumes<vector<l1t::PFCandidate>>(cfg.getParameter<edm::InputTag>("L1PFObjects"))) {
   std::string lNNFile = cfg.getParameter<std::string>("NNFileName");  //,"L1Trigger/Phase2L1Taus/data/tau_3layer.pb");
   fTauNNId_ = std::make_unique<TauNNId>(
-      lNNFile.find("v0") == std::string::npos ? "input_1:0" : "dense_1_input:0", cache, lNNFile, fNParticles_);
+      lNNFile.find("v0") == std::string::npos ? "input_1:0" : "dense_1_input:0", cache->getSession(), lNNFile, fNParticles_);
   produces<l1t::PFTauCollection>("L1PFTausNN");
 }
-std::unique_ptr<TauNNTFCache> L1NNTauProducer::initializeGlobalCache(const edm::ParameterSet& cfg) {
+
+std::unique_ptr<tensorflow::SessionCache> L1NNTauProducer::initializeGlobalCache(const edm::ParameterSet& cfg) {
   tensorflow::setLogging("3");
-  std::string lNNFile = cfg.getParameter<std::string>("NNFileName");
-  edm::FileInPath fp(lNNFile);
-  TauNNTFCache* cache = new TauNNTFCache();
-  cache->graphDef = tensorflow::loadGraphDef(fp.fullPath());
-  return std::unique_ptr<TauNNTFCache>(cache);
-}
-void L1NNTauProducer::globalEndJob(const TauNNTFCache* cache) {
-  if (cache->graphDef != nullptr) {
-    delete cache->graphDef;
-  }
+  std::string graphPath = edm::FileInPath(cfg.getParameter<std::string>("NNFileName")).fullPath();
+  return std::make_unique<tensorflow::SessionCache>(graphPath);
 }
+
 void L1NNTauProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
   edm::Handle<l1t::PFCandidateCollection> l1PFCandidates;
   iEvent.getByToken(fL1PFToken_, l1PFCandidates);
diff --git a/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc b/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc
index 2458abe181816..6c28bebf6cab8 100644
--- a/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc
+++ b/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc
@@ -5,10 +5,9 @@
 
 static constexpr unsigned int n_particles_max = 10;
 
-TauNNId::TauNNId(const std::string &iInput, const TauNNTFCache *cache, const std::string &iWeightFile, int iNParticles) {
+TauNNId::TauNNId(const std::string &iInput, const tensorflow::Session * session, const std::string &iWeightFile, int iNParticles) : session_(session) {
   NNvectorVar_.clear();
   edm::FileInPath fp(iWeightFile);
-  session_ = tensorflow::createSession(cache->graphDef);
   fNParticles_ = iNParticles;
 
   fPt_ = std::make_unique<float[]>(fNParticles_);
@@ -19,6 +18,7 @@ TauNNId::TauNNId(const std::string &iInput, const TauNNTFCache *cache, const std
 }
 
 TauNNId::~TauNNId() { tensorflow::closeSession(session_); }
+
 void TauNNId::setNNVectorVar() {
   NNvectorVar_.clear();
   for (int i0 = 0; i0 < fNParticles_; i0++) {

From cb4c5175785b24d01a80ab2de1078354092d6a1f Mon Sep 17 00:00:00 2001
From: Marcel R <github.riga@icloud.com>
Date: Thu, 15 Dec 2022 18:08:41 +0100
Subject: [PATCH 2/3] Code format.

---
 L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h        | 5 ++++-
 L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc | 6 ++++--
 L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc             | 6 +++++-
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h b/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h
index 8613c023648a5..d4d6372368259 100644
--- a/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h
+++ b/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h
@@ -7,7 +7,10 @@
 
 class TauNNId {
 public:
-  TauNNId(const std::string &iInput, const tensorflow::Session *session, const std::string &iWeightFile, int iNParticles);
+  TauNNId(const std::string &iInput,
+          const tensorflow::Session *session,
+          const std::string &iWeightFile,
+          int iNParticles);
   ~TauNNId();
 
   void setNNVectorVar();
diff --git a/L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc b/L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc
index ef04f91444e46..7ac2206e0cedc 100644
--- a/L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc
+++ b/L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc
@@ -48,8 +48,10 @@ L1NNTauProducer::L1NNTauProducer(const edm::ParameterSet& cfg, const tensorflow:
       fNParticles_(cfg.getParameter<int>("nparticles")),
       fL1PFToken_(consumes<vector<l1t::PFCandidate>>(cfg.getParameter<edm::InputTag>("L1PFObjects"))) {
   std::string lNNFile = cfg.getParameter<std::string>("NNFileName");  //,"L1Trigger/Phase2L1Taus/data/tau_3layer.pb");
-  fTauNNId_ = std::make_unique<TauNNId>(
-      lNNFile.find("v0") == std::string::npos ? "input_1:0" : "dense_1_input:0", cache->getSession(), lNNFile, fNParticles_);
+  fTauNNId_ = std::make_unique<TauNNId>(lNNFile.find("v0") == std::string::npos ? "input_1:0" : "dense_1_input:0",
+                                        cache->getSession(),
+                                        lNNFile,
+                                        fNParticles_);
   produces<l1t::PFTauCollection>("L1PFTausNN");
 }
 
diff --git a/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc b/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc
index 6c28bebf6cab8..392f3d2e46e68 100644
--- a/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc
+++ b/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc
@@ -5,7 +5,11 @@
 
 static constexpr unsigned int n_particles_max = 10;
 
-TauNNId::TauNNId(const std::string &iInput, const tensorflow::Session * session, const std::string &iWeightFile, int iNParticles) : session_(session) {
+TauNNId::TauNNId(const std::string &iInput,
+                 const tensorflow::Session *session,
+                 const std::string &iWeightFile,
+                 int iNParticles)
+    : session_(session) {
   NNvectorVar_.clear();
   edm::FileInPath fp(iWeightFile);
   fNParticles_ = iNParticles;

From ce1a2240bcc58fd2449ce21629910f979968491b Mon Sep 17 00:00:00 2001
From: Marcel R <github.riga@icloud.com>
Date: Fri, 16 Dec 2022 10:33:29 +0100
Subject: [PATCH 3/3] Remove obsolete destructor.

---
 L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h | 2 +-
 L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc      | 2 --
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h b/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h
index d4d6372368259..ea01a3e2c578d 100644
--- a/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h
+++ b/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h
@@ -11,7 +11,7 @@ class TauNNId {
           const tensorflow::Session *session,
           const std::string &iWeightFile,
           int iNParticles);
-  ~TauNNId();
+  ~TauNNId(){};
 
   void setNNVectorVar();
   float EvaluateNN();
diff --git a/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc b/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc
index 392f3d2e46e68..3cce460c0337c 100644
--- a/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc
+++ b/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc
@@ -21,8 +21,6 @@ TauNNId::TauNNId(const std::string &iInput,
   fInput_ = iInput;
 }
 
-TauNNId::~TauNNId() { tensorflow::closeSession(session_); }
-
 void TauNNId::setNNVectorVar() {
   NNvectorVar_.clear();
   for (int i0 = 0; i0 < fNParticles_; i0++) {