text: # Copyright (c) Microsoft Corporation.
  # Licensed under the MIT license.
  
  import gc
  import logging
  import os
  from pathlib import Path
  from typing import List, Dict, Tuple, Literal, Optional
  
  import json_tricks
  import torch
  from torch.nn import Module
  from torch.tensor import Tensor
  
  _logger = logging.getLogger(__name__)
  
  
  class Task:
      # NOTE: If we want to support multi-thread, this part need to refactor, maybe use file and lock to sync.
      _reference_counter = {}
  
      def __init__(self, task_id: int, model_path: str, masks_path: str, config_list_path: str) -> None:
          """
          Parameters
          ----------
          task_id
              The unique id of task.
          model_path
              The path of the unwrapped pytorch model that will be pruned in this task.
          masks_path
              The path of the masks that applied on the model before pruning.
          config_list_path
              The path of the config list that used in this task.
          """
          self.task_id = task_id
          self.model_path = model_path
          self.masks_path = masks_path
          self.config_list_path = config_list_path
  
          self.status: Literal['Pending', 'Running', 'Finished'] = 'Pending'
          self.score: Optional[float] = None
  
          self.state = {}
  
          for ref in self.referenced_paths():
              self._reference_counter.setdefault(ref, 0)
              self._reference_counter[ref] += 1
  
          self._cleaned = False
  
      def to_dict(self) -> Dict:
          return {
              'task_id': self.task_id,
              'model_path': str(self.model_path),
              'masks_path': str(self.masks_path),
              'config_list_path': str(self.config_list_path),
              'status': self.status,
              'score': self.score,
              'state': self.state
          }
  
      def load_data(self) -> Tuple[Module, Dict[str, Dict[str, Tensor]], List[Dict]]:
          """
          Returns
          -------
          Tuple[Module, Dict[str, Dict[str, Tensor]], List[Dict]]
              Return the model pruning in this task, the masks of the model before pruning,
              the config list used in this task.
          """
          model = torch.load(self.model_path)
          masks = torch.load(self.masks_path)
          with Path(self.config_list_path).open('r') as f:
              config_list = json_tricks.load(f)
          return model, masks, config_list
  
      def referenced_paths(self) -> List[str]:
          """
          Return the path list that need to count reference in this task.
          """
          return [self.model_path, self.masks_path, self.config_list_path]
  
      def clean_up(self):
          """
          Counter of referenced file paths subtract 1. If the counter reach 0, then delete the file.
          """
          if not self._cleaned:
              for ref in self.referenced_paths():
                  self._reference_counter[ref] -= 1
                  if self._reference_counter[ref] <= 0:
                      os.remove(ref)
                      if self._reference_counter[ref] < 0:
                          _logger.warning('Referance counter error, the number of %s is %d',
                                          ref, self._reference_counter[ref])
              self._cleaned = True
          else:
              _logger.warning('Already clean up task %d', self.task_id)
  
  
  class TaskResult:
      def __init__(self, task_id: int, compact_model: Module, compact_model_masks: Dict[str, Dict[str, Tensor]],
                   pruner_generated_masks: Dict[str, Dict[str, Tensor]], score: Optional[float]) -> None:
          """
          Parameters
          ----------
          task_id
              The unique id of task.
          compact_model
              The unwrapped compact pytorch model after pruning. If the compact model has been speeduped during the pruning process,
              it will have a smaller structure compare with the model before pruning.
              If the compact model has not been speeduped, it will have the same structure with the model before pruning.
          compact_model_masks
              The masks on the compact model. If the compact model has been speeduped during the pruning process,
              the `compact_model_masks` is always an empty dict. If the compact model has not been speeduped,
              the `compact_model_masks` is same as `pruner_generated_masks`.
          pruner_generated_masks
              The masks that can apply on the before pruning model. It is always the output of `pruner.compress()`.
              TODO: If the compact model has been speeduped, the auto infer masks maybe also need.
          score
              The score of the pruning effect. i.e., the accuracy or latency after pruning.
          """
          self.task_id = task_id
          self.compact_model = compact_model
          self.compact_model_masks = compact_model_masks
          self.pruner_generated_masks = pruner_generated_masks
          self.score = score
  
  
  class BasePruningScheduler:
      def generate_task(self) -> Optional[Task]:
          """
          Returns
          -------
          Optional[Task]
              Return the next pruning task.
          """
          raise NotImplementedError()
  
      def record_task_result(self, task_result: TaskResult):
          """
          Parameters
          ----------
          task_result
              The result of the task
          """
          raise NotImplementedError()
  
      def pruning_one_step(self, task: Task) -> TaskResult:
          """
          Pruning the model defined in task.
  
          Parameters
          ----------
          task
              The pruning task in this step.
  
          Returns
          -------
          TaskResult
              Return the result of the task in this step.
          """
          raise NotImplementedError()
  
      def get_best_result(self) -> Tuple[int, Module, Dict[str, Dict[str, Tensor]], float, List[Dict]]:
          """
          Returns
          -------
          Tuple[int, Module, Dict[str, Dict[str, Tensor]], float, List[Dict]]
              Return the task result that has the best performance,
              inculde task id, the compact model, the masks on the compact model, score and config list used in this task.
          """
          raise NotImplementedError()
  
      def compress(self):
          """
          The pruning schedule main loop.
          """
          task = self.generate_task()
  
          while task is not None:
              task_result = self.pruning_one_step(task)
              self.record_task_result(task_result)
              del task_result
              gc.collect()
              task = self.generate_task()
  