[PROTON-DEV] Improve profile interface (#5793)

Author: Jokeren

third_party/proton/csrc/Proton.cpp

@@ -13,16 +13,20 @@ void initProton(pybind11::module &&m) {
   using ret = pybind11::return_value_policy;
   using namespace pybind11::literals;
 
-  m.def("start",
-        [](const std::string &path, const std::string &contextSourceName,
-           const std::string &dataName, const std::string &profilerName,
-           const std::string &profilerPath, const std::string &mode) {
-          auto sessionId = SessionManager::instance().addSession(
-              path, profilerName, profilerPath, contextSourceName, dataName,
-              mode);
-          SessionManager::instance().activateSession(sessionId);
-          return sessionId;
-        });
+  m.def(
+      "start",
+      [](const std::string &path, const std::string &contextSourceName,
+         const std::string &dataName, const std::string &profilerName,
+         const std::string &mode, const std::string &profilerPath) {
+        auto sessionId = SessionManager::instance().addSession(
+            path, profilerName, profilerPath, contextSourceName, dataName,
+            mode);
+        SessionManager::instance().activateSession(sessionId);
+        return sessionId;
+      },
+      pybind11::arg("path"), pybind11::arg("contextSourceName"),
+      pybind11::arg("dataName"), pybind11::arg("profilerName"),
+      pybind11::arg("mode") = "", pybind11::arg("profilerPath") = "");
 
   m.def("activate", [](size_t sessionId) {
     SessionManager::instance().activateSession(sessionId);

third_party/proton/proton/profile.py

@@ -46,9 +46,9 @@ def _check_env(backend: str) -> None:
 def _check_mode(backend: str, mode: Optional[str]) -> None:
     # TODO(Keren): Need a better mode registration mechanism
     backend_modes = {
-        "cupti": [None, "pcsampling"],
-        "roctracer": [None],
-        "instrumentation": [None],
+        "cupti": ["", "pcsampling"],
+        "roctracer": [""],
+        "instrumentation": [""],
     }
 
     if mode not in backend_modes[backend]:
@@ -102,28 +102,21 @@ def start(
         # Ignore the start() call if the script is run from the command line.
         return
 
-    if name is None:
-        name = DEFAULT_PROFILE_NAME
-
-    if backend is None:
-        backend = _select_backend()
+    name = DEFAULT_PROFILE_NAME if name is None else name
+    backend = _select_backend() if backend is None else backend
+    mode = "" if mode is None else mode
 
     _check_env(backend)
     _check_mode(backend, mode)
 
     backend_path = _get_backend_default_path(backend)
 
-    _check_mode(backend, mode)
-
-    if mode is None:
-        mode = ""
-
-    backend_path = _get_backend_default_path(backend)
-
     set_profiling_on()
-    if hook and hook == "triton":
+
+    if hook == "triton":
         register_triton_hook()
-    return libproton.start(name, context, data, backend, backend_path, mode)
+
+    return libproton.start(name, context, data, backend, mode, backend_path)
 
 
 def activate(session: Optional[int] = None) -> None:

third_party/proton/test/test_lib.py

@@ -32,7 +32,7 @@ def test_op():
 
 def test_session(tmp_path: pathlib.Path):
     temp_file = tmp_path / "test_session.hatchet"
-    session_id = libproton.start(str(temp_file.with_suffix("")), "shadow", "tree", _select_backend(), "", "")
+    session_id = libproton.start(str(temp_file.with_suffix("")), "shadow", "tree", _select_backend())
     libproton.deactivate(session_id)
     libproton.activate(session_id)
     libproton.finalize(session_id, "hatchet")
@@ -42,7 +42,7 @@ def test_session(tmp_path: pathlib.Path):
 
 def test_add_metrics(tmp_path: pathlib.Path):
     temp_file = tmp_path / "test_add_metrics.hatchet"
-    libproton.start(str(temp_file.with_suffix("")), "shadow", "tree", _select_backend(), "", "")
+    libproton.start(str(temp_file.with_suffix("")), "shadow", "tree", _select_backend())
     id1 = libproton.record_scope()
     libproton.enter_scope(id1, "one")
     libproton.add_metrics(id1, {"a": 1.0, "b": 2.0})