Coverage for core / src / sensorkit / cli / go.py: 0%

1# SPDX-License-Identifier: Apache-2.0 

2import asyncio 

3 

4import asyncclick as click 

5from loguru import logger 

6from pydantic import BaseModel 

7 

8import sensorkit.api as sk 

9from sensorkit.cli.config import config_load 

10from sensorkit.cli.utils import report_errors 

11from sensorkit.common.logging import add_debug_logger, configure_logging 

12 

13 

14class ServiceDefinition(BaseModel): 

15 """Definition of a service to launch.""" 

16 name: str 

17 module: str 

18 func: str | None = None 

19 

20 @classmethod 

21 def from_shorthand(cls, in_str: str): 

22 """Parse a `name:module[:entrypoint]` shorthand string into a ServiceDefinition.""" 

23 name, module, *func = in_str.split(":") 

24 func = func[0] if func else None 

25 return cls(name=name, module=module, func=func) 

26 

27 

28async def read_config_file(config_file: str | None): 

29 """Load the unified config file, if any.""" 

30 sk.set_config_location(config_file) 

31 config = await sk.load_config() 

32 return [ 

33 ServiceDefinition.from_shorthand(f"{svc.id}:{svc.python_module}") 

34 for svc in config.services 

35 if svc.python_module 

36 ] 

37 

38 

39def _logger_formatter(record): 

40 color = record["extra"].get("entity_color", "magenta") 

41 return ( 

42 "{extra[entity_type]} " 

43 f" <{color}>{{extra[entity]}}</> " 

44 " <level>{message}</level>\n" 

45 ) 

46 

47 

48def _logger_patcher(record): 

49 current = sk.entity() 

50 

51 if current: 

52 entity = current.entity 

53 name = str(entity).lower() 

54 

55 if name == "agent": 

56 entity_type = "🤖 " 

57 entity_color = "fg 84" 

58 elif name == "burr": 

59 entity_type = "☕️ " 

60 entity_color = "fg 130" 

61 elif name == "otto": 

62 entity_type = "🎶 " 

63 entity_color = "fg 213" 

64 elif name == "senpai": 

65 entity_type = "🥷 " 

66 entity_color = "fg 60" 

67 elif "sky_transmission" in record["name"]: 

68 entity_type = "☁️ " 

69 entity_color = "fg 153" 

70 elif sk.device(): 

71 entity_type = "💠 " 

72 entity_color = "fg 57" 

73 elif sk.controller(): 

74 entity_type = "📡 " 

75 entity_color = "fg 39" 

76 elif sk.program(): 

77 entity_type = "⏭️ " 

78 entity_color = "fg 36" 

79 else: 

80 entity_type = "🟣 " 

81 entity_color = "fg 144" 

82 else: 

83 entity = ".".join(record["name"].split(".")[1:]) 

84 entity_type = "⚙️ " 

85 entity_color = "dim" 

86 

87 record["extra"].update( 

88 entity=entity, 

89 entity_type=entity_type, 

90 entity_color=entity_color, 

91 ) 

92 

93 

94@click.command("go") 

95@click.option( 

96 "--config", "-c", "config_file", 

97 metavar="[config_file]", 

98 required=False, 

99 type=click.Path(exists=True), 

100) 

101@click.option( 

102 "--load-config", "-l", 

103 is_flag=True, 

104 help="Automatically load configuration (unified config format only)", 

105) 

106@click.option( 

107 "--log-file", 

108 type=click.Path(dir_okay=False, writable=True, resolve_path=True), 

109 default=None, 

110 help="If set, debug log output goes here instead of the default location.", 

111) 

112@click.option( 

113 "--log-file-append", 

114 default=True, 

115 type=bool, 

116 help="When using --log-file, append to the file (default) or overwrite it.", 

117) 

118@click.option( 

119 "--log-level", 

120 default="INFO", 

121 help="Log level for console output." 

122) 

123@click.option( 

124 "--shutdown-timeout", 

125 type=float, 

126 default=180.0, 

127 help=( 

128 "Max seconds per service for graceful shutdown. Bounds the entire deinit " 

129 "phase (entity @sk.on_detach callbacks) plus backend teardown. Pick a value " 

130 "that comfortably covers your slowest hardware deinit (e.g. dome close, " 

131 "mount park)." 

132 ), 

133) 

134@click.pass_context 

135async def go_command( 

136 ctx: click.Context, 

137 load_config: bool, 

138 config_file: str | None, 

139 log_file: str | None, 

140 log_file_append: bool, 

141 log_level: str, 

142 shutdown_timeout: float, 

143): 

144 """Launch and manage services.""" 

145 from sensorkit.api.entrypoint import ServiceEntrypoint, run_services 

146 

147 # Read service configuration from file. 

148 try: 

149 services = await read_config_file(config_file) 

150 except Exception as e: 

151 logger.opt(exception=e).debug("could not load config") 

152 click.secho(f"Could not read configuration:\n{e}", fg="red", err=True) 

153 return 

154 

155 if load_config: 

156 print() 

157 print(" ▶️ Loading configuration...") 

158 print() 

159 

160 # Run the 'config load' subcommand to load config. 

161 try: 

162 async with asyncio.timeout(5.0): 

163 await ctx.invoke(config_load, file=config_file, verbose=1) 

164 except Exception as e: 

165 click.secho(f"{type(e).__name__} while loading configuration", fg="red", err=True) 

166 return 

167 

168 # Find entrypoint objects. 

169 entrypoints: dict[str, ServiceEntrypoint] = {} 

170 

171 for svc in services: 

172 try: 

173 entrypoints[svc.name] = ServiceEntrypoint.from_spec( 

174 svc.module + (f":{svc.func}" if svc.func else ""), 

175 load_file=True, 

176 ) 

177 except Exception as e: 

178 logger.opt(exception=e).debug("error loading entrypoint") 

179 click.secho( 

180 f"Failed to find entrypoint for '{svc.name}': {e}", 

181 fg="red", 

182 err=True, 

183 ) 

184 return 

185 

186 if not entrypoints: 

187 click.secho("No services defined!", fg="red", err=True) 

188 return 

189 

190 configure_logging(level=log_level, format=_logger_formatter) 

191 logger.configure(patcher=_logger_patcher) 

192 debug_log_dest = add_debug_logger(file=log_file, append=log_file_append) 

193 

194 print() 

195 print(" ▶️ Logging to", debug_log_dest) 

196 print() 

197 print(" ▶️ Starting services...") 

198 print() 

199 

200 with report_errors(): 

201 await run_services(entrypoints, max_restarts=None, shutdown_timeout=shutdown_timeout)