change: seperate logs() from attach()

Author: Chuyang Deng

src/sagemaker/estimator.py

@@ -630,9 +630,15 @@ def attach(cls, training_job_name, sagemaker_session=None, model_channel_name="m
             sagemaker_session=sagemaker_session, job_name=training_job_name
         )
         estimator._current_job_name = estimator.latest_training_job.name
-        estimator.latest_training_job.wait()
+        estimator.latest_training_job.wait(logs="None")
         return estimator
 
+    def logs(self):
+        """Display the logs for Estimator's training job. If the output is a tty or a Jupyter
+        cell, it will be color-coded based on which instance the log entry is from.
+        """
+        self.sagemaker_session.logs_for_job(self.latest_training_job, wait=True)
+
     def deploy(
         self,
         initial_instance_count,

tests/unit/test_estimator.py

@@ -681,6 +681,27 @@ def test_attach_framework(sagemaker_session):
     assert framework_estimator.enable_network_isolation() is True
 
 
+def test_attach_no_logs(sagemaker_session):
+    returned_job_description = RETURNED_JOB_DESCRIPTION.copy()
+    mock_describe_training_job = Mock(
+        name="describe_training_job", return_value=returned_job_description
+    )
+    sagemaker_session.sagemaker_client.describe_training_job = mock_describe_training_job
+    Estimator.attach(training_job_name="job", sagemaker_session=sagemaker_session)
+    sagemaker_session.logs_for_job.assert_not_called()
+
+
+def test_logs(sagemaker_session):
+    returned_job_description = RETURNED_JOB_DESCRIPTION.copy()
+    mock_describe_training_job = Mock(
+        name="describe_training_job", return_value=returned_job_description
+    )
+    sagemaker_session.sagemaker_client.describe_training_job = mock_describe_training_job
+    estimator = Estimator.attach(training_job_name="job", sagemaker_session=sagemaker_session)
+    estimator.logs()
+    sagemaker_session.logs_for_job.assert_called_with(estimator.latest_training_job, wait=True)
+
+
 def test_attach_without_hyperparameters(sagemaker_session):
     returned_job_description = RETURNED_JOB_DESCRIPTION.copy()
     del returned_job_description["HyperParameters"]