br: get_devie_and_memory_allocated() move to torch.py [skip ci]"

Signed-off-by: Brian Roland <broland@nvidia.com>

Author: broland-hat

sub-packages/bionemo-evo2/tests/bionemo/evo2/conftest.py

@@ -20,7 +20,8 @@
 import pytest
 import torch
 
-from bionemo.testing.test_tools import get_device_and_memory_usage
+from bionemo.testing.torch import get_device_and_memory_allocated
+
 
 def pytest_sessionstart(session):
     """Called at the start of the test session."""
@@ -29,10 +30,9 @@ def pytest_sessionstart(session):
         print(
             f"""
             sub-packages/bionemo-evo2/tests/bionemoe/evo2: Starting test session
-            {get_device_and_memory_usage()}
+            {get_device_and_memory_allocated()}
             """
         )
-        
 
 
 def pytest_sessionfinish(session, exitstatus):
@@ -43,19 +43,10 @@ def pytest_sessionfinish(session, exitstatus):
         print(
             f"""
             sub-packages/bionemo-evo2/tests/bionemoe/evo2: Test session complete
-            current_device={torch.cuda.current_device()}
-            current device name: {torch.cuda.get_device_name(torch.cuda.current_device())}
-            gb_available={torch.cuda.mem_get_info()[0] / 1024**3}
-            Peak GPU memory: {peak_memory / 1024**3:.3f} GB
-            Final GPU memory: {final_memory / 1024**3:.3f} GB
+            {get_device_and_memory_allocated()}
             """
         )
 
- 
-
-     
- 
-        
 
 @pytest.fixture(autouse=True)
 def cleanup_after_test():

sub-packages/bionemo-testing/src/bionemo/testing/test_tools.py

@@ -61,3 +61,20 @@ def recursive_assert_approx_equal(x, y, atol=1e-4, rtol=1e-4):
             recursive_assert_approx_equal(x[key], y[key], atol=atol, rtol=rtol)
     else:
         assert x == y
+
+
+def get_device_and_memory_allocated() -> str:
+    """Get the current device index, name, and memory usage."""
+    current_device_index = torch.cuda.current_device()
+    props = torch.cuda.get_device_properties(current_device_index)
+    message = (
+        f"""
+        current device index: {current_device_index}
+        current device uuid: {props.uuid}
+        current device name: {props.name}
+        memory available: {torch.cuda.mem_get_info()[0] / 1024**3:.3f} GB
+        memory allocated: {torch.cuda.memory_allocated() / 1024**3:.3f} GB
+        max memory allocated: {torch.cuda.max_memory_allocated() / 1024**3:.3f} GB
+        """
+    )
+    return message