From cb5a6fbd87b1c52c0f66edbada0e2606b5365acd Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Sat, 21 Mar 2026 16:53:07 +0100 Subject: [PATCH] refactor: convert jax verification script to pytest and update docs --- docs/DEVELOPMENT.md | 6 +++--- tests/test_jax_init.py | 23 ++++------------------- uv.lock | 36 ++++++++++++++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 22 deletions(-) diff --git a/docs/DEVELOPMENT.md b/docs/DEVELOPMENT.md index a88997e..6ec220b 100644 --- a/docs/DEVELOPMENT.md +++ b/docs/DEVELOPMENT.md @@ -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. diff --git a/tests/test_jax_init.py b/tests/test_jax_init.py index 3d19422..e672b55 100644 --- a/tests/test_jax_init.py +++ b/tests/test_jax_init.py @@ -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" diff --git a/uv.lock b/uv.lock index 219f775..ce5ecc2 100644 --- a/uv.lock +++ b/uv.lock @@ -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"