Developed Dockerfile for TPU build job in actions.yml

keras-team · kharshith-k · Jun 16, 2025 · Jun 16, 2025 · Jun 18, 2025 · Jun 18, 2025
commit 09161d716a3de5b6bfecb7a7da9633246dd6a99e
diff --git a/.github/workflows/actions.yml b/.github/workflows/actions.yml
@@ -24,14 +24,17 @@ jobs:
       KERAS_HOME: .github/workflows/config/${{ matrix.backend }}
       KERAS_BACKEND: tensorflow
     steps:
-    - uses: actions/checkout@v4 # Checks-out your repository under $GITHUB_WORKSPACE, so your workflow can access it.
+      - uses: actions/checkout@v4
 
-    - name: Make script executable
-      run: chmod +x .github/workflows/tpu/build.sh # Assuming your script is named my-script.sh
+      - name: Build and run Docker image for TPU tests
+        run: |
+          docker build -f .github/workflows/tpu/Dockerfile -t keras-tpu-test .
+          docker run --rm \
+            -e PYTHON=${{ matrix.python-version }} \
+            -e KERAS_HOME=.github/workflows/config/${{ matrix.backend }} \
+            -e KERAS_BACKEND=tensorflow \
+            keras-tpu-test
 
-    - name: Run my shell script
-      run: .github/workflows/tpu/build.sh # Execute the script
-
   build:
     strategy:
       fail-fast: false
@@ -148,4 +151,4 @@ jobs:
           pip uninstall -y keras keras-nightly
           pip install -e "." --progress-bar off --upgrade
       - name: Run pre-commit
-        run: pre-commit run --all-files --hook-stage manual
+        run: pre-commit run --all-files --hook-stage manual
diff --git a/.github/workflows/tpu/Dockerfile b/.github/workflows/tpu/Dockerfile
@@ -0,0 +1,28 @@
+FROM python:3.10-slim
+
+ENV KERAS_HOME=/github/workspace/.github/workflows/config/tensorflow \
+    KERAS_BACKEND=tensorflow
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git \
+    sudo \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy the entire codebase into the container
+COPY . /github/workspace
+WORKDIR /github/workspace
+
+# Create and activate venv, install pip/setuptools/psutil, then run tests
+RUN cd src/github/keras && \
+    pip install -U pip setuptools && \
+    pip install -U psutil && \
+    pip install -r requirements-tensorflow-tpu.txt && \
+    pip uninstall -y keras keras-nightly && \
+    python3 -c 'import tensorflow as tf;print(tf.__version__);print(tf.config.list_physical_devices("TPU"))' && \
+    python3 -c 'import tensorflow as tf;assert len(tf.config.list_physical_devices("TPU")) > 0' && \
+    pytest keras --ignore keras/src/applications \
+                 --ignore keras/src/layers/merging/merging_test.py \
+                 --cov=keras \
+                 --cov-config=pyproject.toml
+
+CMD ["bash"]