A Manager class for managing the running state of Explorer and Trainer.
| 9 | |
| 10 | |
| 11 | class StateManager: |
| 12 | """A Manager class for managing the running state of Explorer and Trainer.""" |
| 13 | |
| 14 | def __init__( |
| 15 | self, |
| 16 | path: str, |
| 17 | trainer_name: Optional[str] = None, |
| 18 | explorer_name: Optional[str] = None, |
| 19 | config: Optional[Config] = None, |
| 20 | check_config: bool = False, |
| 21 | ): |
| 22 | self.logger = get_logger(__name__, in_ray_actor=True) |
| 23 | self.cache_dir = path |
| 24 | os.makedirs(self.cache_dir, exist_ok=True) |
| 25 | self.stage_state_path = os.path.join(self.cache_dir, "stage_meta.json") |
| 26 | self.explorer_state_path = os.path.join(self.cache_dir, f"{explorer_name}_meta.json") |
| 27 | self.trainer_state_path = os.path.join(self.cache_dir, f"{trainer_name}_meta.json") |
| 28 | self.explorer_server_url_path = os.path.join( |
| 29 | self.cache_dir, f"{explorer_name}_server_url.txt" |
| 30 | ) |
| 31 | if check_config and config is not None: |
| 32 | self._check_config_consistency(config) |
| 33 | |
| 34 | def _check_config_consistency(self, config: Config) -> None: |
| 35 | """Check if the config is consistent with the cache dir backup.""" |
| 36 | backup_config_path = os.path.join(self.cache_dir, "config.json") |
| 37 | if not os.path.exists(backup_config_path): |
| 38 | config.save(backup_config_path) |
| 39 | else: |
| 40 | backup_config = load_config(backup_config_path) |
| 41 | if backup_config != config: |
| 42 | self.logger.warning( |
| 43 | f"The current config is inconsistent with the backup config in {backup_config_path}." |
| 44 | ) |
| 45 | raise ValueError( |
| 46 | f"The current config is inconsistent with the backup config in {backup_config_path}." |
| 47 | ) |
| 48 | |
| 49 | def save_explorer( |
| 50 | self, |
| 51 | current_step: int, |
| 52 | taskset_states: List[Dict], |
| 53 | ) -> None: |
| 54 | with open(self.explorer_state_path, "w", encoding="utf-8") as f: |
| 55 | json.dump( |
| 56 | { |
| 57 | "latest_iteration": current_step, |
| 58 | "taskset_states": taskset_states, |
| 59 | }, |
| 60 | f, |
| 61 | indent=2, |
| 62 | ) |
| 63 | |
| 64 | def load_explorer(self) -> dict: |
| 65 | if os.path.exists(self.explorer_state_path): |
| 66 | try: |
| 67 | with open(self.explorer_state_path, "r", encoding="utf-8") as f: |
| 68 | explorer_meta = json.load(f) |
no outgoing calls