Skip to content

Commit d0d12a8

Browse files
hcho3hcho3
committed
Update lambda functions to latest runtime; reformat
1 parent 9e81b72 commit d0d12a8

File tree

5 files changed

+146
-91
lines changed

5 files changed

+146
-91
lines changed

‎cost_watcher/lambda_function.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
"p4d.24xlarge",
2929
"g4ad.xlarge",
3030
"c5a.2xlarge",
31-
"c6g.medium"
31+
"c6g.medium",
3232
]
3333
no_launch_policy_arn = "arn:aws:iam::492475357299:policy/EC2AccessNoRunInstances"
3434

‎cost_watcher/local_runner.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
1-
import lambda_function
21
import logging
32

3+
import lambda_function
4+
45
logging.basicConfig(level=logging.INFO)
56
logging.info(lambda_function.lambda_handler(None, None))

‎deploy_scripts/deploy_lambda.py

+52-23
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,62 @@
1+
import argparse
2+
import glob
13
import os
24
import shutil
3-
import glob
4-
import argparse
55

66
import boto3
77

8+
89
def deploy_lambda(lambda_function, code_dir):
9-
os.makedirs('codepkg')
10-
for p in glob.glob(os.path.join(code_dir, '*.py')):
11-
print(f'Packaging {p}')
12-
shutil.copy(p, 'codepkg')
13-
print(f'Packaging metadata.ini')
14-
shutil.copy('metadata.ini', 'codepkg')
15-
shutil.make_archive('codepkg', format='zip', root_dir='codepkg', base_dir='.')
16-
shutil.rmtree('codepkg')
17-
print(f'Deploying to Lambda function {lambda_function}...')
18-
client = boto3.client('lambda', region_name='us-west-2')
19-
with open('codepkg.zip', 'rb') as f:
20-
r = client.update_function_code(
10+
# Package code
11+
os.makedirs("codepkg")
12+
for p in glob.glob(os.path.join(code_dir, "*.py")):
13+
print(f"Packaging {p}")
14+
shutil.copy(p, "codepkg")
15+
print(f"Packaging metadata.ini")
16+
shutil.copy("metadata.ini", "codepkg")
17+
shutil.make_archive("codepkg", format="zip", root_dir="codepkg", base_dir=".")
18+
shutil.rmtree("codepkg")
19+
20+
# Re-create lambda and deploy code package
21+
client = boto3.client("lambda", region_name="us-west-2")
22+
print(f"Re-creating Lambda function {lambda_function}...")
23+
try:
24+
client.get_function(FunctionName=lambda_function)
25+
client.delete_function(FunctionName=lambda_function)
26+
except client.exceptions.ResourceNotFoundException:
27+
pass
28+
print(f"Deploying to Lambda function {lambda_function}...")
29+
with open("codepkg.zip", "rb") as f:
30+
client.create_function(
2131
FunctionName=lambda_function,
22-
ZipFile=f.read(),
23-
Publish=False)
24-
print(f'Finished deploying to Lambda function {lambda_function}.')
32+
Runtime="python3.12",
33+
Role="arn:aws:iam::492475357299:role/XGBoost-CI-Lambda",
34+
Handler="lambda_function.lambda_handler",
35+
Code={"ZipFile": f.read()},
36+
Timeout=300,
37+
MemorySize=128,
38+
Publish=False,
39+
PackageType="Zip",
40+
Architectures=["x86_64"],
41+
)
42+
waiter = client.get_waiter("function_exists")
43+
waiter.wait(FunctionName=lambda_function)
44+
print(f"Finished deploying to Lambda function {lambda_function}.")
45+
2546

26-
if __name__ == '__main__':
27-
parser = argparse.ArgumentParser(description='Script to deploy a Lambda function')
28-
parser.add_argument('--lambda-function', type=str, help='Name of the AWS Lambda function',
29-
required=True)
30-
parser.add_argument('--code-dir', type=str,
31-
help='Directory containing Python code for the Lambda function', required=True)
47+
if __name__ == "__main__":
48+
parser = argparse.ArgumentParser(description="Script to deploy a Lambda function")
49+
parser.add_argument(
50+
"--lambda-function",
51+
type=str,
52+
help="Name of the AWS Lambda function",
53+
required=True,
54+
)
55+
parser.add_argument(
56+
"--code-dir",
57+
type=str,
58+
help="Directory containing Python code for the Lambda function",
59+
required=True,
60+
)
3261
args = parser.parse_args()
3362
deploy_lambda(args.lambda_function, args.code_dir)

‎ec2_monitor/lambda_function.py

