mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
[Reland] fix c10:TempFile APIs on Windows (#108508)
PR #106656 was reverted due to IOS failures. It seems that IOS builds don't have full support of std::filesystem. This PR discards std::filesystem changes and add temp file creation on Windows. It also moves the platform syscalls into a separate cpp file. Pull Request resolved: https://github.com/pytorch/pytorch/pull/108508 Approved by: https://github.com/ezyang
This commit is contained in:
cyy
PyTorch MergeBot
@ -1,22 +1,39 @@
|
||||
#include <c10/util/tempfile.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <sys/stat.h>
|
||||
#include <sys/types.h>
|
||||
#include <optional>
|
||||
|
||||
#if !defined(_WIN32)
|
||||
TEST(TempFileTest, MatchesExpectedPattern) {
|
||||
c10::TempFile pattern = c10::make_tempfile("test-pattern-");
|
||||
ASSERT_NE(pattern.name.find("test-pattern-"), std::string::npos);
|
||||
static bool file_exists(const char* path) {
|
||||
struct stat st;
|
||||
return stat(path, &st) == 0 && S_ISREG(st.st_mode);
|
||||
}
|
||||
static bool directory_exists(const char* path) {
|
||||
struct stat st;
|
||||
return stat(path, &st) == 0 && S_ISDIR(st.st_mode);
|
||||
}
|
||||
#else
|
||||
static bool file_exists(const char* path) {
|
||||
struct _stat st;
|
||||
return _stat(path, &st) == 0 && ((st.st_mode & _S_IFMT) == _S_IFREG);
|
||||
}
|
||||
static bool directory_exists(const char* path) {
|
||||
struct _stat st;
|
||||
return _stat(path, &st) == 0 && ((st.st_mode & _S_IFMT) == _S_IFDIR);
|
||||
}
|
||||
#endif // !defined(_WIN32)
|
||||
|
||||
static bool directory_exists(const char* path) {
|
||||
struct stat st;
|
||||
return (stat(path, &st) == 0 && (st.st_mode & S_IFDIR));
|
||||
TEST(TempFileTest, MatchesExpectedPattern) {
|
||||
c10::TempFile file = c10::make_tempfile("test-pattern-");
|
||||
|
||||
ASSERT_TRUE(file_exists(file.name.c_str()));
|
||||
#if !defined(_WIN32)
|
||||
ASSERT_NE(file.name.find("test-pattern-"), std::string::npos);
|
||||
#endif // !defined(_WIN32)
|
||||
}
|
||||
|
||||
TEST(TempDirTest, tryMakeTempdir) {
|
||||
c10::optional<c10::TempDir> tempdir = c10::make_tempdir("test-dir-");
|
||||
std::optional<c10::TempDir> tempdir = c10::make_tempdir("test-dir-");
|
||||
std::string tempdir_name = tempdir->name;
|
||||
|
||||
// directory should exist while tempdir is alive
|
||||
|
||||
162
c10/util/tempfile.cpp
Normal file
162
c10/util/tempfile.cpp
Normal file
@ -0,0 +1,162 @@
|
||||
#include <c10/util/Exception.h>
|
||||
#include <c10/util/tempfile.h>
|
||||
#include <fmt/format.h>
|
||||
|
||||
#if !defined(_WIN32)
|
||||
#include <unistd.h>
|
||||
#include <cerrno>
|
||||
#else // defined(_WIN32)
|
||||
#include <Windows.h>
|
||||
#include <fcntl.h>
|
||||
#include <fileapi.h>
|
||||
#include <io.h>
|
||||
#endif // defined(_WIN32)
|
||||
|
||||
// Creates the filename pattern passed to and completed by `mkstemp`.
|
||||
#if !defined(_WIN32)
|
||||
static std::string make_filename(std::string_view name_prefix) {
|
||||
// The filename argument to `mkstemp` needs "XXXXXX" at the end according to
|
||||
// http://pubs.opengroup.org/onlinepubs/009695399/functions/mkstemp.html
|
||||
constexpr const char* kRandomPattern = "XXXXXX";
|
||||
|
||||
// We see if any of these environment variables is set and use their value, or
|
||||
// else default the temporary directory to `/tmp`.
|
||||
|
||||
const char* tmp_directory = "/tmp";
|
||||
for (const char* variable : {"TMPDIR", "TMP", "TEMP", "TEMPDIR"}) {
|
||||
if (const char* path = getenv(variable)) {
|
||||
tmp_directory = path;
|
||||
break;
|
||||
}
|
||||
}
|
||||
return fmt::format("{}/{}{}", tmp_directory, name_prefix, kRandomPattern);
|
||||
}
|
||||
#else
|
||||
static std::string make_filename() {
|
||||
char name[L_tmpnam_s]{};
|
||||
auto res = tmpnam_s(name, L_tmpnam_s);
|
||||
if (res != 0) {
|
||||
TORCH_WARN("Error generating temporary file");
|
||||
return "";
|
||||
}
|
||||
return name;
|
||||
}
|
||||
#endif // !defined(_WIN32)
|
||||
|
||||
namespace c10 {
|
||||
/// Attempts to return a temporary file or returns `nullopt` if an error
|
||||
/// occurred.
|
||||
std::optional<TempFile> try_make_tempfile(std::string_view name_prefix) {
|
||||
#if defined(_WIN32)
|
||||
auto filename = make_filename();
|
||||
#else
|
||||
auto filename = make_filename(name_prefix);
|
||||
#endif
|
||||
if (filename.empty()) {
|
||||
return std::nullopt;
|
||||
}
|
||||
#if defined(_WIN32)
|
||||
return TempFile(std::move(filename));
|
||||
#else
|
||||
const int fd = mkstemp(filename.data());
|
||||
if (fd == -1) {
|
||||
return std::nullopt;
|
||||
}
|
||||
return TempFile(std::move(filename), fd);
|
||||
#endif // defined(_WIN32)
|
||||
}
|
||||
|
||||
/// Like `try_make_tempfile`, but throws an exception if a temporary file could
|
||||
/// not be returned.
|
||||
TempFile make_tempfile(std::string_view name_prefix) {
|
||||
if (auto tempfile = try_make_tempfile(name_prefix)) {
|
||||
return std::move(*tempfile);
|
||||
}
|
||||
TORCH_CHECK(false, "Error generating temporary file: ", std::strerror(errno));
|
||||
}
|
||||
|
||||
/// Attempts to return a temporary directory or returns `nullopt` if an error
|
||||
/// occurred.
|
||||
std::optional<TempDir> try_make_tempdir(std::string_view name_prefix) {
|
||||
#if defined(_WIN32)
|
||||
for (int i = 0; i < 10; i++) {
|
||||
auto dirname = make_filename();
|
||||
if (dirname.empty()) {
|
||||
return std::nullopt;
|
||||
}
|
||||
if (CreateDirectoryA(dirname.c_str(), nullptr)) {
|
||||
return TempDir(dirname);
|
||||
}
|
||||
if (GetLastError() == ERROR_SUCCESS) {
|
||||
return std::nullopt;
|
||||
}
|
||||
}
|
||||
return std::nullopt;
|
||||
#else
|
||||
auto filename = make_filename(name_prefix);
|
||||
const char* dirname = mkdtemp(filename.data());
|
||||
if (!dirname) {
|
||||
return std::nullopt;
|
||||
}
|
||||
return TempDir(dirname);
|
||||
#endif // defined(_WIN32)
|
||||
}
|
||||
|
||||
#if defined(_WIN32)
|
||||
bool TempFile::open() {
|
||||
if (fd != -1) {
|
||||
return false;
|
||||
}
|
||||
auto err = _sopen_s(
|
||||
&fd,
|
||||
name.c_str(),
|
||||
_O_CREAT | _O_TEMPORARY | _O_EXCL | _O_BINARY | _O_RDWR,
|
||||
_SH_DENYNO,
|
||||
_S_IREAD | _S_IWRITE);
|
||||
if (err != 0) {
|
||||
fd = -1;
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
#endif
|
||||
|
||||
TempFile::~TempFile() {
|
||||
if (!name.empty()) {
|
||||
#if !defined(_WIN32)
|
||||
if (fd >= 0) {
|
||||
unlink(name.c_str());
|
||||
close(fd);
|
||||
}
|
||||
#else
|
||||
if (fd >= 0) {
|
||||
_close(fd);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
TempDir::~TempDir() {
|
||||
if (!name.empty()) {
|
||||
#if !defined(_WIN32)
|
||||
rmdir(name.c_str());
|
||||
#else // defined(_WIN32)
|
||||
RemoveDirectoryA(name.c_str());
|
||||
#endif // defined(_WIN32)
|
||||
}
|
||||
}
|
||||
|
||||
/// Like `try_make_tempdir`, but throws an exception if a temporary directory
|
||||
/// could not be returned.
|
||||
TempDir make_tempdir(std::string_view name_prefix) {
|
||||
if (auto tempdir = try_make_tempdir(name_prefix)) {
|
||||
return std::move(*tempdir);
|
||||
}
|
||||
#if !defined(_WIN32)
|
||||
TORCH_CHECK(
|
||||
false, "Error generating temporary directory: ", std::strerror(errno));
|
||||
#else // defined(_WIN32)
|
||||
TORCH_CHECK(false, "Error generating temporary directory");
|
||||
#endif // defined(_WIN32)
|
||||
}
|
||||
} // namespace c10
|
||||
@ -1,70 +1,17 @@
|
||||
#pragma once
|
||||
|
||||
#include <c10/util/Exception.h>
|
||||
#include <c10/util/Optional.h>
|
||||
|
||||
#include <cerrno>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <c10/macros/Export.h>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#if !defined(_WIN32)
|
||||
#include <unistd.h>
|
||||
#else // defined(_WIN32)
|
||||
#include <Windows.h>
|
||||
#include <fileapi.h>
|
||||
#endif // defined(_WIN32)
|
||||
|
||||
namespace c10 {
|
||||
namespace detail {
|
||||
// Creates the filename pattern passed to and completed by `mkstemp`.
|
||||
// Returns std::vector<char> because `mkstemp` needs a (non-const) `char*` and
|
||||
// `std::string` only provides `const char*` before C++17.
|
||||
#if !defined(_WIN32)
|
||||
inline std::vector<char> make_filename(std::string name_prefix) {
|
||||
// The filename argument to `mkstemp` needs "XXXXXX" at the end according to
|
||||
// http://pubs.opengroup.org/onlinepubs/009695399/functions/mkstemp.html
|
||||
static const std::string kRandomPattern = "XXXXXX";
|
||||
|
||||
// We see if any of these environment variables is set and use their value, or
|
||||
// else default the temporary directory to `/tmp`.
|
||||
static const char* env_variables[] = {"TMPDIR", "TMP", "TEMP", "TEMPDIR"};
|
||||
|
||||
std::string tmp_directory = "/tmp";
|
||||
for (const char* variable : env_variables) {
|
||||
if (const char* path = getenv(variable)) {
|
||||
tmp_directory = path;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<char> filename;
|
||||
filename.reserve(
|
||||
tmp_directory.size() + name_prefix.size() + kRandomPattern.size() + 2);
|
||||
|
||||
filename.insert(filename.end(), tmp_directory.begin(), tmp_directory.end());
|
||||
filename.push_back('/');
|
||||
filename.insert(filename.end(), name_prefix.begin(), name_prefix.end());
|
||||
filename.insert(filename.end(), kRandomPattern.begin(), kRandomPattern.end());
|
||||
filename.push_back('\0');
|
||||
|
||||
return filename;
|
||||
}
|
||||
#endif // !defined(_WIN32)
|
||||
} // namespace detail
|
||||
|
||||
struct TempFile {
|
||||
#if !defined(_WIN32)
|
||||
TempFile() : fd(-1) {}
|
||||
TempFile(std::string name, int fd) : fd(fd), name(std::move(name)) {}
|
||||
struct C10_API TempFile {
|
||||
TempFile(std::string_view name, int fd = -1) noexcept : fd(fd), name(name) {}
|
||||
TempFile(const TempFile&) = delete;
|
||||
TempFile(TempFile&& other) noexcept
|
||||
: fd(other.fd), name(std::move(other.name)) {
|
||||
other.fd = -1;
|
||||
other.name.clear();
|
||||
}
|
||||
|
||||
TempFile& operator=(const TempFile&) = delete;
|
||||
@ -72,26 +19,22 @@ struct TempFile {
|
||||
fd = other.fd;
|
||||
name = std::move(other.name);
|
||||
other.fd = -1;
|
||||
other.name.clear();
|
||||
return *this;
|
||||
}
|
||||
#if defined(_WIN32)
|
||||
bool open();
|
||||
#endif
|
||||
|
||||
~TempFile() {
|
||||
if (fd >= 0) {
|
||||
unlink(name.c_str());
|
||||
close(fd);
|
||||
}
|
||||
}
|
||||
~TempFile();
|
||||
|
||||
int fd;
|
||||
#endif // !defined(_WIN32)
|
||||
|
||||
std::string name;
|
||||
};
|
||||
|
||||
struct TempDir {
|
||||
TempDir() = default;
|
||||
explicit TempDir(std::string name) : name(std::move(name)) {}
|
||||
struct C10_API TempDir {
|
||||
TempDir() = delete;
|
||||
explicit TempDir(std::string_view name) noexcept : name(name) {}
|
||||
TempDir(const TempDir&) = delete;
|
||||
TempDir(TempDir&& other) noexcept : name(std::move(other.name)) {
|
||||
other.name.clear();
|
||||
@ -100,19 +43,10 @@ struct TempDir {
|
||||
TempDir& operator=(const TempDir&) = delete;
|
||||
TempDir& operator=(TempDir&& other) noexcept {
|
||||
name = std::move(other.name);
|
||||
other.name.clear();
|
||||
return *this;
|
||||
}
|
||||
|
||||
~TempDir() {
|
||||
if (!name.empty()) {
|
||||
#if !defined(_WIN32)
|
||||
rmdir(name.c_str());
|
||||
#else // defined(_WIN32)
|
||||
RemoveDirectoryA(name.c_str());
|
||||
#endif // defined(_WIN32)
|
||||
}
|
||||
}
|
||||
~TempDir();
|
||||
|
||||
std::string name;
|
||||
};
|
||||
@ -126,31 +60,14 @@ struct TempDir {
|
||||
/// `"TEMPDIR"` environment variable if any is set, or otherwise `/tmp`;
|
||||
/// `<name-prefix>` is the value supplied to this function, and
|
||||
/// `<random-pattern>` is a random sequence of numbers.
|
||||
/// On Windows, `name_prefix` is ignored and `tmpnam` is used.
|
||||
inline c10::optional<TempFile> try_make_tempfile(
|
||||
std::string name_prefix = "torch-file-") {
|
||||
#if defined(_WIN32)
|
||||
return TempFile{std::tmpnam(nullptr)};
|
||||
#else
|
||||
std::vector<char> filename = detail::make_filename(std::move(name_prefix));
|
||||
const int fd = mkstemp(filename.data());
|
||||
if (fd == -1) {
|
||||
return c10::nullopt;
|
||||
}
|
||||
// Don't make the string from string(filename.begin(), filename.end(), or
|
||||
// there will be a trailing '\0' at the end.
|
||||
return TempFile(filename.data(), fd);
|
||||
#endif // defined(_WIN32)
|
||||
}
|
||||
/// On Windows, `name_prefix` is ignored and `tmpnam_s` is used,
|
||||
/// and no temporary file is opened.
|
||||
C10_API std::optional<TempFile> try_make_tempfile(
|
||||
std::string_view name_prefix = "torch-file-");
|
||||
|
||||
/// Like `try_make_tempfile`, but throws an exception if a temporary file could
|
||||
/// not be returned.
|
||||
inline TempFile make_tempfile(std::string name_prefix = "torch-file-") {
|
||||
if (auto tempfile = try_make_tempfile(std::move(name_prefix))) {
|
||||
return std::move(*tempfile);
|
||||
}
|
||||
TORCH_CHECK(false, "Error generating temporary file: ", std::strerror(errno));
|
||||
}
|
||||
C10_API TempFile make_tempfile(std::string_view name_prefix = "torch-file-");
|
||||
|
||||
/// Attempts to return a temporary directory or returns `nullopt` if an error
|
||||
/// occurred.
|
||||
@ -161,40 +78,11 @@ inline TempFile make_tempfile(std::string name_prefix = "torch-file-") {
|
||||
/// `"TEMPDIR"` environment variable if any is set, or otherwise `/tmp`;
|
||||
/// `<name-prefix>` is the value supplied to this function, and
|
||||
/// `<random-pattern>` is a random sequence of numbers.
|
||||
/// On Windows, `name_prefix` is ignored and `tmpnam` is used.
|
||||
inline c10::optional<TempDir> try_make_tempdir(
|
||||
std::string name_prefix = "torch-dir-") {
|
||||
#if defined(_WIN32)
|
||||
while (true) {
|
||||
const char* dirname = std::tmpnam(nullptr);
|
||||
if (!dirname) {
|
||||
return c10::nullopt;
|
||||
}
|
||||
if (CreateDirectoryA(dirname, NULL)) {
|
||||
return TempDir(dirname);
|
||||
}
|
||||
if (GetLastError() != ERROR_ALREADY_EXISTS) {
|
||||
return c10::nullopt;
|
||||
}
|
||||
}
|
||||
return c10::nullopt;
|
||||
#else
|
||||
std::vector<char> filename = detail::make_filename(std::move(name_prefix));
|
||||
const char* dirname = mkdtemp(filename.data());
|
||||
if (!dirname) {
|
||||
return c10::nullopt;
|
||||
}
|
||||
return TempDir(dirname);
|
||||
#endif // defined(_WIN32)
|
||||
}
|
||||
/// On Windows, `name_prefix` is ignored.
|
||||
C10_API std::optional<TempDir> try_make_tempdir(
|
||||
std::string_view name_prefix = "torch-dir-");
|
||||
|
||||
/// Like `try_make_tempdir`, but throws an exception if a temporary directory
|
||||
/// could not be returned.
|
||||
inline TempDir make_tempdir(std::string name_prefix = "torch-dir-") {
|
||||
if (auto tempdir = try_make_tempdir(std::move(name_prefix))) {
|
||||
return std::move(*tempdir);
|
||||
}
|
||||
TORCH_CHECK(
|
||||
false, "Error generating temporary directory: ", std::strerror(errno));
|
||||
}
|
||||
C10_API TempDir make_tempdir(std::string_view name_prefix = "torch-dir-");
|
||||
} // namespace c10
|
||||
|
||||
@ -83,7 +83,7 @@ int main(int argc, char* argv[]) {
|
||||
setsid(); // Daemonize the process
|
||||
|
||||
std::unique_ptr<ManagerServerSocket> srv_socket;
|
||||
c10::optional<c10::TempDir> tempdir;
|
||||
std::optional<c10::TempDir> tempdir;
|
||||
try {
|
||||
tempdir = c10::try_make_tempdir(/*name_prefix=*/"torch-shm-dir-");
|
||||
if (!tempdir.has_value()) {
|
||||
|
||||
Reference in New Issue
Block a user