[benchmarking] Add configurable GPU memory warning threshold

benchmarking/nightly-benchmark.yaml

@@ -133,6 +133,10 @@ delete_scratch: true
 # for "ray start" to determine object store size.
 object_store_size: 536870912000  # 500GB
 
+# Fraction (0.0-1.0) of GPU memory that may be in use without emitting a
+# warning. Checked both before and after each benchmark run.
+gpu_mem_use_warning_threshold: 0.01
+
 # Global ray settings inherited by all entries. Per-entry ray sections override these values.
 ray:
   num_cpus: 64

benchmarking/run.py

@@ -205,7 +205,12 @@ def run_entry(
         ray_cluster_data = get_ray_cluster_data()
         gpu_stats_before = get_gpu_stats()
         logger.info("\tGPU stats (before):")
-        log_gpu_stats(gpu_stats_before, warn_if_in_use=True)
+        warnings = log_gpu_stats(
+            gpu_stats_before,
+            warn_if_in_use=True,
+            warning_threshold=entry.gpu_mem_use_warning_threshold,
+            warning_threshold_msg="used before benchmark started",
+        )
         logger.info(f"\tRunning command {' '.join(cmd) if isinstance(cmd, list) else cmd}")
         started_exec = time.time()
         run_data = run_command_with_timeout(
@@ -217,7 +222,12 @@ def run_entry(
         )
         ended_exec = time.time()
         logger.info("\tGPU stats (after):")
-        log_gpu_stats(get_gpu_stats())
+        warnings += log_gpu_stats(
+            get_gpu_stats(),
+            warn_if_in_use=True,
+            warning_threshold=entry.gpu_mem_use_warning_threshold,
+            warning_threshold_msg="left in use after benchmark ended",
+        )
         duration = ended_exec - started_exec
 
         # Update result_data
@@ -231,6 +241,7 @@ def run_entry(
                 "logs_dir": logs_path,
                 "ray_cluster_data": ray_cluster_data,
                 "gpu_stats": gpu_stats_before,
+                "warnings": warnings,
             }
         )
         # script_persisted_data is a dictionary with keys "params" and "metrics"

benchmarking/runner/entry.py

@@ -47,13 +47,23 @@ class Entry:
     object_store_size: int | float | str | None = None
     # If set, overrides the session-level delete_scratch setting for this entry
     delete_scratch: bool | None = None
+    # If set, overrides the session-level gpu_mem_use_warning_threshold for this entry
+    gpu_mem_use_warning_threshold: float | None = None
 
     def __post_init__(self) -> None:  # noqa: C901, PLR0912
         """Post-initialization checks and updates for dataclass."""
         # Process object_store_size by converting values representing fractions of system memory to bytes.
         if isinstance(self.object_store_size, float):
             self.object_store_size = int(get_total_memory_bytes() * self.object_store_size)
 
+        # Validate the warning threshold range, if set.
+        if self.gpu_mem_use_warning_threshold is not None and not (0 <= self.gpu_mem_use_warning_threshold <= 1):
+            msg = (
+                f"Invalid gpu_mem_use_warning_threshold for entry '{self.name}': "
+                f"{self.gpu_mem_use_warning_threshold}; must be between 0 and 1 inclusive."
+            )
+            raise ValueError(msg)
+
         # Convert the sink_data list of dicts to a dict of dicts for easier lookup with key from "name".
         # sink_data typically starts as a list of dicts from reading YAML, like this:
         # sink_data:

benchmarking/runner/session.py

@@ -46,12 +46,16 @@ class Session:
     object_store_size: int | float | str | None = 0.5
     # Whether to delete the entry's scratch directory after completion by default
     delete_scratch: bool = True
+    # Fraction of total GPU memory (0.0-1.0) above which a warning is emitted, both
+    # before and after each benchmark run. If None, any usage > 0 triggers a warning.
+    # Entries can override this value.
+    gpu_mem_use_warning_threshold: float | None = None
     # Global ray settings inherited by all entries; per-entry ray sections override these values.
     ray: dict = field(default_factory=dict)
     path_resolver: PathResolver = None
     dataset_resolver: DatasetResolver = None
 
-    def __post_init__(self) -> None:
+    def __post_init__(self) -> None:  # noqa: C901
         """Post-initialization checks and updates for dataclass."""
         names = [entry.name for entry in self.entries]
         if len(names) != len(set(names)):
@@ -63,6 +67,14 @@ def __post_init__(self) -> None:
         if isinstance(self.object_store_size, float):
             self.object_store_size = int(get_total_memory_bytes() * self.object_store_size)
 
+        # Validate the session-level warning threshold range, if set.
+        if self.gpu_mem_use_warning_threshold is not None and not (0 <= self.gpu_mem_use_warning_threshold <= 1):
+            msg = (
+                f"Invalid session-level gpu_mem_use_warning_threshold: "
+                f"{self.gpu_mem_use_warning_threshold}; must be between 0 and 1 inclusive."
+            )
+            raise ValueError(msg)
+
         # Update delete_scratch for each entry that has not been set to the session-level delete_scratch setting
         for entry in self.entries:
             if entry.delete_scratch is None:
@@ -78,6 +90,11 @@ def __post_init__(self) -> None:
             if entry.object_store_size is None:
                 entry.object_store_size = self.object_store_size
 
+        # Update gpu_mem_use_warning_threshold for each entry that has not been set.
+        for entry in self.entries:
+            if entry.gpu_mem_use_warning_threshold is None:
+                entry.gpu_mem_use_warning_threshold = self.gpu_mem_use_warning_threshold
+
         # Apply global ray defaults to each entry, with per-entry ray values taking precedence.
         for entry in self.entries:
             entry.ray = {**self.ray, **entry.ray}

benchmarking/runner/sinks/slack_sink.py

@@ -575,15 +575,7 @@ def register_benchmark_entry_finished(self, result_dict: dict[str, Any], benchma
         pings = [] if result_dict["success"] else sink_data.get("ping_on_failure", [])
         status_text = "✅ success" if result_dict["success"] else "❌ FAILED"
 
-        # Warn if any GPU had memory in use before the benchmark started.
-        # TODO: This could be made into a more generic check that can be configured in the sink config.
-        warnings = []
-        for gpu_id, stats in result_dict.get("gpu_stats", {}).items():
-            if stats.get("memory_used", 0) > 0:
-                pct_used = stats["memory_used"] / stats["memory_total"] * 100
-                warnings.append(
-                    f"GPU {gpu_id} had {stats['memory_used']} MiB ({pct_used:.1f}% of total) used before benchmark started"
-                )
+        warnings = result_dict.get("warnings", [])
 
         # Create a new message for the entry to post in the thread.
         msg = self._create_benchmark_entry_message(

benchmarking/runner/utils.py

@@ -218,20 +218,39 @@ def get_gpu_stats() -> dict:
     return query_data
 
 
-def log_gpu_stats(gpu_stats: dict, warn_if_in_use: bool = False) -> None:
+def log_gpu_stats(
+    gpu_stats: dict,
+    warn_if_in_use: bool = False,
+    warning_threshold: float | None = None,
+    warning_threshold_msg: str = "still in use",
+) -> list[str]:
     """Log GPU memory usage for each GPU as a percentage of total memory.
 
     Args:
         gpu_stats: Dictionary as returned by get_gpu_stats().
-        warn_if_in_use: If True, emit a warning for any GPU with memory_used > 0.
+        warn_if_in_use: If True, emit a warning for any GPU that exceeds the usage threshold.
+        warning_threshold: Fraction of total GPU memory (0.0-1.0) above which a warning is
+            emitted. If None and warn_if_in_use is True, any usage > 0 triggers a warning.
+        warning_threshold_msg: Trailing context phrase appended to each warning message
+            (e.g. "used before benchmark started"). Defaults to "still in use".
+
+    Returns:
+        List of warning strings for any GPUs that triggered a warning.
     """
+    warnings = []
     for gpu_id, stats in gpu_stats.items():
         pct_used = stats["memory_used"] / stats["memory_total"] * 100
         logger.info(f"GPU {gpu_id} : {pct_used:.1f}%")
-        if warn_if_in_use and stats["memory_used"] > 0:
-            logger.warning(
-                f"GPU {gpu_id} has {stats['memory_used']} MiB ({pct_used:.1f}% of total) used before benchmark started"
+        if warn_if_in_use:
+            fraction_used = stats["memory_used"] / stats["memory_total"]
+            threshold_exceeded = (
+                fraction_used > warning_threshold if warning_threshold is not None else stats["memory_used"] > 0
             )
+            if threshold_exceeded:
+                msg = f"GPU {gpu_id}: {stats['memory_used']} MiB ({pct_used:.1f}% of total) {warning_threshold_msg}"
+                logger.warning(msg)
+                warnings.append(msg)
+    return warnings
 
 
 _LEGACY_PATH_FIELDS = ["results_path", "datasets_path", "model_weights_path"]