From 0cbb0796a9d74cc95dfc05f2551fba7aae8a13c3 Mon Sep 17 00:00:00 2001
From: Orel Yehuda <orel.yehuda@intel.com>
Date: Thu, 4 Jul 2024 13:05:10 -0700
Subject: [PATCH 1/2] Expose profiler to python

---
 python/cpp/module.cc            |  1 +
 python/cpp/module.h             |  2 ++
 python/cpp/profiling.cc         | 19 +++++++++++++++++++
 python/ctranslate2/__init__.py  |  1 +
 python/ctranslate2/profiling.py | 10 ++++++++++
 5 files changed, 33 insertions(+)
 create mode 100644 python/cpp/profiling.cc
 create mode 100644 python/ctranslate2/profiling.py

diff --git a/python/cpp/module.cc b/python/cpp/module.cc
index 4489d5314..0d31666eb 100644
--- a/python/cpp/module.cc
+++ b/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);
diff --git a/python/cpp/module.h b/python/cpp/module.h
index 9c9a9a2ff..207e0abf6 100644
--- a/python/cpp/module.h
+++ b/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);
diff --git a/python/cpp/profiling.cc b/python/cpp/profiling.cc
new file mode 100644
index 000000000..4b8dc4192
--- /dev/null
+++ b/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();
+	});
+	}
+
+  }
+}
diff --git a/python/ctranslate2/__init__.py b/python/ctranslate2/__init__.py
index 88da68aec..c783d4ec0 100644
--- a/python/ctranslate2/__init__.py
+++ b/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
diff --git a/python/ctranslate2/profiling.py b/python/ctranslate2/profiling.py
new file mode 100644
index 000000000..5f1eeef41
--- /dev/null
+++ b/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)

From 7d9537384b0ed8c06c46b162474570509562742f Mon Sep 17 00:00:00 2001
From: Orel Yehuda <orel.yehuda@intel.com>
Date: Mon, 8 Jul 2024 09:48:05 -0700
Subject: [PATCH 2/2] flake8 & black & isort reformatting

---
 python/ctranslate2/__init__.py  | 2 +-
 python/ctranslate2/profiling.py | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/python/ctranslate2/__init__.py b/python/ctranslate2/__init__.py
index c783d4ec0..46ada00d5 100644
--- a/python/ctranslate2/__init__.py
+++ b/python/ctranslate2/__init__.py
@@ -42,7 +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
+    from ctranslate2.profiling import dump_profiler, init_profiler
 
     register_extensions()
     del register_extensions
diff --git a/python/ctranslate2/profiling.py b/python/ctranslate2/profiling.py
index 5f1eeef41..6663a0ea5 100644
--- a/python/ctranslate2/profiling.py
+++ b/python/ctranslate2/profiling.py
@@ -1,7 +1,9 @@
-from ctranslate2 import _ext, Device
 import sys
 
-def init_profiler(device = Device.cpu, num_threads = 1):
+from ctranslate2 import Device, _ext
+
+
+def init_profiler(device=Device.cpu, num_threads=1):
     _ext.init_profiling(device, num_threads)