diff --git a/src/sagemaker/algorithm.py b/src/sagemaker/algorithm.py index 05c27a99a5..c7114aee2b 100644 --- a/src/sagemaker/algorithm.py +++ b/src/sagemaker/algorithm.py @@ -229,13 +229,13 @@ def hyperparameters(self): """ return self.hyperparam_dict - def train_image(self): + def training_image_uri(self): """Returns the docker image to use for training. The fit() method, that does the model training, calls this method to find the image to use for model training. """ - raise RuntimeError("train_image is never meant to be called on Algorithm Estimators") + raise RuntimeError("training_image_uri is never meant to be called on Algorithm Estimators") def enable_network_isolation(self): """Return True if this Estimator will need network isolation to run. diff --git a/src/sagemaker/amazon/amazon_estimator.py b/src/sagemaker/amazon/amazon_estimator.py index dae361174c..307ac97f6d 100644 --- a/src/sagemaker/amazon/amazon_estimator.py +++ b/src/sagemaker/amazon/amazon_estimator.py @@ -91,7 +91,7 @@ def __init__( ) self._data_location = data_location - def train_image(self): + def training_image_uri(self): """Placeholder docstring""" return image_uris.retrieve( self.repo_name, self.sagemaker_session.boto_region_name, version=self.repo_version, diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py index 1e1387e62e..c05f51f319 100644 --- a/src/sagemaker/estimator.py +++ b/src/sagemaker/estimator.py @@ -285,7 +285,7 @@ def __init__( self._enable_network_isolation = enable_network_isolation @abstractmethod - def train_image(self): + def training_image_uri(self): """Return the Docker image to use for training. The :meth:`~sagemaker.estimator.EstimatorBase.fit` method, which does @@ -329,7 +329,7 @@ def _ensure_base_job_name(self): """Set ``self.base_job_name`` if it is not set already.""" # honor supplied base_job_name or generate it if self.base_job_name is None: - self.base_job_name = base_name_from_image(self.train_image()) + self.base_job_name = base_name_from_image(self.training_image_uri()) def _get_or_create_name(self, name=None): """Generate a name based on the base job name or training image if needed. @@ -507,7 +507,7 @@ def fit(self, inputs=None, wait=True, logs="All", job_name=None, experiment_conf def _compilation_job_name(self): """Placeholder docstring""" - base_name = self.base_job_name or base_name_from_image(self.train_image()) + base_name = self.base_job_name or base_name_from_image(self.training_image_uri()) return name_from_base("compilation-" + base_name) def compile_model( @@ -1083,7 +1083,7 @@ def start_new(cls, estimator, inputs, experiment_config): if isinstance(estimator, sagemaker.algorithm.AlgorithmEstimator): train_args["algorithm_arn"] = estimator.algorithm_arn else: - train_args["image_uri"] = estimator.train_image() + train_args["image_uri"] = estimator.training_image_uri() if estimator.debugger_rule_configs: train_args["debugger_rule_configs"] = estimator.debugger_rule_configs @@ -1350,7 +1350,7 @@ def __init__( enable_network_isolation=enable_network_isolation, ) - def train_image(self): + def training_image_uri(self): """Returns the docker image to use for training. The fit() method, that does the model training, calls this method to @@ -1424,7 +1424,7 @@ def predict_wrapper(endpoint, session): kwargs["enable_network_isolation"] = self.enable_network_isolation() return Model( - image_uri or self.train_image(), + image_uri or self.training_image_uri(), self.model_data, role, vpc_config=self.get_vpc_config(vpc_config_override), @@ -1826,7 +1826,7 @@ class constructor return init_params - def train_image(self): + def training_image_uri(self): """Return the Docker image to use for training. The :meth:`~sagemaker.estimator.EstimatorBase.fit` method, which does diff --git a/src/sagemaker/rl/estimator.py b/src/sagemaker/rl/estimator.py index 8b319b4b8c..0bc21c6ef7 100644 --- a/src/sagemaker/rl/estimator.py +++ b/src/sagemaker/rl/estimator.py @@ -268,7 +268,7 @@ def create_model( "An unknown RLFramework enum was passed in. framework: {}".format(self.framework) ) - def train_image(self): + def training_image_uri(self): """Return the Docker image to use for training. The :meth:`~sagemaker.estimator.EstimatorBase.fit` method, which does diff --git a/src/sagemaker/tuner.py b/src/sagemaker/tuner.py index 36f9cc69fc..6de34d6aa5 100644 --- a/src/sagemaker/tuner.py +++ b/src/sagemaker/tuner.py @@ -317,7 +317,7 @@ def _prepare_job_name_for_tuning(self, job_name=None): estimator = ( self.estimator or self.estimator_dict[sorted(self.estimator_dict.keys())[0]] ) - base_name = base_name_from_image(estimator.train_image()) + base_name = base_name_from_image(estimator.training_image_uri()) self._current_job_name = name_from_base( base_name, max_length=self.TUNING_JOB_NAME_MAX_LENGTH, short=True ) @@ -1527,7 +1527,7 @@ def _prepare_training_config( if isinstance(estimator, sagemaker.algorithm.AlgorithmEstimator): training_config["algorithm_arn"] = estimator.algorithm_arn else: - training_config["image_uri"] = estimator.train_image() + training_config["image_uri"] = estimator.training_image_uri() training_config["enable_network_isolation"] = estimator.enable_network_isolation() training_config[ diff --git a/src/sagemaker/workflow/airflow.py b/src/sagemaker/workflow/airflow.py index 93ed1c22cb..7b706f8ade 100644 --- a/src/sagemaker/workflow/airflow.py +++ b/src/sagemaker/workflow/airflow.py @@ -149,7 +149,9 @@ def training_base_config(estimator, inputs=None, job_name=None, mini_batch_size= if job_name is not None: estimator._current_job_name = job_name else: - base_name = estimator.base_job_name or utils.base_name_from_image(estimator.train_image()) + base_name = estimator.base_job_name or utils.base_name_from_image( + estimator.training_image_uri() + ) estimator._current_job_name = utils.name_from_base(base_name) if estimator.output_path is None: @@ -164,7 +166,7 @@ def training_base_config(estimator, inputs=None, job_name=None, mini_batch_size= train_config = { "AlgorithmSpecification": { - "TrainingImage": estimator.train_image(), + "TrainingImage": estimator.training_image_uri(), "TrainingInputMode": estimator.input_mode, }, "OutputDataConfig": job_config["output_config"], diff --git a/tests/integ/test_byo_estimator.py b/tests/integ/test_byo_estimator.py index e8a64da77c..475c6a2cdc 100644 --- a/tests/integ/test_byo_estimator.py +++ b/tests/integ/test_byo_estimator.py @@ -151,4 +151,4 @@ def test_async_byo_estimator(sagemaker_session, region, cpu_instance_type, train for prediction in result["predictions"]: assert prediction["score"] is not None - assert estimator.train_image() == image_uri + assert estimator.training_image_uri() == image_uri diff --git a/tests/unit/sagemaker/tensorflow/test_estimator.py b/tests/unit/sagemaker/tensorflow/test_estimator.py index 06c0021a61..1f5a7701ac 100644 --- a/tests/unit/sagemaker/tensorflow/test_estimator.py +++ b/tests/unit/sagemaker/tensorflow/test_estimator.py @@ -521,7 +521,7 @@ def test_hyperparameters_no_model_dir( assert "model_dir" not in hyperparameters -def test_train_image_custom_image(sagemaker_session): +def test_custom_image(sagemaker_session): custom_image = "tensorflow:latest" tf = _build_tf(sagemaker_session, image_uri=custom_image) - assert custom_image == tf.train_image() + assert custom_image == tf.training_image_uri() diff --git a/tests/unit/sagemaker/tensorflow/test_estimator_attach.py b/tests/unit/sagemaker/tensorflow/test_estimator_attach.py index 80d68d7d74..1d263d16c0 100644 --- a/tests/unit/sagemaker/tensorflow/test_estimator_attach.py +++ b/tests/unit/sagemaker/tensorflow/test_estimator_attach.py @@ -98,7 +98,7 @@ def test_attach(sagemaker_session, tensorflow_training_version, tensorflow_train assert estimator.hyperparameters() is not None assert estimator.source_dir == "s3://some/sourcedir.tar.gz" assert estimator.entry_point == "iris-dnn-classifier.py" - assert estimator.train_image() == training_image + assert estimator.training_image_uri() == training_image @patch("sagemaker.utils.create_tar_file", MagicMock()) @@ -207,4 +207,4 @@ def test_attach_custom_image(sagemaker_session): estimator = TensorFlow.attach(training_job_name="neo", sagemaker_session=sagemaker_session) assert estimator.image_uri == training_image - assert estimator.train_image() == training_image + assert estimator.training_image_uri() == training_image diff --git a/tests/unit/test_chainer.py b/tests/unit/test_chainer.py index 7a26d5216e..5679bb5284 100644 --- a/tests/unit/test_chainer.py +++ b/tests/unit/test_chainer.py @@ -415,7 +415,7 @@ def test_model_prepare_container_def_no_instance_type_or_image(chainer_version, assert expected_msg in str(e) -def test_train_image_default(sagemaker_session, chainer_version, chainer_py_version): +def test_training_image_default(sagemaker_session, chainer_version, chainer_py_version): chainer = Chainer( entry_point=SCRIPT_PATH, role=ROLE, @@ -426,7 +426,9 @@ def test_train_image_default(sagemaker_session, chainer_version, chainer_py_vers py_version=chainer_py_version, ) - assert _get_full_cpu_image_uri(chainer_version, chainer_py_version) == chainer.train_image() + assert ( + _get_full_cpu_image_uri(chainer_version, chainer_py_version) == chainer.training_image_uri() + ) def test_attach(sagemaker_session, chainer_version, chainer_py_version): @@ -545,7 +547,7 @@ def test_attach_custom_image(sagemaker_session): estimator = Chainer.attach(training_job_name="neo", sagemaker_session=sagemaker_session) assert estimator.image_uri == training_image - assert estimator.train_image() == training_image + assert estimator.training_image_uri() == training_image @patch("sagemaker.chainer.estimator.python_deprecation_warning") diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py index ecdccdb185..f40fc0aa34 100644 --- a/tests/unit/test_estimator.py +++ b/tests/unit/test_estimator.py @@ -109,7 +109,7 @@ class DummyFramework(Framework): _framework_name = "dummy" - def train_image(self): + def training_image_uri(self): return IMAGE_URI def create_model( diff --git a/tests/unit/test_fm.py b/tests/unit/test_fm.py index 929e8f7e0b..b2f96b1c41 100644 --- a/tests/unit/test_fm.py +++ b/tests/unit/test_fm.py @@ -147,7 +147,7 @@ def test_all_hyperparameters(sagemaker_session): def test_image(sagemaker_session): fm = FactorizationMachines(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) - assert image_uris.retrieve("factorization-machines", REGION) == fm.train_image() + assert image_uris.retrieve("factorization-machines", REGION) == fm.training_image_uri() @pytest.mark.parametrize( diff --git a/tests/unit/test_ipinsights.py b/tests/unit/test_ipinsights.py index 1ad3f6a6cd..0673d375b4 100644 --- a/tests/unit/test_ipinsights.py +++ b/tests/unit/test_ipinsights.py @@ -120,7 +120,7 @@ def test_all_hyperparameters(sagemaker_session): def test_image(sagemaker_session): ipinsights = IPInsights(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) - assert image_uris.retrieve("ipinsights", REGION) == ipinsights.train_image() + assert image_uris.retrieve("ipinsights", REGION) == ipinsights.training_image_uri() @pytest.mark.parametrize( diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 20a0561bb6..bf763b4310 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -83,7 +83,7 @@ def sagemaker_session(): class DummyFramework(Framework): _framework_name = "dummy" - def train_image(self): + def training_image_uri(self): return IMAGE_NAME def create_model(self, role=None, model_server_workers=None): diff --git a/tests/unit/test_kmeans.py b/tests/unit/test_kmeans.py index ea4e6bcb8e..ec5a8476a7 100644 --- a/tests/unit/test_kmeans.py +++ b/tests/unit/test_kmeans.py @@ -112,7 +112,7 @@ def test_all_hyperparameters(sagemaker_session): def test_image(sagemaker_session): kmeans = KMeans(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) - assert image_uris.retrieve("kmeans", REGION) == kmeans.train_image() + assert image_uris.retrieve("kmeans", REGION) == kmeans.training_image_uri() @pytest.mark.parametrize("required_hyper_parameters, value", [("k", "string")]) diff --git a/tests/unit/test_knn.py b/tests/unit/test_knn.py index b22d071833..999de92f5c 100644 --- a/tests/unit/test_knn.py +++ b/tests/unit/test_knn.py @@ -146,7 +146,7 @@ def test_all_hyperparameters_classifier(sagemaker_session): def test_image(sagemaker_session): knn = KNN(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) - assert image_uris.retrieve("knn", REGION) == knn.train_image() + assert image_uris.retrieve("knn", REGION) == knn.training_image_uri() @pytest.mark.parametrize( diff --git a/tests/unit/test_lda.py b/tests/unit/test_lda.py index 467ab0ee60..5e8579c595 100644 --- a/tests/unit/test_lda.py +++ b/tests/unit/test_lda.py @@ -96,7 +96,7 @@ def test_all_hyperparameters(sagemaker_session): def test_image(sagemaker_session): lda = LDA(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) - assert image_uris.retrieve("lda", REGION) == lda.train_image() + assert image_uris.retrieve("lda", REGION) == lda.training_image_uri() @pytest.mark.parametrize("required_hyper_parameters, value", [("num_topics", "string")]) diff --git a/tests/unit/test_linear_learner.py b/tests/unit/test_linear_learner.py index 955bddbef1..4754f6c845 100644 --- a/tests/unit/test_linear_learner.py +++ b/tests/unit/test_linear_learner.py @@ -179,7 +179,7 @@ def test_all_hyperparameters(sagemaker_session): def test_image(sagemaker_session): lr = LinearLearner(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) - assert image_uris.retrieve("linear-learner", REGION) == lr.train_image() + assert image_uris.retrieve("linear-learner", REGION) == lr.training_image_uri() @pytest.mark.parametrize("required_hyper_parameters, value", [("predictor_type", 0)]) diff --git a/tests/unit/test_mxnet.py b/tests/unit/test_mxnet.py index 94a08caac2..069c9f4315 100644 --- a/tests/unit/test_mxnet.py +++ b/tests/unit/test_mxnet.py @@ -670,7 +670,7 @@ def test_attach_custom_image(sagemaker_session): estimator = MXNet.attach(training_job_name="neo", sagemaker_session=sagemaker_session) assert estimator.image_uri == training_image - assert estimator.train_image() == training_image + assert estimator.training_image_uri() == training_image def test_estimator_script_mode_dont_launch_parameter_server(sagemaker_session): diff --git a/tests/unit/test_ntm.py b/tests/unit/test_ntm.py index 5035fbf4cc..6a3aeeabbe 100644 --- a/tests/unit/test_ntm.py +++ b/tests/unit/test_ntm.py @@ -115,7 +115,7 @@ def test_all_hyperparameters(sagemaker_session): def test_image(sagemaker_session): ntm = NTM(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) - assert image_uris.retrieve("ntm", REGION) == ntm.train_image() + assert image_uris.retrieve("ntm", REGION) == ntm.training_image_uri() @pytest.mark.parametrize("required_hyper_parameters, value", [("num_topics", "string")]) diff --git a/tests/unit/test_object2vec.py b/tests/unit/test_object2vec.py index cf5e621587..7a390787ad 100644 --- a/tests/unit/test_object2vec.py +++ b/tests/unit/test_object2vec.py @@ -144,7 +144,7 @@ def test_all_hyperparameters(sagemaker_session): def test_image(sagemaker_session): object2vec = Object2Vec(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) - assert image_uris.retrieve("object2vec", REGION) == object2vec.train_image() + assert image_uris.retrieve("object2vec", REGION) == object2vec.training_image_uri() @pytest.mark.parametrize("required_hyper_parameters, value", [("epochs", "string")]) diff --git a/tests/unit/test_pca.py b/tests/unit/test_pca.py index 43cc5e9705..d6295b5554 100644 --- a/tests/unit/test_pca.py +++ b/tests/unit/test_pca.py @@ -101,7 +101,7 @@ def test_all_hyperparameters(sagemaker_session): def test_image(sagemaker_session): pca = PCA(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) - assert image_uris.retrieve("pca", REGION) == pca.train_image() + assert image_uris.retrieve("pca", REGION) == pca.training_image_uri() @pytest.mark.parametrize("required_hyper_parameters, value", [("num_components", "string")]) diff --git a/tests/unit/test_pytorch.py b/tests/unit/test_pytorch.py index 202b06633b..25c6e0ea11 100644 --- a/tests/unit/test_pytorch.py +++ b/tests/unit/test_pytorch.py @@ -537,7 +537,7 @@ def test_attach_custom_image(sagemaker_session): estimator = PyTorch.attach(training_job_name="neo", sagemaker_session=sagemaker_session) assert estimator.latest_training_job.job_name == "neo" assert estimator.image_uri == training_image - assert estimator.train_image() == training_image + assert estimator.training_image_uri() == training_image @patch("sagemaker.pytorch.estimator.python_deprecation_warning") diff --git a/tests/unit/test_randomcutforest.py b/tests/unit/test_randomcutforest.py index b7e6ef1e49..1b020f2393 100644 --- a/tests/unit/test_randomcutforest.py +++ b/tests/unit/test_randomcutforest.py @@ -107,7 +107,7 @@ def test_all_hyperparameters(sagemaker_session): def test_image(sagemaker_session): randomcutforest = RandomCutForest(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS) - assert image_uris.retrieve("randomcutforest", REGION) == randomcutforest.train_image() + assert image_uris.retrieve("randomcutforest", REGION) == randomcutforest.training_image_uri() @pytest.mark.parametrize("iterable_hyper_parameters, value", [("eval_metrics", 0)]) diff --git a/tests/unit/test_rl.py b/tests/unit/test_rl.py index 73e9e4357b..bfdab596d6 100644 --- a/tests/unit/test_rl.py +++ b/tests/unit/test_rl.py @@ -395,7 +395,7 @@ def test_deploy_ray(sagemaker_session, ray_tensorflow_version): @patch("sagemaker.image_uris.retrieve") -def test_train_image(retrieve_image_uri, sagemaker_session, ray_tensorflow_version): +def test_training_image_uri(retrieve_image_uri, sagemaker_session, ray_tensorflow_version): toolkit = RLToolkit.RAY framework = RLFramework.TENSORFLOW @@ -408,13 +408,13 @@ def test_train_image(retrieve_image_uri, sagemaker_session, ray_tensorflow_versi instance_type=CPU, image_uri=image, ) - assert image == rl.train_image() + assert image == rl.training_image_uri() retrieve_image_uri.assert_not_called() rl = _rl_estimator( sagemaker_session, toolkit, ray_tensorflow_version, framework, instance_type=CPU ) - assert retrieve_image_uri.return_value == rl.train_image() + assert retrieve_image_uri.return_value == rl.training_image_uri() retrieve_image_uri.assert_called_with( "ray-tensorflow", REGION, version=ray_tensorflow_version, instance_type=CPU @@ -540,7 +540,7 @@ def test_attach_custom_image(sagemaker_session): estimator = RLEstimator.attach(training_job_name="neo", sagemaker_session=sagemaker_session) assert estimator.latest_training_job.job_name == "neo" assert estimator.image_uri == training_image - assert estimator.train_image() == training_image + assert estimator.training_image_uri() == training_image def test_wrong_framework_format(sagemaker_session): diff --git a/tests/unit/test_sklearn.py b/tests/unit/test_sklearn.py index c3bcb62426..dc37d599b2 100644 --- a/tests/unit/test_sklearn.py +++ b/tests/unit/test_sklearn.py @@ -140,7 +140,7 @@ def _create_train_job(version): } -def test_train_image(sagemaker_session, sklearn_version): +def test_training_image_uri(sagemaker_session, sklearn_version): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" sklearn = SKLearn( @@ -155,7 +155,7 @@ def test_train_image(sagemaker_session, sklearn_version): source_dir=source_dir, ) - assert _get_full_cpu_image_uri(sklearn_version) == sklearn.train_image() + assert _get_full_cpu_image_uri(sklearn_version) == sklearn.training_image_uri() def test_create_model(sagemaker_session, sklearn_version): @@ -525,7 +525,7 @@ def test_attach_custom_image(sagemaker_session): estimator = SKLearn.attach(training_job_name="neo", sagemaker_session=sagemaker_session) assert estimator.image_uri == training_image - assert estimator.train_image() == training_image + assert estimator.training_image_uri() == training_image def test_estimator_py2_raises(sagemaker_session, sklearn_version): diff --git a/tests/unit/test_tuner.py b/tests/unit/test_tuner.py index 49e1c7ffdd..0e61ec70eb 100644 --- a/tests/unit/test_tuner.py +++ b/tests/unit/test_tuner.py @@ -386,7 +386,7 @@ def test_fit_multi_estimators(sagemaker_session): assert training_config_one["objective_type"] == "Minimize" assert training_config_one["objective_metric_name"] == OBJECTIVE_METRIC_NAME assert training_config_one["input_config"] is None - assert training_config_one["image_uri"] == estimator_one.train_image() + assert training_config_one["image_uri"] == estimator_one.training_image_uri() assert training_config_one["metric_definitions"] == METRIC_DEFINITIONS assert ( training_config_one["static_hyperparameters"]["sagemaker_estimator_module"] @@ -403,7 +403,7 @@ def test_fit_multi_estimators(sagemaker_session): assert training_config_two["objective_metric_name"] == OBJECTIVE_METRIC_NAME_TWO assert len(training_config_two["input_config"]) == 1 assert training_config_two["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] == INPUTS - assert training_config_two["image_uri"] == estimator_two.train_image() + assert training_config_two["image_uri"] == estimator_two.training_image_uri() assert training_config_two["metric_definitions"] is None assert training_config_two["static_hyperparameters"]["mini_batch_size"] == "4000" _assert_parameter_ranges( diff --git a/tests/unit/test_xgboost.py b/tests/unit/test_xgboost.py index 42c47ccbbe..0c0b64f2b9 100644 --- a/tests/unit/test_xgboost.py +++ b/tests/unit/test_xgboost.py @@ -380,7 +380,7 @@ def test_model(sagemaker_session, xgboost_framework_version): assert isinstance(predictor, XGBoostPredictor) -def test_train_image(sagemaker_session, xgboost_framework_version): +def test_training_image_uri(sagemaker_session, xgboost_framework_version): xgboost = XGBoost( entry_point=SCRIPT_PATH, role=ROLE, @@ -391,7 +391,7 @@ def test_train_image(sagemaker_session, xgboost_framework_version): py_version=PYTHON_VERSION, ) - assert _get_full_cpu_image_uri(xgboost_framework_version) in xgboost.train_image() + assert _get_full_cpu_image_uri(xgboost_framework_version) in xgboost.training_image_uri() def test_attach(sagemaker_session, xgboost_framework_version):