aind-disrnn-pipeline

Core pipeline to scale experiments for the AIND-disRNN project.

AIND-disRNN Dispatcher: GitHub, CO capsule

AIND-disRNN Wrapper: GitHub, CO capsule

Credentials

You'll need a W&B API key to push your results to W&B:

1. Ask Han to invite you to the AIND-disRNN team on W&B.
2. Copy your API key from W&B - Account - API key
3. First add your key to Code Ocean account: Account - User Secrets - Add Secret - Custom key
   
4. Then attach the key to this pipeline: Pipeline settings - Credentials
   
5. You may also need to attach your Code Ocean API on the same page.

Usage

In the pipeline, enter the Hydra CLI override in the app panel, for example
-m model.penalties.beta=0.0001,0.001,0.01 model.training.n_steps=100 model.training.n_warmup_steps=50

Hit "Run", and you'll get results in CO asset (computation id c4808a41-d686-4458-830e-a11d4d19ddb7)

Logging and visualization in W&B

Switching between CPU and GPU machines for the wrapper

Based on some benchmark tests, we decided to use CPU machine for the wrapper by default.

To migrate from GPU to CPU, you should pull the latest wrapper and pipeline after this PR and this PR are merged.

If you want to switch between CPU/GPU machines in the future, follow these steps:

1. In the DockerFile of the wrapper capsule, uncomment the correct base image and jax version.

# --- Base image and Jax for CPU (preferred) ---
FROM $REGISTRY_HOST/codeocean/mambaforge3:24.5.0-0-python3.12.4-ubuntu22.04
ARG JAX_PKG="jax==0.6.1"

# --- Base image and Jax for GPU ---
# FROM $REGISTRY_HOST/codeocean/pytorch:2.4.0-cuda12.4.0-mambaforge24.5.0-0-python3.12.4-ubuntu22.04
# ARG JAX_PKG="jax[cuda12]==0.6.1"

2. In the pipeline, you should also check the machine resources.
   
   For example, if you are using CPU machines, you should not see a "GPU" option there, and vice versa.

   However, there seems to be a bug that CO cannot automatically update the "Resources" field after you change the DockerFile of the capsule. A workaround is to add the same wrapper capsule to the pipeline like this
   

   and it somehow forces CO to rethink whether it will need a GPU based on the DockerFile. Once this is done, you can remove the dummy wrapper capsule.

3. You'll need to do a Reproducible Run of the wrapper before you're able to trigger the pipeline with the new CPU/GPU setting.