Coverage for src/par_run/web_cli.py: 0%

98 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-03-29 18:20 -0400

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

2 

3import enum 

4import os 

5import signal 

6import subprocess 

7import sys 

8import time 

9from pathlib import Path 

10from typing import Optional 

11 

12import psutil 

13import typer 

14 

15 

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

17 

18cli_web = typer.Typer() 

19 

20 

21def clean_up(): 

22 """ 

23 Clean up by removing the PID file. 

24 """ 

25 os.remove(PID_FILE) 

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

27 

28 

29def start_web_server(port: int): 

30 """Start the web server""" 

31 if os.path.isfile(PID_FILE): 

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

33 sys.exit(1) 

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

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

36 uvicorn_command = [ 

37 "uvicorn", 

38 "par_run.web:ws_app", 

39 "--host", 

40 "0.0.0.0", 

41 "--port", 

42 str(port), 

43 ] 

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

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

46 

47 # Wait for UVicorn to start 

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

49 start_time = time.time_ns() 

50 

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

52 test_port = get_process_port(process.pid) 

53 if port == test_port: 

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

55 break 

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

57 

58 else: 

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

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

61 

62 

63def stop_web_server(): 

64 """ 

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

66 """ 

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

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

69 sys.exit(1) 

70 

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

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

73 

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

75 os.kill(pid, signal.SIGTERM) 

76 clean_up() 

77 

78 

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

80 process = psutil.Process(pid) 

81 connections = process.connections() 

82 if connections: 

83 port = connections[0].laddr.port 

84 return port 

85 return None 

86 

87 

88def list_uvicorn_processes(): 

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

90 uvicorn_processes = [] 

91 for process in psutil.process_iter(): 

92 try: 

93 process_name = process.name() 

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

95 uvicorn_processes.append(process) 

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

97 pass 

98 if uvicorn_processes: 

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

100 for process in uvicorn_processes: 

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

102 else: 

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

104 

105 

106def get_web_server_status(): 

107 """ 

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

109 """ 

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

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

112 list_uvicorn_processes() 

113 return 

114 

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

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

117 if psutil.pid_exists(pid): 

118 port = get_process_port(pid) 

119 if port: 

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

121 else: 

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

123 else: 

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

125 clean_up() 

126 

127 

128class WebCommand(enum.Enum): 

129 """Web command enumeration.""" 

130 

131 START = "start" 

132 STOP = "stop" 

133 RESTART = "restart" 

134 STATUS = "status" 

135 

136 def __str__(self): 

137 return self.value 

138 

139 

140@cli_web.command() 

141def web( 

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

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

144): 

145 """Run the web server""" 

146 if command == WebCommand.START: 

147 start_web_server(port) 

148 elif command == WebCommand.STOP: 

149 stop_web_server() 

150 elif command == WebCommand.RESTART: 

151 stop_web_server() 

152 start_web_server(port) 

153 elif command == WebCommand.STATUS: 

154 get_web_server_status() 

155 else: 

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

157 raise typer.Abort()