Unverified Commit 2a36d8fb authored by Ziming Qi's avatar Ziming Qi Committed by GitHub
Browse files

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


Signed-off-by: default avatarZiming (2imi9) <148090931+2imi9@users.noreply.github.com>
parent 93726b2a
...@@ -381,7 +381,12 @@ class CpuPlatform(Platform): ...@@ -381,7 +381,12 @@ class CpuPlatform(Platform):
def get_global_cpu_mask(cls) -> set[int]: def get_global_cpu_mask(cls) -> set[int]:
# get global cpu mask # get global cpu mask
if cls.global_cpu_mask is None: if cls.global_cpu_mask is None:
if hasattr(os, "sched_getaffinity"):
cls.global_cpu_mask = os.sched_getaffinity(0) 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 return cls.global_cpu_mask
@classmethod @classmethod
......
...@@ -7,6 +7,7 @@ Copyright (c) 2026 Cambridge Greys Ltd ...@@ -7,6 +7,7 @@ Copyright (c) 2026 Cambridge Greys Ltd
import json import json
import os import os
import platform
import subprocess import subprocess
...@@ -38,10 +39,37 @@ def parse_mask(mask): ...@@ -38,10 +39,37 @@ def parse_mask(mask):
return set(result) 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): def enumerate_resources(resource_map, mask=None, allowed=None):
"""Enumerate system resources""" """Enumerate system resources"""
if allowed is None: if allowed is None:
allowed = os.sched_getaffinity(0) allowed = _get_default_affinity()
if mask is not None: if mask is not None:
allowed = allowed & mask allowed = allowed & mask
...@@ -140,9 +168,7 @@ class OMPProcessManager: ...@@ -140,9 +168,7 @@ class OMPProcessManager:
else: else:
masks = [None] masks = [None]
if mock is None: if mock is None:
data = subprocess.run( data = _get_cpu_topology_json()
["lscpu", "-Je"], check=True, capture_output=True
).stdout
else: else:
with open(mock, mode="rb") as jf: with open(mock, mode="rb") as jf:
data = jf.read() data = jf.read()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment