Run tests in CI on GPU [in progress]

.bashrc

@@ -1,3 +1,2 @@
 export XLA_PYTHON_CLIENT_PREALLOCATE="false"
 export XLA_PYTHON_CLIENT_ALLOCATOR="platform"
-

.github/workflows/ci.yml

@@ -0,0 +1,74 @@
+name: Codebase tests
+
+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"
+
+jobs:
+  test:
+    strategy:
+      matrix:
+        runner: [ParallelHoss]
+    runs-on: ${{ matrix.runner }}
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install apt packages
+        uses: awalsh128/cache-apt-pkgs-action@latest
+        with:
+          packages: mesa-common-dev libegl1-mesa-dev libglfw3-dev libgl1-mesa-dev libglu1-mesa-dev ffmpeg
+          version: 1.0
+
+      # See https://github.com/Jimver/cuda-toolkit/blob/master/src/links/linux-links.ts
+      # for available versions
+      - uses: Jimver/[email protected]
+        id: cuda-toolkit
+        with:
+          cuda: '12.4.0'
+          linux-local-args: '["--toolkit"]'
+          method: 'network'
+
+      - uses: actions/setup-python@v4
+        with:
+          python-version: 3.11.5
+          cache: 'pip'
+          cache-dependency-path: setup.py
+
+      - id: auth
+        uses: google-github-actions/auth@v2
+        with:
+          credentials_json: "${{ secrets.ARTIFACT_REGISTRY_KEY }}"
+
+      - name: Set up Cloud SDK
+        uses: google-github-actions/setup-gcloud@v2
+
+      - name: Cache b3d assets
+        uses: actions/cache@v4
+        with:
+          path: assets
+          key: ${{ runner.os }}-${{ matrix.runner }}-assets
+          restore-keys: ${{ runner.os }}-${{ matrix.runner }}-assets
+
+      - name: Install Python dependencies
+        run: |
+          pip install keyring keyrings.google-artifactregistry-auth
+          pip install -e . --extra-index-url https://us-west1-python.pkg.dev/probcomp-caliban/probcomp/simple/
+
+      - name: Download test assets
+        run: b3d_pull
+
+      - name: Create test results directory
+        run: mkdir -p assets/test_results
+
+      - name: Run Tests
+        run: ./run_tests.sh

.gitignore

@@ -4,4 +4,4 @@
 assets/shared_data_bucket/*
 assets/test_results/*
 assets/bop/*
-**/.ipynb_checkpoints
+**/.ipynb_checkpoints

.pre-commit-config.yaml

@@ -0,0 +1,8 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      - id: check-yaml
+        args: [--unsafe]
+      - id: end-of-file-fixer
+      - id: trailing-whitespace

README.md

@@ -5,55 +5,68 @@
 </p>
 
 # b3d repository
-This repository contains code for Bayesian 3D inverse graphics.
-
-The `b3d/bayes3d/` subdirectory contains code for the `bayes3d` project,
-and the `b3d/chisight/` subdirectory contains code for post-bayes3D
-ChiSight systems (currently, SAMA4D).
 
-# Installing b3d
+This repository contains code for Bayesian 3D inverse graphics.
 
-Create `b3d` conda environment: 
-```
-conda create -n b3d python=3.12
-```
+The `b3d/bayes3d/` subdirectory contains code for the `bayes3d` project, and the `b3d/chisight/` subdirectory contains code for post-bayes3D ChiSight systems (currently, SAMA4D).
 
-Run install script:
-```
-bash -i install.sh
-```
+## Installing CUDA 12.3+
 
-
-### Installing CUDA 12.3+
 Run `nvidia-smi`
-```
+
+```sh
 +-----------------------------------------------------------------------------------------+
 | NVIDIA-SMI 550.54.15              Driver Version: 550.54.15      CUDA Version: 12.4     |
 |-----------------------------------------+------------------------+----------------------+
 ```
 and verify that `CUDA Version: xx.x` is >= 12.3.
 
 If not, ensure your NVIDIA GPU supports CUDA 12.3+ and install a NVIDIA driver version compatible with 12.3+. For GCP instances, use:
-```
+
+```sh
 sudo sh -c "echo 'export DRIVER_VERSION=550.54.15' > /opt/deeplearning/driver-version.sh"
 /opt/deeplearning/install-driver.sh
 ```
 If the above commands fail, then first uninstall the existing driver by running the below, then try again.
