diff --git a/au_test_harness/cros_au_test_harness.py b/au_test_harness/cros_au_test_harness.py
index df6ab93..0745267 100755
--- a/au_test_harness/cros_au_test_harness.py
+++ b/au_test_harness/cros_au_test_harness.py
@@ -28,6 +28,7 @@
 from chromite.lib import cros_build_lib
 from chromite.lib import parallel
 from chromite.lib import sudo
+from chromite.lib import timeout_util
 from crostestutils.au_test_harness import au_test
 from crostestutils.au_test_harness import au_worker
 from crostestutils.lib import dev_server_wrapper
@@ -66,7 +67,7 @@
 
     If the test suite fails, raise a BackgroundFailure.
     """
-    with cros_build_lib.SubCommandTimeout(constants.MAX_TIMEOUT_SECONDS):
+    with timeout_util.Timeout(constants.MAX_TIMEOUT_SECONDS):
       test_result = super(_LessBacktracingTestRunner, self).run(*args, **kwargs)
       if test_result is None or not test_result.wasSuccessful():
         msg = 'Test harness failed. See logs for details.'
diff --git a/generate_test_payloads/cros_generate_test_payloads.py b/generate_test_payloads/cros_generate_test_payloads.py
index faf782e..4703ef1 100755
--- a/generate_test_payloads/cros_generate_test_payloads.py
+++ b/generate_test_payloads/cros_generate_test_payloads.py
@@ -38,6 +38,7 @@
 from chromite.lib import osutils
 from chromite.lib import parallel
 from chromite.lib import sudo
+from chromite.lib import timeout_util
 from crostestutils.au_test_harness import cros_au_test_harness
 from crostestutils.generate_test_payloads import payload_generation_exception
 from crostestutils.lib import dev_server_wrapper
@@ -255,12 +256,12 @@
 
       logging.info(debug_message)
       try:
-        with cros_build_lib.SubCommandTimeout(constants.MAX_TIMEOUT_SECONDS):
+        with timeout_util.Timeout(constants.MAX_TIMEOUT_SECONDS):
           cros_build_lib.SudoRunCommand(command, log_stdout_to_file=log_file,
                                         combine_stdout_stderr=True,
                                         enter_chroot=True, print_cmd=False,
                                         cwd=constants.SOURCE_ROOT)
-      except (cros_build_lib.TimeoutError, cros_build_lib.RunCommandError):
+      except (timeout_util.TimeoutError, cros_build_lib.RunCommandError):
         # Print output first, then re-raise the exception.
         if os.path.isfile(log_file):
           logging.error(osutils.ReadFile(log_file))