Add Thermal Device Selection and Current Context Thermovision mapping (#51) - Fixes #35

* add thermal_device selection and current-context thermovision mapping

* refactor thermovision CUDA-device mapping and add verbose for default current-device monitoring

* Fixes #35

Author: hsahuja0805

nsight/analyze.py

@@ -50,6 +50,7 @@ def kernel(
     thermal_wait: int | None = None,
     thermal_cont: int | None = None,
     thermal_timeout: int | None = None,
+    thermal_device: int | None = None,
     combine_kernel_metrics: Callable[[float, float], float] | None = None,
     output_prefix: str | None = None,
     output_csv: bool = False,
@@ -79,6 +80,7 @@ def kernel(
     thermal_wait: int | None = None,
     thermal_cont: int | None = None,
     thermal_timeout: int | None = None,
+    thermal_device: int | None = None,
     combine_kernel_metrics: Callable[[float, float], float] | None = None,
     output_prefix: str | None = None,
     output_csv: bool = False,
@@ -220,6 +222,14 @@ def wrapped_function(*args, configs=None, **kwargs) -> ProfileResults
         thermal_timeout: Maximum wait time in seconds for GPU to cool down.
             Default: 180 seconds
 
+        thermal_device: CUDA device ordinal to monitor for thermal throttling.
+
+            If set, monitor that CUDA device. If ``None``, monitor the current
+            CUDA device context. Device numbering follows CUDA runtime ordinals and
+            honors ``CUDA_VISIBLE_DEVICES``.
+
+            Default: ``None``
+
         verbosity: Controls the verbosity level of the output. Selecting level N
             enables output at all levels with value <= N. Default: ``VerbosityLevel.INFO``.
 
@@ -313,6 +323,7 @@ def _create_profiler() -> collection.core.NsightProfiler:
             thermal_wait=thermal_wait,
             thermal_cont=thermal_cont,
             thermal_timeout=thermal_timeout,
+            thermal_device=thermal_device,
             output_prefix=prefix,
             output_csv=output_csv,
         )

nsight/collection/core.py

@@ -241,6 +241,7 @@ def run_profile_session(
     thermal_wait: int | None = None,
     thermal_cont: int | None = None,
     thermal_timeout: int | None = None,
+    thermal_device: int | None = None,
 ) -> None:
 
     if verbosity >= VerbosityLevel.INFO:
@@ -255,6 +256,7 @@ def run_profile_session(
             thermal_wait=thermal_wait,
             thermal_cont=thermal_cont,
             thermal_timeout=thermal_timeout,
+            thermal_device=thermal_device,
             verbose=verbosity >= VerbosityLevel.DEBUG,
         )
         thermovision_initialized = thermal_controller.init()
@@ -454,6 +456,16 @@ class ProfileSettings:
     If None, uses default value (180 seconds).
     """
 
+    thermal_device: int | None
+    """
+    CUDA device ordinal to monitor for thermal throttling.
+
+    If set, monitor that CUDA device. If None, monitor the current CUDA device
+    context.
+    Device numbering follows CUDA runtime ordinals and honors
+    ``CUDA_VISIBLE_DEVICES``.
+    """
+
     output_prefix: str | None
     """
     All intermediate profiler files are created with this prefix

nsight/collection/ncu.py

@@ -410,6 +410,7 @@ def collect(
                 settings.thermal_wait,
                 settings.thermal_cont,
                 settings.thermal_timeout,
+                settings.thermal_device,
             )
         finally:
             end_profiling()

nsight/thermovision.py

@@ -17,6 +17,7 @@
 
 # Guard NVML imports
 try:
+    from cuda.core import Device as CudaDevice
     from cuda.core import system
 
     CUDA_CORE_AVAILABLE = True
@@ -52,6 +53,7 @@ def __init__(
         thermal_wait: int | None = None,
         thermal_cont: int | None = None,
         thermal_timeout: int | None = None,
+        thermal_device: int | None = None,
         verbose: bool = False,
     ):
         """Initialize thermal controller.
@@ -71,10 +73,17 @@ def __init__(
             thermal_timeout: Maximum wait time in seconds for GPU to cool down.
                 If None, uses default (180 seconds).
                 Default: None
+            thermal_device: CUDA device ordinal to monitor for thermal throttling.
+                If set, monitor that CUDA device. If None, monitor the current
+                CUDA device context. Device numbering follows CUDA runtime ordinals and
+                honors ``CUDA_VISIBLE_DEVICES``.
+                Default: None
             verbose: Whether to print thermal messages.
                 Default: False
         """
         self.device: Any = None
+        self.thermal_device = thermal_device
+        self._current_cuda_device_id: int | None = None
         self.thermal_mode = thermal_mode
 
         # Set adaptive_mode based on thermal_mode
@@ -101,17 +110,111 @@ def __init__(
     def init(self) -> bool:
         """Initialize NVML and get GPU handle.
 
+        Resolves the NVML device that corresponds to the profiled CUDA device
+        (honoring ``CUDA_VISIBLE_DEVICES``).
+
         Returns:
             True if temperature retrieval is supported, False otherwise.
         """
         if not CUDA_CORE_AVAILABLE:
             return False
 
         if self.device is None:
-            self.device = system.Device(index=0)
+            resolved_device = self._resolve_device()
+            if resolved_device is None:
+                return False
+            self.device = resolved_device
 
         return self._is_temp_retrieval_supported()
 
+    def _map_cuda_to_system_device(self, cuda_device: Any) -> Any:
+        """Map a CUDA device object to its corresponding NVML system device.
+
+        Args:
+            cuda_device: CUDA device object to map.
+
+        Returns:
+            The corresponding NVML ``system.Device``.
+        """
+        system_device = cuda_device.to_system_device()
+        self._current_cuda_device_id = cuda_device.device_id
+        return system_device
+
+    def _resolve_device(self) -> Any:
+        """Resolve the NVML device to monitor.
+
+        The CUDA device (explicitly requested via ``thermal_device`` or the
+        current CUDA device context) is mapped to its underlying physical
+        NVML device by UUID. This ensures the GPU actually running the profiled
+        kernels is monitored, even when ``CUDA_VISIBLE_DEVICES`` remaps device
+        ordinals.
+
+        Returns:
+            The resolved NVML ``system.Device`` to monitor, or ``None`` if no
+            device could be resolved.
+        """
+        cuda_ordinal = self.thermal_device
+        try:
+            # ``None`` selects the current CUDA device and already reflects
+            # CUDA_VISIBLE_DEVICES. Mapping to system device is done by UUID.
+            if self.verbose and cuda_ordinal is None:
+                print(
+                    "[Thermovision] thermal_device not set; using current CUDA device context "
+                    "(set thermal_device to pin monitoring)."
+                )
+            cuda_device = CudaDevice(cuda_ordinal)
+            system_device = self._map_cuda_to_system_device(cuda_device)
+            if self.verbose:
+                print(
+                    f"[Thermovision] Monitoring CUDA device {self._current_cuda_device_id} "
+                    f"(NVML index {system_device.index}, {system_device.name}, UUID {system_device.uuid})"
+                )
+            return system_device
+        except Exception as e:
+            # Fall back to physical GPU 0 so thermal protection still works on
+            # single-GPU systems and when CUDA cannot be initialized here.
+            if cuda_ordinal is not None:
+                print(
+                    f"Warning: Could not resolve thermal_device={cuda_ordinal} "
+                    f"to an NVML device ({e}). Falling back to physical GPU 0."
+                )
+            self._current_cuda_device_id = None
+            try:
+                return system.Device(index=0)
+            except Exception:
+                return None
+
+    def _refresh_device_for_current_cuda_context(self) -> None:
+        """Refresh monitored NVML device when current CUDA device changes.
+
+        This applies only when ``thermal_device`` is not explicitly set. It keeps
+        Thermovision aligned with runtime CUDA context switches in user code
+        (for example via ``torch.cuda.set_device``).
+        """
+        if self.thermal_device is not None:
+            return
+
+        try:
+            cuda_device = CudaDevice(None)
+            current_cuda_id = cuda_device.device_id
+            if (
+                self.device is not None
+                and self._current_cuda_device_id == current_cuda_id
+            ):
+                return
+
+            self.device = self._map_cuda_to_system_device(cuda_device)
+            if self.verbose:
+                print(
+                    f"[Thermovision] Switched monitoring to CUDA device {self._current_cuda_device_id} "
+                    f"(NVML index {self.device.index}, {self.device.name}, UUID {self.device.uuid})"
+                )
+        except Exception:
+            # Keep existing device if refresh fails. If no device exists yet,
+            # try the regular resolution path (including fallback).
+            if self.device is None:
+                self.device = self._resolve_device()
+
     def throttle_guard(self) -> None:
         """Check thermal state and pause if GPU is too hot.
 
@@ -129,6 +232,7 @@ def throttle_guard(self) -> None:
            - Many iterations (>TARGET_MAX_ITERATIONS): GPU heats slowly → decrease thermal_cont (cool less)
         4. Wait until GPU cools back to thermal_cont, then repeat
         """
+        self._refresh_device_for_current_cuda_context()
         tlimit = self._get_gpu_tlimit()
         if tlimit is None:
             return

tests/test_thermovision.py

@@ -1,7 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from unittest.mock import patch
+import os
+from typing import Any
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
@@ -186,6 +188,121 @@ def mock_temp() -> int:
     ), f"Expected thermal_cont=17, got {controller.thermal_cont}"
 
 
+def test_resolve_device_uses_current_cuda_device() -> None:
+    """Without thermal_device, the current CUDA device is monitored (honors CUDA_VISIBLE_DEVICES)."""
+    from nsight.thermovision import ThermalController
+
+    controller = ThermalController(thermal_mode="auto", verbose=False)
+
+    system_device = MagicMock(name="system_device")
+    cuda_device = MagicMock(name="cuda_device")
+    cuda_device.to_system_device.return_value = system_device
+    cuda_device_cls = MagicMock(return_value=cuda_device)
+
+    with patch("nsight.thermovision.CudaDevice", cuda_device_cls):
+        controller.init()
+
+    # CudaDevice(None) selects the CUDA runtime's current device.
+    cuda_device_cls.assert_called_once_with(None)
+    cuda_device.to_system_device.assert_called_once_with()
+    assert controller.device is system_device
+
+
+def test_resolve_device_honors_explicit_thermal_device() -> None:
+    """An explicit thermal_device ordinal selects that CUDA device for monitoring."""
+    from nsight.thermovision import ThermalController
+
+    controller = ThermalController(thermal_mode="auto", thermal_device=3, verbose=False)
+
+    system_device = MagicMock(name="system_device")
+    cuda_device = MagicMock(name="cuda_device")
+    cuda_device.to_system_device.return_value = system_device
+    cuda_device_cls = MagicMock(return_value=cuda_device)
+
+    with patch("nsight.thermovision.CudaDevice", cuda_device_cls):
+        controller.init()
+
+    cuda_device_cls.assert_called_once_with(3)
+    assert controller.device is system_device
+
+
+def test_resolve_device_falls_back_to_physical_gpu_zero() -> None:
+    """If CUDA-to-NVML mapping fails, monitoring falls back to physical GPU 0."""
+    from nsight.thermovision import ThermalController
+
+    controller = ThermalController(thermal_mode="auto", verbose=False)
+
+    fallback_device = MagicMock(name="fallback_device")
+    cuda_device_cls = MagicMock(side_effect=RuntimeError("no cuda"))
+
+    with (
+        patch("nsight.thermovision.CudaDevice", cuda_device_cls),
+        patch(
+            "nsight.thermovision.system.Device", return_value=fallback_device
+        ) as mock_system_device,
+    ):
+        controller.init()
+
+    mock_system_device.assert_called_once_with(index=0)
+    assert controller.device is fallback_device
+
+
+def test_throttle_guard_refreshes_when_current_cuda_device_changes() -> None:
+    """Without explicit thermal_device, guard refresh follows CUDA device switches."""
+    from nsight.thermovision import ThermalController
+
+    controller = ThermalController(thermal_mode="auto", verbose=False)
+
+    system_device_0 = MagicMock(name="system_device_0")
+    system_device_1 = MagicMock(name="system_device_1")
+
+    cuda_device_0 = MagicMock(name="cuda_device_0")
+    cuda_device_0.device_id = 0
+    cuda_device_0.to_system_device.return_value = system_device_0
+
+    cuda_device_1 = MagicMock(name="cuda_device_1")
+    cuda_device_1.device_id = 1
+    cuda_device_1.to_system_device.return_value = system_device_1
+
+    with (
+        patch(
+            "nsight.thermovision.CudaDevice",
+            side_effect=[cuda_device_0, cuda_device_1, cuda_device_1],
+        ),
+        patch(
+            "nsight.thermovision.ThermalController._get_gpu_tlimit", return_value=100
+        ),
+    ):
+        controller.init()
+        assert controller.device is system_device_0
+        controller.throttle_guard()
+
+    assert controller.device is system_device_1
+
+
+def test_throttle_guard_keeps_explicit_thermal_device_binding() -> None:
+    """With explicit thermal_device, guard should not refresh from current device."""
+    from nsight.thermovision import ThermalController
+
+    controller = ThermalController(thermal_mode="auto", thermal_device=2, verbose=False)
+
+    system_device = MagicMock(name="system_device")
+    cuda_device = MagicMock(name="cuda_device")
+    cuda_device.to_system_device.return_value = system_device
+
+    with (
+        patch("nsight.thermovision.CudaDevice", return_value=cuda_device) as mock_cuda,
+        patch(
+            "nsight.thermovision.ThermalController._get_gpu_tlimit", return_value=100
+        ),
+    ):
+        controller.init()
+        controller.throttle_guard()
+
+    assert mock_cuda.call_count == 1
+    assert controller.device is system_device
+
+
 def test_fixed_mode_no_adaptation() -> None:
     """Test that fixed mode does not adapt thermal_cont."""
     from nsight.thermovision import ThermalController