Add gpu processing (#440)

* Add gpu label * activate scanvi * Change compute env * update gpu config * Add Jax * Add jaxx install * Update comp env * Update src/wf_utils/labels_tw.config Co-authored-by: Robrecht Cannoodt <[email protected]> --------- Co-authored-by: Robrecht Cannoodt <[email protected]>
openproblems-bio · Apr 29, 2024 · 6db0442 · 6db0442
1 parent 22ceaf0
commit 6db0442
Show file tree

Hide file tree

Showing 34 changed files with 51 additions and 28 deletions.
diff --git a/src/common/process_task_results/run/run_nf_tower_test.sh b/src/common/process_task_results/run/run_nf_tower_test.sh
@@ -33,6 +33,6 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/common/workflows/transform_meta/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --config /tmp/nextflow.config
diff --git a/src/datasets/resource_scripts/cellxgene_census.sh b/src/datasets/resource_scripts/cellxgene_census.sh
@@ -147,7 +147,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/datasets/workflows/process_cellxgene_census/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file "/tmp/params.yaml" \
   --config /tmp/nextflow.config \
   --labels cellxgene_census,dataset_loader
diff --git a/src/datasets/resource_scripts/dataset_info.sh b/src/datasets/resource_scripts/dataset_info.sh
@@ -37,7 +37,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/datasets/workflows/extract_dataset_info/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file "/tmp/params.yaml" \
   --config /tmp/nextflow.config
 

diff --git a/src/datasets/resource_scripts/openproblems_neurips2021_multimodal.sh b/src/datasets/resource_scripts/openproblems_neurips2021_multimodal.sh
@@ -49,7 +49,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/datasets/workflows/process_openproblems_neurips2021_bmmc/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file "$params_file" \
   --config /tmp/nextflow.config \
   --labels openproblems_neurips2021_bmmc,dataset_loader \
diff --git a/src/datasets/resource_scripts/openproblems_v1.sh b/src/datasets/resource_scripts/openproblems_v1.sh
@@ -176,7 +176,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/datasets/workflows/process_openproblems_v1/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file "$params_file" \
   --config /tmp/nextflow.config \
   --labels openproblems_v1,dataset_loader
diff --git a/src/datasets/resource_scripts/openproblems_v1_multimodal.sh b/src/datasets/resource_scripts/openproblems_v1_multimodal.sh
@@ -79,7 +79,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/datasets/workflows/process_openproblems_v1_multimodal/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file "$params_file" \
   --labels openproblems_v1_multimodal,dataset_loader \
   --config /tmp/nextflow.config
diff --git a/src/datasets/resource_test_scripts/neurips2021_bmmc.sh b/src/datasets/resource_test_scripts/neurips2021_bmmc.sh
@@ -62,7 +62,7 @@ nextflow run . \
 #   --revision main_build \
 #   --main-script target/nextflow/datasets/workflows/process_openproblems_neurips2021_bmmc/main.nf \
 #   --workspace 53907369739130 \
-#   --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+#   --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
 #   --params-file "$params_file" \
 #   --config /tmp/nextflow.config \
 #   --labels predict_modality

diff --git a/src/tasks/batch_integration/methods/scanvi/config.vsh.yaml b/src/tasks/batch_integration/methods/scanvi/config.vsh.yaml
@@ -51,6 +51,9 @@ platforms:
       - type: python
         pypi:
           - scvi-tools>=1.1.0
+      - type: docker
+        run: |
+          pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
   - type: nextflow
     directives:
       label: [midtime, lowmem, lowcpu, gpu]
diff --git a/src/tasks/batch_integration/methods/scvi/config.vsh.yaml b/src/tasks/batch_integration/methods/scvi/config.vsh.yaml
@@ -49,6 +49,9 @@ platforms:
       - type: python
         pypi:
           - scvi-tools>=1.1.0
+      - type: docker
+        run: |
+          pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
   - type: nextflow
     directives: 
       label: [midtime, midmem, lowcpu, gpu]
diff --git a/src/tasks/batch_integration/resources_scripts/process_datasets.sh b/src/tasks/batch_integration/resources_scripts/process_datasets.sh
@@ -26,7 +26,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/batch_integration/workflows/process_datasets/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/batch_integration/resources_scripts/run_benchmark.sh b/src/tasks/batch_integration/resources_scripts/run_benchmark.sh
@@ -15,7 +15,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/batch_integration/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config src/wf_utils/labels_tw.config \

diff --git a/src/tasks/batch_integration/resources_scripts/run_benchmark_test.sh b/src/tasks/batch_integration/resources_scripts/run_benchmark_test.sh
@@ -18,7 +18,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/batch_integration/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/denoising/resources_scripts/process_datasets.sh b/src/tasks/denoising/resources_scripts/process_datasets.sh
@@ -27,7 +27,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/denoising/workflows/process_datasets/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/denoising/resources_scripts/run_benchmark.sh b/src/tasks/denoising/resources_scripts/run_benchmark.sh
@@ -16,7 +16,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/denoising/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config src/wf_utils/labels_tw.config \

diff --git a/src/tasks/denoising/resources_scripts/run_benchmark_test.sh b/src/tasks/denoising/resources_scripts/run_benchmark_test.sh
@@ -18,7 +18,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/denoising/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/dimensionality_reduction/resources_scripts/process_datasets.sh b/src/tasks/dimensionality_reduction/resources_scripts/process_datasets.sh
@@ -27,7 +27,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/dimensionality_reduction/workflows/process_datasets/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/dimensionality_reduction/resources_scripts/run_benchmark.sh b/src/tasks/dimensionality_reduction/resources_scripts/run_benchmark.sh
@@ -15,7 +15,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/dimensionality_reduction/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config src/wf_utils/labels_tw.config \

diff --git a/src/tasks/dimensionality_reduction/resources_scripts/run_benchmark_test.sh b/src/tasks/dimensionality_reduction/resources_scripts/run_benchmark_test.sh
@@ -18,7 +18,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/dimensionality_reduction/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/label_projection/methods/scanvi/config.vsh.yaml b/src/tasks/label_projection/methods/scanvi/config.vsh.yaml
@@ -38,6 +38,9 @@ platforms:
         packages: 
           - scarches
           - scvi-tools>=1.1.0
+      - type: docker
+        run: |
+          pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
   - type: nextflow
     directives: 
       label: [midtime, midmem, highcpu, gpu]
diff --git a/src/tasks/label_projection/methods/scanvi_scarches/config.vsh.yaml b/src/tasks/label_projection/methods/scanvi_scarches/config.vsh.yaml
@@ -45,6 +45,9 @@ platforms:
     setup:
       - type: python
         pypi: scvi-tools>=1.1.0
+      - type: docker
+        run: |
+          pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
   - type: nextflow
     directives:
       label: [midtime, midmem, midcpu, gpu]
diff --git a/src/tasks/label_projection/resources_scripts/process_datasets.sh b/src/tasks/label_projection/resources_scripts/process_datasets.sh
@@ -27,7 +27,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/label_projection/workflows/process_datasets/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/label_projection/resources_scripts/run_benchmark.sh b/src/tasks/label_projection/resources_scripts/run_benchmark.sh
@@ -7,6 +7,7 @@ cat > /tmp/params.yaml << HERE
 input_states: s3://openproblems-data/resources/label_projection/datasets/**/state.yaml
 rename_keys: 'input_train:output_train,input_test:output_test,input_solution:output_solution'
 output_state: "state.yaml"
+settings: '{"method_ids": "scanvi_scarches"}'
 publish_dir: "$publish_dir"
 HERE
 
@@ -15,7 +16,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/label_projection/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config src/wf_utils/labels_tw.config \

diff --git a/src/tasks/label_projection/resources_scripts/run_benchmark_test.sh b/src/tasks/label_projection/resources_scripts/run_benchmark_test.sh
@@ -18,7 +18,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/label_projection/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/label_projection/workflows/run_benchmark/config.vsh.yaml b/src/tasks/label_projection/workflows/run_benchmark/config.vsh.yaml
@@ -68,7 +68,7 @@ functionality:
     - name: label_projection/methods/knn
     - name: label_projection/methods/logistic_regression
     - name: label_projection/methods/mlp
-    # - name: label_projection/methods/scanvi
+    - name: label_projection/methods/scanvi
     - name: label_projection/methods/scanvi_scarches
     - name: label_projection/methods/xgboost
     - name: label_projection/metrics/accuracy

diff --git a/src/tasks/label_projection/workflows/run_benchmark/main.nf b/src/tasks/label_projection/workflows/run_benchmark/main.nf
@@ -19,7 +19,7 @@ workflow run_wf {
     knn,
     logistic_regression,
     mlp,
-    // scanvi,
+    scanvi,
     scanvi_scarches,
     // seurat_transferdata,
     xgboost

diff --git a/src/tasks/match_modalities/resources_scripts/process_datasets.sh b/src/tasks/match_modalities/resources_scripts/process_datasets.sh
@@ -27,7 +27,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/match_modalities/workflows/process_datasets/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/match_modalities/resources_scripts/run_benchmark.sh b/src/tasks/match_modalities/resources_scripts/run_benchmark.sh
@@ -16,7 +16,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/match_modalities/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config src/wf_utils/labels_tw.config \

diff --git a/src/tasks/predict_modality/resources_scripts/process_datasets.sh b/src/tasks/predict_modality/resources_scripts/process_datasets.sh
@@ -25,7 +25,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/predict_modality/workflows/process_datasets/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/predict_modality/resources_scripts/run_benchmark.sh b/src/tasks/predict_modality/resources_scripts/run_benchmark.sh
@@ -17,7 +17,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/predict_modality/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config src/wf_utils/labels_tw.config \

diff --git a/src/tasks/spatial_decomposition/methods/destvi/config.vsh.yaml b/src/tasks/spatial_decomposition/methods/destvi/config.vsh.yaml
@@ -34,7 +34,9 @@ platforms:
       - type: python
         packages: 
           - scvi-tools>=1.1.0
-
+      - type: docker
+        run: |
+          pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
   - type: native
   - type: nextflow
     directives:

diff --git a/src/tasks/spatial_decomposition/methods/stereoscope/config.vsh.yaml b/src/tasks/spatial_decomposition/methods/stereoscope/config.vsh.yaml
@@ -34,6 +34,9 @@ platforms:
       - type: python
         packages: 
           - scvi-tools>=1.1.0
+      - type: docker
+        run: |
+          pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
   - type: native
   - type: nextflow
     directives:

diff --git a/src/tasks/spatial_decomposition/resources_scripts/process_datasets.sh b/src/tasks/spatial_decomposition/resources_scripts/process_datasets.sh
@@ -29,7 +29,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/spatial_decomposition/workflows/process_datasets/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config /tmp/nextflow.config \

diff --git a/src/tasks/spatial_decomposition/resources_scripts/run_benchmark.sh b/src/tasks/spatial_decomposition/resources_scripts/run_benchmark.sh
@@ -15,7 +15,7 @@ tw launch https://github.com/openproblems-bio/openproblems-v2.git \
   --pull-latest \
   --main-script target/nextflow/spatial_decomposition/workflows/run_benchmark/main.nf \
   --workspace 53907369739130 \
-  --compute-env 1pK56PjjzeraOOC2LDZvN2 \
+  --compute-env 6TeIFgV5OY4pJCk8I0bfOh \
   --params-file /tmp/params.yaml \
   --entry-name auto \
   --config src/wf_utils/labels_tw.config \

diff --git a/src/wf_utils/labels_tw.config b/src/wf_utils/labels_tw.config
@@ -34,6 +34,11 @@ process {
   withLabel: highsharedmem {
     containerOptions = { workflow.containerEngine != 'singularity' ? "--shm-size ${String.format("%.0f",task.memory.mega * 0.25)}" : ""}
   }
+  withLabel: gpu {
+    accelerator = 1
+    containerOptions = { workflow.containerEngine == "singularity" ? '--nv':
+       ( workflow.containerEngine == "docker" ? '--gpus all': null ) }
+  }
 }
 
 def get_memory(to_compare) {