diff --git a/tests/slurm_command_gen_strategy/test_nemo_run_slurm_command_gen_strategy.py b/tests/slurm_command_gen_strategy/test_nemo_run_slurm_command_gen_strategy.py index 5bfcd4f6..f98cf3b0 100644 --- a/tests/slurm_command_gen_strategy/test_nemo_run_slurm_command_gen_strategy.py +++ b/tests/slurm_command_gen_strategy/test_nemo_run_slurm_command_gen_strategy.py @@ -59,7 +59,7 @@ def cmd_gen_strategy(self, slurm_system: SlurmSystem) -> NeMoRunSlurmCommandGenS "cmd_args, expected_cmd", [ ( - {"task": "fine_tune", "recipe_name": "llama7_13b"}, + {"docker_image_url": "nvcr.io/nvidia/nemo:24.09", "task": "fine_tune", "recipe_name": "llama7_13b"}, ["nemo", "llm", "fine_tune", "--factory", "llama7_13b", "-y", "trainer.num_nodes=2", "extra_args"], ), ], @@ -83,7 +83,7 @@ def test_generate_test_command( @pytest.mark.parametrize( "cmd_args, expected_exception", [ - ({"recipe_name": "llama7_13b"}, ValueError), + ({"docker_image_url": "nvcr.io/nvidia/nemo:24.09", "recipe_name": "llama7_13b"}, ValueError), ({"task": "fine_tune"}, ValueError), ], )