[FEA][webdatamodule]: support webdataset invocable

[DejunL]

Summary

The webdataset invocables are member functions of the webdataset or webloader class that return the the same instance their are invoked from. Previously webdatamodule does its own computation on the epoch length but now can rely on the user input directly using webdataset/webloader.with_epoch() due to this new additional feature

Details

Describe your changes. You can be more detailed and descriptive here. If it is a code change, Be sure to answer:

• What is changing?

1. Add user input webdataset invocables and apply them to the webdataset/webloader objects upon construction.
2. Remove the computation related to epoch length and update the tests correspondingly.
3. Update the doc to reflect the change

• What is the new or fixed functionality?

(see the summary above)

• Why or when would someone want to use these changes?

These invocables are parts of the webdataset usage and we have been using them in diffdock to set up the epoch length and other properties of the dataset and dataloader objects.

• How can someone use these changes?

(see the updated README.md regarding the invoke_wds and invoke_wld arguments on webdatamodule)

Usage

How does a user interact with the changed code?
(see the updated README.md regarding the invoke_wds and invoke_wld arguments on webdatamodule)

Testing

Tests for these changes can be run via:

pytest -v subpackages/bionemo-webdatamodule/tests

Most of the changes to files with extensions *.py, *.yaml, *.yml, Dockerfile* or requirements.txt DO REQUIRE both pytest- and jet- CI stages.

• Did you review the Before your PR is "Ready for review" section before asking for review?
• Did you make sure your changes have tests? Did you test your changes locally?
• Can you add the SKIP_CI label to your PR?
• Can you add the PYTEST_NOT_REQUIRED label to your PR?
• Can you add the JET_NOT_REQUIRED label to your PR?

[DejunL]

/build-ci

[DejunL]

/build-ci

[DejunL]

/build-ci

[DejunL]

/build-ci

[DejunL]

keep getting CI failure in docker image building:

 build_bionemo_image_multiarch job Failed 

which seems to result from:

10:08:48  2024-12-06T18:08:48.658Z [ERROR] agent.auth.handler: error authenticating: error="Error making API request | Namespace: swgpu-bionemo | URL: PUT https://stg.internal.vault.nvidia.com/v1/auth/jwt/nvidia/jenkins/bionemo-external-bionemo-fw/login | Code: 400 | Request Id: 27716135-c8cf-05ce-6397-90153d35d777 | Errors: * error validating claims: token claim \"job_name\" with value \"build_bionemo_image_multiarch\" not included in the bound_claims allowed by role " backoff=10s
10:08:48  2024-12-06T18:08:48.658Z [ERROR] agent.auth.handler: encountered 400 status code error w/ exit-after-auth-failure enabled
10:08:48  2024-12-06T18:08:48.659Z [ERROR] agent: runtime error encountered: error="Error making API request | Namespace: swgpu-bionemo | URL: PUT https://stg.internal.vault.nvidia.com/v1/auth/jwt/nvidia/jenkins/bionemo-external-bionemo-fw/login | Code: 400 | Request Id: 27716135-c8cf-05ce-6397-90153d35d777 | Errors: * error validating claims: token claim \"job_name\" with value \"build_bionemo_image_multiarch\" not included in the bound_claims allowed by role "
10:08:48  Error encountered during run, refer to logs for more details.

[DejunL]

/build-ci

[DejunL]

/build-ci

[DejunL]

/build-ci

[DejunL]

/build-ci

[DejunL]

@pstjohn CI failed due to esm2 tests:

