Spaces:
Running on L40S
Running on L40S
File size: 6,880 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
import os
from abc import ABC, abstractmethod
from typing import Optional
import torch
from cosmos_framework.utils.config import CheckpointConfig, JobConfig
from cosmos_framework.utils.flags import INTERNAL
from cosmos_framework.model._base import ImaginaireModel
from cosmos_framework.utils import callback
from cosmos_framework.utils.easy_io import easy_io
class AbstractCheckpointer(ABC):
"""The checkpointer class. Supports checkpoint saving/loading to both local disk or object store."""
def __init__(
self,
config_checkpoint: CheckpointConfig,
config_job: JobConfig,
callbacks: Optional[callback.CallBackGroup] = None,
):
"""Constructor of the checkpointer.
Args:
config_checkpoint (CheckpointConfig): The config object for the checkpointer.
"""
self.config_checkpoint = config_checkpoint
# Set the callback functions.
self.callbacks = callbacks
self.save_to_object_store = config_checkpoint.save_to_object_store.enabled
self.load_from_object_store = config_checkpoint.load_from_object_store.enabled
# Set checkpoint directories for local and object store paths
self._local_dirname = os.path.join(config_job.path_local, "checkpoints")
self._object_store_dirname = os.path.join(config_job.path, "checkpoints")
self.strict_resume = config_checkpoint.strict_resume
load_path = config_checkpoint.load_path or None
if not INTERNAL:
from cosmos_framework.utils.checkpoint_db import download_checkpoint_v2
if load_path:
load_path = download_checkpoint_v2(load_path)
self.load_path = load_path
self.load_training_state = config_checkpoint.load_training_state
self.only_load_scheduler_state = config_checkpoint.only_load_scheduler_state
self.save_thread = None
self.verbose = config_checkpoint.verbose
self.keys_not_to_resume = config_checkpoint.keys_not_to_resume
self.keys_to_skip_loading = getattr(config_checkpoint, "keys_to_skip_loading", [])
self.broadcast_via_filesystem = config_checkpoint.broadcast_via_filesystem
# Create the object store client interface.
if config_checkpoint.load_from_object_store.enabled:
self.load_s3_backend_key = "_ckpt_s3_loader"
easy_io.set_s3_backend(
key="_ckpt_s3_loader",
backend_args={
"backend": "s3",
"path_mapping": {
"s3://ckpt/": f"s3://{config_checkpoint.load_from_object_store.bucket}/",
},
"s3_credential_path": config_checkpoint.load_from_object_store.credentials,
},
)
else:
self.load_s3_backend_key = None
if config_checkpoint.save_to_object_store.enabled:
self.save_s3_backend_key = "_ckpt_s3_saver"
easy_io.set_s3_backend(
key="_ckpt_s3_saver",
backend_args={
"backend": "s3",
"path_mapping": {
"s3://ckpt/": f"s3://{config_checkpoint.save_to_object_store.bucket}/",
},
"s3_credential_path": config_checkpoint.save_to_object_store.credentials,
},
)
else:
self.save_s3_backend_key = None
@abstractmethod
def save(
self,
model: ImaginaireModel,
optimizer: torch.optim.Optimizer,
scheduler: torch.optim.lr_scheduler.LRScheduler,
grad_scaler: torch.amp.GradScaler,
iteration: int,
) -> None:
pass
@abstractmethod
def load(
self,
model: ImaginaireModel,
optimizer: Optional[torch.optim.Optimizer] = None,
scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
grad_scaler: Optional[torch.amp.GradScaler] = None,
) -> int:
pass
@property
def save_bucket(self):
"""Get the bucket name for saving checkpoints."""
return self.config_checkpoint.save_to_object_store.bucket if self.save_to_object_store else None
@property
def load_bucket(self):
"""Get the bucket name for loading checkpoints."""
return self.config_checkpoint.load_from_object_store.bucket if self.load_from_object_store else None
@property
def save_dirname(self):
return (
f"s3://{self.save_bucket}/{self._object_store_dirname}"
if self.save_to_object_store
else self._local_dirname
)
@property
def load_dirname(self):
return (
f"s3://{self.load_bucket}/{self._object_store_dirname}"
if self.load_from_object_store
else self._local_dirname
)
def finalize(self) -> None:
"""Finalize the checkpointer."""
if self.save_thread:
self.save_thread.join()
def _read_latest_checkpoint_file(self) -> str | None:
"""Get the file name of the latest saved checkpoint. If it doesn't exist, return None.
Returns:
checkpoint_file (str | None): file name of the latest saved checkpoint.
"""
checkpoint_file = None
checkpoint_path = os.path.join(self.load_dirname, "latest_checkpoint.txt")
if easy_io.exists(f"{checkpoint_path}", backend_key=self.load_s3_backend_key):
checkpoint_file = easy_io.load(f"{checkpoint_path}", backend_key=self.load_s3_backend_key).strip()
return checkpoint_file
def has_resumable_checkpoint(self) -> bool:
"""True iff a ``latest_checkpoint.txt`` exists in the load directory."""
return self._read_latest_checkpoint_file() is not None
def _write_latest_checkpoint_file(self, checkpoint_file: str) -> None:
"""Track the file name of the latest saved checkpoint.
Args:
checkpoint_file (str): file name of the latest saved checkpoint.
"""
content = f"{checkpoint_file}\n"
checkpoint_path = os.path.join(self.save_dirname, "latest_checkpoint.txt")
easy_io.dump(
content,
checkpoint_path,
backend_key=self.save_s3_backend_key,
)
def _check_checkpoint_exists(self, checkpoint_path: str) -> None:
"""If the file checkpoint_path does not exist, raise an error.
Args:
checkpoint_path (str): full path to the checkpoint.
"""
if not easy_io.exists(f"{checkpoint_path}", backend_key=self.load_s3_backend_key):
raise FileNotFoundError(f"File not found (object store): {checkpoint_path}")
|