File size: 6,167 Bytes
8c9ba62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | # -*- coding: utf-8 -*-
"""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 {}
|