Harden macOS metallib discovery

Author: KSemenenko

.github/workflows/ci.yml

@@ -207,60 +207,94 @@ jobs:
           import pathlib
           import shutil
           import sys
+          from typing import Iterable, Optional
 
           try:
               import mlx  # type: ignore
           except ImportError:
               print("::error::The 'mlx' Python package is not installed; cannot locate mlx.metallib.")
               sys.exit(1)
 
-          kernels_dir = None
+          search_dirs: list[pathlib.Path] = []
+          package_dir: Optional[pathlib.Path] = None
+
+          try:
+              spec = importlib.util.find_spec("mlx.backend.metal.kernels")
+          except ModuleNotFoundError:
+              spec = None
 
-          spec = importlib.util.find_spec("mlx.backend.metal.kernels")
           if spec and spec.origin:
-              candidate_origin = spec.origin
-              if candidate_origin:
-                  candidate = pathlib.Path(candidate_origin).resolve().parent
-                  if candidate.exists():
-                      kernels_dir = candidate
+              candidate = pathlib.Path(spec.origin).resolve().parent
+              if candidate.exists():
+                  search_dirs.append(candidate)
 
-          if kernels_dir is None:
+          try:
+              resource = resources.files("mlx.backend.metal") / "kernels"
+          except (ModuleNotFoundError, AttributeError):
+              resource = None
+          else:
               try:
-                  resource = resources.files("mlx.backend.metal") / "kernels"
-              except (ModuleNotFoundError, AttributeError):
-                  resource = None
-              if resource is not None:
-                  try:
-                      with resources.as_file(resource) as extracted:
-                          if extracted is not None:
-                              extracted_path = pathlib.Path(extracted)
-                              if extracted_path.exists():
-                                  kernels_dir = extracted_path
-                  except FileNotFoundError:
-                      pass
-
-          if kernels_dir is None:
-              package_file = getattr(mlx, "__file__", None)
-              if package_file:
-                  package_dir = pathlib.Path(package_file).resolve().parent
-                  candidate = package_dir / "backend" / "metal" / "kernels"
-                  if candidate.exists():
-                      kernels_dir = candidate
-
-          if kernels_dir is None or not kernels_dir.exists():
-              print("::error::Could not locate the MLX metal kernels directory; checked module spec and importlib resources.")
+                  with resources.as_file(resource) as extracted:
+                      if extracted:
+                          extracted_path = pathlib.Path(extracted).resolve()
+                          if extracted_path.exists():
+                              search_dirs.append(extracted_path)
+              except (FileNotFoundError, RuntimeError):
+                  pass
+
+          package_file = getattr(mlx, "__file__", None)
+          if package_file:
+              package_dir = pathlib.Path(package_file).resolve().parent
+              search_dirs.extend(
+                  [
+                      package_dir / "backend" / "metal" / "kernels",
+                      package_dir / "backend" / "metal",
+                      package_dir,
+                  ]
+              )
+
+          ordered_dirs: list[pathlib.Path] = []
+          seen: set[pathlib.Path] = set()
+          for candidate in search_dirs:
+              if not candidate:
+                  continue
+              candidate = candidate.resolve()
+              if candidate in seen:
+                  continue
+              seen.add(candidate)
+              ordered_dirs.append(candidate)
+
+          def iter_metallibs(dirs: Iterable[pathlib.Path]):
+              for directory in dirs:
+                  if not directory.exists():
+                      continue
+                  preferred = directory / "mlx.metallib"
+                  if preferred.exists():
+                      yield preferred
+                      continue
+                  for alternative in sorted(directory.glob("*.metallib")):
+                      yield alternative
+
+          src = next(iter_metallibs(ordered_dirs), None)
+
+          if src is None and package_dir and package_dir.exists():
+              for candidate in package_dir.rglob("mlx.metallib"):
+                  src = candidate
+                  print(f"::warning::Resolved metallib via recursive search under {package_dir}")
+                  break
+
+          if src is None and package_dir and package_dir.exists():
+              for candidate in sorted(package_dir.rglob("*.metallib")):
+                  src = candidate
+                  print(f"::warning::Using metallib {candidate.name} discovered via package-wide search")
+                  break
+
+          if src is None:
+              print("::error::Could not locate any mlx.metallib artifacts within the installed mlx package.")
               sys.exit(1)
 
-          preferred = kernels_dir / "mlx.metallib"
-          if preferred.exists():
-              src = preferred
-          else:
-              candidates = sorted(kernels_dir.glob("*.metallib"))
-              if not candidates:
-                  print(f"::error::No metallib files were found under {kernels_dir}")
-                  sys.exit(1)
-              src = candidates[0]
-              print(f"::warning::Defaulting to metallib file {src.name}")
+          if src.name != "mlx.metallib":
+              print(f"::warning::Using metallib {src.name} from {src.parent}")
 
           dest = pathlib.Path("artifacts/native/osx-arm64/mlx.metallib").resolve()
           dest.parent.mkdir(parents=True, exist_ok=True)