Coverage for src/par_run/cli.py: 91%
162 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-31 17:05 -0400
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-31 17:05 -0400
1"""CLI for running commands in parallel"""
3from collections import OrderedDict
4from pathlib import Path
5from typing import Optional
6import enum
8import rich
9import typer
11from .executor import Command, CommandStatus, ProcessingStrategy, read_commands_ini
13PID_FILE = ".par-run.uvicorn.pid"
15cli_app = typer.Typer()
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.")
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))
45 # Wait for UVicorn to start
46 wait_time = 3 * 10**9 # 3 seconds
47 start_time = time.time_ns()
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
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.")
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
69 with open(PID_FILE, "r") as pid_file:
70 pid = int(pid_file.read().strip())
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()
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
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.")
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
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()
129class WebCommand(enum.Enum):
130 """Web command enumeration."""
132 START = "start"
133 STOP = "stop"
134 RESTART = "restart"
135 STATUS = "status"
137 def __str__(self):
138 return self.value
141class CLICommandCBOnComp:
142 def on_start(self, cmd: Command):
143 rich.print(f"[blue bold]Completed command {cmd.name}[/]")
145 def on_recv(self, _: Command, output: str):
146 rich.print(output)
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=:}[/]")
156class CLICommandCBOnRecv:
157 def on_start(self, cmd: Command):
158 rich.print(f"[blue bold]{cmd.name}: Started[/]")
160 def on_recv(self, cmd: Command, output: str):
161 rich.print(f"{cmd.name}: {output}")
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=:}[/]")
171@cli_app.command()
172def run(
173 style: ProcessingStrategy = typer.Option(help="Processing strategy", default="comp"),
174 show: bool = typer.Option(help="Show available groups and commands", default=False),
175 file: Path = typer.Option(help="The commands.ini file to use", default=Path("commands.ini")),
176 groups: Optional[str] = typer.Option(None, help="Run a specific group of commands, comma spearated"),
177 cmds: Optional[str] = typer.Option(None, help="Run a specific commands, comma spearated"),
178):
179 """Run commands in parallel"""
180 # Overall exit code, need to track all command exit codes to update this
181 exit_code = 0
183 master_groups = read_commands_ini(file)
184 if show:
185 for grp in master_groups:
186 rich.print(f"[blue bold]Group: {grp.name}[/]")
187 for _, cmd in grp.cmds.items():
188 rich.print(f"[green bold]{cmd.name}[/]: {cmd.cmd}")
189 return
191 if groups:
192 master_groups = [grp for grp in master_groups if grp.name in [g.strip() for g in groups.split(",")]]
194 if cmds:
195 for grp in master_groups:
196 grp.cmds = OrderedDict(
197 {
198 cmd_name: cmd
199 for cmd_name, cmd in grp.cmds.items()
200 if cmd_name in [c.strip() for c in cmds.split(",")]
201 }
202 )
203 master_groups = [grp for grp in master_groups if grp.cmds]
205 if not master_groups:
206 rich.print("[blue]No groups or commands found.[/]")
207 raise typer.Exit(0)
209 for grp in master_groups:
210 if style == ProcessingStrategy.ON_COMP:
211 exit_code = exit_code or grp.run(style, CLICommandCBOnComp())
212 elif style == ProcessingStrategy.ON_RECV:
213 exit_code = exit_code or grp.run(style, CLICommandCBOnRecv())
214 else:
215 raise typer.BadParameter("Invalid processing strategy")
217 # Summarise the results
218 for grp in master_groups:
219 rich.print(f"[blue bold]Group: {grp.name}[/]")
220 for _, cmd in grp.cmds.items():
221 if cmd.status == CommandStatus.SUCCESS:
222 rich.print(f"[green bold]Command {cmd.name} succeeded ({cmd.num_non_empty_lines})[/]")
223 else:
224 rich.print(f"[red bold]Command {cmd.name} failed ({cmd.num_non_empty_lines})[/]")
226 raise typer.Exit(exit_code)
229try:
230 import os
231 import signal
232 import subprocess
233 import sys
234 import time
235 from pathlib import Path
236 from typing import Optional
238 import psutil
239 import typer
241 rich.print("[blue]Web commands loaded[/]")
243 PID_FILE = ".par-run.uvicorn.pid"
245 @cli_app.command()
246 def web(
247 command: WebCommand = typer.Argument(..., help="command to control/interract with the web server"),
248 port: int = typer.Option(8001, help="Port to run the web server"),
249 ):
250 """Run the web server"""
251 if command == WebCommand.START:
252 start_web_server(port)
253 elif command == WebCommand.STOP:
254 stop_web_server()
255 elif command == WebCommand.RESTART:
256 stop_web_server()
257 start_web_server(port)
258 elif command == WebCommand.STATUS:
259 get_web_server_status()
260 else:
261 typer.echo(f"Not a valid command '{command}'", err=True)
262 raise typer.Abort()
264except ImportError: # pragma: no cover
265 pass # pragma: no cover