1 #pragma once 2 3 #include <thread> 4 5 #include <torch/csrc/distributed/c10d/TCPStore.hpp> 6 #include <torch/csrc/distributed/c10d/socket.h> 7 8 #ifdef _WIN32 9 #include <io.h> 10 #include <winsock2.h> 11 #else 12 #include <poll.h> 13 #include <unistd.h> 14 #endif 15 16 namespace c10d::detail { 17 18 // Magic number for client validation. 19 static const uint32_t validationMagicNumber = 0x3C85F7CE; 20 21 enum class QueryType : uint8_t { 22 VALIDATE, 23 SET, 24 COMPARE_SET, 25 GET, 26 ADD, 27 CHECK, 28 WAIT, 29 GETNUMKEYS, 30 DELETE_KEY, 31 APPEND, 32 MULTI_GET, 33 MULTI_SET, 34 CANCEL_WAIT, 35 PING, 36 }; 37 38 enum class CheckResponseType : uint8_t { READY, NOT_READY }; 39 40 enum class WaitResponseType : uint8_t { STOP_WAITING, WAIT_CANCELED }; 41 42 // Abstract base class to handle thread state for TCPStoreMasterDaemon. 43 // Contains the windows/unix implementations to signal a 44 // shutdown sequence for the thread 45 class BackgroundThread { 46 public: 47 explicit BackgroundThread(); 48 49 virtual ~BackgroundThread() = 0; 50 virtual std::uint16_t port() const = 0; 51 52 void start(); 53 bool stop_requested(); 54 55 protected: 56 void dispose(); 57 virtual void run() = 0; 58 virtual void stop() = 0; is_running()59 bool is_running() { 60 return is_running_.load(); 61 } 62 63 private: 64 std::atomic<bool> is_running_{false}; 65 std::thread daemonThread_{}; 66 }; 67 68 std::unique_ptr<BackgroundThread> create_tcpstore_backend( 69 const TCPStoreOptions& opts); 70 std::unique_ptr<BackgroundThread> create_libuv_tcpstore_backend( 71 const TCPStoreOptions& opts); 72 bool is_libuv_tcpstore_backend_available(); 73 74 } // namespace c10d::detail 75