-```
+
+```sh
 /opt/deeplearning/uninstall-driver.sh
 sudo reboot
 ```
 
+## Installing b3d
+
+`b3d`'s GenJAX dependency is hosted in Google Artifact Registry. To configure your machine to access the package:
+
+- Check if you can access the [GenJAX Users Group](https://groups.google.com/u/1/a/chi-fro.org/g/genjax-users), and if not ask @sritchie (on Slack) to add you. Send him your Google-Cloud-associated email address.
+- [Install the Google Cloud command line tools](https://cloud.google.com/sdk/docs/install).
+- Run `gcloud auth application-default login`.  (This command needs to be rerun every time your machine reboots.)
+
+When that completes, return to the `hgps` directory and create the `b3d` conda environment:
+
+```sh
+conda create -n b3d python=3.11
+```
+
+Then run the install script:
+
+```sh
+bash -i install.sh
+```
+
 ### Environment Variables
+
 Add the following to your bash_rc:
-```
+
+```sh
 export XLA_PYTHON_CLIENT_PREALLOCATE="false"
 export XLA_PYTHON_CLIENT_ALLOCATOR="platform"
 ```
 
+## Visualizer
 
-### Visualizer
 Tunnel port `8812` for Rerun visualization by adding the `RemoteForward`line to your ssh config:
+
 ```
 Host xx.xx.xxx.xxx
     HostName xx.xx.xxx.xxx
@@ -65,4 +78,4 @@ Host xx.xx.xxx.xxx
 Install rerun on local machine `pip install rerun-sdk` and open viewer:
 ```
 rerun --port 8812
-```
+```

b3d/__init__.py

@@ -27,4 +27,4 @@
     "Rot",
     "MeshLibrary",
     "Mesh"
-]
+]

b3d/bayes3d/__init__.py

@@ -1,4 +1,4 @@
 from .mesh_library import *
 from .enumerative_proposals import *
 from .model import *
-from .scene_graph import *
+from .scene_graph import *

b3d/bayes3d/model.py

@@ -25,7 +25,7 @@ def get_rgb_depth_inliers_from_observed_rendered_args(observed_rgb, rendered_rgb
     observed_lab = b3d.colors.rgb_to_lab(observed_rgb)
     rendered_lab = b3d.colors.rgb_to_lab(rendered_rgb)
     error = (
-        jnp.linalg.norm(observed_lab[...,1:3] - rendered_lab[...,1:3], axis=-1) + 
+        jnp.linalg.norm(observed_lab[...,1:3] - rendered_lab[...,1:3], axis=-1) +
         jnp.abs(observed_lab[...,0] - rendered_lab[...,0])
     )
 
@@ -63,22 +63,22 @@ def logpdf(self, observed, rendered_rgb, rendered_depth, model_args, fx, fy, far
             1.0 * jnp.sum(undecided * areas)  +
             outlier_prob * jnp.sum(outliers * areas)
         ) * multiplier
-    
+
 rgbd_sensor_model = RGBDSensorModel()
 
 def model_multiobject_gl_factory(renderer, image_likelihood=rgbd_sensor_model):
     @genjax.gen
     def model(
-        _num_obj_arr, # new 
+        _num_obj_arr, # new
         model_args,
         object_library
     ):
 
         object_poses = Pose(jnp.zeros((0,3)), jnp.zeros((0,4)))
         object_indices = jnp.empty((0,), dtype=int)
         camera_pose = uniform_pose(jnp.ones(3)*-100.0, jnp.ones(3)*100.0) @ "camera_pose"
-            
-        for i in range(_num_obj_arr.shape[0]):        
+
+        for i in range(_num_obj_arr.shape[0]):
             object_identity = uniform_discrete(jnp.arange(-1, len(object_library.ranges))) @ f"object_{i}"
             object_indices = jnp.concatenate((object_indices, jnp.array([object_identity])))
 
@@ -156,4 +156,4 @@ def rerun_visualize_trace_t(trace, t, modes=["rgb", "depth", "inliers"]):
                     pose.apply(vertices),
                     colors=(attributes * 255).astype(jnp.uint8),
                 ),
