Downloads the pretrained checkpoints to start from. Currently only supports S3 and URLs
(starting_checkpoint_load_path: str, local_pretrain_checkpoints_folder: str)
| 241 | |
| 242 | |
| 243 | def download_starting_checkpoint(starting_checkpoint_load_path: str, local_pretrain_checkpoints_folder: str) -> str: |
| 244 | """Downloads the pretrained checkpoints to start from. |
| 245 | |
| 246 | Currently only supports S3 and URLs |
| 247 | """ |
| 248 | load_object_store = None |
| 249 | parsed_path = urlparse(starting_checkpoint_load_path) |
| 250 | if parsed_path.scheme == "s3": |
| 251 | load_object_store = S3ObjectStore(bucket=parsed_path.netloc) |
| 252 | |
| 253 | download_path = parsed_path.path if parsed_path.scheme == "s3" else starting_checkpoint_load_path |
| 254 | os.makedirs(local_pretrain_checkpoints_folder, exist_ok=True) |
| 255 | local_path = os.path.join( |
| 256 | local_pretrain_checkpoints_folder, |
| 257 | get_checkpoint_name_from_path(parsed_path.path), |
| 258 | ) |
| 259 | if not os.path.exists(local_path): |
| 260 | get_file( |
| 261 | destination=local_path, |
| 262 | path=download_path.lstrip("/"), |
| 263 | object_store=load_object_store, |
| 264 | progress_bar=True, |
| 265 | ) |
| 266 | |
| 267 | return local_path |
| 268 | |
| 269 | |
| 270 | def _setup_gpu_queue(num_gpus: int, manager: SyncManager): |
no test coverage detected