|
1 | 1 | #include "core/util/include/util.hpp" |
2 | 2 |
|
3 | | -#include <cstdlib> |
4 | | -#include <cstdint> |
5 | | -#include <iostream> |
6 | | -#include <memory> |
7 | | -#include <vector> |
8 | 3 | #include <algorithm> |
9 | 4 | #include <array> |
10 | 5 | #include <filesystem> |
11 | | -#include <ranges> |
12 | | -#include <string> |
13 | | -#include <iostream> |
14 | 6 | #include <libenvpp/env.hpp> |
15 | | - |
16 | | -//namespace { |
17 | | -// |
18 | | -//struct MpiEnvVar { |
19 | | -// std::string_view name; |
20 | | -// const char* description; |
21 | | -//}; |
22 | | -// |
23 | | -//} // namespace |
| 7 | +#include <string> |
24 | 8 |
|
25 | 9 | std::string ppc::util::GetAbsolutePath(const std::string& relative_path) { |
26 | 10 | const std::filesystem::path path = std::string(PPC_PATH_TO_PROJECT) + "/tasks/" + relative_path; |
27 | 11 | return path.string(); |
28 | 12 | } |
29 | 13 |
|
30 | 14 | int ppc::util::GetNumThreads() { |
31 | | - env::prefix pre("PPC"); |
32 | | - |
33 | | - const auto num_threads_id = pre.register_variable<unsigned>("NUM_THREADS"); |
34 | | - |
35 | | - auto validated = std::move(pre).parse_and_validate(); |
36 | | - |
37 | | - if (!validated.ok()) { |
38 | | - std::cerr << validated.warning_message() << "\n"; |
39 | | -// throw validated.error_message(); |
| 15 | + const auto num_threads = env::get<int>("PPC_NUM_THREADS"); |
| 16 | + if (num_threads.has_value()) { |
| 17 | + return num_threads.value(); |
40 | 18 | } |
41 | | - |
42 | | - const auto num_threads = validated.get(num_threads_id); |
43 | | - return num_threads ? static_cast<int>(*num_threads) : 1; |
| 19 | + return 1; |
44 | 20 | } |
45 | 21 |
|
46 | | -//const std::array<MpiEnvVar, 13> kMpiEnvVars = {{{.name = "OMPI_COMM_WORLD_SIZE", .description = "OpenMPI"}, |
47 | | -// {.name = "OMPI_UNIVERSE_SIZE", .description = "OpenMPI"}, |
48 | | -// {.name = "PMI_SIZE", .description = "MPICH/Intel MPI"}, |
49 | | -// {.name = "PMI_RANK", .description = "MPICH/Intel MPI"}, |
50 | | -// {.name = "PMI_FD", .description = "MPICH/Intel MPI"}, |
51 | | -// {.name = "HYDRA_CONTROL_FD", .description = "Hydra"}, |
52 | | -// {.name = "PMIX_RANK", .description = "Cray PMI/PMIx"}, |
53 | | -// {.name = "PMIX_NAMESPACE", .description = "Cray PMI/PMIx"}, |
54 | | -// {.name = "SLURM_PROCID", .description = "SLURM"}, |
55 | | -// {.name = "MPI_LOCALRANKID", .description = "IBM Spectrum MPI"}, |
56 | | -// {.name = "MSMPI_RANK", .description = "Microsoft MPI"}, |
57 | | -// {.name = "MSMPI_NODE_ID", .description = "Microsoft MPI"}, |
58 | | -// {.name = "MSMPI_LOCALRANK", .description = "Microsoft MPI"}}}; |
| 22 | +constexpr std::array<std::string_view, 13> kMpiEnvVars = { |
| 23 | + "OMPI_COMM_WORLD_SIZE", "OMPI_UNIVERSE_SIZE", "PMI_SIZE", "PMI_RANK", "PMI_FD", |
| 24 | + "HYDRA_CONTROL_FD", "PMIX_RANK", "SLURM_PROCID", "MSMPI_RANK", "MSMPI_LOCALRANK"}; |
59 | 25 |
|
60 | | -//bool ppc::util::IsUnderMpirun() { |
61 | | -// auto pre = env::prefix(""); |
62 | | -// const auto parsed_and_validated_pre = pre.parse_and_validate(); |
63 | | -// return std::ranges::any_of(kMpiEnvVars, [&](const auto& env_var) { |
64 | | -// const auto mpi_env_id = pre.register_variable<std::string>(env_var.name); |
65 | | -// |
66 | | -// if (!parsed_and_validated_pre.ok()) { |
67 | | -// std::cerr << parsed_and_validated_pre.warning_message(); |
68 | | -// throw parsed_and_validated_pre.error_message(); |
69 | | -// } |
70 | | -// |
71 | | -// if (parsed_and_validated_pre.get(mpi_env_id)) { |
72 | | -// return true; |
73 | | -// } else { |
74 | | -// return false; |
75 | | -// } |
76 | | -// }); |
77 | | -//} |
| 26 | +bool ppc::util::IsUnderMpirun() { |
| 27 | + return std::ranges::any_of(kMpiEnvVars, [&](const auto& env_var) { |
| 28 | + const auto mpi_env = env::get<int>(env_var); |
| 29 | + return static_cast<bool>(mpi_env.has_value()); |
| 30 | + }); |
| 31 | +} |
0 commit comments