Skip to content

Commit 3602e9e

Browse files
committed
refactor(test): Eliminate homemade CPU affinity getter
Python psutil.Process.cpu_affinity() returns the current CPU affinity as a list of integers. So we don't need to have our own CPU affinity getter implementation. Signed-off-by: Takahiro Itazuri <[email protected]>
1 parent 8d2cf59 commit 3602e9e

File tree

1 file changed

+6
-66
lines changed

1 file changed

+6
-66
lines changed

tests/framework/utils.py

Lines changed: 6 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
# SPDX-License-Identifier: Apache-2.0
33
"""Generic utility functions that are used in the framework."""
44
import errno
5-
import functools
65
import json
76
import logging
87
import os
@@ -15,7 +14,6 @@
1514
import typing
1615
from collections import defaultdict, namedtuple
1716
from contextlib import contextmanager
18-
from pathlib import Path
1917
from typing import Dict
2018

2119
import packaging.version
@@ -162,72 +160,14 @@ def _cpus(cls):
162160
163161
See this issue for details:
164162
https://github.com/moby/moby/issues/20770.
165-
"""
166-
# The real processor map is found at different paths based on cgroups version:
167-
# - cgroupsv1: /cpuset.cpus
168-
# - cgroupsv2: /cpuset.cpus.effective
169-
# For more details, see https://docs.kernel.org/admin-guide/cgroup-v2.html#cpuset-interface-files
170-
for path in [
171-
Path("/sys/fs/cgroup/cpuset/cpuset.cpus"),
172-
Path("/sys/fs/cgroup/cpuset.cpus.effective"),
173-
]:
174-
if path.exists():
175-
return ListFormatParser(path.read_text("ascii").strip()).parse()
176-
177-
raise RuntimeError("Could not find cgroups cpuset")
178-
179-
180-
class ListFormatParser:
181-
"""Parser class for LIST FORMAT strings."""
182-
183-
def __init__(self, content):
184-
"""Initialize the parser with the content."""
185-
self._content = content.strip()
186-
187-
@classmethod
188-
def _is_range(cls, rng):
189-
"""Return true if the parser content is a range.
190-
191-
E.g ranges: 0-10.
192-
"""
193-
match = re.search("([0-9][1-9]*)-([0-9][1-9]*)", rng)
194-
# Group is a singular value.
195-
return match is not None
196163
197-
@classmethod
198-
def _range_to_list(cls, rng):
199-
"""Return a range of integers based on the content.
200-
201-
The content respects the LIST FORMAT defined in the
202-
cpuset documentation.
203-
See: https://man7.org/linux/man-pages/man7/cpuset.7.html.
164+
Note that this method is called only once when `CpuMap.arr` is
165+
initialized.
204166
"""
205-
ends = rng.split("-")
206-
if len(ends) != 2:
207-
return []
208-
209-
return list(range(int(ends[0]), int(ends[1]) + 1))
210-
211-
def parse(self):
212-
"""Parse list formats for cpuset and mems.
213-
214-
See LIST FORMAT here:
215-
https://man7.org/linux/man-pages/man7/cpuset.7.html.
216-
"""
217-
if len(self._content) == 0:
218-
return []
219-
220-
groups = self._content.split(",")
221-
arr = set()
222-
223-
def func(acc, cpu):
224-
if ListFormatParser._is_range(cpu):
225-
acc.update(ListFormatParser._range_to_list(cpu))
226-
else:
227-
acc.add(int(cpu))
228-
return acc
229-
230-
return list(functools.reduce(func, groups, arr))
167+
# https://psutil.readthedocs.io/en/latest/#psutil.Process.cpu_affinity
168+
# > If no argument is passed it returns the current CPU affinity as a
169+
# > list of intergers.
170+
return psutil.Process().cpu_affinity()
231171

232172

233173
class CmdBuilder:

0 commit comments

Comments
 (0)