runai_model_streamer_s3_pytorch_model_loading_quickstart.py

1import torch
2from runai_model_streamer_s3 import S3ModelStreamer
3
4# Initialize the S3 Model Streamer
5# Ensure your environment variables (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, etc.) 
6# or AWS config files are set up for authentication.
7streamer = S3ModelStreamer(
8    bucket_name="my-model-bucket",
9    region_name="us-east-1"
10)
11
12# Stream the model directly into a torch-compatible format
13# 'path/to/model.pt' is the key within your S3 bucket
14model_path = "models/llama-7b/model.pt"
15
16with streamer.stream(model_path) as stream:
17    # Use torch.load to load the state dict directly from the stream
18    state_dict = torch.load(stream, map_location="cpu")
19
20print("Model successfully streamed and loaded from S3.")