Update

[ghstack-poisoned]

Author: fegin

tests/unit_tests/test_checkpoint.py

@@ -111,13 +111,16 @@ def tearDown(self):
     def test_save(self, *_):
         """Test that calling save() writes a checkpoint file to disk."""
         job_config = DummyJobConfig(job=self.dummy_job)
+        ft_manager = mock.Mock()
+        ft_manager.enabled = False
         manager = CheckpointManager(
             dummy_dataloader,
             dummy_model_parts,
             dummy_optimizers,
             dummy_lr_schedulers,
             {"trainer": self.trainer_state},
             job_config,
+            ft_manager,
         )
         step = 20
         manager.save(curr_step=step, force=True)
@@ -141,13 +144,16 @@ def test_save(self, *_):
     def test_load(self, *_):
         """Test that load() properly reads the checkpoint file from disk and restores state."""
         job_config = DummyJobConfig(job=self.dummy_job)
+        ft_manager = mock.Mock()
+        ft_manager.enabled = False
         manager = CheckpointManager(
             dummy_dataloader,
             dummy_model_parts,
             dummy_optimizers,
             dummy_lr_schedulers,
             {"trainer": self.trainer_state},
             job_config,
+            ft_manager,
         )
         step = 30
         manager.save(curr_step=step, force=True)
@@ -179,13 +185,16 @@ def test_purge_stale_checkpoints_rank_zero(self, *_):
         """
         job_config = DummyJobConfig(job=self.dummy_job)
         job_config.checkpoint.keep_latest_k = 3
+        ft_manager = mock.Mock()
+        ft_manager.enabled = False
         manager = CheckpointManager(
             dummy_dataloader,
             dummy_model_parts,
             dummy_optimizers,
             dummy_lr_schedulers,
             {"trainer": self.trainer_state},
             job_config,
+            ft_manager,
         )
         steps = [10, 20, 30, 40, 50]
         for s in steps:
@@ -223,13 +232,16 @@ def test_purge_stale_checkpoints_rank_nonzero(self, *_):
         """
         job_config = DummyJobConfig(job=self.dummy_job)
         job_config.checkpoint.keep_latest_k = 3
+        ft_manager = mock.Mock()
+        ft_manager.enabled = False
         manager = CheckpointManager(
             dummy_dataloader,
             dummy_model_parts,
             dummy_optimizers,
             dummy_lr_schedulers,
             {"trainer": self.trainer_state},
             job_config,
+            ft_manager,
         )
         steps = [10, 20, 30, 40, 50]
         for s in steps:
@@ -260,13 +272,16 @@ def test_async_save_calls_async_wait(self, *_):
         # Set async_mode to "async" in the job configuration.
         job_config = DummyJobConfig(job=self.dummy_job)
         job_config.checkpoint.async_mode = "async"
+        ft_manager = mock.Mock()
+        ft_manager.enabled = False
         manager = CheckpointManager(
             dummy_dataloader,
             dummy_model_parts,
             dummy_optimizers,
             dummy_lr_schedulers,
             {"trainer": self.trainer_state},
             job_config,
+            ft_manager,
         )
         # First save: should schedule an async save.
         manager.save(curr_step=10, force=False)