Increase Timeout

kubeflow · Jan 10, 2025 · d2b7314 · d2b7314
1 parent 145d7fc
commit d2b7314
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/examples/jax/jax-dist-spmd-mnist/Dockerfile b/examples/jax/jax-dist-spmd-mnist/Dockerfile
@@ -1,7 +1,7 @@
 FROM python:3.13
 
 RUN pip install --upgrade pip
-RUN pip install --upgrade jax absl-py
+RUN pip install --upgrade jax[k8s] absl-py
 
 RUN apt-get update && apt-get install -y \
     build-essential \
@@ -24,6 +24,6 @@ RUN git clone https://github.com/facebookincubator/gloo.git \
 
 WORKDIR /app
 
-ADD datasets.py spmd_mnist_classifier_fromscratch.py /app
+ADD datasets.py spmd_mnist_classifier_fromscratch.py /app/
 
 ENTRYPOINT ["python3", "spmd_mnist_classifier_fromscratch.py"]
diff --git a/sdk/python/test/e2e/test_e2e_jaxjob.py b/sdk/python/test/e2e/test_e2e_jaxjob.py
@@ -88,7 +88,7 @@ def test_sdk_e2e_with_gang_scheduling(job_namespace):
     logging.info(TRAINING_CLIENT.list_jobs(job_namespace))
 
     try:
-        utils.verify_job_e2e(TRAINING_CLIENT, JOB_NAME, job_namespace, wait_timeout=3000)
+        utils.verify_job_e2e(TRAINING_CLIENT, JOB_NAME, job_namespace, wait_timeout=9000)
     except Exception as e:
         utils.print_job_results(TRAINING_CLIENT, JOB_NAME, job_namespace)
         TRAINING_CLIENT.delete_job(JOB_NAME, job_namespace)
@@ -123,7 +123,7 @@ def test_sdk_e2e(job_namespace):
     logging.info(TRAINING_CLIENT.list_jobs(job_namespace))
 
     try:
-        utils.verify_job_e2e(TRAINING_CLIENT, JOB_NAME, job_namespace, wait_timeout=3000)
+        utils.verify_job_e2e(TRAINING_CLIENT, JOB_NAME, job_namespace, wait_timeout=9000)
     except Exception as e:
         utils.print_job_results(TRAINING_CLIENT, JOB_NAME, job_namespace)
         TRAINING_CLIENT.delete_job(JOB_NAME, job_namespace)