Add the ability to pass UUID prefixes as a way of specifying what GPUs to track..

This is to address #873
Prior to this change you could only pass an index into the number of GPUs on the system.
Now you can pass a UUID prefix, including the 'MIG-' prefix per
https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#cuda-environment-variables
if desired.

Note that I have not been able to test this on real life repo. The reporter of
#873 was not able to provide a repro.

Author: cianc

codecarbon/core/config.py

@@ -1,7 +1,7 @@
 import configparser
 import os
 from pathlib import Path
-from typing import List
+from typing import List, Union
 
 from codecarbon.external.logger import logger
 
@@ -44,24 +44,33 @@ def parse_env_config() -> dict:
     }
 
 
-def parse_gpu_ids(gpu_ids_str: str) -> List[int]:
+def parse_gpu_ids(gpu_ids: Union[str, List[int]]) -> List[str]:
     """
-    Transforms the potential gpu_ids string into a list of int values
+    Transforms the potential gpu_ids into a list of string id values.
+
 
     Args:
-        gpu_ids_str (str): The config file or environment variable value for `gpu_ids`
-        which is read as a string and should be parsed into a list of ints
+        gpu_ids: The config file or environment variable value for `gpu_ids`
 
     Returns:
-        list[int]: The list of GPU ids available declared by the user.
+        list[str]: The list of GPU ids available.
             Potentially empty.
     """
-    if not isinstance(gpu_ids_str, str):
-        return gpu_ids_str
-
-    gpu_ids_str = "".join(c for c in gpu_ids_str if (c.isalnum() or c == ","))
-    str_ids = [gpu_id for gpu_id in gpu_ids_str.split(",") if gpu_id]
-    return list(map(int, str_ids))
+    if isinstance(gpu_ids, str):
+        # Allow '-' in id strings since UUIDs may include them.
+        gpu_ids = "".join(c for c in gpu_ids if (c.isalnum() or c in ("-", ",")))
+        str_ids = [gpu_id for gpu_id in gpu_ids.split(",") if gpu_id]
+        return str_ids
+
+    elif isinstance(gpu_ids, list) and all(
+        isinstance(gpu_id, int) for gpu_id in gpu_ids
+    ):
+        return list(map(str, gpu_ids))
+
+    else:
+        logger.warning(
+            "Invalid gpu_ids format. Expected a string or a list of ints/strings."
+        )
 
 
 def get_hierarchical_config():

codecarbon/core/resource_tracker.py

@@ -183,18 +183,11 @@ def set_CPU_tracking(self):
     def set_GPU_tracking(self):
         logger.info("[setup] GPU Tracking...")
         if self.tracker._gpu_ids:
-            # If _gpu_ids is a string or a list of int, parse it to a list of ints
-            if isinstance(self.tracker._gpu_ids, str) or (
-                isinstance(self.tracker._gpu_ids, list)
-                and all(isinstance(gpu_id, int) for gpu_id in self.tracker._gpu_ids)
-            ):
-                self.tracker._gpu_ids: List[int] = parse_gpu_ids(self.tracker._gpu_ids)
+            self.tracker._gpu_ids = parse_gpu_ids(self.tracker._gpu_ids)
+            if self.tracker._gpu_ids:
                 self.tracker._conf["gpu_ids"] = self.tracker._gpu_ids
                 self.tracker._conf["gpu_count"] = len(self.tracker._gpu_ids)
-            else:
-                logger.warning(
-                    "Invalid gpu_ids format. Expected a string or a list of ints."
-                )
+
         if gpu.is_gpu_details_available():
             logger.info("Tracking Nvidia GPU via pynvml")
             gpu_devices = GPU.from_utils(self.tracker._gpu_ids)

codecarbon/external/hardware.py

