From 18a320b2100980b2b7007ca57ea26519b0262001 Mon Sep 17 00:00:00 2001 From: Dan Choi Date: Fri, 17 Aug 2018 14:18:14 -0700 Subject: [PATCH 1/2] Increase read timeout for sagemaker runtime --- src/sagemaker/session.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py index ec62e09ac0..0eac651b3d 100644 --- a/src/sagemaker/session.py +++ b/src/sagemaker/session.py @@ -29,7 +29,6 @@ from sagemaker.utils import name_from_image, secondary_training_status_message, secondary_training_status_changed import sagemaker.logs - logging.basicConfig() LOGGER = logging.getLogger('sagemaker') LOGGER.setLevel(logging.INFO) @@ -93,7 +92,9 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client): self.sagemaker_client = sagemaker_client or self.boto_session.client('sagemaker') prepend_user_agent(self.sagemaker_client) - self.sagemaker_runtime_client = sagemaker_runtime_client or self.boto_session.client('runtime.sagemaker') + config = botocore.config.Config(read_timeout=80) + self.sagemaker_runtime_client = sagemaker_runtime_client or self.boto_session.client('runtime.sagemaker', + config=config) prepend_user_agent(self.sagemaker_runtime_client) self.local_mode = False @@ -1018,7 +1019,6 @@ def _deployment_entity_exists(describe_fn): def _train_done(sagemaker_client, job_name, last_desc): - in_progress_statuses = ['InProgress', 'Created'] desc = sagemaker_client.describe_training_job(TrainingJobName=job_name) From 65b9446a7345322e3e9318e81120e5a1f7e358c0 Mon Sep 17 00:00:00 2001 From: Dan Choi Date: Tue, 21 Aug 2018 15:37:42 -0700 Subject: [PATCH 2/2] Address comments --- src/sagemaker/session.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py index b3ae9ad783..e904d0531c 100644 --- a/src/sagemaker/session.py +++ b/src/sagemaker/session.py @@ -92,9 +92,12 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client): self.sagemaker_client = sagemaker_client or self.boto_session.client('sagemaker') prepend_user_agent(self.sagemaker_client) - config = botocore.config.Config(read_timeout=80) - self.sagemaker_runtime_client = sagemaker_runtime_client or self.boto_session.client('runtime.sagemaker', - config=config) + if sagemaker_runtime_client is not None: + self.sagemaker_runtime_client = sagemaker_runtime_client + else: + config = botocore.config.Config(read_timeout=80) + self.sagemaker_runtime_client = self.boto_session.client('runtime.sagemaker', config=config) + prepend_user_agent(self.sagemaker_runtime_client) self.local_mode = False