[Deps] Split up of dependencies

.github/workflows/build.yml

@@ -115,22 +115,38 @@ jobs:
         create-symlink: true
         key: forge-build-${{ matrix.build.build-type }}
 
-    - name: Build (creates tvm and tt-forge-fe wheels)
+    - name: Build (creates tvm and both variant tt-forge-fe wheels)
       shell: bash
       run: |
         source env/activate
-        python3 setup.py bdist_wheel
+        # Build both wheel variants (compiler and dev)
+        python3 build_wheels.py
+        # Build TVM wheel
         cd third_party/tvm/python
         python3 setup.py bdist_wheel
         cd ../../..
         cp third_party/tvm/python/dist/*.whl dist/
 
-    - name: Upload Build
+    - name: Upload Compiler Wheel
       uses: actions/upload-artifact@v4
       if: success() || failure()
       with:
-        name: forge-wheel
-        path: dist/*.whl
+        name: forge-wheel-compiler
+        path: dist/compiler/*.whl
+
+    - name: Upload Dev Wheel
+      uses: actions/upload-artifact@v4
+      if: success() || failure()
+      with:
+        name: forge-wheel-dev
+        path: dist/dev/*.whl
+
+    - name: Upload TVM Wheel
+      uses: actions/upload-artifact@v4
+      if: success() || failure()
+      with:
+        name: tvm-wheel
+        path: dist/tvm*.whl
 
 
   unit-test-build:
@@ -246,6 +262,8 @@ jobs:
           disable_search: true
           token: ${{ secrets.CODECOV_TOKEN }}
 
+
+
   build-conclusion:
     needs: [forge-build, unit-test-build]
     if: always()
@@ -258,9 +276,13 @@ jobs:
         shell: bash
         run: |
           if [ "${{ needs.forge-build.result }}" == "success" ]; then
-            echo "Forge wheel is available"
+            echo "Forge wheels are available and built successfully"
+            echo "- Forge build result: PASSED"
+            echo "- Unit test result: ${{ needs.unit-test-build.result }}"
             echo "run_id=${{ github.run_id }}" >> $GITHUB_OUTPUT
           else
-            echo "Forge wheel is NOT available"
+            echo "Build failed"
+            echo "- Forge build result: ${{ needs.forge-build.result }}"
+            echo "- Unit test result: ${{ needs.unit-test-build.result }}"
             echo "run_id=" >> $GITHUB_OUTPUT
           fi

.github/workflows/on-nightly.yml

@@ -79,6 +79,17 @@ jobs:
       docker-image: ${{ needs.docker-build.outputs.docker-image }}
       build: ${{ needs.set-inputs.outputs.rebuild }}
 
+  sanity-test:
+    if: success() || failure()
+    needs:
+      - docker-build
+      - build
+    uses: ./.github/workflows/test-sanity.yml
+    secrets: inherit
+    with:
+      docker-image: ${{ needs.docker-build.outputs.docker-image }}
+      run_id: ${{ needs.build.outputs.run_id }}
+
   test_full_model_passing:
     if: success() || failure()
     needs:

.github/workflows/on-push.yml

@@ -32,3 +32,13 @@ jobs:
       docker-image: ${{ needs.docker-build.outputs.docker-image }}
       run_id: ${{ github.run_id }}
       runs-on: '[{"runs-on": "n150"}, {"runs-on": "n300"}]'
+
+  sanity-test:
+    needs:
+      - docker-build
+      - build
+    uses: ./.github/workflows/test-sanity.yml
+    secrets: inherit
+    with:
+      docker-image: ${{ needs.docker-build.outputs.docker-image }}
+      run_id: ${{ github.run_id }}

.github/workflows/test-sanity.yml

@@ -0,0 +1,140 @@
+name: Sanity Test - Forge Add
+
+on:
+  workflow_dispatch:
+    inputs:
+      run_id:
+        description: 'Run id to download wheels from (or leave empty to search latest)'
+        required: false
+        type: string
+  workflow_call:
+    inputs:
+      run_id:
+        description: 'Run id to download wheels from'
+        required: false
+        type: string
+      docker-image:
+        description: 'Docker image to use for test'
+        required: true
+        type: string
+
+permissions:
+  packages: read
+  actions: read
+  checks: write
+
+jobs:
+  sanity-forge-test:
+    strategy:
+      fail-fast: false
+      matrix:
+        wheel: [
+          { name: "compiler", artifact: "forge-wheel-compiler" }
+        ]
+
+    runs-on:
+      - tt-beta-ubuntu-2204-large
+
+    container:
+      image: ${{ inputs.docker-image }}
+
+    steps:
+      - name: Fetch job id
+        id: fetch-job-id
+        uses: tenstorrent/tt-github-actions/.github/actions/job_id@main
+        with:
+          job_name: "sanity-forge-test (${{ matrix.wheel.name }})"
+
+      - name: Set reusable strings
+        id: strings
+        shell: bash
+        env:
+          JOB_ID: ${{ steps.fetch-job-id.outputs.job_id }}
+        run: |
+          echo "work-dir=$(pwd)" >> "$GITHUB_OUTPUT"
+          echo "test_report_path=reports/report_$JOB_ID.xml" >> "$GITHUB_OUTPUT"
+
+      - name: Git safe dir
+        run: git config --global --add safe.directory ${{ steps.strings.outputs.work-dir }}
+
+      - uses: actions/checkout@v4
+        with:
+          sparse-checkout: |
+            env/
+            forge/test
+            pytest.ini
+            conftest.py
+          clean: true
+          fetch-depth: 1
+
+      - name: Download wheel
+        if: ${{ inputs.run_id }}
+        continue-on-error: true
+        uses: tenstorrent/tt-forge/.github/actions/download-artifact@main
+        with:
+          name: ${{ matrix.wheel.artifact }}
+          run_id: ${{ inputs.run_id }}
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Find and download forge wheel
+        if: ${{ !inputs.run_id }}
+        uses: dawidd6/action-download-artifact@v9
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          workflow_conclusion: success
+          workflow_search: true
+          workflow: build.yml
+          name: ${{ matrix.wheel.artifact }}
+          repo: tenstorrent/tt-forge-fe
+          check_artifacts: true
+          search_artifacts: true
+
+      - name: Install wheel
+        shell: bash
+        run: |
+          source env/activate
+          # Validate wheel file exists and is secure
+          wheel_file=$(find . -name "*.whl" | head -1)
+          if [ -z "$wheel_file" ]; then
+            echo "Error: No wheel file found"
+            exit 1
+          fi
+          echo "Installing wheel: $wheel_file"
+          # Validate wheel file is from expected source
+          if [[ "$wheel_file" != *"forge"* ]]; then
+            echo "Error: Wheel file does not appear to be a forge wheel"
+            exit 1
+          fi
+          pip install "$wheel_file" --force-reinstall
+
+      - name: Run sanity test
+        shell: bash
+        run: |
+          source env/activate
+          echo "Running forge add sanity test for ${{ matrix.wheel.name }} wheel..."
+          pytest forge/test/test_sanity_forge_add.py -v -m "sanity" --tb=short \
+            --junit-xml=${{ steps.strings.outputs.test_report_path }}
+
+      - name: Upload Test Report
+        uses: actions/upload-artifact@v4
+        if: success() || failure()
+        with:
+          name: test-reports-${{ matrix.wheel.name }}-${{ steps.fetch-job-id.outputs.job_id }}
+          path: ${{ steps.strings.outputs.test_report_path }}
+
+      - name: Show Test Report
+        uses: mikepenz/action-junit-report@v5
+        if: success() || failure()
+        with:
+          report_paths: ${{ steps.strings.outputs.test_report_path }}
+          check_name: TT-Forge-FE Sanity Tests (${{ matrix.wheel.name }})
+          comment: false
+          updateComment: false
+          detailed_summary: true
+          group_suite: true
+
+      - name: Cleanup
+        if: always()
+        shell: bash
+        run: |
+          rm -rf *.whl

build_wheels.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import subprocess
+import os
+import shutil
+
+
+def manage_pyproject_toml(disable=False):
+    """Temporarily disable/enable pyproject.toml to allow setup.py full control"""
+    if disable:
+        if os.path.exists("pyproject.toml"):
+            if os.path.exists("pyproject.toml.backup"):
+                os.remove("pyproject.toml.backup")
+            shutil.move("pyproject.toml", "pyproject.toml.backup")
+    else:
+        if os.path.exists("pyproject.toml.backup"):
+            if os.path.exists("pyproject.toml"):
+                os.remove("pyproject.toml")
+            shutil.move("pyproject.toml.backup", "pyproject.toml")
+
+
+def clean_build_dir():
+    """Clean build directory to prevent conflicts between builds"""
+    if os.path.exists("build"):
+        shutil.rmtree("build")
+
+
+def build_and_collect_wheels():
+
+    try:
+        # Temporarily disable pyproject.toml
+        manage_pyproject_toml(disable=True)
+
+        # Build compiler variant
+        clean_build_dir()
+        os.environ["FORGE_VARIANT"] = "compiler"
+        subprocess.run(["python3", "setup.py", "bdist_wheel"], check=True)
+
+        # Build dev variant
+        clean_build_dir()
+        os.environ["FORGE_VARIANT"] = "dev"
+        subprocess.run(["python3", "setup.py", "bdist_wheel"], check=True)
+
+        built_files = os.listdir("dist")
+
+        # Collect compiler wheel
+        compiler_wheel = None
+        dev_wheel = None
+
+        for fname in built_files:
+            if fname.startswith("tt_forge_fe_compiler") and fname.endswith(".whl"):
+                compiler_wheel = fname
+            elif fname.startswith("tt_forge_fe_dev") and fname.endswith(".whl"):
+                dev_wheel = fname
+
+        # Move compiler wheel
+        if compiler_wheel:
+            os.makedirs("dist/compiler", exist_ok=True)
+            shutil.move(os.path.join("dist", compiler_wheel), os.path.join("dist/compiler", compiler_wheel))
+
+        # Move dev wheel
+        if dev_wheel:
+            os.makedirs("dist/dev", exist_ok=True)
+            shutil.move(os.path.join("dist", dev_wheel), os.path.join("dist/dev", dev_wheel))
+
+    finally:
+        # Always restore pyproject.toml
+        manage_pyproject_toml(disable=False)
+
+
+# Make sure clean slate
+os.makedirs("dist", exist_ok=True)
+
+# Build both wheels
+build_and_collect_wheels()

env/core_requirements.txt

@@ -1,47 +1,40 @@
+#core requirements
 setuptools>=61
-# This is needed to avoid issue https://yyz-gitlab.local.tenstorrent.com/devops/devops/-/issues/95
-# jax requires any version of optax which requires any version of chex which in turn
-# requires jax>=0.4.6 which conflicts with our jax == 0.3.16
-# TODO: Remove when jax library is upgraded
-chex==0.1.6
-dataclasses-json==0.5.7
-datasets==2.14.6
+pybind11==2.6.2
+numpy==1.26.4
+scipy>=1.8.0
+pandas==1.5.3
 decorator==5.1.1
-flatbuffers==23.5.26
-# This is needed to prevent AttributeError: module 'ml_dtypes' has no attribute 'float8_e4m3b11'
-flax==0.9.0
-jax==0.4.30
 loguru==0.5.3
-networkx==2.8.5
-numpy==1.26.4
+flatbuffers==23.5.26
+openpyxl==3.1.5
+GitPython==3.1.44
+pyinstrument>=4.1.1
+pytest==6.2.4
+pytest-timeout==2.0.1
+pytest-xdist==2.5.0
+pytorchcv==0.0.67
+pytest-split
+
+#framework requirements
 onnx>=1.15.0
 onnxruntime>=1.16.3
-opencv-python-headless==4.11.0.86
-# This is needed to avoid issue https://yyz-gitlab.local.tenstorrent.com/devops/devops/-/issues/95
-pandas==1.5.3
-pybind11==2.6.2
-pyinstrument>=4.1.1
-scipy>=1.8.0
+tf2onnx==1.15.1
 tensorflow==2.15
 tensorboard==2.15
-tf2onnx==1.15.1
-transformers==4.47.0
-# To avoid warning during the import
-requests==2.28.2
 tflite==2.10.0
-ultralytics==8.3.91
+torch @ https://download.pytorch.org/whl/cpu/torch-2.7.0%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version == "3.10"
+torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.22.0%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version == "3.10"
+timm==1.0.9
+transformers==4.47.0
+dataclasses-json==0.5.7
+
+# Paddle stack
 paddlepaddle==2.6.2
 paddlenlp==2.8.1
 aistudio-sdk==0.2.6
-pytorch_forecasting==1.0.0
-patool
-openpyxl==3.1.5
-GitPython==3.1.44
-mlp-mixer-pytorch==0.2.0
-gliner==0.2.7
-ase==3.24.0
-hippynn==0.0.3
-bi-lstm-crf==0.2.1
-peft
-pyclipper==1.3.0
-shapely==2.1.1
+
+# JAX stack — note version constraints from chex
+jax==0.4.30
+chex==0.1.6
+flax==0.9.0  # Needs compatibility with jax