Skip to content

Commit 081c68a

Browse files
Increase read timeout for sagemaker runtime (#355)
1 parent 5e61265 commit 081c68a

File tree

1 file changed

+6
-3
lines changed

1 file changed

+6
-3
lines changed

src/sagemaker/session.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
from sagemaker.utils import name_from_image, secondary_training_status_message, secondary_training_status_changed
3030
import sagemaker.logs
3131

32-
3332
logging.basicConfig()
3433
LOGGER = logging.getLogger('sagemaker')
3534
LOGGER.setLevel(logging.INFO)
@@ -93,7 +92,12 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
9392
self.sagemaker_client = sagemaker_client or self.boto_session.client('sagemaker')
9493
prepend_user_agent(self.sagemaker_client)
9594

96-
self.sagemaker_runtime_client = sagemaker_runtime_client or self.boto_session.client('runtime.sagemaker')
95+
if sagemaker_runtime_client is not None:
96+
self.sagemaker_runtime_client = sagemaker_runtime_client
97+
else:
98+
config = botocore.config.Config(read_timeout=80)
99+
self.sagemaker_runtime_client = self.boto_session.client('runtime.sagemaker', config=config)
100+
97101
prepend_user_agent(self.sagemaker_runtime_client)
98102

99103
self.local_mode = False
@@ -1028,7 +1032,6 @@ def _deployment_entity_exists(describe_fn):
10281032

10291033

10301034
def _train_done(sagemaker_client, job_name, last_desc):
1031-
10321035
in_progress_statuses = ['InProgress', 'Created']
10331036

10341037
desc = sagemaker_client.describe_training_job(TrainingJobName=job_name)

0 commit comments

Comments
 (0)