1 #pragma once 2 3 #include <functional> 4 #include <map> 5 #include <string> 6 #include <utility> 7 8 #include <c10/macros/Export.h> 9 10 namespace c10d::control_plane { 11 12 // Request represents a request to the handler. This conceptually maps to an 13 // HTTP request but could be called via other transports. 14 class TORCH_API Request { 15 public: 16 virtual ~Request() = default; 17 18 virtual const std::string& body() const = 0; 19 20 virtual const std::multimap<std::string, std::string>& params() const = 0; 21 }; 22 23 // Response represents a response to the handler. This conceptually maps to an 24 // HTTP response but could be called via other transports. 25 class TORCH_API Response { 26 public: 27 virtual ~Response() = default; 28 29 // Set the response body to the provided string. 30 // TODO: add support for chunked responses 31 virtual void setContent( 32 std::string&& content, 33 const std::string& content_type) = 0; 34 35 // Set the response status code. 36 // These should match standard HTTP status codes. 37 virtual void setStatus(int status) = 0; 38 }; 39 40 using HandlerFunc = std::function<void(const Request&, Response&)>; 41 42 // Registers a handler. The name needs to be unique and can be called by using 43 // getHandler directly or via WorkerServer for remote requests. 44 // These handlers are called from a background C++ thread concurrently with the 45 // main thread. These handlers need to be thread safe and not cause issues 46 // during Python training. 47 TORCH_API void registerHandler(const std::string& name, HandlerFunc f); 48 49 // Fetches a handler by name. 50 TORCH_API HandlerFunc getHandler(const std::string& name); 51 52 TORCH_API std::vector<std::string> getHandlerNames(); 53 54 // Registers a handler statically. 55 // See registerHandler for more details. 56 class TORCH_API RegisterHandler { 57 public: RegisterHandler(const std::string & name,HandlerFunc f)58 RegisterHandler(const std::string& name, HandlerFunc f) { 59 registerHandler(name, std::move(f)); 60 } 61 62 // disable move, copy 63 RegisterHandler(const RegisterHandler&) = delete; 64 RegisterHandler(RegisterHandler&&) = delete; 65 RegisterHandler& operator=(const RegisterHandler&) = delete; 66 RegisterHandler& operator=(RegisterHandler&&) = delete; 67 }; 68 69 } // namespace c10d::control_plane 70