evo2 train unit test

[dorotat-nv]

Description

Slightly refactoring train script for evo2 to better handle unit testing and a bug fix

Type of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Refactor
• Documentation update
• Other (please describe):

CI Pipeline Configuration

Configure CI behavior by applying the relevant labels:

• SKIP_CI - Skip all continuous integration tests
• INCLUDE_NOTEBOOKS_TESTS - Execute notebook validation tests in pytest
• INCLUDE_SLOW_TESTS - Execute tests labelled as slow in pytest for extensive testing

Note

By default, the notebooks validation tests are skipped unless explicitly enabled.

Usage

TODO: Add code snippet

Pre-submit Checklist

• I have tested these changes locally
• I have updated the documentation accordingly
• I have added/updated tests as needed
• All existing tests pass successfully

[codecov-commenter]

❌ 5 Tests Failed:

Tests completed	Failed	Passed	Skipped
939	5	934	18

View the top 3 failed test(s) by shortest run time

sub-packages/bionemo-evo2/tests/bionemo/evo2/test_evo2.py::test_golden_values_top_k_logits_and_cosine_similarity[16384]

Stack Traces | 0.001s run time

seq_len = 16384

    @pytest.mark.parametrize("seq_len", [8_192, 16_384])
    def test_golden_values_top_k_logits_and_cosine_similarity(seq_len: int):
        try:
            # TODO (dorotat) remove PBSS source once the model is available on NGC
>           evo2_1b_checkpoint_weights: Path = load("evo2/1b-8k:1.0") / "weights"

.../bionemo/evo2/test_evo2.py:71: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

model_or_data_tag = 'evo2/1b-8k:1.0', source = 'ngc'
resources = {'esm2/3b:2.0': Resource(tag='esm2/3b:2.0', ngc='nvidia/clara/esm2nv3b:2.0', ngc_registry='model', pbss=AnyUrl('s3://b...horbani', email='farhadr@nvidia.com'), description='Test data for ESM2 inference.', unpack=None, decompress=None), ...}
cache_dir = PosixPath('.../github/home/.cache/bionemo')

    def load(
        model_or_data_tag: str,
        source: SourceOptions = DEFAULT_SOURCE,
        resources: dict[str, Resource] | None = None,
        cache_dir: Path | None = None,
    ) -> Path:
        """Download a resource from PBSS or NGC.
    
        Args:
            model_or_data_tag: A pointer to the desired resource. Must be a key in the resources dictionary.
            source: Either "pbss" (NVIDIA-internal download) or "ngc" (NVIDIA GPU Cloud). Defaults to "pbss".
            resources: A custom dictionary of resources. If None, the default resources will be used. (Mostly for testing.)
            cache_dir: The directory to store downloaded files. Defaults to BIONEMO_CACHE_DIR. (Mostly for testing.)
    
        Raises:
            ValueError: If the desired tag was not found, or if an NGC url was requested but not provided.
    
        Returns:
            A Path object pointing either at the downloaded file, or at a decompressed folder containing the
            file(s).
    
        Examples:
            For a resource specified in 'filename.yaml' with tag 'tag', the following will download the file:
            >>> load("filename/tag")
            PosixPath(.../tmp/bionemo/downloaded-file-name)
        """
        if resources is None:
            resources = get_all_resources()
    
        if cache_dir is None:
            cache_dir = BIONEMO_CACHE_DIR
    
        if model_or_data_tag not in resources:
            raise ValueError(f"Resource '{model_or_data_tag}' not found.")
    
        if source == "ngc" and resources[model_or_data_tag].ngc is None:
>           raise ValueError(f"Resource '{model_or_data_tag}' does not have an NGC URL.")
E           ValueError: Resource 'evo2/1b-8k:1.0' does not have an NGC URL.

.../local/lib/python3.12.../core/data/load.py:175: ValueError

During handling of the above exception, another exception occurred:

seq_len = 16384

    @pytest.mark.parametrize("seq_len", [8_192, 16_384])
    def test_golden_values_top_k_logits_and_cosine_similarity(seq_len: int):
        try:
            # TODO (dorotat) remove PBSS source once the model is available on NGC
            evo2_1b_checkpoint_weights: Path = load("evo2/1b-8k:1.0") / "weights"
            # TODO (dorotat) remove PBSS source once the model is available on NGC
            gold_standard_no_fp8 = load("evo2/1b-8k-nofp8-te-goldvalue-testdata-A6000:1.0")
        except ValueError as e:
            if e.args[0].endswith("does not have an NGC URL."):
>               raise ValueError(
                    "Please re-run test with `BIONEMO_DATA_SOURCE=pbss py.test ...`, "
                    "one or more files are missing from ngc."
                )
