Evo2

.github/workflows/unit-tests.yml

@@ -119,7 +119,9 @@ jobs:
           (github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'INCLUDE_SLOW_TESTS'))
         env:
           BIONEMO_DATA_SOURCE: ngc
-        run: pytest -v -m "slow" sub-packages/
+        # Not every sub-package has slow tests, and since some sub-packages have tests under the same name we need
+        #  to run package by package like we do with the fast tests.
+        run: ./ci/scripts/run_pytest.sh --no-nbval --only-slow --allow-no-tests
 
       - name: Run notebook tests
         if: |

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
         exclude: package.lock.json
       - id: detect-secrets
         name: detect-secrets (notebooks only)
-        args: ['--baseline', '.secrets-nb.baseline', '--exclude-files', '^.(?!.*\.ipynb)', '--exclude-lines', '"(hash|id|image/\w+)":.*', ]
+        args: ['--baseline', '.secrets-nb.baseline', '--exclude-files', '^.(?!.*\.ipynb)', '--exclude-lines', '"(hash|id|image/\w+)":.*|<.*at 0x[0-9a-f]+>|object at 0x[0-9a-f]+', ]
   - repo: local
     hooks:
       - id: license-header-check

.secrets.baseline

@@ -139,9 +139,9 @@
         "filename": "pyproject.toml",
         "hashed_secret": "79670e9c9d1c7ea5b81a96a2053d81437712c78e",
         "is_verified": false,
-        "line_number": 44
+        "line_number": 45
       }
     ]
   },
-  "generated_at": "2025-01-15T19:06:19Z"
+  "generated_at": "2025-01-30T14:18:42Z"
 }

ci/benchmarks/partial-conv/evo2_pretrain.yaml

@@ -0,0 +1,63 @@
+scope: partial-conv
+time_limit: 14400
+script_args:
+  # All arguments referenced in the script string must be specified here.
+  # Arguments not referenced in the script string must have the 'arg' field specified.
+  # See jet/core/configs.py for the specification of the configuration class
+  workspace:
+    value: /workspace/bionemo2
+    key_segment: False
+  data_path:
+    value: /data/evo2
+    key_segment: False
+  model:
+    value: evo2
+  variant:
+    value: train
+  config_name:
+    value: 7b
+  precision:
+    value: fp8
+  nodes:
+    value: 4
+  gpus:
+    value: 8
+  batch_size:
+    value: 2
+  pp:
+    value: 1
+  tp:
+    value: 8
+  cp:
+    value: 1
+  acc_grad:
+    value: 1
+  max_steps:
+    value: 20000
+script: |-
+  WANDB_API_KEY=$BIONEMO_WANDB_API_KEY python ${workspace}/sub-packages/bionemo-evo2/src/bionemo/evo2/run/train.py \
+  -d ${workspace}/sub-packages/bionemo-evo2/tests/config/test_dataset_config.yaml \
+  --dataset-dir ${data_path} \
+  --grad-acc-batches ${acc_grad} \
+  --fp8 \
+  --enable-preemption \
+  --ckpt-async-save \
+  --seq-length=8192 \
+  --tensor-parallel-size=${tp} \
+  --context-parallel-size=${cp} \
+  --pipeline-model-parallel-size=${pp} \
+  --workers 8 \
+  --num-nodes=${nodes} \
+  --devices=${gpus} \
+  --micro-batch-size=${batch_size} \
+  --model-size=${config_name} \
+  --max-steps=${max_steps} \
+  --limit-val-batches=20 \
+  --log-every-n-steps=50 \
+  --val-check-interval=500 \
+  --tflops-callback \
+  --experiment-dir=${tensorboard_dir}/${batch_size}bs_${nodes}node_${gpus}gpu_${max_steps}s_${precision}prec \
+  --wandb-project=${wandb_project_name} \
+  --wandb-group=${model}_${variant}_${config_name}__${target} \
+  --wandb-job-type=${pipeline_label} \
+  --disable-checkpointing;

ci/benchmarks/perf/evo2_pretrain.yaml

@@ -0,0 +1,67 @@
+scope: perf
+time_limit: 1800
+script_args:
+  # All arguments referenced in the script string must be specified here.
+  # Arguments not referenced in the script string must have the 'arg' field specified.
+  # See jet/core/configs.py for the specification of the configuration class
+  workspace:
+    value: /workspace/bionemo2
+    key_segment: False
+  data_path:
+    value: /data/evo2
+    key_segment: False
+  model:
+    value: evo2
+  variant:
+    value: train
+  precision:
+    value: fp8
+  gpus:
+    value: 8
+  batch_size:
+    value: 2
+  max_steps:
+    value: 100
+  tp:
+    value: 8
+  cp:
+    value: 1
+  pp:
+    value: 1
+  acc_grad:
+    value: 1
+  products:
+    - nodes: 1
+      config_name: 7b
+    - nodes: 2
+      config_name: 7b
+    - nodes: 8
+      config_name: 40b
+script: |-
+  WANDB_API_KEY=$BIONEMO_WANDB_API_KEY python ${workspace}/sub-packages/bionemo-evo2/src/bionemo/evo2/run/${variant}.py \
+  -d ${workspace}/sub-packages/bionemo-evo2/tests/config/test_dataset_config.yaml \
+  --dataset-dir ${data_path} \
+  --grad-acc-batches ${acc_grad} \
+  --fp8 \
+  --enable-preemption \
+  --ckpt-async-save \
+  --use-megatron-comm-overlap-llama3-8k \
+  --seq-length=8192 \
+  --tensor-parallel-size=${tp} \
+  --context-parallel-size=${cp} \
+  --pipeline-model-parallel-size=${pp} \
+  --workers 8 \
+  --num-nodes=${nodes} \
+  --devices=${gpus} \
+  --micro-batch-size=${batch_size} \
+  --model-size=${config_name} \
+  --max-steps=${max_steps} \
+  --limit-val-batches=20 \
+  --log-every-n-steps=50 \
+  --val-check-interval=${max_steps} \
+  --tflops-callback \
+  --experiment-dir=${tensorboard_dir}/${batch_size}bs_${nodes}node_${gpus}gpu_${max_steps}s_${precision}prec \
+  --wandb-project=${wandb_project_name} \
+  --wandb-group=${model}_${variant}_${config_name}__${target} \
+  --wandb-job-type=${pipeline_label} \
+  --disable-checkpointing;

