From 58594593d1b36d3d9f2db9c539032cfb7dac71c7 Mon Sep 17 00:00:00 2001
From: Yifei Teng <yifeit@google.com>
Date: Wed, 8 Jan 2025 15:24:41 -0800
Subject: [PATCH] Add a CPU check and CI

Similar to https://github.com/AI-Hypercomputer/torchprime/pull/13, we
also add a CPU github action. This action will run `pytest` on the repo.

Currently there is only one test, which is the Llama test in
torch_xla_models.

In order to run the test today, we need a HF_TOKEN. I created a personal
read only token and
https://github.com/AI-Hypercomputer/torchprime/issues/14 tracks avoiding
the need for HF_TOKEN, after which I'll need to remember to invalidate
the token.
---
 .github/workflows/cpu_test.yml | 29 +++++++++++++++++++
 .github/workflows/lint.yml     | 51 ++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+)
 create mode 100644 .github/workflows/cpu_test.yml
 create mode 100644 .github/workflows/lint.yml

diff --git a/.github/workflows/cpu_test.yml b/.github/workflows/cpu_test.yml
new file mode 100644
index 0000000..cc015c6
--- /dev/null
+++ b/.github/workflows/cpu_test.yml
@@ -0,0 +1,29 @@
+name: CPU tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  schedule:
+  - cron: "0 8 * * *"
+
+jobs:
+  pytest:
+    runs-on: ubuntu-22.04
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11"]
+    container:
+      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_${{ matrix.python-version }}_tpuvm
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install dev dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -e '.[dev]'
+    - name: Run PyTest
+      run: |
+        # TODO(https://github.com/AI-Hypercomputer/torchprime/issues/14): Remove and burn the token.
+        export HF_TOKEN=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
+        pytest
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
new file mode 100644
index 0000000..ff26dde
--- /dev/null
+++ b/.github/workflows/lint.yml
@@ -0,0 +1,51 @@
+name: Lint
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  schedule:
+  - cron: "0 10 * * *"
+
+jobs:
+  lint:
+    name: Check format and lint
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v3
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dev dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -e '.[dev]'
+    - name: Lint the code with ruff
+      run: |
+        ruff check
+    - name: Check formatting with yapf
+      shell: bash
+      run: |
+        git_status=$(git status --porcelain)
+        if [[ $git_status ]]; then
+          echo "Checkout code is not clean"
+          echo "${git_status}"
+          exit 1
+        fi
+
+        yapf --recursive -i '*.py' torchprime launcher
+        git_status=$(git status --porcelain)
+        if [[ $git_status ]]; then
+          git diff
+          echo "yapf recommends the changes above, please manually apply them OR automatically apply the changes "
+          echo "by running `yapf -i /PATH/TO/foo.py` to the following files"
+          echo "${git_status}"
+          exit 1
+        else
+          echo "PASSED Python format"
+        fi