From bc9d72c37a3e2f5e004ea22d61de258307e6e7bf Mon Sep 17 00:00:00 2001 From: hebinbin1 Date: Tue, 28 Nov 2023 16:11:01 +0800 Subject: [PATCH] feat: add debug mode support for python backend --- src/pb_stub.cc | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/pb_stub.cc b/src/pb_stub.cc index 3d473101..5da805ac 100644 --- a/src/pb_stub.cc +++ b/src/pb_stub.cc @@ -447,6 +447,21 @@ Stub::StubSetup() deserialize_bytes_ = python_backend_utils.attr("deserialize_bytes_tensor"); serialize_bytes_ = python_backend_utils.attr("serialize_byte_tensor"); + // add a debug mode for triton python backend + py::module os = py::module_::import("os"); + py::object env = os.attr("environ"); + std::string triton_debug = py::str(env.attr("get")("TRITON_DEBUG", "0")); + if (triton_debug == "1") { + py::module debugpy = py::module_::import("debugpy"); + std::string debug_port = py::str(env.attr("get")("TRITON_DEBUG_PORT", "8003")); + int triton_debug_port = std::stoi(debug_port); + LOG_INFO << "Running with debugpy mode on 0.0.0.0:" << triton_debug_port; + py::tuple listen_args(2); + listen_args[0] = "0.0.0.0"; + listen_args[1] = triton_debug_port; + debugpy.attr("listen")(listen_args); + } + return sys; }