Create job to build torch-xla wheel and publish to tt-pypi

Author: jazpurTT

.github/workflows/_build_torch_xla_3.10.yml

@@ -0,0 +1,76 @@
+name: build-torch-xla
+on:
+  workflow_call:
+    inputs:
+      torch_version:
+        description: 'Torch version to build (default: 2.7.0)'
+        required: false
+        type: string
+        default: '2.7.0'
+    outputs:
+      artifact_name:
+        description: 'Name of uploaded wheels artifact'
+        value: ${{ jobs.build-wheels.outputs.artifact_name }}
+  workflow_dispatch:
+jobs:
+  build-wheels:
+    runs-on: ubuntu-latest
+    env:
+      ARTIFACT_NAME: install-artifact-torch-xla-release
+      GIT_VERSIONED_XLA_BUILD: 1
+    container:
+      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:tpu
+      options: --user root
+    outputs:
+      artifact_name: ${{ steps.set_upload_name.outputs.artifact_name }}
+    steps:
+      - name: "Build Torch/XLA wheel"
+        id: build_wheels
+        run: |
+          cmake --version
+          apt-get update && apt-get install -y curl git build-essential
+
+          # Clean up any existing pyenv installation
+          rm -rf $HOME/.pyenv
+
+          curl https://pyenv.run | bash
+          export PATH="$HOME/.pyenv/bin:$PATH"
+          eval "$(pyenv init -)"
+          pyenv install 3.10.12
+          pyenv global 3.10.12
+          ln -sf $HOME/.pyenv/versions/3.10.12/bin/python3.10 /usr/local/bin/python3.10
+
+          # Install essential packages for Python 3.10
+          python3.10 -m pip install --upgrade pip
+          python3.10 -m pip install pyyaml setuptools wheel numpy typing_extensions requests
+
+          cd /tmp
+          git clone --recursive --branch v${{ inputs.torch_version || '2.7.0' }} https://github.com/pytorch/pytorch.git
+          cd pytorch/
+          git clone --recursive https://github.com/tenstorrent/pytorch-xla.git xla
+
+          # copy pre-built wheels from cache
+          python3.10 setup.py bdist_wheel
+          python3.10 setup.py develop
+
+          # Build PyTorch/XLA
+          cd xla/
+          python3.10 setup.py bdist_wheel
+
+          # Collect wheels
+          mkdir -p /dist
+          cp dist/*.whl /dist/
+
+          # Clean up any existing pyenv installation
+          rm -rf $HOME/.pyenv
+
+      - name: "Upload Wheels Artifact"
+        id: upload
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ env.ARTIFACT_NAME }}
+          path: /dist/*.whl
+      
+      - name: Set artifact name output
+        id: set_upload_name
+        run: echo "artifact_name=${{ env.ARTIFACT_NAME }}" >> $GITHUB_OUTPUT

.github/workflows/_publish_torch_xla.yml

@@ -0,0 +1,63 @@
+name: publish-wheel
+on:
+  workflow_call:
+    inputs:
+      artifact_name:
+        required: true
+        type: string
+        description: 'Name of the artifact containing the wheel'
+
+jobs:
+
+  publish-wheels:
+    name: "Publish wheels to internal PyPI"
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+    steps:
+      - name: Validate inputs
+        run: |
+          if [ -z "${{ inputs.artifact_name }}" ]; then
+            echo "ERROR: artifact_name input is empty or not provided!"
+            exit 1
+          fi
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.PYPI_ROLE }}
+          aws-region: ${{ secrets.PYPI_REGION }}
+
+      - name: Install s3pypi
+        run: |
+          pip install s3pypi
+
+      - name: Download wheel artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: ${{ inputs.artifact_name }}
+          path: ./dist
+
+      - name: Publish wheels to internal PyPI
+        run: |
+          wheel_count=$(find ./dist -type f -name "torch_xla*.whl" | wc -l)
+          if [ "$wheel_count" -ne 1 ]; then
+            echo "ERROR: Expected exactly 1 wheel file, but found $wheel_count!"
+            exit 1
+          fi
+ 
+          wheel_file=$(find ./dist -type f -name "torch_xla*.whl" -exec realpath {} \;)
+          wheel_basename=$(basename "$wheel_file")
+          echo "Wheel file found, publishing $wheel_basename to PyPi server"
+
+          s3pypi upload "$wheel_file" --bucket ${{ secrets.PYPI_BUCKET }} --put-root-index --force
+          if [ $? -ne 0 ]; then
+            echo "ERROR: Failed to upload $wheel_basename to S3 PyPI"
+            exit 1
+          fi
+          echo "Successfully uploaded $wheel_basename"

.github/workflows/build_and_publish.yml

@@ -0,0 +1,33 @@
+name: Build and Publish PyTorch/XLA 
+on:
+  push:
+    branches:
+      - master
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
+jobs:
+  check_code_changes:
+    name: Check Code Changes
+    uses: ./.github/workflows/_check_code_changes.yml
+    with:
+      event_name: ${{ github.event_name }}
+      base_sha: ${{ github.event.before }}
+      head_sha: ${{ github.sha }}
+
+  build-torch-xla:
+    name: "Build PyTorch/XLA for Python 3.10"
+    if: needs.check_code_changes.outputs.has_code_changes == 'true'
+    uses: ./.github/workflows/_build_torch_xla_3.10.yml
+    needs: check_code_changes
+
+  publish-torch-xla:
+    name: "Publish PyTorch/XLA"
+    uses: ./.github/workflows/_publish_torch_xla.yml
+    needs: build-torch-xla
+    secrets: inherit
+    with:
+      artifact_name: ${{ needs.build-torch-xla.outputs.artifact_name }}