mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
fe284b0d97
This moves BackgroundThread to TCPStoreBackend.hpp. This will eventually be the interface shared between the current TCPStore backend and the new libuv one. Pull Request resolved: https://github.com/pytorch/pytorch/pull/105163 Approved by: https://github.com/fduwjj, https://github.com/H-Huang
55 lines
1.1 KiB
C++
55 lines
1.1 KiB
C++
#pragma once
|
|
|
|
#include <chrono>
|
|
#include <thread>
|
|
#include <vector>
|
|
|
|
#include <torch/csrc/distributed/c10d/socket.h>
|
|
|
|
#ifdef _WIN32
|
|
#include <io.h>
|
|
#include <winsock2.h>
|
|
#else
|
|
#include <poll.h>
|
|
#include <unistd.h>
|
|
#endif
|
|
|
|
namespace c10d {
|
|
namespace detail {
|
|
|
|
// Abstract base class to handle thread state for TCPStoreMasterDaemon.
|
|
// Contains the windows/unix implementations to signal a
|
|
// shutdown sequence for the thread
|
|
class BackgroundThread {
|
|
public:
|
|
explicit BackgroundThread(Socket&& storeListenSocket);
|
|
|
|
virtual ~BackgroundThread() = 0;
|
|
|
|
protected:
|
|
void dispose();
|
|
|
|
Socket storeListenSocket_;
|
|
std::thread daemonThread_{};
|
|
std::vector<Socket> sockets_{};
|
|
#ifdef _WIN32
|
|
const std::chrono::milliseconds checkTimeout_ = std::chrono::milliseconds{10};
|
|
HANDLE ghStopEvent_{};
|
|
#else
|
|
std::array<int, 2> controlPipeFd_{{-1, -1}};
|
|
#endif
|
|
|
|
private:
|
|
// Initialization for shutdown signal
|
|
void initStopSignal();
|
|
// Triggers the shutdown signal
|
|
void stop();
|
|
// Joins the thread
|
|
void join();
|
|
// Clean up the shutdown signal
|
|
void closeStopSignal();
|
|
};
|
|
|
|
} // namespace detail
|
|
} // namespace c10d
|