| 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() |