+70-49
Original file line numberDiff line numberDiff line change
@@ -1,79 +1,100 @@
1+
import base64
12
import gzip
23
import json
3-
import base64
44
import logging
5-
from datetime import datetime
65
import re
6+
from datetime import datetime
77
from typing import Any
8+
89
import boto3
910

1011
# Set up logging
1112
logger = logging.getLogger(__name__)
1213
logger.setLevel(logging.DEBUG)
1314

14-
ec2_resource = boto3.resource('ec2', region_name='us-west-2')
15-
table = boto3.resource('dynamodb', region_name='us-west-2').Table('XGBoostCIWorkerProvisionRecord')
15+
ec2_resource = boto3.resource("ec2", region_name="us-west-2")
16+
table = boto3.resource("dynamodb", region_name="us-west-2").Table(
17+
"XGBoostCIWorkerProvisionRecord"
18+
)
19+
1620

1721
def get_os_of_ami(image_id: str) -> str:
1822
image = ec2_resource.Image(image_id)
1923
platform_details = image.platform_details
20-
assert platform_details in ['Linux/UNIX', 'Windows']
21-
if platform_details == 'Linux/UNIX':
22-
return 'Linux'
24+
assert platform_details in ["Linux/UNIX", "Windows"]
25+
if platform_details == "Linux/UNIX":
26+
return "Linux"
2327
return platform_details
2428

29+
2530
def lambda_handler(event: Any, context: Any):
26-
cw_data = event['awslogs']['data']
31+
cw_data = event["awslogs"]["data"]
2732
compressed_payload = base64.b64decode(cw_data)
2833
uncompressed_payload = gzip.decompress(compressed_payload)
2934
payload = json.loads(uncompressed_payload)
30-
log_events = payload['logEvents']
35+
log_events = payload["logEvents"]
3136
for log_event in log_events:
32-
message = json.loads(log_event['message'])
33-
if ('eventType' not in message):
34-
logger.debug(f'Message not well-formed: {message}')
37+
message = json.loads(log_event["message"])
38+
if "eventType" not in message:
39+
logger.debug(f"Message not well-formed: {message}")
3540
continue
36-
if message['eventType'] != 'AwsApiCall':
41+
if message["eventType"] != "AwsApiCall":
3742
# Skip events that are not API calls, such as Insights
3843
continue
39-
if ('eventName' not in message) or ('eventTime' not in message):
40-
logger.debug(f'Message not well-formed: {message}')
44+
if ("eventName" not in message) or ("eventTime" not in message):
45+
logger.debug(f"Message not well-formed: {message}")
4146
continue
42-
event_name = message['eventName']
43-
event_time = message['eventTime']
44-
if event_name == 'RunInstances':
45-
if (('responseElements' not in message)
46-
or (not message['responseElements'])
47-
or ('instancesSet' not in message['responseElements'])
48-
or (not message['responseElements']['instancesSet'])
49-
or ('items' not in message['responseElements']['instancesSet'])):
47+
event_name = message["eventName"]
48+
event_time = message["eventTime"]
49+
if event_name == "RunInstances":
50+
if (
51+
("responseElements" not in message)
52+
or (not message["responseElements"])
53+
or ("instancesSet" not in message["responseElements"])
54+
or (not message["responseElements"]["instancesSet"])
55+
or ("items" not in message["responseElements"]["instancesSet"])
56+
):
5057
# RunInstance that did not succeed
5158
continue
52-
for ordinal, ec2 in enumerate(message['responseElements']['instancesSet']['items']):
53-
ec2_id = ec2['instanceId']
54-
ec2_type = ec2['instanceType']
55-
ec2_os = get_os_of_ami(ec2['imageId'])
56-
logger.info(f'RunInstances, InstanceID = {ec2_id} @ {event_time}')
57-
table.put_item(Item={
58-
'Date': event_time.split(sep='T', maxsplit=1)[0],
59-
'Timestamp-Ordinal': f'{event_time}#{ordinal}',
60-
'EventName': 'RunInstances',
61-
'InstanceID': ec2_id,
62-
'InstanceType': ec2_type,
63-
'InstanceOS': ec2_os})
64-
elif event_name == 'TerminateInstances':
65-
if (('responseElements' not in message)
66-
or (not message['responseElements'])
67-
or ('instancesSet' not in message['responseElements'])
68-
or (not message['responseElements']['instancesSet'])
69-
or ('items' not in message['responseElements']['instancesSet'])):
59+
for ordinal, ec2 in enumerate(
60+
message["responseElements"]["instancesSet"]["items"]
61+
):
62+
ec2_id = ec2["instanceId"]
63+
ec2_type = ec2["instanceType"]
64+
ec2_os = get_os_of_ami(ec2["imageId"])
65+
logger.info(f"RunInstances, InstanceID = {ec2_id} @ {event_time}")
66+
table.put_item(
67+
Item={
68+
"Date": event_time.split(sep="T", maxsplit=1)[0],
69+
"Timestamp-Ordinal": f"{event_time}#{ordinal}",
70+
"EventName": "RunInstances",
71+
"InstanceID": ec2_id,
72+
"InstanceType": ec2_type,
73+
"InstanceOS": ec2_os,
74+
}
75+
)
76+
elif event_name == "TerminateInstances":
77+
if (
78+
("responseElements" not in message)
79+
or (not message["responseElements"])
80+
or ("instancesSet" not in message["responseElements"])
81+
or (not message["responseElements"]["instancesSet"])
82+
or ("items" not in message["responseElements"]["instancesSet"])
83+
):
7084
# TerminateInstances that did not succeed
7185
continue
72-
for ordinal, ec2 in enumerate(message['responseElements']['instancesSet']['items']):
73-
ec2_id = ec2['instanceId']
74-
logger.info(f'TerminateInstances, InstanceID = {ec2_id} @ {event_time}')
75-
table.put_item(Item={
76-
'Date': event_time.split(sep='T', maxsplit=1)[0],
77-
'Timestamp-Ordinal': f'{event_time}#{ordinal}',
78-
'EventName': 'TerminateInstances',
79-
'InstanceID': ec2_id})
86+
for ordinal, ec2 in enumerate(
87+
message["responseElements"]["instancesSet"]["items"]
88+
):
89+
ec2_id = ec2["instanceId"]
90+
logger.info(
91+
f"TerminateInstances, InstanceID = {ec2_id} @ {event_time}"
92+
)
93+
table.put_item(
94+
Item={
95+
"Date": event_time.split(sep="T", maxsplit=1)[0],
96+
"Timestamp-Ordinal": f"{event_time}#{ordinal}",
97+
"EventName": "TerminateInstances",
98+
"InstanceID": ec2_id,
99+
}
100+
)

