test_dist_checkpoint_utils.py

import unittest
from unittest.mock import patch, MagicMock, mock_open, call
import os
import shutil
import json
import torch # Keep torch for torch.device and potentially other utilities
import sys

# Import the functions to be tested
from training.utils.dist_checkpoint_utils import (
    load_checkpoint,
    save_checkpoint,
    save_stream_dataloader_state_dict,
    load_stream_dataloader_state_dict
)

# Mock get_pipeline_parallel_rank as it's used to construct file names
@patch('training.utils.dist_checkpoint_utils.get_pipeline_parallel_rank', MagicMock(return_value=0))
class TestDistCheckpointUtils(unittest.TestCase):

    def setUp(self):
        self.test_checkpoint_dir = "test_checkpoints"
        # Ensure a clean state for each test
        if os.path.exists(self.test_checkpoint_dir):
            shutil.rmtree(self.test_checkpoint_dir)
        os.makedirs(self.test_checkpoint_dir, exist_ok=True)

        self.args = MagicMock()
        self.args.checkpoint_path = self.test_checkpoint_dir

        self.pipe = MagicMock()
        self.pipe.global_step = 10
        self.pipe.model = MagicMock()
        self.pipe.model.model = MagicMock() # Mocking model.model attribute for state_dict
        self.pipe.optimizer = MagicMock()
        self.pipe.optimizer.state_dict = MagicMock(return_value={"opt_param": 1})
        self.pipe.scheduler = MagicMock()
        self.pipe.scheduler.state_dict = MagicMock(return_value={"sched_param": 1})
        self.pipe.model.model.state_dict = MagicMock(return_value={"model_param": 1})
        
        # Suppress print statements from the module
        self.original_stdout = sys.stdout
        sys.stdout = MagicMock()

    def tearDown(self):
        if os.path.exists(self.test_checkpoint_dir):
            shutil.rmtree(self.test_checkpoint_dir)
        sys.stdout = self.original_stdout # Restore stdout

    def _create_dummy_checkpoint_files(self, step, create_meta=True, create_model=True, create_optimizer=True, create_scheduler=True):
        checkpoint_step_path = os.path.join(self.test_checkpoint_dir, f"checkpoint_{step}")
        os.makedirs(checkpoint_step_path, exist_ok=True)

        if create_meta:
            with open(os.path.join(checkpoint_step_path, 'meta.json'), 'w') as f:
                json.dump({'step': step}, f)
        
        if create_model:
            torch.save({}, os.path.join(checkpoint_step_path, 'prank_0_checkpoint.pt'))
        if create_optimizer:
            torch.save({}, os.path.join(checkpoint_step_path, 'prank_0_optimizer.pt'))
        if create_scheduler:
            torch.save({}, os.path.join(checkpoint_step_path, 'prank_0_scheduler.pt'))
        
        with open(os.path.join(self.test_checkpoint_dir, 'latest'), 'w') as f:
            f.write(str(step))
        return checkpoint_step_path

    @patch('torch.load')
    def test_load_checkpoint_success(self, mock_torch_load):
        step = 10
        self._create_dummy_checkpoint_files(step)
        mock_torch_load.return_value = {"dummy_state": "value"}

        load_checkpoint(self.pipe, self.args)

        self.assertEqual(self.pipe.global_step, step)
        self.pipe.model.model.load_state_dict.assert_called_once_with({"dummy_state": "value"})
        self.pipe.optimizer.load_state_dict.assert_called_once_with({"dummy_state": "value"})
        self.pipe.scheduler.load_state_dict.assert_called_once_with({"dummy_state": "value"})

    def test_load_checkpoint_no_latest_file(self):
        # No 'latest' file
        load_checkpoint(self.pipe, self.args)
        self.assertTrue(any("no checkpoint available, skipping" in call_args[0][0] for call_args in sys.stdout.write.call_args_list))

    def test_load_checkpoint_meta_file_not_found(self):
        step = 5
        self._create_dummy_checkpoint_files(step, create_meta=False)
        
        load_checkpoint(self.pipe, self.args)
        expected_msg = f"Checkpoint metadata file not found at {os.path.join(self.test_checkpoint_dir, f'checkpoint_{step}', 'meta.json')}"
        self.assertTrue(any(expected_msg in call_args[0][0] for call_args in sys.stdout.write.call_args_list))

    @patch('torch.load', side_effect=FileNotFoundError("File not found"))
    def test_load_checkpoint_model_file_not_found(self, mock_torch_load):
        step = 15
        self._create_dummy_checkpoint_files(step, create_model=False) # Model file won't be there but mock matters more
        
        # We need meta.json to proceed to loading attempts
        checkpoint_step_path = os.path.join(self.test_checkpoint_dir, f"checkpoint_{step}")
        with open(os.path.join(checkpoint_step_path, 'meta.json'), 'w') as f:
            json.dump({'step': step}, f)
        with open(os.path.join(self.test_checkpoint_dir, 'latest'), 'w') as f:
            f.write(str(step))

        load_checkpoint(self.pipe, self.args)
        model_path = os.path.join(checkpoint_step_path, 'prank_0_checkpoint.pt')
        self.assertTrue(any(f"Model checkpoint file not found: {model_path}" in call_args[0][0] for call_args in sys.stdout.write.call_args_list))

    @patch('torch.load', side_effect=RuntimeError("Torch load error"))
    def test_load_checkpoint_torch_load_generic_error(self, mock_torch_load):
        step = 20
        self._create_dummy_checkpoint_files(step) # All files are there

        load_checkpoint(self.pipe, self.args)
        checkpoint_step_path = os.path.join(self.test_checkpoint_dir, f"checkpoint_{step}")
        model_path = os.path.join(checkpoint_step_path, 'prank_0_checkpoint.pt')
        self.assertTrue(any(f"Failed to load model params from {model_path}: Torch load error" in call_args[0][0] for call_args in sys.stdout.write.call_args_list))


    @patch('torch.save')
    @patch('os.makedirs')
    def test_save_checkpoint_directory_creation_and_save(self, mock_os_makedirs, mock_torch_save):
        self.pipe.global_step = 25
        
        save_checkpoint(self.pipe, self.args)

        checkpoint_step_path = os.path.join(self.test_checkpoint_dir, "checkpoint_25")
        mock_os_makedirs.assert_called_once_with(checkpoint_step_path, exist_ok=True)
        
        self.assertEqual(mock_torch_save.call_count, 3) # model, optimizer, scheduler
        
        # Check meta.json content
        meta_path = os.path.join(checkpoint_step_path, 'meta.json')
        self.assertTrue(os.path.exists(meta_path))
        with open(meta_path, 'r') as f:
            meta = json.load(f)
        self.assertEqual(meta['step'], 25)

        # Check latest file content
        latest_path = os.path.join(self.test_checkpoint_dir, 'latest')
        self.assertTrue(os.path.exists(latest_path))
        with open(latest_path, 'r') as f:
            latest_step_str = f.read()
        self.assertEqual(latest_step_str, "25")


    @patch('torch.save')
    @patch('os.makedirs')
    def test_save_stream_dataloader_state_dict_creation_and_save(self, mock_os_makedirs, mock_torch_save):
        self.pipe.global_step = 30
        mock_dataloader = MagicMock()
        mock_dataloader.dataset.state_dict.return_value = {"dataset_state": "some_state"}

        save_stream_dataloader_state_dict(mock_dataloader, self.pipe, self.args)

        checkpoint_step_path = os.path.join(self.test_checkpoint_dir, "checkpoint_30")
        mock_os_makedirs.assert_called_once_with(checkpoint_step_path, exist_ok=True)
        
        dataset_state_path = os.path.join(checkpoint_step_path, 'dataset_state_dict.pt')
        mock_torch_save.assert_called_once_with({"dataset_state": "some_state"}, dataset_state_path)

    @patch('torch.save', side_effect=Exception("Failed to save dataset"))
    @patch('os.makedirs')
    def test_save_stream_dataloader_state_dict_save_error(self, mock_os_makedirs, mock_torch_save):
        self.pipe.global_step = 35
        mock_dataloader = MagicMock()
        mock_dataloader.dataset.state_dict.return_value = {"dataset_state": "some_state"}

        save_stream_dataloader_state_dict(mock_dataloader, self.pipe, self.args)
        
        checkpoint_step_path = os.path.join(self.test_checkpoint_dir, "checkpoint_35")
        dataset_state_path = os.path.join(checkpoint_step_path, 'dataset_state_dict.pt')
        self.assertTrue(any(f"Failed to save dataset state_dict to {dataset_state_path}: Failed to save dataset" in call_args[0][0] for call_args in sys.stdout.write.call_args_list))


    @patch('torch.load')
    def test_load_stream_dataloader_state_dict_success(self, mock_torch_load):
        self.pipe.global_step = 40
        # We need to ensure the checkpoint directory for this step exists for path construction
        checkpoint_step_path = os.path.join(self.test_checkpoint_dir, f"checkpoint_{self.pipe.global_step}")
        os.makedirs(checkpoint_step_path, exist_ok=True)
        # Dummy file for torch.load to be "successful"
        dataset_state_dict_path = os.path.join(checkpoint_step_path, 'dataset_state_dict.pt')
        with open(dataset_state_dict_path, 'w') as f: f.write("dummy data")


        mock_dataloader = MagicMock()
        mock_dataloader.dataset = MagicMock() # Ensure dataset attribute exists
        
        mock_torch_load.return_value = {"loaded_state": "value"}

        load_stream_dataloader_state_dict(mock_dataloader, self.pipe, self.args)

        mock_torch_load.assert_called_once_with(dataset_state_dict_path)
        mock_dataloader.dataset.load_state_dict.assert_called_once_with({"loaded_state": "value"})

    @patch('torch.load', side_effect=FileNotFoundError("Dataset state file not found"))
    def test_load_stream_dataloader_state_dict_file_not_found(self, mock_torch_load):
        self.pipe.global_step = 45
        checkpoint_step_path = os.path.join(self.test_checkpoint_dir, f"checkpoint_{self.pipe.global_step}")
        # No need to create the dummy file, as we are testing FileNotFoundError

        mock_dataloader = MagicMock()
        mock_dataloader.dataset = MagicMock()

        load_stream_dataloader_state_dict(mock_dataloader, self.pipe, self.args)
        
        dataset_state_dict_path = os.path.join(checkpoint_step_path, 'dataset_state_dict.pt')
        self.assertTrue(any(f"Dataset state_dict file not found: {dataset_state_dict_path}" in call_args[0][0] for call_args in sys.stdout.write.call_args_list))

    @patch('torch.load', side_effect=RuntimeError("Torch load error for dataset"))
    def test_load_stream_dataloader_state_dict_generic_error(self, mock_torch_load):
        self.pipe.global_step = 50
        checkpoint_step_path = os.path.join(self.test_checkpoint_dir, f"checkpoint_{self.pipe.global_step}")
        os.makedirs(checkpoint_step_path, exist_ok=True)
        dataset_state_dict_path = os.path.join(checkpoint_step_path, 'dataset_state_dict.pt')
        with open(dataset_state_dict_path, 'w') as f: f.write("dummy data") # File exists

        mock_dataloader = MagicMock()
        mock_dataloader.dataset = MagicMock()

        load_stream_dataloader_state_dict(mock_dataloader, self.pipe, self.args)
        
        self.assertTrue(any(f"Failed to load dataset state_dict from {dataset_state_dict_path}: Torch load error for dataset" in call_args[0][0] for call_args in sys.stdout.write.call_args_list))


if __name__ == '__main__':
    unittest.main()