def _azure_upload_file(self, credentials, local_file, artifact_path):
"""
Uploads a file to a given Azure storage location.
The function uses a file chunking generator with 100 MB being the size limit for each chunk.
This limit is imposed by the stage_block API in azure-storage-blob.
In the case the file size is large and the upload takes longer than the validity of the
given credentials, a new set of credentials are generated and the operation continues. This
is the reason for the first nested try-except block
Finally, since the prevailing credentials could expire in the time between the last
stage_block and the commit, a second try-except block refreshes credentials if needed.
"""
try:
headers = self._extract_headers_from_credentials(credentials.headers)
uploading_block_list = list()
for chunk in yield_file_in_chunks(local_file, _AZURE_MAX_BLOCK_CHUNK_SIZE):
# Base64-encode a UUID, producing a UTF8-encoded bytestring. Then, decode
# the bytestring for compliance with Azure Blob Storage API requests
block_id = base64.b64encode(uuid.uuid4().hex.encode()).decode("utf-8")
try:
put_block(credentials.signed_uri, block_id, chunk, headers=headers)
except requests.HTTPError as e:
if e.response.status_code in [401, 403]:
_logger.info(
"Failed to authorize request, possibly due to credential expiration."
" Refreshing credentials and trying again..."
)
credential_info = self._get_write_credential_infos(
run_id=self.run_id, paths=[artifact_path]
)[0]
put_block(credential_info.signed_uri, block_id, chunk, headers=headers)
else:
raise e
uploading_block_list.append(block_id)
try:
put_block_list(credentials.signed_uri, uploading_block_list, headers=headers)
except requests.HTTPError as e:
if e.response.status_code in [401, 403]:
_logger.info(
"Failed to authorize request, possibly due to credential expiration."
" Refreshing credentials and trying again..."
)
credential_info = self._get_write_credential_infos(
run_id=self.run_id, paths=[artifact_path]
)[0]
put_block_list(
credential_info.signed_uri, uploading_block_list, headers=headers
)
else:
raise e
except Exception as err:
raise MlflowException(err)
def _azure_upload_file(self, credentials, local_file, artifact_path):
"""
Uploads a file to a given Azure storage location.
The function uses a file chunking generator with 100 MB being the size limit for each chunk.
This limit is imposed by the stage_block API in azure-storage-blob.
In the case the file size is large and the upload takes longer than the validity of the
given credentials, a new set of credentials are generated and the operation continues. This
is the reason for the first nested try-except block
Finally, since the prevailing credentials could expire in the time between the last
stage_block and the commit, a second try-except block refreshes credentials if needed.
"""
from azure.core.exceptions import ClientAuthenticationError
from azure.storage.blob import BlobClient
try:
headers = self._extract_headers_from_credentials(credentials.headers)
service = BlobClient.from_blob_url(
blob_url=credentials.signed_uri, credential=None, headers=headers
)
uploading_block_list = list()
for chunk in yield_file_in_chunks(local_file, _AZURE_MAX_BLOCK_CHUNK_SIZE):
block_id = base64.b64encode(uuid.uuid4().hex.encode())
try:
service.stage_block(block_id, chunk, headers=headers)
except ClientAuthenticationError:
_logger.warning(
"Failed to authorize request, possibly due to credential expiration."
"Refreshing credentials and trying again.."
)
credentials = self._get_write_credentials(
self.run_id, artifact_path
).credentials.signed_uri
service = BlobClient.from_blob_url(blob_url=credentials, credential=None)
service.stage_block(block_id, chunk, headers=headers)
uploading_block_list.append(block_id)
try:
service.commit_block_list(uploading_block_list, headers=headers)
except ClientAuthenticationError:
_logger.warning(
"Failed to authorize request, possibly due to credential expiration."
"Refreshing credentials and trying again.."
)
credentials = self._get_write_credentials(
self.run_id, artifact_path
).credentials.signed_uri
service = BlobClient.from_blob_url(blob_url=credentials, credential=None)
service.commit_block_list(uploading_block_list, headers=headers)
except Exception as err:
raise MlflowException(err)