# 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, running on a real
# Apple Silicon runner.
#
# Runner: macos-14 (M1, 3 vCPU / 7 GB / Apple Silicon standard runner
# -- FREE for public repositories per the GitHub Actions billing
# reference; larger variants like macos-14-large/-xlarge are paid so
# we deliberately avoid those).
#
# Why a single Mac job (no Linux+spoof leg): the dispatch tests are
# 100% spoofed monkeypatches and run identically on any host, so the
# Linux leg was duplicating the matrix tests already covered on Mac
# while missing everything Apple-specific. The Mac job runs the SAME
# spoofed matrix PLUS three things only a real Apple Silicon host
# can prove:
#
#   1. unsloth._IS_MLX flips True on Darwin+arm64 with mlx genuinely
#      installed (no spoof).
#   2. Every PR-A MLX-only unsloth_zoo module (mlx_loader, mlx_trainer,
#      mlx_compile, mlx_utils, mlx_cce, gated_delta_vjp) imports
#      against the real `mlx` + `mlx-lm` + `mlx-vlm` PyPI wheels --
#      each does `import mlx.core as mx` at module top level, so this
#      catches a future change that breaks the real wheels without
#      needing a Mac developer in the loop.
#   3. The hardware-dispatch spoofs do not collide with the real
#      environment (the test fixture installs a MetaPathFinder that
#      blocks `import mlx.core` for "no-mlx" profiles, faithfully
#      simulating a Mac without mlx even when mlx IS installed).
#   4. End-to-end MLX training + inference smoke test:
#      run_real_mlx_smoke.py trains unsloth/gemma-3-270m-it for 7
#      deterministic LoRA steps on a single repeated text row, then
#      verifies the trained model can complete the prompt and that
#      losses + grad norms are finite and well-behaved. This is the
#      only place in CI that exercises a real MLX backward pass +
#      optimizer step + inference call.
#
# Three dispatch test files documented in tests/studio/README.md:
#   - 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
#
# Surfaces a single PR check ("MLX CI on Mac M1 / dispatch").
#
# Security audit footprint: every package this workflow installs is
# already covered by .github/workflows/security-audit.yml -- the deps
# come from studio/backend/requirements/studio.txt and unsloth-zoo's
# pyproject (resolved transitively). The git+ install of unsloth-zoo
# is intentionally skipped by the audit (pip-audit cannot resolve a
# git URL through PyPI metadata; the audit comment in security-audit.yml
# documents this). No new package is introduced solely by MLX CI.