-            )
+            )

b3d/bucket_utils/README.md

@@ -13,39 +13,39 @@ By default, the `chi_pull` command pulls _all existing data_ in the GCP bucket i
 By default, files with names that already exist in your local `dcolmap/assets/shared_bucket_data` will not be overwritten with the version on the cloud, i.e. the download will be skipped for those files.
 
 ```
-chi_pull   # pulls all bucket contents into local, 
+chi_pull   # pulls all bucket contents into local,
             with no overwrite on already-existing filenames
 ```
 
 ##### Download all bucket data (enable overwrite)
-To pull _all existing data_ in the GCP bucket _with overwrite_ into your local `dcolmap/assets/shared_bucket_data`, use the `-ow` flag. 
+To pull _all existing data_ in the GCP bucket _with overwrite_ into your local `dcolmap/assets/shared_bucket_data`, use the `-ow` flag.
 Note that other than potential overwriting behavior enabled by `-ow`, `chi_pull` never deletes files on your local directory; be sure to delete deprecated local files before `chi_push`'ing your local contents onto the bucket.
 
 ```
-chi_pull -ow  # pulls all bucket contents into local, 
+chi_pull -ow  # pulls all bucket contents into local,
                 with overwrite on already-existing filenames
 ```
 
 ##### Download select bucket data items
 Finally, to pull select items (files, directories) in the GCP bucket into your local `dcolmap/assets/shared_bucket_data`, use the `-fn` flag, followed by the path of the item relative to the root of the bucket.
-All necessary intermediate directories will automatically be initialized if not already existing in your local `dcomap/assets/shared_bucket_data`. 
+All necessary intermediate directories will automatically be initialized if not already existing in your local `dcomap/assets/shared_bucket_data`.
 
-For example, if you would like to pull `input_data/banana_datas/banana1.npz`, 
+For example, if you would like to pull `input_data/banana_datas/banana1.npz`,
 
 ```
 chi_pull -fn input_data/banana_datas/banana1.npz  # populates dcolmap/assets/shared_bucket_data/input_data/banana_datas/banana1.npz
 ```
 
-For another example, if you would like to pull all items in `input_data/mug_datas`, 
+For another example, if you would like to pull all items in `input_data/mug_datas`,
 
 ```
 chi_pull -fn input_data/mug_datas  # recursively populates dcolmap/assets/shared_bucket_data/input_data/mug_datas
 ```
 
-You can combine these flags. For example, if you would like to pull all items in `input_data/mug_datas` with overwrite enabled, 
+You can combine these flags. For example, if you would like to pull all items in `input_data/mug_datas` with overwrite enabled,
 
 ```
-chi_pull -ow -fn input_data/mug_datas  # recursively populates dcolmap/assets/shared_bucket_data/input_data/mug_datas, 
+chi_pull -ow -fn input_data/mug_datas  # recursively populates dcolmap/assets/shared_bucket_data/input_data/mug_datas,
                                             with overwrite on already-existing filenames
 ```
 
@@ -54,7 +54,7 @@ The `chi_push` command is used for downloading data from your local into the GCP
 
 ##### Upload all local data (disable overwrite)
 ```
-chi_push   # pushes all local content onto bucket, 
+chi_push   # pushes all local content onto bucket,
             with no overwrite on already-existing filenames
 ```
 
@@ -65,6 +65,6 @@ chi_push -ow  # pushes all local content onto bucket, with overwrite. A DANGEROU
 
 ##### Upload select local data items
 ```
-chi_push -ow -fn input_data/mug_datas  # recursively populates gs://hgps_data_bucket/input_data/mug_datas, 
+chi_push -ow -fn input_data/mug_datas  # recursively populates gs://hgps_data_bucket/input_data/mug_datas,
                                             with overwrite on already-existing filenames
 ```

b3d/camera.py

