TEST/CUDA: Restore original cuda device

Author: yosefe

test/gtest/uct/cuda/test_switch_cuda_device.cc

@@ -288,6 +288,9 @@ class test_mem_alloc_device : public test_switch_cuda_device {
     void init() override {
         test_switch_cuda_device::init();
 
+        int current_device;
+        ASSERT_EQ(cudaGetDevice(&current_device), cudaSuccess);
+
         for (auto device = 0; device < m_num_devices; ++device) {
             uct_md_mem_attr_t attr       = {
                 .field_mask = UCT_MD_MEM_ATTR_FIELD_SYS_DEV
@@ -305,6 +308,8 @@ class test_mem_alloc_device : public test_switch_cuda_device {
             m_sys_dev.push_back(attr.sys_dev);
         }
 
+        EXPECT_EQ(cudaSetDevice(current_device), cudaSuccess);
+
         ASSERT_EQ(m_num_devices, m_sys_dev.size());
     }
 
@@ -349,6 +354,9 @@ class test_mem_alloc_device : public test_switch_cuda_device {
     {
         CUdevice current;
 
+        int current_device;
+        ASSERT_EQ(cudaGetDevice(&current_device), cudaSuccess);
+
         // Ensure a valid context for each device
         for (auto device = 0; device < m_num_devices; ++device) {
             ASSERT_EQ(cudaSetDevice(device), cudaSuccess);
@@ -365,11 +373,16 @@ class test_mem_alloc_device : public test_switch_cuda_device {
             EXPECT_EQ(m_num_devices - 1, current);
             ASSERT_UCS_OK(uct_mem_free(&mem));
         }
+
+        EXPECT_EQ(cudaSetDevice(current_device), cudaSuccess);
     }
 
     void test_same_device_alloc(ucs_memory_type_t mem_type,
                                 bool set_sys_dev = true)
     {
+        int current_device;
+        ASSERT_EQ(cudaGetDevice(&current_device), cudaSuccess);
+
         for (auto device = 0; device < m_num_devices; ++device) {
             ASSERT_EQ(cudaSetDevice(device), cudaSuccess);
             ASSERT_UCS_OK(allocate(mem_type,
@@ -379,6 +392,8 @@ class test_mem_alloc_device : public test_switch_cuda_device {
             EXPECT_EQ(m_sys_dev[device], sys_device);
             ASSERT_UCS_OK(uct_mem_free(&mem));
         }
+
+        EXPECT_EQ(cudaSetDevice(current_device), cudaSuccess);
     }
 
     void skip_if_no_fabric(ucs_memory_type_t mem_type)
@@ -526,7 +541,7 @@ class test_p2p_send_on_diff_device : public uct_p2p_test {
         rkey_release(sender(), rkey_dest);
 
         // Set back an original device to avoid test failures
-        ASSERT_EQ(cudaGetDevice(&current_device), cudaSuccess);
+        EXPECT_EQ(cudaSetDevice(current_device), cudaSuccess);
     }
 
     template<typename mem_alloc_t1, typename mem_alloc_t2>