Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(azure): Preliminary support for Azure #118

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

wbrown
Copy link
Collaborator

@wbrown wbrown commented Mar 22, 2024

This PR adds Azure Blob Storage support to tensorizer. Both serialization and deserialization work, very similarly to the S3 support that already exists:

  • For serialization of models to Azure, it presently writes to a temporary file before moving on to uploading it to the Azure Blob Store. This is triggered by a GC cleanup or an explicit close().
  • For deserialization of models from Azure, it generates a SAS token to pass it to the CURLStreamFile class which does an ordinary HTTP pull.

Parameterization is done using the DefaultAzureCredential class. It picks up the credentials via various automated mechanisms. The one that was tested was EnvironmentCredential.

Usage is straightforward -- provide it with an azure:// uri in the form of azure://account/container/blob and it will take it from there.

import tensorizer.serialization as serialization
from transformers import AutoModelForCausalLM
import azure.core.exceptions

model = AutoModelForCausalLM.from_pretrained("eleutherai/gpt-neo-125m")
print("Model loaded.")
serializer = serialization.TensorSerializer(
    "azure://test/data/gpt-neo-125m",
)
try:
    serializer.write_module(model)
    serializer.close()
    print("Done serializing to Azure!")
except azure.core.exceptions.ResourceExistsError:
    print("Resource already exists.")
deserialize = serialization.TensorDeserializer(
    "azure://test/data/gpt-neo-125m",
    verify_hash=True,
)
deserialize.load_into_module(model)
print("Model deserialized from Azure!")

This PR is not complete -- test cases still need to be written, which is vastly complicated by the lack of a library similar to AWS' moto3 for mocking up interfaces.

@wbrown wbrown marked this pull request as draft March 22, 2024 21:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant