(self, *args, **kwargs)
| 320 | return output, self.util_reports |
| 321 | |
| 322 | def checkpoint(self, *args, **kwargs): |
| 323 | if self.monitor_runner is not None: |
| 324 | report: UtilizationReport = self.monitor_runner.stop_and_report( |
| 325 | num_examples=self.num_examples_processed_in_current_job |
| 326 | ) |
| 327 | report.job_type = JobType.GENERATION.value |
| 328 | self.util_reports.append(report) |
| 329 | |
| 330 | self.num_examples_processed_in_current_job = 0 |
| 331 | |
| 332 | if not self.setup_done: |
| 333 | # Nothing got processed yet, so reset and try again |
| 334 | print("Preemted during setup, resetting") |
| 335 | self.reset_state(self.num_checkpoints) |
| 336 | elif ( |
| 337 | len(self.processed_outputs) + len(self.prompts_to_process) |
| 338 | != self.original_num_chunks |
| 339 | ): |
| 340 | print("Race condition detected, starting job over") |
| 341 | # There was an interruption as we were changing state, lets just start over |
| 342 | self.reset_state(self.num_checkpoints) |
| 343 | else: |
| 344 | percent_complete = len(self.processed_outputs) / self.original_num_chunks # type: ignore |
| 345 | print(f"Checkpointing {percent_complete*100:.2f}% complete") |
| 346 | |
| 347 | # Set to none so this can be pickled |
| 348 | self.monitor_runner = None |
| 349 | |
| 350 | return super().checkpoint(*args, **kwargs) |
| 351 | |
| 352 | |
| 353 | class QueryWorkerToServerTask(SubmititWorker): |
nothing calls this directly
no test coverage detected