Merge pull request #6 from tenstorrent/ddilbaz/python_3.11

Uplift wheel python 3.10 to 3.11

Author: ddilbazTT

.github/workflows/_build_torch_xla_3.11.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.11
+          pyenv global 3.11
+          ln -sf $HOME/.pyenv/versions/3.11/bin/python3.11 /usr/local/bin/python3.11
+
+          # Install essential packages for Python 3.11
+          python3.11 -m pip install --upgrade pip
+          python3.11 -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.11 setup.py bdist_wheel
+          python3.11 setup.py develop
+
+          # Build PyTorch/XLA
+          cd xla/
+          python3.11 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

@@ -25,7 +25,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
-          python-version: '3.10'
+          python-version: '3.11'
 
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v4

.github/workflows/build_and_publish.yml

@@ -19,9 +19,9 @@ jobs:
       head_sha: ${{ github.sha }}
 
   build-torch-xla:
-    name: "Build PyTorch/XLA for Python 3.10"
+    name: "Build PyTorch/XLA for Python 3.11"
     if: needs.check_code_changes.outputs.has_code_changes == 'true'
-    uses: ./.github/workflows/_build_torch_xla_3.10.yml
+    uses: ./.github/workflows/_build_torch_xla_3.11.yml
     needs: check_code_changes
 
   publish-torch-xla: