| |
| """State manager.""" |
| import json |
| import os |
| from typing import Dict, List, Optional |
|
|
| from trinity.common.config import Config, load_config |
| from trinity.utils.log import get_logger |
|
|
|
|
| class StateManager: |
| """A Manager class for managing the running state of Explorer and Trainer.""" |
|
|
| def __init__( |
| self, |
| path: str, |
| trainer_name: Optional[str] = None, |
| explorer_name: Optional[str] = None, |
| config: Optional[Config] = None, |
| check_config: bool = False, |
| ): |
| self.logger = get_logger(__name__, in_ray_actor=True) |
| self.cache_dir = path |
| os.makedirs(self.cache_dir, exist_ok=True) |
| self.stage_state_path = os.path.join(self.cache_dir, "stage_meta.json") |
| self.explorer_state_path = os.path.join(self.cache_dir, f"{explorer_name}_meta.json") |
| self.trainer_state_path = os.path.join(self.cache_dir, f"{trainer_name}_meta.json") |
| self.explorer_server_url_path = os.path.join( |
| self.cache_dir, f"{explorer_name}_server_url.txt" |
| ) |
| if check_config and config is not None: |
| self._check_config_consistency(config) |
|
|
| def _check_config_consistency(self, config: Config) -> None: |
| """Check if the config is consistent with the cache dir backup.""" |
| backup_config_path = os.path.join(self.cache_dir, "config.json") |
| if not os.path.exists(backup_config_path): |
| config.save(backup_config_path) |
| else: |
| backup_config = load_config(backup_config_path) |
| if backup_config != config: |
| self.logger.warning( |
| f"The current config is inconsistent with the backup config in {backup_config_path}." |
| ) |
| raise ValueError( |
| f"The current config is inconsistent with the backup config in {backup_config_path}." |
| ) |
|
|
| def save_explorer( |
| self, |
| current_step: int, |
| taskset_states: List[Dict], |
| ) -> None: |
| with open(self.explorer_state_path, "w", encoding="utf-8") as f: |
| json.dump( |
| { |
| "latest_iteration": current_step, |
| "taskset_states": taskset_states, |
| }, |
| f, |
| indent=2, |
| ) |
|
|
| def load_explorer(self) -> dict: |
| if os.path.exists(self.explorer_state_path): |
| try: |
| with open(self.explorer_state_path, "r", encoding="utf-8") as f: |
| explorer_meta = json.load(f) |
| self.logger.info( |
| "----------------------------------\n" |
| "Found existing explorer checkpoint:\n" |
| f" > {explorer_meta}\n" |
| "Continue exploring from this point.\n" |
| "----------------------------------" |
| ) |
| return explorer_meta |
| except Exception as e: |
| self.logger.error(f"Failed to load explore state file: {e}") |
| return {} |
|
|
| def save_explorer_server_url(self, url: str) -> None: |
| with open(self.explorer_server_url_path, "w", encoding="utf-8") as f: |
| f.write(url) |
| self.logger.info(f"Saved explorer server URL to {self.explorer_server_url_path}") |
|
|
| def load_explorer_server_url(self) -> Optional[str]: |
| if os.path.exists(self.explorer_server_url_path): |
| try: |
| with open(self.explorer_server_url_path, "r", encoding="utf-8") as f: |
| url = f.read().strip() |
| self.logger.info( |
| "----------------------------------\n" |
| "Found existing explorer server URL:\n" |
| f" > {url}\n" |
| "----------------------------------" |
| ) |
| return url |
| except Exception as e: |
| self.logger.error(f"Failed to load explorer server URL file: {e}") |
| return None |
|
|
| def save_trainer( |
| self, |
| current_step: int, |
| sample_strategy_state: dict, |
| ) -> None: |
| with open(self.trainer_state_path, "w", encoding="utf-8") as f: |
| json.dump( |
| { |
| "latest_iteration": current_step, |
| "sample_strategy_state": sample_strategy_state, |
| }, |
| f, |
| indent=2, |
| ) |
|
|
| def load_trainer(self) -> dict: |
| if os.path.exists(self.trainer_state_path): |
| try: |
| with open(self.trainer_state_path, "r", encoding="utf-8") as f: |
| trainer_meta = json.load(f) |
| self.logger.info( |
| "----------------------------------\n" |
| "Found existing trainer checkpoint:\n" |
| f" > {trainer_meta}\n" |
| "Continue training from this point.\n" |
| "----------------------------------" |
| ) |
| return trainer_meta |
| except Exception as e: |
| self.logger.warning(f"Failed to load trainer state file: {e}") |
| return {} |
|
|
| def save_stage(self, current_stage: int) -> None: |
| with open(self.stage_state_path, "w", encoding="utf-8") as f: |
| json.dump( |
| { |
| "latest_stage": current_stage, |
| }, |
| f, |
| indent=2, |
| ) |
|
|
| def load_stage(self) -> dict: |
| if os.path.exists(self.stage_state_path): |
| try: |
| with open(self.stage_state_path, "r", encoding="utf-8") as f: |
| stage_meta = json.load(f) |
| self.logger.info( |
| "----------------------------------\n" |
| "Found existing stage checkpoint:\n" |
| f" > {stage_meta}\n" |
| "Continue from this point.\n" |
| "----------------------------------" |
| ) |
| return stage_meta |
| except Exception as e: |
| self.logger.warning(f"Failed to load stage state file: {e}") |
| return {} |
|
|
