Coverage for apps/streamwise_job.py: 90%

1""" 

2StreamWise job to generate a video. 

3""" 

4from __future__ import annotations 

5 

6import os 

7import sys 

8import time 

9import logging 

10import aiofiles 

11import asyncio 

12import traceback 

13 

14from contextlib import asynccontextmanager 

15 

16from enum import Enum 

17from enum import StrEnum 

18from enum import auto 

19 

20from typing import List 

21from typing import Dict 

22from typing import Optional 

23from typing import Any 

24from typing import Union 

25from typing import AsyncIterator 

26from typing import Mapping 

27from typing import Type 

28from typing import Callable 

29 

30from datetime import datetime 

31 

32from client import ServiceRequest 

33 

34from video import VideoQuality 

35from video import QUALITY_TO_NUM_STEPS 

36 

37from resolutions import ASPECT_RATIO 

38from resolutions import RESOLUTIONS 

39 

40 

41# Local relative imports 

42sys.path.append("..") # noqa: E402 

43sys.path.append("../..") # noqa: E402 

44 

45from console_utils import setup_logging 

46 

47from lmm_generator import LMMGenerator 

48from lmm_service_manager import LMMServiceManager 

49 

50from client import ServiceError 

51 

52from character import Characters 

53 

54from console_utils import bytes_to_human 

55 

56from media_utils import get_video_file_info 

57from media_utils import get_video_frames 

58from media_utils import get_font_size 

59from media_utils import add_text_to_frame 

60 

61from k8s_utils import NoActiveContainerError 

62from k8s_utils import NoRunnableContainerError 

63from k8s_utils import ServiceNotFoundError 

64 

65 

66STATUS_EXPIRE_TIME_SECONDS = 10 * 60 # 10 minutes 

67MAX_LOG_TEXT = 100 # Max text length to log 

68SCENE_DEADLINE_INCREMENT_SECS = 5.0 # Seconds added per scene for deadline estimation 

69 

70 

71ExceptionHandler = Callable[[Exception], None] 

72 

73 

74class JobStatus(Enum): 

75 """Status of a StreamWise job.""" 

76 CREATED = auto() 

77 PENDING = auto() 

78 STARTED = auto() 

79 RUNNING = auto() 

80 RETRYING = auto() 

81 COMPLETED = auto() 

82 FAILED = auto() 

83 CANCELLED = auto() 

84 EXPIRED = auto() 

85 UNKNOWN = auto() 

86 

87 

88class OutputMode(StrEnum): 

89 """Output type of the video.""" 

90 AUDIO_ONLY = "audio_only" 

91 VIDEO_AUDIO_SYNCED = "video_audio_synced" 

92 VIDEO_AUDIO_UNSYNCED = "video_audio_unsynced" 

93 UNKNOWN = "unknown" 

94 

95 

96def get_job_id() -> str: 

97 """Generate a unique job ID based on the current timestamp: 20240605T153000123.""" 

98 return datetime.now().strftime("%Y%m%dT%H%M%S%f")[:-3] 

99 

100 

101def is_job_id(job_id: str) -> bool: 

102 """Check if the given string is a valid job ID.""" 

103 try: 

104 datetime.strptime(job_id, "%Y%m%dT%H%M%S%f") 

105 return True 

106 except ValueError: 

107 return False 

108 

109 

110def is_status_terminal(status: JobStatus) -> bool: 

111 """Check if the job status is a terminal status.""" 

112 return status in ( 

113 JobStatus.COMPLETED, 

114 JobStatus.FAILED, 

115 JobStatus.CANCELLED, 

116 ) 

117 

118 

119def is_status_expired(last_modified_time: float) -> bool: 

120 """Check if the job status has expired based on the last modified time.""" 

121 return time.time() - last_modified_time > STATUS_EXPIRE_TIME_SECONDS 

122 

123 

124class StreamWiseJob: 

125 """A generic StreamWise job to generate a video with images, audio, and text.""" 

