diff --git a/sdk/python/kubeflow/training/api/training_client.py b/sdk/python/kubeflow/training/api/training_client.py
index c3ee98b267..782863586b 100644
--- a/sdk/python/kubeflow/training/api/training_client.py
+++ b/sdk/python/kubeflow/training/api/training_client.py
@@ -871,7 +871,8 @@ def wait_for_job_conditions(
         Raises:
             ValueError: Invalid input parameters.
             TimeoutError: Timeout to get Job.
-            RuntimeError: Failed to get Job or Job reaches unexpected Failed condition.
+            RuntimeError: Failed to get Job, or Job reaches Failed condition and
+                Failed is not in `expected_conditions` set.
         """
 
         namespace = namespace or self.namespace
@@ -904,9 +905,9 @@ def wait_for_job_conditions(
             if callback:
                 callback(job)
 
-            # Raise an exception if Job is Failed and Failed is not expected condition.
+            # Raise an exception if Job is Failed and Failed is not the expected condition.
             if (
-                constants.JOB_CONDITION_FAILED not in conditions
+                constants.JOB_CONDITION_FAILED not in expected_conditions
                 and utils.has_condition(conditions, constants.JOB_CONDITION_FAILED)
             ):
                 raise RuntimeError(