Reapply "distributed debug handlers (#126601)" (#127805)

d4l3k · pytorchmergebot · commit 597922ba21b2 · 2024-06-04T19:44:30.000Z
This reverts commit 7646825. Pull Request resolved: #127805 Approved by: https://github.com/PaliC
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -769,6 +769,7 @@ cc_library(
         ":caffe2",
         ":torch_headers",
         "@kineto",
+        "@cpp-httplib",
     ] + if_cuda([
         "@cuda//:nvToolsExt",
         "@cutlass",
diff --git a/WORKSPACE b/WORKSPACE
@@ -168,6 +168,12 @@ new_local_repository(
     path = "third_party/opentelemetry-cpp",
 )
 
+new_local_repository(
+    name = "cpp-httplib",
+    build_file = "//third_party:cpp-httplib.BUILD",
+    path = "third_party/cpp-httplib",
+)
+
 new_local_repository(
     name = "tensorpipe",
     build_file = "//third_party:tensorpipe.BUILD",
diff --git a/build_variables.bzl b/build_variables.bzl
@@ -515,6 +515,8 @@ libtorch_distributed_base_sources = [
     "torch/csrc/distributed/c10d/sequence_num.cpp",
     "torch/csrc/distributed/c10d/socket.cpp",
     "torch/csrc/distributed/c10d/Work.cpp",
+    "torch/csrc/distributed/c10d/control_plane/Handlers.cpp",
+    "torch/csrc/distributed/c10d/control_plane/WorkerServer.cpp",
 ]
 
 # These files are only supported on Linux (and others) but not on Windows.
diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
@@ -1179,6 +1179,9 @@ if(USE_KINETO)
     ${TORCH_ROOT}/third_party/kineto/libkineto/src)
 endif()
 
+target_include_directories(torch_cpu PRIVATE
+  ${TORCH_ROOT}/third_party/cpp-httplib)
+
 install(DIRECTORY "${TORCH_SRC_DIR}/csrc"
   DESTINATION ${TORCH_INSTALL_INCLUDE_DIR}/torch
   FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp")
diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
@@ -1681,3 +1681,7 @@ endif()
 
 # Include google/FlatBuffers
 include(${CMAKE_CURRENT_LIST_DIR}/FlatBuffers.cmake)
+
+# Include cpp-httplib
+add_library(httplib INTERFACE IMPORTED)
+target_include_directories(httplib SYSTEM INTERFACE ${PROJECT_SOURCE_DIR}/third_party/cpp-httplib)
diff --git a/docs/source/distributed.elastic.rst b/docs/source/distributed.elastic.rst
@@ -29,6 +29,7 @@ Documentation
    elastic/metrics
    elastic/events
    elastic/subprocess_handler
+   elastic/control_plane
 
 .. toctree::
    :maxdepth: 1
diff --git a/docs/source/elastic/control_plane.rst b/docs/source/elastic/control_plane.rst
@@ -0,0 +1,10 @@
+Control Plane
+=============
+
+.. automodule:: torch.distributed.elastic.control_plane
+.. currentmodule:: torch.distributed.elastic.control_plane
+
+This module contains optional helpers that add extra debug and control handlers
+into your application.
+
+.. autofunction:: torch.distributed.elastic.control_plane.worker_main
diff --git a/test/distributed/elastic/test_control_plane.py b/test/distributed/elastic/test_control_plane.py
@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+# Owner(s): ["oncall: distributed"]
+
+import json
+import os
+import pickle
+import socket
+import tempfile
+from contextlib import contextmanager
+
+from urllib3.connection import HTTPConnection
+from urllib3.connectionpool import HTTPConnectionPool
+
+from torch.distributed.elastic.control_plane import (
+    TORCH_WORKER_SERVER_SOCKET,
+    worker_main,
+)
+from torch.testing._internal.common_utils import requires_cuda, run_tests, TestCase
+
+
+class UnixHTTPConnection(HTTPConnection):
+    def __init__(self, socket_path: str) -> None:
+        super().__init__("localhost")
+
+        self.socket_path = socket_path
+
+    def connect(self) -> None:
+        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        self.sock.connect(self.socket_path)
+
+
+class UnixHTTPConnectionPool(HTTPConnectionPool):
+    def __init__(self, socket_path: str) -> None:
+        super().__init__("localhost")
+
+        self.socket_path = socket_path
+
+    def _new_conn(self):
+        return UnixHTTPConnection(self.socket_path)
+
+
+@contextmanager
+def local_worker_server() -> None:
+    with tempfile.TemporaryDirectory() as tmpdir:
+        socket_path = os.path.join(tmpdir, "socket.sock")
+        os.environ[TORCH_WORKER_SERVER_SOCKET] = socket_path
+
+        with worker_main():
+            pool = UnixHTTPConnectionPool(socket_path)
+            yield pool
+
+
+class WorkerServerTest(TestCase):
+    def test_worker_server(self) -> None:
+        with local_worker_server() as pool:
+            resp = pool.request("GET", "/")
+            self.assertEqual(resp.status, 200)
+            self.assertEqual(
+                resp.data,
+                b"""<h1>torch.distributed.WorkerServer</h1>
+<a href="/handler/">Handler names</a>
+""",
+            )
+
+            resp = pool.request("POST", "/handler/ping")
+            self.assertEqual(resp.status, 200)
+            self.assertEqual(resp.data, b"pong")
+
+            resp = pool.request("GET", "/handler/")
+            self.assertEqual(resp.status, 200)
+            self.assertIn("ping", json.loads(resp.data))
+
+            resp = pool.request("POST", "/handler/nonexistant")
+            self.assertEqual(resp.status, 404)
+            self.assertIn(b"Handler nonexistant not found:", resp.data)
+
+    @requires_cuda
+    def test_dump_nccl_trace_pickle(self) -> None:
+        with local_worker_server() as pool:
+            resp = pool.request("POST", "/handler/dump_nccl_trace_pickle")
+            self.assertEqual(resp.status, 200)
+            out = pickle.loads(resp.data)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/third_party/cpp-httplib.BUILD b/third_party/cpp-httplib.BUILD
@@ -0,0 +1,10 @@
+load("@rules_cc//cc:defs.bzl", "cc_library")
+
+cc_library(
+    name = "cpp-httplib",
+    hdrs = ["httplib.h"],
+    includes = [
+        "/",
+    ],
+    visibility = ["//visibility:public"],
+)
diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt
@@ -68,6 +68,7 @@ set(TORCH_PYTHON_INCLUDE_DIRECTORIES
     ${TORCH_ROOT}/third_party/onnx
     ${TORCH_ROOT}/third_party/flatbuffers/include
     ${TORCH_ROOT}/third_party/kineto/libkineto/include
+    ${TORCH_ROOT}/third_party/cpp-httplib
 
     ${TORCH_SRC_DIR}/csrc
     ${TORCH_SRC_DIR}/csrc/api/include
@@ -80,6 +81,7 @@ set(TORCH_PYTHON_LINK_LIBRARIES
     Python::Module
     pybind::pybind11
     opentelemetry::api
+    httplib
     shm
     fmt::fmt-header-only
     ATEN_CPU_FILES_GEN_LIB)
diff --git a/torch/_C/_distributed_c10d.pyi b/torch/_C/_distributed_c10d.pyi
@@ -94,6 +94,10 @@ class Logger:
     def _set_uneven_input_join(self) -> None: ...
     def _set_static_graph(self) -> None: ...
 
+class _WorkerServer:
+    def __init__(self, socket_path: str) -> None: ...
+    def shutdown(self) -> None: ...
+
 def get_debug_level(): ...
 def set_debug_level(): ...
 def set_debug_level_from_env(): ...
diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
@@ -28,6 +28,7 @@
 #include <torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp>
 #include <torch/csrc/distributed/c10d/TraceUtils.h>
 #include <torch/csrc/distributed/c10d/Utils.hpp>
+#include <torch/csrc/distributed/c10d/control_plane/Handlers.hpp>
 #include <torch/csrc/distributed/c10d/logger.hpp>
 #include <torch/torch.h>
 
@@ -369,6 +370,13 @@ std::string dump_nccl_trace() {
 }
 #endif
 
+// TODO(c-p-i-o): add a JSON endpoint.
+control_plane::RegisterHandler dumpHandler{
+    "dump_nccl_trace_pickle",
+    [](const control_plane::Request&, control_plane::Response& res) {
+      res.setContent(dump_nccl_trace(), "application/octet-stream");
+    }};
+
 std::optional<std::function<void(std::function<void(const std::string&)>)>>&
 get_cpp_trace_dumper() {
   static std::optional<
diff --git a/torch/csrc/distributed/c10d/control_plane/Handlers.cpp b/torch/csrc/distributed/c10d/control_plane/Handlers.cpp
@@ -0,0 +1,75 @@
+#include <torch/csrc/distributed/c10d/control_plane/Handlers.hpp>
+
+#include <fmt/format.h>
+#include <mutex>
+#include <shared_mutex>
+#include <stdexcept>
+
+namespace c10d {
+namespace control_plane {
+
+namespace {
+
+class HandlerRegistry {
+ public:
+  void registerHandler(const std::string& name, HandlerFunc f) {
+    std::unique_lock<std::shared_mutex> lock(handlersMutex_);
+
+    if (handlers_.find(name) != handlers_.end()) {
+      throw std::runtime_error(
+          fmt::format("Handler {} already registered", name));
+    }
+
+    handlers_[name] = f;
+  }
+
+  HandlerFunc getHandler(const std::string& name) {
+    std::shared_lock<std::shared_mutex> lock(handlersMutex_);
+
+    auto it = handlers_.find(name);
+    if (it == handlers_.end()) {
+      throw std::runtime_error(fmt::format("Failed to find handler {}", name));
+    }
+    return handlers_[name];
+  }
+
+  std::vector<std::string> getHandlerNames() {
+    std::shared_lock<std::shared_mutex> lock(handlersMutex_);
+
+    std::vector<std::string> names;
+    for (const auto& [name, _] : handlers_) {
+      names.push_back(name);
+    }
+    return names;
+  }
+
+ private:
+  std::shared_mutex handlersMutex_{};
+  std::unordered_map<std::string, HandlerFunc> handlers_{};
+};
+
+HandlerRegistry& getHandlerRegistry() {
+  static HandlerRegistry registry;
+  return registry;
+}
+
+RegisterHandler pingHandler{"ping", [](const Request&, Response& res) {
+                              res.setContent("pong", "text/plain");
+                            }};
+
+} // namespace
+
+void registerHandler(const std::string& name, HandlerFunc f) {
+  return getHandlerRegistry().registerHandler(name, f);
+}
+
+HandlerFunc getHandler(const std::string& name) {
+  return getHandlerRegistry().getHandler(name);
+}
+
+std::vector<std::string> getHandlerNames() {
+  return getHandlerRegistry().getHandlerNames();
+}
+
+} // namespace control_plane
+} // namespace c10d
diff --git a/torch/csrc/distributed/c10d/control_plane/Handlers.hpp b/torch/csrc/distributed/c10d/control_plane/Handlers.hpp
@@ -0,0 +1,67 @@
+#pragma once
+
+#include <functional>
+#include <string>
+
+#include <c10/macros/Export.h>
+
+namespace c10d {
+namespace control_plane {
+
+// Request represents a request to the handler. This conceptually maps to an
+// HTTP request but could be called via other transports.
+class TORCH_API Request {
+ public:
+  virtual ~Request() = default;
+
+  virtual const std::string& body() = 0;
+};
+
+// Response represents a response to the handler. This conceptually maps to an
+// HTTP response but could be called via other transports.
+class TORCH_API Response {
+ public:
+  virtual ~Response() = default;
+
+  // Set the response body to the provided string.
+  // TODO: add support for chunked responses
+  virtual void setContent(
+      std::string&& content,
+      const std::string& content_type) = 0;
+
+  // Set the response status code.
+  // These should match standard HTTP status codes.
+  virtual void setStatus(int status) = 0;
+};
+
+using HandlerFunc = std::function<void(const Request&, Response&)>;
+
+// Registers a handler. The name needs to be unique and can be called by using
+// getHandler directly or via WorkerServer for remote requests.
+// These handlers are called from a background C++ thread concurrently with the
+// main thread. These handlers need to be thread safe and not cause issues
+// during Python training.
+TORCH_API void registerHandler(const std::string& name, HandlerFunc f);
+
+// Fetches a handler by name.
+TORCH_API HandlerFunc getHandler(const std::string& name);
+
+TORCH_API std::vector<std::string> getHandlerNames();
+
+// Registers a handler statically.
+// See registerHandler for more details.
+class TORCH_API RegisterHandler {
+ public:
+  RegisterHandler(const std::string& name, HandlerFunc f) {
+    registerHandler(name, f);
+  }
+
+  // disable move, copy
+  RegisterHandler(const RegisterHandler&) = delete;
+  RegisterHandler(RegisterHandler&&) = delete;
+  RegisterHandler& operator=(const RegisterHandler&) = delete;
+  RegisterHandler& operator=(RegisterHandler&&) = delete;
+};
+
+} // namespace control_plane
+} // namespace c10d
diff --git a/torch/csrc/distributed/c10d/control_plane/WorkerServer.cpp b/torch/csrc/distributed/c10d/control_plane/WorkerServer.cpp
diff --git a/torch/csrc/distributed/c10d/control_plane/WorkerServer.hpp b/torch/csrc/distributed/c10d/control_plane/WorkerServer.hpp
diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp
diff --git a/torch/distributed/elastic/control_plane.py b/torch/distributed/elastic/control_plane.py
