| 159 | |
| 160 | |
| 161 | class SaveCheckpointCallback(Callback): |
| 162 | def __init__(self, cache_dir, bucket, local_rank=0, has_trained_epoch=0, has_trained_step=0, syn_times=100): |
| 163 | self.cache_dir = os.path.join(cache_dir, f"rank_{local_rank}") |
| 164 | self.local_rank = local_rank |
| 165 | self.has_trained_epoch = has_trained_epoch |
| 166 | self.has_trained_step = has_trained_step |
| 167 | |
| 168 | self.bucket = os.path.join(bucket, f"rank_{local_rank}") |
| 169 | self.syn_times = syn_times |
| 170 | |
| 171 | if not mox.file.exists(self.bucket): |
| 172 | print("Creating checkpoint bucket dir {}".format(self.bucket)) |
| 173 | mox.file.make_dirs(self.bucket) |
| 174 | |
| 175 | def step_end(self, run_context): |
| 176 | cb_params = run_context.original_args() |
| 177 | cur_step = cb_params.cur_step_num + self.has_trained_step |
| 178 | if cur_step % self.syn_times == 0: |
| 179 | print("Copying checkpoint to the buckets start", flush=True) |
| 180 | self.syn_files() |
| 181 | print("Copying checkpoint to the buckets ends", flush=True) |
| 182 | |
| 183 | def syn_files(self): |
| 184 | process = Process(target=mox.file.copy_parallel, args=(self.cache_dir, self.bucket), name="checkpoint_sync") |
| 185 | process.start() |
no outgoing calls
no test coverage detected