Saving Checkpoints to S3 bucket

[mhmgad]

Problem Statement

While training on AWS sagemaker, it is expensive to keep the checkpoints on the notebook. Being able to upload checkpoints directly to s3 bucket will save time.

Describe the solution you'd like

in check points, allow s3 bucket urls

from pykeen.pipeline import pipeline

result = pipeline(
    model="transe",
    dataset="nations",
    training_kwargs=dict(
        num_epochs=10,
        checkpoint_name="test_checkpoint.pt",
        checkpoint_directory="s3://bucket/checkpoints/",
    ),
)

Describe alternatives you've considered

Uploading checkpoints at the end of training for backup.

Additional information

No response

Issue Template Checks

• This is not a bug report (use a different issue template if it is)
• This is not a question (use the discussions forum instead)

[mberr]

Hi @mhmgad ,

we currently have pykeen.pipeline.PipelineResult.save_to_s3 to upload the final results to s3, but no support for checkpoints yet.

We would be happy to see a PR for this. The place to start from would be this method pykeen.training.training_loop.py:TrainingLoop._save_state.