primary_container_image_uri/deployment_image_uri --> image_uri

Author: laurenyu

src/sagemaker/session.py

@@ -2156,7 +2156,7 @@ def create_model_from_job(
         training_job_name,
         name=None,
         role=None,
-        primary_container_image_uri=None,
+        image_uri=None,
         model_data_url=None,
         env=None,
         vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT,
@@ -2171,7 +2171,7 @@ def create_model_from_job(
             role (str): The ``ExecutionRoleArn`` IAM Role ARN for the ``Model``, specified either
                 by an IAM role name or role ARN. If None, the ``RoleArn`` from the SageMaker
                 Training Job will be used.
-            primary_container_image_uri (str): The Docker image URI (default: None). If None, it
+            image_uri (str): The Docker image URI (default: None). If None, it
                 defaults to the training image URI from ``training_job_name``.
             model_data_url (str): S3 location of the model data (default: None). If None, defaults
                 to the ``ModelS3Artifacts`` of ``training_job_name``.
@@ -2194,7 +2194,7 @@ def create_model_from_job(
         role = role or training_job["RoleArn"]
         env = env or {}
         primary_container = container_def(
-            primary_container_image_uri or training_job["AlgorithmSpecification"]["TrainingImage"],
+            image_uri or training_job["AlgorithmSpecification"]["TrainingImage"],
             model_data_url=model_data_url or training_job["ModelArtifacts"]["S3ModelArtifacts"],
             env=env,
         )
@@ -2700,7 +2700,7 @@ def endpoint_from_job(
         job_name,
         initial_instance_count,
         instance_type,
-        deployment_image_uri=None,
+        image_uri=None,
         name=None,
         role=None,
         wait=True,
@@ -2725,7 +2725,7 @@ def endpoint_from_job(
                 autoscaling.
             instance_type (str): Type of EC2 instance to deploy to an endpoint for prediction,
                 for example, 'ml.c4.xlarge'.
-            deployment_image_uri (str): The Docker image which defines the inference code to be used
+            image_uri (str): The Docker image which defines the inference code to be used
                 as the entry point for accepting prediction requests. If not specified, uses the
                 image used for the training job.
             name (str): Name of the ``Endpoint`` to create. If not specified, uses the training job
@@ -2755,16 +2755,14 @@ def endpoint_from_job(
         """
         job_desc = self.sagemaker_client.describe_training_job(TrainingJobName=job_name)
         output_url = job_desc["ModelArtifacts"]["S3ModelArtifacts"]
-        deployment_image_uri = (
-            deployment_image_uri or job_desc["AlgorithmSpecification"]["TrainingImage"]
-        )
+        image_uri = image_uri or job_desc["AlgorithmSpecification"]["TrainingImage"]
         role = role or job_desc["RoleArn"]
         name = name or job_name
         vpc_config_override = _vpc_config_from_training_job(job_desc, vpc_config_override)
 
         return self.endpoint_from_model_data(
             model_s3_location=output_url,
-            deployment_image_uri=deployment_image_uri,
+            image_uri=image_uri,
             initial_instance_count=initial_instance_count,
             instance_type=instance_type,
             name=name,

tests/unit/test_endpoint_from_job.py

@@ -67,7 +67,7 @@ def test_all_defaults_no_existing_entities(sagemaker_session):
     expected_args = original_args.copy()
     expected_args.pop("job_name")
     expected_args["model_s3_location"] = S3_MODEL_ARTIFACTS
-    expected_args["deployment_image_uri"] = IMAGE
+    expected_args["image_uri"] = IMAGE
     expected_args["role"] = TRAIN_ROLE
     expected_args["name"] = JOB_NAME
     expected_args["model_environment_vars"] = None
@@ -85,7 +85,7 @@ def test_no_defaults_no_existing_entities(sagemaker_session):
         "job_name": JOB_NAME,
         "initial_instance_count": INITIAL_INSTANCE_COUNT,
         "instance_type": INSTANCE_TYPE,
-        "deployment_image_uri": DEPLOY_IMAGE,
+        "image_uri": DEPLOY_IMAGE,
         "role": DEPLOY_ROLE,
         "name": NEW_ENTITY_NAME,
         "model_environment_vars": ENV_VARS,

tests/unit/test_endpoint_from_model_data.py

@@ -56,7 +56,7 @@ def sagemaker_session():
 def test_all_defaults_no_existing_entities(name_from_image_mock, sagemaker_session):
     returned_name = sagemaker_session.endpoint_from_model_data(
         model_s3_location=S3_MODEL_ARTIFACTS,
-        deployment_image_uri=DEPLOY_IMAGE,
+        image_uri=DEPLOY_IMAGE,
         initial_instance_count=INITIAL_INSTANCE_COUNT,
         instance_type=INSTANCE_TYPE,
         role=DEPLOY_ROLE,
@@ -96,7 +96,7 @@ def test_no_defaults_no_existing_entities(name_from_image_mock, sagemaker_sessio
 
     returned_name = sagemaker_session.endpoint_from_model_data(
         model_s3_location=S3_MODEL_ARTIFACTS,
-        deployment_image_uri=DEPLOY_IMAGE,
+        image_uri=DEPLOY_IMAGE,
         initial_instance_count=INITIAL_INSTANCE_COUNT,
         instance_type=INSTANCE_TYPE,
         role=DEPLOY_ROLE,
@@ -145,7 +145,7 @@ def test_model_and_endpoint_config_exist(name_from_image_mock, sagemaker_session
 
     sagemaker_session.endpoint_from_model_data(
         model_s3_location=S3_MODEL_ARTIFACTS,
-        deployment_image_uri=DEPLOY_IMAGE,
+        image_uri=DEPLOY_IMAGE,
         initial_instance_count=INITIAL_INSTANCE_COUNT,
         instance_type=INSTANCE_TYPE,
         wait=False,

tests/unit/test_session.py

@@ -1697,7 +1697,7 @@ def test_create_model_from_job_with_tags(sagemaker_session):
 def test_create_model_from_job_with_image(sagemaker_session):
     ims = sagemaker_session
     ims.sagemaker_client.describe_training_job.return_value = COMPLETED_DESCRIBE_JOB_RESULT
-    ims.create_model_from_job(JOB_NAME, primary_container_image_uri="some-image")
+    ims.create_model_from_job(JOB_NAME, image_uri="some-image")
     [create_model_call] = ims.sagemaker_client.create_model.call_args_list
     assert dict(create_model_call[1]["PrimaryContainer"])["Image"] == "some-image"
 
@@ -1706,10 +1706,7 @@ def test_create_model_from_job_with_container_def(sagemaker_session):
     ims = sagemaker_session
     ims.sagemaker_client.describe_training_job.return_value = COMPLETED_DESCRIBE_JOB_RESULT
     ims.create_model_from_job(
-        JOB_NAME,
-        primary_container_image_uri="some-image",
-        model_data_url="some-data",
-        env={"a": "b"},
+        JOB_NAME, image_uri="some-image", model_data_url="some-data", env={"a": "b"},
     )
     [create_model_call] = ims.sagemaker_client.create_model.call_args_list
     c_def = create_model_call[1]["PrimaryContainer"]