1
- import yaml
2
- import os
3
- import sys
4
1
import json
2
+ import os
5
3
import re
6
4
5
+ import yaml
6
+
7
+
7
8
def generate_intelligent_default_config (metadata : str ) -> dict :
8
9
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" ],
16
17
},
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" ]},
23
21
},
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" ]},
31
26
},
32
- 'Serve' : {
33
- 'S3ModelDataUri' : f"s3://{ metadata ['S3Bucket' ]} /{ metadata ['S3ObjectKeyPrefix' ]} "
34
- }
27
+ "Serve" : {"S3ModelDataUri" : f"s3://{ metadata ['S3Bucket' ]} /{ metadata ['S3ObjectKeyPrefix' ]} " },
35
28
}
36
29
},
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" ],
44
37
}
45
38
}
46
39
}
47
40
}
48
41
},
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" ]}
56
46
}
57
47
}
58
48
},
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" ]}
65
52
}
66
53
},
67
- 'CompilationJob' : {
68
- 'VpcConfig' : {
69
- 'SecurityGroupIds' : metadata ['SecurityGroupIds' ],
70
- 'Subnets' : metadata ['Subnets' ]
71
- }
54
+ "CompilationJob" : {
55
+ "VpcConfig" : {"SecurityGroupIds" : metadata ["SecurityGroupIds" ], "Subnets" : metadata ["Subnets" ]}
72
56
},
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" ],
75
61
},
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" ]}
80
66
},
81
- 'ExecutionRoleArn' : metadata ['UserRoleArn' ]
82
- },
83
- 'ModelPackage' : {
84
- 'ValidationSpecification' : {
85
- 'ValidationRole' : metadata ['UserRoleArn' ]
86
- }
67
+ "RoleArn" : metadata ["UserRoleArn" ],
87
68
},
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" ]},
96
72
},
97
- 'TrainingJob' : {
98
- 'RoleArn' : metadata ['UserRoleArn' ],
99
- 'VpcConfig' : {
100
- 'SecurityGroupIds' : metadata ['SecurityGroupIds' ],
101
- 'Subnets' : metadata ['Subnets' ]
102
- }
103
- }
104
- }
73
+ },
105
74
}
106
75
return config
107
76
108
- if __name__ == '__main__' :
109
- try :
77
+
78
+ if __name__ == "__main__" :
79
+ try :
110
80
config = {}
111
- resource_metadata = ' /opt/ml/metadata/resource-metadata.json'
81
+ resource_metadata = " /opt/ml/metadata/resource-metadata.json"
112
82
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/"
120
90
121
91
if os .path .exists (resource_metadata ):
122
- with open (resource_metadata , 'r' ) as file :
92
+ with open (resource_metadata , "r" ) as file :
123
93
data = json .load (file )
124
94
125
- s3_path = data [META_DATA ].get (PROJECT_S3_PATH , '' )
95
+ s3_path = data [META_DATA ].get (PROJECT_S3_PATH , "" )
126
96
metadata = {
127
97
# user provided bucket
128
- ' S3Bucket' : re .search (r"s3://([^/]+)/" , s3_path ).group (1 ),
98
+ " S3Bucket" : re .search (r"s3://([^/]+)/" , s3_path ).group (1 ),
129
99
# ${datazoneEnvironmentDomainId}/${datazoneEnvironmentProjectId}/${datazoneScopeName}/
130
- ' S3ObjectKeyPrefix' : s3_path .split ("//" )[1 ].split ("/" , 1 )[1 ],
100
+ " S3ObjectKeyPrefix" : s3_path .split ("//" )[1 ].split ("/" , 1 )[1 ],
131
101
# TODO: Is this a billing concern if set default
132
102
# '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 ],
136
106
}
137
-
107
+
138
108
# 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 == ["" ]]
140
110
if empty_values :
141
111
raise AttributeError (f"There are empty values in the metadata: { empty_values } " )
142
112
143
113
config = generate_intelligent_default_config (metadata )
144
114
else :
145
- raise FileNotFoundError (' No resource-metadata.json exists on host!' )
146
-
115
+ raise FileNotFoundError (" No resource-metadata.json exists on host!" )
116
+
147
117
# 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 :
149
119
yaml .dump (config , f , default_flow_style = False , sort_keys = False )
150
-
120
+
151
121
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!" )
0 commit comments