146 lines
5.0 KiB
Python
146 lines
5.0 KiB
Python
"""현재 state부터 다음 게이트까지 단계를 이어 돌림
|
|
어느 잡을 언제 돌릴지는 worker가 정하고 여기는 실행만 함
|
|
"""
|
|
import re
|
|
import time
|
|
import traceback
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from models.pipeline_state import PlayreelState, PosterAliveState, stage_keys
|
|
from pipelines import playreel, poster_alive
|
|
from pipelines.gate import PLAYREEL_GATE_AFTER, POSTER_ALIVE_GATE_AFTER
|
|
from tables.task import PlayreelTask, PosterAliveTask
|
|
|
|
QUEUED = "queued"
|
|
RUNNING = "running"
|
|
AWAITING_REVIEW = "awaiting_review"
|
|
FAILED = "failed"
|
|
DONE = "done"
|
|
|
|
POSTER_ALIVE_RUNNERS = {
|
|
PosterAliveState.DETECT: poster_alive.run_detect,
|
|
PosterAliveState.NARRATION_TEXT: poster_alive.run_narration_text,
|
|
PosterAliveState.MOTION: poster_alive.run_motion,
|
|
PosterAliveState.TTS: poster_alive.run_tts,
|
|
PosterAliveState.BGM: poster_alive.run_bgm,
|
|
PosterAliveState.I2V: poster_alive.run_i2v,
|
|
PosterAliveState.RENDER: poster_alive.run_render,
|
|
}
|
|
|
|
PLAYREEL_RUNNERS = {
|
|
PlayreelState.FETCH: playreel.run_fetch,
|
|
PlayreelState.SPLIT: playreel.run_split,
|
|
PlayreelState.UPSCALE: playreel.run_upscale,
|
|
PlayreelState.ANALYZE: playreel.run_analyze,
|
|
PlayreelState.MOTION: playreel.run_motion,
|
|
PlayreelState.NARRATION: playreel.run_narration,
|
|
PlayreelState.TTS: playreel.run_tts,
|
|
PlayreelState.BGM: playreel.run_bgm,
|
|
PlayreelState.I2V: playreel.run_i2v,
|
|
PlayreelState.HYBRID: playreel.run_hybrid,
|
|
PlayreelState.COMPOSE: playreel.run_compose,
|
|
# ⑫ review는 대응하는 서비스가 없어 통과만 시킨다
|
|
}
|
|
|
|
STATES = {PosterAliveTask: PosterAliveState, PlayreelTask: PlayreelState}
|
|
RUNNERS = {PosterAliveTask: POSTER_ALIVE_RUNNERS, PlayreelTask: PLAYREEL_RUNNERS}
|
|
|
|
|
|
def mark(task, stage: str, status: str) -> None:
|
|
timings = dict(task.stage_timings or {})
|
|
entry = dict(timings.get(stage) or {"started": None, "ended": None})
|
|
entry["status"] = status
|
|
if status == RUNNING:
|
|
entry["started"] = time.time()
|
|
elif status in ("done", FAILED):
|
|
entry["ended"] = time.time()
|
|
timings[stage] = entry
|
|
task.stage_timings = timings
|
|
|
|
|
|
def reset_from(task, stage: str) -> None:
|
|
"""되돌리기·재시도에서 그 단계부터 다시 돌게 만든다"""
|
|
state_type = STATES[type(task)]
|
|
keys = stage_keys(state_type)
|
|
timings = dict(task.stage_timings or {})
|
|
for key in keys[keys.index(stage):]:
|
|
timings[key] = {"status": "idle", "started": None, "ended": None}
|
|
task.stage_timings = timings
|
|
task.state = state_type(stage)
|
|
|
|
|
|
URL_IN_TEXT = re.compile(r"https?://\S+")
|
|
MAX_ERROR_DETAIL = 500
|
|
|
|
|
|
def safe_detail(failure: BaseException) -> str:
|
|
"""검수 화면까지 가는 값이라 트레이스백과 스토리지 주소는 뺀다"""
|
|
message = f"{type(failure).__name__}: {failure}"
|
|
return URL_IN_TEXT.sub("<url>", message)[:MAX_ERROR_DETAIL]
|
|
|
|
|
|
def stops_here(task, stage) -> str | None:
|
|
"""이 단계 뒤에 사람이 볼 게이트가 서는가. playreel은 게이트 이름을 돌려준다"""
|
|
if isinstance(task, PlayreelTask):
|
|
gate = PLAYREEL_GATE_AFTER.get(stage)
|
|
return gate.value if gate else None
|
|
if stage == POSTER_ALIVE_GATE_AFTER and not task.skip_review:
|
|
return AWAITING_REVIEW
|
|
return None
|
|
|
|
|
|
async def step(session: AsyncSession, task) -> bool:
|
|
"""한 단계를 돌린다. 게이트에 닿으면 False를 돌려 멈춤을 알린다"""
|
|
state_type = STATES[type(task)]
|
|
stage = state_type(task.state)
|
|
runner = RUNNERS[type(task)].get(stage)
|
|
|
|
if runner is None: # 러너 없는 마지막 단계
|
|
mark(task, stage.value, "done")
|
|
task.state = state_type.COMPLETED
|
|
await session.commit()
|
|
return True
|
|
|
|
mark(task, stage.value, RUNNING)
|
|
await session.commit()
|
|
|
|
await runner(session, task) # state를 다음 단계로 올리고 커밋한다
|
|
|
|
mark(task, stage.value, "done")
|
|
gate = stops_here(task, stage)
|
|
if gate:
|
|
task.status = AWAITING_REVIEW
|
|
if isinstance(task, PlayreelTask):
|
|
task.gate = gate
|
|
await session.commit()
|
|
return False
|
|
await session.commit()
|
|
return True
|
|
|
|
|
|
async def run_until_gate(session: AsyncSession, task) -> None:
|
|
"""다음 게이트나 끝까지 민다. 실패는 잡에 기록하고 조용히 돌아온다"""
|
|
state_type = STATES[type(task)]
|
|
try:
|
|
while task.state != state_type.COMPLETED:
|
|
if not await step(session, task):
|
|
return
|
|
except Exception as failure:
|
|
failed_stage = str(task.state)
|
|
traceback.print_exc() # 전문은 서버 로그에만 남긴다
|
|
await session.rollback()
|
|
await session.refresh(task)
|
|
mark(task, failed_stage, FAILED)
|
|
task.status = FAILED
|
|
task.error_stage = failed_stage
|
|
task.error_detail = safe_detail(failure)
|
|
await session.commit()
|
|
return
|
|
|
|
task.status = DONE
|
|
if isinstance(task, PlayreelTask):
|
|
task.gate = None
|
|
task.version += 1
|
|
await session.commit()
|