Skip to content

Commit 290be6d

Browse files
aws-tianquawaws-tianquaw
committed
Fix code style
1 parent 9e9d8cc commit 290be6d

5 files changed

Lines changed: 115 additions & 142 deletions

File tree

‎template/v2/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,10 @@
11
import os
22

33
c.ServerApp.terminado_settings = {"shell_command": ["/bin/bash"]}
4-
region = os.environ.get('AWS_REGION')
5-
csp_rule = os.environ.get('JUPYTERSERVER_CSP_RULE')
4+
region = os.environ.get("AWS_REGION")
5+
csp_rule = os.environ.get("JUPYTERSERVER_CSP_RULE")
66

7-
c.ServerApp.tornado_settings = {
8-
"compress_response": True,
9-
'headers': {
10-
'Content-Security-Policy': csp_rule
11-
}
12-
}
7+
c.ServerApp.tornado_settings = {"compress_response": True, "headers": {"Content-Security-Policy": csp_rule}}
138

149
# Do not delete files to trash. Instead, permanently delete files.
1510
c.FileContentsManager.delete_to_trash = False

‎template/v2/dirs/etc/sagemaker-ui/workflows/sample_dag.py‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,18 @@
11
from airflow.decorators import dag
2-
from airflow.utils.dates import days_ago
32
from airflow.operators.bash import BashOperator
43

5-
64
default_args = {
7-
'owner': 'airflow',
5+
"owner": "airflow",
86
}
97

10-
@dag(default_args=default_args, tags=['sample'])
8+
9+
@dag(default_args=default_args, tags=["sample"])
1110
def sample_dag():
1211
def sample_task():
13-
_task = BashOperator(
14-
task_id="hello_world",
15-
bash_command="echo 'hello world!'"
16-
)
12+
_task = BashOperator(task_id="hello_world", bash_command="echo 'hello world!'")
1713
return _task
1814

1915
task = sample_task()
2016

17+
2118
sample_dag = sample_dag()

‎template/v2/dirs/etc/sagemaker-ui/workflows/workflow_client.py‎

Lines changed: 26 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,23 @@
11
import argparse
2-
import boto3
2+
from datetime import datetime, timezone
33
from typing import Optional
4+
5+
import boto3
46
import requests
5-
from datetime import datetime, timezone
67

78
JUPYTERLAB_URL = "http://default:8888/jupyterlab/default/"
89
WORKFLOWS_API_ENDPOINT = "api/sagemaker/workflows"
910
TIMESTAMP_FORMAT = "%Y-%m-%d %H:%M:%S.%f%z"
1011

12+
1113
def _validate_response(function_name: str, response: requests.Response):
1214
if response.status_code == 200:
1315
return response
1416
else:
15-
raise RuntimeError(
16-
f"{function_name} returned {response.status_code}: {str(response.content)}"
17-
)
17+
raise RuntimeError(f"{function_name} returned {response.status_code}: {str(response.content)}")
1818

1919

20-
def update_local_runner_status(
21-
session: requests.Session, status: str, detailed_status: Optional[str] = None, **kwargs
22-
):
20+
def update_local_runner_status(session: requests.Session, status: str, detailed_status: Optional[str] = None, **kwargs):
2321
response = session.post(
2422
url=JUPYTERLAB_URL + WORKFLOWS_API_ENDPOINT + "/update-local-runner-status",
2523
headers={"X-Xsrftoken": session.cookies.get_dict()["_xsrf"]},
@@ -49,46 +47,53 @@ def stop_local_runner(session: requests.Session, **kwargs):
4947
)
5048
return _validate_response("StopLocalRunner", response)
5149

52-
def check_blueprint(region: str, domain_id: str, endpoint:str, **kwargs):
50+
51+
def check_blueprint(region: str, domain_id: str, endpoint: str, **kwargs):
5352
DZ_CLIENT = boto3.client("datazone")
5453
# add correct endpoint for gamma env
5554
if endpoint != "":
5655
DZ_CLIENT = boto3.client("datazone", endpoint_url=endpoint)
5756
try:
58-
blueprint_id = DZ_CLIENT.list_environment_blueprints(managed=True, domainIdentifier=domain_id, name='Workflows')['items'][0]['id']
59-
blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration(domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id)
60-
enabled_regions = blueprint_config['enabledRegions']
57+
blueprint_id = DZ_CLIENT.list_environment_blueprints(
58+
managed=True, domainIdentifier=domain_id, name="Workflows"
59+
)["items"][0]["id"]
60+
blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration(
61+
domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id
62+
)
63+
enabled_regions = blueprint_config["enabledRegions"]
6164
print(str(region in enabled_regions))
6265
except:
6366
print("False")
6467

