text: # Copyright (c) Microsoft Corporation. All rights reserved.
  # Licensed under the MIT License.
  # orttraining_test_checkpoint_storage.py
  
  import pytest
  import torch
  import numpy as np
  import os
  import shutil
  import pickle
  
  from onnxruntime.training import _checkpoint_storage
  
  # Helper functions
  
  
  def _equals(a, b):
      """Checks recursively if two dictionaries are equal"""
      if isinstance(a, dict):
          for key in a:
              if key not in b or not _equals(a[key], b[key]):
                  return False
          return True
      else:
          if isinstance(a, bytes):
              a = a.decode()
          if isinstance(b, bytes):
              b = b.decode()
          are_equal = a == b
          return are_equal if isinstance(are_equal, bool) else are_equal.all()
  
      return False
  
  
  def _numpy_types(obj_value):
      """Return a bool indicating whether or not the input obj_value is a numpy type object
  
      Recursively checks if the obj_value (could be a dictionary) is a numpy type object.
      Exceptions are str and bytes.
  
      Returns true if object is numpy type, str, or bytes
      False if any other type
      """
      if not isinstance(obj_value, dict):
          return type(obj_value).__module__ == np.__name__ or isinstance(obj_value, str) or isinstance(obj_value, bytes)
  
      for _, value in obj_value.items():
          if not _numpy_types(value):
              return False
      return True
  
  
  def _get_dict(separated_key):
      """Create dummy dictionary with different datatypes
  
      Returns the tuple of the entire dummy dictionary created, key argument as a dictionary for _checkpoint_storage.load
      function and the value for that key in the original dictionary
  
      For example the complete dictionary is represented by:
      {
          'int1':1,
          'int2': 2,
          'int_list': [1,2,3,5,6],
          'dict1': {
              'np_array': np.arange(100),
              'dict2': {'int3': 3, 'int4': 4},
              'str1': "onnxruntime"
          },
          'bool1': bool(True),
          'int5': 5,
          'float1': 2.345,
          'np_array_float': np.array([1.234, 2.345, 3.456]),
          'np_array_float_3_dim': np.array([[[1,2],[3,4]], [[5,6],[7,8]]])
      }
  
      if the input key is ['dict1', 'str1'], then the key argument returned is 'dict1/str1'
      and the value corresponding to that is "onnxruntime"
  
      so, for the above example, the returned tuple is:
      (original_dict, {'key': 'dict1/str1', "onnxruntime")
      """
      test_dict = {
          "int1": 1,
          "int2": 2,
          "int_list": [1, 2, 3, 5, 6],
          "dict1": {"np_array": np.arange(100), "dict2": {"int3": 3, "int4": 4}, "str1": "onnxruntime"},
          "bool1": bool(True),
          "int5": 5,
          "float1": 2.345,
          "np_array_float": np.array([1.234, 2.345, 3.456]),
          "np_array_float_3_dim": np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
      }
      key = ""
      expected_val = test_dict
      for single_key in separated_key:
          key += single_key + "/"
          expected_val = expected_val[single_key]
      return test_dict, {"key": key} if len(separated_key) > 0 else dict(), expected_val
  
  
  class _CustomClass(object):
      """Custom object that encpsulates dummy values for loss, epoch and train_step"""
  
      def __init__(self):
          self._loss = 1.23
          self._epoch = 12000
          self._train_step = 25
  
      def __eq__(self, other):
          if isinstance(other, _CustomClass):
              return self._loss == other._loss and self._epoch == other._epoch and self._train_step == other._train_step
  
  
  # Test fixtures
  
  
  @pytest.yield_fixture(scope="function")
  def checkpoint_storage_test_setup():
      checkpoint_dir = os.path.abspath("checkpoint_dir/")
      if not os.path.exists(checkpoint_dir):
          os.makedirs(checkpoint_dir, exist_ok=True)
      pytest.checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.ortcp")
      yield "checkpoint_storage_test_setup"
      shutil.rmtree(checkpoint_dir)
  
  
  @pytest.yield_fixture(scope="function")
  def checkpoint_storage_test_parameterized_setup(request, checkpoint_storage_test_setup):
      yield request.param
  
  
  # Tests
  
  
  @pytest.mark.parametrize(
      "checkpoint_storage_test_parameterized_setup",
      [
          _get_dict([]),
          _get_dict(["int1"]),
          _get_dict(["dict1"]),
          _get_dict(["dict1", "dict2"]),
          _get_dict(["dict1", "dict2", "int4"]),
          _get_dict(["dict1", "str1"]),
          _get_dict(["bool1"]),
          _get_dict(["float1"]),
          _get_dict(["np_array_float"]),
      ],
      indirect=True,
  )
  def test_checkpoint_storage_saved_dict_matches_loaded(checkpoint_storage_test_parameterized_setup):
      to_save = checkpoint_storage_test_parameterized_setup[0]
      key_arg = checkpoint_storage_test_parameterized_setup[1]
      expected = checkpoint_storage_test_parameterized_setup[2]
      _checkpoint_storage.save(to_save, pytest.checkpoint_path)
      loaded = _checkpoint_storage.load(pytest.checkpoint_path, **key_arg)
      assert _equals(loaded, expected)
      assert _numpy_types(loaded)
  
  
  @pytest.mark.parametrize(
      "checkpoint_storage_test_parameterized_setup",
      [{"int_set": {1, 2, 3, 4, 5}}, {"str_set": {"one", "two"}}, [1, 2, 3], 2.352],
      indirect=True,
  )
  def test_checkpoint_storage_saving_non_supported_types_fails(checkpoint_storage_test_parameterized_setup):
      to_save = checkpoint_storage_test_parameterized_setup
      with pytest.raises(Exception):
          _checkpoint_storage.save(to_save, pytest.checkpoint_path)
  
  
  @pytest.mark.parametrize(
      "checkpoint_storage_test_parameterized_setup",
      [
          ({"int64_tensor": torch.tensor(np.arange(100))}, "int64_tensor", torch.int64, np.int64),
          ({"int32_tensor": torch.tensor(np.arange(100), dtype=torch.int32)}, "int32_tensor", torch.int32, np.int32),
          ({"int16_tensor": torch.tensor(np.arange(100), dtype=torch.int16)}, "int16_tensor", torch.int16, np.int16),
          ({"int8_tensor": torch.tensor(np.arange(100), dtype=torch.int8)}, "int8_tensor", torch.int8, np.int8),
          ({"float64_tensor": torch.tensor(np.array([1.0, 2.0]))}, "float64_tensor", torch.float64, np.float64),
          (
              {"float32_tensor": torch.tensor(np.array([1.0, 2.0]), dtype=torch.float32)},
              "float32_tensor",
              torch.float32,
              np.float32,
          ),
          (
              {"float16_tensor": torch.tensor(np.array([1.0, 2.0]), dtype=torch.float16)},
              "float16_tensor",
              torch.float16,
              np.float16,
          ),
      ],
      indirect=True,
  )
  def test_checkpoint_storage_saving_tensor_datatype(checkpoint_storage_test_parameterized_setup):
      tensor_dict = checkpoint_storage_test_parameterized_setup[0]
      tensor_name = checkpoint_storage_test_parameterized_setup[1]
      tensor_dtype = checkpoint_storage_test_parameterized_setup[2]
      np_dtype = checkpoint_storage_test_parameterized_setup[3]
  
      _checkpoint_storage.save(tensor_dict, pytest.checkpoint_path)
  
      loaded = _checkpoint_storage.load(pytest.checkpoint_path)
      assert isinstance(loaded[tensor_name], np.ndarray)
      assert tensor_dict[tensor_name].dtype == tensor_dtype
      assert loaded[tensor_name].dtype == np_dtype
      assert (tensor_dict[tensor_name].numpy() == loaded[tensor_name]).all()
  
  
  @pytest.mark.parametrize(
      "checkpoint_storage_test_parameterized_setup",
      [
          ({"two_dim": torch.ones([2, 4], dtype=torch.float64)}, "two_dim"),
          ({"three_dim": torch.ones([2, 4, 6], dtype=torch.float64)}, "three_dim"),
          ({"four_dim": torch.ones([2, 4, 6, 8], dtype=torch.float64)}, "four_dim"),
      ],
      indirect=True,
  )
  def test_checkpoint_storage_saving_multiple_dimension_tensors(checkpoint_storage_test_parameterized_setup):
      tensor_dict = checkpoint_storage_test_parameterized_setup[0]
      tensor_name = checkpoint_storage_test_parameterized_setup[1]
  
      _checkpoint_storage.save(tensor_dict, pytest.checkpoint_path)
  
      loaded = _checkpoint_storage.load(pytest.checkpoint_path)
      assert isinstance(loaded[tensor_name], np.ndarray)
      assert (tensor_dict[tensor_name].numpy() == loaded[tensor_name]).all()
  
  
  @pytest.mark.parametrize(
      "checkpoint_storage_test_parameterized_setup", [{}, {"a": {}}, {"a": {"b": {}}}], indirect=True
  )
  def test_checkpoint_storage_saving_and_loading_empty_dictionaries_succeeds(checkpoint_storage_test_parameterized_setup):
      saved = checkpoint_storage_test_parameterized_setup
      _checkpoint_storage.save(saved, pytest.checkpoint_path)
  
      loaded = _checkpoint_storage.load(pytest.checkpoint_path)
      assert _equals(saved, loaded)
  
  
  def test_checkpoint_storage_load_file_that_does_not_exist_fails(checkpoint_storage_test_setup):
      with pytest.raises(Exception):
          _checkpoint_storage.load(pytest.checkpoint_path)
  
  
  def test_checkpoint_storage_for_custom_user_dict_succeeds(checkpoint_storage_test_setup):
      custom_class = _CustomClass()
      user_dict = {"tensor1": torch.tensor(np.arange(100), dtype=torch.float32), "custom_class": custom_class}
  
      pickled_bytes = pickle.dumps(user_dict).hex()
      to_save = {"a": torch.tensor(np.array([1.0, 2.0]), dtype=torch.float32), "user_dict": pickled_bytes}
      _checkpoint_storage.save(to_save, pytest.checkpoint_path)
  
      loaded_dict = _checkpoint_storage.load(pytest.checkpoint_path)
      assert (loaded_dict["a"] == to_save["a"].numpy()).all()
      try:
          loaded_dict["user_dict"] = loaded_dict["user_dict"].decode()
      except AttributeError:
          pass
      loaded_obj = pickle.loads(bytes.fromhex(loaded_dict["user_dict"]))
  
      assert torch.all(loaded_obj["tensor1"].eq(user_dict["tensor1"]))
      assert loaded_obj["custom_class"] == custom_class
  