Merge pull request #2 from thunderock/add_gpu_tests

Add gpu tests

Author: thunderock

.pre-commit-config.yaml

@@ -18,3 +18,4 @@ repos:
     rev: v0.0.291
     hooks:
     -   id: ruff
+        args: [ --fix ]

Makefile

@@ -30,4 +30,11 @@ format:
 	@echo "formatting..."
 	@poetry install --only lint
 	@poetry run black .
-	@poetry run pre-commit run --all-files
+	# ruff fix
+	@poetry run pre-commit run --all-files --config .pre-commit-config.yaml
+
+.PHONY: run_tests
+run_tests:
+	@echo "running tests..."
+	@poetry install --only main --only test -vvv
+	@poetry run pytest -q tests

graph_ml/utils/config.py

@@ -0,0 +1,15 @@
+import os
+import torch
+
+
+OS = os.name
+DEVICE_TYPE = "cpu"
+if OS == "posix":
+    if torch.cuda.is_available():
+        DEVICE_TYPE = "cuda"
+    elif torch.backends.mps.is_available():
+        DEVICE_TYPE = "mps"
+
+DEVICE = torch.device(DEVICE_TYPE)
+
+GPU_AVAILABLE = DEVICE_TYPE in ["cuda", "mps"]

graph_ml/utils/gpu_utils.py

@@ -64,3 +64,8 @@ pytest = "^8.0.2"
 [tool.poetry.group.lint.dependencies]
 pre-commit = "^3.6.2"
 black = "^24.2.0"
+
+
+
+# poetry add git+ssh://git@github.com/thunderock/graph_ml.git
+# poetry add whl_url

pyproject.toml

@@ -0,0 +1,25 @@
+from __future__ import annotations
+import os
+
+from graph_ml.utils import config
+
+OS = os.name
+
+
+def test_target_os():
+    assert OS == "posix"
+
+
+def test_device_type():
+    assert config.DEVICE_TYPE in ["cpu", "cuda", "mps"]
+
+
+def test_gpu_available():
+    if config.DEVICE_TYPE == "cuda":
+        assert config.GPU_AVAILABLE
+    elif config.DEVICE_TYPE == "mps":
+        assert config.GPU_AVAILABLE
+    elif config.DEVICE_TYPE == "cpu":
+        assert not config.GPU_AVAILABLE
+    else:
+        assert False