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

[Tests] fix smoke test race condition on first run #4494

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 34 additions & 4 deletions sky/provision/aws/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,36 @@ def _get_role(role_name: str):
f'{role_name}{colorama.Style.RESET_ALL} from AWS.')
raise exc

def _create_instance_profile_if_not_exists(instance_profile_name: str):
try:
iam.meta.client.create_instance_profile(
InstanceProfileName=instance_profile_name)
except aws.botocore_exceptions().ClientError as exc:
if exc.response.get('Error',
{}).get('Code') == 'EntityAlreadyExists':
return
else:
utils.handle_boto_error(
exc,
f'Failed to create instance profile {colorama.Style.BRIGHT}'
f'{instance_profile_name}{colorama.Style.RESET_ALL} in AWS.'
)
raise exc

def _create_role_if_not_exists(role_name: str, policy_doc: Dict[str, Any]):
try:
iam.create_role(RoleName=role_name,
AssumeRolePolicyDocument=json.dumps(policy_doc))
except aws.botocore_exceptions().ClientError as exc:
if exc.response.get('Error',
{}).get('Code') == 'EntityAlreadyExists':
return
else:
utils.handle_boto_error(
exc, f'Failed to create role {colorama.Style.BRIGHT}'
f'{role_name}{colorama.Style.RESET_ALL} in AWS.')
raise exc

instance_profile_name = DEFAULT_SKYPILOT_INSTANCE_PROFILE
profile = _get_instance_profile(instance_profile_name)

Expand All @@ -158,8 +188,7 @@ def _get_role(role_name: str):
f'Creating new IAM instance profile {colorama.Style.BRIGHT}'
f'{instance_profile_name}{colorama.Style.RESET_ALL} for use as the '
'default.')
iam.meta.client.create_instance_profile(
InstanceProfileName=instance_profile_name)
_create_instance_profile_if_not_exists(instance_profile_name)
profile = _get_instance_profile(instance_profile_name)
time.sleep(15) # wait for propagation
assert profile is not None, 'Failed to create instance profile'
Expand All @@ -186,8 +215,9 @@ def _get_role(role_name: str):
'arn:aws:iam::aws:policy/AmazonS3FullAccess',
]

iam.create_role(RoleName=role_name,
AssumeRolePolicyDocument=json.dumps(policy_doc))
# TODO(aylei): need to check and reconcile the role permission
# in case of external modification.
_create_role_if_not_exists(role_name, policy_doc)
role = _get_role(role_name)
assert role is not None, 'Failed to create role'

Expand Down
Loading