Merge pull request #10 from aws/mvs-poll-every-30-s

describing training job call every 30 seconds

Author: mvsusp

src/sagemaker/session.py

@@ -605,7 +605,7 @@ def logs_for_job(self, job_name, wait=False, poll=5):  # noqa: C901 - suppress c
         # Notes:
         # - The JOB_COMPLETE state forces us to do an extra pause and read any items that got to Cloudwatch after
         #   the job was marked complete.
-
+        last_describe_job_call = time.time()
         while True:
             if len(stream_names) < instance_count:
                 # Log streams are created whenever a container starts writing to stdout/err, so this list
@@ -645,8 +645,10 @@ def logs_for_job(self, job_name, wait=False, poll=5):  # noqa: C901 - suppress c
 
             if state == LogState.JOB_COMPLETE:
                 state = LogState.COMPLETE
-            else:
+            elif time.time() - last_describe_job_call >= 30:
                 description = self.sagemaker_client.describe_training_job(TrainingJobName=job_name)
+                last_describe_job_call = time.time()
+
                 status = description['TrainingJobStatus']
 
                 if status == 'Completed' or status == 'Failed':

tests/unit/test_session.py

@@ -311,7 +311,8 @@ def test_logs_for_job_no_wait_on_running(cw, sagemaker_session_ready_lifecycle):
 
 
 @patch('sagemaker.logs.ColorWrap')
-def test_logs_for_job_full_lifecycle(cw, sagemaker_session_full_lifecycle):
+@patch('time.time', side_effect=[0, 30, 60, 90, 120, 150, 180])
+def test_logs_for_job_full_lifecycle(time, cw, sagemaker_session_full_lifecycle):
     ims = sagemaker_session_full_lifecycle
     ims.logs_for_job(JOB_NAME, wait=True, poll=0)
     assert ims.sagemaker_client.describe_training_job.call_args_list == [call(TrainingJobName=JOB_NAME,)] * 3