Run trainer.
(config: Config, *, timeout: Optional[float] = None)
| 137 | |
| 138 | |
| 139 | def train(config: Config, *, timeout: Optional[float] = None) -> StageStatus: |
| 140 | """Run trainer.""" |
| 141 | from trinity.trainer.trainer import Trainer |
| 142 | |
| 143 | trainer = Trainer.get_actor(config) |
| 144 | startup_started_at = time.perf_counter() |
| 145 | startup_time_sec: Optional[float] = None |
| 146 | run_started_at: Optional[float] = None |
| 147 | |
| 148 | try: |
| 149 | ray.get(trainer.prepare.remote(), timeout=timeout) |
| 150 | startup_time_sec = time.perf_counter() - startup_started_at |
| 151 | |
| 152 | run_started_at = time.perf_counter() |
| 153 | ray.get(trainer.sync_weight.remote(), timeout=timeout) |
| 154 | ray.get(trainer.train.remote(), timeout=timeout) |
| 155 | execution_time_sec = time.perf_counter() - run_started_at |
| 156 | return StageStatus( |
| 157 | stage="train", |
| 158 | success=True, |
| 159 | startup_time_sec=startup_time_sec, |
| 160 | execution_time_sec=execution_time_sec, |
| 161 | total_time_sec=time.perf_counter() - startup_started_at, |
| 162 | ) |
| 163 | except Exception as exc: |
| 164 | error = _build_stage_error(exc) |
| 165 | logger.error(f"Trainer failed:\n{error.traceback_text}") |
| 166 | execution_time_sec = ( |
| 167 | time.perf_counter() - run_started_at if run_started_at is not None else None |
| 168 | ) |
| 169 | return StageStatus( |
| 170 | stage="train", |
| 171 | success=False, |
| 172 | startup_time_sec=startup_time_sec, |
| 173 | execution_time_sec=execution_time_sec, |
| 174 | total_time_sec=time.perf_counter() - startup_started_at, |
| 175 | error=error, |
| 176 | ) |
| 177 | finally: |
| 178 | ray.get(trainer.shutdown.remote(), timeout=timeout) |
| 179 | |
| 180 | |
| 181 | def serve(config: Config, *, timeout: Optional[float] = None) -> StageStatus: |