Make start time, end time and period configurable in analytics.TrainingJobAnalytics

Creating an TrainingJobAnalytics object fails if the training job has too many
data points in the specified metrics. Make start time, end time and period
configurable so the caller can get around this limit -
https://docs.aws.amazon.com/AmazonCloudWatch/latest/APIReference/API_GetMetricStatistics.html

Original issue:
aws#701

Author: icywang86rui

src/sagemaker/analytics.py

@@ -201,7 +201,8 @@ class TrainingJobAnalytics(AnalyticsMetricsBase):
 
     CLOUDWATCH_NAMESPACE = '/aws/sagemaker/TrainingJobs'
 
-    def __init__(self, training_job_name, metric_names=None, sagemaker_session=None):
+    def __init__(self, training_job_name, metric_names=None, sagemaker_session=None,
+                 start_time=None, end_time=None, period=None):
         """Initialize a ``TrainingJobAnalytics`` instance.
 
         Args:
@@ -216,6 +217,10 @@ def __init__(self, training_job_name, metric_names=None, sagemaker_session=None)
         self._sage_client = sagemaker_session.sagemaker_client
         self._cloudwatch = sagemaker_session.boto_session.client('cloudwatch')
         self._training_job_name = training_job_name
+        self._start_time = start_time
+        self._end_time = end_time
+        self._period = period if period else 60
+
         if metric_names:
             self._metric_names = metric_names
         else:
@@ -245,13 +250,15 @@ def _determine_timeinterval(self):
         covering the interval of the training job
         """
         description = self._sage_client.describe_training_job(TrainingJobName=self.name)
-        start_time = description[u'TrainingStartTime']  # datetime object
+        start_time = self._start_time if self._start_time else description[u'TrainingStartTime']  # datetime object
         # Incrementing end time by 1 min since CloudWatch drops seconds before finding the logs.
         # This results in logs being searched in the time range in which the correct log line was not present.
         # Example - Log time - 2018-10-22 08:25:55
         #           Here calculated end time would also be 2018-10-22 08:25:55 (without 1 min addition)
         #           CW will consider end time as 2018-10-22 08:25 and will not be able to search the correct log.
-        end_time = description.get(u'TrainingEndTime', datetime.datetime.utcnow()) + datetime.timedelta(minutes=1)
+        end_time = self._end_time if self._end_time else description.get(
+            u'TrainingEndTime', datetime.datetime.utcnow()) + datetime.timedelta(minutes=1)
+
         return {
             'start_time': start_time,
             'end_time': end_time,
@@ -276,7 +283,7 @@ def _fetch_metric(self, metric_name):
             ],
             'StartTime': self._time_interval['start_time'],
             'EndTime': self._time_interval['end_time'],
-            'Period': 60,
+            'Period': self._period,
             'Statistics': ['Average'],
         }
         raw_cwm_data = self._cloudwatch.get_metric_statistics(**request)['Datapoints']

tests/unit/test_analytics.py

@@ -197,7 +197,8 @@ def test_trainer_name():
         'TrainingEndTime': datetime.datetime(2018, 5, 16, 5, 6, 7),
     }
     session = create_sagemaker_session(describe_training_result)
-    trainer = TrainingJobAnalytics("my-training-job", ["metric"], sagemaker_session=session)
+    trainer = TrainingJobAnalytics(training_job_name="my-training-job", metric_names=["metric"],
+                                   sagemaker_session=session)
     assert trainer.name == "my-training-job"
     assert str(trainer).find("my-training-job") != -1
 
@@ -231,7 +232,8 @@ def _metric_stats_results():
 def test_trainer_dataframe():
     session = create_sagemaker_session(describe_training_result=_describe_training_result(),
                                        metric_stats_results=_metric_stats_results())
-    trainer = TrainingJobAnalytics("my-training-job", ["train:acc"], sagemaker_session=session)
+    trainer = TrainingJobAnalytics(training_job_name="my-training-job", metric_names=["train:acc"],
+                                   sagemaker_session=session)
 
     df = trainer.dataframe()
     assert df is not None
@@ -245,3 +247,20 @@ def test_trainer_dataframe():
     trainer.export_csv(tmp_name)
     assert os.path.isfile(tmp_name)
     os.unlink(tmp_name)
+
+
+def test_start_time_end_time_and_period_specified():
+    describe_training_result = {
+        'TrainingStartTime': datetime.datetime(2018, 5, 16, 1, 2, 3),
+        'TrainingEndTime': datetime.datetime(2018, 5, 16, 5, 6, 7),
+    }
+    session = create_sagemaker_session(describe_training_result)
+    start_time = datetime.datetime(2018, 5, 16, 1, 3, 4)
+    end_time = datetime.datetime(2018, 5, 16, 5, 1, 1)
+    period = 300
+    trainer = TrainingJobAnalytics(training_job_name="my-training-job", metric_names=["metric"],
+                                   sagemaker_session=session, start_time=start_time, end_time=end_time, period=period)
+
+    assert trainer._time_interval['start_time'] == start_time
+    assert trainer._time_interval['end_time'] == end_time
+    assert trainer._period == period