Add support for multiple protocols in client API (#251)

Summary:
Pull Request resolved: #251

Add a way to add support for multiple protocols and associated logger objects.

Reviewed By: ilia-cher

Differential Revision: D28579660

fbshipit-source-id: 497c45eca0529ba0620d635c6605fc0f20172d9c

Author: gdankel

libkineto/src/AbstractConfig.h

@@ -52,6 +52,7 @@ class AbstractConfig {
     return timestamp_;
   }
 
+  // Source config string that this was parsed from
   const std::string& source() const {
     return source_;
   }

libkineto/src/ActivityLoggerFactory.h

@@ -0,0 +1,64 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ * All rights reserved.
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <algorithm>
+#include <fmt/format.h>
+#include <functional>
+#include <map>
+#include <string>
+
+namespace KINETO_NAMESPACE {
+
+class ActivityLogger;
+
+class ActivityLoggerFactory {
+
+ public:
+  using FactoryFunc =
+    std::function<std::unique_ptr<ActivityLogger>(const std::string& url)>;
+
+  // Add logger factory for a protocol prefix
+  void addProtocol(const std::string& protocol, FactoryFunc f) {
+    factories_[tolower(protocol)] = f;
+  }
+
+  // Create a logger, invoking the factory for the protocol specified in url
+  std::unique_ptr<ActivityLogger> makeLogger(const std::string& url) const {
+    std::string protocol = extractProtocol(url);
+    auto it = factories_.find(tolower(protocol));
+    if (it != factories_.end()) {
+      return it->second(stripProtocol(url));
+    }
+    throw std::invalid_argument(fmt::format(
+        "No logger registered for the {} protocol prefix",
+        protocol));
+    return nullptr;
+  }
+
+ private:
+  static std::string tolower(std::string s) {
+    std::transform(s.begin(), s.end(), s.begin(),
+        [](unsigned char c) { return std::tolower(c); }
+    );
+    return s;
+  }
+
+  static std::string extractProtocol(std::string url) {
+    return url.substr(0, url.find("://"));
+  }
+
+  static std::string stripProtocol(std::string url) {
+    size_t pos = url.find("://");
+    return pos == url.npos ? url : url.substr(pos + 3);
+  }
+
+  std::map<std::string, FactoryFunc> factories_;
+};
+
+} // namespace KINETO_NAMESPACE

libkineto/src/ActivityProfilerController.cpp

@@ -10,6 +10,7 @@
 #include <chrono>
 #include <thread>
 
+#include "ActivityLoggerFactory.h"
 #include "ActivityTrace.h"
 #include "CuptiActivityInterface.h"
 #include "ThreadUtil.h"
@@ -39,25 +40,29 @@ ActivityProfilerController::~ActivityProfilerController() {
   VLOG(0) << "Stopped activity profiler";
 }
 
+static ActivityLoggerFactory initLoggerFactory() {
+  ActivityLoggerFactory factory;
+  factory.addProtocol("file", [](const std::string& url) {
+      return std::unique_ptr<ActivityLogger>(new ChromeTraceLogger(url));
+  });
+  return factory;
+}
+
 static ActivityLoggerFactory& loggerFactory() {
-  static ActivityLoggerFactory factory{nullptr};
+  static ActivityLoggerFactory factory = initLoggerFactory();
   return factory;
 }
 
-void ActivityProfilerController::setLoggerFactory(
-    const ActivityLoggerFactory& factory) {
-  loggerFactory() = factory;
+void ActivityProfilerController::addLoggerFactory(
+    const std::string& protocol, ActivityLoggerFactory::FactoryFunc factory) {
+  loggerFactory().addProtocol(protocol, factory);
 }
 
 static std::unique_ptr<ActivityLogger> makeLogger(const Config& config) {
   if (config.activitiesLogToMemory()) {
     return std::make_unique<MemoryTraceLogger>(config);
   }
-  if (loggerFactory()) {
-    return loggerFactory()(config);
-  }
-  return std::make_unique<ChromeTraceLogger>(
-      config.activitiesLogFile());
+  return loggerFactory().makeLogger(config.activitiesLogUrl());
 }
 
 void ActivityProfilerController::profilerLoop() {
@@ -137,7 +142,7 @@ std::unique_ptr<ActivityTraceInterface> ActivityProfilerController::stopTrace()
   auto logger = std::make_unique<MemoryTraceLogger>(profiler_->config());
   profiler_->processTrace(*logger);
   profiler_->reset();
-  return std::make_unique<ActivityTrace>(std::move(logger));
+  return std::make_unique<ActivityTrace>(std::move(logger), loggerFactory());
 }
 
 void ActivityProfilerController::addMetadata(

libkineto/src/ActivityProfilerController.h

@@ -11,6 +11,7 @@
 #include <memory>
 #include <thread>
 
+#include "ActivityLoggerFactory.h"
 #include "ActivityProfiler.h"
 #include "ActivityProfilerInterface.h"
 #include "ActivityTraceInterface.h"
@@ -20,9 +21,6 @@ namespace KINETO_NAMESPACE {
 
 class Config;
 
-using ActivityLoggerFactory =
-    std::function<std::unique_ptr<ActivityLogger>(const Config&)>;
-
 class ActivityProfilerController {
  public:
   explicit ActivityProfilerController(bool cpuOnly);
@@ -32,7 +30,9 @@ class ActivityProfilerController {
 
   ~ActivityProfilerController();
 
-  static void setLoggerFactory(const ActivityLoggerFactory& factory);
+  static void addLoggerFactory(
+      const std::string& protocol,
+      ActivityLoggerFactory::FactoryFunc factory);
 
   void scheduleTrace(const Config& config);
 

libkineto/src/ActivityTrace.h

@@ -10,6 +10,7 @@
 #include <memory>
 #include <string>
 
+#include "ActivityLoggerFactory.h"
 #include "ActivityTraceInterface.h"
 #include "output_json.h"
 #include "output_membuf.h"
@@ -18,20 +19,32 @@ namespace libkineto {
 
 class ActivityTrace : public ActivityTraceInterface {
  public:
-  explicit ActivityTrace(std::unique_ptr<MemoryTraceLogger> logger)
-    : logger_(std::move(logger)) {}
+  ActivityTrace(
+      std::unique_ptr<MemoryTraceLogger> tmpLogger,
+      const ActivityLoggerFactory& factory)
+    : memLogger_(std::move(tmpLogger)),
+      loggerFactory_(factory) {
+  }
 
   const std::vector<std::unique_ptr<TraceActivity>>* activities() override {
-    return logger_->traceActivities();
+    return memLogger_->traceActivities();
   };
 
-  void save(const std::string& path) override {
-    ChromeTraceLogger chrome_logger(path);
-    logger_->log(chrome_logger);
+  void save(const std::string& url) override {
+    std::string prefix;
+    // if no protocol is specified, default to file
+    if (url.find("://") == url.npos) {
+      prefix = "file://";
+    }
+    memLogger_->log(*loggerFactory_.makeLogger(prefix + url));
   };
 
  private:
-  std::unique_ptr<MemoryTraceLogger> logger_;
+  // Activities are logged into a buffer
+  std::unique_ptr<MemoryTraceLogger> memLogger_;
+
+  // Alternative logger used by save() if protocol prefix is specified
+  const ActivityLoggerFactory& loggerFactory_;
 };
 
 } // namespace libkineto

libkineto/src/Config.cpp

@@ -61,6 +61,7 @@ const string kHeartbeatMonitorPeriodKey =
 const string kActivitiesEnabledKey = "ACTIVITIES_ENABLED";
 const string kActivityTypesKey = "ACTIVITY_TYPES";
 const string kActivitiesLogFileKey = "ACTIVITIES_LOG_FILE";
+const string kActivitiesLogUrlKey = "ACTIVITIES_LOG_URL";
 const string kActivitiesDurationKey = "ACTIVITIES_DURATION_SECS";
 const string kActivitiesDurationMsecsKey = "ACTIVITIES_DURATION_MSECS";
 const string kActivitiesIterationsKey = "ACTIVITIES_ITERATIONS";
@@ -118,16 +119,16 @@ const string kConfigFile = "/etc/libkineto.conf";
 // Max devices supported on any system
 constexpr uint8_t kMaxDevices = 8;
 
-static std::map<std::string, std::function<AbstractConfig*(const Config&)>>&
+static std::map<std::string, std::function<AbstractConfig*(Config&)>>&
 configFactories() {
-  static std::map<std::string, std::function<AbstractConfig*(const Config&)>>
+  static std::map<std::string, std::function<AbstractConfig*(Config&)>>
       factories;
   return factories;
 }
 
 void Config::addConfigFactory(
     std::string name,
-    std::function<AbstractConfig*(const Config&)> factory) {
+    std::function<AbstractConfig*(Config&)> factory) {
   configFactories()[name] = factory;
 }
 
@@ -286,6 +287,7 @@ bool Config::handleOption(const std::string& name, std::string& val) {
     activityProfilerEnabled_ = toBool(val);
   } else if (name == kActivitiesLogFileKey) {
     activitiesLogFile_ = val;
+    activitiesLogUrl_ = fmt::format("file://{}", val);
     activitiesOnDemandTimestamp_ = timestamp();
   } else if (name == kActivitiesMaxGpuBufferSizeKey) {
     activitiesMaxGpuBufferSize_ = toInt32(val) * 1024 * 1024;

libkineto/src/Config.h

@@ -53,6 +53,15 @@ class Config : public AbstractConfig {
     return activitiesLogFile_;
   }
 
+  // Log activitiy trace to this url
+  const std::string& activitiesLogUrl() const {
+    return activitiesLogUrl_;
+  }
+
+  void setActivitiesLogUrl(const std::string& url) {
+    activitiesLogUrl_ = url;
+  }
+
   bool activitiesLogToMemory() const {
     return activitiesLogToMemory_;
   }
@@ -272,7 +281,7 @@ class Config : public AbstractConfig {
 
   static void addConfigFactory(
       std::string name,
-      std::function<AbstractConfig*(const Config&)> factory);
+      std::function<AbstractConfig*(Config&)> factory);
 
   void print(std::ostream& s) const;
 
@@ -338,6 +347,8 @@ class Config : public AbstractConfig {
   // The activity profiler settings are all on-demand
   std::string activitiesLogFile_;
 
+  std::string activitiesLogUrl_;
+
   // Log activities to memory buffer
   bool activitiesLogToMemory_{false};
 

libkineto/src/init.cpp

@@ -146,7 +146,8 @@ int InitializeInjection(void) {
 }
 
 void suppressLibkinetoLogMessages() {
-  SET_LOG_SEVERITY_LEVEL(ERROR);
+  //SET_LOG_SEVERITY_LEVEL(ERROR);
+  SET_LOG_VERBOSITY_LEVEL(2, {});
 }
 
 } // extern C

libkineto/src/output_json.cpp

@@ -30,6 +30,12 @@ using namespace libkineto;
 namespace KINETO_NAMESPACE {
 
 static constexpr int kSchemaVersion = 1;
+#ifdef __linux__
+static const std::string kDefaultLogFileFmt =
+    "/tmp/libkineto_activities_{}.json";
+#else
+static const std::string kDefaultLogFileFmt = "libkineto_activities_{}.json";
+#endif
 
 void ChromeTraceLogger::handleTraceStart(
     const std::unordered_map<std::string, std::string>& metadata) {
@@ -52,6 +58,10 @@ void ChromeTraceLogger::handleTraceStart(
   "traceEvents": [)JSON";
 }
 
+static std::string defaultFileName() {
+  return fmt::format(kDefaultLogFileFmt, processId());
+}
+
 void ChromeTraceLogger::openTraceFile() {
   traceOf_.open(fileName_, std::ofstream::out | std::ofstream::trunc);
   if (!traceOf_) {
@@ -61,8 +71,8 @@ void ChromeTraceLogger::openTraceFile() {
   }
 }
 
-ChromeTraceLogger::ChromeTraceLogger(const std::string& traceFileName)
-    : fileName_(traceFileName) {
+ChromeTraceLogger::ChromeTraceLogger(const std::string& traceFileName) {
+  fileName_ = traceFileName.empty() ? defaultFileName() : traceFileName;
   traceOf_.clear(std::ios_base::badbit);
   openTraceFile();
 }

libkineto/src/output_json.h

@@ -68,6 +68,10 @@ class ChromeTraceLogger : public libkineto::ActivityLogger {
       std::unique_ptr<ActivityBuffers> buffers,
       int64_t endTime) override;
 
+  std::string traceFileName() const {
+    return fileName_;
+  }
+
  private:
 
 #ifdef HAS_CUPTI

libkineto/test/ActivityProfilerTest.cpp

@@ -164,11 +164,15 @@ class ActivityProfilerTest : public ::testing::Test {
     profiler_ = std::make_unique<ActivityProfiler>(
         cuptiActivities_, /*cpu only*/ false);
     cfg_ = std::make_unique<Config>();
+    loggerFactory.addProtocol("file", [](const std::string& url) {
+        return std::unique_ptr<ActivityLogger>(new ChromeTraceLogger(url));
+    });
   }
 
   std::unique_ptr<Config> cfg_;
   MockCuptiActivities cuptiActivities_;
   std::unique_ptr<ActivityProfiler> profiler_;
+  ActivityLoggerFactory loggerFactory;
 };
 
 
@@ -282,7 +286,7 @@ TEST_F(ActivityProfilerTest, SyncTrace) {
   profiler_->reset();
 
   // Wrapper that allows iterating over the activities
-  ActivityTrace trace(std::move(logger));
+  ActivityTrace trace(std::move(logger), loggerFactory);
   EXPECT_EQ(trace.activities()->size(), 9);
   std::map<std::string, int> activityCounts;
   std::map<int64_t, int> resourceIds;
@@ -362,7 +366,7 @@ TEST_F(ActivityProfilerTest, CorrelatedTimestampTest) {
   auto logger = std::make_unique<MemoryTraceLogger>(*cfg_);
   profiler.processTrace(*logger);
 
-  ActivityTrace trace(std::move(logger));
+  ActivityTrace trace(std::move(logger), loggerFactory);
   std::map<std::string, int> counts;
   for (auto& activity : *trace.activities()) {
     counts[activity->name()]++;