1#!/usr/bin/env python3

2# This file is part of Checkbox.

3#

4# Copyright 2023 Canonical Ltd.

5# Written by:

6# Hanhsuan Lee <hanhsuan.lee@canonical.com>

7#

8# Checkbox is free software: you can redistribute it and/or modify

9# it under the terms of the GNU General Public License version 3,

10# as published by the Free Software Foundation.

11#

12# Checkbox is distributed in the hope that it will be useful,

13# but WITHOUT ANY WARRANTY; without even the implied warranty of

14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

15# GNU General Public License for more details.

16#

17# You should have received a copy of the GNU General Public License

18# along with Checkbox. If not, see <http://www.gnu.org/licenses/>.

19

20 from checkbox_support.helpers.timeout import run_with_timeout

21 import subprocess

22 import threading

23 import argparse

24 import logging

25 import fnmatch

26 import time

27 import json

28 import sys

29 import re

30 import os

31 import typing as T

32 from shlex import split as sh_split

33

34

35 class PrimeOffloader:

36 """

37 A class used to execute process to specific GPU.

38 Have to run this as root.

39

40 :attr logger: console logger

41 :type logger: RootLogger

42

43 :attr check_result:

44 store the result of checking offloading is ok or not.

45 :type check_result: bool

46 """

47

48 logger = logging.getLogger()

49 check_result = False

50

51 def find_file_containing_string(

52 self, search_directory: str, filename_pattern: str, search_string: str

53 ) -> T.Optional[str]:

54 """

55 Search for a file matching a specific pattern

56 that contains a given string.

57

58 :param search_directory: The directory to search through.

59

60 :param filename_pattern: The pattern that filenames should match.

61

62 :param search_string: The string to search for within

63 the file's contents.

64

65 :returns: The full path of the file that contains the string,

66 or None if no file is found.

67 """

68 for root, dirs, files in os.walk(search_directory):

69 for file_name in fnmatch.filter(files, filename_pattern):

70 file_path = os.path.join(root, file_name)

71 # Check if the search string is in the file

72 with open(

73 file_path, "r", encoding="utf-8", errors="ignore"

74 ) as file:

75 if search_string in file.read():

76 return file_path

77

78 def find_card_id(self, pci_bdf: str) -> str:

79 """

80 use pci BDF to find card id under /sys/kernel/debug/dri

81

82 :param pci_bdf: pci device BDF in NNNN:NN:NN.N format

83

84 :returns: card id

85 """

86 pci_bdf_format = "[0-9]{4}:[0-9,a-f]{2}:[0-9,a-f]{2}.[0-9]"

87 if not re.match(pci_bdf_format, pci_bdf.lower()):

88 raise SystemExit("pci BDF format error")

89

90 try:

91 card_path = self.find_file_containing_string(

92 "/sys/kernel/debug/dri", "name", pci_bdf

93 )

94 assert card_path, "Couldn't find a card named: {}".format(pci_bdf)

95 return card_path.split("/")[5]

96 except IndexError as e:

97 raise SystemExit("return value format error {}".format(repr(e)))

98

99 def find_card_name(self, pci_bdf: str) -> str:

100 """

101 use pci BDF to find card name by lshw

102

103 :param pci_bdf: pci device BDF in NNNN:NN:NN.N format

104

105 :returns: card name

106 """

107 cmd = ["lshw", "-c", "display", "-numeric", "-json"]

108 try:

109 card_infos = subprocess.check_output(

110 cmd, shell=False, universal_newlines=True

111 )

112 infos = json.loads(card_infos)

113 for info in infos:

114 if pci_bdf in info["businfo"]:

115 return info["product"]

116 raise SystemExit("Card name not found")

117 except (KeyError, TypeError, json.decoder.JSONDecodeError) as e:

118 raise SystemExit("return value format error {}".format(e))

119 except (subprocess.CalledProcessError, FileNotFoundError) as e:

120 raise SystemExit(

121 "Running command:{} failed due to {}".format(cmd, repr(e))

122 )

123

124 def get_clients(self, card_id: str) -> str:

125 """

126 Use to get clients that running on specific GPU

127 by reading debugfs.

128

129 .. note::

130 While setting prime offload environment such as DRI_PRIME,

131 the process will be listed under kernel debug interface.

132 The location of kernel debug interface is

133 /sys/kernel/debug/dri/<card id>,

134 and the process could be found in

135 /sys/kernel/debug/dri/<card id>/clients

136

137 :param card_id: card id shows in debugfs

138 """

139 filename = "/sys/kernel/debug/dri/{}/clients".format(card_id)

140 with open(filename, "r") as f:

141 return f.read()

142 return ""

143

144 def check_offload(

145 self, cmd: T.List[str], card_id: str, card_name: str, timeout: int

146 ):

147 """

148 Used to check if the provided command is executed on a specific GPU.

149

150 :param cmd: command to be run under prime offload

151

152 :param card_id: card ID of the DRI device

153

154 :param card_name: card name of the DRI device

155

156 :param timeout: timeout for the offloaded command

157 """

158 delay = timeout / 10

159

160 deadline = time.time() + timeout

161

162 while time.time() < deadline:

163 time.sleep(delay)

164 clients = self.get_clients(card_id)

165 # The command shows in /sys/kernel/debug/dri/<card_id>/clients

166 # doesn't include arguments. Therefore cmd[0] is used to search

167 if clients and cmd[0] in clients:

168 self.logger.info("Checking success:")

169 self.logger.info(" Offload process:[{}]".format(cmd))

170 self.logger.info(" Card ID:[{}]".format(card_id))

171 self.logger.info(" Device Name:[{}]".format(card_name))

172 return

173 self.logger.info("Checking fail:")

174 self.logger.info(" Couldn't find process {}".format(cmd))

175 self.check_result = True

176

177 def _find_bdf(self, card_id: str):

178 """

179 Use the /sys/kernel/debug/dri/<card id>/name to get pci BDF.

180

181 :param card_id: card id shows in debugfs

182 """

183 filename = "/sys/kernel/debug/dri/{}/name".format(card_id)

184 with open(filename, "r") as f:

185 data_in_name = f.read()

186 return data_in_name.split()[1].split("=")[1]

187

188 def find_offload(self, cmd_str: str, timeout: int):

189 """

190 Find the card that the command is running on.

191 This script looks for the card on which a specific command is running.

192 It checks ten times in regular intervals if the process is running

193 by looking for it in the /sys/kernel/debug/dri/<card id>/clients.

194 If the timeout is reached, the function will fail

195

196 :param cmd: command that is running

197

198 :param timeout: timeout for command

199 """

200 directory = "/sys/kernel/debug/dri"

201

202 delay = timeout / 10

203

204 deadline = time.time() + timeout

205

206 cmd = sh_split(cmd_str)

207

208 while time.time() < deadline:

209 time.sleep(delay)

210 # The command shows in /sys/kernel/debug/dri/<card_id>/clients

211 # doesn't include arguments. Therefore cmd[0] is used to search

212 card_path = self.find_file_containing_string(

213 directory, "clients", cmd[0]

214 )

215

216 if card_path and directory in card_path:

217 try:

218 # The graphic will be shown such as 0 and 128

219 # at the same time. Therefore, pick up the first one

220 first_card = card_path.splitlines()[0]

221 card_id = first_card.split("/")[5]

222 bdf = self._find_bdf(card_id)

223 self.logger.info("Process is running on:")

224 self.logger.info(" process:[{}]".format(cmd[0]))

225 self.logger.info(

226 " Card ID:[{}]".format(self.find_card_id(bdf))

227 )

228 self.logger.info(

229 " Device Name:[{}]".format(self.find_card_name(bdf))

230 )

231 return

232 except IndexError as e:

233 self.logger.info(

234 "Finding card information failed {}".format(repr(e))

235 )

236

237 self.logger.info("Checking fail:")

238 self.logger.info(" Couldn't find process {}".format(cmd))

239 self.check_result = True

240

241 def check_nv_offload_env(self):

242 """

243 prime offload of nvidia driver is limited.

244 Only on-demand mode is supported.

245 """

246 # nvidia-smi ship with NVIDIA GPU display drivers on Linux

247 # https://developer.nvidia.com/nvidia-system-management-interface

248 # check prime-select to make sure the nv driver is included.

249 # If there is no nv driver, prime offload is fine for other drivers.

250 try:

251 if "on-demand" not in subprocess.check_output(

252 ["prime-select", "query"], universal_newlines=True

253 ):

254 raise SystemExit("System isn't on-demand mode")

255

256 # prime offload couldn't running on nvlink active or inactive

257 # Therefore, only return empty string is supported environment.

258 nvlink = subprocess.check_output(

259 ["nvidia-smi", "nvlink", "-s"], universal_newlines=True

260 )

261 if nvlink:

262 if "error" in nvlink.lower():

263 raise SystemExit("nvidia driver error")

264 raise SystemExit("NVLINK detected")

265 except FileNotFoundError:

266 self.logger.info(

267 "No prime-select, it should be ok to run prime offload"

268 )

269

270 def cmd_runner(self, cmd: T.List[str], env: T.Optional[T.Dict] = None):

271 """

272 use to execute command and piping the output to the screen.

273

274 :param cmd: the command will be executed

275

276 :param env: the environment variables for executing command

277

278 """

279 try:

280 with subprocess.Popen(

281 cmd,

282 env=env,

283 stdout=subprocess.PIPE,

284 universal_newlines=True,

285 ) as runner:

286

287 self.logger.info("running command:[{}]".format(cmd))

288

289 # redirect command output real time

290 while runner.poll() is None: 290 ↛ 292 line 290 didn't jump to line 292 because the condition on line 290 was never true

291 # when stdout=subprocess.PIPE, stdout is not None

292 line = runner.stdout.readline().strip() # type: ignore

293 self.logger.info(line)

294 except subprocess.CalledProcessError as e:

295 raise SystemExit("run command failed {}".format(repr(e)))

296

297 def cmd_finder(self, cmd: str, timeout: int):

298 """

299 run offload command and find it runs on which GPU

300

301 :param cmd: command that running under prime offload

302

303 :param timeout: timeout for offloaded command

304 """

305 if "timeout" in cmd:

306 raise SystemExit("Put timeout in command isn't allowed")

307

308 # use other thread to find offload

309 find_thread = threading.Thread(

310 target=self.find_offload, args=(cmd, timeout)

311 )

312 find_thread.start()

313 try:

314 run_with_timeout(self.cmd_runner, timeout, cmd.split())

315 except TimeoutError:

316 self.logger.info("Test finished")

317 find_thread.join()

318

319 if self.check_result:

320 raise SystemExit("Couldn't find process running on GPU")

321

322 def cmd_checker(

323 self, cmd_str: str, pci_bdf: str, driver: str, timeout: int

324 ):

325 """

326 run offload command and check it runs on correct GPU

327

328 :param cmd: command that running under prime offload

329

330 :param pci_bdf: pci device name in NNNN:NN:NN.N format

331

332 :param driver: GPU driver, such as i915, amdgpu, nvidia

333

334 :param timeout: timeout for offloaded command

335 """

336 card_id = self.find_card_id(pci_bdf)

337 card_name = self.find_card_name(pci_bdf)

338

339 # run offload command in other process

340 dri_pci_bdf_format = re.sub("[:.]", "_", pci_bdf)

341

342 if "timeout" in cmd_str:

343 raise SystemExit("Put timeout in command isn't allowed")

344

345 env = os.environ.copy()

346 if driver in ("nvidia", "pcieport"):

347 offload_env = {

348 "__NV_PRIME_RENDER_OFFLOAD": "1",

349 "__GLX_VENDOR_LIBRARY_NAME": "nvidia",

350 }

351 else:

352 offload_env = {"DRI_PRIME": "pci-{}".format(dri_pci_bdf_format)}

353

354 env.update(offload_env)

355 self.logger.info("prime offload env: {}".format(offload_env))

356

357 # if nv driver under nvidia mode, prime/reverse prime couldn't work.

358 self.check_nv_offload_env()

359

360 cmd = sh_split(cmd_str)

361 # use other thread to check offload is correctly or not

362 check_thread = threading.Thread(

363 target=self.check_offload, args=(cmd, card_id, card_name, timeout)

364 )

365 check_thread.start()

366 try:

367 run_with_timeout(self.cmd_runner, timeout, cmd, env)

368 except TimeoutError:

369 self.logger.info("Test finished")

370 check_thread.join()

371

372 if self.check_result:

373 raise SystemExit(

374 "offload to specific GPU: {} failed".format(pci_bdf)

375 )

376

377 def parse_args(self, args=sys.argv[1:]):

378 """

379 command line arguments parsing

380

381 :param args: arguments from sys

382 :type args: sys.argv

383 """

384 parser = argparse.ArgumentParser(

385 prog="Prime offload tester",

386 description="Test prime offload feature",

387 )

388

389 parser.add_argument(

390 "-c",

391 "--command",

392 type=str,

393 default="glxgears",

394 help="command to offload to specific GPU (default: %(default)s)",

395 )

396 parser.add_argument(

397 "-p",

398 "--pci",

399 type=str,

400 help="pci device bdf in NNNN:NN:NN.N format, such as 0000:00:02.0",

401 )

402 parser.add_argument(

403 "-d",

404 "--driver",

405 type=str,

406 help="Type of GPU driver, such as i915",

407 )

408 parser.add_argument(

409 "-t",

410 "--timeout",

411 type=int,

412 default=20,

413 help="executing command duration in second (default: %(default)s).", # noqa: E501

414 )

415 return parser.parse_args(args)

416

417 def main(self):

418 args = self.parse_args()

419

420 # create self.logger.formatter

421 log_formatter = logging.Formatter(fmt="%(message)s")

422

423 # create logger

424 self.logger.setLevel(logging.INFO)

425

426 # create console handler

427 console_handler = logging.StreamHandler()

428 console_handler.setFormatter(log_formatter)

429

430 # Add console handler to logger

431 self.logger.addHandler(console_handler)

432

433 if args.pci and args.driver:

434 # cmd_checker("glxgears", "0000:00:02.0", "i915", 0)

435 self.cmd_checker(args.command, args.pci, args.driver, args.timeout)

436 else:

437 # cmd_finder("glxgears", 0)

438 self.cmd_finder(args.command, args.timeout)

439

440

441 if __name__ == "__main__":

442 assert os.getuid() == 0, "This test must be run as root"

443 PrimeOffloader().main()