From 7d174673bde0ffd078bf053f8dd49703060591af Mon Sep 17 00:00:00 2001
From: Tomasz Szpartaluk <tomeksz@onspecta.com>
Date: Fri, 3 Sep 2021 17:22:46 +0200
Subject: [PATCH] add profiler wrapper

---
 utils/benchmark.py        |  5 +++--
 utils/profiler_wrapper.py | 25 +++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)
 create mode 100644 utils/profiler_wrapper.py

diff --git a/utils/benchmark.py b/utils/benchmark.py
index 86ba1124..54a2a049 100644
--- a/utils/benchmark.py
+++ b/utils/benchmark.py
@@ -3,6 +3,7 @@
 import utils.misc as utils
 import utils.dataset as utils_ds
 from tqdm.auto import tqdm
+from utils.profiler_wrapper import TBTracer
 
 intra_op_parallelism_threads = None
 
@@ -111,7 +112,7 @@ def run_model(single_pass_func, runner, dataset, batch_size, num_of_runs, timeou
         if dataset.available_instances < num_of_runs * batch_size:
             utils.print_goodbye_message_and_die(
                 f"Number of runs requested exceeds number of instances available in dataset!")
-
+    tracer = TBTracer()
     try:
         if num_of_runs is None:
             single_pass_func(runner, dataset)
@@ -123,7 +124,7 @@ def run_model(single_pass_func, runner, dataset, batch_size, num_of_runs, timeou
                 single_pass_func(runner, dataset)
     except utils_ds.OutOfInstances:
         pass
-
+    tracer.write()
     return dataset.summarize_accuracy(), runner.print_performance_metrics(batch_size)
 
 
diff --git a/utils/profiler_wrapper.py b/utils/profiler_wrapper.py
new file mode 100644
index 00000000..fa6dcf1e
--- /dev/null
+++ b/utils/profiler_wrapper.py
@@ -0,0 +1,25 @@
+import tensorflow as tf
+import os
+from datetime import datetime
+
+def print_prof():
+  if os.getenv("DLS_PROFILER", 0):
+    try:
+      tf.DLS.print_profile_data()  
+    except AttributeError:
+      print("Non dls tf")
+
+class TBTracer:
+  def __init__(self):
+    self.should_trace = os.getenv("TRACE", 0)
+    if self.should_trace: 
+      # Set up logging.
+      options = tf.profiler.experimental.ProfilerOptions()
+      stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+      self.logdir = 'logs/func/%s' % stamp
+      tf.profiler.experimental.start(self.logdir, options = options)
+
+  def write(self):
+    if self.should_trace:
+      tf.profiler.experimental.stop()
+    print_prof()