Skip to content

Commit 20e8587

Browse files
committed
refactor(download_blob_and_return_object): no longer store downloaded model locally
utilize python temp direcory when downloading model pickle, to ensure that model artifact is deleted from local storage
1 parent f98547a commit 20e8587

File tree

1 file changed

+8
-6
lines changed

1 file changed

+8
-6
lines changed

utils.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
import logging
2+
import tempfile
3+
import time
24
from typing import Any
35

46
import dill
@@ -44,22 +46,22 @@ def load_object(file_path: str) -> Any:
4446
raise e
4547

4648

49+
@timing_decorator
4750
def download_blob_and_return_object(
4851
bucket_name: str = EnvironmentVariables().GCP_MLFLOW_MODEL_ARTIFACT_BUCKET_NAME,
4952
source_blob_name: str = "model.pkl",
50-
destination_file_name: str = "gcp_model.pkl",
5153
) -> Any:
5254
"""
5355
Downloads a blob from the bucket and return serialized object
5456
55-
:param destination_file_name:
5657
:param bucket_name: The ID of your GCS bucket
5758
:param source_blob_name: The ID of your GCS object
5859
:return: serialized object
5960
"""
60-
6161
storage_client = storage.Client()
6262
bucket = storage_client.bucket(bucket_name)
63-
blob = bucket.blob(source_blob_name)
64-
blob.download_to_filename(destination_file_name)
65-
return load_object(destination_file_name)
63+
with tempfile.NamedTemporaryFile() as file_location:
64+
blob = bucket.blob(source_blob_name)
65+
blob.download_to_filename(file_location.file.name)
66+
serialized_model = load_object(file_location.file.name)
67+
return serialized_model

0 commit comments

Comments
 (0)