@@ -31,7 +31,7 @@ def from_array(cls, intrinsics: Array):
     def as_array(self):
         """Returns intrinsics as a float array."""
         return jnp.array(self)
-    
+
     def downscale(self, factor):
         return Intrinsics(
             self.width // factor,

b3d/chisight/dense/__init__.py

@@ -1,3 +1,3 @@
 # from . import model, patch_tracking, likelihoods, differentiable_renderer
 
-# __all__ = ["model", "patch_tracking", "likelihoods", "differentiable_renderer"]
+# __all__ = ["model", "patch_tracking", "likelihoods", "differentiable_renderer"]

b3d/chisight/dense/differentiable_renderer.py

@@ -148,7 +148,7 @@ def render_to_average(
     - image (H, W, A)
     """
     weights, attributes = render_to_dist_params(renderer, vertices, faces, vertex_attributes, hyperparams=hyperparams)
-    return dist_params_to_average(weights, attributes, background_attribute)    
+    return dist_params_to_average(weights, attributes, background_attribute)
 
 def dist_params_to_average(weights, attributes, background_attribute):
     extended_attributes = jnp.concatenate([jnp.tile(background_attribute, (attributes.shape[0], attributes.shape[1], 1, 1)), attributes], axis=2)
@@ -216,13 +216,13 @@ def get_weights_and_barycentric_coords(ij, vertices, faces, triangle_intersected
         fill_value = -1
     ) - 1
     unique_triangle_values_safe = jnp.where(unique_triangle_values < 0, unique_triangle_values[0], unique_triangle_values)
-    
+
     signed_dist_values, barycentric_coords = get_signed_dists_and_barycentric_coords(
         ij, unique_triangle_values_safe, vertices, faces, hyperparams_and_intrinsics
     )
     z_values = get_z_values(ij, unique_triangle_values_safe, vertices, faces, hyperparams_and_intrinsics)
     z_values = jnp.where(unique_triangle_values >= 0, z_values, z_values.max())
-    
+
     # Math from the softras paper
     signed_dist_scores = jax.nn.sigmoid(jnp.sign(signed_dist_values) * signed_dist_values ** 2 / SIGMA)
 
@@ -238,7 +238,7 @@ def get_weights_and_barycentric_coords(ij, vertices, faces, triangle_intersected
     unnorm_weights = jnp.where(unique_triangle_values >= 0, unnorm_weights, 0.0)
     unnorm_weights = jnp.concatenate([jnp.array([jnp.exp(jnp.clip(EPSILON/GAMMA, -20, 20))]), unnorm_weights])
     weights = unnorm_weights / jnp.sum(unnorm_weights)
-    
+
     extended_triangle_indices = jnp.concatenate([jnp.array([-10]), unique_triangle_values])
     barycentric_coords = jnp.where(unique_triangle_values[:, None] >= 0, barycentric_coords, -jnp.ones_like(barycentric_coords))
     return (extended_triangle_indices, weights, barycentric_coords)
@@ -295,7 +295,7 @@ def get_signed_dist_and_barycentric_coords(ij, triangle_idx, vertices, faces, hy
     bary = jnp.array([a, b, c]) / (a + b + c + 1e-6)
     bary = jnp.clip(bary, 0., 1.)
     bary = bary / (jnp.sum(bary) + 1e-6)
-    
+
     return (signed_distance, bary)
 
 def pt_is_in_plane(triangle, point_on_plane):
@@ -353,7 +353,7 @@ def project_ray_to_plane(ray_origin, ray_dir, triangle):
 
     t = jnp.dot(normal, vertex1 - ray_origin) / (denom + 1e-5)
     intersection_point = ray_origin + t * ray_dir
-    
+
     return jnp.where(
         jnp.abs(denom) < 1e-6, -jnp.ones(3), intersection_point
     )
@@ -370,4 +370,4 @@ def dist_to_line_seg(a, b, p):
 def _signed_area_to_point(a, b, p):
     Va = b - a
     area = jnp.cross(Va, p - a)[2] / 2
-    return area
+    return area

b3d/chisight/dense/likelihoods/__init__.py

@@ -1,2 +1,2 @@
 from .other_likelihoods import *
-from .krays import *
+from .krays import *