@@ -102,25 +102,40 @@ def _get_gpu_ids(self) -> Iterable[int]:
         Get the Ids of the GPUs that we will monitor
         :return: list of ids
         """
-        gpu_ids = []
         if self.gpu_ids is not None:
-            # Check that the provided GPU ids are valid
-            if not set(self.gpu_ids).issubset(set(range(self.num_gpus))):
-                logger.warning(
-                    f"Unknown GPU ids {gpu_ids}, only {self.num_gpus} GPUs available."
-                )
-            # Keep only the GPUs that are in the provided list
-            for gpu_id in range(self.num_gpus):
-                if gpu_id in self.gpu_ids:
-                    gpu_ids.append(gpu_id)
+            uuids_to_ids = {
+                gpu.get("uuid"): gpu.get("gpu_index")
+                for gpu in self.devices.get_gpu_static_info()
+            }
+            monitored_gpu_ids = []
+
+            for gpu_id in self.gpu_ids:
+                found_gpu_id = False
+                # Is it an index into the number of GPUs on the system?
+                if isinstance(gpu_id, int) or gpu_id.isdigit():
+                    gpu_id = int(gpu_id)
+                    if 0 <= gpu_id < self.num_gpus:
+                        monitored_gpu_ids.append(gpu_id)
+                        found_gpu_id = True
+                # Check if it matches a prefix of any UUID on the system after stripping any 'MIG-'
+                # id prefix per https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#cuda-environment-variables
                 else:
-                    logger.info(
-                        f"GPU number {gpu_id} will not be monitored, at your request."
+                    stripped_gpu_id_str = gpu_id.lstrip("MIG-")
+                    for uuid, id in uuids_to_ids.items():
+                        if uuid.startswith(stripped_gpu_id_str):
+                            monitored_gpu_ids.append(id)
+                            found_gpu_id = True
+                            break
+                if not found_gpu_id:
+                    logger.warning(
+                        f"GPU with ID '{gpu_id}' not found or invalid. It will be ignored."
                     )
-            self.gpu_ids = gpu_ids
+
+            monitored_gpu_ids = sorted(list(set(monitored_gpu_ids)))
+            self.gpu_ids = monitored_gpu_ids
+            return monitored_gpu_ids
         else:
-            gpu_ids = set(range(self.num_gpus))
-        return gpu_ids
+            return list(range(self.num_gpus))
 
     def total_power(self) -> Power:
         return self._total_power
@@ -135,7 +150,7 @@ def from_utils(cls, gpu_ids: Optional[List] = None) -> "GPU":
         new_gpu_ids = gpus._get_gpu_ids()
         if len(new_gpu_ids) < gpus.num_gpus:
             logger.warning(
-                f"You have {gpus.num_gpus} GPUs but we will monitor only {len(gpu_ids)} of them. Check your configuration."
+                f"You have {gpus.num_gpus} GPUs but we will monitor only {len(new_gpu_ids)} ({new_gpu_ids}) of them. Check your configuration."
             )
         return cls(gpu_ids=new_gpu_ids)
 

tests/test_config.py

@@ -32,16 +32,16 @@ def test_clean_env_key(self):
 
     def test_parse_gpu_ids(self):
         for ids, target in [
-            ("0,1,2", [0, 1, 2]),
-            ("[0, 1, 2", [0, 1, 2]),
-            ("(0, 1, 2)", [0, 1, 2]),
-            ("[1]", [1]),
-            ("1", [1]),
-            ("0", [0]),
+            ("0,1,2", ["0", "1", "2"]),
+            ("[0, 1, 2", ["0", "1", "2"]),
+            ("(0, 1, 2)", ["0", "1", "2"]),
+            ("[1]", ["1"]),
+            ("1", ["1"]),
+            ("0", ["0"]),
+            ("MIG-f1e", ["MIG-f1e"]),
             ("", []),
             ([], []),
-            ([1, 2, 3], [1, 2, 3]),
-            (1, 1),
+            ([1, 2, 3], ["1", "2", "3"]),
         ]:
             self.assertEqual(parse_gpu_ids(ids), target)
 
@@ -101,6 +101,7 @@ def test_read_confs(self):
             "USER": "useless key",
             "CODECARBON_ENV_OVERWRITE": "SUCCESS:overwritten",
             "CODECARBON_ENV_NEW_KEY": "cool value",
+            "CODECARBON_ALLOW_MULTIPLE_RUNS": "True",
         },
     )
     def test_read_confs_and_parse_envs(self):
@@ -145,9 +146,8 @@ def test_empty_conf(self):
             "builtins.open", new_callable=get_custom_mock_open(global_conf, local_conf)
         ):
             conf = dict(get_hierarchical_config())
-            target = {
-                "allow_multiple_runs": "True"
-            }  # allow_multiple_runs is a default value
+            # allow_multiple_runs is set in pytest.ini and not mocked, so it's visible here.
+            target = {"allow_multiple_runs": "True"}
             self.assertDictEqual(conf, target)
 
     @mock.patch.dict(
@@ -190,7 +190,7 @@ def test_full_hierarchy(self):
             self.assertEqual(tracker._force_ram_power, 50.5)
             self.assertEqual(tracker._output_dir, "/success/overwritten")
             self.assertEqual(tracker._emissions_endpoint, "http://testhost:2000")
-            self.assertEqual(tracker._gpu_ids, [0, 1])
+            self.assertEqual(tracker._gpu_ids, ["0", "1"])
             self.assertEqual(tracker._co2_signal_api_token, "signal-token")
             self.assertEqual(tracker._project_name, "test-project")
             self.assertTrue(tracker._save_to_file)
@@ -206,7 +206,7 @@ def test_gpu_ids_from_env(self):
             tracker = EmissionsTracker(
                 project_name="test-project", allow_multiple_runs=True
             )
-        self.assertEqual(tracker._gpu_ids, [2, 3])
+        self.assertEqual(tracker._gpu_ids, ["2", "3"])
 
     @mock.patch.dict(
         os.environ,
@@ -220,7 +220,7 @@ def test_too_much_gpu_ids_in_env(self):
             tracker = EmissionsTracker(
                 project_name="test-project", allow_multiple_runs=True
             )
-        self.assertEqual(tracker._gpu_ids, [99])
+        self.assertEqual(tracker._gpu_ids, ["99"])
         gpu_count = 0
         for hardware in tracker._hardware:
             if isinstance(hardware, GPU):

tests/test_gpu.py

@@ -50,7 +50,7 @@ def setup_method(self):
         self.DETAILS = {
             "handle_0": {
                 "name": b"GeForce GTX 1080",
-                "uuid": b"uuid#1",
+                "uuid": b"uuid-1",
                 "memory": real_pynvml.c_nvmlMemory_t(1024, 100, 924),
                 "temperature": 75,
                 "power_usage": 26,
@@ -66,7 +66,7 @@ def setup_method(self):
             },
             "handle_1": {
                 "name": b"GeForce GTX 1080",
-                "uuid": b"uuid#2",
+                "uuid": b"uuid-2",
                 "memory": real_pynvml.c_nvmlMemory_t(1024, 200, 824),
                 "temperature": 79,
                 "power_usage": 29,
@@ -84,7 +84,7 @@ def setup_method(self):
         self.expected = [
             {
                 "name": "GeForce GTX 1080",
-                "uuid": "uuid#1",
+                "uuid": "uuid-1",
                 "total_memory": 1024,
                 "free_memory": 100,
                 "used_memory": 924,
@@ -102,7 +102,7 @@ def setup_method(self):
             },
             {
                 "name": "GeForce GTX 1080",
-                "uuid": "uuid#2",
+                "uuid": "uuid-2",
                 "total_memory": 1024,
                 "free_memory": 200,
                 "used_memory": 824,
@@ -146,14 +146,14 @@ def test_static_gpu_info(self):
         expected = [
             {
                 "name": "GeForce GTX 1080",
-                "uuid": "uuid#1",
+                "uuid": "uuid-1",
                 "total_memory": 1024,
                 "power_limit": 149,
                 "gpu_index": 0,
             },
             {
                 "name": "GeForce GTX 1080",
-                "uuid": "uuid#2",
+                "uuid": "uuid-2",
                 "total_memory": 1024,
                 "power_limit": 149,
                 "gpu_index": 1,
@@ -311,6 +311,33 @@ def mock_nvmlDeviceGetTotalEnergyConsumption(handle):
         expected_power = gpu2_power
         tc.assertAlmostEqual(expected_power.kW, gpu.total_power().kW)
 
+    def test_get_gpu_ids(self):
+        """
+        Check parsing of gpu_ids in various forms.
+        """
+        # Prepare
+        from codecarbon.external.hardware import GPU
+
+        for test_ids, expected_ids in [
+            ([0, 1], [0, 1]),
+            ([0, 1, 2], [0, 1]),
+            ([2], []),
+            (["0", "1"], [0, 1]),
+            # Only two GPUS in the system, so ignore the third (index 2)
+            (["0", "1", "2"], [0, 1]),
+            (["2"], []),
+            # Check UUID-to-index mapping
+            (["uuid-1"], [0]),
+            (["uuid-1", "uuid-2"], [0, 1]),
+            (["uuid-3"], []),
+            # Check UUID-to-index mapping when we need to strip the prefix
+            (["MIG-uuid-1"], [0]),
+            (["MIG-uuid-3"], []),
+        ]:
+            gpu = GPU(test_ids)
+            result = gpu._get_gpu_ids()
+            assert result == expected_ids
+
 
 class TestGpuNotAvailable:
     def setup_method(self):