diff --git a/google/cloud/aiplatform/docker_utils/run.py b/google/cloud/aiplatform/docker_utils/run.py index 4b42ab3643..f7336eb7c2 100644 --- a/google/cloud/aiplatform/docker_utils/run.py +++ b/google/cloud/aiplatform/docker_utils/run.py @@ -138,10 +138,10 @@ def run_prediction_container( raise ValueError(f'artifact_uri must be a GCS path but it is "{artifact_uri}".') envs[prediction.AIP_STORAGE_URI] = artifact_uri if artifact_uri is not None else "" - command = ( + entrypoint = ( serving_container_command[:] if serving_container_command is not None else [] ) - command += serving_container_args if serving_container_args is not None else [] + command = serving_container_args[:] if serving_container_args is not None else [] credential_from_adc_env = credential_path is None credential_path = credential_path or _get_adc_environment_variable() @@ -165,6 +165,7 @@ def run_prediction_container( container = client.containers.run( serving_container_image_uri, command=command if len(command) > 0 else None, + entrypoint=entrypoint if len(entrypoint) > 0 else None, ports={port: host_port}, environment=envs, volumes=volumes, diff --git a/tests/unit/aiplatform/test_docker_utils.py b/tests/unit/aiplatform/test_docker_utils.py index 26a99c47e3..140fa1fe45 100644 --- a/tests/unit/aiplatform/test_docker_utils.py +++ b/tests/unit/aiplatform/test_docker_utils.py @@ -95,6 +95,7 @@ def test_run_prediction_container(self, docker_client_mock): docker_client_mock.containers.run.assert_called_once_with( self.IMAGE_URI, command=None, + entrypoint=None, ports={prediction.DEFAULT_AIP_HTTP_PORT: None}, environment={ prediction.AIP_HTTP_PORT: prediction.DEFAULT_AIP_HTTP_PORT, @@ -145,7 +146,8 @@ def test_run_prediction_container_with_all_parameters( docker_client_mock.containers.run.assert_called_once_with( self.IMAGE_URI, - command=serving_container_command + serving_container_args, + command=serving_container_args, + entrypoint=serving_container_command, ports={serving_container_ports[0]: host_port}, environment=environment, volumes=volumes, @@ -167,6 +169,7 @@ def test_run_prediction_container_credential_from_adc( docker_client_mock.containers.run.assert_called_once_with( self.IMAGE_URI, command=None, + entrypoint=None, ports={prediction.DEFAULT_AIP_HTTP_PORT: None}, environment={ prediction.AIP_HTTP_PORT: prediction.DEFAULT_AIP_HTTP_PORT,