Skip to content

Commit

Permalink
format
Browse files Browse the repository at this point in the history
  • Loading branch information
@landscapepainter
landscapepainter committed Jul 3, 2024
1 parent 33fff63 commit 78c5fd3
Showing 1 changed file with 29 additions and 35 deletions.
64 changes: 29 additions & 35 deletions sky/data/storage.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@

_WAIT_FOR_STORAGE_ACCOUNT_CREATION = 60


def get_cached_enabled_storage_clouds_or_refresh(
raise_if_no_cloud_access: bool = False) -> List[str]:
# This is a temporary solution until https://github.com/skypilot-org/skypilot/issues/1943 # pylint: disable=line-too-long
Expand Down Expand Up @@ -2146,7 +2147,6 @@ def initialize(self):
# If is_sky_managed is specified, then we take no action.
self.is_sky_managed = is_new_bucket


def _get_storage_account_and_resource_group(
self) -> Tuple[str, Optional[str]]:
"""Get storage account and resource group to be used for AzureBlobStore
Expand Down Expand Up @@ -2231,9 +2231,10 @@ def _get_storage_account_and_resource_group(
region=self.region,
#user_hash=common_utils.get_user_hash()))
user_hash=f'dy{common_utils.get_user_hash()}'))
resource_group_name = (self.DEFAULT_RESOURCE_GROUP_NAME.format(
#user_hash=common_utils.get_user_hash()))
user_hash=f'dy{common_utils.get_user_hash()}'))
resource_group_name = (
self.DEFAULT_RESOURCE_GROUP_NAME.format(
#user_hash=common_utils.get_user_hash()))
user_hash=f'dy{common_utils.get_user_hash()}'))
try:
# obtains detailed information about resource group under
# the user's subscription. Used to check if the name
Expand All @@ -2256,15 +2257,16 @@ def _get_storage_account_and_resource_group(
'[bold cyan]Setting up storage account: '
f'{storage_account_name}'):
self._create_storage_account(resource_group_name,
storage_account_name)
storage_account_name)
# wait until new resource creation propagates to Azure.
time.sleep(1)

return storage_account_name, resource_group_name

def _create_storage_account(self, resource_group_name: str, storage_account_name: str) -> None:
def _create_storage_account(self, resource_group_name: str,
storage_account_name: str) -> None:
"""Creates new storage account and assign Storage Blob Data Owner role.
Args:
resource_group_name: str; Name of the resource group which the
storage account will be created under.
Expand All @@ -2277,8 +2279,7 @@ def _create_storage_account(self, resource_group_name: str, storage_account_name
"""
try:
creation_response = (
self.storage_client.storage_accounts.
begin_create(
self.storage_client.storage_accounts.begin_create(
resource_group_name, storage_account_name, {
'sku': {
'name': 'Standard_GRS'
Expand All @@ -2301,14 +2302,12 @@ def _create_storage_account(self, resource_group_name: str, storage_account_name
'Failed to create storage account '
f'{storage_account_name!r}. You may be '
'attempting to create a storage account '
'already being in use. '
f'Details: {common_utils.format_exception(error, use_bracket=True)}'
)
'already being in use. Details: '
f'{common_utils.format_exception(error, use_bracket=True)}')
# Assigning Storage Blob Data Owner role to the
# storage account.
# Reference: https://github.com/Azure/azure-sdk-for-python/issues/35573 # pylint: disable=line-too-long
authorization_client = data_utils.create_az_client(
'authorization')
authorization_client = data_utils.create_az_client('authorization')
graph_client = data_utils.create_az_client('graph')

async def get_object_id():
Expand All @@ -2317,8 +2316,7 @@ async def get_object_id():
# silencing the INFO level response log from httpx request
httpx_logger.setLevel(logging.WARNING)
user = await graph_client.users.with_url(
'https://graph.microsoft.com/v1.0/me'
).get()
'https://graph.microsoft.com/v1.0/me').get()
httpx_logger.setLevel(original_level)
object_id = str(user.additional_data['id'])
return object_id
Expand All @@ -2327,24 +2325,22 @@ async def get_object_id():
# Defintion ID of Storage Blob Data Owner role.
# Reference: https://learn.microsoft.com/en-us/azure/role-based-access-control/built-in-roles/storage#storage-blob-data-owner # pylint: disable=line-too-long
storage_blob_data_owner_role_id = 'b7e6dc6d-f1e8-4753-8033-0f276bb0955b'
role_definition_id = (
'/subscriptions'
f'/{azure.get_subscription_id()}'
'/providers/Microsoft.Authorization'
'/roleDefinitions'
f'/{storage_blob_data_owner_role_id}')

role_definition_id = ('/subscriptions'
f'/{azure.get_subscription_id()}'
'/providers/Microsoft.Authorization'
'/roleDefinitions'
f'/{storage_blob_data_owner_role_id}')

# It may take some time for the created storage account to propagate
# to Azure, we retry to assign the role for several times until
# it propagates
# to Azure, we reattempt to assign the role for several times until
# it fully propagates
role_assignment_start = time.time()
retry = 0
role_assignment_failure_error_msg = (
'Failed to assign Storage Blob Data Owner role to the '
f'storage account {storage_account_name!r}. ')
while (
time.time() - role_assignment_start
< _WAIT_FOR_STORAGE_ACCOUNT_CREATION):
while (time.time() - role_assignment_start <
_WAIT_FOR_STORAGE_ACCOUNT_CREATION):
try:
authorization_client.role_assignments.create(
scope=creation_response.id,
Expand All @@ -2363,15 +2359,14 @@ async def get_object_id():
logger.info(
'Retrying role assignment due to propagation '
'delay of the newly created storage account. '
f'Retry count: {retry}.'
)
f'Retry count: {retry}.')
time.sleep(1)
retry += 1
continue
with ux_utils.print_exception_no_traceback():
raise exceptions.StorageBucketCreateError(
f'{role_assignment_failure_error_msg}'
f'Details: {common_utils.format_exception(error, use_bracket=True)}'
f'Details: {common_utils.format_exception(e, use_bracket=True)}'
)
except azure.exceptions().HttpResponseError as e:
if 'AuthorizationFailed' in str(e):
Expand All @@ -2380,15 +2375,14 @@ async def get_object_id():
f'{role_assignment_failure_error_msg}'
'Please check to see if you have the authorization'
' "Microsoft.Authorization/roleAssignments/write" '
'to assign the role.'
)
'to assign the role to the newly created storage '
'account.')
else:
with ux_utils.print_exception_no_traceback():
raise exceptions.StorageBucketCreateError(
f'{role_assignment_failure_error_msg}'
f'Details: {common_utils.format_exception(error, use_bracket=True)}'
f'Details: {common_utils.format_exception(e, use_bracket=True)}'
)


def upload(self):
"""Uploads source to store bucket.
Expand Down

0 comments on commit 78c5fd3

Please sign in to comment.