[2024-12-07T00:54:00.823Z] FAILED sub-packages/bionemo-esm2/tests/bionemo/esm2/scripts/test_pydantic_train.py::test_pretrain_pydantic_cli - Exception: Pretrain script failed:
[2024-12-07T00:54:00.823Z] cmd_str='bionemo-esm2-train --conf /tmp/pytest-of-root/pytest-1/test_pretrain_pydantic_cli0/results/test_config.yaml'
[2024-12-07T00:54:00.823Z] result.stdout=b'[NeMo I 2024-12-07 00:53:37 config_models:186] Mutating apply_query_key_layer_scaling and core_attention_override based on biobert_spec_option..\n[NeMo I 2024-12-07 00:53:37 nemo_logger:145] Experiments will be logged at /tmp/pytest-of-root/pytest-1/test_pretrain_pydantic_cli0/results/default_experiment/dev\n[NeMo I 2024-12-07 00:53:37 megatron_strategy:315] Fixing mis-match between ddp-config & mcore-optimizer config\n[NeMo I 2024-12-07 00:53:37 megatron_init:396] Rank 0 has data parallel group : [0]\n[NeMo I 2024-12-07 00:53:37 megatron_init:402] Rank 0 has combined group of data parallel and context parallel : [0]\n[NeMo I 2024-12-07 00:53:37 megatron_init:407] All data parallel group ranks with context parallel combined: [[0]]\n[NeMo I 2024-12-07 00:53:37 megatron_init:410] Ranks 0 has data parallel rank: 0\n[NeMo I 2024-12-07 00:53:37 megatron_init:418] Rank 0 has context parallel group: [0]\n[NeMo I 2024-12-07 00:53:37 megatron_init:421] All context parallel group ranks: [[0]]\n[NeMo I 2024-12-07 00:53:37 megatron_init:422] Ranks 0 has context parallel rank: 0\n[NeMo I 2024-12-07 00:53:37 megatron_init:429] Rank 0 has model parallel group: [0]\n[NeMo I 2024-12-07 00:53:37 megatron_init:430] All model parallel group ranks: [[0]]\n[NeMo I 2024-12-07 00:53:37 megatron_init:439] Rank 0 has tensor model parallel group: [0]\n[NeMo I 2024-12-07 00:53:37 megatron_init:443] All tensor model parallel group ranks: [[0]]\n[NeMo I 2024-12-07 00:53:37 megatron_init:444] Rank 0 has tensor model parallel rank: 0\n[NeMo I 2024-12-07 00:53:37 megatron_init:464] Rank 0 has pipeline model parallel group: [0]\n[NeMo I 2024-12-07 00:53:37 megatron_init:476] Rank 0 has embedding group: [0]\n[NeMo I 2024-12-07 00:53:37 megatron_init:482] All pipeline model parallel group ranks: [[0]]\n[NeMo I 2024-12-07 00:53:37 megatron_init:483] Rank 0 has pipeline model parallel rank 0\n[NeMo I 2024-12-07 00:53:37 megatron_init:484] All embedding group ranks: [[0]]\n[NeMo I 2024-12-07 00:53:37 megatron_init:485] Rank 0 has embedding rank: 0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11448 [0] NCCL INFO Bootstrap : Using eth0:****0<0>\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11448 [0] NCCL INFO NET/Plugin: Failed to find ncclNetPlugin_v7 symbol.\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11448 [0] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v6 (v6)\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11448 [0] NCCL INFO NET/Plugin: Failed to find ncclCollNetPlugin_v7 symbol.\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11448 [0] NCCL INFO NET/Plugin: Loaded coll plugin SHARP (v6)\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11448 [0] NCCL INFO cudaDriverVersion 12060\nNCCL version 2.19.4+cuda12.3\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO P2P plugin IBext\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO NET/IB : No device found.\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO NET/IB : No device found.\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO NET/Socket : Using [0]eth0:****0<0>\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Using non-device net plugin version 0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Using network Socket\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO comm 0x55f8130c3220 rank 0 nranks 1 cudaDev 0 nvmlDev 0 busId c1000 commId 0x3c65e5dcaec8812e - Init START\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Setting affinity for GPU 0 to ffffffff\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 00/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 01/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 02/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 03/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 04/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 05/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 06/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 07/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 08/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 09/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 10/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 11/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 12/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 13/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 14/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 15/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 16/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 17/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 18/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 19/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 20/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 21/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 22/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 23/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 24/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 25/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 26/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 27/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 28/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 29/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 30/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Channel 31/32 :    0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Trees [0] -1/-1/-1->0->-1 [1] -1/-1/-1->0->-1 [2] -1/-1/-1->0->-1 [3] -1/-1/-1->0->-1 [4] -1/-1/-1->0->-1 [5] -1/-1/-1->0->-1 [6] -1/-1/-1->0->-1 [7] -1/-1/-1->0->-1 [8] -1/-1/-1->0->-1 [9] -1/-1/-1->0->-1 [10] -1/-1/-1->0->-1 [11] -1/-1/-1->0->-1 [12] -1/-1/-1->0->-1 [13] -1/-1/-1->0->-1 [14] -1/-1/-1->0->-1 [15] -1/-1/-1->0->-1 [16] -1/-1/-1->0->-1 [17] -1/-1/-1->0->-1 [18] -1/-1/-1->0->-1 [19] -1/-1/-1->0->-1 [20] -1/-1/-1->0->-1 [21] -1/-1/-1->0->-1 [22] -1/-1/-1->0->-1 [23] -1/-1/-1->0->-1 [24] -1/-1/-1->0->-1 [25] -1/-1/-1->0->-1 [26] -1/-1/-1->0->-1 [27] -1/-1/-1->0->-1 [28] -1/-1/-1->0->-1 [29] -1/-1/-1->0->-1 [30] -1/-1/-1->0->-1 [31] -1/-1/-1->0->-1\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO P2P Chunksize set to 131072\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Connected all rings\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO Connected all trees\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO 32 coll channels, 0 nvls channels, 32 p2p channels, 32 p2p channels per peer\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11639 [0] NCCL INFO comm 0x55f8130c3220 rank 0 nranks 1 cudaDev 0 nvmlDev 0 busId c1000 commId 0x3c65e5dcaec8812e - Init COMPLETE\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11640 [0] NCCL INFO [Service thread] Connection closed by localRank 0\ntest-pytest-1004-qdghp-fpcn0-rjwsh:11448:11448 [0] NCCL INFO comm 0x55f8130c3220 rank 0 nranks 1 cudaDev 0 busId c1000 - Abort COMPLETE\n'
[2024-12-07T00:54:00.824Z] result.stderr=b'2024-12-07 00:53:35 - faiss.loader - INFO - Loading faiss with AVX2 support.\n2024-12-07 00:53:35 - faiss.loader - INFO - Successfully loaded faiss with AVX2 support.\n[NeMo W 2024-12-07 00:53:36 nemo_logging:361] /usr/local/lib/python3.10/dist-packages/pyannote/core/notebook.py:134: MatplotlibDeprecationWarning: The get_cmap function was deprecated in Matplotlib 3.7 and will be removed two minor releases later. Use ``matplotlib.colormaps[name]`` or ``matplotlib.colormaps.get_cmap(obj)`` instead.\n      cm = get_cmap("Set1")\n    \n[NeMo W 2024-12-07 00:53:36 nemo_logging:361] /usr/local/lib/python3.10/dist-packages/pydantic/_internal/_config.py:291: PydanticDeprecatedSince20: Support for class-based `config` is deprecated, use ConfigDict instead. Deprecated in Pydantic V2.0 to be removed in V3.0. See Pydantic V2 Migration Guide at [https://errors.pydantic.dev/2.9/migration/\n](https://errors.pydantic.dev/2.9/migration//n)      warnings.warn(DEPRECATION_MESSAGE, DeprecationWarning)\n    \n[NeMo W 2024-12-07 00:53:37 ssm:31] The package `megatron.core` was not imported in this environment which is needed for SSMs.\n[NeMo W 2024-12-07 00:53:37 train:187] Mutating training_config.save_every_n_steps to be equal to val_check_interval.\nWARNING: Logging before flag parsing goes to stderr.\nI1207 00:53:37.614806 140431552213440 rank_zero.py:63] Trainer already configured with model summary callbacks: [<class \'lightning.pytorch.callbacks.rich_model_summary.RichModelSummary\'>]. Skipping setting a default `ModelSummary` callback.\nI1207 00:53:37.645136 140431552213440 rank_zero.py:63] GPU available: True (cuda), used: True\nI1207 00:53:37.645528 140431552213440 rank_zero.py:63] TPU available: False, using: 0 TPU cores\nI1207 00:53:37.645570 140431552213440 rank_zero.py:63] HPU available: False, using: 0 HPUs\n[NeMo W 2024-12-07 00:53:37 logger_utils:90] User-set tensorboard is currently turned off. Internally one may still be set by NeMo2.\n[NeMo W 2024-12-07 00:53:37 nemo_logger:173] "update_logger_directory" is True. Overwriting tensorboard logger "save_dir" to /tmp/pytest-of-root/pytest-1/test_pretrain_pydantic_cli0/results\n[NeMo W 2024-12-07 00:53:37 nemo_logger:180] "update_logger_directory" is True. Overwriting wandb logger "save_dir" to /tmp/pytest-of-root/pytest-1/test_pretrain_pydantic_cli0/results/default_experiment\n[NeMo W 2024-12-07 00:53:37 nemo_logger:189] The Trainer already contains a ModelCheckpoint callback. This will be overwritten.\n[NeMo W 2024-12-07 00:53:37 nemo_logger:212] The checkpoint callback was told to monitor a validation value and trainer\'s max_steps was set to 10. Please ensure that max_steps will run for at least 1 epochs to ensure that checkpointing will not error out.\nInitializing distributed: GLOBAL_RANK: 0, MEMBER: 1/1\nI1207 00:53:37.659220 140431552213440 rank_zero.py:63] ----------------------------------------------------------------------------------------------------\ndistributed_backend=nccl\nAll distributed processes registered. Starting with 1 processes\n----------------------------------------------------------------------------------------------------\n\nTraceback (most recent call last):\n  File "/usr/local/bin/bionemo-esm2-train", line 8, in <module>\n    sys.exit(main())\n  File "/usr/local/lib/python3.10/dist-packages/bionemo/esm2/run/main.py", line 126, in main\n    train(\n  File "/usr/local/lib/python3.10/dist-packages/bionemo/llm/train.py", line 241, in train\n    llm.train(\n  File "/usr/local/lib/python3.10/dist-packages/nemo/collections/llm/api.py", line 106, in train\n    trainer.fit(model, data)\n  File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py", line 538, in fit\n    call._call_and_handle_interrupt(\n  File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt\n    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)\n  File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch\n    return function(*args, **kwargs)\n  File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py", line 574, in _fit_impl\n    self._run(model, ckpt_path=ckpt_path)\n  File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py", line 943, in _run\n    call._call_setup_hook(self)  # allow user to set up LightningModule in accelerator environment\n  File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/call.py", line 96, in _call_setup_hook\n    if hasattr(logger, "experiment"):\n  File "/usr/local/lib/python3.10/dist-packages/lightning/fabric/loggers/logger.py", line 118, in experiment\n    return fn(self)\n  File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/loggers/wandb.py", line 406, in experiment\n    self._experiment = wandb.init(**self._wandb_init)\n  File "/usr/local/lib/python3.10/dist-packages/wandb/sdk/wandb_init.py", line 1258, in init\n    init_settings.root_dir = dir  # type: ignore\n  File "/usr/local/lib/python3.10/dist-packages/pydantic/main.py", line 881, in __setattr__\n    self.__pydantic_validator__.validate_assignment(self, name, value)\npydantic_core._pydantic_core.ValidationError: 1 validation error for Settings\nroot_dir\n  Input should be a valid string [type=string_type, input_value=PosixPath(\'/tmp/pytest-of...lts/default_experiment\'), input_type=PosixPath]\n    For further information visit [https://](https://errors.pydantic.dev/2.9/v/string_type/n)

[DejunL]

/build-ci

[DejunL]

/build-ci

[DejunL]

CI still fails with:

Fix End of Files.........................................................Failed
- hook id: end-of-file-fixer
- exit code: 1
- files were modified by this hook

Fixing ci/docker/clobber_dependencies_into_requirements_txt.sh

[DejunL]

/build-ci