Add test coverage for fine-tuning

Author: jstjohn

sub-packages/bionemo-evo2/tests/bionemo/evo2/run/test_train.py

@@ -61,6 +61,17 @@ def small_training_cmd(path, max_steps, val_check, devices: int = 1, additional_
     return cmd
 
 
+def small_training_finetune_cmd(path, max_steps, val_check, prev_ckpt, devices: int = 1, additional_args: str = ""):
+    cmd = (
+        f"train_evo2 --mock-data --result-dir {path} --devices {devices} "
+        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* --limit-val-batches 1 "
+        "--no-activation-checkpointing --add-bias-output --create-tensorboard-logger --create-tflops-callback "
+        f"--max-steps {max_steps} --warmup-steps 1 --val-check-interval {val_check} --limit-val-batches 1 "
+        f"--seq-length 16 --hidden-dropout 0.1 --attention-dropout 0.1 {additional_args} --ckpt-dir {prev_ckpt}"
+    )
+    return cmd
+
+
 def small_training_mamba_cmd(path, max_steps, val_check, devices: int = 1, additional_args: str = ""):
     cmd = (
         f"train_evo2 --mock-data --result-dir {path} --devices {devices} "
@@ -72,9 +83,22 @@ def small_training_mamba_cmd(path, max_steps, val_check, devices: int = 1, addit
     return cmd
 
 
+def small_training_mamba_finetune_cmd(
+    path, max_steps, val_check, prev_ckpt, devices: int = 1, additional_args: str = ""
+):
+    cmd = (
+        f"train_evo2 --mock-data --result-dir {path} --devices {devices} "
+        "--model-size hybrid_mamba_8b --num-layers 2 --hybrid-override-pattern M- --limit-val-batches 1 "
+        "--no-activation-checkpointing --create-tensorboard-logger --create-tflops-callback "
+        f"--max-steps {max_steps} --warmup-steps 1 --val-check-interval {val_check} --limit-val-batches 1 "
+        f"--seq-length 16 --hidden-dropout 0.1 --attention-dropout 0.1 {additional_args} --ckpt-dir {prev_ckpt}"
+    )
+    return cmd
+
+
 @pytest.mark.timeout(256)  # Optional: fail if the test takes too long.
 @pytest.mark.slow
-def test_train_evo2_runs(tmp_path):
+def test_train_evo2_finetune_runs(tmp_path):
     """
     This test runs the `train_evo2` command with mock data in a temporary directory.
     It uses the temporary directory provided by pytest as the working directory.
@@ -83,7 +107,8 @@ def test_train_evo2_runs(tmp_path):
     num_steps = 2
     # Note: The command assumes that `train_evo2` is in your PATH.
     command = small_training_cmd(tmp_path, max_steps=num_steps, val_check=num_steps)
-    run_command_in_subprocess(command=command, path=str(tmp_path))
+    stdout_pretrain: str = run_command_in_subprocess(command=command, path=str(tmp_path))
+    assert "Restoring model weights from RestoreConfig(path='" not in stdout_pretrain
 
     log_dir = tmp_path / "evo2"
     checkpoints_dir = log_dir / "checkpoints"
@@ -109,22 +134,56 @@ def test_train_evo2_runs(tmp_path):
     # Recursively search for files with tensorboard logger
     event_files = list(tensorboard_dir.rglob("events.out.tfevents*"))
     assert event_files, f"No TensorBoard event files found under {tensorboard_dir}"
+    assert len(matching_subfolders) == 1, "Only one checkpoint subfolder should be found."
+    command_finetune = small_training_finetune_cmd(
+        tmp_path / "finetune", max_steps=num_steps, val_check=num_steps, prev_ckpt=matching_subfolders[0]
+    )
+    stdout_finetune: str = run_command_in_subprocess(command=command_finetune, path=str(tmp_path))
+    assert "Restoring model weights from RestoreConfig(path='" in stdout_finetune
+
+    log_dir_ft = tmp_path / "finetune" / "evo2"
+    checkpoints_dir_ft = log_dir_ft / "checkpoints"
+    tensorboard_dir_ft = log_dir_ft / "dev"
+
+    # Check if logs dir exists
+    assert log_dir_ft.exists(), "Logs folder should exist."
+    # Check if checkpoints dir exists
+    assert checkpoints_dir_ft.exists(), "Checkpoints folder does not exist."
+
+    expected_checkpoint_suffix = f"{num_steps}.0-last"
+    # Check if any subfolder ends with the expected suffix
+    matching_subfolders_ft = [
+        p for p in checkpoints_dir_ft.iterdir() if p.is_dir() and (expected_checkpoint_suffix in p.name)
+    ]
+
+    assert matching_subfolders_ft, (
+        f"No checkpoint subfolder ending with '{expected_checkpoint_suffix}' found in {checkpoints_dir_ft}."
+    )
+
+    # Check if directory with tensorboard logs exists
+    assert tensorboard_dir_ft.exists(), "TensorBoard logs folder does not exist."
+    # Recursively search for files with tensorboard logger
+    event_files = list(tensorboard_dir_ft.rglob("events.out.tfevents*"))
+    assert event_files, f"No TensorBoard event files found under {tensorboard_dir_ft}"
+
+    assert len(matching_subfolders_ft) == 1, "Only one checkpoint subfolder should be found."
 
 
 @pytest.mark.timeout(256)  # Optional: fail if the test takes too long.
 @pytest.mark.slow
-def test_train_evo2_mamba_runs(tmp_path):
+def test_train_evo2_mamba_finetune_runs(tmp_path):
     """
     This test runs the `train_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.
     """
     num_steps = 2
     # Note: The command assumes that `train_evo2` is in your PATH.
-    command = small_training_mamba_cmd(tmp_path, max_steps=num_steps, val_check=num_steps)
-    run_command_in_subprocess(command=command, path=str(tmp_path))
+    command = small_training_mamba_cmd(tmp_path / "pretrain", max_steps=num_steps, val_check=num_steps)
+    stdout_pretrain: str = run_command_in_subprocess(command=command, path=str(tmp_path))
+    assert "Restoring model weights from RestoreConfig(path='" not in stdout_pretrain
 
-    log_dir = tmp_path / "evo2"
+    log_dir = tmp_path / "pretrain" / "evo2"
     checkpoints_dir = log_dir / "checkpoints"
     tensorboard_dir = log_dir / "dev"
 
@@ -149,6 +208,40 @@ def test_train_evo2_mamba_runs(tmp_path):
     event_files = list(tensorboard_dir.rglob("events.out.tfevents*"))
     assert event_files, f"No TensorBoard event files found under {tensorboard_dir}"
 
+    assert len(matching_subfolders) == 1, "Only one checkpoint subfolder should be found."
+    command_finetune = small_training_mamba_finetune_cmd(
+        tmp_path / "finetune", max_steps=num_steps, val_check=num_steps, prev_ckpt=matching_subfolders[0]
+    )
+    stdout_finetune: str = run_command_in_subprocess(command=command_finetune, path=str(tmp_path))
+    assert "Restoring model weights from RestoreConfig(path='" in stdout_finetune
+
+    log_dir_ft = tmp_path / "finetune" / "evo2"
+    checkpoints_dir_ft = log_dir_ft / "checkpoints"
+    tensorboard_dir_ft = log_dir_ft / "dev"
+
+    # Check if logs dir exists
+    assert log_dir_ft.exists(), "Logs folder should exist."
+    # Check if checkpoints dir exists
+    assert checkpoints_dir_ft.exists(), "Checkpoints folder does not exist."
+
+    expected_checkpoint_suffix = f"{num_steps}.0-last"
+    # Check if any subfolder ends with the expected suffix
+    matching_subfolders_ft = [
+        p for p in checkpoints_dir_ft.iterdir() if p.is_dir() and (expected_checkpoint_suffix in p.name)
+    ]
+
+    assert matching_subfolders_ft, (
+        f"No checkpoint subfolder ending with '{expected_checkpoint_suffix}' found in {checkpoints_dir_ft}."
+    )
+
+    # Check if directory with tensorboard logs exists
+    assert tensorboard_dir_ft.exists(), "TensorBoard logs folder does not exist."
+    # Recursively search for files with tensorboard logger
+    event_files = list(tensorboard_dir_ft.rglob("events.out.tfevents*"))
+    assert event_files, f"No TensorBoard event files found under {tensorboard_dir_ft}"
+
+    assert len(matching_subfolders_ft) == 1, "Only one checkpoint subfolder should be found."
+
 
 @pytest.mark.timeout(256)  # Optional: fail if the test takes too long.
 @pytest.mark.slow