Repository: tensorflow/tfx

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.


Discussion:
Thanks ConverJens. Supporting custom docker images (based on official tensorflow/serving or equivalents) is definitely in our TODO list. Your suggestion would help us a lot :) I'll keep this issue for tracking the feature development.
@chongkong Sounds great! Are you actively working on this? Otherwise I could create a pull request and give it a try. From my looking through the code, the changes needed seems to be minor. What do you think?
Hi ConverJens. I'm back to this issue :) It can be as simple as adding a single line in the config, but I'm also curious on the requirements for the authentication on fetching private images. Would you share your use case on this?

Also would you tell us how the custom image is different from the public TF serving image? For example is the model baked into the image?
Hi @chongkong.

Of course. We are running TFX pipelines in an on-prem KubeFlow cluster. For security reasons, the k8s cluster is blocked off from internet access and hence, we can't pull the images directly from the external source (GCR). We rely on pulling images in our build system and then pushing them to our own internal registry, were we can also change image names and tags if needed.

What we need is the possibility to specify the image name and registry path as well as the tag. At the moment, we have no need for any authentication options.

As of now we use the vanilla TFServing images so there are no changes to the public one.

Let me know if you have any other questions and how to proceed on this.

@ConverJens 

> Hi @chongkong.
> 
> Of course. We are running TFX pipelines in an on-prem KubeFlow cluster. For security reasons, the k8s cluster is blocked off from internet access and hence, we can't pull the images directly from the external source (GCR). We rely on pulling images in our build system and then pushing them to our own internal registry, were we can also change image names and tags if needed.
> 
> What we need is the possibility to specify the image name and registry path as well as the tag. At the moment, we have no need for any authentication options.
> 
> As of now we use the vanilla TFServing images so there are no changes to the public one.




Resolving patch:
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 @@ def parse_serving_binaries(  # pylint: disable=invalid-name
   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 @@ class TensorFlowServing(ServingBinary):
   def __init__(
       self,
       model_name: Text,
+      image_name: Optional[Text] = None,
       tag: Optional[Text] = None,
       digest: Optional[Text] = None,
   ):
@@ -133,10 +137,11 @@ def __init__(
     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:


Test patch:
diff --git a/tfx/components/infra_validator/serving_bins_test.py b/tfx/components/infra_validator/serving_bins_test.py
new file mode 100644
--- /dev/null
+++ b/tfx/components/infra_validator/serving_bins_test.py
@@ -0,0 +1,54 @@
+# Copyright 2021 Google LLC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for tfx.components.infra_validator.serving_bins."""
+
+import tensorflow as tf
+
+from tfx.components.infra_validator import serving_bins
+from tfx.proto import infra_validator_pb2
+
+
+class ServingBinsTest(tf.test.TestCase):
+
+  def testParseServingBinaries_TensorFlowServing(self):
+    spec = infra_validator_pb2.ServingSpec(
+        tensorflow_serving=infra_validator_pb2.TensorFlowServing(
+            image_name='gcr.io/my_project/my_serving_image',
+            tags=['t1', 't2'],
+            digests=['sha256:d1', 'sha256:d2']))
+    result = serving_bins.parse_serving_binaries(spec)
+
+    self.assertLen(result, 4)
+    for item in result:
+      self.assertIsInstance(item, serving_bins.TensorFlowServing)
+    self.assertCountEqual([item.image for item in result], [
+        'gcr.io/my_project/my_serving_image:t1',
+        'gcr.io/my_project/my_serving_image:t2',
+        'gcr.io/my_project/my_serving_image@sha256:d1',
+        'gcr.io/my_project/my_serving_image@sha256:d2',
+    ])
+
+  def testParseServingBinaries_TensorFlowServing_DefaultImageName(self):
+    spec = infra_validator_pb2.ServingSpec(
+        tensorflow_serving=infra_validator_pb2.TensorFlowServing(
+            tags=['latest']))
+    result = serving_bins.parse_serving_binaries(spec)
+
+    self.assertLen(result, 1)
+    self.assertIsInstance(result[0], serving_bins.TensorFlowServing)
+    self.assertEqual(result[0].image, 'tensorflow/serving:latest')
+
+
+if __name__ == '__main__':
+  tf.test.main()