Update values and disable workflow from running

Author: nitins17

.github/workflows/pytest_cpu_reuse.yml

@@ -1,9 +1,18 @@
 name: Run Pytest CPU tests (resuable workflow)
 
 on:
-  pull_request:
-    branches:
-      - main
+  # pull_request:
+  #   branches:
+  #     - main
+  workflow_dispatch:
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+        - 'yes'
 
 jobs:
   build_jaxlib_artifacts:
@@ -22,32 +31,28 @@ jobs:
         python: ["3.10"]
 
     runs-on: ${{ matrix.runner }}
-    container: ${{ (contains(matrix.runner, 'linux-x86') && 'gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest') ||
+    container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
                 (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest') ||
                 (contains(matrix.runner, 'windows-x86') && null) }}
 
     env:
-      # Disable running `run_docker_container.sh`. JAX self-hosted runners already run in a Docker
-      # container.
-      JAXCI_RUN_DOCKER_CONTAINER: 0
-      # Use RBE to build the artifacts where possible (Linux x86 and Windows).
-      JAXCI_BUILD_ARTIFACT_WITH_RBE: 1
-      # Setup the test environment (disable x64 mode and clone XLA at HEAD)
-      JAXCI_SETUP_TEST_ENVIRONMENT: 1
+      JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
 
     steps:
       - uses: actions/checkout@v3
       # Halt for testing
       - name: Wait For Connection
-        uses: ./actions/ci_connection/
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
       - name: Install pytest
         env:
           JAXCI_PYTHON: python${{ matrix.python }}
         run: $JAXCI_PYTHON -m pip install pytest
       - name: Install dependencies
         env:
           JAXCI_PYTHON: python${{ matrix.python }}
-        run: $JAXCI_PYTHON -m pip install --upgrade     numpy=="2.0.0" scipy=="1.13.1"     wheel build pytest-xdist absl-py opt-einsum colorama portpicker     matplotlib 'importlib_metadata>=4.6' hypothesis flatbuffers filelock
+        run: $JAXCI_PYTHON -m pip install -r build/requirements.in
       - name: Run Pytest CPU tests
         run: ./ci/run_pytest.sh "ci/envs/run_tests/pytest_cpu"