Wandb integration

[olachinkei]

Hi @jstjohn
Please review my PR and test this instead of me (sorry for the inconvenience)

Summary

High level summary of changes. Try to keep this as short and informative as possible: less is more.

• Add wandb integration

Details

Describe your changes. You can be more detailed and descriptive here.

• Added more wandb arguments so that wandb users can user wandb more flexibly (such as resume)

Usage

How does a user interact with the changed code?

TEST_DATA_DIR=$(bionemo_test_data_path single_cell/testdata-20240506 --source pbss); \
python  \
    scripts/singlecell/geneformer/train.py     \
    --data-dir ${TEST_DATA_DIR}/cellxgene_2023-12-15_small/processed_data    \
    --result-dir ./results     \
    --experiment-name test_experiment     \
    --num-gpus 1  \
    --num-nodes 1 \
    --val-check-interval 10 \
    --num-dataset-workers 0 \
    --num-steps 55 \
    --seq-length 128 \
    --limit-val-batches 2 \
    --micro-batch-size 2 \
    --wandb_entity: <your entity> \
    --wandb_project: "test-project" \
    --wandb_tags: "test" \
    --wandb_log_model

After running the first run, please test resume function with the following command

TEST_DATA_DIR=$(bionemo_test_data_path single_cell/testdata-20240506 --source pbss); \
python  \
    scripts/singlecell/geneformer/train.py     \
    --data-dir ${TEST_DATA_DIR}/cellxgene_2023-12-15_small/processed_data    \
    --result-dir ./results     \
    --experiment-name test_experiment     \
    --num-gpus 1  \
    --num-nodes 1 \
    --val-check-interval 10 \
    --num-dataset-workers 0 \
    --num-steps 55 \
    --seq-length 128 \
    --limit-val-batches 2 \
    --micro-batch-size 2 \
    --wandb_entity <your entity> \
    --wandb_project test-project \
    --wandb_tags test \
    --wandb_id <first run_id> \
    --wandb_log_model

[jstjohn]

/build-ci

[jstjohn]

@olachinkei Sorry Kei, I did not get a chance to run this test today but I did just kick off our CI on this PR so we can see how that does.

[olachinkei]

Thank you, @jstjohn !
Just to confirm the next steps, is it now my turn to take action? Or should I wait for you to run it?

[jstjohn]

/build-ci

[tshimko-nv]

Nice to see the increased functionality! Approved.

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[jstjohn]

/build-ci