Spaces:
Running on Zero
Running on Zero
File size: 2,956 Bytes
a1833c8 | 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""Export all nodes from the config store to yaml files."""
from cosmos_framework.inference.common.init import init_script
init_script(env={"COSMOS_TRAINING": "1"})
import importlib
from typing import Annotated, Any, Generator
import pydantic
import tyro
import yaml
from hydra.core.config_store import ConfigNode, ConfigStore
from cosmos_framework.inference.common.args import ResolvedPath, tyro_cli
from cosmos_framework.inference.common.config import _InvalidMode, apply_config_replacements, unstructure_config
from cosmos_framework.utils import config_helper, log
class Args(pydantic.BaseModel):
output_dir: Annotated[ResolvedPath, tyro.conf.arg(aliases=("-o",))]
config_file: str = "configs/base/config.py"
"""Config module path."""
invalid: _InvalidMode = "error"
"""How to handle unknown field types."""
def _iter_config_nodes(repo: dict) -> Generator[ConfigNode]:
"""Iteratively yield all config store nodes."""
# Depth-First Search.
stack = [repo]
while stack:
current_dict = stack.pop()
for value in current_dict.values():
if isinstance(value, dict):
stack.append(value)
else:
yield value
def _save_config(
config: Any, name: str, group: str | None, package: str | None, output_dir: ResolvedPath, invalid: _InvalidMode
):
config_dict = unstructure_config(config, invalid=invalid)
config_str = yaml.safe_dump(config_dict)
config_str = apply_config_replacements(config_str)
node_file = output_dir.joinpath(*filter(None, [group, name])).with_suffix(".yaml")
node_file.parent.mkdir(parents=True, exist_ok=True)
with node_file.open("w") as f:
if package:
f.write(f"# @package {package}\n")
f.write(config_str)
def export_config_store(args: Args):
config_module = importlib.import_module(config_helper.get_config_module(args.config_file))
config = config_module.make_config()
cs = ConfigStore.instance()
_save_config(
config, "base_config.yaml", group=None, package="_global_", output_dir=args.output_dir, invalid=args.invalid
)
for node in _iter_config_nodes(cs.repo):
if node.name.startswith("_"):
continue
if node.group in ["hydra"]:
continue
try:
_save_config(
node.node,
name=node.name,
group=node.group,
package=node.package,
output_dir=args.output_dir,
invalid=args.invalid,
)
except Exception as e:
log.error(f"Error saving config '{node.name}': {e}")
continue
def main():
args = tyro_cli(Args, description=__doc__)
export_config_store(args)
if __name__ == "__main__":
main()
|