AsyncNet: option for time based tracing and trace path (pytorch#13440)

Summary:
Pull Request resolved: pytorch#13440

Time based tracing is easier to look at when multiple nets are running asynchronously.
This diff also adds an option to change the path to dump trace files.

Reviewed By: aazzolini, ilia-cher

Differential Revision: D12479259

fbshipit-source-id: 94d379634ba7b90c111c92b1136ffa4226b8bb8c

Author: jspark1105

caffe2/core/net_async_tracing.cc

@@ -53,11 +53,18 @@ int getCounterForNetName(const std::string& net_name) {
   return counter;
 }
 
-Tracer::Tracer(const NetBase* net, const std::string& net_name)
-    : net_(net), filename_(net_name), iter_(0) {
+Tracer::Tracer(
+    const NetBase* net,
+    const std::string& net_name,
+    TracingConfig config)
+    : net_(net),
+      filename_(net_name),
+      iter_(0),
+      dumping_iter_(0),
+      config_(config) {
   std::replace(filename_.begin(), filename_.end(), '/', '_');
-  filename_ = FLAGS_caffe2_net_async_tracing_filepath + "/" + filename_ +
-      +"_id_" + caffe2::to_string(getCounterForNetName(net_name));
+  filename_ = this->config().filepath + "/" + filename_ + "_id_" +
+      caffe2::to_string(getCounterForNetName(net_name));
   timer_.Start();
 }
 
@@ -251,6 +258,10 @@ int Tracer::bumpIter() {
   return iter_++;
 }
 
+int Tracer::bumpDumpingIter() {
+  return dumping_iter_++;
+}
+
 void Tracer::dumpTracingResultAndClearEvents(const std::string& file_suffix) {
   if (events_.empty() || filename_.empty()) {
     return;
@@ -294,7 +305,9 @@ void TracerGuard::addArgument(TracingField field, const char* value) {
       event_.category_ = value;
       break;
     }
-    default: { CAFFE_THROW("Unexpected tracing string field ", field); }
+    default: {
+      CAFFE_THROW("Unexpected tracing string field ", field);
+    }
   }
 }
 
@@ -316,7 +329,9 @@ void TracerGuard::addArgument(TracingField field, int value) {
       event_.thread_label_ = value;
       break;
     }
-    default: { CAFFE_THROW("Unexpected tracing int field ", field); }
+    default: {
+      CAFFE_THROW("Unexpected tracing int field ", field);
+    }
   }
 }
 
@@ -388,25 +403,66 @@ bool hasEnableTracingFlag(const NetBase* net) {
   return GetFlagArgument(net->debug_def(), "enable_tracing", false);
 }
 
