# Path: vmoe/checkpoints/periodic_actions.py
# class PeriodicSaveCheckpoint(periodic_actions.PeriodicCallback):
#   def __init__(
#       self,
#       *,
#       prefix: str,
#       state_axis_resources: PyTree,
#       mesh: Optional[Mesh] = None,
#       num_shards: int = 0,
#       num_threads: Optional[int] = None,
#       wait_seconds: Optional[int] = None,
#       every_steps: Optional[int] = None,
#       every_secs: Optional[float] = None,
#       on_steps: Optional[Iterable[int]] = None,
#       keep_last: Optional[int] = None,
#       keep_steps_multiple_of: Optional[int] = None,
#       execute_async: bool = True,
#       report_progress: Optional[periodic_actions.ReportProgress] = None,
#       report_progress_name: str = 'ckpt'):
#   def __del__(self):
#   def _makedirs(cls, workdir: str):
#   def _remove_old_checkpoints(cls, prefix: str, keep_last: int,
#                               keep_multiple: int, thread_pool: ThreadPool):
#     def _parse_step_from_filepath(filepath):
#     def _find_step_numbers(filepaths):
#     def _remove():
#       def match_remove_fn(filepath):
#   def _block_async_result(self, wait_seconds: Optional[int]):
#   def _make_callback_fn(self, prefix, state_axis_resources, mesh, num_shards,
#                         wait_seconds, keep_last, keep_multiple, execute_async,
#                         thread_pool, report_progress, report_progress_name):
#     def callback_fn(step: int, t: float, state: PyTree):


import os
import random
import time
from unittest import mock
from absl.testing import absltest
from vmoe.checkpoints import periodic_actions
    # When calling save_checkpoint, we do nothing but we'll wait a few seconds.
    def _save_checkpoint_side_effect(*args, thread_pool, **kwargs):
      del args
      del kwargs
      wait = lambda: time.sleep(random.randint(1, 3))
      return thread_pool.apply_async(wait)
    mock_save_checkpoint.side_effect = _save_checkpoint_side_effect
    # Run a few steps, calling saver on each step.
    prefix = os.path.join(self.create_tempdir().full_path, 'ckpt')
    saver = periodic_actions.PeriodicSaveCheckpoint(
        prefix=prefix,
        state_axis_resources={},
        every_steps=4)
    for step in range(1, 10):
      saver(step=step, state={})
    # Check that the saver was called twice, on steps 4 and 8.
    call_args_list = mock_save_checkpoint.call_args_list
    self.assertLen(call_args_list, 2)
    self.assertEqual(call_args_list[0],
                     mock.call(prefix=prefix + '_4', num_shards=0, mesh=None,
                               makedirs=False, overwrite=True, tree={},
                               axis_resources={}, thread_pool=mock.ANY))
    self.assertEqual(call_args_list[1],
                     mock.call(prefix=prefix + '_8', num_shards=0, mesh=None,
                               makedirs=False, overwrite=True, tree={},
                               axis_resources={}, thread_pool=mock.ANY))
    saver.__del__()

  def test_report_progress(self):
    mock_report_progress = mock.MagicMock(

# Path: vmoe/checkpoints/periodic_actions.py
# class PeriodicSaveCheckpoint(periodic_actions.PeriodicCallback):
#   def __init__(
#       self,
#       *,
#       prefix: str,
#       state_axis_resources: PyTree,
#       mesh: Optional[Mesh] = None,
#       num_shards: int = 0,
#       num_threads: Optional[int] = None,
#       wait_seconds: Optional[int] = None,
#       every_steps: Optional[int] = None,
#       every_secs: Optional[float] = None,
#       on_steps: Optional[Iterable[int]] = None,
#       keep_last: Optional[int] = None,
#       keep_steps_multiple_of: Optional[int] = None,
#       execute_async: bool = True,
#       report_progress: Optional[periodic_actions.ReportProgress] = None,
#       report_progress_name: str = 'ckpt'):
#   def __del__(self):
#   def _makedirs(cls, workdir: str):
#   def _remove_old_checkpoints(cls, prefix: str, keep_last: int,
#                               keep_multiple: int, thread_pool: ThreadPool):
#     def _parse_step_from_filepath(filepath):
#     def _find_step_numbers(filepaths):
#     def _remove():
#       def match_remove_fn(filepath):
#   def _block_async_result(self, wait_seconds: Optional[int]):
#   def _make_callback_fn(self, prefix, state_axis_resources, mesh, num_shards,
#                         wait_seconds, keep_last, keep_multiple, execute_async,
#                         thread_pool, report_progress, report_progress_name):
#     def callback_fn(step: int, t: float, state: PyTree):


# Path: vmoe/checkpoints/periodic_actions_test.py
import os
import random
import time
from unittest import mock
from absl.testing import absltest
from vmoe.checkpoints import periodic_actions

    # When calling save_checkpoint, we do nothing but we'll wait a few seconds.
    def _save_checkpoint_side_effect(*args, thread_pool, **kwargs):
      del args
      del kwargs
      wait = lambda: time.sleep(random.randint(1, 3))
      return thread_pool.apply_async(wait)
    mock_save_checkpoint.side_effect = _save_checkpoint_side_effect
    # Run a few steps, calling saver on each step.
    prefix = os.path.join(self.create_tempdir().full_path, 'ckpt')
    saver = periodic_actions.PeriodicSaveCheckpoint(
        prefix=prefix,
        state_axis_resources={},
        every_steps=4)
    for step in range(1, 10):
      saver(step=step, state={})
    # Check that the saver was called twice, on steps 4 and 8.
    call_args_list = mock_save_checkpoint.call_args_list
    self.assertLen(call_args_list, 2)
    self.assertEqual(call_args_list[0],
                     mock.call(prefix=prefix + '_4', num_shards=0, mesh=None,
                               makedirs=False, overwrite=True, tree={},
                               axis_resources={}, thread_pool=mock.ANY))
    self.assertEqual(call_args_list[1],
                     mock.call(prefix=prefix + '_8', num_shards=0, mesh=None,
                               makedirs=False, overwrite=True, tree={},
                               axis_resources={}, thread_pool=mock.ANY))
    saver.__del__()

  def test_report_progress(self):
    mock_report_progress = mock.MagicMock(

        periodic_actions.periodic_actions.ReportProgress)