xgboost/tests/cpp/plugin/federated/test_federated_tracker.cc
Jiaming Yuan 3fbb221fec
[coll] Implement shutdown for tracker and comm. (#10208)
- Force shutdown the tracker.
- Implement shutdown notice for error handling thread in comm.
2024-04-20 04:08:17 +08:00

36 lines
973 B
C++

/**
* Copyright 2023-2024, XGBoost Contributors
*/
#include <gtest/gtest.h>
#include <memory> // for make_unique
#include <string> // for string
#include "../../../../src/collective/tracker.h" // for GetHostAddress
#include "federated_tracker.h"
#include "xgboost/json.h" // for Json
namespace xgboost::collective {
TEST(FederatedTrackerTest, Basic) {
Json config{Object()};
config["federated_secure"] = Boolean{false};
config["n_workers"] = Integer{3};
auto tracker = std::make_unique<FederatedTracker>(config);
ASSERT_FALSE(tracker->Ready());
auto fut = tracker->Run();
auto args = tracker->WorkerArgs();
ASSERT_TRUE(tracker->Ready());
ASSERT_GE(tracker->Port(), 1);
std::string host;
auto rc = GetHostAddress(&host);
ASSERT_EQ(get<String const>(args["dmlc_tracker_uri"]), host);
rc = tracker->Shutdown();
ASSERT_TRUE(rc.OK());
ASSERT_TRUE(fut.get().OK());
ASSERT_FALSE(tracker->Ready());
}
} // namespace xgboost::collective