126 

127 def __init__( 

128 self, 

129 app_name: str, 

130 job_id: str, 

131 service_manager: LMMServiceManager, 

132 config: Dict[str, Any] = {}, 

133 ) -> None: 

134 self.app_name = app_name 

135 self.job_id = job_id 

136 self.job_path = f"/tmp/{self.app_name}/{self.job_id}" 

137 

138 self.logger = self._get_logger() 

139 

140 self.service_manager = service_manager 

141 self.gen = LMMGenerator( 

142 self.app_name, 

143 self.job_id, 

144 self.service_manager) 

145 

146 """ 

147 max_tokens: int = 5 * 1024, 

148 num_characters: int = 2, 

149 max_dialogues: int = 10, 

150 max_words_per_dialogue: int = 50, 

151 edit_image: bool = True, 

152 debug_image: bool = True, 

153 upscaling: bool = True, 

154 output_mode: str = "video_audio_synced", 

155 resolution: str = "high", 

156 speech_speed: float = 1.1, 

157 """ 

158 self.config = config 

159 

160 # Useful to setup deadlines for each request 

161 self.submission_time = time.time() 

162 

163 # TODO self.get_config_str("aspect_ratio", ASPECT_RATIO) 

164 self.aspect_ratio = ASPECT_RATIO 

165 resolution_str = self.get_config_str("resolution", "high") 

166 if resolution_str == "adaptive": 

167 resolution_str = "medium" 

168 self.width, self.height = RESOLUTIONS[self.aspect_ratio][resolution_str] 

169 

170 self.characters = Characters() 

171 self.transcript_scenes: List[Any] = [] 

172 

173 # Async task for the job processing, initialized to None until the job is started 

174 self.task: Optional[asyncio.Task] = None 

175 

176 async def generate( 

177 self, 

178 job_config: Dict[str, Any], 

179 ) -> None: 

180 """Generate the video based on the job configuration.""" 

181 raise NotImplementedError("Subclasses must implement generate.") 

182 

183 def _handle_service_not_found(self, ex: Exception) -> None: 

184 self.logger.error(f"Service not found for {self.job_id}: {ex}") 

185 

186 def _handle_service_error(self, ex: Exception) -> None: 

187 self.logger.error(f"Service error for {self.job_id}: {ex}") 

188 

189 def _handle_container_error(self, ex: Exception) -> None: 

190 self.logger.error(str(ex)) 

191 

192 def _handle_value_error(self, ex: Exception) -> None: 

193 self.logger.error(f"Value error for {self.job_id}: {ex}") 

194 

195 def _handle_file_error( 

196 self, 

197 ex: Exception 

198 ) -> None: 

199 self.logger.error(f"File not found for {self.job_id}: {ex}") 

200 

201 def _handle_unknown_error(self, ex: Exception) -> None: 

202 self.logger.error(f"Error for {self.job_id} [{type(ex).__name__}]: {ex}") 

203 self.logger.error(traceback.format_exc()) 

204 

205 def _default_exception_handlers( 

206 self, 

207 ) -> dict[Type[Exception], ExceptionHandler]: 

208 """ 

209 Base exception -> handler mapping. 

210 Subclasses may extend or override this. 

211 """ 

212 return { 

213 ServiceNotFoundError: self._handle_service_not_found, 

214 ServiceError: self._handle_service_error, 

215 NoActiveContainerError: self._handle_container_error, 

216 NoRunnableContainerError: self._handle_container_error, 

217 ValueError: self._handle_value_error, 

218 FileNotFoundError: self._handle_file_error, 

219 } 

220 

221 @asynccontextmanager 

222 async def job_status_handler( 

223 self, 

224 extra_handlers: Mapping[Type[Exception], ExceptionHandler] | None = None, 

225 ) -> AsyncIterator[None]: 

226 """ 

227 Async context manager that: 

228 * Logs wall-clock execution time. 

229 * Saves job status transitions. 

230 * Catches and handles exceptions with appropriate handlers. 

231 """ 