68+
6569
COMMAND_REGISTRY = {
6670
"update-local-runner-status": update_local_runner_status,
6771
"start-local-runner": start_local_runner,
6872
"stop-local-runner": stop_local_runner,
69-
"check-blueprint": check_blueprint
73+
"check-blueprint": check_blueprint,
7074
}
7175

76+
7277
def main():
7378
parser = argparse.ArgumentParser(description="Workflow local runner client")
7479
subparsers = parser.add_subparsers(dest="command", help="Available commands")
7580

76-
update_status_parser = subparsers.add_parser(
77-
"update-local-runner-status", help="Update status of local runner"
78-
)
81+
update_status_parser = subparsers.add_parser("update-local-runner-status", help="Update status of local runner")
7982
update_status_parser.add_argument("--status", type=str, required=True, help="Status to update")
80-
update_status_parser.add_argument(
81-
"--detailed-status", type=str, required=False, help="Detailed status text"
82-
)
83+
update_status_parser.add_argument("--detailed-status", type=str, required=False, help="Detailed status text")
8384

8485
start_parser = subparsers.add_parser("start-local-runner", help="Start local runner")
8586

8687
stop_parser = subparsers.add_parser("stop-local-runner", help="Stop local runner")
8788

8889
check_blueprint_parser = subparsers.add_parser("check-blueprint", help="Check Workflows blueprint")
89-
check_blueprint_parser.add_argument("--domain-id", type=str, required=True, help="Datazone Domain ID for blueprint check")
90+
check_blueprint_parser.add_argument(
91+
"--domain-id", type=str, required=True, help="Datazone Domain ID for blueprint check"
92+
)
9093
check_blueprint_parser.add_argument("--region", type=str, required=True, help="Datazone Domain region")
91-
check_blueprint_parser.add_argument("--endpoint", type=str, required=True, help="Datazone endpoint for blueprint check")
94+
check_blueprint_parser.add_argument(
95+
"--endpoint", type=str, required=True, help="Datazone endpoint for blueprint check"
96+
)
9297

9398
args = parser.parse_args()
9499

‎template/v2/dirs/etc/sagemaker/sm_pysdk_default_config.py‎

Lines changed: 73 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -1,152 +1,122 @@
1-
import yaml
2-
import os
3-
import sys
41
import json
2+
import os
53
import re
64

