| 50 | |
| 51 | |
| 52 | class UploadToLocalFile(UploadTaskGroupBuilder): |
| 53 | def __init__(self, dest_dir): |
| 54 | self.dest_dir = dest_dir |
| 55 | |
| 56 | def build(self, epoch, checkpoint_manager): |
| 57 | with TaskGroup(WorkspaceType.GLOBAL) as upload_task_group: |
| 58 | for node, manager in checkpoint_manager._node_managers: |
| 59 | with Node(str(node)), Task(): |
| 60 | src_path = db_name(epoch, manager._node_name, manager._db_prefix) |
| 61 | dest_path = os.path.join(self.dest_dir, str(node)) |
| 62 | ops.Python((local_copy_op, |
| 63 | [src_path, dest_path], {}))([], []) |
| 64 | return upload_task_group |
| 65 | |
| 66 | |
| 67 | class TestCheckpoint(TestCase): |
no outgoing calls
searching dependent graphs…