+TracingConfig getTracingConfigFromNet(const NetBase* net) {
+  ArgumentHelper arg_helper(net->debug_def());
+  TracingConfig cfg;
+
+  cfg.mode = (arg_helper.GetSingleArgument<std::string>("tracing_mode", "") ==
+              "GLOBAL_TIMESLICE")
+      ? TracingMode::GLOBAL_TIMESLICE
+      : TracingMode::EVERY_K_ITERATIONS;
+
+  cfg.filepath = arg_helper.GetSingleArgument<std::string>(
+      "tracing_filepath", FLAGS_caffe2_net_async_tracing_filepath);
+
+  cfg.trace_every_nth_batch = arg_helper.GetSingleArgument<int>(
+      "trace_every_nth_batch", FLAGS_caffe2_net_async_tracing_nth);
+  cfg.dump_every_nth_batch = arg_helper.GetSingleArgument<int>(
+      "dump_every_nth_batch", FLAGS_caffe2_net_async_tracing_dumping_nth);
+
+  cfg.trace_for_n_ms =
+      arg_helper.GetSingleArgument<int>("trace_for_n_ms", cfg.trace_for_n_ms);
+  cfg.trace_every_n_ms = arg_helper.GetSingleArgument<int>(
+      "trace_every_n_ms", cfg.trace_every_n_ms);
+
+  return cfg;
+};
+
 std::shared_ptr<Tracer> create(
     const NetBase* net,
     const std::string& net_name) {
   // Enable the tracer if the net has the "enable_tracing" argument set OR
   // if the command line option includes the net name option in the list of
   // tracable nets.
   bool trace_net = hasEnableTracingFlag(net) || isTraceableNetName(net_name);
-  return trace_net ? std::make_shared<Tracer>(net, net_name) : nullptr;
+  return trace_net
+      ? std::make_shared<Tracer>(net, net_name, getTracingConfigFromNet(net))
+      : nullptr;
 }
 
 bool startIter(const std::shared_ptr<Tracer>& tracer) {
   if (!tracer) {
     return false;
   }
   auto iter = tracer->bumpIter();
-  auto is_enabled = iter % FLAGS_caffe2_net_async_tracing_nth == 0;
+  bool is_enabled;
+  bool should_dump;
+  if (tracer->config().mode == TracingMode::EVERY_K_ITERATIONS) {
+    is_enabled = iter % tracer->config().trace_every_nth_batch == 0;
+    should_dump = iter % tracer->config().dump_every_nth_batch == 0;
+  } else {
+    using namespace std::chrono;
+    auto ms =
+        duration_cast<milliseconds>(system_clock::now().time_since_epoch())
+            .count();
+    is_enabled = (ms % tracer->config().trace_every_n_ms) <
+        tracer->config().trace_for_n_ms;
+    // dump just after disabled tracing
+    should_dump = tracer->isEnabled() && !is_enabled;
+  }
   tracer->setEnabled(is_enabled);
-  if (iter % FLAGS_caffe2_net_async_tracing_dumping_nth == 0) {
-    int dumping_iter = iter / FLAGS_caffe2_net_async_tracing_dumping_nth;
+  if (should_dump) {
+    int dumping_iter = tracer->bumpDumpingIter();
     tracer->dumpTracingResultAndClearEvents(caffe2::to_string(dumping_iter));
   }
   return is_enabled;

caffe2/core/net_async_tracing.h

@@ -50,9 +50,30 @@ enum TracingField {
   TRACE_CATEGORY,
 };
 
+enum class TracingMode {
+  EVERY_K_ITERATIONS,
+  GLOBAL_TIMESLICE,
+};
+
+struct TracingConfig {
+  TracingMode mode{TracingMode::EVERY_K_ITERATIONS};
+  std::string filepath{"/tmp"};
+
+  // for TracingMode::EVERY_K_ITERATIONS
+  int64_t trace_every_nth_batch = 100;
+  int64_t dump_every_nth_batch = 10000;
+
+  // for TracingMode::GLOBAL_TIMESLICE
+  int64_t trace_every_n_ms = 2 * 60 * 1000; // 2min
+  int64_t trace_for_n_ms = 1000; // 1sec
+};
+
 class CAFFE2_API Tracer {
  public:
-  Tracer(const NetBase* net, const std::string& net_name);
+  Tracer(
+      const NetBase* net,
+      const std::string& net_name,
+      TracingConfig = TracingConfig{});
 
   void recordEvent(const TracerEvent& event);
   std::string opTraceName(const OperatorBase* op);
@@ -62,7 +83,11 @@ class CAFFE2_API Tracer {
   void renameThreads();
   void setEnabled(bool enabled);
   bool isEnabled() const;
+  const TracingConfig& config() {
+    return config_;
+  }
   int bumpIter();
+  int bumpDumpingIter();
   // Dump the tracing result to file with given suffix, and then
   // clear current events.
   void dumpTracingResultAndClearEvents(const std::string& file_suffix);
@@ -77,6 +102,8 @@ class CAFFE2_API Tracer {
   bool enabled_ = false;
   Timer timer_;
   int iter_;
+  int dumping_iter_;
+  TracingConfig config_;
 
   friend class TracerGuard;
 };
@@ -115,7 +142,9 @@ CAFFE2_API int extractShardId(const std::string& name);
 // line flag)
 CAFFE2_API bool isTraceableNetName(const std::string& net_name);
 
-CAFFE2_API std::shared_ptr<Tracer> create(const NetBase* net, const std::string& net_name);
+CAFFE2_API std::shared_ptr<Tracer> create(
+    const NetBase* net,
+    const std::string& net_name);
 CAFFE2_API bool startIter(const std::shared_ptr<Tracer>& tracer);
 
 } // namespace tracing

caffe2/core/net_async_tracing_test.cc

@@ -37,6 +37,78 @@ TEST(NetAsyncTracingTest, ExtractShardId) {
   testExtractShardId("FC:shard:15", 15);
 }
 
+TEST(NetAsyncTracingTest, EveryKIteration) {
+  const auto spec = R"DOC(
+      name: "example"
+      type: "async_scheduling"
+      arg {
+        name: "enable_tracing"
+        i: 1
+      }
+      arg {
+        name: "tracing_mode"
+        s: "EVERY_K_ITERATIONS"
+      }
+      arg {
+        name: "tracing_filepath"
+        s: "/tmp"
+      }
+      arg {
+        name: "trace_every_nth_batch"
+        i: 1
+      }
+      arg {
+        name: "dump_every_nth_batch"
+        i: 1
+      }
+      op {
+        output: "out"
+        type: "UniformFill"
+      }
+)DOC";
+
+  NetDef net_def;
+  CAFFE_ENFORCE(TextFormat::ParseFromString(spec, &net_def));
+
+  Workspace ws;
+  std::unique_ptr<NetBase> net(CreateNet(net_def, &ws));
+  net->Run();
+}
+
+TEST(NetAsyncTracingTest, GlobalTimeSlice) {
+  const auto spec = R"DOC(
+      name: "example"
+      type: "async_scheduling"
+      arg {
+        name: "enable_tracing"
+        i: 1
+      }
+      arg {
+        name: "tracing_filepath"
+        s: "/tmp"
+      }
+      arg {
+        name: "trace_for_n_ms"
+        i: 1
+      }
+      arg {
+        name: "trace_every_n_ms"
+        i: 1
+      }
+      op {
+        output: "out"
+        type: "UniformFill"
+      }
+)DOC";
+
+  NetDef net_def;
+  CAFFE_ENFORCE(TextFormat::ParseFromString(spec, &net_def));
+
+  Workspace ws;
+  std::unique_ptr<NetBase> net(CreateNet(net_def, &ws));
+  net->Run();
+}
+
 } // namespace tracing
 
 } // namespace caffe2