Update jax-array-api.yml

Author: quoctruong

.github/workflows/jax-array-api.yml

@@ -13,16 +13,24 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  build:
-    runs-on: linux-x86-ct5lp-x4
-    container: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:infrastructure-public-image-2d2a7b1e6e2e
-    steps:
-    - name: Set up Python 3.11
-      uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b  # v5.3.0
-      with:
-        python-version: 3.11
-    - name: Use the value
-      id: step_two
-      run: |
-        pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-        python -c 'import jax; print("Global device count:", jax.device_count()); print("local devices =", jax.local_devices())'
+  # Label of the container job
+  container-job:
+    # Containers must run in Linux based operating systems
+    runs-on: linux-x86-n2-32
+    # Docker Hub image that `container-job` executes in
+    container: node:20-bookworm-slim
+
+    # Service containers to run with `container-job`
+    services:
+      # Label used to access the service container
+      redis:
+        # Docker Hub image
+        image: redis
+        # Set health checks to wait until redis has started
+        options: >-
+          --health-cmd "redis-cli ping"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+
+