# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.

# Focused PR gate for the MLX dispatch surface. Two jobs:
#
#   1. mlx-dispatch (Linux+CPU)
#      Runs the three test files documented in tests/studio/README.md with
#      hardware probes spoofed (no Apple Silicon, no real GPU, no real MLX
#      install required):
#        - test_hardware_dispatch_matrix.py    parametrized 7-profile matrix
#                                              + 2 dispatch-priority canaries
#        - test_is_mlx_dispatch_gate.py        AST + runtime guard on
#                                              unsloth._IS_MLX
#        - test_mlx_training_worker_behaviors.py  AST contract checks on
#                                                 studio/backend/core/training/worker.py
#
#   2. mlx-real-apple-silicon (macOS-14, M1)
#      Runs the SAME tests on a real Apple Silicon runner with a real mlx
#      install. macos-14 is GitHub-hosted standard (3 vCPU / 7GB / M1) and
#      is FREE for public repositories per the GitHub Actions billing
#      docs, so adding this job costs zero minutes. Catches the residual
#      gap the spoofed Linux job cannot: real `import mlx`, real
#      `_IS_MLX = True` activation, real `unsloth_zoo.mlx_loader`
#      submodule load, and a sanity import of every PR-A MLX module.
#
# Surfaces "MLX dispatch broke" as its own check in the PR UI rather than
# burying it inside the broader Backend CI. The Backend CI Repo tests
# (CPU) job already covers test_mlx_training_worker_behaviors.py via auto
# discovery, and runs the two state-sensitive matrix files in their own
# pytest invocation; this workflow re-runs them in isolation as a focused
# guard.

name: MLX CI

on:
  pull_request:
    paths:
      - 'unsloth/__init__.py'
      - 'unsloth/_gpu_init.py'
      - 'studio/backend/utils/hardware/**'
      - 'studio/backend/core/training/worker.py'
      - 'studio/backend/core/inference/mlx_inference.py'
      - 'tests/studio/test_hardware_dispatch_matrix.py'
      - 'tests/studio/test_is_mlx_dispatch_gate.py'
      - 'tests/studio/test_mlx_training_worker_behaviors.py'
      - 'tests/conftest.py'
      - '.github/workflows/mlx-ci.yml'
  push:
    branches: [main, pip]

concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: true

permissions:
  contents: read

jobs:
  mlx-dispatch:
    name: MLX dispatch matrix + gates
    runs-on: ubuntu-latest
    timeout-minutes: 10
    steps:
      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5  # v4.3.1

      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
        with:
          python-version: '3.12'
          cache: 'pip'

      # The dispatch matrix and gate tests load unsloth/__init__.py and
      # studio.backend.utils.hardware under monkeypatched probes, so the
      # full unsloth import chain has to succeed on CPU. This mirrors the
      # Repo tests (CPU) dep set in studio-backend-ci.yml.
      - name: Install deps (CPU import chain for unsloth + studio backend)
        run: |
          python -m pip install --upgrade pip
          pip install -r studio/backend/requirements/studio.txt
          pip install \
            python-multipart aiofiles sqlalchemy cryptography \
            pyyaml jinja2 mammoth unpdf requests typer \
            'numpy<3' pytest pytest-asyncio httpx
          pip install --index-url https://download.pytorch.org/whl/cpu \
            'torch>=2.4,<2.11' 'torchvision<0.26'
          pip install 'transformers>=4.51,<5.5'
          pip install 'bitsandbytes>=0.45'
          pip install 'unsloth_zoo>=2026.5.1'
          pip install -e . --no-deps

      - name: MLX dispatch tests (3 files, ~36 tests, hardware spoofed)
        env:
          PYTHONPATH: ${{ github.workspace }}/studio
          UNSLOTH_COMPILE_DISABLE: '1'
        run: |
          python -m pytest -v --tb=short \
            tests/studio/test_hardware_dispatch_matrix.py \
            tests/studio/test_is_mlx_dispatch_gate.py \
            tests/studio/test_mlx_training_worker_behaviors.py

  mlx-real-apple-silicon:
    name: MLX real Apple Silicon (macos-14, M1)
    runs-on: macos-14
    timeout-minutes: 15
    steps:
      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5  # v4.3.1

      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
        with:
          python-version: '3.12'
          cache: 'pip'

      # Real mlx install. The GitHub-hosted macos-14 runner is M1 / arm64
      # so platform.system() == 'Darwin' and platform.machine() == 'arm64'
      # without any spoofing -- exactly the conditions unsloth._IS_MLX
      # checks for. Once mlx is on the path the gate flips True and the
      # MLX import branch in unsloth/__init__.py runs.
      #
      # Install unsloth_zoo from git main rather than PyPI: the PR-A
      # MLX support (#620) merged AFTER the most recent unsloth_zoo
      # release, so the PyPI wheel still raises NotImplementedError on
      # Apple Silicon at `device_type.get_device_type()`. Studio's own
      # install.sh overlays unsloth-zoo from git main for the same
      # reason; we mirror that here.
      - name: Install deps (real mlx + unsloth + unsloth-zoo from git)
        run: |
          python -m pip install --upgrade pip
          pip install mlx mlx-lm
          pip install pytest pytest-asyncio
          pip install --no-deps "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"
          pip install -e . --no-deps

      - name: Verify _IS_MLX flips True on real Apple Silicon
        run: |
          python -c "
          import platform
          assert platform.system() == 'Darwin', platform.system()
          assert platform.machine() == 'arm64', platform.machine()
          import unsloth
          assert unsloth._IS_MLX is True, f'expected _IS_MLX=True on real Apple Silicon, got {unsloth._IS_MLX}'
          print('OK: _IS_MLX activated on real Apple Silicon')
          "

      - name: Smoke-import every MLX-only unsloth_zoo module
        run: |
          python -c "
          import importlib
          for name in [
              'unsloth_zoo.mlx_loader',
              'unsloth_zoo.mlx_trainer',
              'unsloth_zoo.mlx_compile',
              'unsloth_zoo.mlx_utils',
              'unsloth_zoo.mlx_cce',
              'unsloth_zoo.gated_delta_vjp',
          ]:
              importlib.import_module(name)
              print('OK:', name)
          from unsloth_zoo.mlx_loader import FastMLXModel
          from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig
          assert hasattr(FastMLXModel, 'from_pretrained')
          print('OK: FastMLXModel + MLXTrainer surface present')
          "

      - name: MLX dispatch tests on real Apple Silicon (spoofs still apply)
        env:
          PYTHONPATH: ${{ github.workspace }}/studio
          UNSLOTH_COMPILE_DISABLE: '1'
        run: |
          python -m pytest -v --tb=short \
            tests/studio/test_hardware_dispatch_matrix.py \
            tests/studio/test_is_mlx_dispatch_gate.py \
            tests/studio/test_mlx_training_worker_behaviors.py