5+
import yaml
6+
7+
78
def generate_intelligent_default_config(metadata: str) -> dict:
89
config = {
9-
'SchemaVersion': '1.0',
10-
'SageMaker': {
11-
'PythonSDK': {
12-
'Modules': {
13-
'Session': {
14-
'DefaultS3Bucket': metadata['S3Bucket'],
15-
'DefaultS3ObjectKeyPrefix': metadata['S3ObjectKeyPrefix']
10+
"SchemaVersion": "1.0",
11+
"SageMaker": {
12+
"PythonSDK": {
13+
"Modules": {
14+
"Session": {
15+
"DefaultS3Bucket": metadata["S3Bucket"],
16+
"DefaultS3ObjectKeyPrefix": metadata["S3ObjectKeyPrefix"],
1617
},
17-
'RemoteFunction': {
18-
'IncludeLocalWorkDir': True,
19-
'VpcConfig': {
20-
'SecurityGroupIds': metadata['SecurityGroupIds'],
21-
'Subnets': metadata['Subnets']
22-
}
18+
"RemoteFunction": {
19+
"IncludeLocalWorkDir": True,
20+
"VpcConfig": {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]},
2321
},
24-
'NotebookJob': {
25-
'RoleArn': metadata['UserRoleArn'],
26-
'S3RootUri': f"s3://{metadata['S3Bucket']}/{metadata['S3ObjectKeyPrefix']}",
27-
'VpcConfig': {
28-
'SecurityGroupIds': metadata['SecurityGroupIds'],
29-
'Subnets': metadata['Subnets']
30-
}
22+
"NotebookJob": {
23+
"RoleArn": metadata["UserRoleArn"],
24+
"S3RootUri": f"s3://{metadata['S3Bucket']}/{metadata['S3ObjectKeyPrefix']}",
25+
"VpcConfig": {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]},
3126
},
32-
'Serve': {
33-
'S3ModelDataUri': f"s3://{metadata['S3Bucket']}/{metadata['S3ObjectKeyPrefix']}"
34-
}
27+
"Serve": {"S3ModelDataUri": f"s3://{metadata['S3Bucket']}/{metadata['S3ObjectKeyPrefix']}"},
3528
}
3629
},
37-
'MonitoringSchedule': {
38-
'MonitoringScheduleConfig': {
39-
'MonitoringJobDefinition': {
40-
'NetworkConfig': {
41-
'VpcConfig': {
42-
'SecurityGroupIds': metadata['SecurityGroupIds'],
43-
'Subnets': metadata['Subnets']
30+
"MonitoringSchedule": {
31+
"MonitoringScheduleConfig": {
32+
"MonitoringJobDefinition": {
33+
"NetworkConfig": {
34+
"VpcConfig": {
35+
"SecurityGroupIds": metadata["SecurityGroupIds"],
36+
"Subnets": metadata["Subnets"],
4437
}
4538
}
4639
}
4740
}
4841
},
49-
'AutoMLJob': {
50-
'AutoMLJobConfig': {
51-
'SecurityConfig': {
52-
'VpcConfig': {
53-
'SecurityGroupIds': metadata['SecurityGroupIds'],
54-
'Subnets': metadata['Subnets']
55-
}
42+
"AutoMLJob": {
43+
"AutoMLJobConfig": {
44+
"SecurityConfig": {
45+
"VpcConfig": {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]}
5646
}
5747
}
5848
},
59-
'AutoMLJobV2': {
60-
'SecurityConfig': {
61-
'VpcConfig': {
62-
'SecurityGroupIds': metadata['SecurityGroupIds'],
63-
'Subnets': metadata['Subnets']
64-
}
49+
"AutoMLJobV2": {
50+
"SecurityConfig": {
51+
"VpcConfig": {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]}
6552
}
6653
},
67-
'CompilationJob': {
68-
'VpcConfig': {
69-
'SecurityGroupIds': metadata['SecurityGroupIds'],
70-
'Subnets': metadata['Subnets']
71-
}
54+
"CompilationJob": {
55+
"VpcConfig": {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]}
7256
},
73-
'Pipeline': {
74-
'RoleArn': metadata['UserRoleArn']
57+
"Pipeline": {"RoleArn": metadata["UserRoleArn"]},
58+
"Model": {
59+
"VpcConfig": {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]},
60+
"ExecutionRoleArn": metadata["UserRoleArn"],
7561
},
76-
'Model': {
77-
'VpcConfig': {
78-
'SecurityGroupIds': metadata['SecurityGroupIds'],
79-
'Subnets': metadata['Subnets']
62+
"ModelPackage": {"ValidationSpecification": {"ValidationRole": metadata["UserRoleArn"]}},
63+
"ProcessingJob": {
64+
"NetworkConfig": {
65+
"VpcConfig": {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]}
8066
},
81-
'ExecutionRoleArn': metadata['UserRoleArn']
82-
},
83-
'ModelPackage': {
84-
'ValidationSpecification': {
85-
'ValidationRole': metadata['UserRoleArn']
86-
}
67+
"RoleArn": metadata["UserRoleArn"],
8768
},
88-
'ProcessingJob': {
89-
'NetworkConfig': {
90-
'VpcConfig': {
91-
'SecurityGroupIds': metadata['SecurityGroupIds'],
92-
'Subnets': metadata['Subnets']
93-
}
94-
},
95-
'RoleArn': metadata['UserRoleArn']
69+
"TrainingJob": {
70+
"RoleArn": metadata["UserRoleArn"],
71+
"VpcConfig": {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]},
9672
},
97-
'TrainingJob': {
98-
'RoleArn': metadata['UserRoleArn'],
99-
'VpcConfig': {
100-
'SecurityGroupIds': metadata['SecurityGroupIds'],
101-
'Subnets': metadata['Subnets']
102-
}
103-
}
104-
}
73+
},
10574
}
10675
return config
10776

108-
if __name__=='__main__':
109-
try:
77+
78+
if __name__ == "__main__":
79+
try:
11080
config = {}
111-
resource_metadata = '/opt/ml/metadata/resource-metadata.json'
81+
resource_metadata = "/opt/ml/metadata/resource-metadata.json"
11282

113-
PROJECT_S3_PATH = 'ProjectS3Path'
114-
SECURITY_GROUP = 'SecurityGroup'
115-
PRIVATE_SUBNETS = 'PrivateSubnets'
116-
META_DATA = 'AdditionalMetadata'
117-
EXECUTION_ROLE_ARN = 'ExecutionRoleArn'
118-
CONFIG_FILE_NAME = 'config.yaml'
119-
CONFIG_DIR = '/etc/xdg/sagemaker/'
83+
PROJECT_S3_PATH = "ProjectS3Path"
84+
SECURITY_GROUP = "SecurityGroup"
85+
PRIVATE_SUBNETS = "PrivateSubnets"
86+
META_DATA = "AdditionalMetadata"
87+
EXECUTION_ROLE_ARN = "ExecutionRoleArn"
88+
CONFIG_FILE_NAME = "config.yaml"
89+
CONFIG_DIR = "/etc/xdg/sagemaker/"
12090

12191
if os.path.exists(resource_metadata):
122-
with open(resource_metadata, 'r') as file:
92+
with open(resource_metadata, "r") as file:
12393
data = json.load(file)
12494

125-
s3_path = data[META_DATA].get(PROJECT_S3_PATH, '')
95+
s3_path = data[META_DATA].get(PROJECT_S3_PATH, "")
12696
metadata = {
12797
# user provided bucket
128-
'S3Bucket': re.search(r"s3://([^/]+)/", s3_path).group(1),
98+
"S3Bucket": re.search(r"s3://([^/]+)/", s3_path).group(1),
12999
# ${datazoneEnvironmentDomainId}/${datazoneEnvironmentProjectId}/${datazoneScopeName}/
130-
'S3ObjectKeyPrefix': s3_path.split("//")[1].split("/", 1)[1],
100+
"S3ObjectKeyPrefix": s3_path.split("//")[1].split("/", 1)[1],
131101
# TODO: Is this a billing concern if set default
132102
# 'InstanceType': 'ml.m5.xlarge',
133-
'SecurityGroupIds': data[META_DATA].get(SECURITY_GROUP, '').split(','),
134-
'Subnets': data[META_DATA].get(PRIVATE_SUBNETS, '').split(','),
135-
'UserRoleArn': data[EXECUTION_ROLE_ARN]
103+
"SecurityGroupIds": data[META_DATA].get(SECURITY_GROUP, "").split(","),
104+
"Subnets": data[META_DATA].get(PRIVATE_SUBNETS, "").split(","),
105+
"UserRoleArn": data[EXECUTION_ROLE_ARN],
136106
}
137-
107+
138108
# Not create config file when invalid value exists in metadata
139-
empty_values = [key for key, value in metadata.items() if value == "" or value == ['']]
109+
empty_values = [key for key, value in metadata.items() if value == "" or value == [""]]
140110
if empty_values:
141111
raise AttributeError(f"There are empty values in the metadata: {empty_values}")
142112

143113
config = generate_intelligent_default_config(metadata)
144114
else:
145-
raise FileNotFoundError('No resource-metadata.json exists on host!')
146-
115+
raise FileNotFoundError("No resource-metadata.json exists on host!")
116+
147117
# Write the config YAML file to default location of the admin config file
148-
with open(os.path.join(CONFIG_DIR, CONFIG_FILE_NAME), 'w') as f:
118+
with open(os.path.join(CONFIG_DIR, CONFIG_FILE_NAME), "w") as f:
149119
yaml.dump(config, f, default_flow_style=False, sort_keys=False)
150-
120+
151121
except Exception as e:
152-
print(f"Error: {e}, SageMaker PySDK intelligent config file is not valid!")
122+
print(f"Error: {e}, SageMaker PySDK intelligent config file is not valid!")

‎test/test_dockerfile_based_harness.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,10 @@
7373
("sagemaker_studio_cli.test.Dockerfile", ["sagemaker_studio_cli"]),
7474
("sm_spark_cli.test.Dockerfile", []),
7575
("sagemaker_studio_dataengineering_sessions.test.Dockerfile", ["sagemaker-studio-dataengineering-sessions"]),
76-
("sagemaker_studio_dataengineering_extensions.test.Dockerfile", ["sagemaker-studio-dataengineering-extensions"]),
76+
(
77+
"sagemaker_studio_dataengineering_extensions.test.Dockerfile",
78+
["sagemaker-studio-dataengineering-extensions"],
79+
),
7780
("sagemaker_studio.integ.Dockerfile", ["sagemaker_studio"]),
7881
],
7982
)
@@ -149,7 +152,10 @@ def test_dockerfiles_for_cpu(
149152
("sagemaker_studio_cli.test.Dockerfile", ["sagemaker_studio_cli"]),
150153
("sm_spark_cli.test.Dockerfile", []),
151154
("sagemaker_studio_dataengineering_sessions.test.Dockerfile", ["sagemaker-studio-dataengineering-sessions"]),
152-
("sagemaker_studio_dataengineering_extensions.test.Dockerfile", ["sagemaker-studio-dataengineering-extensions"]),
155+
(
156+
"sagemaker_studio_dataengineering_extensions.test.Dockerfile",
157+
["sagemaker-studio-dataengineering-extensions"],
158+
),
153159
("sagemaker_studio.integ.Dockerfile", ["sagemaker_studio"]),
154160
],
155161
)

0 commit comments

Comments
 (0)