File: TestUtils.hpp

package info (click to toggle)
pytorch 1.13.1%2Bdfsg-4
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 139,252 kB
  • sloc: cpp: 1,100,274; python: 706,454; ansic: 83,052; asm: 7,618; java: 3,273; sh: 2,841; javascript: 612; makefile: 323; xml: 269; ruby: 185; yacc: 144; objc: 68; lex: 44
file content (134 lines) | stat: -rw-r--r-- 2,392 bytes parent folder | download | duplicates (4)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
#pragma once

#ifndef _WIN32
#include <signal.h>
#include <sys/wait.h>
#include <unistd.h>
#endif

#include <sys/types.h>
#include <cstring>

#include <condition_variable>
#include <mutex>
#include <string>
#include <system_error>
#include <vector>

namespace c10d {
namespace test {

class Semaphore {
 public:
  void post(int n = 1) {
    std::unique_lock<std::mutex> lock(m_);
    n_ += n;
    cv_.notify_all();
  }

  void wait(int n = 1) {
    std::unique_lock<std::mutex> lock(m_);
    while (n_ < n) {
      cv_.wait(lock);
    }
    n_ -= n;
  }

 protected:
  int n_ = 0;
  std::mutex m_;
  std::condition_variable cv_;
};

#ifdef _WIN32
std::string autoGenerateTmpFilePath() {
  char tmp[L_tmpnam_s];
  errno_t err;
  err = tmpnam_s(tmp, L_tmpnam_s);
  if (err != 0)
  {
    throw std::system_error(errno, std::system_category());
  }
  return std::string(tmp);
}

std::string tmppath() {
  const char* tmpfile = getenv("TMPFILE");
  if (tmpfile) {
    return std::string(tmpfile);
  }
  else {
    return autoGenerateTmpFilePath();
  }
}
#else
std::string tmppath() {
  // TMPFILE is for manual test execution during which the user will specify
  // the full temp file path using the environmental variable TMPFILE
  const char* tmpfile = getenv("TMPFILE");
  if (tmpfile) {
    return std::string(tmpfile);
  }

  const char* tmpdir = getenv("TMPDIR");
  if (tmpdir == nullptr) {
    tmpdir = "/tmp";
  }

  // Create template
  std::vector<char> tmp(256);
  auto len = snprintf(tmp.data(), tmp.size(), "%s/testXXXXXX", tmpdir);
  tmp.resize(len);

  // Create temporary file
  auto fd = mkstemp(&tmp[0]);
  if (fd == -1) {
    throw std::system_error(errno, std::system_category());
  }
  close(fd);
  return std::string(tmp.data(), tmp.size());
}
#endif

bool isTSANEnabled() {
  auto s = std::getenv("PYTORCH_TEST_WITH_TSAN");
  return s && strcmp(s, "1") == 0;
}
struct TemporaryFile {
  std::string path;

  TemporaryFile() {
    path = tmppath();
  }

  ~TemporaryFile() {
    unlink(path.c_str());
  }
};

#ifndef _WIN32
struct Fork {
  pid_t pid;

  Fork() {
    pid = fork();
    if (pid < 0) {
      throw std::system_error(errno, std::system_category(), "fork");
    }
  }

  ~Fork() {
    if (pid > 0) {
      kill(pid, SIGKILL);
      waitpid(pid, nullptr, 0);
    }
  }

  bool isChild() {
    return pid == 0;
  }
};
#endif

} // namespace test
} // namespace c10d