|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. |
| 3 | +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. |
| 4 | +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved |
| 5 | +# SPDX-License-Identifier: LicenseRef-Apache2 |
| 6 | +# |
| 7 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 8 | +# you may not use this file except in compliance with the License. |
| 9 | +# You may obtain a copy of the License at |
| 10 | +# |
| 11 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 12 | +# |
| 13 | +# Unless required by applicable law or agreed to in writing, software |
| 14 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 15 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 16 | +# See the License for the specific language governing permissions and |
| 17 | +# limitations under the License. |
| 18 | +import pytest |
| 19 | + |
| 20 | +from bionemo.evo2.run.common import small_training_cmd, small_training_finetune_cmd |
| 21 | +from bionemo.testing.subprocess_utils import run_command_in_subprocess |
| 22 | + |
| 23 | + |
| 24 | +@pytest.mark.timeout(512) # Optional: fail if the test takes too long. |
| 25 | +@pytest.mark.slow |
| 26 | +@pytest.mark.parametrize("with_peft", [True, False]) |
| 27 | +def test_train_evo2_finetune_runs_lora(tmp_path, with_peft: bool): |
| 28 | + """ |
| 29 | + This test runs the `train_evo2` command with mock data in a temporary directory. |
| 30 | + It uses the temporary directory provided by pytest as the working directory. |
| 31 | + The command is run in a subshell, and we assert that it returns an exit code of 0. |
| 32 | + """ |
| 33 | + num_steps = 2 |
| 34 | + # Note: The command assumes that `train_evo2` is in your PATH. |
| 35 | + command = small_training_cmd(tmp_path / "pretrain", max_steps=num_steps, val_check=num_steps) |
| 36 | + stdout_pretrain: str = run_command_in_subprocess(command=command, path=str(tmp_path)) |
| 37 | + assert "Restoring model weights from RestoreConfig(path='" not in stdout_pretrain |
| 38 | + |
| 39 | + log_dir = tmp_path / "pretrain" / "evo2" |
| 40 | + checkpoints_dir = log_dir / "checkpoints" |
| 41 | + tensorboard_dir = log_dir / "dev" |
| 42 | + |
| 43 | + # Check if logs dir exists |
| 44 | + assert log_dir.exists(), "Logs folder should exist." |
| 45 | + # Check if checkpoints dir exists |
| 46 | + assert checkpoints_dir.exists(), "Checkpoints folder does not exist." |
| 47 | + |
| 48 | + expected_checkpoint_suffix = f"{num_steps}.0-last" |
| 49 | + # Check if any subfolder ends with the expected suffix |
| 50 | + matching_subfolders = [ |
| 51 | + p for p in checkpoints_dir.iterdir() if p.is_dir() and (expected_checkpoint_suffix in p.name) |
| 52 | + ] |
| 53 | + |
| 54 | + assert matching_subfolders, ( |
| 55 | + f"No checkpoint subfolder ending with '{expected_checkpoint_suffix}' found in {checkpoints_dir}." |
| 56 | + ) |
| 57 | + |
| 58 | + # Check if directory with tensorboard logs exists |
| 59 | + assert tensorboard_dir.exists(), "TensorBoard logs folder does not exist." |
| 60 | + # Recursively search for files with tensorboard logger |
| 61 | + event_files = list(tensorboard_dir.rglob("events.out.tfevents*")) |
| 62 | + assert event_files, f"No TensorBoard event files found under {tensorboard_dir}" |
| 63 | + assert len(matching_subfolders) == 1, "Only one checkpoint subfolder should be found." |
| 64 | + if with_peft: |
| 65 | + result_dir = tmp_path / "lora_finetune" |
| 66 | + additional_args = "--lora-finetune" |
| 67 | + else: |
| 68 | + result_dir = tmp_path / "finetune" |
| 69 | + additional_args = "" |
| 70 | + |
| 71 | + command_finetune = small_training_finetune_cmd( |
| 72 | + result_dir, |
| 73 | + max_steps=num_steps, |
| 74 | + val_check=num_steps, |
| 75 | + prev_ckpt=matching_subfolders[0], |
| 76 | + create_tflops_callback=not with_peft, |
| 77 | + additional_args=additional_args, |
| 78 | + ) |
| 79 | + stdout_finetune: str = run_command_in_subprocess(command=command_finetune, path=str(tmp_path)) |
| 80 | + assert "Restoring model weights from RestoreConfig(path='" in stdout_finetune |
| 81 | + |
| 82 | + log_dir_ft = result_dir / "evo2" |
| 83 | + checkpoints_dir_ft = log_dir_ft / "checkpoints" |
| 84 | + tensorboard_dir_ft = log_dir_ft / "dev" |
| 85 | + |
| 86 | + # Check if logs dir exists |
| 87 | + assert log_dir_ft.exists(), "Logs folder should exist." |
| 88 | + # Check if checkpoints dir exists |
| 89 | + assert checkpoints_dir_ft.exists(), "Checkpoints folder does not exist." |
| 90 | + |
| 91 | + expected_checkpoint_suffix = f"{num_steps}.0-last" |
| 92 | + # Check if any subfolder ends with the expected suffix |
| 93 | + matching_subfolders_finetune = [ |
| 94 | + p for p in checkpoints_dir_ft.iterdir() if p.is_dir() and (expected_checkpoint_suffix in p.name) |
| 95 | + ] |
| 96 | + |
| 97 | + assert matching_subfolders_finetune, ( |
| 98 | + f"No checkpoint subfolder ending with '{expected_checkpoint_suffix}' found in {checkpoints_dir_ft}." |
| 99 | + ) |
| 100 | + |
| 101 | + # Check if directory with tensorboard logs exists |
| 102 | + assert tensorboard_dir_ft.exists(), "TensorBoard logs folder does not exist." |
| 103 | + # Recursively search for files with tensorboard logger |
| 104 | + event_files = list(tensorboard_dir_ft.rglob("events.out.tfevents*")) |
| 105 | + assert event_files, f"No TensorBoard event files found under {tensorboard_dir_ft}" |
| 106 | + |
| 107 | + assert len(matching_subfolders_finetune) == 1, "Only one checkpoint subfolder should be found." |
| 108 | + |
| 109 | + # With LoRA, test resuming from a saved LoRA checkpoint |
| 110 | + if with_peft: |
| 111 | + result_dir = tmp_path / "lora_finetune_resume" |
| 112 | + |
| 113 | + # Resume from LoRA checkpoint |
| 114 | + command_resume_finetune = small_training_finetune_cmd( |
| 115 | + result_dir, |
| 116 | + max_steps=num_steps, |
| 117 | + val_check=num_steps, |
| 118 | + prev_ckpt=matching_subfolders[0], |
| 119 | + create_tflops_callback=False, |
| 120 | + additional_args=f"--lora-finetune --lora-checkpoint-path {matching_subfolders_finetune[0]}", |
| 121 | + ) |
| 122 | + stdout_finetune: str = run_command_in_subprocess(command=command_resume_finetune, path=str(tmp_path)) |
| 123 | + |
| 124 | + log_dir_ft = result_dir / "evo2" |
| 125 | + checkpoints_dir_ft = log_dir_ft / "checkpoints" |
| 126 | + tensorboard_dir_ft = log_dir_ft / "dev" |
| 127 | + |
| 128 | + # Check if logs dir exists |
| 129 | + assert log_dir_ft.exists(), "Logs folder should exist." |
| 130 | + # Check if checkpoints dir exists |
| 131 | + assert checkpoints_dir_ft.exists(), "Checkpoints folder does not exist." |
0 commit comments