Address issues #255 and #295

brian2cuda/cuda_prefs.py

@@ -205,6 +205,15 @@ def validate_bundle_size_expression(string):
         validator=lambda v: v is None or isinstance(v, int)
     ),
 
+    gpu_selection_strategy=BrianPreference(
+        docs='''Strategy used to select a GPU automatically when ``gpu_id`` is not set.
+        ``performance`` uses an ``nvidia-smi``-based estimate derived from
+        ``free_memory_mb * compute_capability * (1 - utilization_gpu / 100)``.
+        ``legacy`` selects the GPU with the highest compute capability.''',
+        default='performance',
+        validator=lambda v: v in ['legacy', 'performance']
+    ),
+
     extra_compile_args_nvcc=BrianPreference(
         docs='Extra compile arguments (a list of strings) to pass to the nvcc compiler.',
         default=['-w', '-use_fast_math']

brian2cuda/tests/test_gpu_detection.py

@@ -16,6 +16,8 @@
     reset_gpu_selection,
     get_gpu_selection,
     restore_gpu_selection,
+    _parse_nvidia_smi_gpu_metrics,
+    get_best_gpu,
 )
 
 # Only catch our own log messages
@@ -177,3 +179,64 @@ def test_no_gpu_detection_preference(reset_gpu_detection, use_default_prefs):
     prefs.devices.cuda_standalone.cuda_backend.gpu_id = 0
     prefs.devices.cuda_standalone.cuda_backend.compute_capability = device.minimal_compute_capability
     run(0*ms)
+
+
+def test_parse_nvidia_smi_gpu_metrics():
+    parsed = _parse_nvidia_smi_gpu_metrics(
+        "0, 24564, 20000, 5\n1, 24564, 1000, 97\n"
+    )
+    assert parsed == [
+        {
+            "gpu_id": 0,
+            "memory_total_mb": 24564.0,
+            "memory_free_mb": 20000.0,
+            "utilization_percent": 5.0,
+        },
+        {
+            "gpu_id": 1,
+            "memory_total_mb": 24564.0,
+            "memory_free_mb": 1000.0,
+            "utilization_percent": 97.0,
+        },
+    ]
+
+
+def test_performance_gpu_selection_prefers_higher_cuda_performance(monkeypatch, use_default_prefs):
+    prefs.devices.cuda_standalone.cuda_backend.gpu_selection_strategy = "performance"
+
+    monkeypatch.setattr(
+        "brian2cuda.utils.gputools.get_available_gpus",
+        lambda: ["GPU0", "GPU1"],
+    )
+    monkeypatch.setattr(
+        "brian2cuda.utils.gputools.get_gpu_performance",
+        lambda gpu_id: {
+            0: {"compute_capability": 8.6, "performance": 20000.0 * 8.6 * 0.95},
+            1: {"compute_capability": 9.0, "performance": 1000.0 * 9.0 * 0.03},
+        }[gpu_id],
+    )
+
+    gpu_id, compute_capability = get_best_gpu()
+    assert gpu_id == 0
+    assert compute_capability == 8.6
+
+
+def test_performance_gpu_selection_falls_back_to_legacy(monkeypatch, use_default_prefs):
+    prefs.devices.cuda_standalone.cuda_backend.gpu_selection_strategy = "performance"
+
+    monkeypatch.setattr(
+        "brian2cuda.utils.gputools.get_available_gpus",
+        lambda: ["GPU0", "GPU1"],
+    )
+    monkeypatch.setattr(
+        "brian2cuda.utils.gputools.get_compute_capability",
+        lambda gpu_id: {0: 7.0, 1: 8.0}[gpu_id],
+    )
+    monkeypatch.setattr(
+        "brian2cuda.utils.gputools.get_gpu_performance",
+        lambda gpu_id: (_ for _ in ()).throw(RuntimeError("nvidia-smi unavailable")),
+    )
+
+    gpu_id, compute_capability = get_best_gpu()
+    assert gpu_id == 1
+    assert compute_capability == 8.0

brian2cuda/tools/test_suite/run_test_suite.py

@@ -54,8 +54,9 @@
 from io import StringIO
 
 import brian2
-from brian2 import test, prefs
+from brian2 import prefs
 import brian2cuda
+from brian2cuda.tests import run as run_brian2cuda_tests
 
 bot = None
 if args.notify_slack:
