Declare coremltools and scikit-learn only where they work (pytorch#22090)

### Summary

Two dependency markers in the wheel do not match where the packages they
name are usable.

**1. coremltools on Linux aarch64.** coremltools is declared on Darwin
and on Linux. Linux is deliberate: the Core ML export flow runs there,
so a model can be lowered for Apple hardware from a Linux machine. The
problem is that the marker cannot tell the two Linux architectures
apart, and coremltools publishes no build for Linux aarch64.

On that architecture pip falls back to the source archive and produces a
pure-Python install with none of the compiled extensions. Measured with
coremltools 9.0 on a Linux aarch64 machine:

```
Tag: py3-none-any
Root-Is-Purelib: true
$ find <site-packages>/coremltools -name '*.so'      # no output at all
IMPORT FAIL coremltools.libmilstoragepython ModuleNotFoundError
IMPORT FAIL coremltools.libcoremlpython     ModuleNotFoundError
IMPORT FAIL coremltools.libmodelpackage     ModuleNotFoundError
```

It imports, it reports version 9.0, and then it cannot write a model,
because `libmilstoragepython` is what stores the weights of an
mlprogram. The same version on Linux x86_64 installs
`coremltools-9.0-cp3XX-none-manylinux1_x86_64.whl`, which carries
`libmilstoragepython.so` and `libmodelpackage.so`. So the Linux half of
the marker is narrowed to x86_64, and users on other Linux architectures
are no longer handed a package that looks installed and cannot do the
job.

**2. scikit-learn has no marker at all.** It is there for Core ML
palettization, so it should follow coremltools, but instead it installs
everywhere, including Windows, where coremltools is already excluded. A
Windows install therefore carries scikit-learn and scipy for a backend
it cannot use. This gives it the same marker.

### Effect by platform

For coremltools, macOS and Linux x86_64 are unchanged at every supported
Python version. The Python 3.14 exclusion was already in the marker
before this change. What changes is Linux on any architecture other than
x86_64.

For scikit-learn, it is no longer installed on Windows, on non-x86_64
Linux, or on Python 3.14, and scipy no longer arrives behind it there.
Nothing in the shipped library imports either one at module scope. The
one library module that uses scipy, the TurboQuant codebook solver in
`extension/llm/modules/turboquant/codebook.py`, already imports it
inside the function and raises a message telling the user to install it.
That module is not example-only:
`extension/llm/modules/turboquant/kv_cache.py` imports it and calls it
from `TurboQuantKVCache.__init__`, and
`backends/mlx/llm/turboquant_cache.py` re-exports that class. See the
install-step note below.

There is no new extra. There is one new install step, on the platforms
where scipy
stops arriving: a wheel user on Windows, non-x86_64 Linux, or Python
3.14 who
constructs `TurboQuantKVCache` now needs `pip install scipy`, because
`TurboQuantKVCache.__init__` calls the codebook solver unconditionally.
scipy was
never declared in `setup.py`; it only ever arrived as scikit-learn's
transitive
dependency, so gating scikit-learn gates it too. The guarded import
names scipy and
the install command, so the failure is actionable rather than a bare
`ModuleNotFoundError`, but it is a real change for those users.

### Follow-on changes

- `conftest.py` drops `backends/apple/coreml` from collection wherever
coremltools is not declared. Every test file there that defines tests
imports coremltools at module scope, so collection would otherwise fail.
Windows was already covered by an existing rule; this extends it to the
other Linux architectures.

- A new test, `.ci/scripts/tests/test_coreml_markers.py`. The rule for
where coremltools can be installed is now written twice, once as a PEP
508 marker in `setup.py` and once as a Python condition in
`conftest.py`, in two different languages, with only a comment holding
them together. The test reads both out of their real files with `ast`
rather than restating them a third time, and checks that they agree
across nine platforms and six Python versions. If one is changed and the
other is not, this fails in CI instead of failing at collection with a
bare `ModuleNotFoundError`.

- scikit-learn and scipy move to `requirements-examples.txt`. Two
example scripts import scikit-learn,
`examples/openvino/aot_optimize_and_infer.py` and
`examples/qualcomm/scripts/mobilebert_fine_tune.py`, and three Qualcomm
example scripts import scipy directly. Both used to arrive by accident,
scikit-learn as an unmarked base dependency and scipy as its transitive,
and neither was ever declared where it is used.
`requirements-examples.txt` is where example-only dependencies already
live, next to timm and transformers, and it is what
`install_requirements.py` and the example CI jobs install. Putting
scikit-learn in the `openvino` extra instead would not have been enough,
since that script also needs timm and transformers, which the extra does
not carry.

### Test plan

Evaluated both markers with `packaging` across macOS, Linux x86_64,
Linux aarch64, Windows and Python 3.10 through 3.15:

```
coremltools
   linux x86_64 py3.10    -> True
   linux aarch64 py3.10   -> False
   linux x86_64 py3.13    -> True
   linux x86_64 py3.14    -> False
   darwin arm64 py3.12    -> True
   windows AMD64 py3.12   -> False
```

scikit-learn now evaluates identically, and the `conftest.py` condition
agrees with the coremltools marker on every one of those combinations.
That is what the new test covers, and it passes:

```
167 passed
```

Confirmed the test can fail. Widening the `conftest.py` condition to
also accept Linux aarch64 turns four cases red with a message naming the
platform and the Python version:

```
AssertionError: on Linux/aarch64 with Python 3.13 setup.py does not declare
coremltools while conftest.py collects the Core ML tests, so pytest will either
fail at collection on a missing import or silently skip tests that could have run
```

Confirmed the crippled install by hand on a Linux aarch64 machine, as
shown above, and confirmed on Linux x86_64 that the published wheel for
the same version ships the compiled extensions.

The unit test jobs that collect the Core ML tests are the Linux x86_64
job on Python 3.10 and the macOS arm64 job on Python 3.11. Both are
platforms where coremltools is still declared, so those tests collect
and run exactly as before. The Windows job already skipped the whole
Apple directory.

---------

Co-authored-by: r <r@e>

Author: shoumikhin

.ci/scripts/tests/test_coreml_markers.py

@@ -0,0 +1,322 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Unit tests that the Core ML availability rules agree.
+
+Where coremltools can be installed is written down three times: as a dependency marker in
+setup.py, as a condition in conftest.py that drops the Core ML tests from collection, and
+as `is_supported_platform_for_coreml_lowering` in export/utils.py, which callers use as an
+import guard. They are written in different languages, PEP 508 and Python, so nothing but a
+comment keeps them together. When they drift, the wheel declares a package on a platform
+the test run assumes is absent, or the other way round, and pytest fails at collection with
+a bare ModuleNotFoundError.
+
+The rules are read out of the three files rather than restated here, so a change to any one
+of them is exercised by this test instead of being duplicated a fourth time.
+"""
+
+import ast
+import fnmatch
+import re
+from pathlib import Path
+from types import SimpleNamespace
+from typing import NamedTuple
+
+import pytest
+from packaging.markers import Marker
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+
+# platform_system, platform_machine, sys_platform
+PLATFORMS = [
+    ("Darwin", "arm64", "darwin"),
+    ("Darwin", "x86_64", "darwin"),
+    ("Linux", "x86_64", "linux"),
+    ("Linux", "aarch64", "linux"),
+    ("Linux", "armv7l", "linux"),
+    ("Linux", "ppc64le", "linux"),
+    ("Linux", "s390x", "linux"),
+    ("Windows", "AMD64", "win32"),
+    ("Windows", "ARM64", "win32"),
+]
+
+# Every version in requires-python, plus one past the end.
+PYTHONS = [(3, 10), (3, 11), (3, 12), (3, 13), (3, 14), (3, 15)]
+
+
+class _VersionInfo(NamedTuple):
+    """Stands in for sys.version_info: compares as a tuple, has .major/.minor."""
+
+    major: int
+    minor: int
+
+
+def _normalise(name: str) -> str:
+    # PEP 503: runs of -, _ and . collapse to a single -, then casefold.
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+def _project_name(requirement: str) -> str:
+    # Strip everything a PEP 508 requirement can carry after the project name: an extras
+    # list, a version specifier, an environment marker, or a URL.
+    head = re.split(r"[\[<>=!~;@\s]", requirement.strip(), maxsplit=1)[0]
+    return _normalise(head)
+
+
+def _base_dependency(name: str) -> str:
+    """Return the requirement string for `name` from setup.py's _base_dependencies()."""
+    # Compare PEP 503 normalised project names rather than using startswith, which would
+    # match a longer name that merely begins with this one (scikit-learn against
+    # scikit-learn-intelex) and would also match the function's own docstring, since that
+    # is just another string constant in the body.
+    wanted = _normalise(name)
+    tree = ast.parse((REPO_ROOT / "setup.py").read_text())
+    for node in ast.walk(tree):
+        if not (
+            isinstance(node, ast.FunctionDef) and node.name == "_base_dependencies"
+        ):
+            continue
+        for constant in ast.walk(node):
+            if not (
+                isinstance(constant, ast.Constant) and isinstance(constant.value, str)
+            ):
+                continue
+            if _project_name(constant.value) == wanted:
+                return constant.value
+    raise AssertionError(f"no {name} requirement in setup.py _base_dependencies()")
+
+
+def _conftest_condition() -> str:
+    """Return the source of conftest.py's _coremltools_is_declared expression."""
+    tree = ast.parse((REPO_ROOT / "conftest.py").read_text())
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Assign) and any(
+            isinstance(target, ast.Name) and target.id == "_coremltools_is_declared"
+            for target in node.targets
+        ):
+            return ast.unparse(node.value)
+    raise AssertionError("no _coremltools_is_declared assignment in conftest.py")
+
+
+def _coreml_ignore_globs() -> list[str]:
+    """Return the globs conftest.py adds when _coremltools_is_declared is false.
+
+    Reads the `if not _coremltools_is_declared:` statement rather than the module's
+    resulting state, because importing conftest.py here would evaluate the condition
+    for the interpreter running this test instead of for the platforms under test.
+    """
+    tree = ast.parse((REPO_ROOT / "conftest.py").read_text())
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.If):
+            continue
+        test = node.test
+        # `if not _coremltools_is_declared:`
+        if not (
+            isinstance(test, ast.UnaryOp)
+            and isinstance(test.op, ast.Not)
+            and isinstance(test.operand, ast.Name)
+            and test.operand.id == "_coremltools_is_declared"
+        ):
+            continue
+        globs: list[str] = []
+        for stmt in node.body:
+            # `collect_ignore_glob += [...]` or `collect_ignore_glob = [...]`
+            targets = (
+                [stmt.target]
+                if isinstance(stmt, ast.AugAssign)
+                else stmt.targets if isinstance(stmt, ast.Assign) else []
+            )
+            if not any(
+                isinstance(t, ast.Name) and t.id == "collect_ignore_glob"
+                for t in targets
+            ):
+                continue
+            for element in ast.walk(stmt.value):
+                if isinstance(element, ast.Constant) and isinstance(element.value, str):
+                    globs.append(element.value)
+        if globs:
+            return globs
+    raise AssertionError(
+        "conftest.py does not gate collect_ignore_glob on `not _coremltools_is_declared`, "
+        "so the Core ML tests are collected even where coremltools is not installed"
+    )
+
+
+def _marker_says_installed(
+    requirement: str, system, machine, sys_platform, python
+) -> bool:
+    _, _, marker = requirement.partition(";")
+    assert marker.strip(), f"{requirement} has no environment marker"
+    return Marker(marker).evaluate(
+        {
+            "platform_system": system,
+            "platform_machine": machine,
+            "sys_platform": sys_platform,
+            "python_version": f"{python[0]}.{python[1]}",
+            "python_full_version": f"{python[0]}.{python[1]}.0",
+        }
+    )
+
+
+def _conftest_says_installed(machine, sys_platform, python) -> bool:
+    return bool(
+        eval(  # the expression is read out of conftest.py, not taken from input
+            compile(_conftest_condition(), "<conftest.py>", "eval"),
+            {
+                "sys": SimpleNamespace(platform=sys_platform, version_info=python),
+                "platform": SimpleNamespace(machine=lambda: machine),
+            },
+        )
+    )
+
+
+def _support_helper_says_supported(system, machine, python) -> bool:
+    """Evaluate export/utils.py's is_supported_platform_for_coreml_lowering().
+
+    The function is extracted and run against stub platform/sys modules rather than
+    imported, because importing export.utils pulls in torch and would answer for the
+    interpreter running this test instead of for the platform under test.
+    """
+    tree = ast.parse((REPO_ROOT / "export" / "utils.py").read_text())
+    for node in tree.body:
+        if (
+            isinstance(node, ast.FunctionDef)
+            and node.name == "is_supported_platform_for_coreml_lowering"
+        ):
+            namespace = {
+                "platform": SimpleNamespace(
+                    system=lambda: system, machine=lambda: machine
+                ),
+                # A namedtuple-like value so both `>= (3, 14)` and `.major`/`.minor`
+                # work, matching what sys.version_info supports.
+                "sys": SimpleNamespace(version_info=_VersionInfo(*python)),
+                "logging": SimpleNamespace(info=lambda *a, **k: None),
+            }
+            module = ast.Module(body=[node], type_ignores=[])
+            exec(compile(module, "<export/utils.py>", "exec"), namespace)
+            return bool(namespace["is_supported_platform_for_coreml_lowering"]())
+    raise AssertionError(
+        "no is_supported_platform_for_coreml_lowering in export/utils.py"
+    )
+
+
+@pytest.mark.parametrize("system,machine,sys_platform", PLATFORMS)
+@pytest.mark.parametrize("python", PYTHONS)
+def test_conftest_matches_the_coremltools_marker(system, machine, sys_platform, python):
+    declared = _marker_says_installed(
+        _base_dependency("coremltools"), system, machine, sys_platform, python
+    )
+    collected = _conftest_says_installed(machine, sys_platform, python)
+    assert declared == collected, (
+        f"on {system}/{machine} with Python {python[0]}.{python[1]} setup.py "
+        f"{'declares' if declared else 'does not declare'} coremltools while conftest.py "
+        f"{'collects' if collected else 'skips'} the Core ML tests, so pytest will either fail "
+        "at collection on a missing import or silently skip tests that could have run"
+    )
+
+
+@pytest.mark.parametrize("system,machine,sys_platform", PLATFORMS)
+@pytest.mark.parametrize("python", PYTHONS)
+def test_scikit_learn_follows_coremltools(system, machine, sys_platform, python):
+    # scikit-learn is there for coremltools' palettization, so it is only useful where
+    # coremltools is. Nothing else in the wheel imports scikit-learn itself; scipy, which
+    # used to arrive as its transitive dependency, is declared in requirements-examples.txt
+    # for the code that does import it.
+    assert _marker_says_installed(
+        _base_dependency("scikit-learn"), system, machine, sys_platform, python
+    ) == _marker_says_installed(
+        _base_dependency("coremltools"), system, machine, sys_platform, python
+    )
+
+
+def test_the_platforms_core_ml_is_used_on_are_still_covered():
+    # A guard on the parametrized tests above: they would pass just as happily if both
+    # rules said "never", which would quietly stop running the Core ML tests everywhere.
+    requirement = _base_dependency("coremltools")
+    assert _marker_says_installed(requirement, "Darwin", "arm64", "darwin", (3, 12))
+    assert _marker_says_installed(requirement, "Linux", "x86_64", "linux", (3, 12))
+    assert not _marker_says_installed(requirement, "Linux", "aarch64", "linux", (3, 12))
+    assert not _marker_says_installed(requirement, "Windows", "AMD64", "win32", (3, 12))
+
+
+def test_conftest_acts_on_the_condition():
+    # The tests above only compare the two rules as expressions. They would still pass if
+    # conftest.py computed _coremltools_is_declared and then never used it, which is the
+    # one failure that reaches contributors as a ModuleNotFoundError at collection.
+    #
+    # Membership, not equality: adding a second correct glob is a legitimate change, and
+    # test_the_ignore_globs_cover_every_core_ml_test_module below already catches a glob
+    # that stops matching.
+    assert "backends/apple/coreml/**" in _coreml_ignore_globs()
+
+
+@pytest.mark.parametrize("system,machine,sys_platform", PLATFORMS)
+@pytest.mark.parametrize("python", PYTHONS)
+def test_export_support_helper_never_claims_more_than_the_marker(
+    system, machine, sys_platform, python
+):
+    # export/utils.py is the third place this rule is written down. Callers use it as an
+    # import guard (export/target_recipes.py raises ValueError when it says no), so it must
+    # never claim support where coremltools is not declared, or the guard is bypassed and
+    # the import raises ModuleNotFoundError instead.
+    #
+    # The reverse is allowed: the helper is deliberately narrower than the marker, because
+    # coremltools installs on Darwin x86_64 but lowering is only supported on Apple silicon.
+    declared = _marker_says_installed(
+        _base_dependency("coremltools"), system, machine, sys_platform, python
+    )
+    supported = _support_helper_says_supported(system, machine, python)
+    assert not (supported and not declared), (
+        f"on {system}/{machine} with Python {python[0]}.{python[1]} "
+        "is_supported_platform_for_coreml_lowering() reports support while setup.py does "
+        "not declare coremltools, so the import guard in export/target_recipes.py is "
+        "bypassed and importing coremltools raises ModuleNotFoundError"
+    )
+
+
+def test_the_support_helper_still_says_yes_where_core_ml_is_used():
+    # A guard on the one-directional test above, which a helper that always returned False
+    # would satisfy trivially.
+    assert _support_helper_says_supported("Darwin", "arm64", (3, 12))
+    assert _support_helper_says_supported("Linux", "x86_64", (3, 12))
+    assert not _support_helper_says_supported("Windows", "AMD64", (3, 12))
+
+
+def test_the_support_helper_tracks_the_marker_on_python_version():
+    # The one-directional test above is deliberately loose on architecture, because
+    # coremltools installs on Darwin x86_64 while lowering needs Apple silicon. It must not
+    # be loose on the Python version too: if the marker starts declaring coremltools on a
+    # version the helper still refuses, every caller keeps taking the "not supported" branch
+    # on a platform where Core ML now works, and nothing else in this file notices.
+    #
+    # So pin the helper to the marker's own bound rather than to a literal, and read the
+    # bound out of setup.py so widening the marker fails here until the helper follows.
+    marker = _base_dependency("coremltools")
+    for python in PYTHONS:
+        declared = _marker_says_installed(marker, "Darwin", "arm64", "darwin", python)
+        supported = _support_helper_says_supported("Darwin", "arm64", python)
+        assert supported == declared, (
+            f"on Darwin/arm64 with Python {python[0]}.{python[1]} setup.py "
+            f"{'declares' if declared else 'does not declare'} coremltools but "
+            f"is_supported_platform_for_coreml_lowering() reports "
+            f"{'support' if supported else 'no support'}; the version bound in "
+            "export/utils.py has drifted from the marker in setup.py"
+        )
+
+
+def test_the_ignore_globs_cover_every_core_ml_test_module():
+    # A glob that no longer matches would silence nothing. Checked against the files on
+    # disk so that moving the Core ML tests without updating conftest.py fails here.
+    coreml_tests = sorted(
+        path.relative_to(REPO_ROOT).as_posix()
+        for path in (REPO_ROOT / "backends/apple/coreml").rglob("test_*.py")
+    )
+    assert coreml_tests, "no Core ML test modules found, so this guard proves nothing"
+    globs = _coreml_ignore_globs()
+    for test in coreml_tests:
+        assert any(
+            fnmatch.fnmatch(test, glob) for glob in globs
+        ), f"{test} is not covered by conftest.py's ignore globs {globs}"

.wiki/backends/coreml/overview.md

@@ -273,4 +273,4 @@ The `to_edge_with_preserved_ops` API (experimental) allows preserving ops like `
 
 ## CoreML Export on Linux
 
-CoreML export now works on Linux (as of v0.6+). The `coremltools` package can run on Linux for AOT compilation, though runtime execution still requires macOS/iOS. [Source: #9800]
+CoreML export now works on Linux x86_64 (as of v0.6+). The `coremltools` package can run there for AOT compilation, though runtime execution still requires macOS/iOS. It is x86_64 only because coremltools publishes no build for any other Linux architecture. [Source: #9800]

conftest.py

@@ -1,4 +1,5 @@
 import hashlib
+import platform
 import sys
 
 import torch
@@ -14,6 +15,22 @@
         "backends/apple/**",
     ]
 
+# Every test file under backends/apple/coreml that defines tests imports coremltools at module
+# scope, so collection fails wherever the wheel does not declare it. Windows is already covered
+# above; what is left is any Linux that is not x86_64, since coremltools publishes no build for
+# those, and Python 3.14, for which it publishes no build on any platform. Keep this condition in
+# sync with the coremltools marker in setup.py; .ci/scripts/tests/test_coreml_markers.py checks
+# that the two agree.
+_coremltools_is_declared = (
+    sys.platform == "darwin"
+    or (sys.platform.startswith("linux") and platform.machine() == "x86_64")
+) and sys.version_info < (3, 14)
+
+if not _coremltools_is_declared:
+    collect_ignore_glob += [
+        "backends/apple/coreml/**",
+    ]
+
 
 def pytest_runtest_setup(item):
     # Set a stable seed for each test based on a hash of the test name.

export/utils.py

@@ -7,6 +7,7 @@
 # pyre-strict
 import logging
 import platform
+import sys
 
 import torch
 
@@ -20,6 +21,17 @@ def is_supported_platform_for_coreml_lowering() -> bool:
     system = platform.system()
     machine = platform.machine().lower()
 
+    # coremltools has no wheel for 3.14 yet, so setup.py does not declare it
+    # there. Callers use this as an import guard, so reporting the platform as
+    # supported would turn a handled "not supported" into a ModuleNotFoundError.
+    # Keep this in step with the coremltools marker in setup.py.
+    if sys.version_info >= (3, 14):
+        logging.info(
+            f"Unsupported Python for CoreML: {sys.version_info.major}."
+            f"{sys.version_info.minor}"
+        )
+        return False
+
     # Check for Linux x86_64
     if system == "Linux" and machine == "x86_64":
         return True

requirements-examples.txt

@@ -1,6 +1,13 @@
 # pip packages needed to run examples.
 # TODO: Make each example publish its own requirements.txt
 datasets == 3.6.0 # 4.0.0 deprecates trust_remote_code and load scripts. For now pin to 3.6.0
+# examples/openvino/aot_optimize_and_infer.py scores accuracy with scikit-learn and
+# examples/qualcomm/scripts/mobilebert_fine_tune.py uses it too. Three Qualcomm example scripts
+# import scipy directly. Both used to arrive by accident, scikit-learn as an unmarked base
+# dependency and scipy as its transitive, and neither was ever declared. Now that scikit-learn
+# follows coremltools they have to be named where they are actually used.
+scikit-learn >= 1.7.1
+scipy
 timm == 1.0.7
 torchsr == 1.0.4
 # torchtune caps pyarrow below 21, and no pyarrow in that range ships a cp314 wheel,