env activation

probcomp · Jul 16, 2024 · 3b5e1ea · 3b5e1ea
1 parent 9bd9242
commit 3b5e1ea
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/pixi.lock b/pixi.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -17,6 +17,9 @@ platforms = ["linux-64"]
 [tool.pixi.system-requirements]
 cuda = "12"
 
+[tool.pixi.activation]
+scripts = ["scripts/env-activation.sh"]
+
 [tool.pixi.pypi-options]
 index-url = "https://pypi.org/simple"
 extra-index-urls = ["https://[email protected]/probcomp-caliban/probcomp/simple"]
@@ -60,10 +63,10 @@ ninja = "*"
 pytest = "*"
 
 [tool.pixi.tasks]
-egl = "cd .pixi/envs/gpu/x86_64-conda-linux-gnu/sysroot/usr/lib64/ && ln -s libEGL_mesa.so.0.0.0 libEGL.so"
+foo = "export FOO=8"
 rerun = "rerun --port 8812"
 b3d-pull = {cmd = "python b3d_pull.py -ow", cwd = "b3d/bucket_utils" }
-test = { cmd = "pytest tests/dense_model_unit_tests/triangle_depth_posterior/test_triangle_depth_posterior.py", env = { XLA_PYTHON_CLIENT_PREALLOCATE = "false", XLA_PYTHON_CLIENT_ALLOCATOR = "platform", TORCH_CUDA_ARCH_LIST = "8.5", CPLUS_INCLUDE_PATH = "$CONDA_PREFIX/targets/x86_64-linux/include"} }
+test = { cmd = "pytest tests/dense_model_unit_tests/triangle_depth_posterior/test_triangle_depth_posterior.py", env = { XLA_PYTHON_CLIENT_PREALLOCATE = "false", XLA_PYTHON_CLIENT_ALLOCATOR = "platform", CPLUS_INCLUDE_PATH = "$CONDA_PREFIX/targets/x86_64-linux/include"} }
 
 [tool.pytest.ini_options]
 pythonpath = ["src"]
diff --git a/scripts/env-activation.sh b/scripts/env-activation.sh
@@ -0,0 +1,7 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+capability=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader)
+
+export TORCH_CUDA_ARCH_LIST="$capability"