E               ValueError: Please re-run test with `BIONEMO_DATA_SOURCE=pbss py.test ...`, one or more files are missing from ngc.

.../bionemo/evo2/test_evo2.py:76: ValueError

sub-packages/bionemo-evo2/tests/bionemo/evo2/test_evo2.py::test_golden_values_top_k_logits_and_cosine_similarity[8192]

Stack Traces | 0.001s run time

seq_len = 8192

    @pytest.mark.parametrize("seq_len", [8_192, 16_384])
    def test_golden_values_top_k_logits_and_cosine_similarity(seq_len: int):
        try:
            # TODO (dorotat) remove PBSS source once the model is available on NGC
>           evo2_1b_checkpoint_weights: Path = load("evo2/1b-8k:1.0") / "weights"

.../bionemo/evo2/test_evo2.py:71: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

model_or_data_tag = 'evo2/1b-8k:1.0', source = 'ngc'
resources = {'esm2/3b:2.0': Resource(tag='esm2/3b:2.0', ngc='nvidia/clara/esm2nv3b:2.0', ngc_registry='model', pbss=AnyUrl('s3://b...horbani', email='farhadr@nvidia.com'), description='Test data for ESM2 inference.', unpack=None, decompress=None), ...}
cache_dir = PosixPath('.../github/home/.cache/bionemo')

    def load(
        model_or_data_tag: str,
        source: SourceOptions = DEFAULT_SOURCE,
        resources: dict[str, Resource] | None = None,
        cache_dir: Path | None = None,
    ) -> Path:
        """Download a resource from PBSS or NGC.
    
        Args:
            model_or_data_tag: A pointer to the desired resource. Must be a key in the resources dictionary.
            source: Either "pbss" (NVIDIA-internal download) or "ngc" (NVIDIA GPU Cloud). Defaults to "pbss".
            resources: A custom dictionary of resources. If None, the default resources will be used. (Mostly for testing.)
            cache_dir: The directory to store downloaded files. Defaults to BIONEMO_CACHE_DIR. (Mostly for testing.)
    
        Raises:
            ValueError: If the desired tag was not found, or if an NGC url was requested but not provided.
    
        Returns:
            A Path object pointing either at the downloaded file, or at a decompressed folder containing the
            file(s).
    
        Examples:
            For a resource specified in 'filename.yaml' with tag 'tag', the following will download the file:
            >>> load("filename/tag")
            PosixPath(.../tmp/bionemo/downloaded-file-name)
        """
        if resources is None:
            resources = get_all_resources()
    
        if cache_dir is None:
            cache_dir = BIONEMO_CACHE_DIR
    
        if model_or_data_tag not in resources:
            raise ValueError(f"Resource '{model_or_data_tag}' not found.")
    
        if source == "ngc" and resources[model_or_data_tag].ngc is None:
>           raise ValueError(f"Resource '{model_or_data_tag}' does not have an NGC URL.")
E           ValueError: Resource 'evo2/1b-8k:1.0' does not have an NGC URL.

.../local/lib/python3.12.../core/data/load.py:175: ValueError

During handling of the above exception, another exception occurred:

seq_len = 8192

    @pytest.mark.parametrize("seq_len", [8_192, 16_384])
    def test_golden_values_top_k_logits_and_cosine_similarity(seq_len: int):
        try:
            # TODO (dorotat) remove PBSS source once the model is available on NGC
            evo2_1b_checkpoint_weights: Path = load("evo2/1b-8k:1.0") / "weights"
            # TODO (dorotat) remove PBSS source once the model is available on NGC
            gold_standard_no_fp8 = load("evo2/1b-8k-nofp8-te-goldvalue-testdata-A6000:1.0")
        except ValueError as e:
            if e.args[0].endswith("does not have an NGC URL."):
>               raise ValueError(
                    "Please re-run test with `BIONEMO_DATA_SOURCE=pbss py.test ...`, "
                    "one or more files are missing from ngc."
                )
E               ValueError: Please re-run test with `BIONEMO_DATA_SOURCE=pbss py.test ...`, one or more files are missing from ngc.

.../bionemo/evo2/test_evo2.py:76: ValueError

sub-packages/bionemo-evo2/tests/bionemo/evo2/run/test_predict.py::test_train_evo2_runs

Stack Traces | 0.002s run time

