Upload folder using huggingface_hub (part 2)

1f88cea verified about 2 months ago

8.52 kB

import datetime
import os
import warnings
from pathlib import Path
import time
from typing import Optional
import pytorch_lightning as pl
import torch
from omegaconf import OmegaConf, listconfig
from pytorch_lightning import LightningModule
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.strategies import DDPStrategy

from boltzgen.task.task import Task
from boltzgen.task.train.data import DataConfig, TrainingDataModule


class Training(Task):
    """Training configuration."""

    def __init__(
        self,
        data: DataConfig,
        model: LightningModule,
        output: str,
        name: str,
        trainer: Optional[dict] = None,
        resume: Optional[str] = None,
        pretrained: Optional[str] = None,
        wandb: Optional[dict] = None,
        disable_checkpoint: bool = False,
        slurm: bool = False,
        matmul_precision: Optional[str] = None,
        find_unused_parameters: Optional[bool] = False,
        ddp_timeout_seconds: Optional[int] = 6400,
        save_every_n_train_steps: Optional[int] = None,
        save_top_k: Optional[int] = 1,
        validation_only: bool = False,
        debug: bool = False,
        strict_loading: bool = True,
        metric_mode: str = "max",
    ) -> None:
        """Initialize training configuration.
        Parameters
        ----------
        data : DataConfig
            The data configuration.
        model : LightningModule
            The model.
        output : str
            The output directory.
        name : str
            The output subdirectory and wandb run name.
        trainer : Optional[dict], optional
            The trainer configuration, by default None.
        resume : Optional[str], optional
            The resume checkpoint, by default None
        pretrained : Optional[str], optional
            The pretrained model, by default None
        wandb : Optional[dict], optional
            The wandb configuration, by default None
        disable_checkpoint : bool, optional
            Disable checkpoint, by default False
        """
        if not isinstance(data, DataConfig):
            data = DataConfig(**data)

        self.data = data
        self.model = model
        self.output = output
        self.name = name
        self.trainer = trainer
        self.resume = resume
        self.pretrained = pretrained
        self.wandb = wandb
        self.disable_checkpoint = disable_checkpoint
        self.slurm = slurm
        self.matmul_precision = matmul_precision
        self.find_unused_parameters = find_unused_parameters
        self.ddp_timeout_seconds = ddp_timeout_seconds
        self.save_every_n_train_steps = save_every_n_train_steps
        self.save_top_k = save_top_k
        self.validation_only = validation_only
        self.debug = debug
        self.strict_loading = strict_loading
        self.metric_mode = metric_mode

    def run(self, config: OmegaConf) -> None:
        """Run training.
        Parameters
        ----------
        config : OmegaConf
            The configuration for the task, for bookkeeping.
        """
        # Disable some warnings
        warnings.filterwarnings(
            "ignore", ".*when logging on epoch level in distributed setting.*"
        )

        # Experiment with this during training (high or medium)
        if self.matmul_precision is not None:
            torch.set_float32_matmul_precision(self.matmul_precision)

        # Create trainer dict
        if self.trainer is None:
            self.trainer = {}

        # Flip some arguments in debug mode
        devices = self.trainer.get("devices", 1)

        if self.debug:
            self.slurm = False
            if isinstance(devices, int):
                devices = 1
            elif isinstance(devices, (list, listconfig.ListConfig)):
                devices = [devices[0]]
            self.trainer["devices"] = devices
            self.data.num_workers = 0
            if self.wandb:
                self.wandb = None

        # slurm
        if self.slurm:
            self.trainer["devices"] = int(
                os.environ.get("SLURM_NTASKS_PER_NODE", "auto")
            )
            self.trainer["num_nodes"] = int(os.environ.get("SLURM_NNODES", 1))

        # Create objects
        data_module = TrainingDataModule(self.data)
        model_module = self.model

        if self.pretrained and not self.resume:
            file_path = self.pretrained

            print(f"Loading model from {file_path}")
            model_module = type(model_module).load_from_checkpoint(
                file_path, map_location="cpu", strict=False, weights_only=False, **(model_module.hparams)
            )

        # Create checkpoint callback
        callbacks = self.trainer.get("callbacks", [])
        dirpath = f"{self.output}/{self.name}"
        os.makedirs(dirpath, exist_ok=True)

        if not self.disable_checkpoint:
            if self.slurm:
                jobid = os.environ.get("SLURM_JOB_ID", "")
                dirpath = Path(self.output) / jobid
                dirpath.mkdir(parents=True, exist_ok=True)
                print(
                    "Configuring ModelCheckpoint for SLURM to directory: ", str(dirpath)
                )
            mc = ModelCheckpoint(
                monitor="val/lddt",
                save_top_k=self.save_top_k,
                save_last=True,
                mode=self.metric_mode,
                every_n_epochs=1,
            )
            callbacks.append(mc)

        # Create wandb logger
        loggers = []
        if self.wandb:
            extra = {}
            if self.slurm:
                ckpt_path = (
                    "hpc"
                    if len([f for f in os.listdir(dirpath) if f.startswith("hpc")])
                    else None
                )
                extra = {"resume": "must"} if ckpt_path else {}

            wdb_logger = WandbLogger(
                name=self.name,
                group=self.wandb["group"],
                save_dir=dirpath,
                dir=f".tmp_wandb/{time.time()}",
                project=self.wandb["project"],
                entity=self.wandb["entity"],
                log_model=False,
                id=os.environ.get("SLURM_JOB_ID", None) if self.slurm else None,
                **extra,
            )
            loggers.append(wdb_logger)
            # Save the config to wandb
            # This fails when not on rank0, so just catching
            try:
                config_out = Path(wdb_logger.experiment.dir) / "run.yaml"
                with Path.open(config_out, "w") as f:
                    OmegaConf.save(config, f)
                wdb_logger.experiment.save(str(config_out))
            except Exception:  # noqa: BLE001, S110
                pass

        # Set up trainer
        strategy = "auto"
        if (isinstance(devices, int) and devices > 1) or (
            isinstance(devices, (list, listconfig.ListConfig)) and len(devices) > 1
        ):
            strategy = DDPStrategy(
                find_unused_parameters=self.find_unused_parameters,
                timeout=datetime.timedelta(seconds=self.ddp_timeout_seconds),
            )

        trainer = pl.Trainer(
            default_root_dir=str(dirpath),
            strategy=strategy,
            callbacks=callbacks,
            logger=loggers,
            enable_checkpointing=not self.disable_checkpoint,
            reload_dataloaders_every_n_epochs=1,
            **self.trainer,
        )

        # Run training
        if self.slurm:
            # If we are on SLURM, we need to check if we have a checkpoint for this job
            ckpt_path = (
                "hpc"
                if len([f for f in os.listdir(dirpath) if f.startswith("hpc")])
                else None
            )
            # If none check with we have a resume checkpoint passed in
            if ckpt_path is None:
                ckpt_path = self.resume
        else:
            ckpt_path = self.resume

        if not self.strict_loading:
            model_module.strict_loading = False

        if self.validation_only:
            trainer.validate(
                model_module,
                datamodule=data_module,
                ckpt_path=ckpt_path,
            )
        else:
            trainer.fit(
                model_module,
                datamodule=data_module,
                ckpt_path=ckpt_path,
            )

	import datetime
	import os
	import warnings
	from pathlib import Path
	import time
	from typing import Optional
	import pytorch_lightning as pl
	import torch
	from omegaconf import OmegaConf, listconfig
	from pytorch_lightning import LightningModule
	from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
	from pytorch_lightning.loggers import WandbLogger
	from pytorch_lightning.strategies import DDPStrategy

	from boltzgen.task.task import Task
	from boltzgen.task.train.data import DataConfig, TrainingDataModule


	class Training(Task):
	"""Training configuration."""

	def __init__(
	self,
	data: DataConfig,
	model: LightningModule,
	output: str,
	name: str,
	trainer: Optional[dict] = None,
	resume: Optional[str] = None,
	pretrained: Optional[str] = None,
	wandb: Optional[dict] = None,
	disable_checkpoint: bool = False,
	slurm: bool = False,
	matmul_precision: Optional[str] = None,
	find_unused_parameters: Optional[bool] = False,
	ddp_timeout_seconds: Optional[int] = 6400,
	save_every_n_train_steps: Optional[int] = None,
	save_top_k: Optional[int] = 1,
	validation_only: bool = False,
	debug: bool = False,
	strict_loading: bool = True,
	metric_mode: str = "max",
	) -> None:
	"""Initialize training configuration.

	Parameters
	----------
	data : DataConfig
	The data configuration.
	model : LightningModule
	The model.
	output : str
	The output directory.
	name : str
	The output subdirectory and wandb run name.
	trainer : Optional[dict], optional
	The trainer configuration, by default None.
	resume : Optional[str], optional
	The resume checkpoint, by default None
	pretrained : Optional[str], optional
	The pretrained model, by default None
	wandb : Optional[dict], optional
	The wandb configuration, by default None
	disable_checkpoint : bool, optional
	Disable checkpoint, by default False

	"""
	if not isinstance(data, DataConfig):
	data = DataConfig(**data)

	self.data = data
	self.model = model
	self.output = output
	self.name = name
	self.trainer = trainer
	self.resume = resume
	self.pretrained = pretrained
	self.wandb = wandb
	self.disable_checkpoint = disable_checkpoint
	self.slurm = slurm
	self.matmul_precision = matmul_precision
	self.find_unused_parameters = find_unused_parameters
	self.ddp_timeout_seconds = ddp_timeout_seconds
	self.save_every_n_train_steps = save_every_n_train_steps
	self.save_top_k = save_top_k
	self.validation_only = validation_only
	self.debug = debug
	self.strict_loading = strict_loading
	self.metric_mode = metric_mode

	def run(self, config: OmegaConf) -> None:
	"""Run training.

	Parameters
	----------
	config : OmegaConf
	The configuration for the task, for bookkeeping.

	"""
	# Disable some warnings
	warnings.filterwarnings(
	"ignore", ".when logging on epoch level in distributed setting."
	)

	# Experiment with this during training (high or medium)
	if self.matmul_precision is not None:
	torch.set_float32_matmul_precision(self.matmul_precision)

	# Create trainer dict
	if self.trainer is None:
	self.trainer = {}

	# Flip some arguments in debug mode
	devices = self.trainer.get("devices", 1)

	if self.debug:
	self.slurm = False
	if isinstance(devices, int):
	devices = 1
	elif isinstance(devices, (list, listconfig.ListConfig)):
	devices = [devices[0]]
	self.trainer["devices"] = devices
	self.data.num_workers = 0
	if self.wandb:
	self.wandb = None

	# slurm
	if self.slurm:
	self.trainer["devices"] = int(
	os.environ.get("SLURM_NTASKS_PER_NODE", "auto")
	)
	self.trainer["num_nodes"] = int(os.environ.get("SLURM_NNODES", 1))

	# Create objects
	data_module = TrainingDataModule(self.data)
	model_module = self.model

	if self.pretrained and not self.resume:
	file_path = self.pretrained

	print(f"Loading model from {file_path}")
	model_module = type(model_module).load_from_checkpoint(
	file_path, map_location="cpu", strict=False, weights_only=False, **(model_module.hparams)
	)

	# Create checkpoint callback
	callbacks = self.trainer.get("callbacks", [])
	dirpath = f"{self.output}/{self.name}"
	os.makedirs(dirpath, exist_ok=True)

	if not self.disable_checkpoint:
	if self.slurm:
	jobid = os.environ.get("SLURM_JOB_ID", "")
	dirpath = Path(self.output) / jobid
	dirpath.mkdir(parents=True, exist_ok=True)
	print(
	"Configuring ModelCheckpoint for SLURM to directory: ", str(dirpath)
	)
	mc = ModelCheckpoint(
	monitor="val/lddt",
	save_top_k=self.save_top_k,
	save_last=True,
	mode=self.metric_mode,
	every_n_epochs=1,
	)
	callbacks.append(mc)

	# Create wandb logger
	loggers = []
	if self.wandb:
	extra = {}
	if self.slurm:
	ckpt_path = (
	"hpc"
	if len([f for f in os.listdir(dirpath) if f.startswith("hpc")])
	else None
	)
	extra = {"resume": "must"} if ckpt_path else {}

	wdb_logger = WandbLogger(
	name=self.name,
	group=self.wandb["group"],
	save_dir=dirpath,
	dir=f".tmp_wandb/{time.time()}",
	project=self.wandb["project"],
	entity=self.wandb["entity"],
	log_model=False,
	id=os.environ.get("SLURM_JOB_ID", None) if self.slurm else None,
	**extra,
	)
	loggers.append(wdb_logger)
	# Save the config to wandb
	# This fails when not on rank0, so just catching
	try:
	config_out = Path(wdb_logger.experiment.dir) / "run.yaml"
	with Path.open(config_out, "w") as f:
	OmegaConf.save(config, f)
	wdb_logger.experiment.save(str(config_out))
	except Exception: # noqa: BLE001, S110
	pass

	# Set up trainer
	strategy = "auto"
	if (isinstance(devices, int) and devices > 1) or (
	isinstance(devices, (list, listconfig.ListConfig)) and len(devices) > 1
	):
	strategy = DDPStrategy(
	find_unused_parameters=self.find_unused_parameters,
	timeout=datetime.timedelta(seconds=self.ddp_timeout_seconds),
	)

	trainer = pl.Trainer(
	default_root_dir=str(dirpath),
	strategy=strategy,
	callbacks=callbacks,
	logger=loggers,
	enable_checkpointing=not self.disable_checkpoint,
	reload_dataloaders_every_n_epochs=1,
	**self.trainer,
	)

	# Run training
	if self.slurm:
	# If we are on SLURM, we need to check if we have a checkpoint for this job
	ckpt_path = (
	"hpc"
	if len([f for f in os.listdir(dirpath) if f.startswith("hpc")])
	else None
	)
	# If none check with we have a resume checkpoint passed in
	if ckpt_path is None:
	ckpt_path = self.resume
	else:
	ckpt_path = self.resume

	if not self.strict_loading:
	model_module.strict_loading = False

	if self.validation_only:
	trainer.validate(
	model_module,
	datamodule=data_module,
	ckpt_path=ckpt_path,
	)
	else:
	trainer.fit(
	model_module,
	datamodule=data_module,
	ckpt_path=ckpt_path,
	)