name: MLX CI on Mac M1

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/studio/run_real_mlx_smoke.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:
  dispatch:
    name: dispatch
    runs-on: macos-14
    # 25 min: dispatch + spoofed matrix + 7-step real LoRA training is
    # under 2 min; GGUF export builds llama.cpp via cmake on Apple
    # Silicon (~5-7 min), so we budget headroom.
    timeout-minutes: 25
    steps:
      # harden-runner audit mode: macOS runners cannot use blocking mode
      # today (eBPF egress enforcement is Linux-only), but audit mode is
      # supported cross-platform and surfaces the egress destinations in
      # the runner log. This produces the data needed to graduate this
      # job to a block-mode allowlist once macOS support lands.
      - name: Harden runner (audit)
        uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450  # v2.19.1
        with:
          egress-policy: audit

      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
        with:
          persist-credentials: false

      - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405  # v6.2.0
        with:
          python-version: '3.12'
          cache: 'pip'

      # macOS install ladder, validated locally against a Linux
      # mac-sim venv (platform spoofed + mlx_simulation shim + real
      # datasets/transformers/structlog).
      #
      # 1. studio/backend/requirements/studio.txt brings structlog,
      #    fastapi, etc. The hardware probe imports structlog at
      #    module top level.
      # 2. Same pytest / numpy / httpx stack the rest of the repo CI
      #    uses.
      # 3. torch is explicitly installed: unsloth-zoo's pyproject
      #    deliberately excludes torch on darwin+arm64 (mlx replaces
      #    it for runtime use), but the dispatch tests spoof
      #    torch.cuda / torch.xpu / torch.backends.mps via monkeypatch
      #    and so the test process needs torch importable. We pull
      #    from the PyTorch CPU index so Apple Silicon gets the
      #    explicit cpu+MPS arm64 wheel rather than something the
      #    default PyPI resolver might pick up. The CPU index hosts
      #    macosx_*_arm64 wheels alongside the Linux x86_64 ones.
      # 4. unsloth-zoo from git main (NOT PyPI), WITH deps. PR-A's
      #    MLX support landed after the most recent unsloth-zoo PyPI
      #    release; the wheel still raises NotImplementedError on
      #    Apple Silicon when device_type.get_device_type() runs
      #    unguarded. Studio's own install.sh overlays unsloth-zoo
      #    from git main for the same reason. Pulling deps lets pip
      #    resolve the platform-conditional MLX-only wheels (mlx,
      #    mlx-lm, mlx-vlm gated on darwin+arm64 in unsloth-zoo's
      #    pyproject) AND the shared deps (datasets, transformers,
      #    sentencepiece, ...) that unsloth's MLX branch loads via
      #    dataprep/raw_text.py.
      # 5. unsloth -e . --no-deps so the editable install does not
      #    fight the unsloth-zoo dep set.
      #
      # All explicit pip installs are version-pinned to a single
      # released version (the latest as of 2026-05-07 within each
      # project's existing constraint range). bump alongside the rest
      # of the security audit when a new release lands.
      - name: Install deps
        run: |
          python -m pip install --upgrade pip
          pip install -r studio/backend/requirements/studio.txt
          pip install \
            'python-multipart==0.0.27' \
            'aiofiles==25.1.0' \
            'sqlalchemy==2.0.49' \
            'cryptography==48.0.0' \
            'pyyaml==6.0.3' \
            'jinja2==3.1.6' \
            'mammoth==1.12.0' \
            'unpdf==1.0.0' \
            'requests==2.33.1' \
            'typer==0.25.1' \
            'numpy==2.4.4' \
            'pytest==9.0.3' \
            'pytest-asyncio==1.3.0' \
            'httpx==0.28.1'
          pip install --index-url https://download.pytorch.org/whl/cpu \
            'torch==2.10.0'
          # github.com occasionally 500s on the git fetch; retry the
          # zoo install so a single upstream blip does not fail CI.
          for attempt in 1 2 3; do
            if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then
              break
            fi
            if [ "$attempt" -eq 3 ]; then
              echo "::error::pip install unsloth_zoo failed after 3 attempts"
              exit 1
            fi
            delay=$((5 * attempt))
            echo "::warning::unsloth_zoo install failed (attempt $attempt/3), retrying in ${delay}s..."
            sleep "$delay"
          done
          pip install -e . --no-deps

      # Real Apple Silicon sanity: confirm _IS_MLX activates on real
      # hardware with no platform spoof.
      - 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')
          "

      # Real Apple Silicon sanity: confirm every PR-A MLX-only module
      # loads against real mlx + mlx-lm + mlx-vlm wheels.
      - 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')
          "

      # Spoofed dispatch matrix. Runs on the real Mac too -- the
      # test fixture installs a MetaPathFinder that blocks
      # `import mlx.core` for "no-mlx" profiles, so the spoofs
      # faithfully simulate every supported hardware combo regardless
      # of whether mlx is installed for real.
      - name: MLX dispatch tests (3 files, 36 tests)
        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

      # Studio prebuilt llama.cpp install + GGUF inference. Drives the
      # exact path Studio's setup.sh takes on macOS: invokes
      # studio/install_llama_prebuilt.py with --published-repo
      # ggml-org/llama.cpp and --published-release-tag b9049 (the
      # latest llama.cpp release at the time this step was added; bump
      # via UNSLOTH_LLAMA_TAG / DEFAULT_LLAMA_TAG when refreshing).
      # The installer downloads llama-b9049-bin-macos-arm64.tar.gz,
      # which is the universal Apple Silicon (arm64) build -- the
      # same artifact works on M1/M2/M3/M4 because llama.cpp compiles
      # against the ARMv8.2 baseline.
      #
      # The b9049 release also publishes:
      #   - llama-b9049-bin-macos-arm64-kleidiai.tar.gz
      #     KleidiAI dispatches at runtime; on M1 it falls back where
      #     ISA features (e.g. I8MM) are missing, so this asset also
      #     runs on M1 -- Studio just doesn't choose it by default.
      #   - llama-b9049-bin-macos-x64.tar.gz
      #     Intel-only; would only run on M1 via Rosetta 2 emulation,
      #     which we explicitly avoid.
      #   - iOS XCFramework
      #     iOS-app build artifact, unrelated to a macOS desktop CI.
      #
      # After install, downloads a small published GGUF
      # (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) from HuggingFace and
      # runs the prebuilt llama-cli on it. Asserts the prompt echo
      # appears in stdout. If the install fails OR the binary exits
      # non-zero, that's an Unsloth/Studio bug.
      - name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1)
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN }}
          # install_llama_prebuilt.py hits the GitHub releases API to
          # resolve the asset URL. Anonymous calls share the runner-IP
          # rate-limit bucket and 403 quickly -- pass the workflow's
          # automatic GITHUB_TOKEN to bump us to the 5000/hr authenticated
          # bucket.
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          set -euo pipefail
          INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
          rm -rf "$INSTALL_DIR"
          # --simple-policy is required when --published-repo points
          # at upstream ggml-org/llama.cpp; that repo doesn't ship the
          # llama-prebuilt-manifest.json asset Studio's default policy
          # expects, so the simple platform-specific policy maps
          # Darwin+arm64 -> bin-macos-arm64 directly. studio/setup.sh
          # passes both --published-repo ggml-org/llama.cpp AND
          # --simple-policy automatically on macOS, so this CI step
          # exercises the same code path users hit when they run
          # `curl -fsSL https://unsloth.ai/install.sh | sh`.
          python studio/install_llama_prebuilt.py \
            --install-dir "$INSTALL_DIR" \
            --published-repo ggml-org/llama.cpp \
            --published-release-tag b9049 \
            --simple-policy

          # Studio bundles only llama-server + llama-quantize from the
          # prebuilt (not llama-cli) -- inference goes through
          # llama-server's HTTP /completion endpoint. Validate both:
          # llama-quantize --help proves the dynamic libs link, then
          # spin up llama-server and POST a /completion request on a
          # tiny published GGUF.
          LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server"
          LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize"
          [ -x "$LLAMA_SERVER" ] || { echo "::error::llama-server missing at $LLAMA_SERVER"; find "$INSTALL_DIR/build" -type f | head -40; exit 1; }
          [ -x "$LLAMA_QUANT" ]  || { echo "::error::llama-quantize missing at $LLAMA_QUANT"; exit 1; }
          echo "llama-server : $LLAMA_SERVER"
          echo "llama-quantize: $LLAMA_QUANT"
          "$LLAMA_QUANT" --help >/dev/null && echo "  llama-quantize loads OK"

          mkdir -p /tmp/ggufs
          bash .github/scripts/hf-download-with-retry.sh \
            'unsloth/gemma-3-270m-it-GGUF' \
            'gemma-3-270m-it-Q4_K_M.gguf' \
            /tmp/ggufs

          PORT=18080
          echo "=== starting llama-server on 127.0.0.1:$PORT ==="
          "$LLAMA_SERVER" \
            -m /tmp/ggufs/gemma-3-270m-it-Q4_K_M.gguf \
            --host 127.0.0.1 \
            --port "$PORT" \
            -c 256 \
            -n 16 \
            --no-warmup \
            > /tmp/llama-server.log 2>&1 &
          SERVER_PID=$!
          trap 'kill "$SERVER_PID" 2>/dev/null || true' EXIT

          # Wait for /health to come up
          for i in $(seq 1 30); do
            if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
              echo "  server up after ${i}s"
              break
            fi
            sleep 1
          done
          if ! curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
            echo "::error::llama-server never became healthy"
            tail -40 /tmp/llama-server.log
            exit 1
          fi

          PROMPT="Hello, my name is"
          echo "=== POST /completion ==="
          RESP=$(curl -sf -X POST "http://127.0.0.1:$PORT/completion" \
            -H 'Content-Type: application/json' \
            -d "{\"prompt\":\"$PROMPT\",\"n_predict\":16,\"temperature\":0,\"seed\":3407}")
          echo "raw response (head): $(echo "$RESP" | head -c 600)"
          CONTENT=$(echo "$RESP" | python -c "import json,sys; print(json.loads(sys.stdin.read()).get('content',''))")
          echo "completion content: $CONTENT"

          if [ -z "$CONTENT" ]; then
            echo "::error::llama-server /completion returned empty content"
            tail -40 /tmp/llama-server.log
            exit 1
          fi
          echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works"

      # Real MLX training + inference smoke test. Trains
      # unsloth/gemma-3-270m-it for 7 deterministic LoRA steps
      # (batch_size=2, gradient_accumulation_steps=3) on a single
      # repeated row ("<<HELLO!!>> My name is Unsloth!"), then saves
      # the trained model in 3 export formats. The `train` subcommand
      # captures per-phase timing + peak GPU + peak RSS into
      # train_metrics.json so we can detect regressions across CI runs.
      - name: MLX export round-trip — TRAIN + SAVE 3 formats
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN }}
          UNSLOTH_COMPILE_DISABLE: '1'
        run: |
          mkdir -p mlx_workdir
          python tests/studio/run_real_mlx_smoke.py train \
            --workdir "$PWD/mlx_workdir"

      # Each reload step runs in a FRESH Python process to confirm
      # the cold-start path users would hit in production also works
      # (not just the in-memory continuation of a still-running
      # trainer). FastMLXModel.from_pretrained gets called from
      # scratch; mx.random is re-seeded; per-step timing + peak
      # memory are emitted to {format}_reload_metrics.json next to
      # the saved dir.
      - name: MLX export round-trip — RELOAD LoRA (fresh process)
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN }}
          UNSLOTH_COMPILE_DISABLE: '1'
        run: |
          python tests/studio/run_real_mlx_smoke.py reload \
            --format lora \
            --dir "$PWD/mlx_workdir/lora"

      - name: MLX export round-trip — RELOAD merged_16bit (fresh process)
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN }}
          UNSLOTH_COMPILE_DISABLE: '1'
        run: |
          python tests/studio/run_real_mlx_smoke.py reload \
            --format merged \
            --dir "$PWD/mlx_workdir/merged_16bit"

      # GGUF reload uses the llama-cli binary that save_pretrained_gguf
      # built. If save_pretrained_gguf was skipped during train (e.g.
      # llama.cpp's convert_hf_to_gguf asserts on the model's tokenizer
      # vocab -- a downstream llama.cpp limitation, not an unsloth_zoo
      # bug), this step emits a workflow warning and exits 0 so the
      # LoRA + merged_16bit assertions remain the gating signal.
      - name: MLX export round-trip — RELOAD GGUF via llama-cli (fresh process)
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN }}
        run: |
          if python -c "import json,sys; m=json.load(open('mlx_workdir/train_metrics.json')); sys.exit(0 if m.get('gguf_supported') else 1)"; then
            python tests/studio/run_real_mlx_smoke.py reload \
              --format gguf \
              --dir "$PWD/mlx_workdir/gguf"
          else
            REASON=$(python -c "import json; m=json.load(open('mlx_workdir/train_metrics.json')); print(m.get('gguf_skip_reason') or 'unknown')")
            echo "::warning title=GGUF round-trip skipped::${REASON}"
            echo "GGUF export was skipped during the train phase. Reason:"
            echo "  ${REASON}"
            echo "Continuing without failing the job; the LoRA + merged_16bit"
            echo "reload assertions are still gating this PR."
          fi

      # Print all metrics JSON files so regressions are visible in the
      # job log. always() so we get telemetry even if a reload step
      # asserted gibberish.
      - name: MLX export round-trip — aggregate metrics
        if: always()
        run: |
          for f in mlx_workdir/train_metrics.json \
                   mlx_workdir/lora_reload_metrics.json \
                   mlx_workdir/merged_reload_metrics.json \
                   mlx_workdir/gguf_reload_metrics.json; do
            echo "=== $f ==="
            cat "$f" 2>/dev/null || echo "(missing)"
            echo
          done