Refactor the rest of the 1P estimators (aws#24)

Author: laurenyu

src/sagemaker/amazon/kmeans.py

@@ -133,8 +133,8 @@ def create_model(self):
         s3 model data produced by this Estimator."""
         return KMeansModel(self.model_data, self.role, self.sagemaker_session)
 
-    def fit(self, records, mini_batch_size=5000, **kwargs):
-        super(KMeans, self).fit(records, mini_batch_size, **kwargs)
+    def prepare_for_training(self, records, mini_batch_size=5000, job_name=None):
+        super(KMeans, self).prepare_for_training(records, mini_batch_size=mini_batch_size, job_name=job_name)
 
     def hyperparameters(self):
         """Return the SageMaker hyperparameters for training this KMeans Estimator"""

src/sagemaker/amazon/lda.py

@@ -93,11 +93,12 @@ def create_model(self):
 
         return LDAModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session)
 
-    def fit(self, records, mini_batch_size, **kwargs):
+    def prepare_for_training(self, records, mini_batch_size, job_name=None):
         # mini_batch_size is required, prevent explicit calls with None
         if mini_batch_size is None:
             raise ValueError("mini_batch_size must be set")
-        super(LDA, self).fit(records, mini_batch_size, **kwargs)
+
+        super(LDA, self).prepare_for_training(records, mini_batch_size=mini_batch_size, job_name=job_name)
 
 
 class LDAPredictor(RealTimePredictor):

src/sagemaker/amazon/linear_learner.py

@@ -228,12 +228,12 @@ def create_model(self):
 
         return LinearLearnerModel(self.model_data, self.role, self.sagemaker_session)
 
-    def fit(self, records, mini_batch_size=None, **kwargs):
+    def prepare_for_training(self, records, mini_batch_size=None, job_name=None):
         # mini_batch_size can't be greater than number of records or training job fails
         default_mini_batch_size = min(self.DEFAULT_MINI_BATCH_SIZE,
                                       max(1, int(records.num_records / self.train_instance_count)))
         use_mini_batch_size = mini_batch_size or default_mini_batch_size
-        super(LinearLearner, self).fit(records, use_mini_batch_size, **kwargs)
+        super(LinearLearner, self).prepare_for_training(records, mini_batch_size=use_mini_batch_size, job_name=job_name)
 
 
 class LinearLearnerPredictor(RealTimePredictor):

src/sagemaker/amazon/ntm.py

@@ -113,10 +113,10 @@ def create_model(self):
 
         return NTMModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session)
 
-    def fit(self, records, mini_batch_size=None, **kwargs):
+    def prepare_for_training(self, records, mini_batch_size, job_name=None):
         if mini_batch_size is not None and (mini_batch_size < 1 or mini_batch_size > 10000):
             raise ValueError("mini_batch_size must be in [1, 10000]")
-        super(NTM, self).fit(records, mini_batch_size, **kwargs)
+        super(NTM, self).prepare_for_training(records, mini_batch_size=mini_batch_size, job_name=job_name)
 
 
 class NTMPredictor(RealTimePredictor):

src/sagemaker/amazon/randomcutforest.py

@@ -87,13 +87,11 @@ def create_model(self):
 
         return RandomCutForestModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session)
 
-    def fit(self, records, mini_batch_size=None, **kwargs):
-        if mini_batch_size is None:
-            mini_batch_size = RandomCutForest.MINI_BATCH_SIZE
-        elif mini_batch_size != RandomCutForest.MINI_BATCH_SIZE:
+    def prepare_for_training(self, records, mini_batch_size=MINI_BATCH_SIZE, job_name=None):
+        if mini_batch_size != self.MINI_BATCH_SIZE:
             raise ValueError("Random Cut Forest uses a fixed mini_batch_size of {}"
-                             .format(RandomCutForest.MINI_BATCH_SIZE))
-        super(RandomCutForest, self).fit(records, mini_batch_size, **kwargs)
+                             .format(self.MINI_BATCH_SIZE))
+        super(RandomCutForest, self).prepare_for_training(records, mini_batch_size=mini_batch_size, job_name=job_name)
 
 
 class RandomCutForestPredictor(RealTimePredictor):

tests/unit/test_fm.py

@@ -210,31 +210,31 @@ def test_call_fit(base_fit, sagemaker_session):
     assert base_fit.call_args[0][1] == MINI_BATCH_SIZE
 
 
-def test_call_fit_none_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_no_mini_batch_size(sagemaker_session):
     fm = FactorizationMachines(base_job_name='fm', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
-    fm.fit(data)
+    fm.prepare_for_training(data)
 
 
-def test_call_fit_wrong_type_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_wrong_type_mini_batch_size(sagemaker_session):
     fm = FactorizationMachines(base_job_name='fm', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
 
     with pytest.raises((TypeError, ValueError)):
-        fm.fit(data, 'some')
+        fm.prepare_for_training(data, 'some')
 
 
-def test_call_fit_wrong_value_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_wrong_value_mini_batch_size(sagemaker_session):
     fm = FactorizationMachines(base_job_name='fm', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
     with pytest.raises(ValueError):
-        fm.fit(data, 0)
+        fm.prepare_for_training(data, 0)
 
 
 def test_model_image(sagemaker_session):

tests/unit/test_kmeans.py

@@ -175,31 +175,33 @@ def test_call_fit(base_fit, sagemaker_session):
     assert base_fit.call_args[0][1] == MINI_BATCH_SIZE
 
 
-def test_call_fit_none_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_no_mini_batch_size(sagemaker_session):
     kmeans = KMeans(base_job_name='kmeans', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
-    kmeans.fit(data)
+    kmeans.prepare_for_training(data)
 
+    assert kmeans.mini_batch_size == 5000
 
-def test_call_fit_wrong_type_mini_batch_size(sagemaker_session):
+
+def test_prepare_for_training_wrong_type_mini_batch_size(sagemaker_session):
     kmeans = KMeans(base_job_name='kmeans', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
 
     with pytest.raises((TypeError, ValueError)):
-        kmeans.fit(data, 'some')
+        kmeans.prepare_for_training(data, 'some')
 
 
-def test_call_fit_wrong_value_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_wrong_value_mini_batch_size(sagemaker_session):
     kmeans = KMeans(base_job_name='kmeans', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
     with pytest.raises(ValueError):
-        kmeans.fit(data, 0)
+        kmeans.prepare_for_training(data, 0)
 
 
 def test_model_image(sagemaker_session):

tests/unit/test_lda.py

@@ -147,32 +147,32 @@ def test_call_fit(base_fit, sagemaker_session):
     assert base_fit.call_args[0][1] == MINI_BATCH_SZIE
 
 
-def test_call_fit_none_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_no_mini_batch_size(sagemaker_session):
     lda = LDA(base_job_name='lda', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
     with pytest.raises(ValueError):
-        lda.fit(data, None)
+        lda.prepare_for_training(data, None)
 
 
-def test_call_fit_wrong_type_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_wrong_type_mini_batch_size(sagemaker_session):
     lda = LDA(base_job_name='lda', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
 
     with pytest.raises(ValueError):
-        lda.fit(data, 'some')
+        lda.prepare_for_training(data, 'some')
 
 
-def test_call_fit_wrong_value_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_wrong_value_mini_batch_size(sagemaker_session):
     lda = LDA(base_job_name='lda', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
     with pytest.raises(ValueError):
-        lda.fit(data, 0)
+        lda.prepare_for_training(data, 0)
 
 
 def test_model_image(sagemaker_session):

tests/unit/test_linear_learner.py

@@ -218,35 +218,27 @@ def test_optional_hyper_parameters_value(sagemaker_session, optional_hyper_param
 DEFAULT_MINI_BATCH_SIZE = 1000
 
 
-@patch('sagemaker.amazon.amazon_estimator.AmazonAlgorithmEstimatorBase.fit')
-def test_call_fit_calculate_batch_size_1(base_fit, sagemaker_session):
+def test_prepare_for_training_calculate_batch_size_1(sagemaker_session):
     lr = LinearLearner(base_job_name='lr', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM, channel='train')
 
-    lr.fit(data)
+    lr.prepare_for_training(data)
 
-    base_fit.assert_called_once()
-    assert len(base_fit.call_args[0]) == 2
-    assert base_fit.call_args[0][0] == data
-    assert base_fit.call_args[0][1] == 1
+    assert lr.mini_batch_size == 1
 
 
-@patch('sagemaker.amazon.amazon_estimator.AmazonAlgorithmEstimatorBase.fit')
-def test_call_fit_calculate_batch_size_2(base_fit, sagemaker_session):
+def test_prepare_for_training_calculate_batch_size_2(sagemaker_session):
     lr = LinearLearner(base_job_name='lr', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX),
                      num_records=10000,
                      feature_dim=FEATURE_DIM,
                      channel='train')
 
-    lr.fit(data)
+    lr.prepare_for_training(data)
 
-    base_fit.assert_called_once()
-    assert len(base_fit.call_args[0]) == 2
-    assert base_fit.call_args[0][0] == data
-    assert base_fit.call_args[0][1] == DEFAULT_MINI_BATCH_SIZE
+    assert lr.mini_batch_size == DEFAULT_MINI_BATCH_SIZE
 
 
 @patch('sagemaker.amazon.amazon_estimator.AmazonAlgorithmEstimatorBase.fit')

tests/unit/test_ntm.py

@@ -193,32 +193,32 @@ def test_call_fit_none_mini_batch_size(sagemaker_session):
     ntm.fit(data)
 
 
-def test_call_fit_wrong_type_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_wrong_type_mini_batch_size(sagemaker_session):
     ntm = NTM(base_job_name="ntm", sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet("s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
 
     with pytest.raises((TypeError, ValueError)):
-        ntm.fit(data, "some")
+        ntm.prepare_for_training(data, "some")
 
 
-def test_call_fit_wrong_value_lower_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_wrong_value_lower_mini_batch_size(sagemaker_session):
     ntm = NTM(base_job_name="ntm", sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet("s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
     with pytest.raises(ValueError):
-        ntm.fit(data, 0)
+        ntm.prepare_for_training(data, 0)
 
 
-def test_call_fit_wrong_value_upper_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_wrong_value_upper_mini_batch_size(sagemaker_session):
     ntm = NTM(base_job_name="ntm", sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet("s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
     with pytest.raises(ValueError):
-        ntm.fit(data, 10001)
+        ntm.prepare_for_training(data, 10001)
 
 
 def test_model_image(sagemaker_session):

tests/unit/test_pca.py

@@ -143,15 +143,17 @@ def test_call_fit(base_fit, sagemaker_session):
     assert base_fit.call_args[0][1] == MINI_BATCH_SIZE
 
 
-def test_call_fit_none_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_no_mini_batch_size(sagemaker_session):
     pca = PCA(base_job_name='pca', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
-    pca.fit(data)
+    pca.prepare_for_training(data)
 
+    assert pca.mini_batch_size == 1
 
-def test_call_fit_wrong_type_mini_batch_size(sagemaker_session):
+
+def test_prepare_for_training_wrong_type_mini_batch_size(sagemaker_session):
     pca = PCA(base_job_name='pca', sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
 
     data = RecordSet('s3://{}/{}'.format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,

tests/unit/test_randomcutforest.py

@@ -141,35 +141,37 @@ def test_call_fit(base_fit, sagemaker_session):
     assert base_fit.call_args[0][1] == MINI_BATCH_SIZE
 
 
-def test_call_fit_none_mini_batch_size(sagemaker_session):
+def test_prepare_for_training_no_mini_batch_size(sagemaker_session):
     randomcutforest = RandomCutForest(base_job_name="randomcutforest", sagemaker_session=sagemaker_session,
                                       **ALL_REQ_ARGS)
 
     data = RecordSet("s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
-    randomcutforest.fit(data)
+    randomcutforest.prepare_for_training(data)
 
+    assert randomcutforest.mini_batch_size == MINI_BATCH_SIZE
 
-def test_call_fit_wrong_type_mini_batch_size(sagemaker_session):
+
+def test_prepare_for_training_wrong_type_mini_batch_size(sagemaker_session):
     randomcutforest = RandomCutForest(base_job_name="randomcutforest", sagemaker_session=sagemaker_session,
                                       **ALL_REQ_ARGS)
 
     data = RecordSet("s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=FEATURE_DIM,
                      channel='train')
 
     with pytest.raises((TypeError, ValueError)):
-        randomcutforest.fit(data, 1234)
+        randomcutforest.prepare_for_training(data, 1234)
 
 
-def test_call_fit_feature_dim_greater_than_max_allowed(sagemaker_session):
+def test_prepare_for_training_feature_dim_greater_than_max_allowed(sagemaker_session):
     randomcutforest = RandomCutForest(base_job_name="randomcutforest", sagemaker_session=sagemaker_session,
                                       **ALL_REQ_ARGS)
 
     data = RecordSet("s3://{}/{}".format(BUCKET_NAME, PREFIX), num_records=1, feature_dim=MAX_FEATURE_DIM + 1,
                      channel='train')
 
     with pytest.raises((TypeError, ValueError)):
-        randomcutforest.fit(data)
+        randomcutforest.prepare_for_training(data)
 
 
 def test_model_image(sagemaker_session):