1111# limitations under the License.
1212
1313
14- import configparser
1514import argparse
16- from email import message_from_string
15+ import base64
16+ import configparser
1717import json
18+ import logging
1819import os
20+ from email import message_from_string
21+
1922import boto3
20- from botocore .config import Config
2123import yaml
22- import base64
23- import logging
24+ from botocore .config import Config
2425from retrying import retry
2526
2627COMPUTE_FLEET_SHARED_LOCATION = "/opt/parallelcluster/shared/"
2728
28- COMPUTE_FLEET_SHARED_DNA_LOCATION = COMPUTE_FLEET_SHARED_LOCATION + ' dna/'
29+ COMPUTE_FLEET_SHARED_DNA_LOCATION = COMPUTE_FLEET_SHARED_LOCATION + " dna/"
2930
30- COMPUTE_FLEET_LAUNCH_TEMPLATE_CONFIG = COMPUTE_FLEET_SHARED_LOCATION + ' launch-templates-config.json'
31+ COMPUTE_FLEET_LAUNCH_TEMPLATE_CONFIG = COMPUTE_FLEET_SHARED_LOCATION + " launch-templates-config.json"
3132
3233logger = logging .getLogger (__name__ )
3334logging .basicConfig (level = logging .INFO )
3435
36+
3537def get_compute_launch_template_ids (lt_config_file_name ):
36- """Load launch-templates-config.json which contains ID, Version number and Logical ID of all queues in Compute Fleet's Launch Template.
37- The format of launch-templates-config.json is
38+ """
39+ Load launch-templates-config.json.
40+
41+ It contains ID, Version number and Logical ID of all queues in Compute Fleet's Launch Template.
42+
43+ The format of launch-templates-config.json:
3844 {
3945 "Queues": {
4046 "queue1": {
@@ -61,24 +67,25 @@ def get_compute_launch_template_ids(lt_config_file_name):
6167 }
6268 }
6369 }
70+
6471 """
6572 lt_config = None
6673 try :
67- with open (lt_config_file_name , 'r' ) as file :
74+ with open (lt_config_file_name , "r" ) as file :
6875 lt_config = json .loads (file .read ())
6976 except Exception as err :
70- logger .warn ("Unable to read %s due to %s" , lt_config_file_name , err )
77+ logger .warning ("Unable to read %s due to %s" , lt_config_file_name , err )
7178
72- return lt_config
79+ return lt_config
7380
7481
7582def share_compute_fleet_dna (args ):
76- """Creates all dna.json for each queue in cluster."""
83+ """Create dna.json for each queue in cluster."""
7784 lt_config = get_compute_launch_template_ids (COMPUTE_FLEET_LAUNCH_TEMPLATE_CONFIG )
7885 if lt_config :
79- all_queues = lt_config .get (' Queues' )
86+ all_queues = lt_config .get (" Queues" )
8087 for _ , queues in all_queues .items ():
81- compute_resources = queues .get (' ComputeResources' )
88+ compute_resources = queues .get (" ComputeResources" )
8289 for _ , compute_res in compute_resources .items ():
8390 get_latest_dna_data (compute_res , COMPUTE_FLEET_SHARED_DNA_LOCATION , args )
8491
@@ -98,63 +105,62 @@ def parse_proxy_config():
98105@retry (stop_max_attempt_number = 5 , wait_fixed = 3000 )
99106def get_user_data (lt_id , lt_version , region_name ):
100107 """
101- Calls EC2 DescribeLaunchTemplateVersions API to get UserData from Launch Template specified.
108+ Get UserData from specified Launch Template using EC2 DescribeLaunchTemplateVersions API.
109+
102110 :param lt_id: Launch Template ID (eg: lt-12345678901234567)
103111 :param lt_version: Launch Template latest Version Number (eg: 2)
104112 :param region_name: AWS region name (eg: us-east-1)
105- :return: User_data in MIME format
113+ :return: string of user_data in MIME format
106114 """
107115 decoded_data = None
108116 try :
109117 proxy_config = parse_proxy_config ()
110118
111119 ec2_client = boto3 .client ("ec2" , region_name = region_name , config = proxy_config )
112120 response = ec2_client .describe_launch_template_versions (
113- LaunchTemplateId = lt_id ,
121+ LaunchTemplateId = lt_id ,
114122 Versions = [
115123 lt_version ,
116124 ],
117- ).get (' LaunchTemplateVersions' )
118- decoded_data = base64 .b64decode (response [0 ][' LaunchTemplateData' ][ ' UserData' ], validate = True ).decode (' utf-8' )
125+ ).get (" LaunchTemplateVersions" )
126+ decoded_data = base64 .b64decode (response [0 ][" LaunchTemplateData" ][ " UserData" ], validate = True ).decode (" utf-8" )
119127 except Exception as err :
120128 if hasattr (err , "message" ):
121129 err = err .message
122130 logger .error (
123- "Unable to get UserData for launch template %s with version %s.\n Exception: %s" ,
124- lt_id , lt_version , err
131+ "Unable to get UserData for launch template %s with version %s.\n Exception: %s" , lt_id , lt_version , err
125132 )
126133
127134 return decoded_data
128135
129136
130137def get_write_directives_section (user_data ):
131- """
132- Parses MIME formatted UserData that we get from EC2 to extract write_files section from cloud-config section.
133- """
138+ """Get write_files section from cloud-config section of MIME formatted UserData."""
134139 write_directives_section = None
135140 try :
136141 data = message_from_string (user_data )
137142 for cloud_config_section in data .walk ():
138- if cloud_config_section .get_content_type () == ' text/cloud-config' :
139- write_directives_section = yaml .safe_load (cloud_config_section ._payload ).get (' write_files' )
143+ if cloud_config_section .get_content_type () == " text/cloud-config" :
144+ write_directives_section = yaml .safe_load (cloud_config_section ._payload ).get (" write_files" )
140145 except Exception as err :
141146 logger .error ("Error occurred while parsing write_files section.\n Exception: %s" , err )
142147 return write_directives_section
143148
144149
145150def write_dna_files (write_files_section , shared_storage_loc ):
146151 """
147- Writes the dna.json in shared location after extracting it from write_files section of UserData.
152+ After extracting dna.json from write_files section of UserData, write it in shared location.
153+
148154 :param write_files_section: Entire write_files section from UserData
149155 :param shared_storage_loc: Shared Storage Location of where to write dna.json
150156 :return: None
151157 """
152158 try :
153- file_path = shared_storage_loc + "-dna.json"
159+ file_path = shared_storage_loc + "-dna.json"
154160 for data in write_files_section :
155- if data [' path' ] in [' /tmp/dna.json' ]:
156- with open (file_path ,"w " ) as file :
157- file .write (json .dumps (json .loads (data [' content' ]),indent = 4 ))
161+ if data [" path" ] in [" /tmp/dna.json" ]: # nosec B108
162+ with open (file_path , "w" , encoding = "utf-8 " ) as file :
163+ file .write (json .dumps (json .loads (data [" content" ]), indent = 4 ))
158164 except Exception as err :
159165 if hasattr (err , "message" ):
160166 err = err .message
@@ -163,16 +169,19 @@ def write_dna_files(write_files_section, shared_storage_loc):
163169
164170def get_latest_dna_data (resource , output_location , args ):
165171 """
166- Function to get latest User Data, extract relevant details and write dna.json.
172+ Get latest User Data, extract relevant details and write dna.json.
173+
167174 :param resource: Resource containing LT ID, Version and Logical id
168175 :param output_location: Shared Storage Location were we want to write dna.json
169176 :param args: Command Line arguments
170177 :rtype: None
171178 """
172- user_data = get_user_data (resource .get ('LaunchTemplate' ).get ('Id' ), resource .get ('LaunchTemplate' ).get ('Version' ), args .region )
179+ user_data = get_user_data (
180+ resource .get ("LaunchTemplate" ).get ("Id" ), resource .get ("LaunchTemplate" ).get ("Version" ), args .region
181+ )
173182 if user_data :
174183 write_directives = get_write_directives_section (user_data )
175- write_dna_files (write_directives , output_location + resource .get (' LaunchTemplate' ).get ("LogicalId" ))
184+ write_dna_files (write_directives , output_location + resource .get (" LaunchTemplate" ).get ("LogicalId" ))
176185
177186
178187def cleanup (directory_loc ):
@@ -183,7 +192,7 @@ def cleanup(directory_loc):
183192 if os .path .isfile (f_path ):
184193 os .remove (f_path )
185194 except Exception as err :
186- logger .warn ( f "Unable to delete %s due to %s" , f_path , err )
195+ logger .warning ( "Unable to delete %s due to %s" , f_path , err )
187196
188197
189198def _parse_cli_args ():
@@ -224,9 +233,11 @@ def main():
224233 except Exception as err :
225234 if hasattr (err , "message" ):
226235 err = err .message
227- logger .exception ("Encountered exception when fetching latest dna.json for ComputeFleet, exiting gracefully: %s" , err )
236+ logger .exception (
237+ "Encountered exception when fetching latest dna.json for ComputeFleet, exiting gracefully: %s" , err
238+ )
228239 raise SystemExit (0 )
229240
230241
231242if __name__ == "__main__" :
232- main ()
243+ main ()
0 commit comments