Skip to content

Commit

Permalink
Expose profiler to python
Browse files Browse the repository at this point in the history
  • Loading branch information
yehudaorel committed Jul 4, 2024
1 parent 39f48f2 commit 0cbb079
Show file tree
Hide file tree
Showing 5 changed files with 33 additions and 0 deletions.
1 change: 1 addition & 0 deletions python/cpp/module.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
2 changes: 2 additions & 0 deletions python/cpp/module.h
Original file line number Diff line number Diff line change
@@ -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);
Expand Down
19 changes: 19 additions & 0 deletions python/cpp/profiling.cc
Original file line number Diff line number Diff line change
@@ -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();
});
}

}
}
1 change: 1 addition & 0 deletions python/ctranslate2/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 10 additions & 0 deletions python/ctranslate2/profiling.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 0cbb079

Please sign in to comment.