Expose profiler to python

Author: yehudaorel

python/cpp/module.cc

@@ -76,6 +76,7 @@ PYBIND11_MODULE(_ext, m)
   m.def("set_random_seed", &ctranslate2::set_random_seed, py::arg("seed"),
         "Sets the seed of random generators.");
 
+  ctranslate2::python::register_profiling(m);
   ctranslate2::python::register_logging(m);
   ctranslate2::python::register_storage_view(m);
   ctranslate2::python::register_translation_stats(m);

python/cpp/module.h

@@ -1,12 +1,14 @@
 #pragma once
 
 #include <pybind11/pybind11.h>
+#include <pybind11/chrono.h>
 
 namespace py = pybind11;
 
 namespace ctranslate2 {
   namespace python {
 
+    void register_profiling(py::module& m);
     void register_encoder(py::module& m);
     void register_generation_result(py::module& m);
     void register_generator(py::module& m);

python/cpp/profiling.cc

@@ -0,0 +1,19 @@
+#include "module.h"
+#include <sstream>
+#include <ctranslate2/profiler.h>
+
+namespace ctranslate2 {
+  namespace python {
+
+    void register_profiling(py::module& m) {
+
+      m.def("init_profiling", &ctranslate2::init_profiling);
+      m.def("dump_profiling", []() {
+        std::ostringstream oss;
+        ctranslate2::dump_profiling(oss);
+        return oss.str();
+	});
+	}
+
+  }
+}

python/ctranslate2/__init__.py

@@ -42,6 +42,7 @@
     )
     from ctranslate2.extensions import register_extensions
     from ctranslate2.logging import get_log_level, set_log_level
+    from ctranslate2.profiling import init_profiler, dump_profiler
 
     register_extensions()
     del register_extensions

python/ctranslate2/profiling.py

@@ -0,0 +1,10 @@
+from ctranslate2 import _ext, Device
+import sys
+
+def init_profiler(device = Device.cpu, num_threads = 1):
+    _ext.init_profiling(device, num_threads)
+
+
+def dump_profiler():
+    profiling_data = _ext.dump_profiling()
+    sys.stdout.write(profiling_data)