from __future__ import annotations import argparse import logging import os import random import tempfile import uuid from dataclasses import dataclass, field from enum import Enum from typing import TYPE_CHECKING, Optional # import clusterscope try: import clusterscope except Exception: clusterscope = None import hydra import numpy as np import torch from omegaconf import OmegaConf from omegaconf.errors import InterpolationKeyError from onescience.utils.uma.common import gp_utils if TYPE_CHECKING: from omegaconf import DictConfig from onescience.utils.uma.components.reducer import Reducer from onescience.utils.uma.components.runner import Runner from submitit import AutoExecutor from submitit.core.utils import JobPaths, cloudpickle_dump from submitit.helpers import Checkpointable, DelayedSubmission from submitit.slurm.slurm import SlurmJobEnvironment from onescience.distributed.manager import DistributedManager from onescience.utils.uma.common import distutils from onescience.monitoring.uma.logger import WandBSingletonLogger from onescience.monitoring.uma.runtime_logging import setup_logging from onescience.utils.uma.common.utils import ( get_commit_hash, get_timestamp_uid, setup_env_vars, ) # this effects the cli only since the actual job will be run in subprocesses or remoe logging.basicConfig(level=logging.INFO) ALLOWED_TOP_LEVEL_KEYS = {"job", "runner", "reducer"} LOG_DIR_NAME = "logs" CHECKPOINT_DIR_NAME = "checkpoints" RESULTS_DIR = "results" CONFIG_FILE_NAME = "canonical_config.yaml" PREEMPTION_STATE_DIR_NAME = "preemption_state" class SchedulerType(str, Enum): LOCAL = "local" SLURM = "slurm" class DeviceType(str, Enum): CPU = "cpu" CUDA = "cuda" class RunType(str, Enum): RUN = "run" REDUCE = "reduce" class DistributedInitMethod(str, Enum): TCP = "tcp" FILE = "file" @dataclass class SlurmConfig: mem_gb: int = 80 timeout_hr: int = 168 cpus_per_task: int = 8 partition: Optional[str] = ( None # omegaconf in python 3.9 does not backport annotations ) qos: Optional[str] = None # omegaconf in python 3.9 does not backport annotations account: Optional[str] = ( None # omegaconf in python 3.9 does not backport annotations ) additional_parameters: Optional[dict] = None # 字典格式，用于存储环境变量和配置 @dataclass class SchedulerConfig: mode: SchedulerType = SchedulerType.LOCAL distributed_init_method: DistributedInitMethod = DistributedInitMethod.TCP ranks_per_node: int = 1 num_nodes: int = 1 num_array_jobs: int = 1 # 新增：仅 LOCAL+elastic 多节点时使用（也可不填，见下文） rdzv_backend: str = "c10d" rdzv_endpoint: Optional[str] = None run_id: str = field(default_factory=lambda: f"run_{uuid.uuid4().hex[:8]}") slurm: SlurmConfig = field(default_factory=lambda: SlurmConfig) @dataclass class SlurmEnv: # reflects the job_id given by submitit (slurm id with array job id and array task id if they exist) job_id: Optional[str] = ( None # omegaconf in python 3.9 does not backport annotations ) # reflects SLURM_JOB_ID only raw_job_id: Optional[str] = ( None # omegaconf in python 3.9 does not backport annotations ) # SLURM_ARRAY_JOB_ID array_job_id: Optional[str] = ( None # omegaconf in python 3.9 does not backport annotations ) # SLURM_ARRAY_TASK_ID array_task_id: Optional[str] = ( None # omegaconf in python 3.9 does not backport annotations ) # reflects SLURM_RESTART_COUNT env variable restart_count: Optional[str] = ( None # omegaconf in python 3.9 does not backport annotations ) @dataclass class Metadata: # read-only metadata about the job, not user inputs commit: str log_dir: str checkpoint_dir: str results_dir: str config_path: str preemption_checkpoint_dir: str cluster_name: str array_job_num: int = 0 slurm_env: SlurmEnv = field(default_factory=lambda: SlurmEnv()) @dataclass class JobConfig: run_name: str = field( default_factory=lambda: get_timestamp_uid() + uuid.uuid4().hex.upper()[0:4] ) timestamp_id: str = field(default_factory=lambda: get_timestamp_uid()) run_dir: str = field(default_factory=lambda: tempfile.TemporaryDirectory().name) device_type: DeviceType = DeviceType.CUDA debug: bool = False scheduler: SchedulerConfig = field(default_factory=lambda: SchedulerConfig) logger: Optional[dict] = ( None # omegaconf in python 3.9 does not backport annotations ) seed: int = 0 deterministic: bool = False runner_state_path: Optional[str] = ( None # omegaconf in python 3.9 does not backport annotations ) # read-only metadata about the job, not user inputs metadata: Optional[Metadata] = ( None # omegaconf in python 3.9 does not backport annotations ) graph_parallel_group_size: Optional[int] = None def __post_init__(self) -> None: self.run_dir = os.path.abspath(self.run_dir) try: try: cluster = clusterscope.cluster() except Exception as e: # DCU 或受限环境下，clusterscope 可能调用 nvidia-smi 失败 cluster = os.environ.get("ONESCIENCE_CLUSTER", "sghpc") if os.environ.get("LOCAL_RANK", "0") == "0": print( f"[WARN] clusterscope.cluster() failed " f"({type(e).__name__}: {e}); fallback to '{cluster}'" ) except RuntimeError: cluster = "" self.metadata = Metadata( commit=get_commit_hash(), log_dir=os.path.join(self.run_dir, self.timestamp_id, LOG_DIR_NAME), checkpoint_dir=os.path.join( self.run_dir, self.timestamp_id, CHECKPOINT_DIR_NAME ), results_dir=os.path.join(self.run_dir, self.timestamp_id, RESULTS_DIR), config_path=os.path.join(self.run_dir, self.timestamp_id, CONFIG_FILE_NAME), preemption_checkpoint_dir=os.path.join( self.run_dir, self.timestamp_id, CHECKPOINT_DIR_NAME, PREEMPTION_STATE_DIR_NAME, ), cluster_name=cluster, ) def _set_seeds(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def _set_deterministic_mode() -> None: # this is required for full cuda deterministic mode logging.info("Setting deterministic mode!") os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False torch.use_deterministic_algorithms(True) def _get_slurm_env() -> SlurmEnv: slurm_job_env = SlurmJobEnvironment() try: slurm_env = SlurmEnv( job_id=slurm_job_env.job_id, raw_job_id=slurm_job_env.raw_job_id, array_job_id=slurm_job_env.array_job_id, array_task_id=slurm_job_env.array_task_id, restart_count=os.environ.get("SLURM_RESTART_COUNT"), ) except KeyError: # slurm environment variables are undefined, running locally slurm_env = SlurmEnv() return slurm_env def remove_runner_state_from_submission(log_folder: str, job_id: str) -> None: # (HACK) Decouple the job from the runner state by manually modifying it # this ensures the saved runner state is not re-submitted in the event of a node failure # ie: if the job was started at state t=T, a requeue during node failure would resubmit the job # starting at state t=T again without calling the checkpoint callback, losing all progress in between. job_path = JobPaths(folder=log_folder, job_id=job_id) if os.path.isfile(job_path.submitted_pickle): submission_obj = DelayedSubmission.load(job_path.submitted_pickle) submission_obj.args[0].job.runner_state_path = None cloudpickle_dump(submission_obj, job_path.submitted_pickle) class Submitit(Checkpointable): def __init__(self) -> None: self.config = None self.runner = None self.reducer = None def __call__( self, dict_config: DictConfig, run_type: RunType = RunType.RUN ) -> None: self.config = dict_config self.run_type = run_type # modify the config metadata to add slurm info if they exist self.config.job.metadata.slurm_env = _get_slurm_env() setup_env_vars() setup_logging() dist_config = map_job_config_to_dist_config(self.config.job) logging.info("Setting up distributed backend...") distutils.setup(dist_config) distutils.synchronize() if ( distutils.is_master() and self.config.job.scheduler.mode == SchedulerType.SLURM ): # this pickle file is shared across all processes so can only modify this on the main rank remove_runner_state_from_submission( dict_config.job.metadata.log_dir, self.config.job.metadata.slurm_env.job_id, ) if self.config.job.graph_parallel_group_size is not None: logging.info("Setting up graph parallel...") gp_utils.setup_graph_parallel_groups( self.config.job.graph_parallel_group_size, dist_config["distributed_backend"], ) self._init_logger() print(f"rank:{distutils.get_rank()}") _set_seeds(self.config.job.seed) if self.config.job.deterministic: _set_deterministic_mode() if run_type == RunType.RUN: logging.info("Calling runner.run() ...") self.runner: Runner = hydra.utils.instantiate(self.config.runner) self.runner.job_config = self.config.job # must call resume state AFTER the runner has been initialized self.runner.load_state(self.config.job.runner_state_path) self.runner.run() elif run_type == RunType.REDUCE: logging.info("Calling reducer.reduce() ...") self.reducer: Reducer = hydra.utils.instantiate(self.config.reducer) self.reducer.job_config = self.config.job self.reducer.runner_config = self.config.runner # must call resume state AFTER the runner has been initialized self.reducer.load_state(self.config.job.runner_state_path) self.reducer.reduce() else: raise ValueError(f"run type {run_type} is not recognized!") distutils.cleanup() def _init_logger(self) -> None: if ( self.config.job.logger and distutils.is_master() and not self.config.job.debug and self.config.job.metadata.array_job_num == 0 ): # get a partial function from the config and instantiate wandb with it # currently code assumes that we only use the WandBSingletonLogger logger_initializer = hydra.utils.instantiate(self.config.job.logger) simple_config = OmegaConf.to_container( self.config, resolve=True, throw_on_missing=True ) logger_initializer( config=simple_config, run_id=self.config.job.timestamp_id, run_name=self.config.job.run_name, log_dir=self.config.job.metadata.log_dir, ) def checkpoint(self, *args, **kwargs) -> DelayedSubmission: logging.error("Submitit checkpointing callback is triggered") save_path = self.config.job.metadata.preemption_checkpoint_dir cfg_copy = self.config.copy() # only assign if the save was successful cfg_copy.job.runner_state_path = None if ( self.run_type == RunType.RUN and self.runner.save_state(save_path, is_preemption=True) ) or ( self.run_type == RunType.REDUCE and self.reducer.save_state(save_path, is_preemption=True) ): cfg_copy.job.runner_state_path = save_path if WandBSingletonLogger.initialized(): WandBSingletonLogger.get_instance().mark_preempting() logging.info( f"Submitit checkpointing callback is completed, resuming with use the following state: {save_path}" ) return DelayedSubmission(Submitit(), cfg_copy) def map_job_config_to_dist_config(job_cfg: JobConfig) -> dict: scheduler_config = job_cfg.scheduler return { "world_size": scheduler_config.num_nodes * scheduler_config.ranks_per_node, "distributed_backend": ( "gloo" if job_cfg.device_type == DeviceType.CPU else "nccl" ), "submit": scheduler_config.mode == SchedulerType.SLURM, "cpu": job_cfg.device_type == DeviceType.CPU, "init_method": scheduler_config.distributed_init_method, # for distributed shared file initialization "shared_file_dir": os.path.join(job_cfg.run_dir, job_cfg.timestamp_id), "array_job_num": job_cfg.metadata.array_job_num, } def get_canonical_config(config: DictConfig) -> DictConfig: # manually initialize metadata, because OmegaConf currently doesn't call __post_init__ on dataclasses job = OmegaConf.to_object(config.job) job.__post_init__() config.job = job # check that each key other than the allowed top level keys are used in config # find all top level keys are not in the allowed set all_keys = set(config.keys()).difference(ALLOWED_TOP_LEVEL_KEYS) used_keys = set() for key in all_keys: # make a copy of all keys except the key in question copy_cfg = OmegaConf.create({k: v for k, v in config.items() if k != key}) try: OmegaConf.resolve(copy_cfg) except InterpolationKeyError: # if this error is thrown, this means the key was actually required used_keys.add(key) unused_keys = all_keys.difference(used_keys) if unused_keys != set(): raise ValueError( f"Found unused keys in the config: {unused_keys}, please remove them!, only keys other than {ALLOWED_TOP_LEVEL_KEYS} or ones that are used as variables are allowed." ) # resolve the config to fully replace the variables and delete all top level keys except for the ALLOWED_TOP_LEVEL_KEYS for _k in ("ONESCIENCE_DATASETS_DIR", "ONESCIENCE_MODELS_DIR"): if OmegaConf.select(config, _k, default=None) is None: _v = os.environ.get(_k) if _v: OmegaConf.update(config, _k, _v, force_add=True) OmegaConf.resolve(config) return OmegaConf.create( {k: v for k, v in config.items() if k in ALLOWED_TOP_LEVEL_KEYS} ) def get_hydra_config_from_yaml( config_yml: str, overrides_args: list[str] ) -> DictConfig: # Load the configuration from the file os.environ["HYDRA_FULL_ERROR"] = "1" config_directory = os.path.dirname(os.path.abspath(config_yml)) config_name = os.path.basename(config_yml) hydra.initialize_config_dir(config_directory, version_base="1.1") cfg = hydra.compose(config_name=config_name, overrides=overrides_args) # merge default structured config with initialized job object cfg = OmegaConf.merge({"job": OmegaConf.structured(JobConfig)}, cfg) # canonicalize config (remove top level keys that just used replacing variables) return get_canonical_config(cfg) def _runner_wrapper(config: DictConfig, run_type: RunType = RunType.RUN): # This is needed when using elastic_launch for local runs since it looks for # the __name__ attribute of the function, Submitit.__call__ does not have one Submitit()(config, run_type) def _run_under_torchrun(cfg: DictConfig) -> None: """在 torchrun / srun+torchrun 启动的 worker 进程里运行训练。 调用方已经是 rank N 之一（LOCAL_RANK 已由 torchrun 注入），因此这里 绝对不能再走 submitit 或 elastic_launch，而是直接： 1) setup_env_vars / setup_logging 2) DistributedManager.initialize() (读取 RANK/LOCAL_RANK/WORLD_SIZE 等 env 变量) 3) gp_utils (可选) 4) _set_seeds + deterministic 5) WandB 日志 (rank0) 6) hydra.utils.instantiate(cfg.runner) + runner.run() 7) DistributedManager.cleanup() """ setup_env_vars() setup_logging() world_size = int(os.environ.get("WORLD_SIZE", "1")) logging.info( f"Init distributed from ENV: RANK={os.environ.get('RANK')}, " f"LOCAL_RANK={os.environ.get('LOCAL_RANK')}, WORLD_SIZE={world_size}, " f"MASTER={os.environ.get('MASTER_ADDR')}:{os.environ.get('MASTER_PORT')}" ) DistributedManager.initialize() if cfg.job.graph_parallel_group_size is not None: backend = "gloo" if cfg.job.device_type == DeviceType.CPU else "nccl" gp_utils.setup_graph_parallel_groups( cfg.job.graph_parallel_group_size, backend ) _set_seeds(cfg.job.seed) if cfg.job.deterministic: _set_deterministic_mode() dm = DistributedManager() if cfg.job.logger and dm.rank == 0 and not cfg.job.debug: logger_initializer = hydra.utils.instantiate(cfg.job.logger) simple_config = OmegaConf.to_container( cfg, resolve=True, throw_on_missing=True ) logger_initializer( config=simple_config, run_id=cfg.job.timestamp_id, run_name=cfg.job.run_name, log_dir=cfg.job.metadata.log_dir, ) runner = hydra.utils.instantiate(cfg.runner) runner.job_config = cfg.job runner.load_state(cfg.job.runner_state_path) runner.run() DistributedManager.cleanup() def main( args: argparse.Namespace | None = None, override_args: list[str] | None = None ): if args is None: parser = argparse.ArgumentParser() parser.add_argument("-c", "--config", type=str, required=True) args, override_args = parser.parse_known_args() cfg = get_hydra_config_from_yaml(args.config, override_args) log_dir = cfg.job.metadata.log_dir os.makedirs(cfg.job.run_dir, exist_ok=True) os.makedirs(log_dir, exist_ok=True) OmegaConf.save(cfg, cfg.job.metadata.config_path) logging.info(f"saved canonical config to {cfg.job.metadata.config_path}") scheduler_cfg = cfg.job.scheduler # ------------------------------------------------------------------ # torchrun / srun+torchrun worker 分支 # LOCAL_RANK 由 torchrun 注入, 说明我们已经在 rank N 的子进程中, # 不应再走 submitit / elastic_launch (否则会递归 spawn N*N 个进程)。 # demo/run.sh 会把 LOCAL 调度与 torchrun 启动器组合使用, 进入这里。 # ------------------------------------------------------------------ if ( os.environ.get("LOCAL_RANK") is not None and scheduler_cfg.mode == SchedulerType.LOCAL ): _run_under_torchrun(cfg) return logging.info(f"Running fairchemv2 cli with {cfg}") if scheduler_cfg.mode == SchedulerType.SLURM: # Run on cluster assert ( os.getenv("SLURM_SUBMIT_HOST") is None ), "SLURM DID NOT SUBMIT JOB!! Please do not submit jobs from an active slurm job (srun or otherwise)" executor = AutoExecutor(folder=log_dir, slurm_max_num_timeout=3) executor.update_parameters( name=cfg.job.run_name, mem_gb=scheduler_cfg.slurm.mem_gb, timeout_min=scheduler_cfg.slurm.timeout_hr * 60, slurm_partition=scheduler_cfg.slurm.partition, # gpus_per_node=scheduler_cfg.ranks_per_node, gpus_per_node=None, # 不设置 gpus-per-node cpus_per_task=scheduler_cfg.slurm.cpus_per_task, tasks_per_node=scheduler_cfg.ranks_per_node, nodes=scheduler_cfg.num_nodes, slurm_qos=scheduler_cfg.slurm.qos, slurm_account=scheduler_cfg.slurm.account, # 新增：把 YAML 里的 additional_parameters 传给 sbatch slurm_additional_parameters=getattr( scheduler_cfg.slurm, "additional_parameters", None ), ) if scheduler_cfg.num_array_jobs == 1: job = executor.submit(Submitit(), cfg) logging.info( f"Submitted job id: {cfg.job.timestamp_id}, slurm id: {job.job_id}, logs: {cfg.job.metadata.log_dir}" ) jobs = [job] elif scheduler_cfg.num_array_jobs > 1: executor.update_parameters( slurm_array_parallelism=scheduler_cfg.num_array_jobs, ) jobs = [] with executor.batch(): for job_number in range(scheduler_cfg.num_array_jobs): _cfg = cfg.copy() _cfg.job.metadata.array_job_num = job_number job = executor.submit(Submitit(), _cfg) jobs.append(job) logging.info(f"Submitted {len(jobs)} jobs: {jobs[0].job_id.split('_')[0]}") if "reducer" in cfg: job_id = jobs[0].job_id.split("_")[0] executor.update_parameters( name=f"{cfg.job.run_name}_reduce", # set a single node, or do we want the same config as the Runner or a separate JobConfig nodes=1, slurm_dependency=f"afterok:{job_id}", slurm_additional_parameters={ "kill-on-invalid-dep": "yes" }, # kill the reducer if run fails ) executor.submit(Submitit(), cfg, RunType.REDUCE) else: from torch.distributed.launcher.api import LaunchConfig, elastic_launch if scheduler_cfg.num_nodes > 1: cfg.job.scheduler.num_nodes = 1 logging.warning( f"You cannot use more than one node (scheduler_cfg.num_nodes={scheduler_cfg.num_nodes}) in LOCAL mode, over-riding to 1 node" ) if scheduler_cfg.ranks_per_node > 1: logging.info( f"Running in local mode with {scheduler_cfg.ranks_per_node} ranks using device_type:{cfg.job.device_type}" ) launch_config = LaunchConfig( min_nodes=1, max_nodes=1, nproc_per_node=scheduler_cfg.ranks_per_node, rdzv_backend="c10d", max_restarts=0, ) elastic_launch(launch_config, _runner_wrapper)(cfg) if "reducer" in cfg: elastic_launch(launch_config, _runner_wrapper)(cfg, RunType.REDUCE) else: logging.info("Running in local mode without elastic launch") distutils.setup_env_local() Submitit()(cfg) if "reducer" in cfg: Submitit()(cfg, RunType.REDUCE) if __name__ == "__main__": main()