Uses NIGHTLY dependency instead for container image build when UNCONS…

…TRAINED was specified.

PiperOrigin-RevId: 360343968

tfx/orchestration/test_utils.py

@@ -75,25 +75,32 @@ def build_docker_image(container_image: str, repo_base: str):
     repo_base: The src path to use to build docker image.
   """
 
-  # Default to NIGHTLY. GIT_MASTER might be better to use the latest source,
-  # But it takes too long (~1h) to build packages from scratch. If some changes
-  # in a dependent package break tests, just run a nightly build of dependent
-  # package again.
-  dependency_selector = os.getenv('TFX_DEPENDENCY_SELECTOR') or 'NIGHTLY'
   [docker_image_repo, docker_image_tag] = container_image.split(':')
-
-  logging.info('Building image %s with %s dependency', container_image,
-               dependency_selector)
+  envs = {
+      'DOCKER_IMAGE_REPO': docker_image_repo,
+      'DOCKER_IMAGE_TAG': docker_image_tag
+  }
+
+  # UNCONSTRAINED means that we will use the latest version which is built from
+  # latest source. But it takes too long (~1h) to build packages from
+  # GIT_MASTER. So we fallback to NIGHTLY here. If some changes in a dependent
+  # package break tests, just run a nightly build of dependent package again.
+  # TODO(b/181290953): Use UNCONSTRAINED as it is if we can supply latest wheel
+  # packages for TFX family libraries to the image build script.
+  dependency_selector = os.getenv('TFX_DEPENDENCY_SELECTOR')
+  if dependency_selector == 'UNCONSTRAINED':
+    dependency_selector = 'NIGHTLY'
+
+  if dependency_selector:
+    envs['TFX_DEPENDENCY_SELECTOR'] = dependency_selector
+
+  logging.info('Building image %s with env:%s', container_image, envs)
   with Timer('BuildingTFXContainerImage'), _chdir(repo_base):
     subprocess.check_call(
         args=[
             os.path.join(repo_base, 'tfx/tools/docker/build_docker_image.sh'),
         ],
-        env={
-            'DOCKER_IMAGE_REPO': docker_image_repo,
-            'DOCKER_IMAGE_TAG': docker_image_tag,
-            'TFX_DEPENDENCY_SELECTOR': dependency_selector,
-        },
+        env=envs,
         shell=True,
     )