Skip to content

Commit 1016eab

Browse files
committed
test: move logging logic to a dedicated file
Signed-off-by: Gabriel Dos Santos <[email protected]>
1 parent 5caaf5d commit 1016eab

File tree

3 files changed

+82
-59
lines changed

3 files changed

+82
-59
lines changed

unit_tests/CMakeLists.txt

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -182,14 +182,18 @@ target_sources(
182182
PRIVATE
183183
test_main.cpp
184184
test_sendrecv.cpp
185+
PRIVATE
186+
FILE_SET HEADERS
187+
BASE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}
188+
FILES logging.hpp
185189
)
186190
target_compile_features(test-main PRIVATE cxx_std_20)
187191
target_link_libraries(
188192
test-main
189193
PRIVATE
194+
GTest::gtest_main
190195
KokkosComm::KokkosComm
191196
MPI::MPI_CXX
192-
GTest::gtest_main
193197
fmt::fmt
194198
)
195199

@@ -219,7 +223,8 @@ if(KokkosComm_ENABLE_NCCL)
219223
test-main
220224
PRIVATE
221225
FILE_SET HEADERS
222-
BASE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/nccl
226+
BASE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}
227+
FILES nccl/utils.hpp
223228
)
224229
# Link against NCCL because we depend on raw calls
225230
target_link_libraries(test-main PRIVATE NCCL::NCCL)

