// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "bindings.h"
#include "tenzr/runtime/pjrt_profiler.h"
namespace tenzr::python {
void BindProfiler(nb::module_& m) {
auto prof_m = m.def_submodule("profiler", "Hardware stream telemetry, DMA tracking, and xprof profiling");
prof_m.def("start", &profiler::Start);
prof_m.def("stop", &profiler::Stop);
prof_m.def("reset", &profiler::Reset);
prof_m.def("is_enabled", &profiler::IsEnabled);
prof_m.def("export_chrome_trace", &profiler::ExportChromeTrace, "path"_a);
prof_m.def(
"record_stream_kernel",
[](std::string name, int64_t stream_id, int64_t dur_us, std::string cat) {
profiler::RecordStreamKernel(std::move(name), stream_id, dur_us, std::move(cat));
},
"name"_a, "stream_id"_a, "dur_us"_a, "cat"_a = "gpu_kernel");
prof_m.def("record_memory_transfer", &profiler::RecordMemoryTransfer, "name"_a, "bytes"_a, "h2d"_a, "dur_us"_a);
prof_m.def("record_collective_comm", &profiler::RecordCollectiveComm, "name"_a, "bytes"_a, "dur_us"_a);
// BottleneckReport
nb::class_<profiler::BottleneckReport>(prof_m, "BottleneckReport")
.def_ro("primary_bottleneck", &profiler::BottleneckReport::primary_bottleneck)
.def_ro("compute_to_comm_ratio", &profiler::BottleneckReport::compute_to_comm_ratio)
.def_ro("memory_transfer_pct", &profiler::BottleneckReport::memory_transfer_pct)
.def_ro("collective_comm_pct", &profiler::BottleneckReport::collective_comm_pct)
.def_ro("recommendations", &profiler::BottleneckReport::recommendations)
.def("format", &profiler::BottleneckReport::Format);
// ProfileSummary
nb::class_<profiler::ProfileSummary>(prof_m, "ProfileSummary")
.def_ro("total_duration_us", &profiler::ProfileSummary::total_duration_us)
.def_ro("total_events", &profiler::ProfileSummary::total_events)
.def("format_report", &profiler::ProfileSummary::FormatReport)
.def("identify_bottlenecks", &profiler::ProfileSummary::IdentifyBottlenecks)
.def(
"check_latency_regression",
[](const profiler::ProfileSummary& s, int64_t baseline_us, double threshold_pct) {
std::string msg;
bool regressed = s.CheckLatencyRegression(baseline_us, threshold_pct, &msg);
return nb::make_tuple(regressed, msg);
},
"baseline_us"_a, "threshold_pct"_a = 10.0);
prof_m.def("analyze_trace", &profiler::AnalyzeTrace);
// TraceSession context manager
nb::class_<profiler::TraceSession>(prof_m, "TraceSession")
.def(nb::init<std::string>(), "path"_a = "tenzr_profile.trace.json")
.def("__enter__", [](profiler::TraceSession& s) -> profiler::TraceSession& { return s; })
.def("__exit__", [](profiler::TraceSession& s, nb::handle, nb::handle, nb::handle) { s.Stop(); });
prof_m.def("trace", [](std::string path) { return profiler::TraceSession(std::move(path)); }, "path"_a = "tenzr_profile.trace.json");
// PjrtProfiler: drives the PJRT plugin profiler extension and captures a serialized XPlane profile.
nb::class_<PjrtProfiler>(prof_m, "PjrtProfiler")
.def(nb::init<>())
.def("supported", &PjrtProfiler::Supported)
.def("start", &PjrtProfiler::Start)
.def("stop", &PjrtProfiler::Stop)
.def("collect_data",
[](const PjrtProfiler& p) -> nb::bytes {
const std::vector<uint8_t> data = p.CollectData();
if (data.empty()) {
return nb::bytes("", 0);
}
return nb::bytes(reinterpret_cast<const char*>(data.data()), data.size());
})
.def("dump", &PjrtProfiler::Dump, "path"_a);
}
} // namespace tenzr::python