Fix NCCLX UnitTest

Summary:
Mocks doesn't return any valid cuda device count. The code was relying on
uninitialized value, at the mercy of compiler. Recently the code started to
segfault when it does `rank_ % device_count` with device_count being 0

Changes
- Explicitly initialize the device_count to 0
- Extend mock to always return `1` by default for `getDeviceCount`
- Fix the UT to setup mock behaviors by default

Test Failure - https://www.internalfb.com/intern/test/562950190647332

```
[ RUN      ] TorchCommNCCLXTest.InitializationFailsWithInvalidDeviceId

I1026 04:38:23.001531 1483368 TorchCommNCCLXBootstrap.cpp:43] [TC] TORCHCOMM_NCCLX_BOOTSTRAP_UNIQUEID_EXCHANGE_METHOD not set, defaulting to auto
*** Aborted at 1761478703 (Unix time, try 'date -d 1761478703') ***
*** Signal 8 (SIGFPE) (0x3170a51f) received by PID 1483368 (pthread TID 0x7fd11c3e3000) (linux TID 1483368) (code: integer divide by zero), stack trace: ***
    @ 000000000ba032ff folly::symbolizer::(anonymous namespace)::signalHandler(int, siginfo_t*, void*)
                       ./fbcode/folly/debugging/symbolizer/SignalHandler.cpp:528
    @ 000000000004455f (unknown)
                       /home/engshare/third-party2/glibc/2.34/src/glibc-2.34/signal/../sysdeps/unix/sysv/linux/libc_sigaction.c:8
                       -> /home/engshare/third-party2/glibc/2.34/src/glibc-2.34/signal/../sysdeps/unix/sysv/linux/x86_64/libc_sigaction.c
    @ 000000003170a51f torch::comms::TorchCommNCCLXBootstrap::TorchCommNCCLXBootstrap(c10::intrusive_ptr<c10d::Store, c10::detail::intrusive_target_default_null_type<c10d::Store> >, c10::Device, std::shared_ptr<torch::comms::NcclxApi>, std::shared_ptr<torch::comms::CudaApi>, std::chrono::duration<long, std::ratio<1l, 1000l> >)
                       ./fbcode/comms/torchcomms/ncclx/TorchCommNCCLXBootstrap.cpp:63
    @ 00000000316fd281 torch::comms::TorchCommNCCLX::init(c10::Device, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, torch::comms::CommOptions const&)
                       ./fbcode/comms/torchcomms/ncclx/TorchCommNCCLX.cpp:74
    @ 000000000b94fb02 torch::comms::test::TorchCommNCCLXTest_InitializationFailsWithInvalidDeviceId_Test::TestBody()
                       ./fbcode/comms/torchcomms/ncclx/tests/unit/cpp/TorchCommNCCLXTest.cpp:155
    @ 000000000b9a586f testing::Test::Run()
                       fbsource/src/gtest.cc:2751
    @ 000000000b9a6cde testing::TestInfo::Run()
                       fbsource/src/gtest.cc:2897
    @ 000000000b9a80f4 testing::TestSuite::Run()
                       fbsource/src/gtest.cc:3075
    @ 000000000b9ba72f testing::internal::UnitTestImpl::RunAllTests()
                       fbsource/src/gtest.cc:6066
    @ 000000000b9b980f testing::UnitTest::Run()
                       fbsource/src/gtest.cc:5606
    @ 000000000b9c977d main
                       fbsource/gtest/gtest.h:2337
    @ 000000000002c656 __libc_start_call_main
                       /home/engshare/third-party2/glibc/2.34/src/glibc-2.34/csu/../sysdeps/nptl/libc_start_call_main.h:58
                       -> /home/engshare/third-party2/glibc/2.34/src/glibc-2.34/csu/../sysdeps/x86/libc-start.c
    @ 000000000002c717 __libc_start_main
                       /home/engshare/third-party2/glibc/2.34/src/glibc-2.34/csu/../csu/libc-start.c:409
                       -> /home/engshare/third-party2/glibc/2.34/src/glibc-2.34/csu/../sysdeps/x86/libc-start.c
    @ 000000000b944220 _start
                       /home/engshare/third-party2/glibc/2.34/src/glibc-2.34/csu/../sysdeps/x86_64/start.S:116
```