‎s3_monitor/lambda_function.py

+21-17
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,44 @@
1+
import base64
12
import gzip
23
import json
3-
import base64
44
import logging
5-
from datetime import datetime
65
import re
6+
from datetime import datetime
77
from typing import Any
8+
89
import boto3
910

1011
# Set up logging
1112
logger = logging.getLogger(__name__)
1213
logger.setLevel(logging.DEBUG)
1314

15+
1416
def lambda_handler(event: Any, context: Any):
1517
"""Hanlder to be called by AWS Lambda"""
16-
cw_client = boto3.client('cloudwatch', region_name='us-west-2')
17-
18-
cw_data = event['awslogs']['data']
18+
cw_client = boto3.client("cloudwatch", region_name="us-west-2")
19+
20+
cw_data = event["awslogs"]["data"]
1921
compressed_payload = base64.b64decode(cw_data)
2022
uncompressed_payload = gzip.decompress(compressed_payload)
2123
payload = json.loads(uncompressed_payload)
22-
log_events = payload['logEvents']
24+
log_events = payload["logEvents"]
2325
for log_event in log_events:
24-
message = json.loads(log_event['message'])
25-
if message['eventName'] == 'GetObject':
26-
nbyte = message['additionalEventData']['bytesTransferredOut']
27-
timestamp = datetime.fromisoformat(re.sub(r'Z$', '+00:00', message['eventTime']))
28-
logger.info(f'GetObject with {nbyte} bytes')
26+
message = json.loads(log_event["message"])
27+
if message["eventName"] == "GetObject":
28+
nbyte = message["additionalEventData"]["bytesTransferredOut"]
29+
timestamp = datetime.fromisoformat(
30+
re.sub(r"Z$", "+00:00", message["eventTime"])
31+
)
32+
logger.info(f"GetObject with {nbyte} bytes")
2933
cw_client.put_metric_data(
3034
MetricData=[
3135
{
32-
'MetricName': 'bytesTransferredOut',
33-
'Dimensions': [],
34-
'Unit': 'None',
35-
'Value': nbyte,
36-
'Timestamp': timestamp
36+
"MetricName": "bytesTransferredOut",
37+
"Dimensions": [],
38+
"Unit": "None",
39+
"Value": nbyte,
40+
"Timestamp": timestamp,
3741
}
3842
],
39-
Namespace='XGBoostCICostWatcher'
43+
Namespace="XGBoostCICostWatcher",
4044
)

0 commit comments

Comments
 (0)