[misc]: Use runner module in all tests and add unit test for runner dynamic instantiation

Author: techverve

fastvideo/tests/training/Vanilla/mfu_calculation.py

@@ -17,7 +17,7 @@
 
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
-from fastvideo.training.wan_training_pipeline import WanTrainingPipeline
+from fastvideo.training.runner import main
 
 MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_PATH = "data/crush-smol_processed_t2v/training_dataset/worker_1/worker_0/"
@@ -123,10 +123,11 @@ def run_worker():
         "1"
     ])
     # Call the main training function
-    pipeline = WanTrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
+    args.pipeline_class = "WanTrainingPipeline"
+    args.pipeline_module = "fastvideo.training.wan_training_pipeline"
+    
+    # Call the main training function
+    main(args)
 
 
 def test_distributed_training():

fastvideo/tests/training/Vanilla/test_training_loss.py

@@ -14,7 +14,7 @@
 
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
-from fastvideo.training.wan_training_pipeline import WanTrainingPipeline
+from fastvideo.training.runner import main
 
 wandb_name = "test_training_loss"
 a40_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/a40_reference_wandb_summary.json"
@@ -50,10 +50,11 @@ def run_worker():
         "--not_apply_cfg_solver", "--dit_precision", "fp32", "--max_grad_norm", "1.0"
     ])
     # Call the main training function
-    pipeline = WanTrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
+    args.pipeline_class = "WanTrainingPipeline"
+    args.pipeline_module = "fastvideo.training.wan_training_pipeline"
+    
+    # Call the main training function
+    main(args)
 
 
 def test_distributed_training():

fastvideo/tests/training/distill/test_distill_dmd.py

@@ -14,7 +14,7 @@
 
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
-from fastvideo.training.wan_distillation_pipeline import WanDistillationPipeline
+from fastvideo.training.runner import main
 
 wandb_name = "test_distill_dmd"
 
@@ -49,10 +49,11 @@ def run_worker():
         "--enable_gradient_checkpointing_type", "full"
     ])
     # Call the main training function
-    pipeline = WanDistillationPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
+    args.pipeline_class = "WanDistillationPipeline"
+    args.pipeline_module = "fastvideo.training.wan_distillation_pipeline"
+    
+    # Call the main training function
+    main(args)
 
 
 def test_distributed_training():

fastvideo/tests/training/self-forcing/test_self_forcing.py

@@ -11,9 +11,10 @@
 from fastvideo.utils import logger
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-from fastvideo.training.wan_self_forcing_distillation_pipeline import WanSelfForcingDistillationPipeline
+
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
+from fastvideo.training.runner import main
 
 wandb_name = "test_self_forcing_distill"
 
@@ -148,10 +149,11 @@ def run_worker():
     ])
 
     # Call the main training function
-    pipeline = WanSelfForcingDistillationPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Self-forcing distillation training pipeline done")
+    args.pipeline_class = "WanSelfForcingDistillationPipeline"
+    args.pipeline_module = "fastvideo.training.wan_self_forcing_distillation_pipeline"
+    
+    # Call the main training function
+    main(args)
 
 
 def test_distributed_training():

fastvideo/tests/training/test_runner.py

@@ -1,18 +1,26 @@
-import subprocess
-import sys
+from unittest.mock import patch, MagicMock
+from fastvideo.training.runner import main
+import argparse
 
-def test_runner_help():
-    # Run the runner with --help to ensure it parses arguments correctly
-    # and doesn't have any syntax errors or missing imports.
-    result = subprocess.run(
-        [sys.executable, "-m", "fastvideo.training.runner", "--help"],
-        capture_output=True,
-        text=True
-    )
-    
-    # Check if the process completed successfully
-    assert result.returncode == 0, f"Runner failed with exit code {result.returncode}.\nStderr: {result.stderr}"
-    
-    # Check if help output is printed
-    assert "--pipeline_class" in result.stdout
-    assert "--pipeline_module" in result.stdout
+def test_runner_invokes_correct_class():
+    with patch('fastvideo.training.runner.importlib.import_module') as mock_import_module:
+        mock_module = MagicMock()
+        mock_import_module.return_value = mock_module
+        
+        mock_pipeline_class = MagicMock()
+        setattr(mock_module, "WanTrainingPipeline", mock_pipeline_class)
+        
+        mock_pipeline_instance = MagicMock()
+        mock_pipeline_class.from_pretrained.return_value = mock_pipeline_instance
+        
+        args = argparse.Namespace(
+            pipeline_module="fastvideo.training.wan_training_pipeline",
+            pipeline_class="WanTrainingPipeline",
+            pretrained_model_name_or_path="test_model"
+        )
+        
+        main(args)
+        
+        mock_import_module.assert_called_once_with("fastvideo.training.wan_training_pipeline")
+        mock_pipeline_class.from_pretrained.assert_called_once_with("test_model", args=args)
+        mock_pipeline_instance.train.assert_called_once()