33
33
from sagemaker .jumpstart .hub .utils import (
34
34
get_hub_model_version ,
35
35
get_info_from_hub_resource_arn ,
36
+ create_hub_bucket_if_it_does_not_exist ,
37
+ generate_default_hub_bucket_name ,
36
38
create_s3_object_reference_from_uri ,
37
39
construct_hub_arn_from_name ,
38
40
)
@@ -88,23 +90,25 @@ def _fetch_hub_bucket_name(self) -> str:
88
90
if hub_output_location :
89
91
location = create_s3_object_reference_from_uri (hub_output_location )
90
92
return location .bucket
93
+ default_bucket_name = generate_default_hub_bucket_name (self ._sagemaker_session )
91
94
JUMPSTART_LOGGER .warning (
92
- "There is not a Hub bucket associated with %s." ,
95
+ "There is not a Hub bucket associated with %s. Using %s " ,
93
96
self .hub_name ,
97
+ default_bucket_name ,
94
98
)
95
- return None
99
+ return default_bucket_name
96
100
except exceptions .ClientError :
101
+ hub_bucket_name = generate_default_hub_bucket_name (self ._sagemaker_session )
97
102
JUMPSTART_LOGGER .warning (
98
- "There is not a Hub bucket associated with %s." ,
103
+ "There is not a Hub bucket associated with %s. Using %s " ,
99
104
self .hub_name ,
105
+ hub_bucket_name ,
100
106
)
101
- return None
107
+ return hub_bucket_name
102
108
103
109
def _generate_hub_storage_location (self , bucket_name : Optional [str ] = None ) -> None :
104
110
"""Generates an ``S3ObjectLocation`` given a Hub name."""
105
111
hub_bucket_name = bucket_name or self ._fetch_hub_bucket_name ()
106
- if hub_bucket_name is None :
107
- return
108
112
curr_timestamp = datetime .now ().timestamp ()
109
113
return S3ObjectLocation (bucket = hub_bucket_name , key = f"{ self .hub_name } -{ curr_timestamp } " )
110
114
@@ -127,18 +131,16 @@ def create(
127
131
) -> Dict [str , str ]:
128
132
"""Creates a hub with the given description"""
129
133
130
- s3_storage_config = (
131
- {"S3OutputPath" : self .hub_storage_location .get_uri ()}
132
- if self .hub_storage_location
133
- else None
134
+ create_hub_bucket_if_it_does_not_exist (
135
+ self .hub_storage_location .bucket , self ._sagemaker_session
134
136
)
135
137
136
138
return self ._sagemaker_session .create_hub (
137
139
hub_name = self .hub_name ,
138
140
hub_description = description ,
139
141
hub_display_name = display_name ,
140
142
hub_search_keywords = search_keywords ,
141
- s3_storage_config = s3_storage_config ,
143
+ s3_storage_config = { "S3OutputPath" : self . hub_storage_location . get_uri ()} ,
142
144
tags = tags ,
143
145
)
144
146
0 commit comments