Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Destroy the ComputationClient when the program exits #5750

Merged
merged 2 commits into from
Oct 31, 2023
Merged

Destroy the ComputationClient when the program exits #5750

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
10d3cdb
Destroy the ComputationClient when the program exits
will-cromar Oct 30, 2023
022d854
Fix extra error when PJRT_DEVICE is not set
will-cromar Oct 31, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion torch_xla/csrc/init_python_bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,6 @@ void PrepareToExit() {
runtime::GetComputationClientIfInitialized();
if (client != nullptr) {
XLAGraphExecutor::Get()->WaitDeviceOps({});
client->PrepareToExit();
}
}

Expand Down
2 changes: 0 additions & 2 deletions torch_xla/csrc/runtime/computation_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -344,8 +344,6 @@ class ComputationClient {

virtual MemoryInfo GetMemoryInfo(const std::string& device) = 0;

virtual void PrepareToExit() = 0;

// Block until pass in devices' async operation are finished. If empty, all
// the local devices will be waited for.
virtual void WaitDeviceOps(const std::vector<std::string>& devices) = 0;
Expand Down
2 changes: 0 additions & 2 deletions torch_xla/csrc/runtime/pjrt_computation_client.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,6 @@ class PjRtComputationClient : public ComputationClient {

std::shared_ptr<std::vector<std::string>> GetReplicationDevices() override;

void PrepareToExit() override { return; };

void WaitDeviceOps(const std::vector<std::string>& devices) override;

std::map<std::string, Metric> GetMetrics() const override;
Expand Down
13 changes: 7 additions & 6 deletions torch_xla/csrc/runtime/runtime.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,11 @@ namespace torch_xla {
namespace runtime {
namespace {

std::atomic<ComputationClient*> g_computation_client(nullptr);
std::once_flag g_computation_client_once;
std::atomic<bool> g_computation_client_initialized(false);

ComputationClient* CreateClient() {
bool was_initialized = g_computation_client_initialized.exchange(true);
XLA_CHECK(!was_initialized) << "ComputationClient already initialized";
if (sys_util::GetEnvBool("XLA_DUMP_FATAL_STACK", false)) {
tsl::testing::InstallStacktraceHandler();
}
Expand All @@ -23,6 +24,7 @@ ComputationClient* CreateClient() {
if (sys_util::GetEnvString(env::kEnvPjRtDevice, "") != "") {
client = new PjRtComputationClient();
} else {
g_computation_client_initialized = false;
XLA_ERROR() << "$PJRT_DEVICE is not set." << std::endl;
}

Expand All @@ -34,13 +36,12 @@ ComputationClient* CreateClient() {
} // namespace

ComputationClient* GetComputationClient() {
std::call_once(g_computation_client_once,
[&]() { g_computation_client = std::move(CreateClient()); });
return g_computation_client.load();
static auto client = std::unique_ptr<ComputationClient>(CreateClient());
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the static variable is guaranteed to be initialized once, do you still need the function GetComputationClientIfInitialized and the variable std::atomic<bool> g_computation_client_initialized(false);?

Copy link
Collaborator Author

@will-cromar will-cromar Oct 31, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The client is guaranteed to be initialized once on the first time we call into GetComputationClient. There are a couple of calls that want to tell if client has been initialized (ie comparing GetComputationClientIfInitialized to nullptr), but there's not a way I can see to tell if client has been initialized. As soon as you call GetComputationClient, it gets created. That's why I added the extra flag outside of GetComputationClient.

On a second look, I think it would actually be clearer to eliminate ComputationClient* GetComputationClientIfInitialized and replace it with bool ComputationClientIsInitialzed().

return client.get();
}

ComputationClient* GetComputationClientIfInitialized() {
return g_computation_client.load();
return g_computation_client_initialized ? GetComputationClient() : nullptr;
}

} // namespace runtime
Expand Down