text: # BSD 3-Clause License
  #
  # Copyright (C) 2021 THL A29 Limited, a Tencent company.  All rights reserved.
  #
  # Redistribution and use in source and binary forms, with or without modification,
  # are permitted provided that the following conditions are met:
  #
  #  * Redistributions of source code must retain the above copyright notice, this
  #    list of conditions and the following disclaimer.
  #
  #  * Redistributions in binary form must reproduce the above copyright notice,
  #    this list of conditions and the following disclaimer in the documentation
  #    and/or other materials provided with the distribution.
  #
  #  * Neither the name of the psutil authors nor the names of its contributors
  #    may be used to endorse or promote products derived from this software without
  #    specific prior written permission.
  #
  # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
  # ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
  # ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  # (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
  # ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  
  import unittest
  from patrickstar.core.comm import CommInfo
  
  import torch
  
  from patrickstar.core import (
      ChunkTensorIndex,
      ChunkType,
      AccessType,
      register_param,
      ParamType,
  )
  
  
  class TestAccess(unittest.TestCase):
      def setUp(self):
          pass
  
      def _check_order(self, chunk_tensor_index, chunk_id):
          start_offset_list = []
          for info in chunk_tensor_index.generate_tensor_info_in_order(chunk_id=chunk_id):
              start_offset = info.start_offset
              if len(start_offset_list) > 0:
                  assert start_offset > start_offset_list[-1]
              start_offset_list.append(start_offset)
  
      def test_add_tensor(self):
          chunk_tensor_index = ChunkTensorIndex(1024)
  
          chunk_tensor_index.add_chunk(
              chunk_id=0,
              comm_info=CommInfo(chunk_type=ChunkType.PARAM_FP32, group_id=0, offset=0),
          )
  
          param_numel_list = [10, 20, 30, 20, 7, 2]
          param_list = []
          offset = 0
          for param_id, numel in enumerate(param_numel_list):
              param = torch.nn.Parameter(torch.zeros(10))
              chunk_tensor_index.add_tensor(
                  chunk_id=0,
                  tensor_id=param_id,
                  start_offset=offset,
                  numel=numel,
                  param=param,
                  access_type=AccessType.DATA,
              )
              offset += numel
              param_list.append(param)
  
          self._check_order(chunk_tensor_index, 0)
  
      def test_append_tensor(self):
          chunk_tensor_index = ChunkTensorIndex(20)
          param_numel_list = [10, 20, 30, 20, 7, 2]
  
          success_list = []
          for param_id, numel in enumerate(param_numel_list):
              param = torch.nn.Parameter(torch.zeros(numel))
              register_param(
                  param, ParamType.CHUNK_BASED, torch.float, f"param_{param_id}"
              )
              is_success = chunk_tensor_index.try_insert_tensor(0, param, AccessType.DATA)
              success_list.append(is_success)
          self.assertEqual(success_list, [True, False, False, False, True, True])
  
      def test_append_tensor_list(self):
          chunk_tensor_index = ChunkTensorIndex(20)
          param_numel_list = [7, 2]
          param_list = []
  
          param_list = []
          for param_id, numel in enumerate(param_numel_list):
              param = torch.nn.Parameter(torch.zeros(numel))
              register_param(
                  param, ParamType.CHUNK_BASED, torch.float, f"param_{param_id}"
              )
              param_list.append(param)
          # 7， 2
          is_success = chunk_tensor_index.try_insert_tensor_list(
              0, param_list, AccessType.DATA
          )
          self.assertTrue(is_success)
  
          # 7， 2， 6， 5
          param_numel_list = [6, 5]
          param_list = []
          for param_id, numel in enumerate(param_numel_list):
              param = torch.nn.Parameter(torch.zeros(numel))
              register_param(
                  param, ParamType.CHUNK_BASED, torch.float, f"param_{param_id}"
              )
              param_list.append(param)
  
          is_success = chunk_tensor_index.try_insert_tensor_list(
              0, param_list, AccessType.DATA
          )
          self.assertTrue(is_success)
  
          # 7， 2，(6), 5
          chunk_tensor_index.delete_tensor(0, param_list[0], AccessType.DATA)
          param_numel_list = [8]
          param_list = []
          for param_id, numel in enumerate(param_numel_list):
              param = torch.nn.Parameter(torch.zeros(numel))
              register_param(
                  param, ParamType.CHUNK_BASED, torch.float, f"param_{param_id}"
              )
              param_list.append(param)
          is_success = chunk_tensor_index.try_insert_tensor_list(
              0, param_list, AccessType.DATA
          )
          self.assertFalse(is_success)
  
          # 7， 2，(6) 5
          param_numel_list = [7]
          param_list = []
          for param_id, numel in enumerate(param_numel_list):
              param = torch.nn.Parameter(torch.zeros(numel))
              register_param(
                  param, ParamType.CHUNK_BASED, torch.float, f"param_{param_id}"
              )
              param_list.append(param)
          is_success = chunk_tensor_index.try_insert_tensor_list(
              0, param_list, AccessType.DATA
          )
          self.assertFalse(is_success)
  
          # 7， 2，(6) 5
          param_numel_list = [1, 2, 3]
          param_list = []
          for param_id, numel in enumerate(param_numel_list):
              param = torch.nn.Parameter(torch.zeros(numel))
              register_param(
                  param, ParamType.CHUNK_BASED, torch.float, f"param_{param_id}"
              )
              param_list.append(param)
          is_success = chunk_tensor_index.try_insert_tensor_list(
              0, param_list, AccessType.DATA
          )
          self.assertTrue(is_success)
  
      def test_chunk_layout_consistency(self):
          r"""
          Check if the chunk layout of optimizer state are aligned to
          param fp16.
          """
          chunk_tensor_index = ChunkTensorIndex(20)
  
          param_numel_list = [10, 5]
          param_list = []
  
          for _, numel in enumerate(param_numel_list):
              param = torch.nn.Parameter(torch.zeros(numel))
              register_param(param, ParamType.CHUNK_BASED, torch.float, f"param_{numel}")
              chunk_tensor_index.try_insert_tensor(0, param, AccessType.DATA)
              param_list.append(param)
  
          param_numel_list = [6, 9]
          for _, numel in enumerate(param_numel_list):
              param = torch.nn.Parameter(torch.zeros(numel))
              register_param(param, ParamType.CHUNK_BASED, torch.float, f"param_{numel}")
              chunk_tensor_index.try_insert_tensor(1, param, AccessType.DATA)
              param_list.append(param)
  
          # Now, we have 2 chunks, (10, 5) (6, 9)
          param_momentum = torch.nn.Parameter(torch.zeros(10))
          register_param(
              param_momentum, ParamType.CHUNK_BASED, torch.float, f"param_{numel}"
          )
          chunk_id = chunk_tensor_index.get_optimizer_state_chunk_id(
              param_list[0], AccessType.DATA, ChunkType.MOMENTUM
          )
          self.assertTrue(chunk_id is None)
  
          chunk_tensor_index.register_optimizer_state_chunk_id(
              param_list[0], AccessType.DATA, ChunkType.MOMENTUM, 3
          )
          chunk_id = chunk_tensor_index.get_optimizer_state_chunk_id(
              param_list[0], AccessType.DATA, ChunkType.MOMENTUM
          )
          self.assertTrue(chunk_id == 3, f"chunk_id is {chunk_id} should be 3")
  
  
  if __name__ == "__main__":
      unittest.main()
  