ci/scripts/run_pytest.sh

@@ -28,6 +28,8 @@ Options:
     --skip-docs    Skip running tests in the docs directory
     --no-nbval     Skip jupyter notebook validation tests
     --skip-slow    Skip tests marked as slow (@pytest.mark.slow)
+    --only-slow    Only run tests marked as slow (@pytest.mark.slow)
+    --allow-no-tests    Allow sub-packages with no found tests (for example no slow tests if --only-slow is set)
 
 Note: Documentation tests (docs/) are only run when notebook validation
       is enabled (--no-nbval not set) and docs are not skipped
@@ -50,6 +52,8 @@ declare -a coverage_files
 SKIP_DOCS=false
 NO_NBVAL=false
 SKIP_SLOW=false
+ONLY_SLOW=false
+ALLOW_NO_TESTS=false
 error=false
 
 # Parse command line arguments
@@ -58,6 +62,8 @@ while (( $# > 0 )); do
         --skip-docs) SKIP_DOCS=true ;;
         --no-nbval) NO_NBVAL=true ;;
         --skip-slow) SKIP_SLOW=true ;;
+        --only-slow) ONLY_SLOW=true ;;
+        --allow-no-tests) ALLOW_NO_TESTS=true ;;
         -h|--help) usage ;;
         *) echo "Unknown option: $1" >&2; usage 1 ;;
     esac
@@ -85,6 +91,7 @@ PYTEST_OPTIONS=(
 )
 [[ "$NO_NBVAL" != true ]] && PYTEST_OPTIONS+=(--nbval-lax)
 [[ "$SKIP_SLOW" == true ]] && PYTEST_OPTIONS+=(-m "not slow")
+[[ "$ONLY_SLOW" == true ]] && PYTEST_OPTIONS+=(-m "slow")
 
 # Define test directories
 TEST_DIRS=(./sub-packages/bionemo-*/)
@@ -94,13 +101,32 @@ fi
 
 echo "Test directories: ${TEST_DIRS[*]}"
 
+clean_pycache() {
+    # Use the provided base directory or default to current directory
+    local base_dir="${1:-.}"
+    echo "Cleaning Python cache files in $base_dir..."
+    find "$base_dir" -regex '^.*\(__pycache__\|\.py[co]\)$' -delete
+}
+
 # Run tests with coverage
 for dir in "${TEST_DIRS[@]}"; do
     echo "Running pytest in $dir"
-
-    if ! pytest "${PYTEST_OPTIONS[@]}" --junitxml=$(basename $dir).junit.xml -o junit_family=legacy "$dir"; then
-        error=true
+    # Run pytest but don't exit on failure - we'll handle the exit code separately. This is needed because our script is
+    #  running in pipefail mode and pytest will exit with a non-zero exit code if it finds no tests.
+    { pytest "${PYTEST_OPTIONS[@]}" --junitxml=$(basename $dir).junit.xml -o junit_family=legacy "$dir"; exit_code=$?; } || true
+
+    if [[ $exit_code -ne 0 ]]; then
+        if [[ "$ALLOW_NO_TESTS" == true && $exit_code -eq 5 ]]; then
+            # Exit code 5 means no tests found, which is allowed if --allow-no-tests is set
+            echo "No tests found in $dir (exit code $exit_code) - continuing as --allow-no-tests is set"
+        else
+            echo "Error: pytest failed with exit code $exit_code"
+            error=true
+        fi
     fi
+
+    # Avoid duplicated pytest cache filenames.
+    clean_pycache "$dir"
 done
 
 # Exit with appropriate status

ci/scripts/utils.sh

@@ -20,10 +20,11 @@ check_git_repository() {
     if ! git diff-index --quiet HEAD --; then
         if [ $? -eq 128 ]; then
             echo "ERROR: Not in a git repository!" >&2
+            return 1
         else
-            echo "ERROR: Repository is dirty! Commit all changes before building the image!" >&2
+            echo "Warning: Repository is dirty! Commit all changes before building the image!" >&2
+            return 0
         fi
-        return 1
     fi
 }
 

docs/docs/user-guide/examples/bionemo-evo2/.gitignore

@@ -0,0 +1,14 @@
+# ignore temp files made by this tutorial
+# chromosome files
+*.fa
+*.fa.gz
+
+# config files
+*.yaml
+
+# directories created during these notebook runs.
+nemo2_evo2_1b_8k/
+preprocessed_data/
+pretraining_demo/
+brca1_fasta_files/
+brca1/