refactor: convert jax verification script to pytest and update docs
This commit is contained in:
parent
741aafab07
commit
cb5a6fbd87
GPG key ID:
38297DE43F75FFE2
3 changed files with 43 additions and 22 deletions
|
|
@ -53,9 +53,9 @@ If you prefer not to use Docker:
|
|||
|
||||
## Hardware Acceleration (JAX)
|
||||
|
||||
Verify your setup by running:
|
||||
Verify your setup by running the JAX initialization test:
|
||||
|
||||
```bash
|
||||
python -c "import jax; print(jax.devices())"
|
||||
uv run pytest tests/test_jax_init.py
|
||||
```
|
||||
In the devcontainer, this will report a `GpuDevice` if a GPU is detected and the `cuda` extra was installed.
|
||||
In the devcontainer, this will succeed on both CPU and GPU. A `GpuDevice` is expected if a GPU is detected and the `cuda` extra was installed.
|
||||
|
|
|
|||
|
|
@ -1,22 +1,7 @@
|
|||
import jax
|
||||
import sys
|
||||
|
||||
def verify_jax():
|
||||
print(f"Python version: {sys.version}")
|
||||
print(f"JAX version: {jax.__version__}")
|
||||
|
||||
|
||||
def test_jax_initializes():
|
||||
"""Verify that JAX initializes and exposes at least one device."""
|
||||
devices = jax.devices()
|
||||
print(f"Available devices: {devices}")
|
||||
|
||||
gpu_found = any(d.device_kind == 'gpu' for d in devices)
|
||||
if gpu_found:
|
||||
print("SUCCESS: GPU detected!")
|
||||
else:
|
||||
print("INFO: Only CPU detected (expected if not in GPU-enabled environment/container).")
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
verify_jax()
|
||||
except Exception as e:
|
||||
print(f"ERROR during JAX initialization: {e}")
|
||||
sys.exit(1)
|
||||
assert len(devices) > 0, "JAX should expose at least one device"
|
||||
|
|
|
|||
36
uv.lock
generated
36
uv.lock
generated
|
|
@ -30,6 +30,7 @@ cuda = [
|
|||
[package.dev-dependencies]
|
||||
dev = [
|
||||
{ name = "pre-commit" },
|
||||
{ name = "pytest" },
|
||||
{ name = "ruff" },
|
||||
]
|
||||
|
||||
|
|
@ -51,6 +52,7 @@ provides-extras = ["cuda"]
|
|||
[package.metadata.requires-dev]
|
||||
dev = [
|
||||
{ name = "pre-commit", specifier = ">=4.0.0" },
|
||||
{ name = "pytest", specifier = ">=8.0.0" },
|
||||
{ name = "ruff", specifier = ">=0.15.2" },
|
||||
]
|
||||
|
||||
|
|
@ -623,6 +625,15 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/a4/ed/1f1afb2e9e7f38a545d628f864d562a5ae64fe6f7a10e28ffb9b185b4e89/importlib_resources-6.5.2-py3-none-any.whl", hash = "sha256:789cfdc3ed28c78b67a06acb8126751ced69a3d5f79c095a98298cd8a760ccec", size = 37461, upload-time = "2025-01-03T18:51:54.306Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iniconfig"
|
||||
version = "2.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ipykernel"
|
||||
version = "7.2.0"
|
||||
|
|
@ -1349,6 +1360,15 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/48/31/05e764397056194206169869b50cf2fee4dbbbc71b344705b9c0d878d4d8/platformdirs-4.9.2-py3-none-any.whl", hash = "sha256:9170634f126f8efdae22fb58ae8a0eaa86f38365bc57897a6c4f781d1f5875bd", size = 21168, upload-time = "2026-02-16T03:56:08.891Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pluggy"
|
||||
version = "1.6.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pre-commit"
|
||||
version = "4.5.1"
|
||||
|
|
@ -1519,6 +1539,22 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest"
|
||||
version = "9.0.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "colorama", marker = "sys_platform == 'win32'" },
|
||||
{ name = "iniconfig" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pluggy" },
|
||||
{ name = "pygments" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "python-dateutil"
|
||||
version = "2.9.0.post0"
|
||||
|
|
|
|||
Reference in a new issue