class RslearnLightningCLI(LightningCLI):
"""LightningCLI that links data.tasks to model.tasks and supports environment variables."""
def add_arguments_to_parser(self, parser: LightningArgumentParser) -> None:
"""Link data.tasks to model.tasks and set sensible defaults.
Args:
parser: the argument parser
"""
parser.link_arguments(
"data.init_args.task", "model.init_args.task", apply_on="instantiate"
)
parser.set_defaults(
{
# Disable the sampler replacement, since the rslearn data module will set the
# sampler as needed.
"trainer.use_distributed_sampler": False,
# Default to DDP with find_unused_parameters, since it works in most settings.
"trainer.strategy": {
"class_path": "lightning.pytorch.strategies.DDPStrategy",
"init_args": {"find_unused_parameters": True},
},
}
)
# Project management option to have rslearn manage checkpoints and W&B run.
parser.add_argument(
"--management_dir",
type=str | None,
help="Enable project management, and use this directory to store checkpoints and configs. If enabled, rslearn will automatically manages checkpoint directory/loading and W&B run",
default=None,
)
parser.add_argument(
"--project_name",
type=str | None,
help="The project name (used with --management_dir)",
default=None,
)
parser.add_argument(
"--run_name",
type=str | None,
help="A unique name for this experiment (used with --management_dir)",
default=None,
)
parser.add_argument(
"--run_description",
type=str,
help="Optional description of this experiment (used with --management_dir)",
default="",
)
parser.add_argument(
"--load_checkpoint_mode",
type=str,
help="Which checkpoint to load, if any (used with --management_dir). 'none' never loads any checkpoint, 'last' loads the most recent checkpoint, and 'best' loads the best checkpoint. 'auto' will use 'last' during fit and 'best' during val/test/predict.",
default="auto",
)
parser.add_argument(
"--load_checkpoint_required",
type=str,
help="Whether to fail if the expected checkpoint based on load_checkpoint_mode does not exist (used with --management_dir). 'yes' will fail while 'no' won't. 'auto' will use 'no' during fit and 'yes' during val/test/predict.",
default="auto",
)
parser.add_argument(
"--log_mode",
type=str,
help="Whether to log to W&B (used with --management_dir). 'yes' will enable logging, 'no' will disable logging, and 'auto' will use 'yes' during fit and 'no' during val/test/predict.",
default="auto",
)
def _get_checkpoint_path(
self,
project_dir: UPath,
load_checkpoint_mode: str,
load_checkpoint_required: str,
stage: str,
) -> str | None:
"""Get path to checkpoint to load from, or None to not restore checkpoint.
Args:
project_dir: the project directory determined from the project management
directory.
load_checkpoint_mode: "none" to not load any checkpoint, "last" to load the
most recent checkpoint, "best" to load the best checkpoint. "auto" to
use "last" during fit and "best" during val/test/predict.
load_checkpoint_required: "yes" to fail if no checkpoint exists, "no" to
ignore. "auto" will use "no" during fit and "yes" during
val/test/predict.
stage: the lightning stage (fit/val/test/predict).
Returns:
the path to the checkpoint for setting c.ckpt_path, or None if no
checkpoint should be restored.
"""
# Resolve auto options if used.
if load_checkpoint_mode == "auto":
if stage == "fit":
load_checkpoint_mode = "last"
else:
load_checkpoint_mode = "best"
if load_checkpoint_required == "auto":
if stage == "fit":
load_checkpoint_required = "no"
else:
load_checkpoint_required = "yes"
if load_checkpoint_required == "yes" and load_checkpoint_mode == "none":
raise ValueError(
"load_checkpoint_required cannot be set when load_checkpoint_mode is none"
)
ckpt_path: str | None = None
if load_checkpoint_mode == "best":
# Checkpoints should be either:
# - last.ckpt
# - of the form "A=B-C=D-....ckpt" with one key being epoch=X
# So we want the one with the highest epoch, and only use last.ckpt if
# it's the only option.
# User should set save_top_k=1 so there's just one, otherwise we won't
# actually know which one is the best.
best_checkpoint = None
best_epochs = None
# Avoid error in case project_dir doesn't exist.
fnames = project_dir.iterdir() if project_dir.exists() else []
for option in fnames:
if not option.name.endswith(".ckpt"):
continue
# Try to see what epochs this checkpoint is at.
# If it is some other format, then set it 0 so we only use it if it's
# the only option.
# If it is last.ckpt then we set it -100 to only use it if there is not
# even another format like "best.ckpt".
extracted_epochs = 0
if option.name == "last.ckpt":
extracted_epochs = -100
parts = option.name.split(".ckpt")[0].split("-")
for part in parts:
kv_parts = part.split("=")
if len(kv_parts) != 2:
continue
if kv_parts[0] != "epoch":
continue
extracted_epochs = int(kv_parts[1])
if best_epochs is None or extracted_epochs > best_epochs:
best_checkpoint = option
best_epochs = extracted_epochs
if best_checkpoint is not None:
# Cache the checkpoint so we only need to download once in case we
# reuse it later.
# We only cache with --load_best since this is the only scenario where we
# expect to keep reusing the same checkpoint.
ckpt_path = get_cached_checkpoint(best_checkpoint)
elif load_checkpoint_mode == "last":
last_checkpoint_path = project_dir / "last.ckpt"
if last_checkpoint_path.exists():
ckpt_path = str(last_checkpoint_path)
else:
raise ValueError(f"unknown load_checkpoint_mode {load_checkpoint_mode}")
if load_checkpoint_required == "yes" and ckpt_path is None:
raise ValueError(
f"load_checkpoint_required is set but no checkpoint was found under {project_dir}"
)
return ckpt_path
def enable_project_management(self, management_dir: str) -> None:
"""Enable project management in the specified directory.
Sets default_root_dir on the trainer (used by ManagedBestLastCheckpoint),
configures W&B logging, and handles checkpoint loading/resuming.
Args:
management_dir: the directory to store checkpoints and W&B.
"""
subcommand = self.config.subcommand
c = self.config[subcommand]
# Project name and run name are required with project management.
if not c.project_name or not c.run_name:
raise ValueError(
"project name and run name must be set when using project management"
)
# Get project directory within the project management directory.
project_dir = UPath(management_dir) / c.project_name / c.run_name
# Set default_root_dir so ManagedBestLastCheckpoint can resolve it.
c.trainer.default_root_dir = str(project_dir)
# Disable Lightning's default ModelCheckpoint since we use our own
# (ManagedBestLastCheckpoint). Without this, Lightning adds a redundant
# ModelCheckpoint that writes to the logger's directory.
c.trainer.enable_checkpointing = False
# Configure W&B logging.
should_log = False
if c.log_mode == "yes":
should_log = True
elif c.log_mode == "auto":
should_log = subcommand == "fit"
if should_log:
if not c.trainer.logger:
c.trainer.logger = jsonargparse.Namespace(
{
"class_path": "lightning.pytorch.loggers.WandbLogger",
"init_args": jsonargparse.Namespace(),
}
)
c.trainer.logger.init_args.project = c.project_name
c.trainer.logger.init_args.name = c.run_name
if c.run_description:
c.trainer.logger.init_args.notes = c.run_description
# Add callback to save W&B run ID for resume.
if "callbacks" not in c.trainer or not c.trainer.callbacks:
c.trainer.callbacks = []
has_wandb_callback = any(
cb.class_path == "SaveWandbRunIdCallback" for cb in c.trainer.callbacks
)
if not has_wandb_callback:
config_str = json.dumps(
c.as_dict(), default=lambda _: "<not serializable>"
)
c.trainer.callbacks.append(
jsonargparse.Namespace(
{
"class_path": "SaveWandbRunIdCallback",
"init_args": jsonargparse.Namespace(
{
"project_dir": str(project_dir),
"config_str": config_str,
}
),
}
)
)
elif c.trainer.logger:
logger.warning(
"Model management is enabled and logging should be off, but the model config specifies a logger. "
"The logger should be removed from the model config, since it will not be automatically disabled."
)
# Load existing checkpoint.
checkpoint_path = self._get_checkpoint_path(
project_dir=project_dir,
load_checkpoint_mode=c.load_checkpoint_mode,
load_checkpoint_required=c.load_checkpoint_required,
stage=subcommand,
)
if checkpoint_path is not None:
logger.info(f"found checkpoint to resume from at {checkpoint_path}")
c.ckpt_path = checkpoint_path
# If we are resuming from a checkpoint for training, we also try to resume the W&B run.
if (
subcommand == "fit"
and (project_dir / WANDB_ID_FNAME).exists()
and should_log
):
with (project_dir / WANDB_ID_FNAME).open("r") as f:
wandb_id = f.read().strip()
c.trainer.logger.init_args.id = wandb_id
c.trainer.logger.init_args.resume = "must"
def before_instantiate_classes(self) -> None:
"""Called before Lightning class initialization."""
if not hasattr(self.config, "subcommand"):
logger.warning(
"Config does not have subcommand attribute, assuming we are in run=False mode"
)
subcommand = None
c = self.config
else:
subcommand = self.config.subcommand
c = self.config[subcommand]
# For predict, make sure that return_predictions is False.
# Otherwise all the predictions would be stored in memory which can lead to
# high memory consumption.
if subcommand == "predict":
c.return_predictions = False
if c.management_dir:
self.enable_project_management(c.management_dir)