Skip to content

Commit 49f0248

Browse files
committed
fix integ test hub
1 parent 567c864 commit 49f0248

File tree

1 file changed

+14
-12
lines changed
  • src/sagemaker/jumpstart/hub

1 file changed

+14
-12
lines changed

src/sagemaker/jumpstart/hub/hub.py

+14-12
Original file line numberDiff line numberDiff line change
@@ -56,8 +56,8 @@ class Hub:
5656
def __init__(
5757
self,
5858
hub_name: str,
59+
sagemaker_session: Session,
5960
bucket_name: Optional[str] = None,
60-
sagemaker_session: Optional[Session] = None,
6161
) -> None:
6262
"""Instantiates a SageMaker ``Hub``.
6363
@@ -94,20 +94,22 @@ def create(
9494
"""Creates a hub with the given description"""
9595
curr_timestamp = datetime.now().timestamp()
9696

97-
return self._sagemaker_session.create_hub(
98-
hub_name=self.hub_name,
99-
hub_description=description,
100-
hub_display_name=display_name,
101-
hub_search_keywords=search_keywords,
102-
s3_storage_config={
97+
request = {
98+
"hub_name": self.hub_name,
99+
"hub_description": description,
100+
"hub_display_name": display_name,
101+
"hub_search_keywords": search_keywords,
102+
"tags": tags
103+
}
104+
105+
if self.bucket_name:
106+
request["s3_storage_config"] = {
103107
"S3OutputPath": (
104108
f"s3://{self.bucket_name}/{self.hub_name}-{curr_timestamp}"
105-
if self.bucket_name
106-
else None
107109
)
108-
},
109-
tags=tags,
110-
)
110+
}
111+
112+
return self._sagemaker_session.create_hub(**request)
111113

112114
def describe(self, hub_name: Optional[str] = None) -> DescribeHubResponse:
113115
"""Returns descriptive information about the Hub"""

0 commit comments

Comments
 (0)