Skip to content

Commit

Permalink
Expand ranges when reading thread_siblings_list
Browse files Browse the repository at this point in the history
Handle both ranges and individual values
  • Loading branch information
vzhurba01 committed Oct 3, 2023
1 parent 0d0d427 commit e212634
Showing 1 changed file with 12 additions and 3 deletions.
15 changes: 12 additions & 3 deletions legate/util/system.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,13 @@ def LIB_PATH(self) -> str:
def cpus(self) -> tuple[CPUInfo, ...]:
"""A list of CPUs on the system."""

def expand_range(value: str) -> tuple[int, ...]:
if "-" not in value:
return tuple((int(value),))
start, stop = value.split("-")

return tuple(x for x in range(int(start), int(stop) + 1))

N = multiprocessing.cpu_count()

if sys.platform == "darwin":
Expand All @@ -98,9 +105,11 @@ def cpus(self) -> tuple[CPUInfo, ...]:
line = open(
f"/sys/devices/system/cpu/cpu{i}/topology/thread_siblings_list" # noqa E501
).read()
sibling_sets.add(
tuple(sorted(int(x) for x in line.strip().split(",")))
)

flattened = []
for x in (expand_range(value) for value in line.strip().split(",")):
flattened += [y for y in x]
sibling_sets.add(tuple(sorted(flattened)))
return tuple(
CPUInfo(siblings) for siblings in sorted(sibling_sets)
)
Expand Down

0 comments on commit e212634

Please sign in to comment.