local tests + ci

.envrc

@@ -0,0 +1,3 @@
+watch_file pyproject.toml
+eval "$(pixi shell-hook)"
+

.github/workflows/ci.yml

@@ -0,0 +1,53 @@
+name: Codebase tests (GPU) 
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+  # Allows you to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+env:
+  XLA_FLAGS: --xla_gpu_enable_command_buffer=
+  XLA_PYTHON_CLIENT_PREALLOCATE: "false"
+  XLA_PYTHON_CLIENT_ALLOCATOR: "platform"
+  CPLUS_INCLUDE_PATH: "$CONDA_PREFIX/targets/x86_64-linux/include"
+
+jobs:
+  test:
+    strategy:
+      matrix:
+        runner: [ParallelHoss]
+    runs-on: ${{ matrix.runner }}
+    steps:
+      - name: Clone repo
+      - uses: actions/checkout@v4
+
+      - id: auth
+        uses: google-github-actions/auth@v2
+        with:
+          credentials_json: "${{ secrets.ARTIFACT_REGISTRY_KEY }}"
+
+      - name: Setup gcloud
+        uses: google-github-actions/setup-gcloud@v2
+
+      - name: Install system environment and project dependencies
+        env:
+          GITHUB_CI: "true"
+        run: ./install.sh
+
+      - name: Setup caching for test assets
+        uses: actions/cache@v4
+        with:
+          path: assets
+          key: ${{ runner.os }}-${{ matrix.runner }}-assets
+          restore-keys: ${{ runner.os }}-${{ matrix.runner }}-assets
+
+      - name: Run tests
+        uses: prefix-dev/[email protected]
+        with:
+          pixi-version: v0.25.0
+          cache: true
+        run: pixi run all-tests-ci

install.sh

@@ -1,3 +1,4 @@
+
 #!/usr/bin/env bash
 
 set -euo pipefail

pyproject.toml

@@ -47,25 +47,31 @@ mesa-libegl-cos7-x86_64 = "*"
 mesa-libglu-devel-cos7-x86_64 = "*"
 mesa-libglu-cos7-x86_64 = "*"
 natsort = ">=8.4.0,<8.5"
+ninja = "*"
 numpy = "==1.26.4"
 optax = ">=0.2.2,<0.3"
 pdoc3 = ">=0.11.1,<0.12"
 pillow = "==10.3.0"
 polars = ">=0.20,<0.21"
+pytest = "*"
 pytorch = {version = "==2.3.0", build = "cuda12*"}
 rerun-sdk = "==0.16.1"
 scikit-learn = ">=1.5.1,<1.6"
 scipy = ">=1.14.0,<1.15"
 torchvision = "*"
-trimesh = ">=4.4.3,<4.5"
+trimesh = "==4.2.4"
 tqdm = "==4.66.2"
-ninja = "*"
-pytest = "*"
 
 [tool.pixi.tasks]
-b3d-pull = {cmd = "python b3d_pull.py -ow", cwd = "b3d/bucket_utils" }
+b3d-pull = {cmd = "python b3d_pull.py -ow", cwd = "src/b3d/bucket_utils" }
+
 rerun = "rerun --port 8812"
+
 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", TORCH_CUDA_ARCH_LIST = "8.9"} }
 
+all-tests-local = { cmd = "pytest --ignore=tests/sama4d/video_to_tracks_and_segmentation/keypoint_tracking_and_segmentation_test.py --ignore=tests/test_feature_track_data.py --ignore=tests/test_image_posterior_resolution_invariance.py --ignore=tests/test_likelihood_invariances.py --ignore=tests/test_mug_handle_posterior.py --ignore=tests/test_mug_smc_pose_inference_synthetic_and_real.py tests", env = { XLA_PYTHON_CLIENT_PREALLOCATE = "false", XLA_PYTHON_CLIENT_ALLOCATOR = "platform", CPLUS_INCLUDE_PATH = "$CONDA_PREFIX/targets/x86_64-linux/include", TORCH_CUDA_ARCH_LIST = "8.9"} }
+
+all-tests-ci = { cmd = "pytest --ignore=tests/sama4d/video_to_tracks_and_segmentation/keypoint_tracking_and_segmentation_test.py --ignore=tests/test_feature_track_data.py --ignore=tests/test_image_posterior_resolution_invariance.py --ignore=tests/test_likelihood_invariances.py --ignore=tests/test_mug_handle_posterior.py --ignore=tests/test_mug_smc_pose_inference_synthetic_and_real.py tests", depends-on = ["b3d-pull"], env = { XLA_PYTHON_CLIENT_PREALLOCATE = "false", XLA_PYTHON_CLIENT_ALLOCATOR = "platform", CPLUS_INCLUDE_PATH = "$CONDA_PREFIX/targets/x86_64-linux/include", TORCH_CUDA_ARCH_LIST = "8.9"} }
+
 [tool.pytest.ini_options]
 pythonpath = ["src"]

scripts/check-auth.sh

@@ -76,7 +76,9 @@ gcloud-global-install() {
 }
 
 gcloud-authenticated() {
-	if gcloud auth application-default print-access-token >/dev/null; then
+  if [ "$GITHUB_CI" = "true" ]; then
+    return 0
+	elif gcloud auth application-default print-access-token >/dev/null; then
 		return 0
 	else
 		return 1
@@ -93,7 +95,7 @@ gcloud-init() {
 
 gcloud-auth-adc() {
 	if ! gcloud auth application-default login; then
-		return 1
+		returny 1
 	else
 		return 0
 	fi

scripts/env-activation.sh

@@ -9,3 +9,7 @@ export XLA_FLAGS=--xla_gpu_enable_command_buffer=
 export XLA_PYTHON_CLIENT_PREALLOCATE=false
 export XLA_PYTHON_CLIENT_ALLOCATOR=platform
 export CPLUS_INCLUDE_PATH="$CONDA_PREFIX/targets/x86_64-linux/include"
+
+if ! [ -e "$CONDA_PREFIX/x86_64-conda-linux-gnu/sysroot/usr/lib64/libEGL.so" ]; then
+  cp assets/system/libEGL.so "$CONDA_PREFIX/x86_64-conda-linux-gnu/sysroot/usr/lib64/"
+fi

src/b3d/utils.py

@@ -25,7 +25,8 @@
 # # # # # # # # # # # # 
 
 def get_root_path() -> Path:
-    return Path(Path(b3d.__file__).parents[1])
+    """Returns the project root directory."""
+    return Path(Path(b3d.__file__).parents[2])
 
 
 def get_assets() -> Path: