runner.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. import io
  13. import sys
  14. import abc
  15. import shlex
  16. import locale
  17. import asyncio
  18. from .utils.arg import CLIArgument
  19. from .utils.subprocess import run_cmd as _run_cmd, CompletedProcess
  20. from ...utils import logging
  21. from ...utils.misc import abspath
  22. from ...utils.flags import DRY_RUN
  23. from ...utils.errors import raise_unsupported_api_error, CalledProcessError
  24. __all__ = ['BaseRunner', 'InferOnlyRunner']
  25. class BaseRunner(metaclass=abc.ABCMeta):
  26. """
  27. Abstract base class of Runner.
  28. Runner is responsible for executing training/inference/compression commands.
  29. """
  30. def __init__(self, runner_root_path):
  31. """
  32. Initialize the instance.
  33. Args:
  34. runner_root_path (str): Path of the directory where the scripts reside.
  35. """
  36. super().__init__()
  37. self.runner_root_path = abspath(runner_root_path)
  38. # Path to python interpreter
  39. self.python = sys.executable
  40. def prepare(self):
  41. """
  42. Make preparations for the execution of commands.
  43. For example, download prerequisites and install dependencies.
  44. """
  45. # By default we do nothing
  46. pass
  47. @abc.abstractmethod
  48. def train(self, config_path, cli_args, device, ips, save_dir, do_eval=True):
  49. """
  50. Execute model training command.
  51. Args:
  52. config_path (str): Path of the configuration file.
  53. cli_args (list[base.utils.arg.CLIArgument]): List of command-line
  54. arguments.
  55. device (str): A string that describes the device(s) to use, e.g.,
  56. 'cpu', 'xpu:0', 'gpu:1,2'.
  57. ips (str|None): Paddle cluster node ips, e.g.,
  58. '192.168.0.16,192.168.0.17'.
  59. save_dir (str): Directory to save log files.
  60. do_eval (bool, optional): Whether to perform model evaluation during
  61. training. Default: True.
  62. Returns:
  63. paddlex.repo_apis.base.utils.subprocess.CompletedProcess
  64. """
  65. raise NotImplementedError
  66. @abc.abstractmethod
  67. def evaluate(self, config_path, cli_args, device, ips):
  68. """
  69. Execute model evaluation command.
  70. Args:
  71. config_path (str): Path of the configuration file.
  72. cli_args (list[base.utils.arg.CLIArgument]): List of command-line
  73. arguments.
  74. device (str): A string that describes the device(s) to use, e.g.,
  75. 'cpu', 'xpu:0', 'gpu:1,2'.
  76. ips (str|None): Paddle cluster node ips, e.g.,
  77. '192.168.0.16,192.168.0.17'.
  78. Returns:
  79. paddlex.repo_apis.base.utils.subprocess.CompletedProcess
  80. """
  81. raise NotImplementedError
  82. @abc.abstractmethod
  83. def predict(self, config_path, cli_args, device):
  84. """
  85. Execute prediction command.
  86. Args:
  87. config_path (str): Path of the configuration file.
  88. cli_args (list[base.utils.arg.CLIArgument]): List of command-line
  89. arguments.
  90. device (str): A string that describes the device(s) to use, e.g.,
  91. 'cpu', 'xpu:0', 'gpu:1,2'.
  92. Returns:
  93. paddlex.repo_apis.base.utils.subprocess.CompletedProcess
  94. """
  95. raise NotImplementedError
  96. @abc.abstractmethod
  97. def export(self, config_path, cli_args, device):
  98. """
  99. Execute model export command.
  100. Args:
  101. config_path (str): Path of the configuration file.
  102. cli_args (list[base.utils.arg.CLIArgument]): List of command-line
  103. arguments.
  104. device (str): A string that describes the device(s) to use, e.g.,
  105. 'cpu', 'xpu:0', 'gpu:1,2'.
  106. Returns:
  107. paddlex.repo_apis.base.utils.subprocess.CompletedProcess
  108. """
  109. raise NotImplementedError
  110. @abc.abstractmethod
  111. def infer(self, config_path, cli_args, device):
  112. """
  113. Execute model inference command.
  114. Args:
  115. config_path (str): Path of the configuration file.
  116. cli_args (list[base.utils.arg.CLIArgument]): List of command-line
  117. arguments.
  118. device (str): A string that describes the device(s) to use, e.g.,
  119. 'cpu', 'xpu:0', 'gpu:1,2'.
  120. Returns:
  121. paddlex.repo_apis.base.utils.subprocess.CompletedProcess
  122. """
  123. raise NotImplementedError
  124. @abc.abstractmethod
  125. def compression(self, config_path, train_cli_args, export_cli_args, device,
  126. train_save_dir):
  127. """
  128. Execute model compression (quantization aware training and model export)
  129. commands.
  130. Args:
  131. config_path (str): Path of the configuration file.
  132. train_cli_args (list[base.utils.arg.CLIArgument]): List of
  133. command-line arguments used for model training.
  134. export_cli_args (list[base.utils.arg.CLIArgument]): List of
  135. command-line arguments used for model export.
  136. device (str): A string that describes the device(s) to use, e.g.,
  137. 'cpu', 'xpu:0', 'gpu:1,2'.
  138. train_save_dir (str): Directory to store model snapshots.
  139. Returns:
  140. tuple[paddlex.repo_apis.base.utils.subprocess.CompletedProcess]
  141. """
  142. raise NotImplementedError
  143. def distributed(self, device, ips=None, log_dir=None):
  144. """ distributed """
  145. # TODO: docstring
  146. args = [self.python]
  147. if device is None:
  148. return args, None
  149. device, dev_ids = self.parse_device(device)
  150. if len(dev_ids) == 0:
  151. return args, None
  152. else:
  153. num_devices = len(dev_ids)
  154. dev_ids = ','.join(dev_ids)
  155. if num_devices > 1:
  156. args.extend(['-m', 'paddle.distributed.launch'])
  157. args.extend(['--devices', dev_ids])
  158. if ips is not None:
  159. args.extend(['--ips', ips])
  160. if log_dir is None:
  161. log_dir = os.getcwd()
  162. args.extend(['--log_dir', self._get_dist_train_log_dir(log_dir)])
  163. elif num_devices == 1:
  164. new_env = os.environ.copy()
  165. if device == 'xpu':
  166. new_env['XPU_VISIBLE_DEVICES'] = dev_ids
  167. else:
  168. new_env['CUDA_VISIBLE_DEVICES'] = dev_ids
  169. return args, new_env
  170. return args, None
  171. def parse_device(self, device):
  172. """ parse_device """
  173. # According to https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/device/set_device_cn.html
  174. if ':' not in device:
  175. device_type, dev_ids = device, []
  176. else:
  177. device_type, dev_ids = device.split(':')
  178. dev_ids = dev_ids.split(',')
  179. if device_type not in ('cpu', 'gpu', 'xpu'):
  180. raise ValueError("Unsupported device type.")
  181. for dev_id in dev_ids:
  182. if not dev_id.isdigit():
  183. raise ValueError("Device ID must be an integer.")
  184. return device_type, dev_ids
  185. def run_cmd(self,
  186. cmd,
  187. env=None,
  188. switch_wdir=True,
  189. silent=False,
  190. echo=True,
  191. capture_output=False,
  192. log_path=None):
  193. """ run_cmd """
  194. def _trans_args(cmd):
  195. out = []
  196. for ele in cmd:
  197. if isinstance(ele, CLIArgument):
  198. out.extend(ele.lst)
  199. else:
  200. out.append(ele)
  201. return out
  202. cmd = _trans_args(cmd)
  203. if DRY_RUN:
  204. # TODO: Accommodate Windows system
  205. logging.info(' '.join(shlex.quote(x) for x in cmd))
  206. # Mock return
  207. return CompletedProcess(
  208. cmd, returncode=0, stdout=str(cmd), stderr=None)
  209. if switch_wdir:
  210. if isinstance(switch_wdir, str):
  211. # In this case `switch_wdir` specifies a relative path
  212. cwd = os.path.join(self.runner_root_path, switch_wdir)
  213. else:
  214. cwd = self.runner_root_path
  215. else:
  216. cwd = None
  217. if not capture_output:
  218. if log_path is not None:
  219. logging.warning(
  220. "`log_path` will be ignored when `capture_output` is False.")
  221. cp = _run_cmd(
  222. cmd,
  223. env=env,
  224. cwd=cwd,
  225. silent=silent,
  226. echo=echo,
  227. pipe_stdout=False,
  228. pipe_stderr=False,
  229. blocking=True)
  230. cp = CompletedProcess(
  231. cp.args, cp.returncode, stdout=cp.stdout, stderr=cp.stderr)
  232. else:
  233. # Refer to
  234. # https://stackoverflow.com/questions/17190221/subprocess-popen-cloning-stdout-and-stderr-both-to-terminal-and-variables/25960956
  235. async def _read_display_and_record_from_stream(in_stream,
  236. out_stream, files):
  237. # According to
  238. # https://docs.python.org/3/library/subprocess.html#frequently-used-arguments
  239. _ENCODING = locale.getpreferredencoding(False)
  240. chars = []
  241. out_stream_is_buffered = hasattr(out_stream, 'buffer')
  242. while True:
  243. flush = False
  244. char = await in_stream.read(1)
  245. if char == b'':
  246. break
  247. if out_stream_is_buffered:
  248. out_stream.buffer.write(char)
  249. chars.append(char)
  250. if char == b'\n':
  251. flush = True
  252. elif char == b'\r':
  253. # NOTE: In order to get tqdm progress bars to produce normal outputs
  254. # we treat '\r' as an ending character of line
  255. flush = True
  256. if flush:
  257. line = b''.join(chars)
  258. line = line.decode(_ENCODING)
  259. if not out_stream_is_buffered:
  260. # We use line buffering
  261. out_stream.write(line)
  262. else:
  263. out_stream.buffer.flush()
  264. for f in files:
  265. f.write(line)
  266. chars.clear()
  267. async def _tee_proc_call(proc_call, out_files, err_files):
  268. proc = await proc_call
  269. await asyncio.gather(
  270. _read_display_and_record_from_stream(proc.stdout,
  271. sys.stdout, out_files),
  272. _read_display_and_record_from_stream(proc.stderr,
  273. sys.stderr, err_files))
  274. # NOTE: https://docs.python.org/3/library/subprocess.html#subprocess.Popen.wait
  275. retcode = await proc.wait()
  276. return retcode
  277. # Non-blocking call with stdout and stderr piped
  278. with io.StringIO() as stdout_buf, io.StringIO() as stderr_buf:
  279. proc_call = _run_cmd(
  280. cmd,
  281. env=env,
  282. cwd=cwd,
  283. echo=echo,
  284. silent=silent,
  285. pipe_stdout=True,
  286. pipe_stderr=True,
  287. blocking=False,
  288. async_run=True)
  289. out_files = [stdout_buf]
  290. err_files = [stderr_buf]
  291. if log_path is not None:
  292. log_dir = os.path.dirname(log_path)
  293. os.makedirs(log_dir, exist_ok=True)
  294. log_file = open(log_path, 'w', encoding='utf-8')
  295. logging.info(f"\nLog path: {os.path.abspath(log_path)} \n")
  296. out_files.append(log_file)
  297. err_files.append(log_file)
  298. try:
  299. retcode = asyncio.run(
  300. _tee_proc_call(proc_call, out_files, err_files))
  301. finally:
  302. if log_path is not None:
  303. log_file.close()
  304. cp = CompletedProcess(cmd, retcode,
  305. stdout_buf.getvalue(),
  306. stderr_buf.getvalue())
  307. if cp.returncode != 0:
  308. raise CalledProcessError(
  309. cp.returncode, cp.args, output=cp.stdout, stderr=cp.stderr)
  310. return cp
  311. def _get_dist_train_log_dir(self, log_dir):
  312. """ _get_dist_train_log_dir """
  313. return os.path.join(log_dir, 'distributed_train_logs')
  314. def _get_train_log_path(self, log_dir):
  315. """ _get_train_log_path """
  316. return os.path.join(log_dir, 'train.log')
  317. class InferOnlyRunner(BaseRunner):
  318. """ InferOnlyRunner """
  319. def train(self, *args, **kwargs):
  320. """ train """
  321. raise_unsupported_api_error(self.__class__, 'train')
  322. def evaluate(self, *args, **kwargs):
  323. """ evaluate """
  324. raise_unsupported_api_error(self.__class__, 'evalaute')
  325. def predict(self, *args, **kwargs):
  326. """ predict """
  327. raise_unsupported_api_error(self.__class__, 'predict')
  328. def export(self, *args, **kwargs):
  329. """ export """
  330. raise_unsupported_api_error(self.__class__, 'export')
  331. def compression(self, *args, **kwargs):
  332. """ compression """
  333. raise_unsupported_api_error(self.__class__, 'compression')