tmp_path = PosixPath('.../pytest-of-root/pytest-2/test_train_evo2_runs0')
num_sequences = 5, target_sequence_lengths = [3149, 3140, 1024, 3149, 3149]

    def test_train_evo2_runs(
        tmp_path, num_sequences: int = 5, target_sequence_lengths: list[int] = [3149, 3140, 1024, 3149, 3149]
    ):
        """
        This test runs the `predict_evo2` command with mock data in a temporary directory.
        It uses the temporary directory provided by pytest as the working directory.
        The command is run in a subshell, and we assert that it returns an exit code of 0.
        """
        fasta_file_path = tmp_path / "test.fasta"
        create_fasta_file(
            fasta_file_path, num_sequences, sequence_lengths=target_sequence_lengths, repeating_dna_pattern=ALU_SEQUENCE
        )
        # Create a mock data directory.
        open_port = find_free_network_port()
        # a local copy of the environment
        env = dict(**os.environ)
        env["MASTER_PORT"] = str(open_port)
        try:
>           checkpoint_path = load("evo2/1b-8k:1.0")

.../evo2/run/test_predict.py:51: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

model_or_data_tag = 'evo2/1b-8k:1.0', source = 'ngc'
resources = {'esm2/3b:2.0': Resource(tag='esm2/3b:2.0', ngc='nvidia/clara/esm2nv3b:2.0', ngc_registry='model', pbss=AnyUrl('s3://b...horbani', email='farhadr@nvidia.com'), description='Test data for ESM2 inference.', unpack=None, decompress=None), ...}
cache_dir = PosixPath('.../github/home/.cache/bionemo')

    def load(
        model_or_data_tag: str,
        source: SourceOptions = DEFAULT_SOURCE,
        resources: dict[str, Resource] | None = None,
        cache_dir: Path | None = None,
    ) -> Path:
        """Download a resource from PBSS or NGC.
    
        Args:
            model_or_data_tag: A pointer to the desired resource. Must be a key in the resources dictionary.
            source: Either "pbss" (NVIDIA-internal download) or "ngc" (NVIDIA GPU Cloud). Defaults to "pbss".
            resources: A custom dictionary of resources. If None, the default resources will be used. (Mostly for testing.)
            cache_dir: The directory to store downloaded files. Defaults to BIONEMO_CACHE_DIR. (Mostly for testing.)
    
        Raises:
            ValueError: If the desired tag was not found, or if an NGC url was requested but not provided.
    
        Returns:
            A Path object pointing either at the downloaded file, or at a decompressed folder containing the
            file(s).
    
        Examples:
            For a resource specified in 'filename.yaml' with tag 'tag', the following will download the file:
            >>> load("filename/tag")
            PosixPath(.../tmp/bionemo/downloaded-file-name)
        """
        if resources is None:
            resources = get_all_resources()
    
        if cache_dir is None:
            cache_dir = BIONEMO_CACHE_DIR
    
        if model_or_data_tag not in resources:
            raise ValueError(f"Resource '{model_or_data_tag}' not found.")
    
        if source == "ngc" and resources[model_or_data_tag].ngc is None:
>           raise ValueError(f"Resource '{model_or_data_tag}' does not have an NGC URL.")
E           ValueError: Resource 'evo2/1b-8k:1.0' does not have an NGC URL.

.../local/lib/python3.12.../core/data/load.py:175: ValueError

During handling of the above exception, another exception occurred:

tmp_path = PosixPath('.../pytest-of-root/pytest-2/test_train_evo2_runs0')
num_sequences = 5, target_sequence_lengths = [3149, 3140, 1024, 3149, 3149]

    def test_train_evo2_runs(
        tmp_path, num_sequences: int = 5, target_sequence_lengths: list[int] = [3149, 3140, 1024, 3149, 3149]
    ):
        """
        This test runs the `predict_evo2` command with mock data in a temporary directory.
        It uses the temporary directory provided by pytest as the working directory.
        The command is run in a subshell, and we assert that it returns an exit code of 0.
        """
        fasta_file_path = tmp_path / "test.fasta"
        create_fasta_file(
            fasta_file_path, num_sequences, sequence_lengths=target_sequence_lengths, repeating_dna_pattern=ALU_SEQUENCE
        )
        # Create a mock data directory.
        open_port = find_free_network_port()
        # a local copy of the environment
        env = dict(**os.environ)
        env["MASTER_PORT"] = str(open_port)
        try:
            checkpoint_path = load("evo2/1b-8k:1.0")
        except ValueError as e:
            if e.args[0].endswith("does not have an NGC URL."):
>               raise ValueError(
                    "Please re-run test with `BIONEMO_DATA_SOURCE=pbss py.test ...`, "
                    "one or more files are missing from ngc."
                )
E               ValueError: Please re-run test with `BIONEMO_DATA_SOURCE=pbss py.test ...`, one or more files are missing from ngc.

.../evo2/run/test_predict.py:54: ValueError

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}