Coverage for src/par_run/cli.py: 91%

1"""CLI for running commands in parallel"""

2

3 from collections import OrderedDict

4 from pathlib import Path

5 from typing import Optional

6 import enum

7

8 import rich

9 import typer

10

11 from .executor import Command, CommandStatus, ProcessingStrategy, read_commands_toml

12

13 PID_FILE = ".par-run.uvicorn.pid"

14

15 cli_app = typer.Typer()

16

17

18# Web only functions

19 def clean_up():

20 """

21 Clean up by removing the PID file.

22 """

23 os.remove(PID_FILE)

24 typer.echo("Cleaned up PID file.")

25

26

27 def start_web_server(port: int):

28 """Start the web server"""

29 if os.path.isfile(PID_FILE):

30 typer.echo("UVicorn server is already running.")

31 sys.exit(1)

32 with open(PID_FILE, "w", encoding="utf-8") as pid_file:

33 typer.echo(f"Starting UVicorn server on port {port}...")

34 uvicorn_command = [

35 "uvicorn",

36 "par_run.web:ws_app",

37 "--host",

38 "0.0.0.0",

39 "--port",

40 str(port),

41 ]

42 process = subprocess.Popen(uvicorn_command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

43 pid_file.write(str(process.pid))

44

45 # Wait for UVicorn to start

46 wait_time = 3 * 10**9 # 3 seconds

47 start_time = time.time_ns()

48

49 while time.time_ns() - start_time < wait_time:

50 test_port = get_process_port(process.pid)

51 if port == test_port:

52 typer.echo(f"UVicorn server is running on port {port} in {(time.time_ns() - start_time)/10**6:.2f} ms.")

53 break

54 time.sleep(0.1) # Poll every 0.1 seconds

55

56 else:

57 typer.echo(f"UVicorn server did not respond within {wait_time} seconds.")

58 typer.echo("run 'par-run web status' to check the status.")

59

60

61 def stop_web_server():

62 """

63 Stop the UVicorn server by reading its PID from the PID file and sending a termination signal.

64 """

65 if not Path(PID_FILE).is_file():

66 typer.echo("UVicorn server is not running.")

67 return

68

69 with open(PID_FILE, "r") as pid_file:

70 pid = int(pid_file.read().strip())

71

72 typer.echo(f"Stopping UVicorn server with {pid=:}...")

73 try:

74 os.kill(pid, signal.SIGTERM)

75 except ProcessLookupError:

76 pass

77 clean_up()

78

79

80 def get_process_port(pid: int) -> Optional[int]:

81 process = psutil.Process(pid)

82 connections = process.connections()

83 if connections:

84 port = connections[0].laddr.port

85 return port

86 return None

87

88

89 def list_uvicorn_processes():

90 """Check for other UVicorn processes and list them"""

91 uvicorn_processes = []

92 for process in psutil.process_iter():

93 try:

94 process_name = process.name()

95 if "uvicorn" in process_name.lower():

96 uvicorn_processes.append(process)

97 except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):

98 pass

99 if uvicorn_processes:

100 typer.echo("Other UVicorn processes:")

101 for process in uvicorn_processes:

102 typer.echo(f"PID: {process.pid}, Name: {process.name()}")

103 else:

104 typer.echo("No other UVicorn processes found.")

105

106

107 def get_web_server_status():

108 """

109 Get the status of the UVicorn server by reading its PID from the PID file.

110 """

111 if not os.path.isfile(PID_FILE):

112 typer.echo("No pid file found. Server likely not running.")

113 list_uvicorn_processes()

114 return

115

116 with open(PID_FILE, "r") as pid_file:

117 pid = int(pid_file.read().strip())

118 if psutil.pid_exists(pid):

119 port = get_process_port(pid)

120 if port:

121 typer.echo(f"UVicorn server is running with {pid=}, {port=}")

122 else:

123 typer.echo(f"UVicorn server is running with {pid=:}, couldn't determine port.")

124 else:

125 typer.echo("UVicorn server is not running but pid files exists, deleting it.")

126 clean_up()

127

128

129 class WebCommand(enum.Enum):

130 """Web command enumeration."""

131

132 START = "start"

133 STOP = "stop"

134 RESTART = "restart"

135 STATUS = "status"

136

137 def __str__(self):

138 return self.value

139

140

141 class CLICommandCBOnComp:

142 def on_start(self, cmd: Command):

143 rich.print(f"[blue bold]Completed command {cmd.name}[/]")

144

145 def on_recv(self, _: Command, output: str):

146 rich.print(output)

147

148 def on_term(self, cmd: Command, exit_code: int):

149 """Callback function for when a command receives output"""

150 if cmd.status == CommandStatus.SUCCESS:

151 rich.print(f"[green bold]Command {cmd.name} finished[/]")

152 elif cmd.status == CommandStatus.FAILURE:

153 rich.print(f"[red bold]Command {cmd.name} failed, {exit_code=:}[/]")

154

155

156 class CLICommandCBOnRecv:

157 def on_start(self, cmd: Command):

158 rich.print(f"[blue bold]{cmd.name}: Started[/]")

159

160 def on_recv(self, cmd: Command, output: str):

161 rich.print(f"{cmd.name}: {output}")

162

163 def on_term(self, cmd: Command, exit_code: int):

164 """Callback function for when a command receives output"""

165 if cmd.status == CommandStatus.SUCCESS:

166 rich.print(f"[green bold]{cmd.name}: Finished[/]")

167 elif cmd.status == CommandStatus.FAILURE:

168 rich.print(f"[red bold]{cmd.name}: Failed, {exit_code=:}[/]")

169

170

171 def format_elapsed_time(seconds: float) -> str:

172 """

173 Converts a number of seconds into a human-readable time format of HH:MM:SS.xxx

174

175 Args:

176 seconds (float): The number of seconds elapsed.

177

178 Returns:

179 str: The formatted time string.

180 """

181 hours = int(seconds) // 3600

182 minutes = (int(seconds) % 3600) // 60

183 seconds = seconds % 60 # Keeping the fractional part of seconds

184

185 # Return formatted string with seconds rounded to 2 d.p.

186 return f"{hours:02}:{minutes:02}:{seconds:06.3f}"

187

188

189@cli_app.command()

190 def run(

191 style: ProcessingStrategy = typer.Option(help="Processing strategy", default="comp"),

192 show: bool = typer.Option(help="Show available groups and commands", default=False),

193 file: Path = typer.Option(help="The commands.ini file to use", default=Path("pyproject.toml")),

194 groups: Optional[str] = typer.Option(None, help="Run a specific group of commands, comma spearated"),

195 cmds: Optional[str] = typer.Option(None, help="Run a specific commands, comma spearated"),

196):

197 """Run commands in parallel"""

198 # Overall exit code, need to track all command exit codes to update this

199 exit_code = 0

200 st_all = time.perf_counter()

201 # console = rich.console.Console()

202 master_groups = read_commands_toml(file)

203 if show:

204 for grp in master_groups:

205 rich.print(f"[blue bold]Group: {grp.name}[/]")

206 for _, cmd in grp.cmds.items():

207 rich.print(f"[green bold]{cmd.name}[/]: {cmd.cmd}")

208 return

209

210 if groups:

211 master_groups = [grp for grp in master_groups if grp.name in [g.strip() for g in groups.split(",")]]

212

213 if cmds:

214 for grp in master_groups:

215 grp.cmds = OrderedDict(

216 {

217 cmd_name: cmd

218 for cmd_name, cmd in grp.cmds.items()

219 if cmd_name in [c.strip() for c in cmds.split(",")]

220 }

221 )

222 master_groups = [grp for grp in master_groups if grp.cmds]

223

224 if not master_groups:

225 rich.print("[blue]No groups or commands found.[/]")

226 raise typer.Exit(0)

227

228 for grp in master_groups:

229 if style == ProcessingStrategy.ON_COMP:

230 exit_code = exit_code or grp.run(style, CLICommandCBOnComp())

231 elif style == ProcessingStrategy.ON_RECV:

232 exit_code = exit_code or grp.run(style, CLICommandCBOnRecv())

233 else:

234 raise typer.BadParameter("Invalid processing strategy")

235

236 # Summarise the results

237 console = rich.console.Console()

238 for grp in master_groups:

239 console.print(f"[blue bold]Group: {grp.name}[/]")

240 for _, cmd in grp.cmds.items():

241 elap_str = ""

242 if cmd.elapsed:

243 elap_str = f", {format_elapsed_time(cmd.elapsed)}"

244 else:

245 elap_str = ", XX:XX:XX.xxx"

246

247 if cmd.status == CommandStatus.SUCCESS:

248 left_seg = f"[green bold]Command {cmd.name} succeeded "

249 else:

250 left_seg = f"[red bold]Command {cmd.name} failed "

251

252 right_seg = f"({cmd.num_non_empty_lines}{elap_str})[/]"

253

254 # Adjust total line width dynamically based on max width and other content

255 pad_length = (

256 100 - len(left_seg) - len(right_seg) - 10

257 if "succeeded" in left_seg

258 else 100 - len(left_seg) - len(right_seg) - 12

259 )

260

261 rich.print(f"{left_seg}{' ' * pad_length}{right_seg}")

262 end_style = "[green bold]" if exit_code == 0 else "[red bold]"

263 rich.print(f"\n{end_style}Total elapsed time: {format_elapsed_time(time.perf_counter() - st_all)}[/]")

264 raise typer.Exit(exit_code)

265

266

267 try:

268 import os

269 import signal

270 import subprocess

271 import sys

272 import time

273 from pathlib import Path

274 from typing import Optional

275

276 import psutil

277 import typer

278

279 rich.print("[blue]Web commands loaded[/]")

280

281 PID_FILE = ".par-run.uvicorn.pid"

282

283 @cli_app.command()

284 def web(

285 command: WebCommand = typer.Argument(..., help="command to control/interract with the web server"),

286 port: int = typer.Option(8001, help="Port to run the web server"),

287 ):

288 """Run the web server"""

289 if command == WebCommand.START:

290 start_web_server(port)

291 elif command == WebCommand.STOP:

292 stop_web_server()

293 elif command == WebCommand.RESTART:

294 stop_web_server()

295 start_web_server(port)

296 elif command == WebCommand.STATUS:

297 get_web_server_status()

298 else:

299 typer.echo(f"Not a valid command '{command}'", err=True)

300 raise typer.Abort()

301

302 except ImportError: # pragma: no cover

303 pass # pragma: no cover

304

305 if __name__ == "__main__":

306 cli_app()