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

172 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-12 08:38 -0400

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

2 

3from collections import OrderedDict 

4from pathlib import Path 

5from typing import Optional 

6import enum 

7 

8import rich 

9import typer 

10 

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

12 

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

14 

15cli_app = typer.Typer() 

16 

17 

18# Web only functions 

19def 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 

27def 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 

61def 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 

80def 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 

89def 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 

107def 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 

129class 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 

141class 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 

156class 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 

171def 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.2f}" 

187 

188 

189@cli_app.command() 

190def 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 

201 master_groups = read_commands_toml(file) 

202 if show: 

203 for grp in master_groups: 

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

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

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

207 return 

208 

209 if groups: 

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

211 

212 if cmds: 

213 for grp in master_groups: 

214 grp.cmds = OrderedDict( 

215 { 

216 cmd_name: cmd 

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

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

219 } 

220 ) 

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

222 

223 if not master_groups: 

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

225 raise typer.Exit(0) 

226 

227 for grp in master_groups: 

228 if style == ProcessingStrategy.ON_COMP: 

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

230 elif style == ProcessingStrategy.ON_RECV: 

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

232 else: 

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

234 

235 # Summarise the results 

236 for grp in master_groups: 

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

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

239 if cmd.status == CommandStatus.SUCCESS: 

240 elap_str = "" 

241 if cmd.elapsed: 

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

243 rich.print(f"[green bold]Command {cmd.name} succeeded ({cmd.num_non_empty_lines}{elap_str})[/]") 

244 else: 

245 rich.print(f"[red bold]Command {cmd.name} failed ({cmd.num_non_empty_lines}{elap_str})[/]") 

246 

247 raise typer.Exit(exit_code) 

248 

249 

250try: 

251 import os 

252 import signal 

253 import subprocess 

254 import sys 

255 import time 

256 from pathlib import Path 

257 from typing import Optional 

258 

259 import psutil 

260 import typer 

261 

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

263 

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

265 

266 @cli_app.command() 

267 def web( 

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

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

270 ): 

271 """Run the web server""" 

272 if command == WebCommand.START: 

273 start_web_server(port) 

274 elif command == WebCommand.STOP: 

275 stop_web_server() 

276 elif command == WebCommand.RESTART: 

277 stop_web_server() 

278 start_web_server(port) 

279 elif command == WebCommand.STATUS: 

280 get_web_server_status() 

281 else: 

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

283 raise typer.Abort() 

284 

285except ImportError: # pragma: no cover 

286 pass # pragma: no cover 

287 

288if __name__ == "__main__": 

289 cli_app()