Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement python bindings for profiler #1739

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 dump_profiler, init_profiler

register_extensions()
del register_extensions
Expand Down
12 changes: 12 additions & 0 deletions python/ctranslate2/profiling.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
import sys

from ctranslate2 import Device, _ext


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)
Loading