Skip to content

Commit

Permalink
Allow to set permissions per ami in cluster.networks.config
Browse files Browse the repository at this point in the history
  • Loading branch information
sidoruka committed Dec 28, 2023
1 parent 4fc6079 commit 3d5814b
Showing 1 changed file with 29 additions and 5 deletions.
34 changes: 29 additions & 5 deletions scripts/autoscaling/aws/nodeup.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,14 @@ def get_preference(preference_name):
pipe_log('An error occured while getting preference {}, empty value is going to be used'.format(preference_name))
return None

def get_run_info(run_id):
pipe_api = PipelineAPI(api_url, None)
try:
return pipe_api.load_run_efficiently(run_id)
except:
pipe_log('An error occured while getting info for run id {}'.format(run_id))
return None

def load_cloud_config():
global __CLOUD_METADATA__
global __CLOUD_TAGS__
Expand Down Expand Up @@ -238,7 +246,7 @@ def get_security_groups(aws_region, security_groups):
def get_well_known_hosts(aws_region):
return get_cloud_config_section(aws_region, "well_known_hosts")

def get_allowed_instance_image(cloud_region, instance_type, default_image, api_token):
def get_allowed_instance_image(cloud_region, instance_type, default_image, api_token, run_id):
default_init_script = os.path.dirname(os.path.abspath(__file__)) + '/init.sh'
default_embedded_scripts = None
default_object = { "instance_mask_ami": default_image, "instance_mask": None, "init_script": default_init_script,
Expand All @@ -263,6 +271,20 @@ def get_allowed_instance_image(cloud_region, instance_type, default_image, api_t
# If something is wrong with the permissions check - do not use a restricted image
continue

docker_image_list = image_config["docker_image"] if "docker_image" in image_config else None
try:
if docker_image_list:
pipe_log('Image config with restricted docker image found ({}), checking for match with a current run'.format(docker_image_list))
run_info = get_run_info(run_id)
if not run_info or \
not 'dockerImage' in run_info or \
not run_info['dockerImage'] in docker_image_list:
continue
except:
# If something is wrong with the permissions check - do not use a restricted image
continue


instance_mask = image_config["instance_mask"]
instance_mask_ami = image_config["ami"]
init_script = image_config.get("init_script", default_object["init_script"])
Expand Down Expand Up @@ -414,7 +436,8 @@ def run_instance(api_url, api_token, bid_price, ec2, aws_region, ins_hdd, kms_en
swap_size = get_swap_size(aws_region, ins_type, is_spot)
user_data_script = get_user_data_script(api_url, api_token, aws_region, ins_type, ins_img, kube_ip,
kubeadm_token,
global_distribution_url, swap_size, pre_pull_images, node_ssh_port)
global_distribution_url, swap_size, pre_pull_images, node_ssh_port,
run_id)
if is_spot:
ins_id, ins_ip = find_spot_instance(ec2, aws_region, bid_price, run_id, pool_id, ins_img, ins_type, ins_key, ins_hdd, kms_encyr_key_id,
user_data_script, num_rep, time_rep, swap_size, kube_client, instance_additional_spec, availability_zone, security_groups, subnet, network_interface, is_dedicated, performance_network)
Expand Down Expand Up @@ -769,8 +792,9 @@ def replace_docker_images(pre_pull_images, user_data_script):

def get_user_data_script(api_url, api_token, aws_region, ins_type, ins_img, kube_ip,
kubeadm_token,
global_distribution_url, swap_size, pre_pull_images, node_ssh_port):
allowed_instance = get_allowed_instance_image(aws_region, ins_type, ins_img, api_token)
global_distribution_url, swap_size, pre_pull_images, node_ssh_port,
run_id):
allowed_instance = get_allowed_instance_image(aws_region, ins_type, ins_img, api_token, run_id)
if allowed_instance and allowed_instance["init_script"]:
init_script = open(allowed_instance["init_script"], 'r')
user_data_script = init_script.read()
Expand Down Expand Up @@ -1516,7 +1540,7 @@ def main():
instance_additional_spec = None
if not ins_img or ins_img == 'null':
# Redefine default instance image if cloud metadata has specific rules for instance type
allowed_instance = get_allowed_instance_image(aws_region, ins_type, ins_img, api_token)
allowed_instance = get_allowed_instance_image(aws_region, ins_type, ins_img, api_token, run_id)
if allowed_instance and allowed_instance["instance_mask"]:
pipe_log('Found matching rule {instance_mask}/{ami} for requested instance type {instance_type}\nImage {ami} will be used'.format(instance_mask=allowed_instance["instance_mask"],
ami=allowed_instance["instance_mask_ami"],
Expand Down

0 comments on commit 3d5814b

Please sign in to comment.