We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
InfraValidator fails for on-prem deploys
The InfraValidator component has no option for specifying from which docker registry to pull the TF serving image. For on-prem deployments (or other isolated k8s clusters) this means that InfraValidator fails with an ImagePullBackError.

Proposed solution: Add option to specify (one) image (not just registry) along with the existing tags and digests options. The user is responsible for ensuring that the image is a working TF serving image.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tfx/components/infra_validator/serving_bins.py`
Content:
```
1 # Lint as: python2, python3
2 # Copyright 2020 Google LLC. All Rights Reserved.
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License");
5 # you may not use this file except in compliance with the License.
6 # You may obtain a copy of the License at
7 #
8 #     http://www.apache.org/licenses/LICENSE-2.0
9 #
10 # Unless required by applicable law or agreed to in writing, software
11 # distributed under the License is distributed on an "AS IS" BASIS,
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 # See the License for the specific language governing permissions and
14 # limitations under the License.
15 """Modules for organizing various model server binaries."""
16 
17 from __future__ import absolute_import
18 from __future__ import division
19 from __future__ import print_function
20 
21 import abc
22 import os
23 from typing import Any, Dict, List, Optional, Text
24 
25 from docker import types as docker_types
26 import six
27 
28 from tfx.components.infra_validator.model_server_clients import base_client
29 from tfx.components.infra_validator.model_server_clients import tensorflow_serving_client
30 from tfx.proto import infra_validator_pb2
31 from tfx.utils.model_paths import tf_serving_flavor
32 
33 
34 def parse_serving_binaries(  # pylint: disable=invalid-name
35     serving_spec: infra_validator_pb2.ServingSpec) -> List['ServingBinary']:
36   """Parse `ServingBinary`s from `ServingSpec`."""
37   result = []
38   serving_binary = serving_spec.WhichOneof('serving_binary')
39   if serving_binary == 'tensorflow_serving':
40     config = serving_spec.tensorflow_serving
41     for tag in config.tags:
42       result.append(TensorFlowServing(model_name=serving_spec.model_name,
43                                       tag=tag))
44     for digest in config.digests:
45       result.append(TensorFlowServing(model_name=serving_spec.model_name,
46                                       digest=digest))
47     return result
48   else:
49     raise ValueError('Invalid serving_binary {}'.format(serving_binary))
50 
51 
52 class ServingBinary(six.with_metaclass(abc.ABCMeta, object)):
53   """Base class for serving binaries."""
54 
55   @abc.abstractproperty
56   def container_port(self) -> int:
57     """Container port of the model server.
58 
59     Only applies to docker compatible serving binaries.
60     """
61     raise NotImplementedError('{} is not docker compatible.'.format(
62         type(self).__name__))
63 
64   @abc.abstractproperty
65   def image(self) -> Text:
66     """Container image of the model server.
67 
68     Only applies to docker compatible serving binaries.
69     """
70     raise NotImplementedError('{} is not docker compatible.'.format(
71         type(self).__name__))
72 
73   @abc.abstractmethod
74   def MakeEnvVars(self, *args: Any) -> Dict[Text, Text]:
75     """Construct environment variables to be used in container image.
76 
77     Only applies to docker compatible serving binaries.
78 
79     Args:
80       *args: List of unresolved variables to configure environment variables.
81 
82     Returns:
83       A dictionary of environment variables inside container.
84     """
85     raise NotImplementedError('{} is not docker compatible.'.format(
86         type(self).__name__))
87 
88   @abc.abstractmethod
89   def MakeDockerRunParams(self, *args: Any) -> Dict[Text, Text]:
90     """Make parameters for docker `client.containers.run`.
91 
92     Only applies to docker compatible serving binaries.
93 
94     Args:
95       *args: List of unresolved variables to configure docker run parameters.
96 
97     Returns:
98       A dictionary of docker run parameters.
99     """
100     raise NotImplementedError('{} is not docker compatible.'.format(
101         type(self).__name__))
102 
103   @abc.abstractmethod
104   def MakeClient(self, endpoint: Text) -> base_client.BaseModelServerClient:
105     """Create a model server client of this serving binary."""
106     raise NotImplementedError('{} does not implement MakeClient.'.format(
107         type(self).__name__))
108 
109 
110 class TensorFlowServing(ServingBinary):
111   """TensorFlow Serving binary."""
112 
113   _BASE_DOCKER_RUN_PARAMS = {
114       # Enable auto-removal of the container on docker daemon after container
115       # process exits.
116       'auto_remove': True,
117       # Run container in the background instead of streaming its output.
118       'detach': True,
119       # Publish all ports to the host.
120       'publish_all_ports': True,
121   }
122   _DEFAULT_IMAGE_NAME = 'tensorflow/serving'
123   _DEFAULT_GRPC_PORT = 8500
124   _DEFAULT_MODEL_BASE_PATH = '/model'
125 
126   def __init__(
127       self,
128       model_name: Text,
129       tag: Optional[Text] = None,
130       digest: Optional[Text] = None,
131   ):
132     super(TensorFlowServing, self).__init__()
133     self._model_name = model_name
134     if (tag is None) == (digest is None):
135       raise ValueError('Exactly one of `tag` or `digest` should be used.')
136     if tag is not None:
137       self._image = '{}:{}'.format(self._DEFAULT_IMAGE_NAME, tag)
138     else:
139       self._image = '{}@{}'.format(self._DEFAULT_IMAGE_NAME, digest)
140 
141   @property
142   def container_port(self) -> int:
143     return self._DEFAULT_GRPC_PORT
144 
145   @property
146   def image(self) -> Text:
147     return self._image
148 
149   def MakeEnvVars(
150       self, model_path: Optional[Text] = None) -> Dict[Text, Text]:
151     if model_path is None:
152       model_base_path = self._DEFAULT_MODEL_BASE_PATH
153     else:
154       model_base_path = tf_serving_flavor.parse_model_base_path(model_path)
155     return {
156         'MODEL_NAME': self._model_name,
157         'MODEL_BASE_PATH': model_base_path
158     }
159 
160   def MakeDockerRunParams(
161       self,
162       model_path: Text,
163       needs_mount: bool) -> Dict[Text, Any]:
164     """Make parameters for docker `client.containers.run`.
165 
166     Args:
167       model_path: A path to the model.
168       needs_mount: If True, model_path will be mounted to the container.
169 
170     Returns:
171       A dictionary of docker run parameters.
172     """
173     result = dict(
174         self._BASE_DOCKER_RUN_PARAMS,
175         image=self._image)
176 
177     if needs_mount:
178       # model_path should be a local directory. In order to make TF Serving see
179       # the host model path, we need to mount model path volume to the
180       # container.
181       assert os.path.isdir(model_path), '{} does not exist'.format(model_path)
182       container_model_path = tf_serving_flavor.make_model_path(
183           model_base_path=self._DEFAULT_MODEL_BASE_PATH,
184           model_name=self._model_name,
185           version=1)
186       result.update(
187           environment=self.MakeEnvVars(),
188           mounts=[
189               docker_types.Mount(
190                   type='bind',
191                   target=container_model_path,
192                   source=model_path,
193                   read_only=True)
194           ])
195     else:
196       # model_path is presumably a remote URI. TF Serving is able to pickup
197       # model in remote directly using gfile, so all we need to do is setting
198       # environment variables correctly.
199       result.update(
200           environment=self.MakeEnvVars(model_path=model_path))
201 
202     return result
203 
204   def MakeClient(self, endpoint: Text) -> base_client.BaseModelServerClient:
205     return tensorflow_serving_client.TensorFlowServingClient(
206         endpoint=endpoint, model_name=self._model_name)
207
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/tfx/components/infra_validator/serving_bins.py b/tfx/components/infra_validator/serving_bins.py
--- a/tfx/components/infra_validator/serving_bins.py
+++ b/tfx/components/infra_validator/serving_bins.py
@@ -38,11 +38,14 @@
   serving_binary = serving_spec.WhichOneof('serving_binary')
   if serving_binary == 'tensorflow_serving':
     config = serving_spec.tensorflow_serving
