s3torchconnector.lightning

Submodules

• s3torchconnector.lightning.s3_lightning_checkpoint

Classes

S3LightningCheckpoint

A checkpoint manager for S3 using the CheckpointIO interface.

Package Contents

class s3torchconnector.lightning.S3LightningCheckpoint(region: str, s3client_config: s3torchconnector._s3client.S3ClientConfig | None = None, endpoint: str | None = None)

Bases: lightning.pytorch.plugins.io.CheckpointIO

A checkpoint manager for S3 using the CheckpointIO interface.

region

save_checkpoint(checkpoint: Dict[str, Any], path: str, storage_options: Any | None = None) → None

Save model/training states as a checkpoint file through state-dump and upload to S3.

Parameters:

• checkpoint (Dict[str, Any]) – Containing model and trainer state

• path (str) – Write-target S3 uri

• storage_options – Optional parameters when saving the model/training states.

load_checkpoint(path: str, map_location: Any | None = None) → Dict[str, Any]

Load checkpoint from an S3 location when resuming or loading ckpt for test/validate/predict stages.

Parameters:

• path (str) – S3 uri to checkpoint

• map_location – A function, torch.device, string or a dict specifying how to remap storage locations.

Returns:

The loaded checkpoint

Return type:

Dict[str, Any]

Raises:

S3Exception – An error occurred accessing S3.

remove_checkpoint(path: str) → None

Remove checkpoint file from the S3 uri.

Parameters:

path (str) – S3 uri to checkpoint

Raises:

S3Exception – An error occurred accessing S3.

teardown() → None

This method is called to teardown the process.