@@ -180,17 +181,19 @@
             buffer.add(print_lines)
             buffer.print_all()
 
-        success = test(codegen_targets=[],
-                       long_tests=args.no_long_tests,
-                       test_codegen_independent=False,
-                       test_standalone=target,
-                       reset_preferences=False,
-                       fail_for_not_implemented=args.fail_not_implemented,
-                       test_in_parallel=test_in_parallel,
-                       extra_test_dirs=extra_test_dirs,
-                       float_dtype=None,
-                       additional_args=additional_args,
-                       build_options=build_options)
+        success = run_brian2cuda_tests(
+            test_standalone=[target],
+            long_tests=args.no_long_tests,
+            reset_preferences=False,
+            fail_for_not_implemented=args.fail_not_implemented,
+            test_in_parallel=bool(test_in_parallel),
+            build_options=build_options,
+            extra_test_dirs=extra_test_dirs,
+            float_dtype=None,
+            quiet=args.quiet,
+            debug=args.debug,
+            additional_args=additional_args,
+        )
 
         successes.append(success)
 

brian2cuda/utils/gputools.py

@@ -38,6 +38,57 @@
 }
 
 
+def _parse_nvidia_smi_gpu_metrics(csv_output):
+    metrics = []
+    for line in csv_output.splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        index_str, memory_total_str, memory_free_str, utilization_str = [
+            part.strip() for part in line.split(",")
+        ]
+        metrics.append(
+            {
+                "gpu_id": int(index_str),
+                "memory_total_mb": float(memory_total_str),
+                "memory_free_mb": float(memory_free_str),
+                "utilization_percent": float(utilization_str),
+            }
+        )
+    return metrics
+
+
+def _parse_device_query_device_blocks(device_query_output):
+    blocks = {}
+    current_gpu_id = None
+    current_lines = []
+
+    for line in device_query_output.splitlines():
+        if line.startswith("Device "):
+            if current_gpu_id is not None:
+                blocks[current_gpu_id] = current_lines
+            current_gpu_id = int(line.split(":")[0].split()[1])
+            current_lines = [line]
+        elif current_gpu_id is not None:
+            current_lines.append(line)
+
+    if current_gpu_id is not None:
+        blocks[current_gpu_id] = current_lines
+
+    return blocks
+
+
+def _parse_device_query_compute_capability(block_lines):
+    for line in block_lines:
+        stripped = line.strip()
+        if stripped.startswith("CUDA Capability Major/Minor version number"):
+            major = int(stripped[-3])
+            minor = int(stripped[-1])
+            return major + 0.1 * minor
+
+    raise RuntimeError("Could not parse compute capability from `deviceQuery` output.")
+
+
 def get_cuda_path():
     """
     Detect the path to the CUDA installation (e.g. '/usr/local/cuda'). This takes into
@@ -498,37 +549,87 @@ def _get_compute_capability_with_device_query(gpu_id):
                 "`prefs.devices.cuda_standalone.cuda_backend.device_query_path`"
             )
     device_query_output = _run_command_with_output(device_query_path)
-    lines = device_query_output.split("\n")
-    compute_capability = None
-    for i, line in enumerate(lines):
-        if line.startswith("Device "):
-            # example line:
-            # `Device 0: "GeForce MX150"`
-            this_gpu_id = int(line[7])  # "Device i ..."  <- i in position 7
-            if this_gpu_id == gpu_id:
-                # Get GPU name: word in quotation
-                gpu_name = re.findall(r'\"(.+?)\"', line)[0]
-                # Make sure we got the right GPU here
-                assert gpu_list[gpu_id] == gpu_name
-                # The compute capability is shown 2 lines after the "Device ..." line
-                # Example line:
-                # `  CUDA Capability Major/Minor version number:    6.1`
-                compute_capability_line = lines[i + 2]
-                assert compute_capability_line.strip().startswith(
-                    "CUDA Capability Major/Minor version number"
-                ), f"Unexpected line parsed: {compute_capability_line}"
-                # Last 3 chars are the compute capability
-                major = int(compute_capability_line[-3])
-                minor = int(compute_capability_line[-1])
-                # Turn into float
-                compute_capability = major + 0.1 * minor
-    return compute_capability
+    blocks = _parse_device_query_device_blocks(device_query_output)
+    try:
+        block_lines = blocks[gpu_id]
+    except KeyError as error:
+        raise RuntimeError(
+            f"Could not find GPU {gpu_id} in `deviceQuery` output."
+        ) from error
+
+    gpu_name = re.findall(r'\"(.+?)\"', block_lines[0])[0]
+    assert gpu_list[gpu_id] == gpu_name
+    return _parse_device_query_compute_capability(block_lines)
+
+
+def _get_nvidia_smi_gpu_metrics():
+    command = (
+        "nvidia-smi --query-gpu=index,memory.total,memory.free,utilization.gpu "
+        "--format=csv,noheader,nounits"
+    )
+    try:
+        return _parse_nvidia_smi_gpu_metrics(
+            _run_command_with_output(command)
+        )
+    except (RuntimeError, FileNotFoundError, ValueError) as error:
+        raise RuntimeError(
+            "Failed to query GPU memory/utilization metrics with `nvidia-smi`."
+        ) from error
+
+
+def get_gpu_performance(gpu_id):
+    """
+    Estimate GPU suitability from free memory, GPU utilization, and compute capability.
+    """
+    gpu_metrics = _get_nvidia_smi_gpu_metrics()
+    try:
+        gpu_metric = next(
+            metric
+            for metric in gpu_metrics
+            if metric["gpu_id"] == gpu_id
+        )
+    except StopIteration as error:
+        raise RuntimeError(
+            f"Could not find GPU {gpu_id} in `nvidia-smi` metrics output."
+        ) from error
+
+    compute_capability = get_compute_capability(gpu_id)
+    memory_free_mb = gpu_metric["memory_free_mb"]
+    utilization_percent = gpu_metric["utilization_percent"]
+    performance = (
+        memory_free_mb
+        * compute_capability
+        * (1.0 - utilization_percent / 100.0)
+    )
+    return {
+        "compute_capability": compute_capability,
+        "memory_free_mb": memory_free_mb,
+        "utilization_percent": utilization_percent,
+        "performance": performance,
+    }
 
 
 def get_best_gpu():
     """