unit_tests/logging.hpp

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
2+
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
3+
4+
#pragma once
5+
6+
#include <array>
7+
#include <cstdlib>
8+
#include <string_view>
9+
10+
#include <cuda.h>
11+
#include <mpi.h>
12+
#if defined(KOKKOSCOMM_ENABLE_NCCL)
13+
#include <nccl.h>
14+
#endif
15+
16+
#include <fmt/core.h>
17+
18+
namespace logging {
19+
20+
enum struct Level {
21+
FATAL,
22+
ERROR,
23+
WARN,
24+
INFO,
25+
TRACE,
26+
};
27+
28+
using namespace std::string_view_literals;
29+
constexpr std::array level_txt{"FATAL"sv, "ERROR"sv, "WARNING"sv, "INFO"sv, "TRACE"sv};
30+
31+
} // namespace logging
32+
33+
#define KC_LOG(lvl, ...) \
34+
fmt::println("[{}] {}:{}: {}", logging::level_txt[static_cast<int>(lvl)], __FILE__, __LINE__, \
35+
fmt::format(__VA_ARGS__))
36+
37+
#define KC_FATAL(...) (KC_LOG(logging::Level::FATAL, __VA_ARGS__), std::exit(EXIT_FAILURE))
38+
39+
#define KC_ERROR(...) KC_LOG(logging::Level::ERROR, __VA_ARGS__)
40+
41+
#define KC_WARN(...) KC_LOG(logging::Level::WARN, __VA_ARGS__)
42+
43+
#define KC_INFO(...) KC_LOG(logging::Level::INFO, __VA_ARGS__)
44+
45+
#define KC_TRACE(...) KC_LOG(logging::Level::TRACE, __VA_ARGS__)
46+
47+
#define KC_CHECK(expr, ...) ((expr) ? void(0) : KC_FATAL(__VA_ARGS__))
48+
49+
#define KC_CUDA_CHECK(expr) \
50+
([&]() { \
51+
cudaError_t kc_res_ = (expr); \
52+
return kc_res_ == cudaSuccess ? void(0) \
53+
: KC_FATAL("CUDA check failed: `" #expr "`: {}", cudaGetErrorString(kc_res_)); \
54+
}())
55+
56+
#define KC_MPI_CHECK(expr) \
57+
([&]() { \
58+
int kc_res_ = (expr); \
59+
return kc_res_ == MPI_SUCCESS ? void(0) : KC_FATAL("MPI check failed: `" #expr "`: {}", kc_res_); \
60+
}())
61+
62+
#if defined(KOKKOSCOMM_ENABLE_NCCL)
63+
#define KC_NCCL_CHECK(expr) \
64+
([&]() { \
65+
ncclResult_t kc_res_ = (expr); \
66+
return kc_res_ == ncclSuccess ? void(0) \
67+
: KC_FATAL("NCCL check failed: `" #expr "`: {}", ncclGetErrorString(kc_res_)); \
68+
}())
69+
#endif

unit_tests/nccl/utils.hpp

Lines changed: 6 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -3,68 +3,17 @@
33

44
#pragma once
55

6-
#include <array>
7-
#include <cstddef>
8-
#include <cstdint>
96
#include <cstdio>
107
#include <cstdlib>
11-
#include <fstream>
12-
#include <string>
13-
#include <string_view>
14-
#include <vector>
158

169
#include <fmt/core.h>
1710
#include <mpi.h>
1811
#include <nccl.h>
1912
#include <cuda_runtime.h>
2013

21-
namespace {
22-
23-
enum struct LogLevel {
24-
FATAL,
25-
ERROR,
26-
WARN,
27-
INFO,
28-
TRACE,
29-
};
30-
31-
using namespace std::string_view_literals;
32-
constexpr std::array level_txt{"FATAL"sv, "ERROR"sv, "WARNING"sv, "INFO"sv, "TRACE"sv};
33-
34-
#define KC_LOG(lvl, ...) \
35-
fmt::println("[{}] {}:{}: {}", level_txt[static_cast<int>(lvl)], __FILE__, __LINE__, fmt::format(__VA_ARGS__))
36-
37-
#define KC_FATAL(...) (KC_LOG(LogLevel::FATAL, __VA_ARGS__), std::exit(EXIT_FAILURE))
38-
39-
#define KC_ERROR(...) KC_LOG(LogLevel::ERROR, __VA_ARGS__)
14+
#include <unit_tests/logging.hpp>
4015

41-
#define KC_WARN(...) KC_LOG(LogLevel::WARN, __VA_ARGS__)
42-
43-
#define KC_INFO(...) KC_LOG(LogLevel::INFO, __VA_ARGS__)
44-
45-
#define KC_TRACE(...) KC_LOG(LogLevel::TRACE, __VA_ARGS__)
46-
47-
#define KC_CHECK(expr, ...) ((expr) ? void(0) : KC_FATAL(__VA_ARGS__))
48-
49-
#define KC_MPI_CHECK(expr) \
50-
([&]() { \
51-
int kc_res_ = (expr); \
52-
return kc_res_ == MPI_SUCCESS ? void(0) : KC_FATAL("MPI check failed: `" #expr "`: {}", kc_res_); \
53-
}())
54-
55-
#define KC_NCCL_CHECK(expr) \
56-
([&]() { \
57-
ncclResult_t kc_res_ = (expr); \
58-
return kc_res_ == ncclSuccess ? void(0) \
59-
: KC_FATAL("NCCL check failed: `" #expr "`: {}", ncclGetErrorString(kc_res_)); \
60-
}())
61-
62-
#define KC_CUDA_CHECK(expr) \
63-
([&]() { \
64-
cudaError_t kc_res_ = (expr); \
65-
return kc_res_ == cudaSuccess ? void(0) \
66-
: KC_FATAL("CUDA check failed: `" #expr "`: {}", cudaGetErrorString(kc_res_)); \
67-
}())
16+
namespace {
6817

6918
[[nodiscard]] auto get_local_rank(MPI_Comm comm, int my_rank) -> int {
7019
MPI_Comm node_comm;
@@ -94,16 +43,15 @@ class Ctx {
9443
int n_ranks, my_rank;
9544
MPI_Comm_size(mpi_comm, &n_ranks);
9645
MPI_Comm_rank(mpi_comm, &my_rank);
97-
KC_INFO("P%d/%d - MPI initialized", n_ranks, my_rank);
9846
int local_rank = get_local_rank(mpi_comm, my_rank);
9947

10048
int n_gpus;
10149
KC_CUDA_CHECK(cudaGetDeviceCount(&n_gpus));
102-
KC_INFO("P%d found %d CUDA devices", my_rank, n_gpus);
50+
KC_INFO("P{} found {} CUDA devices", my_rank, n_gpus);
10351

104-
KC_CHECK(local_rank <= n_gpus, "P%d needs GPU %d but only %d devices available", my_rank, local_rank, n_gpus);
52+
KC_CHECK(local_rank <= n_gpus, "P{} needs device #{} but only {} devices available", my_rank, local_rank, n_gpus);
10553
KC_CUDA_CHECK(cudaSetDevice(local_rank));
106-
KC_INFO("P%d assigned to CUDA device %d", my_rank, local_rank);
54+
KC_INFO("P{} assigned to CUDA device #{}", my_rank, local_rank);
10755

10856
// Get NCCL unique ID at rank 0 and broadcast it to all others
10957
ncclUniqueId nccl_id;
@@ -125,6 +73,7 @@ class Ctx {
12573
}
12674

12775
~Ctx() { KC_NCCL_CHECK(ncclCommDestroy(comm_)); }
76+
// Forbid copies and moves
12877
Ctx(const Ctx &) = delete;
12978
auto operator=(const Ctx &) -> Ctx & = delete;
13079
Ctx(Ctx &&) = delete;

0 commit comments

Comments
 (0)