+    image_name = config.image_name or None
     for tag in config.tags:
-      result.append(TensorFlowServing(model_name=serving_spec.model_name,
+      result.append(TensorFlowServing(image_name=image_name,
+                                      model_name=serving_spec.model_name,
                                       tag=tag))
     for digest in config.digests:
-      result.append(TensorFlowServing(model_name=serving_spec.model_name,
+      result.append(TensorFlowServing(image_name=image_name,
+                                      model_name=serving_spec.model_name,
                                       digest=digest))
     return result
   else:
@@ -126,6 +129,7 @@
   def __init__(
       self,
       model_name: Text,
+      image_name: Optional[Text] = None,
       tag: Optional[Text] = None,
       digest: Optional[Text] = None,
   ):
@@ -133,10 +137,11 @@
     self._model_name = model_name
     if (tag is None) == (digest is None):
       raise ValueError('Exactly one of `tag` or `digest` should be used.')
+    image_name = image_name or self._DEFAULT_IMAGE_NAME
     if tag is not None:
-      self._image = '{}:{}'.format(self._DEFAULT_IMAGE_NAME, tag)
+      self._image = '{}:{}'.format(image_name, tag)
     else:
-      self._image = '{}@{}'.format(self._DEFAULT_IMAGE_NAME, digest)
+      self._image = '{}@{}'.format(image_name, digest)
 
   @property
   def container_port(self) -> int:

{"golden_diff": "diff --git a/tfx/components/infra_validator/serving_bins.py b/tfx/components/infra_validator/serving_bins.py\n--- a/tfx/components/infra_validator/serving_bins.py\n+++ b/tfx/components/infra_validator/serving_bins.py\n@@ -38,11 +38,14 @@\n   serving_binary = serving_spec.WhichOneof('serving_binary')\n   if serving_binary == 'tensorflow_serving':\n     config = serving_spec.tensorflow_serving\n+    image_name = config.image_name or None\n     for tag in config.tags:\n-      result.append(TensorFlowServing(model_name=serving_spec.model_name,\n+      result.append(TensorFlowServing(image_name=image_name,\n+                                      model_name=serving_spec.model_name,\n                                       tag=tag))\n     for digest in config.digests:\n-      result.append(TensorFlowServing(model_name=serving_spec.model_name,\n+      result.append(TensorFlowServing(image_name=image_name,\n+                                      model_name=serving_spec.model_name,\n                                       digest=digest))\n     return result\n   else:\n@@ -126,6 +129,7 @@\n   def __init__(\n       self,\n       model_name: Text,\n+      image_name: Optional[Text] = None,\n       tag: Optional[Text] = None,\n       digest: Optional[Text] = None,\n   ):\n@@ -133,10 +137,11 @@\n     self._model_name = model_name\n     if (tag is None) == (digest is None):\n       raise ValueError('Exactly one of `tag` or `digest` should be used.')\n+    image_name = image_name or self._DEFAULT_IMAGE_NAME\n     if tag is not None:\n-      self._image = '{}:{}'.format(self._DEFAULT_IMAGE_NAME, tag)\n+      self._image = '{}:{}'.format(image_name, tag)\n     else:\n-      self._image = '{}@{}'.format(self._DEFAULT_IMAGE_NAME, digest)\n+      self._image = '{}@{}'.format(image_name, digest)\n \n   @property\n   def container_port(self) -> int:\n", "issue": "InfraValidator fails for on-prem deploys\nThe InfraValidator component has no option for specifying from which docker registry to pull the TF serving image. For on-prem deployments (or other isolated k8s clusters) this means that InfraValidator fails with an ImagePullBackError.\r\n\r\nProposed solution: Add option to specify (one) image (not just registry) along with the existing tags and digests options. The user is responsible for ensuring that the image is a working TF serving image.\n", "before_files": [{"content": "# Lint as: python2, python3\n# Copyright 2020 Google LLC. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Modules for organizing various model server binaries.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport abc\nimport os\nfrom typing import Any, Dict, List, Optional, Text\n\nfrom docker import types as docker_types\nimport six\n\nfrom tfx.components.infra_validator.model_server_clients import base_client\nfrom tfx.components.infra_validator.model_server_clients import tensorflow_serving_client\nfrom tfx.proto import infra_validator_pb2\nfrom tfx.utils.model_paths import tf_serving_flavor\n\n\ndef parse_serving_binaries(  # pylint: disable=invalid-name\n    serving_spec: infra_validator_pb2.ServingSpec) -> List['ServingBinary']:\n  \"\"\"Parse `ServingBinary`s from `ServingSpec`.\"\"\"\n  result = []\n  serving_binary = serving_spec.WhichOneof('serving_binary')\n  if serving_binary == 'tensorflow_serving':\n    config = serving_spec.tensorflow_serving\n    for tag in config.tags:\n      result.append(TensorFlowServing(model_name=serving_spec.model_name,\n                                      tag=tag))\n    for digest in config.digests:\n      result.append(TensorFlowServing(model_name=serving_spec.model_name,\n                                      digest=digest))\n    return result\n  else:\n    raise ValueError('Invalid serving_binary {}'.format(serving_binary))\n\n\nclass ServingBinary(six.with_metaclass(abc.ABCMeta, object)):\n  \"\"\"Base class for serving binaries.\"\"\"\n\n  @abc.abstractproperty\n  def container_port(self) -> int:\n    \"\"\"Container port of the model server.\n\n    Only applies to docker compatible serving binaries.\n    \"\"\"\n    raise NotImplementedError('{} is not docker compatible.'.format(\n        type(self).__name__))\n\n  @abc.abstractproperty\n  def image(self) -> Text:\n    \"\"\"Container image of the model server.\n\n    Only applies to docker compatible serving binaries.\n    \"\"\"\n    raise NotImplementedError('{} is not docker compatible.'.format(\n        type(self).__name__))\n\n  @abc.abstractmethod\n  def MakeEnvVars(self, *args: Any) -> Dict[Text, Text]:\n    \"\"\"Construct environment variables to be used in container image.\n\n    Only applies to docker compatible serving binaries.\n\n    Args:\n      *args: List of unresolved variables to configure environment variables.\n\n    Returns:\n      A dictionary of environment variables inside container.\n    \"\"\"\n    raise NotImplementedError('{} is not docker compatible.'.format(\n        type(self).__name__))\n\n  @abc.abstractmethod\n  def MakeDockerRunParams(self, *args: Any) -> Dict[Text, Text]:\n    \"\"\"Make parameters for docker `client.containers.run`.\n\n    Only applies to docker compatible serving binaries.\n\n    Args:\n      *args: List of unresolved variables to configure docker run parameters.\n\n    Returns:\n      A dictionary of docker run parameters.\n    \"\"\"\n    raise NotImplementedError('{} is not docker compatible.'.format(\n        type(self).__name__))\n\n  @abc.abstractmethod\n  def MakeClient(self, endpoint: Text) -> base_client.BaseModelServerClient:\n    \"\"\"Create a model server client of this serving binary.\"\"\"\n    raise NotImplementedError('{} does not implement MakeClient.'.format(\n        type(self).__name__))\n\n\nclass TensorFlowServing(ServingBinary):\n  \"\"\"TensorFlow Serving binary.\"\"\"\n\n  _BASE_DOCKER_RUN_PARAMS = {\n      # Enable auto-removal of the container on docker daemon after container\n      # process exits.\n      'auto_remove': True,\n      # Run container in the background instead of streaming its output.\n      'detach': True,\n      # Publish all ports to the host.\n      'publish_all_ports': True,\n  }\n  _DEFAULT_IMAGE_NAME = 'tensorflow/serving'\n  _DEFAULT_GRPC_PORT = 8500\n  _DEFAULT_MODEL_BASE_PATH = '/model'\n\n  def __init__(\n      self,\n      model_name: Text,\n      tag: Optional[Text] = None,\n      digest: Optional[Text] = None,\n  ):\n    super(TensorFlowServing, self).__init__()\n    self._model_name = model_name\n    if (tag is None) == (digest is None):\n      raise ValueError('Exactly one of `tag` or `digest` should be used.')\n    if tag is not None:\n      self._image = '{}:{}'.format(self._DEFAULT_IMAGE_NAME, tag)\n    else:\n      self._image = '{}@{}'.format(self._DEFAULT_IMAGE_NAME, digest)\n\n  @property\n  def container_port(self) -> int:\n    return self._DEFAULT_GRPC_PORT\n\n  @property\n  def image(self) -> Text:\n    return self._image\n\n  def MakeEnvVars(\n      self, model_path: Optional[Text] = None) -> Dict[Text, Text]:\n    if model_path is None:\n      model_base_path = self._DEFAULT_MODEL_BASE_PATH\n    else:\n      model_base_path = tf_serving_flavor.parse_model_base_path(model_path)\n    return {\n        'MODEL_NAME': self._model_name,\n        'MODEL_BASE_PATH': model_base_path\n    }\n\n  def MakeDockerRunParams(\n      self,\n      model_path: Text,\n      needs_mount: bool) -> Dict[Text, Any]:\n    \"\"\"Make parameters for docker `client.containers.run`.\n\n    Args:\n      model_path: A path to the model.\n      needs_mount: If True, model_path will be mounted to the container.\n\n    Returns:\n      A dictionary of docker run parameters.\n    \"\"\"\n    result = dict(\n        self._BASE_DOCKER_RUN_PARAMS,\n        image=self._image)\n\n    if needs_mount:\n      # model_path should be a local directory. In order to make TF Serving see\n      # the host model path, we need to mount model path volume to the\n      # container.\n      assert os.path.isdir(model_path), '{} does not exist'.format(model_path)\n      container_model_path = tf_serving_flavor.make_model_path(\n          model_base_path=self._DEFAULT_MODEL_BASE_PATH,\n          model_name=self._model_name,\n          version=1)\n      result.update(\n          environment=self.MakeEnvVars(),\n          mounts=[\n              docker_types.Mount(\n                  type='bind',\n                  target=container_model_path,\n                  source=model_path,\n                  read_only=True)\n          ])\n    else:\n      # model_path is presumably a remote URI. TF Serving is able to pickup\n      # model in remote directly using gfile, so all we need to do is setting\n      # environment variables correctly.\n      result.update(\n          environment=self.MakeEnvVars(model_path=model_path))\n\n    return result\n\n  def MakeClient(self, endpoint: Text) -> base_client.BaseModelServerClient:\n    return tensorflow_serving_client.TensorFlowServingClient(\n        endpoint=endpoint, model_name=self._model_name)\n", "path": "tfx/components/infra_validator/serving_bins.py"}], "after_files": [{"content": "# Lint as: python2, python3\n# Copyright 2020 Google LLC. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Modules for organizing various model server binaries.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport abc\nimport os\nfrom typing import Any, Dict, List, Optional, Text\n\nfrom docker import types as docker_types\nimport six\n\nfrom tfx.components.infra_validator.model_server_clients import base_client\nfrom tfx.components.infra_validator.model_server_clients import tensorflow_serving_client\nfrom tfx.proto import infra_validator_pb2\nfrom tfx.utils.model_paths import tf_serving_flavor\n\n\ndef parse_serving_binaries(  # pylint: disable=invalid-name\n    serving_spec: infra_validator_pb2.ServingSpec) -> List['ServingBinary']:\n  \"\"\"Parse `ServingBinary`s from `ServingSpec`.\"\"\"\n  result = []\n  serving_binary = serving_spec.WhichOneof('serving_binary')\n  if serving_binary == 'tensorflow_serving':\n    config = serving_spec.tensorflow_serving\n    image_name = config.image_name or None\n    for tag in config.tags:\n      result.append(TensorFlowServing(image_name=image_name,\n                                      model_name=serving_spec.model_name,\n                                      tag=tag))\n    for digest in config.digests:\n      result.append(TensorFlowServing(image_name=image_name,\n                                      model_name=serving_spec.model_name,\n                                      digest=digest))\n    return result\n  else:\n    raise ValueError('Invalid serving_binary {}'.format(serving_binary))\n\n\nclass ServingBinary(six.with_metaclass(abc.ABCMeta, object)):\n  \"\"\"Base class for serving binaries.\"\"\"\n\n  @abc.abstractproperty\n  def container_port(self) -> int:\n    \"\"\"Container port of the model server.\n\n    Only applies to docker compatible serving binaries.\n    \"\"\"\n    raise NotImplementedError('{} is not docker compatible.'.format(\n        type(self).__name__))\n\n  @abc.abstractproperty\n  def image(self) -> Text:\n    \"\"\"Container image of the model server.\n\n    Only applies to docker compatible serving binaries.\n    \"\"\"\n    raise NotImplementedError('{} is not docker compatible.'.format(\n        type(self).__name__))\n\n  @abc.abstractmethod\n  def MakeEnvVars(self, *args: Any) -> Dict[Text, Text]:\n    \"\"\"Construct environment variables to be used in container image.\n\n    Only applies to docker compatible serving binaries.\n\n    Args:\n      *args: List of unresolved variables to configure environment variables.\n\n    Returns:\n      A dictionary of environment variables inside container.\n    \"\"\"\n    raise NotImplementedError('{} is not docker compatible.'.format(\n        type(self).__name__))\n\n  @abc.abstractmethod\n  def MakeDockerRunParams(self, *args: Any) -> Dict[Text, Text]:\n    \"\"\"Make parameters for docker `client.containers.run`.\n\n    Only applies to docker compatible serving binaries.\n\n    Args:\n      *args: List of unresolved variables to configure docker run parameters.\n\n    Returns:\n      A dictionary of docker run parameters.\n    \"\"\"\n    raise NotImplementedError('{} is not docker compatible.'.format(\n        type(self).__name__))\n\n  @abc.abstractmethod\n  def MakeClient(self, endpoint: Text) -> base_client.BaseModelServerClient:\n    \"\"\"Create a model server client of this serving binary.\"\"\"\n    raise NotImplementedError('{} does not implement MakeClient.'.format(\n        type(self).__name__))\n\n\nclass TensorFlowServing(ServingBinary):\n  \"\"\"TensorFlow Serving binary.\"\"\"\n\n  _BASE_DOCKER_RUN_PARAMS = {\n      # Enable auto-removal of the container on docker daemon after container\n      # process exits.\n      'auto_remove': True,\n      # Run container in the background instead of streaming its output.\n      'detach': True,\n      # Publish all ports to the host.\n      'publish_all_ports': True,\n  }\n  _DEFAULT_IMAGE_NAME = 'tensorflow/serving'\n  _DEFAULT_GRPC_PORT = 8500\n  _DEFAULT_MODEL_BASE_PATH = '/model'\n\n  def __init__(\n      self,\n      model_name: Text,\n      image_name: Optional[Text] = None,\n      tag: Optional[Text] = None,\n      digest: Optional[Text] = None,\n  ):\n    super(TensorFlowServing, self).__init__()\n    self._model_name = model_name\n    if (tag is None) == (digest is None):\n      raise ValueError('Exactly one of `tag` or `digest` should be used.')\n    image_name = image_name or self._DEFAULT_IMAGE_NAME\n    if tag is not None:\n      self._image = '{}:{}'.format(image_name, tag)\n    else:\n      self._image = '{}@{}'.format(image_name, digest)\n\n  @property\n  def container_port(self) -> int:\n    return self._DEFAULT_GRPC_PORT\n\n  @property\n  def image(self) -> Text:\n    return self._image\n\n  def MakeEnvVars(\n      self, model_path: Optional[Text] = None) -> Dict[Text, Text]:\n    if model_path is None:\n      model_base_path = self._DEFAULT_MODEL_BASE_PATH\n    else:\n      model_base_path = tf_serving_flavor.parse_model_base_path(model_path)\n    return {\n        'MODEL_NAME': self._model_name,\n        'MODEL_BASE_PATH': model_base_path\n    }\n\n  def MakeDockerRunParams(\n      self,\n      model_path: Text,\n      needs_mount: bool) -> Dict[Text, Any]:\n    \"\"\"Make parameters for docker `client.containers.run`.\n\n    Args:\n      model_path: A path to the model.\n      needs_mount: If True, model_path will be mounted to the container.\n\n    Returns:\n      A dictionary of docker run parameters.\n    \"\"\"\n    result = dict(\n        self._BASE_DOCKER_RUN_PARAMS,\n        image=self._image)\n\n    if needs_mount:\n      # model_path should be a local directory. In order to make TF Serving see\n      # the host model path, we need to mount model path volume to the\n      # container.\n      assert os.path.isdir(model_path), '{} does not exist'.format(model_path)\n      container_model_path = tf_serving_flavor.make_model_path(\n          model_base_path=self._DEFAULT_MODEL_BASE_PATH,\n          model_name=self._model_name,\n          version=1)\n      result.update(\n          environment=self.MakeEnvVars(),\n          mounts=[\n              docker_types.Mount(\n                  type='bind',\n                  target=container_model_path,\n                  source=model_path,\n                  read_only=True)\n          ])\n    else:\n      # model_path is presumably a remote URI. TF Serving is able to pickup\n      # model in remote directly using gfile, so all we need to do is setting\n      # environment variables correctly.\n      result.update(\n          environment=self.MakeEnvVars(model_path=model_path))\n\n    return result\n\n  def MakeClient(self, endpoint: Text) -> base_client.BaseModelServerClient:\n    return tensorflow_serving_client.TensorFlowServingClient(\n        endpoint=endpoint, model_name=self._model_name)\n", "path": "tfx/components/infra_validator/serving_bins.py"}]}