-    Get the "best" GPU available. This currently chooses the GPU with highest compute
-    capability and lowest GPU ID (as reported by `nvidia-smi`)
+    Get the "best" GPU available according to the configured selection strategy.
+    """
+    strategy = prefs.devices.cuda_standalone.cuda_backend.gpu_selection_strategy
+    if strategy == "legacy":
+        return get_best_gpu_legacy()
+
+    try:
+        return get_best_gpu_by_performance()
+    except (RuntimeError, FileNotFoundError, ValueError) as error:
+        logger.warn(
+            "Performance-based GPU selection failed, falling back to legacy "
+            f"selection: {error}"
+        )
+        return get_best_gpu_legacy()
+
+
+def get_best_gpu_legacy():
+    """
+    Get the GPU with highest compute capability and lowest GPU ID.
     """
     gpu_list = get_available_gpus()
     best_gpu_id = 0
@@ -542,8 +643,34 @@ def get_best_gpu():
     return best_gpu_id, best_compute_capability
 
 
+def get_best_gpu_by_performance():
+    """
+    Select the GPU with the highest nvidia-smi-based performance estimate.
+    """
+    gpu_list = get_available_gpus()
+    best_gpu_id = 0
+    best_compute_capability = 0
+    best_performance = -1
+
+    for gpu_id, _ in enumerate(gpu_list):
+        metrics = get_gpu_performance(gpu_id)
+        performance = metrics["performance"]
+        compute_capability = metrics["compute_capability"]
+        if (
+            performance > best_performance
+            or (
+                performance == best_performance
+                and compute_capability > best_compute_capability
+            )
+        ):
+            best_gpu_id = gpu_id
+            best_compute_capability = compute_capability
+            best_performance = performance
+
+    return best_gpu_id, best_compute_capability
+
+
 if __name__ == "__main__":
     print(get_best_gpu())
     #a = nvidia_smi()
     #print(a)
-

pyproject.toml

@@ -44,7 +44,7 @@ version_scheme = 'post-release'
 local_scheme = 'no-local-version'
 write_to = 'brian2cuda/_version.py'
 tag_regex = '^(?P<version>v\d+(?:\.\d+){0,2}[^\+]*(?:\+.*)?)$'
-fallback_version = 'unknown'
+fallback_version = '0.0'
 
 [build-system]
 requires = [