Spaces:
Running on L40S
Running on L40S
File size: 3,759 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""Canonical Hydra-group registry for checkpoint SKUs and checkpointer types."""
from typing import Dict
from hydra.core.config_store import ConfigStore
from cosmos_framework.utils import config
from cosmos_framework.checkpoint.dummy import Checkpointer as DummyCheckpointer
from cosmos_framework.utils.config import CheckpointConfig
from cosmos_framework.utils.lazy_config import LazyCall as L
from cosmos_framework.checkpoint.dcp import DistributedCheckpointer
local_object_store = config.ObjectStoreConfig(
enabled=False,
)
pdx_object_store = config.ObjectStoreConfig(
enabled=True,
credentials="credentials/pdx_vfm_checkpoint.secret",
bucket="checkpoints",
)
s3_object_store = config.ObjectStoreConfig(
enabled=True,
credentials="credentials/s3_training.secret",
bucket="bucket4",
)
s3_eu_object_store = config.ObjectStoreConfig(
enabled=True,
credentials="credentials/s3_training_eu.secret",
bucket="checkpoints-eu-west-3",
)
gcp_object_store = config.ObjectStoreConfig(
enabled=True,
credentials="credentials/gcp_checkpoint.secret",
bucket="bucket1",
)
neb_eu_object_store = config.ObjectStoreConfig(
enabled=True,
credentials="credentials/neb_eu.secret",
bucket="nv-01-10206-checkpoint-experiments",
)
CHECKPOINT_LOCAL = CheckpointConfig(
save_to_object_store=local_object_store,
load_from_object_store=local_object_store,
save_iter=5000,
broadcast_via_filesystem=True,
dcp_async_mode_enabled=True,
)
CHECKPOINT_PDX = CheckpointConfig(
save_to_object_store=pdx_object_store,
load_from_object_store=pdx_object_store,
save_iter=5000,
broadcast_via_filesystem=True,
dcp_async_mode_enabled=True,
)
CHECKPOINT_S3 = CheckpointConfig(
save_to_object_store=s3_object_store,
load_from_object_store=s3_object_store,
save_iter=5000,
broadcast_via_filesystem=True,
dcp_async_mode_enabled=True,
)
CHECKPOINT_S3_EU = CheckpointConfig(
save_to_object_store=s3_eu_object_store,
load_from_object_store=s3_eu_object_store,
save_iter=5000,
broadcast_via_filesystem=True,
dcp_async_mode_enabled=True,
)
CHECKPOINT_GCP = CheckpointConfig(
save_to_object_store=gcp_object_store,
save_iter=1000,
load_from_object_store=gcp_object_store,
load_path="",
load_training_state=False,
strict_resume=True,
enable_gcs_patch_in_boto3=True,
dcp_async_mode_enabled=True,
)
CHECKPOINT_NEB_EU = CheckpointConfig(
save_to_object_store=neb_eu_object_store,
load_from_object_store=neb_eu_object_store,
save_iter=2000,
broadcast_via_filesystem=True,
)
def register_checkpoint():
cs = ConfigStore.instance()
cs.store(group="checkpoint", package="checkpoint", name="local", node=CHECKPOINT_LOCAL)
cs.store(group="checkpoint", package="checkpoint", name="pdx", node=CHECKPOINT_PDX)
cs.store(group="checkpoint", package="checkpoint", name="s3", node=CHECKPOINT_S3)
cs.store(group="checkpoint", package="checkpoint", name="s3_eu", node=CHECKPOINT_S3_EU)
cs.store(group="checkpoint", package="checkpoint", name="gcp", node=CHECKPOINT_GCP)
cs.store(group="checkpoint", package="checkpoint", name="neb_eu", node=CHECKPOINT_NEB_EU)
DUMMY_CHECKPOINTER: Dict[str, str] = L(DummyCheckpointer)()
DISTRIBUTED_CHECKPOINTER: Dict[str, str] = L(DistributedCheckpointer)()
def register_ckpt_type():
cs = ConfigStore.instance()
cs.store(group="ckpt_type", package="checkpoint.type", name="dummy", node=DUMMY_CHECKPOINTER)
cs.store(group="ckpt_type", package="checkpoint.type", name="dcp", node=DISTRIBUTED_CHECKPOINTER)
|