fix: prevent DCGM handle/connection leak on connectivity failure (#1078) (#1089)

Author: deesharma24

health-monitors/gpu-health-monitor/gpu_health_monitor/dcgm_watcher/dcgm.py

@@ -259,33 +259,48 @@ def _initialize_dcgm_monitoring(self, dcgm_handle: pydcgm.DcgmHandle) -> tuple:
 
         Returns:
             A tuple of (dcgm_group, gpu_ids, gpu_serials)
+
+        If any step after group creation fails the group is deleted before the
+        exception propagates so that it does not leak on the DCGM server.
         """
         dcgm_group = self._create_dcgm_group_with_all_entities(dcgm_handle)
-        with metrics.dcgm_api_latency.labels("group_health_set").time():
-            dcgm_group.health.Set(dcgm_structs.DCGM_HEALTH_WATCH_ALL)
+        try:
+            with metrics.dcgm_api_latency.labels("group_health_set").time():
+                dcgm_group.health.Set(dcgm_structs.DCGM_HEALTH_WATCH_ALL)
 
-        gpu_ids = dcgm_group.GetGpuIds()
-        gpu_serials = self._get_gpu_serial_numbers(dcgm_handle)
-        log.info(f"dcgm gpu_id are {gpu_ids}")
+            gpu_ids = dcgm_group.GetGpuIds()
+            gpu_serials = self._get_gpu_serial_numbers(dcgm_handle)
+            log.info(f"dcgm gpu_id are {gpu_ids}")
 
-        return dcgm_group, gpu_ids, gpu_serials
+            return dcgm_group, gpu_ids, gpu_serials
+        except Exception as e:
+            log.warning(f"DCGM monitoring initialization failed, rolling back group: {e}")
+            try:
+                dcgm_group.Delete()
+            except Exception as del_err:
+                log.warning(f"Failed to delete DCGM group during init rollback: {del_err}")
+                metrics.dcgm_api_failures.labels("init_group_rollback").inc()
+            raise
 
     def _cleanup_dcgm_resources(
         self,
         dcgm_group: pydcgm.DcgmGroup,
         dcgm_handle: pydcgm.DcgmHandle,
     ):
-        """Clean up DCGM resources safely."""
-        try:
-            if dcgm_group:
+        """Clean up DCGM resources safely.
+
+        Group deletion and handle shutdown are in separate try blocks so that
+        a failure in Delete() does not prevent Shutdown() from running.
+        """
+        if dcgm_group:
+            try:
                 dcgm_group.Delete()
-                dcgm_group = None
-            if dcgm_handle:
-                # Clean up the handle
-                dcgm_handle.Shutdown()
-                del dcgm_handle
-        except Exception as e:
-            log.error(f"Error cleaning up DCGM handle: {e}")
+            except Exception as e:
+                log.warning(f"Error deleting DCGM group (will still shut down handle): {e}")
+                metrics.dcgm_api_failures.labels("group_delete_error").inc()
+
+        if dcgm_handle:
+            dcgm_handle.Shutdown()
 
     def start(self, fields_to_monitor: list[str], exit: Event) -> None:
         dcgm_handle = None

health-monitors/gpu-health-monitor/gpu_health_monitor/tests/test_dcgm_watcher/test_dcgm.py

@@ -17,6 +17,7 @@
 import dcgm_structs, dcgm_errors, dcgm_fields, dcgm_field_helpers
 from threading import Event, Thread
 from ctypes import pointer
+import pytest
 
 
 class FakeEventProcessorInTest(dcgm.types.CallbackInterface):
@@ -505,3 +506,63 @@ def test_initialize_dcgm_monitoring(self, mock_dcgm_group, mock_dcgm_handle):
         assert len(gpu_serials) == 4
         # Verify that health.Set was called on the actual group object
         group.health.Set.assert_called_once()
+
+
+class TestDCGMHandleLeakFix:
+    """Tests for the DCGM handle/connection leak fix (issue #1078).
+
+    Covers: split try-blocks in cleanup and init rollback.
+    """
+
+    def _make_watcher(self):
+        return dcgm.DCGMWatcher(
+            addr="localhost:5555",
+            poll_interval_seconds=10,
+            callbacks=[],
+            dcgm_k8s_service_enabled=False,
+        )
+
+    @pytest.mark.parametrize(
+        "group_is_none, delete_raises",
+        [
+            (False, True),
+            (True, False),
+        ],
+        ids=["delete_throws", "none_group"],
+    )
+    def test_cleanup_dcgm_resources(self, group_is_none, delete_raises):
+        """Shutdown() always runs regardless of Delete() outcome."""
+        watcher = self._make_watcher()
+        dcgm_handle_mock = MagicMock()
+
+        dcgm_group_mock = None
+        if not group_is_none:
+            dcgm_group_mock = MagicMock()
+            if delete_raises:
+                dcgm_group_mock.Delete.side_effect = Exception("Delete failed")
+
+        watcher._cleanup_dcgm_resources(dcgm_group_mock, dcgm_handle_mock)
+
+        dcgm_handle_mock.Shutdown.assert_called_once()
+        if not group_is_none:
+            dcgm_group_mock.Delete.assert_called_once()
+
+    @patch("pydcgm.DcgmGroup.__new__")
+    def test_init_monitoring_rolls_back_group_on_failure(self, mock_dcgm_group):
+        """Group must be deleted if initialization fails after group creation."""
+        watcher = self._make_watcher()
+        dcgm_handle_mock = MagicMock()
+        dcgm_group_mock = MagicMock()
+        mock_dcgm_group.return_value = dcgm_group_mock
+
+        dcgm_system_mock = MagicMock()
+        dcgm_system_mock.discovery.GetEntityGroupEntities.return_value = [0, 1]
+        dcgm_handle_mock.GetSystem.return_value = dcgm_system_mock
+
+        # health.Set() fails after group is created
+        dcgm_group_mock.health.Set.side_effect = Exception("DCGM connection lost")
+
+        with pytest.raises(Exception, match="DCGM connection lost"):
+            watcher._initialize_dcgm_monitoring(dcgm_handle_mock)
+
+        dcgm_group_mock.Delete.assert_called_once()