Reviewed By: pavanbalaji

Differential Revision: D85489376

fbshipit-source-id: 993377d73b05c91cb0df4219e4f8f4e9b78de36b

Author: Saif Hasan

comms/torchcomms/ncclx/TorchCommNCCLXBootstrap.cpp

@@ -1,13 +1,15 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 
-#include "comms/torchcomms/ncclx/TorchCommNCCLXBootstrap.hpp"
+#include <stdexcept>
+
 #include <ATen/cuda/CUDAContext.h>
 #include <dlfcn.h>
 #include <torch/csrc/distributed/c10d/TCPStore.hpp> // @manual
 #include "comms/torchcomms/StoreManager.hpp"
 #include "comms/torchcomms/TorchCommLogging.hpp"
 #include "comms/torchcomms/TorchCommUtils.hpp"
 #include "comms/torchcomms/ncclx/TorchCommNCCLX.hpp"
+#include "comms/torchcomms/ncclx/TorchCommNCCLXBootstrap.hpp"
 #include "nccl.h" // @manual
 
 namespace torch {
@@ -54,11 +56,16 @@ TorchCommNCCLXBootstrap::TorchCommNCCLXBootstrap(
       [](unsigned char c) { return std::tolower(c); });
 
   if (device_.index() == -1) {
-    int device_count;
+    int device_count{0};
     CUDA_CHECK(
         cuda_api_,
         cuda_api_->getDeviceCount(&device_count),
         "Failed to get CUDA device count");
+    if (device_count <= 0) {
+      throw std::invalid_argument(
+          "No CUDA devices found; please check your CUDA installation");
+    }
+    TC_LOG(INFO, nullptr) << "Found " << device_count << " CUDA devices";
 
     device_ = c10::Device(c10::kCUDA, rank_ % device_count);
     TC_LOG(INFO, nullptr)

comms/torchcomms/ncclx/tests/unit/cpp/TorchCommNCCLXTest.cpp

@@ -51,10 +51,10 @@ TEST_F(TorchCommNCCLXTest, InitializationRank0GetUniqueId) {
   // Clear - 1 (destructor)
   setupCCAExpectations(1, 2, 1);
 
-  auto comm = createMockedTorchComm();
-
   cuda_mock_->setupDefaultBehaviors();
 
+  auto comm = createMockedTorchComm();
+
   // Expect rank 0 to get unique ID and store it
   ncclUniqueId expected_id{};
   memset(&expected_id, 0x42, sizeof(expected_id)); // Fill with test pattern
@@ -134,6 +134,8 @@ TEST_F(TorchCommNCCLXTest, InitializationFailsWithInvalidDeviceId) {
   // Setup CCA expectations - first init (device -1) should succeed
   setupCCAExpectations(1, 2, 0);
 
+  cuda_mock_->setupDefaultBehaviors();
+
   // Test with negative device ID
   {
     at::Device invalid_device(at::DeviceType::CUDA, -1);

comms/torchcomms/ncclx/tests/unit/cpp/mocks/CudaMock.cpp

@@ -16,6 +16,9 @@ void CudaMock::setupDefaultBehaviors() {
   // Device management - return success by default
   ON_CALL(*this, setDevice(_)).WillByDefault(Return(cudaSuccess));
 
+  ON_CALL(*this, getDeviceCount(_))
+      .WillByDefault(DoAll(SetArgPointee<0>(1), Return(cudaSuccess)));
+
   ON_CALL(*this, getDeviceProperties(_, _))
       .WillByDefault(
           DoAll(SetArgPointee<0>(cudaDeviceProp{}), Return(cudaSuccess)));