Remove ngcsdk re-installs. Remove Apex dependency.

Signed-off-by: Cory Ye <[email protected]>

Author: cspades

Dockerfile

@@ -140,7 +140,7 @@ RUN rm -rf /opt/pytorch/pytorch/third_party/onnx
 # environment, and does not use the current uv.lock file. Note that with python 3.12, we now need to set
 # UV_BREAK_SYSTEM_PACKAGES, since the pytorch base image has made the decision not to use a virtual environment and UV
 # does not respect the PIP_BREAK_SYSTEM_PACKAGES environment variable set in the base dockerfile.
-COPY --from=ghcr.io/astral-sh/uv:0.4.25 /uv /usr/local/bin/uv
+COPY --from=ghcr.io/astral-sh/uv:0.6.13 /uv /usr/local/bin/uv
 ENV UV_LINK_MODE=copy \
   UV_COMPILE_BYTECODE=1 \
   UV_PYTHON_DOWNLOADS=never \
@@ -175,7 +175,7 @@ uv pip install maturin --no-build-isolation
 git clone https://github.com/NVIDIA/nvidia-resiliency-ext
 uv pip install nvidia-resiliency-ext/
 rm -rf nvidia-resiliency-ext/
-# ngcsdk causes strange dependency conflicts that we will resolve later
+# ngcsdk causes strange dependency conflicts (ngcsdk requires protobuf<4, but nemo_toolkit requires protobuf==4.24.4, deleting it from the uv pip install prevents installation conflicts)
 sed -i "/ngcsdk/d" ./sub-packages/bionemo-core/pyproject.toml
 # Remove llama-index because bionemo doesn't use it and it adds CVEs to container
 sed -i "/llama-index/d" ./3rdparty/NeMo/requirements/requirements_nlp.txt
@@ -185,7 +185,7 @@ uv pip install --no-build-isolation \
 -r /requirements-cve.txt \
 -r /requirements-test.txt
 
-# Install back ngcsdk. Somehow doing it here avoids a large dependency loop
+# Install back ngcsdk, as a WAR for the protobuf version conflict with nemo_toolkit.
 uv pip install ngcsdk
 
 # Addressing security scan issue - CVE vulnerability https://github.com/advisories/GHSA-g4r7-86gm-pgqc The package is a
@@ -234,7 +234,7 @@ USER $USERNAME
 COPY --from=bionemo2-base --chown=$USERNAME:$USERNAME --chmod=777 \
   /usr/local/lib/python3.12/dist-packages /usr/local/lib/python3.12/dist-packages
 
-COPY --from=ghcr.io/astral-sh/uv:0.4.25 /uv /usr/local/bin/uv
+COPY --from=ghcr.io/astral-sh/uv:0.6.13 /uv /usr/local/bin/uv
 ENV UV_LINK_MODE=copy \
   UV_COMPILE_BYTECODE=0 \
   UV_PYTHON_DOWNLOADS=never \

sub-packages/bionemo-amplify/pyproject.toml

@@ -9,7 +9,7 @@ readme = "README.md"
 description = "A BioNeMo sub-package for training AMPLIFY models."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # internal

sub-packages/bionemo-core/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "BioNeMo core interfaces and PyTorch-related code."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # bionemo sub-packages

sub-packages/bionemo-esm2/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "BioNeMo ESM2 model."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # internal

sub-packages/bionemo-evo2/pyproject.toml

@@ -8,12 +8,14 @@ readme = "README.md"
 description = "Library containing data preprocessing, training, and inference tooling for Evo2."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
+    # internal
     "bionemo-noodles",
     "bionemo-core",
     "bionemo-llm",
+    # external
 ]
 
 [project.optional-dependencies]

sub-packages/bionemo-example_model/pyproject.toml

@@ -12,7 +12,7 @@ description = "BioNeMo example_model: Example model for documentation and tutori
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
 classifiers = ["Private :: Do Not Upload", "Programming Language :: Python :: 3.10"]
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     'bionemo-core',

sub-packages/bionemo-fw/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "BioNeMo Framework (FW): Production grade framework for AI-enabled Drug Discovery. Consists of all independently installable bionemo feature packages too."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     'bionemo-core',

sub-packages/bionemo-geneformer/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "BioNeMo Geneformer"
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # bionemo sub-packages

sub-packages/bionemo-geometric/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "BioNeMo component library for graphical neural networks (GNNs) solving drug discovery problems."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     'bionemo-core',

sub-packages/bionemo-llm/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "BioNeMo Large Language Model Components using NeMo and Megatron"
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # bionemo sub-packages

sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/lightning.py

@@ -24,7 +24,6 @@
 from nemo.lightning.megatron_parallel import DataT, MegatronLossReduction
 from nemo.lightning.pytorch.optim import MegatronOptimizerModule
 from torch import Tensor
-from torch.optim import Adam
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 from bionemo.llm.lightning import (
@@ -40,7 +39,6 @@
     "BertBatch",
     "BertModel",
     "SequenceBatch",
-    "bert_default_optimizer",
     "bert_forward_step",
     "biobert_data_step",
     "biobert_lightning_module",
@@ -182,19 +180,6 @@ def biobert_lightning_module(
     )
 
 
-def bert_default_optimizer(model: torch.nn.Module):
-    """Returns the default optimizer for the BERT model.
-
-    Args:
-        model: The BERT model.
-
-    Returns:
-        The default optimizer initialized for this BERT module's parameters.
-        Uses a learning rate of 1e-4 and weight decay of 1e-2.
-    """
-    return Adam(model.parameters(), lr=1e-4, weight_decay=0.01)
-
-
 def get_batch_on_this_context_parallel_rank(batch: Dict[str, Tensor], in_place: bool = True) -> Dict[str, Tensor]:
     """Ensures that the input batch is in the right format for context parallel rank.
 

sub-packages/bionemo-moco/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "BioNeMo Modular Co-Design: Making building Diffusion and Flow Matching generative models easier"
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # bionemo sub-packages

sub-packages/bionemo-noodles/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "Python wrapper around [noodles](https://github.com/zaeleus/noodles)."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # internal

sub-packages/bionemo-scdl/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "SCDL, a Dataset class for Single Cell data."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # bionemo sub-packages

sub-packages/bionemo-size-aware-batching/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "Provides a simple way to create mini-batches in a memory consumption-aware manner, making it useful for tasks like training models on datasets with varying memory requirements."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # bionemo sub-packages

sub-packages/bionemo-testing/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "Utilities aiding test creation for BioNeMo sub-packages."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # bionemo sub-packages

sub-packages/bionemo-webdatamodule/pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 description = "PyTorch Lightning Data Module for WebDataset files."
 authors = [{ name = "BioNeMo Team", email = "[email protected]" }]
 requires-python = ">=3.10"
-license-files = ["LICENSE"]
+license = { file = "LICENSE" }
 dynamic = ["version"]
 dependencies = [
     # bionemo sub-packages