232 await self.save_status(JobStatus.STARTED) 

233 self.logger.info(f"Starting job {self.job_id}.") 

234 

235 t0 = time.time() 

236 await self.save_status(JobStatus.RUNNING) 

237 try: 

238 yield 

239 

240 # Successful case 

241 await self.save_status(JobStatus.COMPLETED) 

242 self.logger.info(f"Job {self.job_id} completed successfully.") 

243 except Exception as ex: 

244 await self.save_status(JobStatus.FAILED) 

245 

246 handlers = self._default_exception_handlers() 

247 if extra_handlers: 

248 handlers = { 

249 **handlers, 

250 **extra_handlers 

251 } 

252 for exc_type, handler in handlers.items(): 

253 if isinstance(ex, exc_type): 

254 handler(ex) 

255 raise 

256 self._handle_unknown_error(ex) 

257 raise 

258 finally: 

259 elapsed = time.time() - t0 

260 self.logger.info(f"Job {self.job_id} finished in {elapsed:.3f} seconds.") 

261 

262 def get_submission_time(self) -> float: 

263 """Get the submission time of the job.""" 

264 return self.submission_time 

265 

266 def get_config_bool(self, config_name: str) -> bool: 

267 """Get a boolean configuration value.""" 

268 if not self.config: 

269 return False 

270 return bool(self.config.get(config_name, False)) 

271 

272 def get_config_int( 

273 self, 

274 config_name: str, 

275 default_value: int = -1 

276 ) -> int: 

277 """Get an int configuration value.""" 

278 if not self.config: 

279 return default_value 

280 return int(self.config.get(config_name, default_value)) 

281 

282 def get_config_float( 

283 self, 

284 config_name: str, 

285 default_value: float = -1.0 

286 ) -> float: 

287 """Get a float configuration value.""" 

288 if not self.config: 

289 return default_value 

290 return float(self.config.get(config_name, default_value)) 

291 

292 def get_config_str( 

293 self, 

294 config_name: str, 

295 default_ret: str = "" 

296 ) -> str: 

297 """Get a string configuration value.""" 

298 if not self.config: 

299 return default_ret 

300 return str(self.config.get(config_name, default_ret)) 

301 

302 def get_config_output_mode(self) -> OutputMode: 

303 output_mode_str = self.get_config_str("output_mode") 

304 if not output_mode_str: 

305 return OutputMode.UNKNOWN 

306 output_mode = OutputMode(output_mode_str) 

307 return output_mode 

308 

309 def get_num_steps(self) -> int: 

310 """Get the number of steps for video generation.""" 

311 quality_str = self.get_config_str("quality", VideoQuality.MEDIUM.value) 

312 num_steps = QUALITY_TO_NUM_STEPS[VideoQuality.MEDIUM.value] 

313 num_steps = QUALITY_TO_NUM_STEPS.get(quality_str, num_steps) # Default to medium 

314 return num_steps 

315 

316 async def close(self) -> None: 

317 """Close all clients.""" 

318 await self.gen.stop() 

319 

320 def _get_logger(self) -> logging.Logger: 

321 logger = setup_logging( 

322 path=self.job_path, 

323 file_name=f"job_{self.job_id}.log", 

324 level=logging.DEBUG, 

325 use_global=False) 

326 return logger 

327 

328 def _log_video_info( 

329 self, 

330 prefix: str, 

331 video_content: Union[bytes, str], 

332 ) -> None: 

333 video_file_info = get_video_file_info(video_content) 

334 video_num_bytes = video_file_info["overall"]["num_bytes"] 

335 

336 video_info = video_file_info["video"] 

337 video_fps = video_info["fps"] 

338 video_duration = video_info["duration_seconds"] 

339 video_num_frames = video_info["num_frames"] 

340 width, height = video_info["width"], video_info["height"] 

341 

