[Bugfix][CPU] Fix macOS compatibility broken by #36487 (#38970)

Signed-off-by: Ziming (2imi9) <148090931+2imi9@users.noreply.github.com>
This commit is contained in:
Ziming Qi
2026-04-04 10:05:58 -04:00
committed by GitHub
parent 93726b2a1c
commit 2a36d8fb72
2 changed files with 36 additions and 5 deletions

View File

@@ -381,7 +381,12 @@ class CpuPlatform(Platform):
def get_global_cpu_mask(cls) -> set[int]:
# get global cpu mask
if cls.global_cpu_mask is None:
cls.global_cpu_mask = os.sched_getaffinity(0)
if hasattr(os, "sched_getaffinity"):
cls.global_cpu_mask = os.sched_getaffinity(0)
else:
# macOS does not support sched_getaffinity
cpu_count = os.cpu_count() or 1
cls.global_cpu_mask = set(range(cpu_count))
return cls.global_cpu_mask
@classmethod

View File

@@ -7,6 +7,7 @@ Copyright (c) 2026 Cambridge Greys Ltd
import json
import os
import platform
import subprocess
@@ -38,10 +39,37 @@ def parse_mask(mask):
return set(result)
def _get_default_affinity() -> set[int]:
"""Get the set of CPUs the process is allowed to run on."""
if hasattr(os, "sched_getaffinity"):
return os.sched_getaffinity(0)
# macOS does not support sched_getaffinity; fall back to cpu_count
cpu_count = os.cpu_count() or 1
return set(range(cpu_count))
def _get_cpu_topology_json() -> bytes:
"""Get CPU topology as JSON.
On Linux this uses ``lscpu -Je``. On other platforms (e.g. macOS) we
synthesize a simple topology where every logical CPU is its own core
on NUMA node 0, which is sufficient for the OMP place-list builder.
"""
if platform.system() == "Linux":
return subprocess.run(["lscpu", "-Je"], check=True, capture_output=True).stdout
# Fallback for non-Linux (macOS, etc.)
cpu_count = os.cpu_count() or 1
cpus = []
for i in range(cpu_count):
cpus.append({"cpu": str(i), "core": str(i), "node": "0"})
return json.dumps({"cpus": cpus}).encode()
def enumerate_resources(resource_map, mask=None, allowed=None):
"""Enumerate system resources"""
if allowed is None:
allowed = os.sched_getaffinity(0)
allowed = _get_default_affinity()
if mask is not None:
allowed = allowed & mask
@@ -140,9 +168,7 @@ class OMPProcessManager:
else:
masks = [None]
if mock is None:
data = subprocess.run(
["lscpu", "-Je"], check=True, capture_output=True
).stdout
data = _get_cpu_topology_json()
else:
with open(mock, mode="rb") as jf:
data = jf.read()