342 self.logger.info( 

343 f"{prefix} with " 

344 f"{video_duration:.3f} seconds, " 

345 f"{video_num_frames} frames, " 

346 f"{video_fps} FPS, " 

347 f"{bytes_to_human(video_num_bytes)}, and " 

348 f"{width}x{height} pixels.") 

349 

350 def get_queued_requests(self) -> List[str]: 

351 return self.gen.get_queued_requests() 

352 

353 def get_requests(self) -> Dict[str, ServiceRequest]: 

354 return self.gen.get_requests() 

355 

356 async def save_status( 

357 self, 

358 status: JobStatus, 

359 ) -> None: 

360 """Save the status of a job asynchronously.""" 

361 if not isinstance(status, JobStatus): 

362 raise ValueError(f"Invalid status type: {status}. Must be a JobStatus.") 

363 await aiofiles.os.makedirs(self.job_path, exist_ok=True) 

364 status_file = os.path.join(self.job_path, "status.txt") 

365 async with aiofiles.open(status_file, mode="a") as file: 

366 await file.write(f"{time.time()},{status.value}\n") 

367 

368 async def get_status(self) -> JobStatus: 

369 """Get the latest status of a job asynchronously.""" 

370 status_file = os.path.join(self.job_path, "status.txt") 

371 if not os.path.exists(status_file): 

372 return JobStatus.UNKNOWN 

373 async with aiofiles.open(status_file, mode="r") as file: 

374 lines = await file.readlines() 

375 if not lines: 

376 return JobStatus.UNKNOWN 

377 last_line = lines[-1].strip() 

378 try: 

379 _, status_value = last_line.split(",", 1) 

380 status = JobStatus(int(status_value)) 

381 return status 

382 except (ValueError, IndexError) as e: 

383 self.logger.error(f"Error parsing status file: {e}") 

384 return JobStatus.UNKNOWN 

385 

386 def _handle_scene_exception( 

387 self, 

388 scene_id: int, 

389 ex: Exception, 

390 ) -> None: 

391 """ 

392 Handle exceptions during scene generation. 

393 Covers ServiceError, container errors (NoRunnableContainerError, 

394 NoActiveContainerError, ServiceNotFoundError), and any other exception. 

395 """ 

396 if isinstance(ex, ServiceError): 

397 self.logger.error(f"[{scene_id}] Service error: {ex}") 

398 elif isinstance(ex, (NoRunnableContainerError, NoActiveContainerError, ServiceNotFoundError)): 

399 self.logger.error(f"[{scene_id}] {ex}") 

400 else: 

401 self.logger.error(f"[{scene_id}] Error ({type(ex).__name__}): {ex}") 

402 

403 def get_scene_deadline( 

404 self, 

405 scene_index: int, 

406 ) -> float: 

407 """Get the deadline for a given scene.""" 

408 return self.get_submission_time() + (scene_index * SCENE_DEADLINE_INCREMENT_SECS) 

409 

410 async def _overlay_subtitles_on_frames( 

411 self, 

412 video_binary: bytes, 

413 text: str, 

414 ) -> List[Any]: 

415 """Overlay subtitle text onto every video frame at the bottom-center. 

416 

417 Args: 

418 video_binary: Raw video bytes from which frames and dimensions are extracted. 

419 text: Subtitle text to render on each frame. 

420 

421 Returns: 

422 A new list of frames with the subtitle burned in. 

423 """ 

424 video_frames = await get_video_frames(video_binary) 

425 video_info = get_video_file_info(video_binary)["video"] 

426 width = video_info["width"] 

427 height = video_info["height"] 

428 assert width is not None and height is not None, "Video dimensions must be available." 

429 font_size = get_font_size(width, height) 

430 font_size = font_size * 2 // 3 # Smaller font for subtitles 

431 return [ 

432 add_text_to_frame(frame, text=text, font_size=font_size, position="bottom-center") 

433 for frame in video_frames 

434 ]