Compare commits
56 Commits
c8793cc3cf
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
| 4fb3d12244 | |||
| 789a807d50 | |||
| ad141a2887 | |||
| 6774067505 | |||
| c451f5313b | |||
| 9241b5cd90 | |||
| 8bfc8e153f | |||
| 232aa52adb | |||
| d2f7030446 | |||
| 43ee610780 | |||
| f79c5c26df | |||
| 7108e5e4f5 | |||
| 1e6638a64b | |||
| 32308bede6 | |||
| ac6409605c | |||
| e4d02b8059 | |||
| 94a034ef38 | |||
| 2a11d05f4a | |||
| c2e77a7310 | |||
| bea27a75cf | |||
| 39adfc5fc5 | |||
| 1a848faac4 | |||
| cb70226f42 | |||
| de24bae984 | |||
| 0e6c893b4e | |||
| fb80973e38 | |||
| 31b0e7dbc4 | |||
| 6169f48eb8 | |||
| 27a6df6cff | |||
| 803fdb6278 | |||
| 77e21b54e6 | |||
| 4d0c89ce79 | |||
| 4b60ab34c3 | |||
| 53a0657027 | |||
| 91f01d126b | |||
| 0702cf052f | |||
| 8aa3f1c3b2 | |||
| 4db0551d33 | |||
| 4d837fdd31 | |||
| 2567a6f10b | |||
| 17ed1943f1 | |||
| 8d246b5b32 | |||
| b4bec9d51b | |||
| f32792e4a9 | |||
| f152545d3b | |||
| bf3d6ee694 | |||
| 44bc065796 | |||
| 9127616669 | |||
| 900f45c2ff | |||
| eb34cbc0f7 | |||
| 0de09613d2 | |||
| a5274a4fa7 | |||
| 4e72f8ca2e | |||
| 44c6811352 | |||
| 9eef2c5015 | |||
| b05e5714e3 |
9
.mcp.json
Normal file
9
.mcp.json
Normal file
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"mcpServers": {
|
||||
"co-gahusb": {
|
||||
"type": "http",
|
||||
"url": "https://gahusb.synology.me/api/co/mcp",
|
||||
"headers": { "Authorization": "Bearer ${CO_BUS_KEY}" }
|
||||
}
|
||||
}
|
||||
}
|
||||
11
CLAUDE.md
11
CLAUDE.md
@@ -137,3 +137,14 @@ cd C:\Users\jaeoh\Desktop\workspace\web-ai\ai_trade
|
||||
- **spec amendment 발생 시**: 코드는 `web-ai` 에 commit, spec 갱신은 `web-ui/docs/superpowers/specs/` 에 commit (Phase 4 spread formula 변경 사례 = web-ui commit `534ded5`)
|
||||
|
||||
자세한 V1 가이드는 `signal_v1/CLAUDE.md` 참조 (있다면).
|
||||
|
||||
---
|
||||
|
||||
## 협업 팀 버스 (co-gahusb) — 이 세션의 역할: **AI**
|
||||
|
||||
이 세션은 AI 리서치(AI) 역할이다. co-gahusb MCP 툴로 다른 세션(FE/BE/Producer)과 협업한다.
|
||||
- **소유권**: 이 세션은 `web-ai` repo만 쓴다(FE=web-ui, BE=web-backend).
|
||||
- **공유 리소스 변경 전 반드시 `acquire_lock(resource, "AI")`**: 대상 = `nas-deploy`, `stock-db-schema`, `lotto-db-schema`, `memory-mirror`, `nginx-conf`, `compose`. 점유 중이면 대기, 긴 작업은 `heartbeat_lock`, 끝나면 `release_lock`.
|
||||
- **모든 툴 호출에 `role="AI"`** (또는 `from_role`/`created_by`에 AI).
|
||||
- **수신**: `/loop`로 주기적으로 `read_inbox("AI", after_id=<last>)` + `list_tasks(assignee_role="AI")` 확인.
|
||||
- 키 `CO_BUS_KEY`는 환경변수로 주입(커밋 금지). `.mcp.json`의 `${CO_BUS_KEY}`가 프로세스 환경변수에서 치환됨 → `setx CO_BUS_KEY "..."` 후 새 터미널에서 `claude` 실행.
|
||||
|
||||
211
README.md
Normal file
211
README.md
Normal file
@@ -0,0 +1,211 @@
|
||||
# web-ai
|
||||
|
||||
Windows AI 머신(AMD 9800X3D + RTX 5070 Ti 16GB)에서 동작하는 두 영역의 서비스:
|
||||
|
||||
1. **ai_trade** — Confidence Signal Pipeline V2. NAS stock 백엔드와 KIS Open API를 결합해 매수/매도 신호를 생성하는 FastAPI 워커.
|
||||
2. **services** — NAS↔Windows 분산 렌더링 워커(인스타 카드 / 음악 / 영상 / 이미지) + task-watcher.
|
||||
|
||||
상위 워크스페이스 컨텍스트는 `../CLAUDE.md`, 본 디렉토리 상세는 `CLAUDE.md`, 운영 체크포인트는 `CHECK_POINT.md` 참조.
|
||||
|
||||
---
|
||||
|
||||
## 디렉토리 구조
|
||||
|
||||
| 경로 | 역할 | 포트 |
|
||||
|------|------|------|
|
||||
| `ai_trade/` | 자동매매 메인. Chronos-bolt(또는 Chronos-2) + 분봉 모멘텀 + KIS WebSocket 호가 + 매수/매도 신호 생성기. | `:8001` |
|
||||
| `services/_shared/` | 4개 render worker 공통 모듈 (`ReliableQueue` — BLMOVE + ack/fail + recovery). | — |
|
||||
| `services/insta-render/` | Instagram 카드 Playwright 렌더 워커. NAS Redis `queue:insta-render` 소비. | `:18710` |
|
||||
| `services/music-render/` | Suno + MusicGen 음악 생성 워커. `queue:music-render` 소비. | `:18711` |
|
||||
| `services/video-render/` | sora / veo / kling / seedance 4 provider 영상 생성 게이트웨이. `queue:video-render` 소비. | `:18712` |
|
||||
| `services/image-render/` | gpt_image / nano_banana / flux(ComfyUI 로컬) 3 provider. `queue:image-render` 소비. | `:18714` |
|
||||
| `services/task-watcher/` | 박재오 작업 시간대에 `queue:paused` 토글 → 워커 일시 정지. | `:18713` |
|
||||
| `legacy/signal_v1/` | ⚠ **DEPRECATED** (2026-05-19). LSTM 봇. 자동 실행 차단됨. | OFF |
|
||||
|
||||
---
|
||||
|
||||
## ai_trade — Confidence Signal Pipeline V2
|
||||
|
||||
NAS stock 백엔드(`:18500`)에서 portfolio / news_sentiment / screener를 pull하고, KIS REST/WebSocket으로 분봉·호가를 보강한 뒤 Chronos 예측과 5분봉 모멘텀 분류로 매수/매도 신호를 생성한다.
|
||||
|
||||
### 매수 (screener Top-N + portfolio)
|
||||
|
||||
모두 충족 시 confidence 계산 → threshold 초과 시 emit:
|
||||
|
||||
1. `chronos.median > 0`
|
||||
2. `chronos.q90 - chronos.q10 < 0.6` (absolute spread)
|
||||
3. `minute_momentum == strong_up`
|
||||
4. `asking_price.bid_ratio >= 0.6`
|
||||
|
||||
종합 confidence = `chronos_conf * 0.5 + minute_score * 0.3 + screener_norm * 0.2`. `> 0.7` 시 emit.
|
||||
|
||||
### 매도 (portfolio only, 우선순위 stop_loss → anomaly → take_profit)
|
||||
|
||||
- **stop_loss**: `pnl_pct < -7%` 즉시 (confidence=1.0)
|
||||
- **anomaly**: `chronos.median < -1%` + `strong_down` + `bid_ratio < 0.4` + 종합 conf > 0.7
|
||||
- **take_profit**: `pnl_pct > 15%` 검토 (confidence=0.6)
|
||||
|
||||
### 핵심 파일
|
||||
|
||||
| 파일 | 책임 |
|
||||
|------|------|
|
||||
| `main.py` | FastAPI app + lifespan (의존성 wiring) + poll_loop task 생성 |
|
||||
| `config.py` | `Settings` dataclass — 환경변수 로드 |
|
||||
| `state.py` | `PollState` (process-wide singleton) — portfolio·screener·signals 등 + `get_active_signals` / `purge_expired_signals` |
|
||||
| `stock_client.py` | NAS stock 백엔드 pull (X-WebAI-Key + 메모리 캐시) |
|
||||
| `kis_client.py` | KIS REST 분봉/호가 + asyncio.Lock 직렬화 + 지수 backoff |
|
||||
| `kis_websocket.py` | KIS WebSocket 호가 + approval_key + 재연결 |
|
||||
| `chronos_predictor.py` | HuggingFace Chronos zero-shot 분위수 예측 (FP32 강제) |
|
||||
| `minute_momentum.py` | 5분봉 → strong_up / weak_up / neutral / weak_down / strong_down |
|
||||
| `signal_generator.py` | 매수/매도 룰 엔진. cycle_id + expires_at 부착 |
|
||||
| `pull_worker.py` | asyncio cron — 시간대별 분기 + post-close 트리거 + signal 생성 + expired purge |
|
||||
| `scheduler.py` | 폴링 윈도우 판정 (KST 캘린더 + 휴장일) |
|
||||
| `rate_limit.py` | 초당 N회 token bucket + `SignalDedup` SQLite WAL |
|
||||
|
||||
### 시작
|
||||
|
||||
```bat
|
||||
cd ai_trade
|
||||
start.bat
|
||||
```
|
||||
|
||||
→ `Uvicorn running on http://0.0.0.0:8001`, `poll_loop started`.
|
||||
|
||||
휴장일/장 외 시간엔 poll_loop만 idle.
|
||||
|
||||
### 헬스 / 로그
|
||||
|
||||
```powershell
|
||||
curl http://localhost:8001/health
|
||||
Get-Content logs\ai_trade.log -Wait
|
||||
nvidia-smi
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## services — NAS↔Windows 분산 워커
|
||||
|
||||
NAS측 lab 서비스(insta-lab / music-lab / video-lab / image-render NAS측)가 `queue:<worker>-render` 에 LPUSH로 작업을 enqueue. Windows worker가 BLMOVE로 atomic dequeue 후 처리, 완료 시 NAS internal webhook으로 결과 통지.
|
||||
|
||||
### 신뢰성 패턴 (`_shared.ReliableQueue`)
|
||||
|
||||
- **dequeue**: `BLMOVE main → processing:<queue>:<worker_id>` (atomic).
|
||||
- **ack**: `LREM processing 1 raw` (성공).
|
||||
- **fail**: `LREM processing` → `attempts++` 후 main 재큐 또는 `max_attempts` 도달 시 `dead_letter:<queue>` 이동.
|
||||
- **recover**: startup 시 자신의 processing list orphan을 main queue로 (attempts 증가).
|
||||
|
||||
### 시작 (NAS, WSL2 Docker)
|
||||
|
||||
```bash
|
||||
cd services
|
||||
docker compose up -d insta-render music-render video-render image-render task-watcher
|
||||
```
|
||||
|
||||
build context는 `services/` 루트. 각 Dockerfile은 `_shared` 모듈을 함께 COPY하고 `PYTHONPATH=/app`.
|
||||
|
||||
### 운영 조작
|
||||
|
||||
```bash
|
||||
# 워커 일시 정지 / 재개
|
||||
redis-cli -h 192.168.45.54 SET queue:paused 1
|
||||
redis-cli -h 192.168.45.54 DEL queue:paused
|
||||
|
||||
# 큐 / dead-letter 점검
|
||||
redis-cli -h 192.168.45.54 LLEN queue:insta-render
|
||||
redis-cli -h 192.168.45.54 LLEN dead_letter:queue:insta-render
|
||||
redis-cli -h 192.168.45.54 KEYS 'processing:*'
|
||||
```
|
||||
|
||||
### 환경 변수
|
||||
|
||||
| 변수 | 용도 |
|
||||
|------|------|
|
||||
| `REDIS_URL` | NAS Redis (`redis://192.168.45.54:6379`) |
|
||||
| `NAS_BASE_URL` | NAS 대상 서비스 URL (insta-lab `:18700`, music-lab `:18600`, video-lab `:18801`, image-render NAS측 `:18802`) |
|
||||
| `INTERNAL_API_KEY` | NAS internal webhook 인증 |
|
||||
| `WORKER_ID` | (권장) `<service>-prod-1` 등 영속 ID. hostname 기반 default는 컨테이너 재기동 시 바뀌어 orphan 추적 불가 |
|
||||
| `OPENAI_API_KEY` / `GEMINI_API_KEY` / `KLING_*` / `SEEDANCE_API_KEY` / `SUNO_API_KEY` | 각 provider 인증 |
|
||||
| `COMFYUI_URL` | image-render FLUX 로컬 ComfyUI (`http://host.docker.internal:8188`) |
|
||||
| `FLUX_BLOCK_TRADING_HOURS` | `1` 이면 장중(09:00~15:30) FLUX 차단 (Chronos GPU 보호) |
|
||||
|
||||
---
|
||||
|
||||
## 환경 변수 (ai_trade)
|
||||
|
||||
| 변수 | 기본 | 설명 |
|
||||
|------|------|------|
|
||||
| `STOCK_API_URL` | (필수) | NAS stock 백엔드 base URL |
|
||||
| `WEBAI_API_KEY` | (필수) | stock 백엔드 호출 시 X-WebAI-Key |
|
||||
| `SIGNAL_V2_PORT` | `8001` | uvicorn 포트 |
|
||||
| `KIS_ENV_TYPE` | `virtual` | `virtual` / `real` |
|
||||
| `KIS_REAL_APP_KEY` / `KIS_REAL_APP_SECRET` / `KIS_REAL_ACCOUNT` | — | KIS 실계좌 |
|
||||
| `KIS_VIRTUAL_APP_KEY` / `KIS_VIRTUAL_APP_SECRET` / `KIS_VIRTUAL_ACCOUNT` | — | KIS 모의계좌 |
|
||||
| `V1_TOKEN_PATH` | `legacy/signal_v1/data/kis_token.json` | KIS 토큰 파일 (V1 토큰 read-only 공유) |
|
||||
| `CHRONOS_MODEL` | `amazon/chronos-2` | Chronos 모델 ID |
|
||||
| `STOP_LOSS_PCT` | `-0.07` | 손절 임계 |
|
||||
| `TAKE_PROFIT_PCT` | `0.15` | 익절 임계 |
|
||||
| `CHRONOS_SPREAD_THRESHOLD` | `0.6` | 매수 hard gate spread 상한 |
|
||||
| `ASKING_BID_RATIO_THRESHOLD` | `0.6` | 매수 hard gate 호가 비율 |
|
||||
| `CONFIDENCE_THRESHOLD` | `0.7` | 매수 종합 confidence 하한 |
|
||||
| `MIN_MOMENTUM_FOR_BUY` | `strong_up` | 매수 hard gate 모멘텀 단계 |
|
||||
| `SIGNAL_TTL_SECONDS` | `300` | emit signal expires_at TTL |
|
||||
|
||||
`.env` 는 web-ai 루트 (이 디렉토리)에 둔다. **절대 커밋 금지.**
|
||||
|
||||
---
|
||||
|
||||
## 테스트
|
||||
|
||||
```bash
|
||||
# ai_trade
|
||||
python -m pytest ai_trade/tests -q
|
||||
|
||||
# services/_shared 공통 모듈
|
||||
cd services/_shared && python -m pytest tests/ -q
|
||||
|
||||
# 각 worker
|
||||
cd services/insta-render && python -m pytest tests/ -q
|
||||
cd services/music-render && python -m pytest tests/ -q
|
||||
cd services/video-render && python -m pytest tests/ -q
|
||||
cd services/image-render && python -m pytest tests/ -q
|
||||
```
|
||||
|
||||
**`.venv` 한글 사용자 경로 깨짐**으로 시스템 Python(`C:\Users\jaeoh\AppData\Local\Programs\Python\Python312\python.exe`) 사용 권장. 또는 `py -3.12 -m pytest …`.
|
||||
|
||||
---
|
||||
|
||||
## 알려진 함정
|
||||
|
||||
1. **KIS rate limit (EGW00201)** — V1+V2 동시 실행 시 충돌. V1은 `legacy/`로 격리. ai_trade는 `asyncio.Lock`으로 throttle 직렬화 (`kis_client.py`).
|
||||
2. **`.venv` 한글 경로** — 시스템 Python 사용.
|
||||
3. **Chronos FP16 overflow** — 한국 주가 5만원+ 시 inf. FP32 강제됨.
|
||||
4. **post-close 트리거** — 상태기반(`last_post_close_date`)으로 변경됨. 16:00 이후 + 오늘 미실행이면 trigger.
|
||||
5. **services worker_id** — env로 명시 권장. hostname 기반 default는 컨테이너 재기동 시 바뀌어 orphan 분실 위험.
|
||||
6. **dead-letter 누적** — `redis-cli LLEN dead_letter:*` 정기 점검 필요.
|
||||
7. **Dockerfile build context** — `services/` 루트 (각 worker 디렉토리 아님). compose 변경 동반.
|
||||
|
||||
---
|
||||
|
||||
## Phase 진행 상태 (Confidence Signal Pipeline V2)
|
||||
|
||||
| Phase | 내용 | 상태 |
|
||||
|-------|------|------|
|
||||
| 0 | Architecture & contract spec | ✅ |
|
||||
| 1 | stock 백엔드 WebAI API 보강 (NAS) | ✅ |
|
||||
| 1.5 | V1 → `signal_v1/` rename → `legacy/` 격리 | ✅ |
|
||||
| 2 | ai_trade pull worker + signal API client + scheduler | ✅ |
|
||||
| 3a | KIS REST 분봉 + WebSocket 호가 + NXT 스케줄 | ✅ |
|
||||
| 3b | Chronos-bolt-base 추론 + 5분봉 모멘텀 분류기 | ✅ |
|
||||
| 4 | Signal Generator + 로깅 | ✅ |
|
||||
| 4.5 | 코드 리뷰 F1-F6 hotfix (토큰 경로 / throttle Lock / post-close 상태기반 / Chronos abs / state.signals lifecycle / render queue 신뢰성) | ✅ |
|
||||
| 5 | agent-office `/signal` + Ollama Qwen3 14B + 이중 텔레그램 | ⏳ |
|
||||
| 6 | signal_v1 deprecation (legacy 완료, 아카이브만 남음) | 일부 ✅ |
|
||||
| 7 | 운영 모니터링 + 4주 IC 검증 | ⏳ |
|
||||
|
||||
상세 spec/plan은 `../web-ui/docs/superpowers/specs/` / `../web-ui/docs/superpowers/plans/` (별도 repo).
|
||||
|
||||
---
|
||||
|
||||
## 라이선스 / 사용
|
||||
|
||||
비공개. 박재오 개인 웹 플랫폼.
|
||||
@@ -10,6 +10,10 @@ import numpy as np
|
||||
logger = logging.getLogger(__name__)
|
||||
KST = ZoneInfo("Asia/Seoul")
|
||||
|
||||
# F4: signal_generator hard gate와 동일한 absolute spread threshold.
|
||||
# zero-shot median≈0에서 conf가 0으로 폭락하던 relative 산식 (spread/abs(median)) 대체.
|
||||
_SPREAD_THRESHOLD = 0.6
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChronosPrediction:
|
||||
@@ -103,8 +107,8 @@ class ChronosPredictor:
|
||||
median = float((q50_price - last_close) / last_close)
|
||||
q10 = float((q10_price - last_close) / last_close)
|
||||
q90 = float((q90_price - last_close) / last_close)
|
||||
spread = (q90 - q10) / max(abs(median), 0.001)
|
||||
conf = float(max(0.0, min(1.0, 1.0 - spread / 2.0)))
|
||||
spread = q90 - q10 # F4: absolute spread
|
||||
conf = float(max(0.0, min(1.0, 1.0 - spread / _SPREAD_THRESHOLD)))
|
||||
results[ticker] = ChronosPrediction(
|
||||
median=median, q10=q10, q90=q90, conf=conf, as_of=now_iso,
|
||||
)
|
||||
@@ -124,8 +128,8 @@ class ChronosPredictor:
|
||||
median = float(np.quantile(returns, 0.5))
|
||||
q10 = float(np.quantile(returns, 0.1))
|
||||
q90 = float(np.quantile(returns, 0.9))
|
||||
spread = (q90 - q10) / max(abs(median), 0.001)
|
||||
conf = float(max(0.0, min(1.0, 1.0 - spread / 2.0)))
|
||||
spread = q90 - q10 # F4: absolute spread
|
||||
conf = float(max(0.0, min(1.0, 1.0 - spread / _SPREAD_THRESHOLD)))
|
||||
results[ticker] = ChronosPrediction(
|
||||
median=median, q10=q10, q90=q90, conf=conf, as_of=now_iso,
|
||||
)
|
||||
|
||||
@@ -31,7 +31,7 @@ class Settings:
|
||||
v1_token_path: Path = field(
|
||||
default_factory=lambda: Path(
|
||||
os.getenv("V1_TOKEN_PATH",
|
||||
str(Path(__file__).parent.parent / "signal_v1" / "data" / "kis_token.json"))
|
||||
str(Path(__file__).parent.parent / "legacy" / "signal_v1" / "data" / "kis_token.json"))
|
||||
)
|
||||
)
|
||||
chronos_model: str = field(default_factory=lambda: os.getenv("CHRONOS_MODEL", "amazon/chronos-2"))
|
||||
@@ -53,6 +53,9 @@ class Settings:
|
||||
min_momentum_for_buy: str = field(
|
||||
default_factory=lambda: os.getenv("MIN_MOMENTUM_FOR_BUY", "strong_up")
|
||||
)
|
||||
signal_ttl_seconds: int = field(
|
||||
default_factory=lambda: int(os.getenv("SIGNAL_TTL_SECONDS", "300"))
|
||||
)
|
||||
|
||||
@property
|
||||
def kis_is_virtual(self) -> bool:
|
||||
|
||||
@@ -38,6 +38,7 @@ class KISClient:
|
||||
self._client = httpx.AsyncClient(timeout=timeout)
|
||||
self._token_cache: tuple[str, float] | None = None # (token, file_mtime)
|
||||
self._last_throttle_at = 0.0
|
||||
self._throttle_lock = asyncio.Lock()
|
||||
|
||||
async def close(self) -> None:
|
||||
await self._client.aclose()
|
||||
@@ -56,10 +57,13 @@ class KISClient:
|
||||
return token
|
||||
|
||||
async def _throttle(self) -> None:
|
||||
elapsed = time.monotonic() - self._last_throttle_at
|
||||
if elapsed < _THROTTLE_INTERVAL:
|
||||
await asyncio.sleep(_THROTTLE_INTERVAL - elapsed)
|
||||
self._last_throttle_at = time.monotonic()
|
||||
# F2: Lock으로 직렬화. 없으면 asyncio.gather 동시 호출 시 race로
|
||||
# 같은 elapsed 계산 후 동시에 깨어나 KIS 초당 2회(EGW00201) 위반.
|
||||
async with self._throttle_lock:
|
||||
elapsed = time.monotonic() - self._last_throttle_at
|
||||
if elapsed < _THROTTLE_INTERVAL:
|
||||
await asyncio.sleep(_THROTTLE_INTERVAL - elapsed)
|
||||
self._last_throttle_at = time.monotonic()
|
||||
|
||||
def _common_headers(self, tr_id: str) -> dict[str, str]:
|
||||
token = self._read_v1_token()
|
||||
|
||||
@@ -24,6 +24,7 @@ async def poll_loop(
|
||||
) -> None:
|
||||
"""FastAPI lifespan 에서 asyncio.create_task 로 시작."""
|
||||
logger.info("poll_loop started")
|
||||
last_post_close_date = None # F3: state-based post-close trigger
|
||||
while not shutdown.is_set():
|
||||
now = datetime.now(KST)
|
||||
if _is_market_day(now) and _is_polling_window(now):
|
||||
@@ -36,10 +37,14 @@ async def poll_loop(
|
||||
update_minute_momentum_for_all(state)
|
||||
except Exception:
|
||||
logger.exception("minute momentum update failed")
|
||||
# Post-close trigger (16:00 KST)
|
||||
if _is_post_close_trigger(now) and chronos is not None and kis_client is not None:
|
||||
# Post-close trigger (F3: 상태기반 — 16:00 이후 + 오늘 미실행)
|
||||
if (
|
||||
_is_post_close_trigger(now, last_post_close_date)
|
||||
and chronos is not None and kis_client is not None
|
||||
):
|
||||
try:
|
||||
await _run_post_close_cycle(kis_client, chronos, state)
|
||||
last_post_close_date = now.date()
|
||||
except Exception:
|
||||
logger.exception("post-close cycle failed")
|
||||
# Phase 4: generate signals
|
||||
@@ -49,6 +54,11 @@ async def poll_loop(
|
||||
generate_signals(state, dedup, settings)
|
||||
except Exception:
|
||||
logger.exception("generate_signals failed")
|
||||
# F5: cycle 끝에 expired signal purge (consumer 미사용 케이스 보호)
|
||||
try:
|
||||
state.purge_expired_signals(datetime.now(KST))
|
||||
except Exception:
|
||||
logger.exception("purge_expired_signals failed")
|
||||
interval = _next_interval(now)
|
||||
try:
|
||||
await asyncio.wait_for(shutdown.wait(), timeout=interval)
|
||||
|
||||
@@ -76,12 +76,21 @@ def _seconds_until_nxt_or_market_open(now: datetime) -> float:
|
||||
return 86400.0
|
||||
|
||||
|
||||
def _is_post_close_trigger(now: datetime) -> bool:
|
||||
"""16:00 KST ±1분 (post-close cycle 트리거). 평일/영업일만."""
|
||||
def _is_post_close_trigger(now: datetime, last_post_close_date) -> bool:
|
||||
"""F3 — 16:00 KST 이후 오늘 아직 post-close cycle 안 돌렸으면 True (상태기반).
|
||||
|
||||
이전엔 16:00:00-16:00:59 1분 윈도우라 5분 sleep + 비결정적 cycle 시작시각
|
||||
조합으로 영영 못 잡는 경우 발생 (예: cycle이 15:31에 시작되면 16:01에 깸).
|
||||
|
||||
Args:
|
||||
now: 현재 KST datetime.
|
||||
last_post_close_date: 마지막 post-close 실행 영업일 date (None=미실행).
|
||||
"""
|
||||
if not _is_market_day(now):
|
||||
return False
|
||||
t = now.time()
|
||||
return time(16, 0) <= t < time(16, 1)
|
||||
if now.time() < time(16, 0):
|
||||
return False
|
||||
return last_post_close_date != now.date()
|
||||
|
||||
|
||||
def _seconds_until_next_market_open(now: datetime) -> float:
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timedelta
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -20,7 +20,12 @@ MOMENTUM_SCORES = {
|
||||
|
||||
|
||||
def generate_signals(state, dedup, settings) -> None:
|
||||
"""Phase 4 entry — state-mutating. Evaluation order: sell first (priority), then buy. A ticker receiving a sell signal in this cycle is excluded from buy evaluation to avoid silent overwrite."""
|
||||
"""Phase 4 entry — state-mutating. F5: cycle_id += 1 (호출마다, emit 여부 무관).
|
||||
|
||||
Evaluation order: sell first (priority), then buy. A ticker receiving a sell
|
||||
signal in this cycle is excluded from buy evaluation to avoid silent overwrite.
|
||||
"""
|
||||
state.signal_cycle_id += 1
|
||||
_evaluate_sell_signals(state, dedup, settings)
|
||||
_evaluate_buy_signals(state, dedup, settings)
|
||||
|
||||
@@ -45,9 +50,10 @@ def _evaluate_buy_signals(state, dedup, settings) -> None:
|
||||
if dedup.is_recent(ticker, "buy", within_hours=24):
|
||||
logger.debug("buy %s skipped: dedup 24h", ticker)
|
||||
continue
|
||||
state.signals[ticker] = _build_buy_signal(state, ticker, name, rank, confidence)
|
||||
state.signals[ticker] = _build_buy_signal(state, ticker, name, rank, confidence, settings)
|
||||
dedup.record(ticker, "buy", confidence=confidence)
|
||||
logger.info("signal emit %s buy conf=%.3f rank=%s", ticker, confidence, rank)
|
||||
logger.info("signal emit %s buy conf=%.3f rank=%s cycle=%d",
|
||||
ticker, confidence, rank, state.signal_cycle_id)
|
||||
|
||||
|
||||
def _buy_candidates(state) -> list[tuple[str, str, int | None]]:
|
||||
@@ -96,8 +102,11 @@ def _compute_buy_confidence(state, ticker: str, rank: int | None) -> float:
|
||||
return chronos_conf * 0.5 + minute_score * 0.3 + screener_norm * 0.2
|
||||
|
||||
|
||||
def _build_buy_signal(state, ticker: str, name: str, rank: int | None, confidence: float) -> dict:
|
||||
def _build_buy_signal(state, ticker: str, name: str, rank: int | None, confidence: float, settings) -> dict:
|
||||
ap = state.asking_price[ticker]
|
||||
as_of_dt = datetime.now(KST)
|
||||
ttl = getattr(settings, "signal_ttl_seconds", 300)
|
||||
expires_at = (as_of_dt + timedelta(seconds=ttl)).isoformat()
|
||||
return {
|
||||
"ticker": ticker,
|
||||
"name": name,
|
||||
@@ -107,7 +116,9 @@ def _build_buy_signal(state, ticker: str, name: str, rank: int | None, confidenc
|
||||
"avg_price": None,
|
||||
"pnl_pct": None,
|
||||
"context": _build_context(state, ticker, rank),
|
||||
"as_of": datetime.now(KST).isoformat(),
|
||||
"as_of": as_of_dt.isoformat(),
|
||||
"cycle_id": state.signal_cycle_id,
|
||||
"expires_at": expires_at,
|
||||
}
|
||||
|
||||
|
||||
@@ -132,23 +143,24 @@ def _evaluate_sell_signals(state, dedup, settings) -> None:
|
||||
continue
|
||||
state.signals[ticker] = sell
|
||||
dedup.record(ticker, "sell", confidence=sell["confidence_webai"])
|
||||
logger.info("signal emit %s sell conf=%.3f reason=%s",
|
||||
logger.info("signal emit %s sell conf=%.3f reason=%s cycle=%d",
|
||||
ticker, sell["confidence_webai"],
|
||||
sell.get("context", {}).get("sell_reason"))
|
||||
sell.get("context", {}).get("sell_reason"),
|
||||
state.signal_cycle_id)
|
||||
|
||||
|
||||
def _try_stop_loss(state, holding: dict, settings) -> dict | None:
|
||||
pnl = holding.get("pnl_pct")
|
||||
if pnl is None or pnl >= settings.stop_loss_pct:
|
||||
return None
|
||||
return _build_sell_signal(state, holding, confidence=1.0, reason="stop_loss")
|
||||
return _build_sell_signal(state, holding, confidence=1.0, reason="stop_loss", settings=settings)
|
||||
|
||||
|
||||
def _try_take_profit(state, holding: dict, settings) -> dict | None:
|
||||
pnl = holding.get("pnl_pct")
|
||||
if pnl is None or pnl <= settings.take_profit_pct:
|
||||
return None
|
||||
return _build_sell_signal(state, holding, confidence=0.6, reason="take_profit")
|
||||
return _build_sell_signal(state, holding, confidence=0.6, reason="take_profit", settings=settings)
|
||||
|
||||
|
||||
def _try_anomaly(state, holding: dict, settings) -> dict | None:
|
||||
@@ -168,11 +180,14 @@ def _try_anomaly(state, holding: dict, settings) -> dict | None:
|
||||
confidence = pred["conf"] * 0.5 + minute_score * 0.3 + 1.0 * 0.2
|
||||
if confidence <= settings.confidence_threshold:
|
||||
return None
|
||||
return _build_sell_signal(state, holding, confidence=confidence, reason="anomaly")
|
||||
return _build_sell_signal(state, holding, confidence=confidence, reason="anomaly", settings=settings)
|
||||
|
||||
|
||||
def _build_sell_signal(state, holding: dict, confidence: float, reason: str) -> dict:
|
||||
def _build_sell_signal(state, holding: dict, confidence: float, reason: str, settings=None) -> dict:
|
||||
ticker = holding["ticker"]
|
||||
as_of_dt = datetime.now(KST)
|
||||
ttl = getattr(settings, "signal_ttl_seconds", 300) if settings else 300
|
||||
expires_at = (as_of_dt + timedelta(seconds=ttl)).isoformat()
|
||||
return {
|
||||
"ticker": ticker,
|
||||
"name": holding.get("name", ticker),
|
||||
@@ -182,7 +197,9 @@ def _build_sell_signal(state, holding: dict, confidence: float, reason: str) ->
|
||||
"avg_price": holding.get("avg_price"),
|
||||
"pnl_pct": holding.get("pnl_pct"),
|
||||
"context": _build_context(state, ticker, rank=None, sell_reason=reason),
|
||||
"as_of": datetime.now(KST).isoformat(),
|
||||
"as_of": as_of_dt.isoformat(),
|
||||
"cycle_id": state.signal_cycle_id,
|
||||
"expires_at": expires_at,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""PollState — process-wide singleton."""
|
||||
from collections import deque
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -15,8 +16,44 @@ class PollState:
|
||||
chronos_predictions: dict[str, dict] = field(default_factory=dict)
|
||||
minute_momentum: dict[str, str] = field(default_factory=dict)
|
||||
signals: dict[str, dict] = field(default_factory=dict)
|
||||
# F5 lifecycle
|
||||
signal_cycle_id: int = 0
|
||||
last_updated: dict[str, str] = field(default_factory=dict)
|
||||
fetch_errors: dict[str, int] = field(default_factory=dict)
|
||||
|
||||
def get_active_signals(self, now: datetime) -> list[dict]:
|
||||
"""expires_at > now 인 신호만 반환. expires_at 없거나 파싱 실패는 expired 취급."""
|
||||
active: list[dict] = []
|
||||
for sig in self.signals.values():
|
||||
expires_at = sig.get("expires_at")
|
||||
if not expires_at:
|
||||
continue
|
||||
try:
|
||||
exp_dt = datetime.fromisoformat(expires_at)
|
||||
except ValueError:
|
||||
continue
|
||||
if exp_dt > now:
|
||||
active.append(sig)
|
||||
return active
|
||||
|
||||
def purge_expired_signals(self, now: datetime) -> int:
|
||||
"""만료된 signal 제거. expires_at 없거나 파싱 실패도 제거. 제거 개수 반환."""
|
||||
to_drop = []
|
||||
for ticker, sig in self.signals.items():
|
||||
expires_at = sig.get("expires_at")
|
||||
if not expires_at:
|
||||
to_drop.append(ticker)
|
||||
continue
|
||||
try:
|
||||
exp_dt = datetime.fromisoformat(expires_at)
|
||||
except ValueError:
|
||||
to_drop.append(ticker)
|
||||
continue
|
||||
if exp_dt <= now:
|
||||
to_drop.append(ticker)
|
||||
for t in to_drop:
|
||||
del self.signals[t]
|
||||
return len(to_drop)
|
||||
|
||||
|
||||
state = PollState()
|
||||
|
||||
@@ -90,3 +90,54 @@ def test_return_computed_from_price_relative_to_last_close(mock_pipeline, mock_t
|
||||
daily = {"005930": _daily_ohlcv(list(range(41, 101)))} # last = 100
|
||||
result = predictor.predict_batch(daily)
|
||||
assert abs(result["005930"].median - 0.10) < 0.001
|
||||
|
||||
|
||||
# ----- F4: absolute spread 기반 confidence -----
|
||||
|
||||
def test_confidence_high_when_spread_near_zero(mock_pipeline, mock_torch_cpu):
|
||||
"""F4 — median≈0 + spread≈0 일 때 conf≈1 (현 relative 산식의 회귀 케이스).
|
||||
|
||||
한국 주가 100000원, q10=q50=q90=100000 → median=0, spread=0.
|
||||
Relative 산식 (spread/abs(median))은 0/0.001 보호선이라 spread=0이면 conf=1로
|
||||
동작하지만, median≈0 + 미세 spread(예 1원) 케이스에서 폭증 → conf=0.
|
||||
Absolute 산식은 그런 폭증 없음.
|
||||
"""
|
||||
quantiles = _mk_quantiles_tensor(100000.0, 100000.0, 100000.0)
|
||||
mock_pipeline.predict_quantiles.return_value = (quantiles, None)
|
||||
|
||||
from ai_trade.chronos_predictor import ChronosPredictor
|
||||
predictor = ChronosPredictor(model_name="mock-model")
|
||||
daily = {"005930": _daily_ohlcv([100000] * 60)}
|
||||
result = predictor.predict_batch(daily)
|
||||
assert result["005930"].conf > 0.95, (
|
||||
f"median≈0 + spread≈0인데 conf={result['005930'].conf} (F4 회귀)"
|
||||
)
|
||||
|
||||
|
||||
def test_confidence_half_at_spread_03(mock_pipeline, mock_torch_cpu):
|
||||
"""F4 — spread 0.30일 때 conf ≈ 0.5 (1 - 0.3/0.6)."""
|
||||
# q10=85000 → -0.15, q90=115000 → 0.15, q50=100000 → 0.0
|
||||
# spread = 0.30, conf = 1 - 0.30/0.60 = 0.50
|
||||
quantiles = _mk_quantiles_tensor(85000.0, 100000.0, 115000.0)
|
||||
mock_pipeline.predict_quantiles.return_value = (quantiles, None)
|
||||
|
||||
from ai_trade.chronos_predictor import ChronosPredictor
|
||||
predictor = ChronosPredictor(model_name="mock-model")
|
||||
daily = {"005930": _daily_ohlcv([100000] * 60)}
|
||||
result = predictor.predict_batch(daily)
|
||||
conf = result["005930"].conf
|
||||
assert 0.45 < conf < 0.55, f"spread=0.30에서 conf={conf} (expected ≈0.5)"
|
||||
|
||||
|
||||
def test_confidence_zero_at_threshold_spread(mock_pipeline, mock_torch_cpu):
|
||||
"""F4 — spread가 _SPREAD_THRESHOLD(0.6)이면 conf=0."""
|
||||
quantiles = _mk_quantiles_tensor(70000.0, 100000.0, 130000.0)
|
||||
mock_pipeline.predict_quantiles.return_value = (quantiles, None)
|
||||
|
||||
from ai_trade.chronos_predictor import ChronosPredictor
|
||||
predictor = ChronosPredictor(model_name="mock-model")
|
||||
daily = {"005930": _daily_ohlcv([100000] * 60)}
|
||||
result = predictor.predict_batch(daily)
|
||||
assert result["005930"].conf < 0.05, (
|
||||
f"spread=threshold에서 conf={result['005930'].conf} (expected ≈0)"
|
||||
)
|
||||
|
||||
22
ai_trade/tests/test_config_token_path.py
Normal file
22
ai_trade/tests/test_config_token_path.py
Normal file
@@ -0,0 +1,22 @@
|
||||
"""F1 — V1_TOKEN_PATH default가 legacy/signal_v1/ 경유인지 검증."""
|
||||
from pathlib import Path
|
||||
|
||||
from ai_trade.config import Settings
|
||||
|
||||
|
||||
def test_v1_token_default_path_uses_legacy_dir(monkeypatch):
|
||||
"""env에 V1_TOKEN_PATH 없으면 legacy/signal_v1/data/kis_token.json"""
|
||||
monkeypatch.delenv("V1_TOKEN_PATH", raising=False)
|
||||
settings = Settings()
|
||||
expected_suffix = Path("legacy") / "signal_v1" / "data" / "kis_token.json"
|
||||
assert str(settings.v1_token_path).endswith(str(expected_suffix)), (
|
||||
f"expected default to end with {expected_suffix}, got {settings.v1_token_path}"
|
||||
)
|
||||
|
||||
|
||||
def test_v1_token_env_override_wins(monkeypatch, tmp_path):
|
||||
"""env로 명시한 경로가 default를 덮어씀."""
|
||||
custom = tmp_path / "custom_token.json"
|
||||
monkeypatch.setenv("V1_TOKEN_PATH", str(custom))
|
||||
settings = Settings()
|
||||
assert settings.v1_token_path == custom
|
||||
@@ -1,5 +1,7 @@
|
||||
"""Tests for KISClient (REST)."""
|
||||
import asyncio
|
||||
import json
|
||||
import time as time_module
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
@@ -159,3 +161,30 @@ async def test_get_daily_ohlcv_returns_60_bars(kis_client_factory):
|
||||
assert "datetime" in bars[0]
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_throttle_serializes_concurrent_gather(kis_client_factory):
|
||||
"""F2 — 5개 동시 요청이 asyncio.gather로 들어와도 0.5초 간격으로 직렬화.
|
||||
|
||||
초당 2회 = 0.5초 간격. 5개 요청 시 최소 (5-1)*0.5 = 2.0초.
|
||||
Lock 없으면 race condition으로 거의 동시에 나가 0.5초대로 끝남.
|
||||
"""
|
||||
sample = {"output2": []}
|
||||
respx.get(
|
||||
"https://openapivts.koreainvestment.com:29443"
|
||||
"/uapi/domestic-stock/v1/quotations/inquire-time-itemchartprice"
|
||||
).mock(return_value=httpx.Response(200, json=sample))
|
||||
|
||||
client = kis_client_factory()
|
||||
try:
|
||||
start = time_module.monotonic()
|
||||
await asyncio.gather(*[client.get_minute_ohlcv(f"00593{i}") for i in range(5)])
|
||||
elapsed = time_module.monotonic() - start
|
||||
# 5 throttle = 최소 (5-1)*0.5 = 2.0s, tolerance 0.3s
|
||||
assert elapsed >= 1.7, (
|
||||
f"throttle race condition: 5 concurrent calls took only {elapsed:.2f}s, "
|
||||
f"expected >=1.7s (0.5s * 4 inter-call gaps)"
|
||||
)
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
@@ -122,6 +122,7 @@ def test_poll_loop_calls_generate_signals_after_cycle(monkeypatch):
|
||||
settings.asking_bid_ratio_threshold = 0.6
|
||||
settings.confidence_threshold = 0.7
|
||||
settings.min_momentum_for_buy = "strong_up"
|
||||
settings.signal_ttl_seconds = 300
|
||||
|
||||
generate_signals(state, dedup, settings)
|
||||
|
||||
@@ -129,3 +130,112 @@ def test_poll_loop_calls_generate_signals_after_cycle(monkeypatch):
|
||||
assert state.signals["005930"]["action"] == "sell"
|
||||
assert state.signals["005930"]["confidence_webai"] == 1.0
|
||||
dedup.record.assert_called_with("005930", "sell", confidence=1.0)
|
||||
|
||||
|
||||
async def test_post_close_fires_at_1601_when_not_yet_today(monkeypatch):
|
||||
"""F3 — 16:01에 깬 cycle도 오늘 post_close 안 돌렸으면 호출됨 (회귀 방지)."""
|
||||
from datetime import datetime as _dt
|
||||
from zoneinfo import ZoneInfo as _ZI
|
||||
import asyncio as _asyncio
|
||||
|
||||
from ai_trade import pull_worker
|
||||
|
||||
_kst = _ZI("Asia/Seoul")
|
||||
now_at_1601 = _dt(2026, 5, 18, 16, 1, tzinfo=_kst)
|
||||
|
||||
class FrozenDateTime:
|
||||
@staticmethod
|
||||
def now(tz=None):
|
||||
return now_at_1601
|
||||
|
||||
monkeypatch.setattr(pull_worker, "datetime", FrozenDateTime)
|
||||
monkeypatch.setattr(pull_worker, "_is_market_day", lambda n: True)
|
||||
monkeypatch.setattr(pull_worker, "_is_polling_window", lambda n: True)
|
||||
monkeypatch.setattr(pull_worker, "_next_interval", lambda n: 0.01)
|
||||
monkeypatch.setattr(pull_worker, "_run_polling_cycle", AsyncMock())
|
||||
monkeypatch.setattr(pull_worker, "update_minute_momentum_for_all", lambda s: None)
|
||||
post_close = AsyncMock()
|
||||
monkeypatch.setattr(pull_worker, "_run_post_close_cycle", post_close)
|
||||
|
||||
state = MagicMock()
|
||||
chronos = MagicMock()
|
||||
kis = MagicMock()
|
||||
shutdown = _asyncio.Event()
|
||||
|
||||
async def _stop_soon():
|
||||
await _asyncio.sleep(0.05)
|
||||
shutdown.set()
|
||||
|
||||
_asyncio.create_task(_stop_soon())
|
||||
await pull_worker.poll_loop(
|
||||
client=MagicMock(),
|
||||
state=state,
|
||||
shutdown=shutdown,
|
||||
kis_client=kis,
|
||||
chronos=chronos,
|
||||
dedup=None,
|
||||
settings=None,
|
||||
)
|
||||
|
||||
assert post_close.await_count >= 1, "post-close가 16:01에 호출되지 않음 (F3 회귀)"
|
||||
|
||||
|
||||
async def test_poll_loop_purges_expired_signals(monkeypatch):
|
||||
"""F5 — 매 cycle 끝에 expired signal이 제거됨."""
|
||||
from datetime import datetime as _dt
|
||||
from zoneinfo import ZoneInfo as _ZI
|
||||
import asyncio as _asyncio
|
||||
|
||||
from ai_trade import pull_worker
|
||||
from ai_trade.state import PollState
|
||||
|
||||
_kst = _ZI("Asia/Seoul")
|
||||
now = _dt(2026, 5, 18, 10, 0, tzinfo=_kst)
|
||||
|
||||
class FrozenDT:
|
||||
@staticmethod
|
||||
def now(tz=None):
|
||||
return now
|
||||
|
||||
state = PollState()
|
||||
state.signals = {
|
||||
"OLD": {
|
||||
"ticker": "OLD",
|
||||
"expires_at": _dt(2026, 5, 18, 9, 0, tzinfo=_kst).isoformat(),
|
||||
"cycle_id": 1,
|
||||
},
|
||||
"FRESH": {
|
||||
"ticker": "FRESH",
|
||||
"expires_at": _dt(2026, 5, 18, 10, 30, tzinfo=_kst).isoformat(),
|
||||
"cycle_id": 1,
|
||||
},
|
||||
}
|
||||
|
||||
monkeypatch.setattr(pull_worker, "datetime", FrozenDT)
|
||||
monkeypatch.setattr(pull_worker, "_is_market_day", lambda n: True)
|
||||
monkeypatch.setattr(pull_worker, "_is_polling_window", lambda n: True)
|
||||
monkeypatch.setattr(pull_worker, "_next_interval", lambda n: 0.01)
|
||||
monkeypatch.setattr(pull_worker, "_run_polling_cycle", AsyncMock())
|
||||
monkeypatch.setattr(pull_worker, "update_minute_momentum_for_all", lambda s: None)
|
||||
monkeypatch.setattr(pull_worker, "_is_post_close_trigger", lambda *a, **k: False)
|
||||
|
||||
shutdown = _asyncio.Event()
|
||||
|
||||
async def stop_soon():
|
||||
await _asyncio.sleep(0.05)
|
||||
shutdown.set()
|
||||
|
||||
_asyncio.create_task(stop_soon())
|
||||
|
||||
await pull_worker.poll_loop(
|
||||
client=MagicMock(),
|
||||
state=state,
|
||||
shutdown=shutdown,
|
||||
kis_client=MagicMock(),
|
||||
chronos=MagicMock(),
|
||||
dedup=None,
|
||||
settings=None,
|
||||
)
|
||||
|
||||
assert "OLD" not in state.signals
|
||||
assert "FRESH" in state.signals
|
||||
|
||||
@@ -79,3 +79,41 @@ def test_next_interval_dead_zone_skip():
|
||||
interval = _next_interval(now)
|
||||
# 02:00 → 04:30 = 2.5h = 9000s
|
||||
assert 9000 - 60 < interval < 9000 + 60
|
||||
|
||||
|
||||
# ----- F3 post-close 상태기반 트리거 -----
|
||||
|
||||
from datetime import date as _date # noqa: E402
|
||||
from ai_trade.scheduler import _is_post_close_trigger # noqa: E402
|
||||
|
||||
|
||||
def test_post_close_trigger_fires_at_1601_if_not_yet_today():
|
||||
"""F3 — 16:01에 깬 cycle도 오늘 아직 안 돌렸으면 trigger."""
|
||||
now = _kst(2026, 5, 18, 16, 1)
|
||||
assert _is_post_close_trigger(now, last_post_close_date=None) is True
|
||||
|
||||
|
||||
def test_post_close_trigger_skips_if_already_today():
|
||||
"""F3 — 이미 오늘 돌렸으면 trigger 안 함."""
|
||||
now = _kst(2026, 5, 18, 16, 5)
|
||||
today = _date(2026, 5, 18)
|
||||
assert _is_post_close_trigger(now, last_post_close_date=today) is False
|
||||
|
||||
|
||||
def test_post_close_trigger_skips_before_1600():
|
||||
"""F3 — 16:00 전에는 trigger 안 함."""
|
||||
now = _kst(2026, 5, 18, 15, 59)
|
||||
assert _is_post_close_trigger(now, last_post_close_date=None) is False
|
||||
|
||||
|
||||
def test_post_close_trigger_fires_next_day_after_reset():
|
||||
"""F3 — 다음 영업일이 되면 다시 trigger."""
|
||||
now = _kst(2026, 5, 19, 16, 0)
|
||||
yesterday = _date(2026, 5, 18)
|
||||
assert _is_post_close_trigger(now, last_post_close_date=yesterday) is True
|
||||
|
||||
|
||||
def test_post_close_trigger_skips_on_holiday():
|
||||
"""F3 — 휴장일에는 trigger 안 함 (2026-05-05 어린이날)."""
|
||||
now = _kst(2026, 5, 5, 16, 30)
|
||||
assert _is_post_close_trigger(now, last_post_close_date=None) is False
|
||||
|
||||
@@ -16,6 +16,7 @@ def _settings(**overrides):
|
||||
asking_bid_ratio_threshold=0.6,
|
||||
confidence_threshold=0.7,
|
||||
min_momentum_for_buy="strong_up",
|
||||
signal_ttl_seconds=300,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
m = MagicMock()
|
||||
@@ -170,3 +171,48 @@ def test_sell_signal_triggers_on_anomaly_path(dedup_mock):
|
||||
assert sig["action"] == "sell"
|
||||
assert sig["context"]["sell_reason"] == "anomaly"
|
||||
assert sig["confidence_webai"] > 0.7
|
||||
|
||||
|
||||
# ----- F5: cycle_id + expires_at 부착 -----
|
||||
|
||||
def test_emit_attaches_cycle_id_and_expires_at(dedup_mock):
|
||||
"""F5 — emit signal에 cycle_id (state.signal_cycle_id) + expires_at 부착."""
|
||||
from datetime import datetime, timedelta
|
||||
from zoneinfo import ZoneInfo
|
||||
_kst = ZoneInfo("Asia/Seoul")
|
||||
|
||||
state = _make_state_with_buy_candidate()
|
||||
before = datetime.now(_kst)
|
||||
generate_signals(state, dedup_mock, _settings(signal_ttl_seconds=300))
|
||||
after = datetime.now(_kst)
|
||||
|
||||
sig = state.signals["005930"]
|
||||
assert sig["cycle_id"] == 1
|
||||
assert "expires_at" in sig
|
||||
exp_dt = datetime.fromisoformat(sig["expires_at"])
|
||||
assert before + timedelta(seconds=295) < exp_dt < after + timedelta(seconds=305)
|
||||
|
||||
|
||||
def test_cycle_id_increments_each_call(dedup_mock):
|
||||
"""F5 — generate_signals 호출마다 cycle_id += 1 (emit 여부 무관)."""
|
||||
state = _make_state_with_buy_candidate()
|
||||
generate_signals(state, dedup_mock, _settings())
|
||||
assert state.signal_cycle_id == 1
|
||||
# 2번째 호출 — dedup이 막아도 cycle_id는 증가
|
||||
dedup_mock.is_recent.return_value = True
|
||||
generate_signals(state, dedup_mock, _settings())
|
||||
assert state.signal_cycle_id == 2
|
||||
|
||||
|
||||
def test_sell_signal_also_carries_cycle_id_and_expires_at(dedup_mock):
|
||||
"""F5 — sell signal도 동일하게 부착."""
|
||||
from datetime import datetime
|
||||
state = _make_state_with_holding(pnl_pct=-0.08, current_price=68000)
|
||||
generate_signals(state, dedup_mock, _settings(signal_ttl_seconds=120))
|
||||
|
||||
assert "005930" in state.signals
|
||||
sig = state.signals["005930"]
|
||||
assert sig["action"] == "sell"
|
||||
assert sig["cycle_id"] == 1
|
||||
# parse expires_at as ISO — must succeed
|
||||
datetime.fromisoformat(sig["expires_at"])
|
||||
|
||||
66
ai_trade/tests/test_state_signals_lifecycle.py
Normal file
66
ai_trade/tests/test_state_signals_lifecycle.py
Normal file
@@ -0,0 +1,66 @@
|
||||
"""F5 — state.signals lifecycle (expires_at + cycle_id)."""
|
||||
from datetime import datetime, timedelta
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
from ai_trade.state import PollState
|
||||
|
||||
KST = ZoneInfo("Asia/Seoul")
|
||||
|
||||
|
||||
def test_initial_signal_cycle_id_is_zero():
|
||||
state = PollState()
|
||||
assert state.signal_cycle_id == 0
|
||||
|
||||
|
||||
def test_get_active_signals_excludes_expired():
|
||||
state = PollState()
|
||||
now = datetime(2026, 5, 25, 10, 0, tzinfo=KST)
|
||||
future = (now + timedelta(seconds=300)).isoformat()
|
||||
past = (now - timedelta(seconds=60)).isoformat()
|
||||
state.signals = {
|
||||
"A": {"ticker": "A", "expires_at": future, "cycle_id": 1, "action": "buy"},
|
||||
"B": {"ticker": "B", "expires_at": past, "cycle_id": 1, "action": "buy"},
|
||||
}
|
||||
active = state.get_active_signals(now)
|
||||
tickers = [s["ticker"] for s in active]
|
||||
assert "A" in tickers
|
||||
assert "B" not in tickers
|
||||
|
||||
|
||||
def test_get_active_signals_treats_missing_expires_as_expired():
|
||||
"""expires_at 없는 legacy 신호는 expired로 간주."""
|
||||
state = PollState()
|
||||
now = datetime(2026, 5, 25, 10, 0, tzinfo=KST)
|
||||
state.signals = {"C": {"ticker": "C", "action": "buy"}}
|
||||
assert state.get_active_signals(now) == []
|
||||
|
||||
|
||||
def test_purge_expired_signals_removes_expired():
|
||||
state = PollState()
|
||||
now = datetime(2026, 5, 25, 10, 0, tzinfo=KST)
|
||||
future = (now + timedelta(seconds=300)).isoformat()
|
||||
past = (now - timedelta(seconds=60)).isoformat()
|
||||
state.signals = {
|
||||
"A": {"ticker": "A", "expires_at": future, "cycle_id": 1},
|
||||
"B": {"ticker": "B", "expires_at": past, "cycle_id": 1},
|
||||
}
|
||||
removed = state.purge_expired_signals(now)
|
||||
assert "A" in state.signals
|
||||
assert "B" not in state.signals
|
||||
assert removed == 1
|
||||
|
||||
|
||||
# ----- SIGNAL_TTL_SECONDS env -----
|
||||
|
||||
def test_signal_ttl_seconds_default(monkeypatch):
|
||||
monkeypatch.delenv("SIGNAL_TTL_SECONDS", raising=False)
|
||||
from ai_trade.config import Settings
|
||||
s = Settings()
|
||||
assert s.signal_ttl_seconds == 300
|
||||
|
||||
|
||||
def test_signal_ttl_seconds_env_override(monkeypatch):
|
||||
monkeypatch.setenv("SIGNAL_TTL_SECONDS", "60")
|
||||
from ai_trade.config import Settings
|
||||
s = Settings()
|
||||
assert s.signal_ttl_seconds == 60
|
||||
0
services/_shared/__init__.py
Normal file
0
services/_shared/__init__.py
Normal file
2
services/_shared/pytest.ini
Normal file
2
services/_shared/pytest.ini
Normal file
@@ -0,0 +1,2 @@
|
||||
[pytest]
|
||||
asyncio_mode = auto
|
||||
135
services/_shared/reliable_queue.py
Normal file
135
services/_shared/reliable_queue.py
Normal file
@@ -0,0 +1,135 @@
|
||||
"""F6 — Reliable Redis queue with processing list + recovery + retry.
|
||||
|
||||
Pattern:
|
||||
- BLMOVE main → processing (atomic dequeue)
|
||||
- ack: LREM processing (1 occurrence)
|
||||
- fail: LREM processing + (re-enqueue with attempts++ OR move to dead-letter)
|
||||
- recover: startup-time orphan recovery (worker's processing list → main queue)
|
||||
|
||||
Producer side stays unchanged: LPUSH queue:<x> <json payload>.
|
||||
Worker side: dequeue() → process → ack(raw) on success or fail(raw, payload) on error.
|
||||
Startup: await queue.recover() to re-enqueue orphans.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def default_worker_id(queue_key: str) -> str:
|
||||
"""env WORKER_ID > hostname-pid."""
|
||||
explicit = os.getenv("WORKER_ID")
|
||||
if explicit:
|
||||
return explicit
|
||||
return f"{queue_key}-{socket.gethostname()}-{os.getpid()}"
|
||||
|
||||
|
||||
class ReliableQueue:
|
||||
"""BLMOVE-backed atomic dequeue + processing list + retry/dead-letter."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis,
|
||||
queue_key: str,
|
||||
worker_id: Optional[str] = None,
|
||||
max_attempts: int = 3,
|
||||
):
|
||||
self._redis = redis
|
||||
self._queue_key = queue_key
|
||||
self._worker_id = worker_id or default_worker_id(queue_key)
|
||||
self._processing_key = f"processing:{queue_key}:{self._worker_id}"
|
||||
self._dead_letter_key = f"dead_letter:{queue_key}"
|
||||
self._max_attempts = max_attempts
|
||||
|
||||
@property
|
||||
def worker_id(self) -> str:
|
||||
return self._worker_id
|
||||
|
||||
@property
|
||||
def processing_key(self) -> str:
|
||||
return self._processing_key
|
||||
|
||||
async def dequeue(self, timeout: int = 5) -> Optional[tuple[dict, bytes]]:
|
||||
"""Atomically move 1 item from main queue tail to processing head.
|
||||
|
||||
Returns (parsed_dict, raw_bytes) or None on timeout/parse-error.
|
||||
Caller MUST call ack(raw) on success or fail(raw, payload) on error.
|
||||
"""
|
||||
raw = await self._redis.blmove(
|
||||
self._queue_key, self._processing_key,
|
||||
timeout, "RIGHT", "LEFT",
|
||||
)
|
||||
if raw is None:
|
||||
return None
|
||||
try:
|
||||
payload = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
logger.error(
|
||||
"invalid payload on dequeue, moving to dead-letter: %r", raw[:200]
|
||||
)
|
||||
await self._redis.lrem(self._processing_key, 1, raw)
|
||||
await self._redis.lpush(self._dead_letter_key, raw)
|
||||
return None
|
||||
return payload, raw
|
||||
|
||||
async def ack(self, raw: bytes) -> None:
|
||||
"""Successful processing — remove from processing list."""
|
||||
removed = await self._redis.lrem(self._processing_key, 1, raw)
|
||||
if removed == 0:
|
||||
logger.warning("ack on missing payload (already removed?): %r", raw[:100])
|
||||
|
||||
async def fail(self, raw: bytes, payload: dict) -> None:
|
||||
"""Failed processing — remove from processing list and re-enqueue or dead-letter."""
|
||||
await self._redis.lrem(self._processing_key, 1, raw)
|
||||
attempts = int(payload.get("attempts", 0)) + 1
|
||||
if attempts >= self._max_attempts:
|
||||
payload["attempts"] = attempts
|
||||
await self._redis.lpush(self._dead_letter_key, json.dumps(payload).encode())
|
||||
logger.error(
|
||||
"task moved to dead-letter after %d attempts: task_id=%s",
|
||||
attempts, payload.get("task_id"),
|
||||
)
|
||||
return
|
||||
payload["attempts"] = attempts
|
||||
await self._redis.lpush(self._queue_key, json.dumps(payload).encode())
|
||||
logger.info(
|
||||
"task re-enqueued (attempt %d/%d): task_id=%s",
|
||||
attempts, self._max_attempts, payload.get("task_id"),
|
||||
)
|
||||
|
||||
async def recover(self) -> int:
|
||||
"""Startup: move all orphans from this worker's processing list back to main queue.
|
||||
|
||||
Increments attempts counter (orphan == implicit failure). Returns count.
|
||||
"""
|
||||
count = 0
|
||||
while True:
|
||||
raw = await self._redis.lpop(self._processing_key)
|
||||
if raw is None:
|
||||
break
|
||||
try:
|
||||
payload = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
await self._redis.lpush(self._dead_letter_key, raw)
|
||||
count += 1
|
||||
continue
|
||||
payload["attempts"] = int(payload.get("attempts", 0)) + 1
|
||||
if payload["attempts"] >= self._max_attempts:
|
||||
await self._redis.lpush(
|
||||
self._dead_letter_key, json.dumps(payload).encode()
|
||||
)
|
||||
else:
|
||||
await self._redis.lpush(
|
||||
self._queue_key, json.dumps(payload).encode()
|
||||
)
|
||||
count += 1
|
||||
if count:
|
||||
logger.info(
|
||||
"recovered %d orphaned items for worker %s", count, self._worker_id
|
||||
)
|
||||
return count
|
||||
1
services/_shared/requirements.txt
Normal file
1
services/_shared/requirements.txt
Normal file
@@ -0,0 +1 @@
|
||||
redis>=5.0.0
|
||||
0
services/_shared/tests/__init__.py
Normal file
0
services/_shared/tests/__init__.py
Normal file
84
services/_shared/tests/test_reliable_queue.py
Normal file
84
services/_shared/tests/test_reliable_queue.py
Normal file
@@ -0,0 +1,84 @@
|
||||
"""F6 — ReliableQueue: atomic dequeue + recovery + retry."""
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import fakeredis.aioredis
|
||||
import pytest
|
||||
|
||||
# Make `_shared` importable when tests run from services/_shared
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
|
||||
|
||||
from _shared.reliable_queue import ReliableQueue
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def redis():
|
||||
r = fakeredis.aioredis.FakeRedis(decode_responses=False)
|
||||
yield r
|
||||
await r.flushall()
|
||||
await r.aclose()
|
||||
|
||||
|
||||
async def test_dequeue_atomically_moves_to_processing(redis):
|
||||
"""BLMOVE: queue → processing 원자적 이동."""
|
||||
q = ReliableQueue(redis, queue_key="queue:test", worker_id="w1")
|
||||
await redis.lpush("queue:test", json.dumps({"task_id": "t1"}).encode())
|
||||
result = await q.dequeue(timeout=1)
|
||||
assert result is not None
|
||||
payload, raw = result
|
||||
assert payload["task_id"] == "t1"
|
||||
assert await redis.llen("queue:test") == 0
|
||||
assert await redis.llen("processing:queue:test:w1") == 1
|
||||
|
||||
|
||||
async def test_dequeue_returns_none_on_timeout(redis):
|
||||
q = ReliableQueue(redis, queue_key="queue:test", worker_id="w1")
|
||||
result = await q.dequeue(timeout=1)
|
||||
assert result is None
|
||||
|
||||
|
||||
async def test_ack_removes_from_processing(redis):
|
||||
q = ReliableQueue(redis, queue_key="queue:test", worker_id="w1")
|
||||
await redis.lpush("queue:test", json.dumps({"task_id": "t1"}).encode())
|
||||
_, raw = await q.dequeue(timeout=1)
|
||||
await q.ack(raw)
|
||||
assert await redis.llen("processing:queue:test:w1") == 0
|
||||
|
||||
|
||||
async def test_recover_returns_orphaned_to_main_queue(redis):
|
||||
"""startup recovery: 잔존 processing list 항목을 main queue로 되돌림."""
|
||||
orphan = json.dumps({"task_id": "t1", "attempts": 0}).encode()
|
||||
await redis.lpush("processing:queue:test:w1", orphan)
|
||||
q = ReliableQueue(redis, queue_key="queue:test", worker_id="w1")
|
||||
recovered = await q.recover()
|
||||
assert recovered == 1
|
||||
assert await redis.llen("processing:queue:test:w1") == 0
|
||||
payload, _ = await q.dequeue(timeout=1)
|
||||
assert payload["task_id"] == "t1"
|
||||
assert payload["attempts"] == 1 # incremented on recover
|
||||
|
||||
|
||||
async def test_fail_below_max_attempts_returns_to_main_queue(redis):
|
||||
q = ReliableQueue(redis, queue_key="queue:test", worker_id="w1", max_attempts=3)
|
||||
await redis.lpush("queue:test", json.dumps({"task_id": "t1", "attempts": 0}).encode())
|
||||
payload, raw = await q.dequeue(timeout=1)
|
||||
await q.fail(raw, payload)
|
||||
assert await redis.llen("processing:queue:test:w1") == 0
|
||||
assert await redis.llen("queue:test") == 1
|
||||
requeued_raw = await redis.lindex("queue:test", 0)
|
||||
requeued = json.loads(requeued_raw)
|
||||
assert requeued["attempts"] == 1
|
||||
|
||||
|
||||
async def test_fail_at_max_attempts_moves_to_dead_letter(redis):
|
||||
q = ReliableQueue(redis, queue_key="queue:test", worker_id="w1", max_attempts=3)
|
||||
await redis.lpush(
|
||||
"queue:test", json.dumps({"task_id": "t1", "attempts": 2}).encode()
|
||||
)
|
||||
payload, raw = await q.dequeue(timeout=1)
|
||||
await q.fail(raw, payload)
|
||||
# attempts 2 → 3 (== max) → dead-letter
|
||||
assert await redis.llen("queue:test") == 0
|
||||
assert await redis.llen("processing:queue:test:w1") == 0
|
||||
assert await redis.llen("dead_letter:queue:test") == 1
|
||||
@@ -3,7 +3,8 @@ name: web-ai-services
|
||||
services:
|
||||
insta-render:
|
||||
build:
|
||||
context: ./insta-render
|
||||
context: .
|
||||
dockerfile: insta-render/Dockerfile
|
||||
container_name: insta-render
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
@@ -13,7 +14,7 @@ services:
|
||||
- REDIS_URL=${REDIS_URL:-redis://192.168.45.54:6379}
|
||||
- NAS_BASE_URL=${NAS_BASE_URL:-http://192.168.45.54:18700}
|
||||
- INTERNAL_API_KEY=${INTERNAL_API_KEY:-}
|
||||
- INSTA_MEDIA_ROOT=${INSTA_MEDIA_ROOT:-/mnt/nas/webpage/data/insta}
|
||||
- INSTA_MEDIA_ROOT=${INSTA_MEDIA_ROOT:-/mnt/nas/webpage/data/insta/insta_cards}
|
||||
- INSTA_MEDIA_URL_PREFIX=${INSTA_MEDIA_URL_PREFIX:-/media/insta}
|
||||
- CARD_TEMPLATE_DIR=/app/templates
|
||||
volumes:
|
||||
@@ -23,3 +24,105 @@ services:
|
||||
interval: 60s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
|
||||
music-render:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: music-render/Dockerfile
|
||||
container_name: music-render
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "18711:8000"
|
||||
environment:
|
||||
- TZ=Asia/Seoul
|
||||
- REDIS_URL=${REDIS_URL:-redis://192.168.45.54:6379}
|
||||
- NAS_BASE_URL=${NAS_BASE_URL:-http://192.168.45.54:18600}
|
||||
- INTERNAL_API_KEY=${INTERNAL_API_KEY:-}
|
||||
- SUNO_API_KEY=${SUNO_API_KEY:-}
|
||||
- MUSIC_AI_SERVER_URL=${MUSIC_AI_SERVER_URL:-http://host.docker.internal:8765}
|
||||
- MUSIC_MEDIA_ROOT=${MUSIC_MEDIA_ROOT:-/mnt/nas/webpage/data/music}
|
||||
- MUSIC_MEDIA_URL_PREFIX=${MUSIC_MEDIA_URL_PREFIX:-/media/music}
|
||||
extra_hosts:
|
||||
- "host.docker.internal:host-gateway"
|
||||
volumes:
|
||||
- /mnt/nas/webpage/data/music:/mnt/nas/webpage/data/music
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"]
|
||||
interval: 60s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
|
||||
video-render:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: video-render/Dockerfile
|
||||
container_name: video-render
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "18712:8000"
|
||||
environment:
|
||||
- TZ=Asia/Seoul
|
||||
- REDIS_URL=${REDIS_URL:-redis://192.168.45.54:6379}
|
||||
- NAS_BASE_URL=${NAS_BASE_URL:-http://192.168.45.54:18801}
|
||||
- INTERNAL_API_KEY=${INTERNAL_API_KEY:-}
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY:-}
|
||||
- GEMINI_API_KEY=${GEMINI_API_KEY:-}
|
||||
- KLING_ACCESS_KEY=${KLING_ACCESS_KEY:-}
|
||||
- KLING_SECRET_KEY=${KLING_SECRET_KEY:-}
|
||||
- SEEDANCE_API_KEY=${SEEDANCE_API_KEY:-}
|
||||
- VIDEO_MEDIA_ROOT=${VIDEO_MEDIA_ROOT:-/mnt/nas/webpage/data/video}
|
||||
- VIDEO_MEDIA_URL_PREFIX=${VIDEO_MEDIA_URL_PREFIX:-/media/video}
|
||||
volumes:
|
||||
- /mnt/nas/webpage/data/video:/mnt/nas/webpage/data/video
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"]
|
||||
interval: 60s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
|
||||
task-watcher:
|
||||
build:
|
||||
context: ./task-watcher
|
||||
container_name: task-watcher
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "18713:8000"
|
||||
environment:
|
||||
- TZ=Asia/Seoul
|
||||
- REDIS_URL=${REDIS_URL:-redis://192.168.45.54:6379}
|
||||
- STOCK_BASE_URL=${STOCK_BASE_URL:-http://192.168.45.54:18500}
|
||||
- TRADING_START=${TRADING_START:-07:00}
|
||||
- TRADING_END=${TRADING_END:-16:30}
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"]
|
||||
interval: 60s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
|
||||
image-render:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: image-render/Dockerfile
|
||||
container_name: image-render
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "18714:8000"
|
||||
environment:
|
||||
- TZ=Asia/Seoul
|
||||
- REDIS_URL=${REDIS_URL:-redis://192.168.45.54:6379}
|
||||
- NAS_BASE_URL=${NAS_BASE_URL:-http://192.168.45.54:18802}
|
||||
- INTERNAL_API_KEY=${INTERNAL_API_KEY:-}
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY:-}
|
||||
- GEMINI_API_KEY=${GEMINI_API_KEY:-}
|
||||
- COMFYUI_URL=${COMFYUI_URL:-http://host.docker.internal:8188}
|
||||
- FLUX_BLOCK_TRADING_HOURS=${FLUX_BLOCK_TRADING_HOURS:-1}
|
||||
- IMAGE_MEDIA_ROOT=${IMAGE_MEDIA_ROOT:-/mnt/nas/webpage/data/image}
|
||||
extra_hosts:
|
||||
- "host.docker.internal:host-gateway"
|
||||
volumes:
|
||||
- /mnt/nas/webpage/data/image:/mnt/nas/webpage/data/image
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"]
|
||||
interval: 60s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
|
||||
19
services/image-render/Dockerfile
Normal file
19
services/image-render/Dockerfile
Normal file
@@ -0,0 +1,19 @@
|
||||
FROM python:3.12-slim-bookworm
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY image-render/requirements.txt /app/
|
||||
RUN pip install --no-cache-dir --timeout 600 --retries 5 -r requirements.txt
|
||||
|
||||
# F6: 공통 ReliableQueue 모듈 (services/_shared)
|
||||
COPY _shared /app/_shared
|
||||
COPY image-render/. /app/
|
||||
ENV PYTHONPATH=/app
|
||||
|
||||
EXPOSE 8000
|
||||
CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "1"]
|
||||
5
services/image-render/conftest.py
Normal file
5
services/image-render/conftest.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""Make services/ root importable so `from _shared.reliable_queue import ...` works during tests."""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
18
services/image-render/env.example
Normal file
18
services/image-render/env.example
Normal file
@@ -0,0 +1,18 @@
|
||||
# Redis (NAS)
|
||||
REDIS_URL=redis://192.168.45.54:6379
|
||||
|
||||
# NAS image-lab webhook
|
||||
NAS_BASE_URL=http://192.168.45.54:18802
|
||||
INTERNAL_API_KEY=replace-me
|
||||
|
||||
# API provider keys (worker reports failed if missing)
|
||||
OPENAI_API_KEY=
|
||||
GEMINI_API_KEY=
|
||||
# Seedance key not used by image-render
|
||||
|
||||
# FLUX local
|
||||
COMFYUI_URL=http://host.docker.internal:8188
|
||||
FLUX_BLOCK_TRADING_HOURS=1
|
||||
|
||||
# NAS SMB mount target (image-render writes to this, NAS reads via /media/image/)
|
||||
IMAGE_MEDIA_ROOT=/mnt/nas/webpage/data/image
|
||||
36
services/image-render/main.py
Normal file
36
services/image-render/main.py
Normal file
@@ -0,0 +1,36 @@
|
||||
"""image-render FastAPI entry — health + lifespan (worker loop spawn)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
import worker
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
worker_task = asyncio.create_task(worker.worker_loop())
|
||||
logger.info("image-render lifespan 시작")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
worker_task.cancel()
|
||||
try:
|
||||
await worker_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
logger.info("image-render lifespan 종료")
|
||||
|
||||
|
||||
app = FastAPI(lifespan=lifespan)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"ok": True, "service": "image-render"}
|
||||
54
services/image-render/nas_client.py
Normal file
54
services/image-render/nas_client.py
Normal file
@@ -0,0 +1,54 @@
|
||||
"""NAS webhook 어댑터 — Windows worker → NAS image-lab HTTP 위임.
|
||||
|
||||
video-render nas_client 복제 (call-time os.getenv으로 테스트 격리).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_TIMEOUT = 10.0
|
||||
|
||||
|
||||
def _post(payload: Dict[str, Any]) -> None:
|
||||
nas_base_url = os.getenv("NAS_BASE_URL", "http://192.168.45.54:18802")
|
||||
internal_api_key = os.getenv("INTERNAL_API_KEY", "")
|
||||
url = f"{nas_base_url}/api/internal/image/update"
|
||||
try:
|
||||
r = httpx.post(
|
||||
url,
|
||||
headers={"X-Internal-Key": internal_api_key},
|
||||
json=payload,
|
||||
timeout=_TIMEOUT,
|
||||
)
|
||||
if r.status_code != 200:
|
||||
logger.error("webhook %s returned %d: %s",
|
||||
payload.get("task_id"), r.status_code, r.text[:200])
|
||||
except Exception:
|
||||
logger.exception("webhook %s 호출 실패", payload.get("task_id"))
|
||||
|
||||
|
||||
def webhook_update_task(
|
||||
task_id: str,
|
||||
status: str,
|
||||
progress: int,
|
||||
message: str = "",
|
||||
image_url: Optional[str] = None,
|
||||
error: Optional[str] = None,
|
||||
) -> None:
|
||||
payload: Dict[str, Any] = {
|
||||
"task_id": task_id,
|
||||
"status": status,
|
||||
"progress": progress,
|
||||
"message": message,
|
||||
}
|
||||
if image_url is not None:
|
||||
payload["image_url"] = image_url
|
||||
if error is not None:
|
||||
payload["error"] = error
|
||||
_post(payload)
|
||||
0
services/image-render/providers/__init__.py
Normal file
0
services/image-render/providers/__init__.py
Normal file
18
services/image-render/providers/_media.py
Normal file
18
services/image-render/providers/_media.py
Normal file
@@ -0,0 +1,18 @@
|
||||
"""b64 이미지 → NAS SMB 경로 저장 → /media/image URL 반환."""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import os
|
||||
import uuid
|
||||
|
||||
IMAGE_MEDIA_ROOT = os.getenv("IMAGE_MEDIA_ROOT", "/mnt/nas/webpage/data/image")
|
||||
IMAGE_MEDIA_URL_PREFIX = os.getenv("IMAGE_MEDIA_URL_PREFIX", "/media/image")
|
||||
|
||||
|
||||
def save_b64_png(task_id: str, b64_data: str) -> str:
|
||||
os.makedirs(IMAGE_MEDIA_ROOT, exist_ok=True)
|
||||
fname = f"{task_id}-{uuid.uuid4().hex[:8]}.png"
|
||||
path = os.path.join(IMAGE_MEDIA_ROOT, fname)
|
||||
with open(path, "wb") as f:
|
||||
f.write(base64.b64decode(b64_data))
|
||||
return f"{IMAGE_MEDIA_URL_PREFIX}/{fname}"
|
||||
79
services/image-render/providers/flux.py
Normal file
79
services/image-render/providers/flux.py
Normal file
@@ -0,0 +1,79 @@
|
||||
"""FLUX 로컬 — ComfyUI HTTP API.
|
||||
|
||||
POST {COMFYUI_URL}/prompt (workflow JSON) → prompt_id
|
||||
GET {COMFYUI_URL}/history/{prompt_id} → outputs → image filename
|
||||
GET {COMFYUI_URL}/view?filename=... → PNG bytes → b64
|
||||
|
||||
워크플로우 JSON은 `flux_workflow.json` (ComfyUI UI에서 "Save (API Format)"로 export, CLIPTextEncode 노드 text를 "%PROMPT%"로 수동 치환). 박재오 산출물.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64, json, logging, os, time
|
||||
from datetime import datetime, timezone, timedelta
|
||||
import requests
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
from providers._media import save_b64_png
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
COMFYUI_URL = os.getenv("COMFYUI_URL", "http://127.0.0.1:8188")
|
||||
WORKFLOW_PATH = os.path.join(os.path.dirname(__file__), "flux_workflow.json")
|
||||
POLL_INTERVAL = 2
|
||||
POLL_MAX = 120
|
||||
|
||||
|
||||
def _is_trading_hours() -> bool:
|
||||
kst = timezone(timedelta(hours=9))
|
||||
now = datetime.now(kst)
|
||||
if now.weekday() >= 5:
|
||||
return False
|
||||
return (now.hour, now.minute) >= (9, 0) and (now.hour, now.minute) <= (15, 30)
|
||||
|
||||
|
||||
def _load_workflow(prompt: str, size: str) -> dict:
|
||||
with open(WORKFLOW_PATH, encoding="utf-8") as f:
|
||||
wf = json.load(f)
|
||||
# CLIPTextEncode 노드의 text를 prompt로 치환 (workflow에 "%PROMPT%" placeholder 사용)
|
||||
raw = json.dumps(wf).replace("%PROMPT%", prompt.replace('"', "'"))
|
||||
return json.loads(raw)
|
||||
|
||||
|
||||
def _submit_prompt(workflow: dict) -> str:
|
||||
r = requests.post(f"{COMFYUI_URL}/prompt", json={"prompt": workflow}, timeout=30)
|
||||
r.raise_for_status()
|
||||
return r.json()["prompt_id"]
|
||||
|
||||
|
||||
def _poll_image_b64(prompt_id: str):
|
||||
for _ in range(POLL_MAX):
|
||||
h = requests.get(f"{COMFYUI_URL}/history/{prompt_id}", timeout=10)
|
||||
data = h.json().get(prompt_id)
|
||||
if data and data.get("outputs"):
|
||||
for node_out in data["outputs"].values():
|
||||
for img in node_out.get("images", []):
|
||||
view = requests.get(f"{COMFYUI_URL}/view",
|
||||
params={"filename": img["filename"], "subfolder": img.get("subfolder", ""), "type": img.get("type", "output")},
|
||||
timeout=30)
|
||||
view.raise_for_status()
|
||||
return base64.b64encode(view.content).decode()
|
||||
time.sleep(POLL_INTERVAL)
|
||||
return None
|
||||
|
||||
|
||||
def run_flux_generation(task_id: str, params: dict) -> None:
|
||||
try:
|
||||
if os.getenv("FLUX_BLOCK_TRADING_HOURS") == "1" and _is_trading_hours():
|
||||
webhook_update_task(task_id, "failed", 0, "", error="장중 GPU 보호 — FLUX 거부 (API provider 사용 권장)")
|
||||
return
|
||||
webhook_update_task(task_id, "processing", 10, "FLUX (ComfyUI) 생성 중...")
|
||||
wf = _load_workflow(params["prompt"], params.get("size") or "1024x1024")
|
||||
pid = _submit_prompt(wf)
|
||||
b64 = _poll_image_b64(pid)
|
||||
if not b64:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="ComfyUI 타임아웃 또는 출력 없음")
|
||||
return
|
||||
url = save_b64_png(task_id, b64)
|
||||
webhook_update_task(task_id, "succeeded", 100, "완료", image_url=url)
|
||||
except Exception as e:
|
||||
logger.exception("flux task=%s 실패", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
83
services/image-render/providers/flux_workflow.json
Normal file
83
services/image-render/providers/flux_workflow.json
Normal file
@@ -0,0 +1,83 @@
|
||||
{
|
||||
"5": {
|
||||
"inputs": {
|
||||
"width": 1024,
|
||||
"height": 1024,
|
||||
"batch_size": 1
|
||||
},
|
||||
"class_type": "EmptyLatentImage",
|
||||
"_meta": {"title": "Empty Latent Image"}
|
||||
},
|
||||
"6": {
|
||||
"inputs": {
|
||||
"text": "%PROMPT%",
|
||||
"clip": ["11", 0]
|
||||
},
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {"title": "Positive Prompt"}
|
||||
},
|
||||
"8": {
|
||||
"inputs": {
|
||||
"samples": ["13", 0],
|
||||
"vae": ["10", 0]
|
||||
},
|
||||
"class_type": "VAEDecode",
|
||||
"_meta": {"title": "VAE Decode"}
|
||||
},
|
||||
"9": {
|
||||
"inputs": {
|
||||
"filename_prefix": "flux",
|
||||
"images": ["8", 0]
|
||||
},
|
||||
"class_type": "SaveImage",
|
||||
"_meta": {"title": "Save Image"}
|
||||
},
|
||||
"10": {
|
||||
"inputs": {
|
||||
"vae_name": "ae.safetensors"
|
||||
},
|
||||
"class_type": "VAELoader",
|
||||
"_meta": {"title": "Load VAE"}
|
||||
},
|
||||
"11": {
|
||||
"inputs": {
|
||||
"clip_name1": "clip_l.safetensors",
|
||||
"clip_name2": "t5xxl_fp8_e4m3fn.safetensors",
|
||||
"type": "flux"
|
||||
},
|
||||
"class_type": "DualCLIPLoader",
|
||||
"_meta": {"title": "Dual CLIP Loader"}
|
||||
},
|
||||
"12": {
|
||||
"inputs": {
|
||||
"unet_name": "flux1-schnell-fp8.safetensors",
|
||||
"weight_dtype": "default"
|
||||
},
|
||||
"class_type": "UNETLoader",
|
||||
"_meta": {"title": "Load Diffusion Model"}
|
||||
},
|
||||
"13": {
|
||||
"inputs": {
|
||||
"seed": 0,
|
||||
"steps": 4,
|
||||
"cfg": 1.0,
|
||||
"sampler_name": "euler",
|
||||
"scheduler": "simple",
|
||||
"denoise": 1.0,
|
||||
"model": ["12", 0],
|
||||
"positive": ["6", 0],
|
||||
"negative": ["33", 0],
|
||||
"latent_image": ["5", 0]
|
||||
},
|
||||
"class_type": "KSampler",
|
||||
"_meta": {"title": "KSampler"}
|
||||
},
|
||||
"33": {
|
||||
"inputs": {
|
||||
"text": "",
|
||||
"clip": ["11", 0]
|
||||
},
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {"title": "Negative Prompt (empty for Schnell)"}
|
||||
}
|
||||
}
|
||||
47
services/image-render/providers/gpt_image.py
Normal file
47
services/image-render/providers/gpt_image.py
Normal file
@@ -0,0 +1,47 @@
|
||||
"""GPT Image 2.0 — OpenAI Images API.
|
||||
|
||||
POST https://api.openai.com/v1/images/generations
|
||||
body {model:"gpt-image-1", prompt, size, n:1} → data[0].b64_json
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
import requests
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
from providers._media import save_b64_png
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
OPENAI_URL = "https://api.openai.com/v1/images/generations"
|
||||
DEFAULT_MODEL = "gpt-image-1"
|
||||
|
||||
|
||||
def run_gpt_image_generation(task_id: str, params: dict) -> None:
|
||||
try:
|
||||
if not os.getenv("OPENAI_API_KEY"):
|
||||
webhook_update_task(task_id, "failed", 0, "", error="OPENAI_API_KEY 미설정 (Windows .env)")
|
||||
return
|
||||
webhook_update_task(task_id, "processing", 10, "GPT Image 호출 중...")
|
||||
body = {
|
||||
"model": params.get("model") or DEFAULT_MODEL,
|
||||
"prompt": params["prompt"],
|
||||
"size": params.get("size") or "1024x1024",
|
||||
"n": 1,
|
||||
}
|
||||
resp = requests.post(
|
||||
OPENAI_URL,
|
||||
headers={"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY')}", "Content-Type": "application/json"},
|
||||
json=body,
|
||||
timeout=120,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"OpenAI {resp.status_code}: {resp.text[:200]}")
|
||||
return
|
||||
b64 = resp.json()["data"][0]["b64_json"]
|
||||
url = save_b64_png(task_id, b64)
|
||||
webhook_update_task(task_id, "succeeded", 100, "완료", image_url=url)
|
||||
except Exception as e:
|
||||
logger.exception("gpt_image task=%s 실패", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
52
services/image-render/providers/nano_banana.py
Normal file
52
services/image-render/providers/nano_banana.py
Normal file
@@ -0,0 +1,52 @@
|
||||
"""Nano Banana — Gemini 2.5 Flash Image (generativelanguage API).
|
||||
|
||||
POST /v1beta/models/{MODEL}:generateContent
|
||||
→ candidates[0].content.parts[*].inlineData.data (b64 png)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging, os
|
||||
import requests
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
from providers._media import save_b64_png
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
GEMINI_BASE = "https://generativelanguage.googleapis.com/v1beta"
|
||||
DEFAULT_MODEL = "gemini-2.5-flash-image"
|
||||
|
||||
|
||||
def _extract_b64(data: dict):
|
||||
for cand in data.get("candidates", []):
|
||||
for part in cand.get("content", {}).get("parts", []):
|
||||
inline = part.get("inlineData") or part.get("inline_data")
|
||||
if inline and inline.get("data"):
|
||||
return inline["data"]
|
||||
return None
|
||||
|
||||
|
||||
def run_nano_banana_generation(task_id: str, params: dict) -> None:
|
||||
try:
|
||||
if not os.getenv("GEMINI_API_KEY"):
|
||||
webhook_update_task(task_id, "failed", 0, "", error="GEMINI_API_KEY 미설정 (Windows .env)")
|
||||
return
|
||||
webhook_update_task(task_id, "processing", 10, "Nano Banana (Gemini) 호출 중...")
|
||||
model_id = params.get("model") or DEFAULT_MODEL
|
||||
body = {"contents": [{"parts": [{"text": params["prompt"]}]}]}
|
||||
resp = requests.post(
|
||||
f"{GEMINI_BASE}/models/{model_id}:generateContent",
|
||||
headers={"x-goog-api-key": os.getenv("GEMINI_API_KEY"), "Content-Type": "application/json"},
|
||||
json=body, timeout=120,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Gemini {resp.status_code}: {resp.text[:200]}")
|
||||
return
|
||||
b64 = _extract_b64(resp.json())
|
||||
if not b64:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Gemini 응답에 이미지 없음")
|
||||
return
|
||||
url = save_b64_png(task_id, b64)
|
||||
webhook_update_task(task_id, "succeeded", 100, "완료", image_url=url)
|
||||
except Exception as e:
|
||||
logger.exception("nano_banana task=%s 실패", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
9
services/image-render/requirements.txt
Normal file
9
services/image-render/requirements.txt
Normal file
@@ -0,0 +1,9 @@
|
||||
fastapi==0.115.6
|
||||
uvicorn[standard]==0.34.0
|
||||
requests==2.32.3
|
||||
redis>=5.0
|
||||
httpx>=0.27
|
||||
openai>=1.50.0
|
||||
pytest>=8.0
|
||||
pytest-asyncio>=0.24
|
||||
respx>=0.21
|
||||
0
services/image-render/tests/__init__.py
Normal file
0
services/image-render/tests/__init__.py
Normal file
21
services/image-render/tests/test_flux.py
Normal file
21
services/image-render/tests/test_flux.py
Normal file
@@ -0,0 +1,21 @@
|
||||
import providers.flux as fx
|
||||
|
||||
def test_blocked_during_trading_hours(monkeypatch):
|
||||
monkeypatch.setenv("FLUX_BLOCK_TRADING_HOURS", "1")
|
||||
monkeypatch.setattr(fx, "_is_trading_hours", lambda: True)
|
||||
calls = []
|
||||
monkeypatch.setattr(fx, "webhook_update_task", lambda *a, **k: calls.append((a, k)))
|
||||
fx.run_flux_generation("t1", {"prompt": "a cat"})
|
||||
assert calls[-1][0][1] == "failed"
|
||||
assert "장중" in calls[-1][1]["error"]
|
||||
|
||||
def test_success_polls_history_and_saves(monkeypatch):
|
||||
monkeypatch.setattr(fx, "_is_trading_hours", lambda: False)
|
||||
calls = []
|
||||
monkeypatch.setattr(fx, "webhook_update_task", lambda *a, **k: calls.append((a, k)))
|
||||
monkeypatch.setattr(fx, "_load_workflow", lambda prompt, size: {"3": {}})
|
||||
monkeypatch.setattr(fx, "_submit_prompt", lambda wf: "pid-1")
|
||||
monkeypatch.setattr(fx, "_poll_image_b64", lambda pid: "ZmFrZQ==")
|
||||
monkeypatch.setattr(fx, "save_b64_png", lambda tid, b64: "/media/image/t1.png")
|
||||
fx.run_flux_generation("t1", {"prompt": "a cat"})
|
||||
assert [c for c in calls if c[0][1] == "succeeded"]
|
||||
32
services/image-render/tests/test_gpt_image.py
Normal file
32
services/image-render/tests/test_gpt_image.py
Normal file
@@ -0,0 +1,32 @@
|
||||
import providers.gpt_image as gi
|
||||
|
||||
|
||||
def test_missing_key_reports_failed(monkeypatch):
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
calls = []
|
||||
monkeypatch.setattr(gi, "webhook_update_task", lambda *a, **k: calls.append((a, k)))
|
||||
gi.run_gpt_image_generation("t1", {"prompt": "a cat"})
|
||||
# 마지막 호출이 failed
|
||||
assert calls[-1][0][1] == "failed"
|
||||
|
||||
|
||||
def test_success_saves_and_reports_url(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
calls = []
|
||||
monkeypatch.setattr(gi, "webhook_update_task", lambda *a, **k: calls.append((a, k)))
|
||||
monkeypatch.setattr(gi, "save_b64_png", lambda tid, b64: "/media/image/t1.png")
|
||||
|
||||
class FakeResp:
|
||||
status_code = 200
|
||||
|
||||
def json(self):
|
||||
return {"data": [{"b64_json": "ZmFrZQ=="}]}
|
||||
|
||||
def raise_for_status(self):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(gi.requests, "post", lambda *a, **k: FakeResp())
|
||||
|
||||
gi.run_gpt_image_generation("t1", {"prompt": "a cat"})
|
||||
succeeded = [c for c in calls if c[0][1] == "succeeded"]
|
||||
assert succeeded and succeeded[-1][1]["image_url"] == "/media/image/t1.png"
|
||||
25
services/image-render/tests/test_nano_banana.py
Normal file
25
services/image-render/tests/test_nano_banana.py
Normal file
@@ -0,0 +1,25 @@
|
||||
import providers.nano_banana as nb
|
||||
|
||||
def test_missing_key_reports_failed(monkeypatch):
|
||||
monkeypatch.delenv("GEMINI_API_KEY", raising=False)
|
||||
calls = []
|
||||
monkeypatch.setattr(nb, "webhook_update_task", lambda *a, **k: calls.append((a, k)))
|
||||
nb.run_nano_banana_generation("t1", {"prompt": "a cat"})
|
||||
assert calls[-1][0][1] == "failed"
|
||||
|
||||
def test_success_extracts_inline_data(monkeypatch):
|
||||
monkeypatch.setenv("GEMINI_API_KEY", "g-test")
|
||||
calls = []
|
||||
monkeypatch.setattr(nb, "webhook_update_task", lambda *a, **k: calls.append((a, k)))
|
||||
monkeypatch.setattr(nb, "save_b64_png", lambda tid, b64: "/media/image/t1.png")
|
||||
|
||||
class FakeResp:
|
||||
status_code = 200
|
||||
def json(self):
|
||||
return {"candidates": [{"content": {"parts": [
|
||||
{"inlineData": {"mimeType": "image/png", "data": "ZmFrZQ=="}}
|
||||
]}}]}
|
||||
monkeypatch.setattr(nb.requests, "post", lambda *a, **k: FakeResp())
|
||||
|
||||
nb.run_nano_banana_generation("t1", {"prompt": "a cat"})
|
||||
assert [c for c in calls if c[0][1] == "succeeded"]
|
||||
20
services/image-render/tests/test_nas_client.py
Normal file
20
services/image-render/tests/test_nas_client.py
Normal file
@@ -0,0 +1,20 @@
|
||||
import nas_client
|
||||
|
||||
|
||||
def test_webhook_includes_image_url(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_post(payload):
|
||||
captured.update(payload)
|
||||
|
||||
monkeypatch.setattr(nas_client, "_post", fake_post)
|
||||
nas_client.webhook_update_task("t1", "succeeded", 100, "done", image_url="/media/image/t1.png")
|
||||
assert captured["task_id"] == "t1"
|
||||
assert captured["image_url"] == "/media/image/t1.png"
|
||||
|
||||
|
||||
def test_webhook_omits_none_fields(monkeypatch):
|
||||
captured = {}
|
||||
monkeypatch.setattr(nas_client, "_post", lambda p: captured.update(p))
|
||||
nas_client.webhook_update_task("t2", "processing", 10, "working")
|
||||
assert "image_url" not in captured and "error" not in captured
|
||||
69
services/image-render/tests/test_worker.py
Normal file
69
services/image-render/tests/test_worker.py
Normal file
@@ -0,0 +1,69 @@
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
import worker
|
||||
|
||||
|
||||
def test_dispatch_routes_to_provider(monkeypatch):
|
||||
called = {}
|
||||
monkeypatch.setattr(worker, "run_gpt_image_generation", lambda tid, p: called.setdefault("gpt", (tid, p)))
|
||||
worker._dispatch({"job_type": "gpt_image_generation", "task_id": "t1", "params": {"prompt": "x"}})
|
||||
assert called["gpt"][0] == "t1"
|
||||
|
||||
|
||||
def test_dispatch_unknown_job_type_reports_failed(monkeypatch):
|
||||
calls = []
|
||||
monkeypatch.setattr(worker, "webhook_update_task", lambda *a, **k: calls.append((a, k)))
|
||||
worker._dispatch({"job_type": "midjourney_generation", "task_id": "t9", "params": {}})
|
||||
assert calls[-1][0][1] == "failed"
|
||||
|
||||
|
||||
# ----- F6: ReliableQueue poll_once -----
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_acks_on_success(monkeypatch):
|
||||
payload = {"task_id": "t1", "job_type": "gpt_image_generation", "params": {}}
|
||||
raw = json.dumps(payload).encode()
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=(payload, raw))
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
monkeypatch.setattr(worker, "_dispatch", MagicMock())
|
||||
handled = await worker.poll_once(fake_queue)
|
||||
assert handled is True
|
||||
fake_queue.ack.assert_awaited_once_with(raw)
|
||||
fake_queue.fail.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_calls_fail_on_dispatch_exception(monkeypatch):
|
||||
payload = {"task_id": "t2", "job_type": "gpt_image_generation", "params": {}}
|
||||
raw = json.dumps(payload).encode()
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=(payload, raw))
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
|
||||
def _boom(p):
|
||||
raise RuntimeError("dispatch crash")
|
||||
|
||||
monkeypatch.setattr(worker, "_dispatch", _boom)
|
||||
handled = await worker.poll_once(fake_queue)
|
||||
assert handled is True
|
||||
fake_queue.fail.assert_awaited_once_with(raw, payload)
|
||||
fake_queue.ack.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_returns_false_on_timeout(monkeypatch):
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=None)
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
monkeypatch.setattr(worker, "_dispatch", MagicMock())
|
||||
handled = await worker.poll_once(fake_queue)
|
||||
assert handled is False
|
||||
fake_queue.ack.assert_not_awaited()
|
||||
fake_queue.fail.assert_not_awaited()
|
||||
102
services/image-render/worker.py
Normal file
102
services/image-render/worker.py
Normal file
@@ -0,0 +1,102 @@
|
||||
"""Redis ReliableQueue worker — F6 신뢰성 패턴 (BLMOVE + ack/fail + recovery).
|
||||
|
||||
queue:paused 가 set이면 대기 (task-watcher가 박재오 활동 감지 시 set).
|
||||
string-based dispatch + getattr (테스트 patch 호환).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
from providers.gpt_image import run_gpt_image_generation
|
||||
from providers.nano_banana import run_nano_banana_generation
|
||||
from providers.flux import run_flux_generation
|
||||
from _shared.reliable_queue import ReliableQueue
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
REDIS_URL = os.getenv("REDIS_URL", "redis://192.168.45.54:6379")
|
||||
QUEUE_KEY = "queue:image-render"
|
||||
PAUSED_KEY = "queue:paused"
|
||||
|
||||
# string names so `unittest.mock.patch` / `monkeypatch.setattr` on `worker.<name>`
|
||||
# is correctly intercepted by getattr(sys.modules[__name__], ...)
|
||||
_DISPATCH_TABLE = {
|
||||
"gpt_image_generation": "run_gpt_image_generation",
|
||||
"nano_banana_generation": "run_nano_banana_generation",
|
||||
"flux_generation": "run_flux_generation",
|
||||
}
|
||||
|
||||
|
||||
def _dispatch(payload: dict) -> None:
|
||||
"""payload[job_type] → provider 함수 호출 (sync, worker_loop에서 asyncio.to_thread로 wrap)."""
|
||||
job_type = payload.get("job_type", "")
|
||||
task_id = payload.get("task_id", "")
|
||||
params = payload.get("params", {})
|
||||
fn_name = _DISPATCH_TABLE.get(job_type)
|
||||
if fn_name is None:
|
||||
logger.error("unknown job_type=%s task=%s", job_type, task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"unknown job_type: {job_type}")
|
||||
return
|
||||
try:
|
||||
fn = getattr(sys.modules[__name__], fn_name)
|
||||
except AttributeError:
|
||||
logger.error("dispatch table typo for job_type=%s name=%s task=%s", job_type, fn_name, task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"internal dispatch error: {fn_name}")
|
||||
return
|
||||
fn(task_id, params)
|
||||
|
||||
|
||||
async def poll_once(queue: ReliableQueue) -> bool:
|
||||
"""F6 — 1 cycle: dequeue → _dispatch → ack/fail. Returns True if a job handled."""
|
||||
result = await queue.dequeue(timeout=5)
|
||||
if result is None:
|
||||
return False
|
||||
payload, raw = result
|
||||
try:
|
||||
await asyncio.to_thread(_dispatch, payload)
|
||||
except Exception:
|
||||
logger.exception("dispatch unhandled exception task_id=%s",
|
||||
payload.get("task_id"))
|
||||
await queue.fail(raw, payload)
|
||||
return True
|
||||
await queue.ack(raw)
|
||||
return True
|
||||
|
||||
|
||||
async def worker_loop():
|
||||
redis = aioredis.from_url(REDIS_URL, decode_responses=False)
|
||||
queue = ReliableQueue(redis, queue_key=QUEUE_KEY)
|
||||
logger.info("image-render worker started worker_id=%s queue=%s",
|
||||
queue.worker_id, QUEUE_KEY)
|
||||
try:
|
||||
recovered = await queue.recover()
|
||||
if recovered:
|
||||
logger.info("recovered %d orphaned items at startup", recovered)
|
||||
except Exception:
|
||||
logger.exception("startup recover failed")
|
||||
|
||||
while True:
|
||||
try:
|
||||
paused = await redis.get(PAUSED_KEY)
|
||||
if paused == b"1":
|
||||
await asyncio.sleep(10)
|
||||
continue
|
||||
await poll_once(queue)
|
||||
except asyncio.CancelledError:
|
||||
logger.info("worker_loop cancelled")
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("worker_loop iteration 실패, 5초 후 재시도")
|
||||
await asyncio.sleep(5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
asyncio.run(worker_loop())
|
||||
@@ -7,8 +7,9 @@ REDIS_URL=redis://192.168.45.54:6379
|
||||
NAS_BASE_URL=http://192.168.45.54:18700
|
||||
INTERNAL_API_KEY=__copy_from_nas_dotenv__
|
||||
|
||||
# NAS SMB mount 안의 미디어 디렉토리 (/mnt/nas/webpage/data/insta/)
|
||||
INSTA_MEDIA_ROOT=/mnt/nas/webpage/data/insta
|
||||
# NAS SMB mount 안의 미디어 디렉토리.
|
||||
# ⚠️ nginx가 /media/insta를 data/insta/insta_cards/로 서빙하므로 반드시 insta_cards까지 포함.
|
||||
INSTA_MEDIA_ROOT=/mnt/nas/webpage/data/insta/insta_cards
|
||||
|
||||
# nginx 서빙 prefix (NAS webhook payload에 보낼 result_path 만들 때)
|
||||
INSTA_MEDIA_URL_PREFIX=/media/insta
|
||||
|
||||
@@ -12,11 +12,14 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
libcairo2 libasound2 libatspi2.0-0 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY requirements.txt .
|
||||
COPY insta-render/requirements.txt /app/
|
||||
RUN pip install --no-cache-dir --timeout 600 --retries 5 -r requirements.txt
|
||||
RUN playwright install chromium
|
||||
|
||||
COPY . .
|
||||
# F6: 공통 ReliableQueue 모듈 (services/_shared)
|
||||
COPY _shared /app/_shared
|
||||
COPY insta-render/. /app/
|
||||
ENV PYTHONPATH=/app
|
||||
|
||||
EXPOSE 8000
|
||||
CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "1"]
|
||||
|
||||
@@ -151,8 +151,11 @@ async def _render_slate_locked(slate: dict, slate_id: int, template: str) -> Lis
|
||||
html_path = f.name
|
||||
try:
|
||||
await page.goto(f"file://{html_path}", wait_until="networkidle")
|
||||
await page.evaluate("document.fonts.ready") # 웹폰트 로딩 완료까지 대기
|
||||
out_path = os.path.join(out_dir, f"{spec['page_no']:02d}.png")
|
||||
await page.screenshot(path=out_path, full_page=False, omit_background=False)
|
||||
if os.path.getsize(out_path) < 1000: # 빈/깨진 PNG 방어
|
||||
raise RuntimeError(f"rendered PNG too small: {out_path}")
|
||||
paths.append(out_path)
|
||||
finally:
|
||||
try:
|
||||
|
||||
5
services/insta-render/conftest.py
Normal file
5
services/insta-render/conftest.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""Make services/ root importable so `from _shared.reliable_queue import ...` works during tests."""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
@@ -3,52 +3,85 @@
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<style>
|
||||
@import url('https://fonts.googleapis.com/css2?family=Noto+Sans+KR:wght@400;700;900&display=swap');
|
||||
@import url('https://cdn.jsdelivr.net/gh/orioncactus/pretendard@v1.3.9/dist/web/static/pretendard.css');
|
||||
* { margin: 0; padding: 0; box-sizing: border-box; }
|
||||
html, body {
|
||||
width: 1080px; height: 1350px;
|
||||
font-family: 'Noto Sans KR', sans-serif;
|
||||
html, body { width: 1080px; height: 1350px; }
|
||||
body {
|
||||
font-family: 'Pretendard', 'Noto Sans KR', sans-serif;
|
||||
background: #F7F7FA; color: #14171A;
|
||||
-webkit-font-smoothing: antialiased;
|
||||
}
|
||||
.card {
|
||||
width: 1080px; height: 1350px;
|
||||
padding: 80px 72px;
|
||||
display: flex; flex-direction: column; justify-content: space-between;
|
||||
background: linear-gradient(180deg, #FFFFFF 0%, #F7F7FA 100%);
|
||||
border-top: 16px solid {{ accent_color }};
|
||||
position: relative; width: 1080px; height: 1350px; overflow: hidden;
|
||||
padding: 96px 84px 72px;
|
||||
display: flex; flex-direction: column;
|
||||
background: #FFFFFF;
|
||||
}
|
||||
.accent-bar { position: absolute; top: 0; left: 0; width: 100%; height: 14px; background: {{ accent_color | safe }}; }
|
||||
.badge {
|
||||
display: inline-block; padding: 8px 20px; border-radius: 999px;
|
||||
background: {{ accent_color }}; color: #fff;
|
||||
font-size: 28px; font-weight: 700; letter-spacing: -0.02em;
|
||||
align-self: flex-start; padding: 10px 24px; border-radius: 999px;
|
||||
background: {{ accent_color | safe }}; color: #fff;
|
||||
font-size: 30px; font-weight: 700; letter-spacing: -0.02em;
|
||||
}
|
||||
.idx { font-size: 120px; font-weight: 800; line-height: 1; color: {{ accent_color | safe }}; letter-spacing: -0.04em; }
|
||||
.content { flex: 1; display: flex; flex-direction: column; justify-content: center; gap: 36px; }
|
||||
.headline {
|
||||
font-size: {{ 96 if page_type == 'cover' else 72 }}px;
|
||||
font-weight: 900; line-height: 1.15; letter-spacing: -0.04em;
|
||||
margin-top: 32px;
|
||||
font-weight: 800; line-height: 1.18; letter-spacing: -0.04em; color: #14171A;
|
||||
display: -webkit-box; -webkit-box-orient: vertical; overflow: hidden;
|
||||
}
|
||||
.body {
|
||||
font-size: 40px; font-weight: 400; line-height: 1.55;
|
||||
margin-top: 40px; color: #2A2F35;
|
||||
.cover .headline { font-size: 104px; -webkit-line-clamp: 4; }
|
||||
.body-page .headline { font-size: 76px; -webkit-line-clamp: 3; }
|
||||
.cta .headline { font-size: 88px; -webkit-line-clamp: 3; }
|
||||
.sub {
|
||||
font-size: 42px; font-weight: 400; line-height: 1.5; color: #3A4047;
|
||||
display: -webkit-box; -webkit-box-orient: vertical; overflow: hidden; -webkit-line-clamp: 8;
|
||||
white-space: pre-wrap;
|
||||
}
|
||||
.cover .sub { -webkit-line-clamp: 5; }
|
||||
.footer {
|
||||
display: flex; justify-content: space-between; align-items: center;
|
||||
font-size: 28px; color: #6B7280; font-weight: 500;
|
||||
font-size: 28px; color: #8A9099; font-weight: 600; margin-top: 40px;
|
||||
}
|
||||
.cta { font-weight: 700; color: {{ accent_color }}; }
|
||||
.cta-pill {
|
||||
align-self: flex-start; margin-top: 8px; padding: 18px 40px; border-radius: 16px;
|
||||
background: {{ accent_color | safe }}; color: #fff; font-size: 40px; font-weight: 700;
|
||||
}
|
||||
.progress { display: flex; gap: 10px; }
|
||||
.progress i { width: 14px; height: 14px; border-radius: 50%; background: #D8DCE0; display: inline-block; }
|
||||
.progress i.on { background: {{ accent_color | safe }}; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="card">
|
||||
<div>
|
||||
<span class="badge">{{ page_type|upper }}</span>
|
||||
<h1 class="headline">{{ headline }}</h1>
|
||||
<p class="body">{{ body }}</p>
|
||||
</div>
|
||||
<div class="card {{ 'cover' if page_type=='cover' else ('cta' if page_type=='cta' else 'body-page') }}">
|
||||
<div class="accent-bar"></div>
|
||||
|
||||
{% if page_type == 'cover' %}
|
||||
<span class="badge">{{ category_label|default('') or '오늘의 이슈' }}</span>
|
||||
<div class="content">
|
||||
<h1 class="headline">{{ headline }}</h1>
|
||||
<p class="sub">{{ body }}</p>
|
||||
</div>
|
||||
{% elif page_type == 'cta' %}
|
||||
<div class="content">
|
||||
<h1 class="headline">{{ headline }}</h1>
|
||||
<p class="sub">{{ body }}</p>
|
||||
{% if cta %}<div class="cta-pill">{{ cta }}</div>{% endif %}
|
||||
</div>
|
||||
{% else %}
|
||||
<span class="idx">{{ '%02d'|format(page_no - 1) }}</span>
|
||||
<div class="content">
|
||||
<h1 class="headline">{{ headline }}</h1>
|
||||
<p class="sub">{{ body }}</p>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<div class="footer">
|
||||
<span>{{ page_no }} / {{ total_pages }}</span>
|
||||
{% if cta %}<span class="cta">{{ cta }}</span>{% endif %}
|
||||
{% if page_type == 'cover' or page_type == 'cta' %}
|
||||
<span>{{ brand_handle|default('') }}</span><span>{{ page_no }} / {{ total_pages }}</span>
|
||||
{% else %}
|
||||
<div class="progress">{% for n in range(2, total_pages) %}<i class="{{ 'on' if n <= page_no }}"></i>{% endfor %}</div>
|
||||
<span>{{ page_no }} / {{ total_pages }}</span>
|
||||
{% endif %}
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
"""worker.py — Redis BLPOP + webhook 단위 테스트."""
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import pytest
|
||||
import httpx
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import worker
|
||||
from card_renderer import render_slate, init_browser, shutdown_browser
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -112,11 +115,118 @@ async def test_process_one_render_failure_reports_failed(monkeypatch, fake_slate
|
||||
worker.NAS_BASE_URL = "http://nas.test"
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
await worker._process_one(client, {
|
||||
"task_id": "t-3",
|
||||
"params": {"slate_id": 99},
|
||||
})
|
||||
# F6: _process_one은 webhook(failed) 호출 후 raise — poll_once가 fail(raw)로 retry/dead-letter.
|
||||
with pytest.raises(RuntimeError, match="Chromium"):
|
||||
await worker._process_one(client, {
|
||||
"task_id": "t-3",
|
||||
"params": {"slate_id": 99},
|
||||
})
|
||||
|
||||
last = calls[-1]
|
||||
assert last["status"] == "failed"
|
||||
assert "Chromium" in last["error"]
|
||||
|
||||
|
||||
# ----- F6: ReliableQueue (ack on success, fail on exception) -----
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_acks_on_success(monkeypatch):
|
||||
"""F6 — 성공 시 queue.ack(raw) 호출 + fail 안 부름."""
|
||||
fake_payload = {
|
||||
"task_id": "t-ok",
|
||||
"params": {"slate_id": 7, "theme": "default"},
|
||||
}
|
||||
fake_raw = json.dumps(fake_payload).encode()
|
||||
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=(fake_payload, fake_raw))
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
|
||||
process_mock = AsyncMock()
|
||||
monkeypatch.setattr(worker, "_process_one", process_mock)
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
handled = await worker.poll_once(fake_queue, client)
|
||||
|
||||
assert handled is True
|
||||
process_mock.assert_awaited_once()
|
||||
fake_queue.ack.assert_awaited_once_with(fake_raw)
|
||||
fake_queue.fail.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_calls_fail_on_exception(monkeypatch):
|
||||
"""F6 — _process_one 예외 시 queue.fail(raw, payload) 호출."""
|
||||
fake_payload = {
|
||||
"task_id": "t-err",
|
||||
"params": {"slate_id": 9, "theme": "default"},
|
||||
}
|
||||
fake_raw = json.dumps(fake_payload).encode()
|
||||
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=(fake_payload, fake_raw))
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
|
||||
async def boom(client, payload):
|
||||
raise RuntimeError("simulated dispatch failure")
|
||||
|
||||
monkeypatch.setattr(worker, "_process_one", boom)
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
handled = await worker.poll_once(fake_queue, client)
|
||||
|
||||
assert handled is True
|
||||
fake_queue.fail.assert_awaited_once_with(fake_raw, fake_payload)
|
||||
fake_queue.ack.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_render_produces_nonempty_1080x1350(tmp_path, monkeypatch):
|
||||
"""Phase 2 — fonts.ready 대기 + PNG 비어있음 검증: 10장 모두 > 1000 bytes."""
|
||||
import card_renderer as _cr
|
||||
templates_dir = str(Path(__file__).resolve().parent.parent / "templates")
|
||||
monkeypatch.setattr(_cr, "CARD_TEMPLATE_DIR", templates_dir)
|
||||
monkeypatch.setattr(_cr, "INSTA_MEDIA_ROOT", str(tmp_path))
|
||||
await init_browser()
|
||||
try:
|
||||
slate = {
|
||||
"cover_copy": {"headline": "헤드라인", "body": "서브", "accent_color": "#0F62FE"},
|
||||
"body_copies": [{"headline": f"포인트{i}", "body": "본문"} for i in range(8)],
|
||||
"cta_copy": {"headline": "요약", "body": "마무리", "cta": "팔로우"},
|
||||
}
|
||||
paths = await render_slate(slate, slate_id=99999)
|
||||
assert len(paths) == 10
|
||||
for p in paths:
|
||||
assert os.path.getsize(p) > 1000 # 비어있지 않음
|
||||
finally:
|
||||
await shutdown_browser()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_returns_false_on_timeout(monkeypatch):
|
||||
"""F6 — dequeue가 None 반환(타임아웃)이면 False 리턴, ack/fail 안 부름."""
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=None)
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
|
||||
process_mock = AsyncMock()
|
||||
monkeypatch.setattr(worker, "_process_one", process_mock)
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
handled = await worker.poll_once(fake_queue, client)
|
||||
|
||||
assert handled is False
|
||||
process_mock.assert_not_awaited()
|
||||
fake_queue.ack.assert_not_awaited()
|
||||
fake_queue.fail.assert_not_awaited()
|
||||
|
||||
|
||||
def test_make_queue_redis_socket_timeout_exceeds_block():
|
||||
"""BLMOVE(블록 5s) dequeue가 read-timeout 경계 경합으로 깨지지 않도록
|
||||
socket_timeout이 블록보다 충분히 커야 한다 (회귀 가드)."""
|
||||
c = worker.make_queue_redis()
|
||||
st = c.connection_pool.connection_kwargs.get("socket_timeout")
|
||||
assert st is not None and st > 5 # blmove 블록(5s)보다 커야 안정
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
"""Redis BLPOP worker — queue:insta-render → render_slate → NAS webhook.
|
||||
"""Redis ReliableQueue worker — F6 신뢰성 패턴 (BLMOVE + ack/fail + recovery).
|
||||
|
||||
queue:paused가 set이면 대기 (task-watcher가 박재오 활동 감지 시 set).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any
|
||||
@@ -14,6 +13,7 @@ import httpx
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
from card_renderer import render_slate
|
||||
from _shared.reliable_queue import ReliableQueue
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -57,7 +57,10 @@ async def _fetch_slate(client: httpx.AsyncClient, slate_id: int) -> dict:
|
||||
|
||||
|
||||
async def _process_one(client: httpx.AsyncClient, payload: dict) -> None:
|
||||
"""단일 작업 처리: fetch slate → render → webhook."""
|
||||
"""단일 작업 처리: fetch slate → render → webhook. 예외 발생 시 webhook(failed) 호출 후 raise.
|
||||
|
||||
F6: webhook 통신 외 예외는 poll_once가 fail(raw, payload)로 retry/dead-letter 처리.
|
||||
"""
|
||||
task_id = payload["task_id"]
|
||||
params = payload.get("params", {})
|
||||
slate_id = params.get("slate_id")
|
||||
@@ -69,7 +72,6 @@ async def _process_one(client: httpx.AsyncClient, payload: dict) -> None:
|
||||
slate = await _fetch_slate(client, slate_id)
|
||||
await _post_update(client, task_id, "processing", 50)
|
||||
paths = await render_slate(slate, slate_id, template=template)
|
||||
# 결과 URL은 첫 페이지의 nginx 경로
|
||||
first_url = f"{INSTA_MEDIA_URL_PREFIX}/{slate_id}/01.png"
|
||||
await _post_update(
|
||||
client, task_id, "succeeded", 100, result_path=first_url
|
||||
@@ -78,29 +80,61 @@ async def _process_one(client: httpx.AsyncClient, payload: dict) -> None:
|
||||
except Exception as e:
|
||||
logger.exception("render task=%s 실패", task_id)
|
||||
await _post_update(client, task_id, "failed", 0, error=str(e))
|
||||
raise
|
||||
|
||||
|
||||
async def poll_once(queue: ReliableQueue, client: httpx.AsyncClient) -> bool:
|
||||
"""1 cycle: dequeue → _process_one → ack/fail. Returns True if a job handled."""
|
||||
result = await queue.dequeue(timeout=5)
|
||||
if result is None:
|
||||
return False
|
||||
payload, raw = result
|
||||
try:
|
||||
await _process_one(client, payload)
|
||||
except Exception:
|
||||
await queue.fail(raw, payload)
|
||||
return True
|
||||
await queue.ack(raw)
|
||||
return True
|
||||
|
||||
|
||||
# 블로킹 dequeue는 BLMOVE(블록 5s)를 쓴다. redis-py 블로킹 read에서 socket_timeout이
|
||||
# 블록(5s) 이하이거나 None이면 read-timeout이 블록 경계와 경합해 간헐적으로
|
||||
# "Timeout reading"이 터져 잡을 못 꺼낸다(슬레이트 draft 정지). 실험상 socket_timeout이
|
||||
# 블록보다 충분히 크면(10/30) 항상 안정. → 블록보다 넉넉히 큰 값을 명시한다.
|
||||
QUEUE_SOCKET_TIMEOUT = 30 # > dequeue blmove 블록(5s)
|
||||
|
||||
|
||||
def make_queue_redis():
|
||||
"""블로킹 dequeue(BLMOVE)용 redis 클라이언트. socket_timeout > 블록(5s) 보장."""
|
||||
return aioredis.from_url(
|
||||
REDIS_URL, decode_responses=False,
|
||||
socket_timeout=QUEUE_SOCKET_TIMEOUT, socket_keepalive=True,
|
||||
)
|
||||
|
||||
|
||||
async def worker_loop():
|
||||
"""무한 루프 — paused 체크 → BLPOP → process_one."""
|
||||
redis = aioredis.from_url(REDIS_URL, decode_responses=False)
|
||||
"""무한 루프 — paused 체크 → ReliableQueue.dequeue → process_one → ack/fail."""
|
||||
redis = make_queue_redis()
|
||||
queue = ReliableQueue(redis, queue_key=QUEUE_KEY)
|
||||
async with httpx.AsyncClient() as client:
|
||||
logger.info("insta-render worker started (queue=%s)", QUEUE_KEY)
|
||||
logger.info("insta-render worker started worker_id=%s queue=%s",
|
||||
queue.worker_id, QUEUE_KEY)
|
||||
# F6: startup recovery — 이전 crash 시 잔존 orphan 재큐
|
||||
try:
|
||||
recovered = await queue.recover()
|
||||
if recovered:
|
||||
logger.info("recovered %d orphaned items at startup", recovered)
|
||||
except Exception:
|
||||
logger.exception("startup recover failed")
|
||||
|
||||
while True:
|
||||
try:
|
||||
paused = await redis.get(PAUSED_KEY)
|
||||
if paused == b"1":
|
||||
await asyncio.sleep(10)
|
||||
continue
|
||||
item = await redis.blpop(QUEUE_KEY, timeout=1)
|
||||
if item is None:
|
||||
continue
|
||||
_, raw = item
|
||||
try:
|
||||
payload = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
logger.error("invalid queue payload: %r", raw[:200])
|
||||
continue
|
||||
await _process_one(client, payload)
|
||||
await poll_once(queue, client)
|
||||
except asyncio.CancelledError:
|
||||
logger.info("worker_loop cancelled")
|
||||
raise
|
||||
|
||||
20
services/music-render/.env.example
Normal file
20
services/music-render/.env.example
Normal file
@@ -0,0 +1,20 @@
|
||||
# Plan-B-Music — Windows music-render worker
|
||||
|
||||
# NAS Redis 큐
|
||||
REDIS_URL=redis://192.168.45.54:6379
|
||||
|
||||
# NAS internal webhook
|
||||
NAS_BASE_URL=http://192.168.45.54:18600
|
||||
INTERNAL_API_KEY=__copy_from_nas_dotenv__
|
||||
|
||||
# Suno API (sunoapi.org 래퍼) — NAS .env에서 옮겨옴
|
||||
SUNO_API_KEY=__paste_suno_key_here__
|
||||
|
||||
# MusicGen 호스트 (Windows native Python — 박재오 PC localhost)
|
||||
MUSIC_AI_SERVER_URL=http://host.docker.internal:8765
|
||||
|
||||
# NAS SMB mount 안의 음악 디렉토리
|
||||
MUSIC_MEDIA_ROOT=/mnt/nas/webpage/data/music
|
||||
|
||||
# nginx 서빙 prefix (NAS webhook payload용)
|
||||
MUSIC_MEDIA_URL_PREFIX=/media/music
|
||||
20
services/music-render/Dockerfile
Normal file
20
services/music-render/Dockerfile
Normal file
@@ -0,0 +1,20 @@
|
||||
FROM python:3.12-slim-bookworm
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# requests SSL 의존성만 필요 (Chromium 불필요)
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY music-render/requirements.txt /app/
|
||||
RUN pip install --no-cache-dir --timeout 600 --retries 5 -r requirements.txt
|
||||
|
||||
# F6: 공통 ReliableQueue 모듈 (services/_shared)
|
||||
COPY _shared /app/_shared
|
||||
COPY music-render/. /app/
|
||||
ENV PYTHONPATH=/app
|
||||
|
||||
EXPOSE 8000
|
||||
CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "1"]
|
||||
5
services/music-render/conftest.py
Normal file
5
services/music-render/conftest.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""Make services/ root importable so `from _shared.reliable_queue import ...` works during tests."""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
88
services/music-render/main.py
Normal file
88
services/music-render/main.py
Normal file
@@ -0,0 +1,88 @@
|
||||
"""music-render FastAPI entry — health + lifespan + sync forward endpoints.
|
||||
|
||||
NAS music-lab이 sync helpers(lyrics, credits, timestamped, style-boost)를
|
||||
httpx로 forward해서 이 endpoint들을 호출.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
import worker
|
||||
from providers.sync_ops import (
|
||||
generate_lyrics, get_credits,
|
||||
get_timestamped_lyrics, generate_style_boost,
|
||||
)
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
worker_task = asyncio.create_task(worker.worker_loop())
|
||||
logger.info("music-render lifespan 시작")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
worker_task.cancel()
|
||||
try:
|
||||
await worker_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
logger.info("music-render lifespan 종료")
|
||||
|
||||
|
||||
app = FastAPI(lifespan=lifespan)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"ok": True, "service": "music-render"}
|
||||
|
||||
|
||||
# ── Sync forward endpoints ──────────────────────────────────────────────
|
||||
# NAS music-lab의 /api/music/lyrics 등 sync helpers가 이 endpoint들로 forward.
|
||||
|
||||
class LyricsRequest(BaseModel):
|
||||
prompt: str
|
||||
|
||||
|
||||
@app.post("/api/music-render/sync/lyrics")
|
||||
def sync_lyrics(req: LyricsRequest):
|
||||
result = generate_lyrics(req.prompt)
|
||||
if not result:
|
||||
raise HTTPException(502, "가사 생성 실패")
|
||||
return result
|
||||
|
||||
|
||||
@app.get("/api/music-render/sync/credits")
|
||||
def sync_credits():
|
||||
result = get_credits()
|
||||
if result is None:
|
||||
raise HTTPException(502, "크레딧 조회 실패")
|
||||
return result
|
||||
|
||||
|
||||
@app.get("/api/music-render/sync/timestamped-lyrics")
|
||||
def sync_timestamped_lyrics(task_id: str, suno_id: str):
|
||||
result = get_timestamped_lyrics(task_id, suno_id)
|
||||
if not result:
|
||||
raise HTTPException(502, "타임스탬프 가사 조회 실패")
|
||||
return result
|
||||
|
||||
|
||||
class StyleBoostRequest(BaseModel):
|
||||
content: str
|
||||
|
||||
|
||||
@app.post("/api/music-render/sync/style-boost")
|
||||
def sync_style_boost(req: StyleBoostRequest):
|
||||
result = generate_style_boost(req.content)
|
||||
if not result:
|
||||
raise HTTPException(502, "스타일 부스트 생성 실패")
|
||||
return result
|
||||
80
services/music-render/nas_client.py
Normal file
80
services/music-render/nas_client.py
Normal file
@@ -0,0 +1,80 @@
|
||||
"""NAS webhook 어댑터 — Windows worker가 NAS DB 직접 접근 못하므로 HTTP로 위임.
|
||||
|
||||
기존 NAS suno_provider/local_provider의 `update_task`, `add_track` 호출을
|
||||
이 모듈의 webhook_update_task/webhook_add_track으로 치환.
|
||||
|
||||
webhook 실패는 raise하지 않고 logger.error로 기록 (provider 로직 흐름 유지).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_TIMEOUT = 10.0
|
||||
|
||||
|
||||
def _post(payload: Dict[str, Any]) -> None:
|
||||
nas_base_url = os.getenv("NAS_BASE_URL", "http://192.168.45.54:18600")
|
||||
internal_api_key = os.getenv("INTERNAL_API_KEY", "")
|
||||
url = f"{nas_base_url}/api/internal/music/update"
|
||||
try:
|
||||
r = httpx.post(
|
||||
url,
|
||||
headers={"X-Internal-Key": internal_api_key},
|
||||
json=payload,
|
||||
timeout=_TIMEOUT,
|
||||
)
|
||||
if r.status_code != 200:
|
||||
logger.error("webhook %s returned %d: %s",
|
||||
payload.get("task_id"), r.status_code, r.text[:200])
|
||||
except Exception:
|
||||
logger.exception("webhook %s 호출 실패", payload.get("task_id"))
|
||||
|
||||
|
||||
def webhook_update_task(
|
||||
task_id: str,
|
||||
status: str,
|
||||
progress: int,
|
||||
message: str = "",
|
||||
audio_url: Optional[str] = None,
|
||||
error: Optional[str] = None,
|
||||
) -> None:
|
||||
"""기존 update_task(task_id, status, progress, message, audio_url=None, error=None) 대체."""
|
||||
payload: Dict[str, Any] = {
|
||||
"task_id": task_id,
|
||||
"status": status,
|
||||
"progress": progress,
|
||||
"message": message,
|
||||
}
|
||||
if audio_url is not None:
|
||||
payload["audio_url"] = audio_url
|
||||
if error is not None:
|
||||
payload["error"] = error
|
||||
_post(payload)
|
||||
|
||||
|
||||
def webhook_add_track(
|
||||
task_id: str,
|
||||
status: str,
|
||||
progress: int,
|
||||
message: str = "",
|
||||
audio_url: Optional[str] = None,
|
||||
track: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""update + add_track을 한 webhook 호출로 결합 (NAS internal_router가 둘 다 처리)."""
|
||||
payload: Dict[str, Any] = {
|
||||
"task_id": task_id,
|
||||
"status": status,
|
||||
"progress": progress,
|
||||
"message": message,
|
||||
}
|
||||
if audio_url is not None:
|
||||
payload["audio_url"] = audio_url
|
||||
if track is not None:
|
||||
payload["track"] = track
|
||||
_post(payload)
|
||||
0
services/music-render/providers/__init__.py
Normal file
0
services/music-render/providers/__init__.py
Normal file
106
services/music-render/providers/local.py
Normal file
106
services/music-render/providers/local.py
Normal file
@@ -0,0 +1,106 @@
|
||||
"""Local MusicGen Provider — Windows AI 머신의 native MusicGen 서버(:8765) 호출.
|
||||
|
||||
NAS music-lab/app/local_provider.py 이식. DB 호출만 webhook으로 변환.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from nas_client import webhook_update_task, webhook_add_track
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MUSIC_AI_SERVER_URL = os.getenv("MUSIC_AI_SERVER_URL", "")
|
||||
MUSIC_MEDIA_ROOT = os.getenv("MUSIC_MEDIA_ROOT", "/mnt/nas/webpage/data/music")
|
||||
MUSIC_MEDIA_BASE = os.getenv("MUSIC_MEDIA_URL_PREFIX", "/media/music")
|
||||
|
||||
|
||||
def run_local_generation(task_id: str, params: dict) -> None:
|
||||
"""MusicGen 생성 → /mnt/nas/.../music/{task_id}.mp3 저장 → add_track."""
|
||||
try:
|
||||
webhook_update_task(task_id, "processing", 10, "AI 서버에 연결 중...")
|
||||
if not MUSIC_AI_SERVER_URL:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="MUSIC_AI_SERVER_URL 미설정")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 30, "음악 생성 중...")
|
||||
resp = requests.post(f"{MUSIC_AI_SERVER_URL}/generate", json=params, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "",
|
||||
error=f"AI 서버 오류: {resp.status_code} {resp.text[:200]}")
|
||||
return
|
||||
|
||||
ai_task_id = resp.json().get("task_id")
|
||||
if not ai_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="AI 서버 응답에 task_id 없음")
|
||||
return
|
||||
|
||||
remote_url = None
|
||||
for _ in range(120):
|
||||
time.sleep(5)
|
||||
sr = requests.get(f"{MUSIC_AI_SERVER_URL}/status/{ai_task_id}", timeout=10)
|
||||
sd = sr.json()
|
||||
st = sd.get("status")
|
||||
prog = sd.get("progress", 0)
|
||||
msg = sd.get("message", "음악 생성 중...")
|
||||
scaled = 30 + int(prog * 0.49)
|
||||
webhook_update_task(task_id, "processing", scaled, msg)
|
||||
|
||||
if st == "succeeded":
|
||||
remote_url = sd.get("audio_url")
|
||||
break
|
||||
elif st == "failed":
|
||||
webhook_update_task(task_id, "failed", 0, "",
|
||||
error=sd.get("error", "AI 서버 생성 실패"))
|
||||
return
|
||||
|
||||
if not remote_url:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="AI 서버 타임아웃 (10분)")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 80, "파일 저장 중...")
|
||||
filename = f"{task_id}.mp3"
|
||||
os.makedirs(MUSIC_MEDIA_ROOT, exist_ok=True)
|
||||
file_path = os.path.join(MUSIC_MEDIA_ROOT, filename)
|
||||
|
||||
dl = requests.get(remote_url, timeout=120, stream=True)
|
||||
dl.raise_for_status()
|
||||
with open(file_path, "wb") as f:
|
||||
for chunk in dl.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
audio_url = f"{MUSIC_MEDIA_BASE}/{filename}"
|
||||
genre = params.get("genre", "")
|
||||
moods = params.get("moods", [])
|
||||
mood_str = moods[0] if moods else "Original"
|
||||
title = params.get("title") or (
|
||||
f"{genre} — {mood_str} Mix" if genre else f"{mood_str} Mix"
|
||||
)
|
||||
|
||||
track = {
|
||||
"title": title,
|
||||
"genre": genre,
|
||||
"moods": moods,
|
||||
"instruments": params.get("instruments", []),
|
||||
"duration_sec": params.get("duration_sec"),
|
||||
"bpm": params.get("bpm"),
|
||||
"key": params.get("key", ""),
|
||||
"scale": params.get("scale", ""),
|
||||
"prompt": params.get("prompt", ""),
|
||||
"audio_url": audio_url,
|
||||
"file_path": f"/app/data/{filename}",
|
||||
"task_id": task_id,
|
||||
"provider": "local",
|
||||
}
|
||||
webhook_add_track(task_id, "succeeded", 100, "생성 완료",
|
||||
audio_url=audio_url, track=track)
|
||||
|
||||
except requests.Timeout:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="AI 서버 타임아웃")
|
||||
except Exception as e:
|
||||
logger.exception("local generation error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
690
services/music-render/providers/suno.py
Normal file
690
services/music-render/providers/suno.py
Normal file
@@ -0,0 +1,690 @@
|
||||
"""Suno API Provider — sunoapi.org 래퍼.
|
||||
|
||||
NAS music-lab/app/suno_provider.py에서 이식. 차이점:
|
||||
- DB 호출(update_task, add_track 등)을 nas_client.webhook_* 으로 변환
|
||||
- 결과 MP3는 MUSIC_MEDIA_ROOT (/mnt/nas/webpage/data/music/)에 직접 저장
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
from nas_client import webhook_update_task, webhook_add_track
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SUNO_BASE_URL = "https://api.sunoapi.org/api/v1"
|
||||
SUNO_API_KEY = os.getenv("SUNO_API_KEY", "")
|
||||
MUSIC_MEDIA_ROOT = os.getenv("MUSIC_MEDIA_ROOT", "/mnt/nas/webpage/data/music")
|
||||
MUSIC_MEDIA_BASE = os.getenv("MUSIC_MEDIA_URL_PREFIX", "/media/music")
|
||||
|
||||
POLL_INTERVAL = 8
|
||||
POLL_MAX_ATTEMPTS = 40
|
||||
|
||||
|
||||
def _headers() -> dict:
|
||||
return {
|
||||
"Authorization": f"Bearer {SUNO_API_KEY}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
|
||||
def _build_suno_payload(params: dict) -> dict:
|
||||
"""프론트엔드 params → sunoapi.org 요청 형식 (NAS 코드 그대로 이식)."""
|
||||
instrumental = params.get("instrumental", False)
|
||||
has_lyrics = bool(params.get("lyrics"))
|
||||
custom_mode = has_lyrics or bool(params.get("genre")) or bool(params.get("moods"))
|
||||
|
||||
payload = {
|
||||
"customMode": custom_mode,
|
||||
"instrumental": instrumental,
|
||||
"model": params.get("model", "V4"),
|
||||
"callBackUrl": "https://example.com/noop",
|
||||
}
|
||||
|
||||
if custom_mode:
|
||||
if instrumental:
|
||||
payload["prompt"] = ""
|
||||
elif has_lyrics:
|
||||
payload["prompt"] = params["lyrics"][:3000]
|
||||
else:
|
||||
prompt_text = params.get("prompt", "")
|
||||
payload["prompt"] = prompt_text[:3000] if prompt_text else ""
|
||||
|
||||
style_parts = []
|
||||
if params.get("genre"):
|
||||
style_parts.append(params["genre"])
|
||||
if params.get("moods"):
|
||||
style_parts.extend(params["moods"])
|
||||
if params.get("instruments"):
|
||||
style_parts.extend(params["instruments"][:3])
|
||||
if style_parts:
|
||||
payload["style"] = ", ".join(style_parts)[:200]
|
||||
|
||||
if params.get("title"):
|
||||
payload["title"] = params["title"][:80]
|
||||
else:
|
||||
parts = []
|
||||
if params.get("prompt"):
|
||||
parts.append(params["prompt"])
|
||||
if params.get("genre"):
|
||||
parts.append(params["genre"])
|
||||
if params.get("moods"):
|
||||
parts.append(", ".join(params["moods"]))
|
||||
payload["prompt"] = " ".join(parts)[:500] if parts else "instrumental music"
|
||||
|
||||
if params.get("vocal_gender"):
|
||||
payload["vocalGender"] = params["vocal_gender"]
|
||||
if params.get("negative_tags"):
|
||||
payload["negativeTags"] = params["negative_tags"]
|
||||
if params.get("style_weight") is not None:
|
||||
payload["styleWeight"] = params["style_weight"]
|
||||
if params.get("audio_weight") is not None:
|
||||
payload["audioWeight"] = params["audio_weight"]
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def _poll_suno_record(
|
||||
record_info_path: str,
|
||||
suno_task_id: str,
|
||||
task_id: str,
|
||||
max_attempts: int = POLL_MAX_ATTEMPTS,
|
||||
interval: int = POLL_INTERVAL,
|
||||
progress_msg_map: dict = None,
|
||||
) -> Optional[dict]:
|
||||
"""범용 Suno 작업 폴링. SUCCESS 시 response 객체 반환."""
|
||||
error_statuses = {
|
||||
"CREATE_TASK_FAILED", "GENERATE_AUDIO_FAILED",
|
||||
"CALLBACK_EXCEPTION", "SENSITIVE_WORD_ERROR",
|
||||
}
|
||||
default_msgs = {
|
||||
"PENDING": "대기열에서 대기 중...",
|
||||
"TEXT_SUCCESS": "가사 생성 완료, 음악 생성 중...",
|
||||
"FIRST_SUCCESS": "첫 번째 트랙 완료, 두 번째 생성 중...",
|
||||
"GENERATING": "생성 중...",
|
||||
}
|
||||
msgs = {**default_msgs, **(progress_msg_map or {})}
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
time.sleep(interval)
|
||||
try:
|
||||
resp = requests.get(
|
||||
f"{SUNO_BASE_URL}{record_info_path}",
|
||||
headers=_headers(),
|
||||
params={"taskId": suno_task_id},
|
||||
timeout=15,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
continue
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
continue
|
||||
data = body.get("data", {})
|
||||
status = data.get("status", "")
|
||||
progress = min(15 + int((attempt / max_attempts) * 65), 79)
|
||||
|
||||
if status == "SUCCESS":
|
||||
return data.get("response", data)
|
||||
elif status in error_statuses:
|
||||
error_msg = data.get("errorMessage") or data.get("msg") or f"Suno 작업 실패 ({status})"
|
||||
webhook_update_task(task_id, "failed", 0, "", error=error_msg)
|
||||
return None
|
||||
else:
|
||||
msg = msgs.get(status, f"처리 중... ({status})")
|
||||
if status == "FIRST_SUCCESS":
|
||||
progress = max(progress, 60)
|
||||
webhook_update_task(task_id, "processing", progress, msg)
|
||||
except Exception as e:
|
||||
logger.warning("Suno poll error (attempt %d): %s", attempt, e)
|
||||
continue
|
||||
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Suno 작업 타임아웃")
|
||||
return None
|
||||
|
||||
|
||||
def _download_and_register(
|
||||
task_id: str, song: dict, params: dict, filename_suffix: str = "",
|
||||
) -> Optional[dict]:
|
||||
"""Suno CDN에서 MP3 다운로드 → /mnt/nas/...에 직접 저장 → webhook으로 add_track."""
|
||||
audio_url_remote = song.get("audioUrl") or song.get("audio_url", "")
|
||||
if not audio_url_remote:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Suno 응답에 audioUrl이 없습니다")
|
||||
return None
|
||||
|
||||
filename = f"{task_id}{filename_suffix}.mp3"
|
||||
os.makedirs(MUSIC_MEDIA_ROOT, exist_ok=True)
|
||||
file_path = os.path.join(MUSIC_MEDIA_ROOT, filename)
|
||||
|
||||
try:
|
||||
dl = requests.get(audio_url_remote, timeout=120, stream=True)
|
||||
dl.raise_for_status()
|
||||
with open(file_path, "wb") as f:
|
||||
for chunk in dl.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
except Exception as e:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"오디오 다운로드 실패: {e}")
|
||||
return None
|
||||
|
||||
local_audio_url = f"{MUSIC_MEDIA_BASE}/{filename}"
|
||||
|
||||
genre = params.get("genre", song.get("tags", ""))
|
||||
moods = params.get("moods", [])
|
||||
mood_str = moods[0] if moods else "Original"
|
||||
title = (
|
||||
song.get("title")
|
||||
or params.get("title")
|
||||
or (f"{genre} — {mood_str} Mix" if genre else f"{mood_str} Mix")
|
||||
)
|
||||
|
||||
track_data = {
|
||||
"title": title,
|
||||
"genre": genre,
|
||||
"moods": moods,
|
||||
"instruments": params.get("instruments", []),
|
||||
"duration_sec": int(song["duration"]) if song.get("duration") else params.get("duration_sec"),
|
||||
"bpm": params.get("bpm"),
|
||||
"key": params.get("key", ""),
|
||||
"scale": params.get("scale", ""),
|
||||
"prompt": song.get("prompt", params.get("prompt", "")),
|
||||
"audio_url": local_audio_url,
|
||||
# NAS file_path는 NAS 관점 — /app/data 안의 경로
|
||||
"file_path": f"/app/data/{filename}",
|
||||
"task_id": task_id,
|
||||
"provider": "suno",
|
||||
"lyrics": song.get("prompt", params.get("lyrics", "")),
|
||||
"image_url": song.get("imageUrl") or song.get("image_url", ""),
|
||||
"suno_id": song.get("id", ""),
|
||||
}
|
||||
return track_data
|
||||
|
||||
|
||||
def run_suno_generation(task_id: str, params: dict) -> None:
|
||||
"""BackgroundTask: Suno API로 곡 생성 → MP3 → NAS SMB 저장 → webhook add_track."""
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정 (Windows .env)")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 5, "Suno API에 연결 중...")
|
||||
payload = _build_suno_payload(params)
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/generate", headers=_headers(), json=payload, timeout=30)
|
||||
|
||||
if resp.status_code != 200:
|
||||
err = resp.text[:300] if resp.text else f"HTTP {resp.status_code}"
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Suno API 오류: {err}")
|
||||
return
|
||||
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Suno API 거부: {body.get('msg', '?')}")
|
||||
return
|
||||
|
||||
suno_task_id = body.get("data", {}).get("taskId", "")
|
||||
if not suno_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Suno 응답에 taskId 없음")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 15, "곡 생성 대기열에 등록됨...")
|
||||
|
||||
response = _poll_suno_record("/generate/record-info", suno_task_id, task_id)
|
||||
if not response:
|
||||
return
|
||||
completed = response.get("sunoData") or []
|
||||
if not completed:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Suno 완료했으나 트랙 데이터 없음")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 80, "오디오 파일 다운로드 중...")
|
||||
track = _download_and_register(task_id, completed[0], params)
|
||||
if not track:
|
||||
return
|
||||
|
||||
webhook_add_track(task_id, "succeeded", 100, "생성 완료",
|
||||
audio_url=track["audio_url"], track=track)
|
||||
|
||||
if len(completed) > 1:
|
||||
try:
|
||||
# 보조 변형은 SMB에 파일만 저장. NAS _sync_library_with_disk가 다음
|
||||
# GET /api/music/library 호출 시 자동으로 라이브러리에 등록.
|
||||
_download_and_register(f"{task_id}_v2", completed[1], params)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
except requests.Timeout:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Suno API 타임아웃")
|
||||
except Exception as e:
|
||||
logger.exception("Suno generation error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
|
||||
|
||||
def run_suno_extend(task_id: str, params: dict) -> None:
|
||||
"""기존 곡을 특정 지점부터 연장."""
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정")
|
||||
return
|
||||
webhook_update_task(task_id, "processing", 5, "곡 연장 요청 중...")
|
||||
payload = {
|
||||
"audioId": params["suno_id"],
|
||||
"defaultParamFlag": not bool(params.get("prompt")),
|
||||
"prompt": params.get("prompt", ""),
|
||||
"continueAt": params.get("continue_at", 0),
|
||||
"model": params.get("model", "V4"),
|
||||
"callBackUrl": "https://example.com/noop",
|
||||
}
|
||||
if params.get("style"):
|
||||
payload["style"] = params["style"]
|
||||
if params.get("title"):
|
||||
payload["title"] = params["title"]
|
||||
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/generate/extend", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Suno Extend 오류: {resp.text[:300]}")
|
||||
return
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Extend 거부: {body.get('msg', '?')}")
|
||||
return
|
||||
suno_task_id = body.get("data", {}).get("taskId", "")
|
||||
if not suno_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Extend 응답에 taskId 없음")
|
||||
return
|
||||
webhook_update_task(task_id, "processing", 15, "곡 연장 대기열에 등록됨...")
|
||||
|
||||
response = _poll_suno_record("/generate/record-info", suno_task_id, task_id)
|
||||
if not response:
|
||||
return
|
||||
completed = response.get("sunoData") or []
|
||||
if not completed:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="연장 완료했으나 트랙 없음")
|
||||
return
|
||||
webhook_update_task(task_id, "processing", 80, "연장된 오디오 다운로드 중...")
|
||||
track = _download_and_register(task_id, completed[0], params)
|
||||
if track:
|
||||
webhook_add_track(task_id, "succeeded", 100, "곡 연장 완료",
|
||||
audio_url=track["audio_url"], track=track)
|
||||
except Exception as e:
|
||||
logger.exception("Suno extend error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
|
||||
|
||||
def run_vocal_removal(task_id: str, params: dict) -> None:
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정")
|
||||
return
|
||||
webhook_update_task(task_id, "processing", 5, "보컬 분리 요청 중...")
|
||||
payload = {"audioId": params["suno_id"], "callBackUrl": "https://example.com/noop"}
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/vocal-removal/generate", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Vocal Removal 오류: {resp.text[:300]}")
|
||||
return
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Vocal Removal 거부: {body.get('msg', '?')}")
|
||||
return
|
||||
suno_task_id = body.get("data", {}).get("taskId", "")
|
||||
if not suno_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="응답에 taskId 없음")
|
||||
return
|
||||
webhook_update_task(task_id, "processing", 15, "보컬 분리 처리 중...")
|
||||
response = _poll_suno_record("/vocal-removal/record-info", suno_task_id, task_id)
|
||||
if not response:
|
||||
return
|
||||
completed = response.get("sunoData") or []
|
||||
if not completed:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="분리 완료했으나 트랙 없음")
|
||||
return
|
||||
webhook_update_task(task_id, "processing", 80, "분리된 오디오 다운로드 중...")
|
||||
vp = {**params, "title": f"{params.get('title', 'Track')} (Vocals)"}
|
||||
track = _download_and_register(task_id, completed[0], vp)
|
||||
if len(completed) > 1:
|
||||
ip = {**params, "title": f"{params.get('title', 'Track')} (Instrumental)"}
|
||||
# Instrumental 변형은 SMB에 파일만 저장. NAS _sync_library_with_disk가 자동 등록.
|
||||
_download_and_register(f"{task_id}_inst", completed[1], ip)
|
||||
if track:
|
||||
webhook_add_track(task_id, "succeeded", 100, "보컬 분리 완료",
|
||||
audio_url=track["audio_url"], track=track)
|
||||
except Exception as e:
|
||||
logger.exception("vocal removal error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
|
||||
|
||||
def run_cover_image(task_id: str, params: dict) -> None:
|
||||
"""Suno 곡의 커버 이미지 2장 (URL JSON 반환)."""
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정"); return
|
||||
webhook_update_task(task_id, "processing", 5, "커버 이미지 생성 요청 중...")
|
||||
suno_task_id = params.get("suno_task_id", "")
|
||||
if not suno_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="suno_task_id 필요"); return
|
||||
payload = {"taskId": suno_task_id, "callBackUrl": "https://example.com/noop"}
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/suno/cover/generate", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Cover API 오류: {resp.text[:300]}"); return
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Cover 거부: {body.get('msg', '?')}"); return
|
||||
cover_task_id = body.get("data", {}).get("taskId", suno_task_id)
|
||||
webhook_update_task(task_id, "processing", 15, "커버 이미지 생성 중...")
|
||||
response = _poll_suno_record(
|
||||
"/suno/cover/record-info", cover_task_id, task_id,
|
||||
max_attempts=30, interval=5,
|
||||
progress_msg_map={"PENDING": "이미지 생성 대기 중...", "GENERATING": "이미지 생성 중..."},
|
||||
)
|
||||
if not response:
|
||||
return
|
||||
images = response.get("images") or response.get("sunoData") or []
|
||||
urls = []
|
||||
if isinstance(images, list):
|
||||
for img in images:
|
||||
if isinstance(img, str):
|
||||
urls.append(img)
|
||||
elif isinstance(img, dict):
|
||||
urls.append(img.get("imageUrl") or img.get("image_url", ""))
|
||||
webhook_update_task(task_id, "succeeded", 100, "커버 완료",
|
||||
audio_url=json.dumps(urls))
|
||||
except Exception as e:
|
||||
logger.exception("cover image error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
|
||||
|
||||
def run_wav_convert(task_id: str, params: dict) -> None:
|
||||
"""곡을 WAV 포맷으로 변환 (URL만)."""
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정"); return
|
||||
webhook_update_task(task_id, "processing", 5, "WAV 변환 요청 중...")
|
||||
payload = {
|
||||
"taskId": params["suno_task_id"],
|
||||
"audioId": params["suno_id"],
|
||||
"callBackUrl": "https://example.com/noop",
|
||||
}
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/wav/generate", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code == 409:
|
||||
body = resp.json()
|
||||
wav_url = body.get("data", {}).get("audioWavUrl", "")
|
||||
if wav_url:
|
||||
webhook_update_task(task_id, "succeeded", 100, "WAV 캐시", audio_url=wav_url)
|
||||
return
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"WAV 오류: {resp.text[:300]}"); return
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"WAV 거부: {body.get('msg', '?')}"); return
|
||||
wav_task_id = body.get("data", {}).get("taskId", params["suno_task_id"])
|
||||
webhook_update_task(task_id, "processing", 15, "WAV 변환 처리 중...")
|
||||
response = _poll_suno_record(
|
||||
"/wav/record-info", wav_task_id, task_id,
|
||||
max_attempts=30, interval=5,
|
||||
progress_msg_map={"PENDING": "WAV 대기 중...", "GENERATING": "WAV 변환 중..."},
|
||||
)
|
||||
if not response:
|
||||
return
|
||||
wav_url = ""
|
||||
sd = response.get("sunoData") or []
|
||||
if sd and isinstance(sd, list) and isinstance(sd[0], dict):
|
||||
wav_url = sd[0].get("audioWavUrl", "")
|
||||
if not wav_url:
|
||||
wav_url = response.get("audioWavUrl", "")
|
||||
webhook_update_task(task_id, "succeeded", 100, "WAV 변환 완료", audio_url=wav_url)
|
||||
except Exception as e:
|
||||
logger.exception("wav convert error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
|
||||
|
||||
def run_stem_split(task_id: str, params: dict) -> None:
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정"); return
|
||||
webhook_update_task(task_id, "processing", 5, "12스템 분리 요청 중...")
|
||||
payload = {
|
||||
"taskId": params["suno_task_id"],
|
||||
"audioId": params["suno_id"],
|
||||
"type": "split_stem",
|
||||
"callBackUrl": "https://example.com/noop",
|
||||
}
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/vocal-removal/generate", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Stem API 오류: {resp.text[:300]}"); return
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Stem 거부: {body.get('msg', '?')}"); return
|
||||
stem_task_id = body.get("data", {}).get("taskId", "")
|
||||
if not stem_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="응답에 taskId 없음"); return
|
||||
webhook_update_task(task_id, "processing", 15, "12스템 분리 처리 중 (2~3분)...")
|
||||
response = _poll_suno_record(
|
||||
"/vocal-removal/record-info", stem_task_id, task_id,
|
||||
max_attempts=40, interval=8,
|
||||
progress_msg_map={"PENDING": "스템 대기 중...", "GENERATING": "스템 분리 중..."},
|
||||
)
|
||||
if not response:
|
||||
return
|
||||
sd = response.get("sunoData") or []
|
||||
stems = {}
|
||||
names = ["vocal", "backing_vocals", "drums", "bass", "guitar", "keyboard",
|
||||
"strings", "brass", "woodwinds", "percussion", "synth", "fx"]
|
||||
for i, item in enumerate(sd):
|
||||
if isinstance(item, dict):
|
||||
nm = names[i] if i < len(names) else f"stem_{i}"
|
||||
stems[nm] = item.get("audioUrl") or item.get("audio_url", "")
|
||||
webhook_update_task(task_id, "succeeded", 100, "12스템 완료",
|
||||
audio_url=json.dumps(stems))
|
||||
except Exception as e:
|
||||
logger.exception("stem split error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
|
||||
|
||||
def run_upload_cover(task_id: str, params: dict) -> None:
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정"); return
|
||||
webhook_update_task(task_id, "processing", 5, "AI Cover 요청 중...")
|
||||
payload = {
|
||||
"uploadUrl": params["upload_url"],
|
||||
"customMode": params.get("custom_mode", True),
|
||||
"instrumental": params.get("instrumental", False),
|
||||
"model": params.get("model", "V4"),
|
||||
"callBackUrl": "https://example.com/noop",
|
||||
}
|
||||
for k, ak in [("prompt", "prompt"), ("style", "style"), ("title", "title"),
|
||||
("vocal_gender", "vocalGender"), ("negative_tags", "negativeTags"),
|
||||
("style_weight", "styleWeight"), ("audio_weight", "audioWeight")]:
|
||||
if params.get(k):
|
||||
payload[ak] = params[k]
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/generate/upload-cover", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Upload Cover 오류: {resp.text[:300]}"); return
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Upload Cover 거부: {body.get('msg', '?')}"); return
|
||||
suno_task_id = body.get("data", {}).get("taskId", "")
|
||||
if not suno_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="응답에 taskId 없음"); return
|
||||
webhook_update_task(task_id, "processing", 15, "AI Cover 생성 중...")
|
||||
response = _poll_suno_record("/generate/record-info", suno_task_id, task_id)
|
||||
if not response:
|
||||
return
|
||||
completed = response.get("sunoData") or []
|
||||
if not completed:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Cover 완료했으나 트랙 없음"); return
|
||||
track = _download_and_register(task_id, completed[0], params)
|
||||
if track:
|
||||
webhook_add_track(task_id, "succeeded", 100, "AI Cover 완료",
|
||||
audio_url=track["audio_url"], track=track)
|
||||
except Exception as e:
|
||||
logger.exception("upload cover error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
|
||||
|
||||
def run_upload_extend(task_id: str, params: dict) -> None:
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정"); return
|
||||
webhook_update_task(task_id, "processing", 5, "Upload Extend 요청 중...")
|
||||
payload = {
|
||||
"uploadUrl": params["upload_url"],
|
||||
"defaultParamFlag": params.get("default_param_flag", True),
|
||||
"model": params.get("model", "V4"),
|
||||
"callBackUrl": "https://example.com/noop",
|
||||
}
|
||||
for k, ak in [("prompt", "prompt"), ("style", "style"), ("title", "title"),
|
||||
("continue_at", "continueAt"), ("instrumental", "instrumental"),
|
||||
("vocal_gender", "vocalGender"), ("negative_tags", "negativeTags")]:
|
||||
if params.get(k) is not None:
|
||||
payload[ak] = params[k]
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/generate/upload-extend", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Upload Extend 오류: {resp.text[:300]}"); return
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Upload Extend 거부: {body.get('msg', '?')}"); return
|
||||
suno_task_id = body.get("data", {}).get("taskId", "")
|
||||
if not suno_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="응답에 taskId 없음"); return
|
||||
webhook_update_task(task_id, "processing", 15, "Upload Extend 생성 중...")
|
||||
response = _poll_suno_record("/generate/record-info", suno_task_id, task_id)
|
||||
if not response:
|
||||
return
|
||||
completed = response.get("sunoData") or []
|
||||
if not completed:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Upload Extend 완료했으나 트랙 없음"); return
|
||||
track = _download_and_register(task_id, completed[0], params)
|
||||
if track:
|
||||
webhook_add_track(task_id, "succeeded", 100, "Upload Extend 완료",
|
||||
audio_url=track["audio_url"], track=track)
|
||||
except Exception as e:
|
||||
logger.exception("upload extend error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
|
||||
|
||||
def run_add_vocals(task_id: str, params: dict) -> None:
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정"); return
|
||||
webhook_update_task(task_id, "processing", 5, "보컬 추가 요청 중...")
|
||||
payload = {
|
||||
"uploadUrl": params["upload_url"],
|
||||
"prompt": params.get("prompt", ""),
|
||||
"title": params.get("title", ""),
|
||||
"style": params.get("style", ""),
|
||||
"negativeTags": params.get("negative_tags", ""),
|
||||
"callBackUrl": "https://example.com/noop",
|
||||
}
|
||||
for k, ak in [("vocal_gender", "vocalGender"), ("model", "model"),
|
||||
("style_weight", "styleWeight"), ("audio_weight", "audioWeight")]:
|
||||
if params.get(k) is not None:
|
||||
payload[ak] = params[k]
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/generate/add-vocals", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Add Vocals 오류: {resp.text[:300]}"); return
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Add Vocals 거부: {body.get('msg', '?')}"); return
|
||||
suno_task_id = body.get("data", {}).get("taskId", "")
|
||||
if not suno_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="응답에 taskId 없음"); return
|
||||
webhook_update_task(task_id, "processing", 15, "AI 보컬 생성 중...")
|
||||
response = _poll_suno_record("/generate/record-info", suno_task_id, task_id)
|
||||
if not response:
|
||||
return
|
||||
completed = response.get("sunoData") or []
|
||||
if not completed:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="보컬 추가 완료했으나 트랙 없음"); return
|
||||
track = _download_and_register(task_id, completed[0], params)
|
||||
if track:
|
||||
webhook_add_track(task_id, "succeeded", 100, "보컬 추가 완료",
|
||||
audio_url=track["audio_url"], track=track)
|
||||
except Exception as e:
|
||||
logger.exception("add vocals error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
|
||||
|
||||
def run_add_instrumental(task_id: str, params: dict) -> None:
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정"); return
|
||||
webhook_update_task(task_id, "processing", 5, "인스트루멘탈 추가 요청 중...")
|
||||
payload = {
|
||||
"uploadUrl": params["upload_url"],
|
||||
"title": params.get("title", ""),
|
||||
"tags": params.get("tags", ""),
|
||||
"negativeTags": params.get("negative_tags", ""),
|
||||
"callBackUrl": "https://example.com/noop",
|
||||
}
|
||||
for k, ak in [("vocal_gender", "vocalGender"), ("model", "model"),
|
||||
("style_weight", "styleWeight"), ("audio_weight", "audioWeight")]:
|
||||
if params.get(k) is not None:
|
||||
payload[ak] = params[k]
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/generate/add-instrumental", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Add Inst 오류: {resp.text[:300]}"); return
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Add Inst 거부: {body.get('msg', '?')}"); return
|
||||
suno_task_id = body.get("data", {}).get("taskId", "")
|
||||
if not suno_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="응답에 taskId 없음"); return
|
||||
webhook_update_task(task_id, "processing", 15, "AI 반주 생성 중...")
|
||||
response = _poll_suno_record("/generate/record-info", suno_task_id, task_id)
|
||||
if not response:
|
||||
return
|
||||
completed = response.get("sunoData") or []
|
||||
if not completed:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Add Inst 완료했으나 트랙 없음"); return
|
||||
track = _download_and_register(task_id, completed[0], params)
|
||||
if track:
|
||||
webhook_add_track(task_id, "succeeded", 100, "Add Instrumental 완료",
|
||||
audio_url=track["audio_url"], track=track)
|
||||
except Exception as e:
|
||||
logger.exception("add instrumental error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
|
||||
|
||||
def run_video_generate(task_id: str, params: dict) -> None:
|
||||
try:
|
||||
if not SUNO_API_KEY:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SUNO_API_KEY 미설정"); return
|
||||
webhook_update_task(task_id, "processing", 5, "뮤직비디오 생성 요청 중...")
|
||||
payload = {
|
||||
"taskId": params["suno_task_id"],
|
||||
"audioId": params["suno_id"],
|
||||
"callBackUrl": "https://example.com/noop",
|
||||
}
|
||||
if params.get("author"):
|
||||
payload["author"] = params["author"][:50]
|
||||
if params.get("domain_name"):
|
||||
payload["domainName"] = params["domain_name"][:50]
|
||||
resp = requests.post(f"{SUNO_BASE_URL}/mp4/generate", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Video 오류: {resp.text[:300]}"); return
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Video 거부: {body.get('msg', '?')}"); return
|
||||
video_task_id = body.get("data", {}).get("taskId", params.get("suno_task_id", ""))
|
||||
webhook_update_task(task_id, "processing", 15, "뮤직비디오 렌더링 중...")
|
||||
response = _poll_suno_record(
|
||||
"/mp4/record-info", video_task_id, task_id,
|
||||
max_attempts=60, interval=10,
|
||||
progress_msg_map={"PENDING": "비디오 대기 중...", "GENERATING": "비디오 렌더링 중..."},
|
||||
)
|
||||
if not response:
|
||||
return
|
||||
video_url = ""
|
||||
sd = response.get("sunoData") or []
|
||||
if sd and isinstance(sd, list) and isinstance(sd[0], dict):
|
||||
video_url = sd[0].get("videoUrl") or sd[0].get("video_url", "")
|
||||
if not video_url:
|
||||
video_url = response.get("video_url") or response.get("videoUrl", "")
|
||||
webhook_update_task(task_id, "succeeded", 100, "뮤직비디오 완료", audio_url=video_url)
|
||||
except Exception as e:
|
||||
logger.exception("video generate error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
131
services/music-render/providers/sync_ops.py
Normal file
131
services/music-render/providers/sync_ops.py
Normal file
@@ -0,0 +1,131 @@
|
||||
"""Sync Suno API helpers — main.py FastAPI sync endpoints에서 호출.
|
||||
|
||||
NAS music-lab/app/suno_provider.py의 sync 함수들 이식.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SUNO_BASE_URL = "https://api.sunoapi.org/api/v1"
|
||||
SUNO_API_KEY = os.getenv("SUNO_API_KEY", "")
|
||||
|
||||
|
||||
def _headers() -> dict:
|
||||
return {
|
||||
"Authorization": f"Bearer {SUNO_API_KEY}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
|
||||
def generate_lyrics(prompt: str) -> Optional[dict]:
|
||||
"""Suno 가사 생성 API — 폴링 결과 반환."""
|
||||
if not SUNO_API_KEY:
|
||||
return None
|
||||
try:
|
||||
resp = requests.post(
|
||||
f"{SUNO_BASE_URL}/lyrics",
|
||||
headers=_headers(),
|
||||
json={"prompt": prompt[:200]},
|
||||
timeout=30,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
return None
|
||||
body = resp.json()
|
||||
if body.get("code") != 200:
|
||||
return body
|
||||
task_id = body.get("data", {}).get("taskId", "")
|
||||
if not task_id:
|
||||
return body
|
||||
return _poll_lyrics(task_id)
|
||||
except Exception as e:
|
||||
logger.warning("Suno lyrics API error: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def _poll_lyrics(lyrics_task_id: str) -> Optional[dict]:
|
||||
for _ in range(15):
|
||||
time.sleep(3)
|
||||
try:
|
||||
resp = requests.get(
|
||||
f"{SUNO_BASE_URL}/lyrics/record-info",
|
||||
headers=_headers(),
|
||||
params={"taskId": lyrics_task_id},
|
||||
timeout=15,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
continue
|
||||
body = resp.json()
|
||||
data = body.get("data", {})
|
||||
if data.get("status") == "complete":
|
||||
items = data.get("data") or data.get("sunoData") or []
|
||||
if items and isinstance(items, list):
|
||||
return {
|
||||
"id": lyrics_task_id,
|
||||
"status": "complete",
|
||||
"text": items[0].get("text", ""),
|
||||
"title": items[0].get("title", ""),
|
||||
}
|
||||
return {"id": lyrics_task_id, "status": "complete", "text": ""}
|
||||
except Exception:
|
||||
continue
|
||||
return None
|
||||
|
||||
|
||||
def get_credits() -> Optional[dict]:
|
||||
if not SUNO_API_KEY:
|
||||
return None
|
||||
for path in ["/generate/credit", "/get-credits"]:
|
||||
try:
|
||||
resp = requests.get(f"{SUNO_BASE_URL}{path}", headers=_headers(), timeout=15)
|
||||
if resp.status_code == 200:
|
||||
body = resp.json()
|
||||
data = body.get("data", body)
|
||||
if isinstance(data, (int, float)):
|
||||
return {"credits_left": int(data)}
|
||||
return data
|
||||
except Exception as e:
|
||||
logger.warning("Suno credits API error (%s): %s", path, e)
|
||||
return None
|
||||
|
||||
|
||||
def get_timestamped_lyrics(suno_task_id: str, suno_id: str) -> Optional[dict]:
|
||||
if not SUNO_API_KEY:
|
||||
return None
|
||||
try:
|
||||
resp = requests.post(
|
||||
f"{SUNO_BASE_URL}/generate/get-timestamped-lyrics",
|
||||
headers=_headers(),
|
||||
json={"taskId": suno_task_id, "audioId": suno_id},
|
||||
timeout=30,
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
body = resp.json()
|
||||
return body.get("data", body)
|
||||
except Exception as e:
|
||||
logger.warning("Timestamped lyrics error: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def generate_style_boost(content: str) -> Optional[dict]:
|
||||
if not SUNO_API_KEY:
|
||||
return None
|
||||
try:
|
||||
resp = requests.post(
|
||||
f"{SUNO_BASE_URL}/style/generate",
|
||||
headers=_headers(),
|
||||
json={"content": content},
|
||||
timeout=30,
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
body = resp.json()
|
||||
return body.get("data", body)
|
||||
except Exception as e:
|
||||
logger.warning("Style boost error: %s", e)
|
||||
return None
|
||||
9
services/music-render/requirements.txt
Normal file
9
services/music-render/requirements.txt
Normal file
@@ -0,0 +1,9 @@
|
||||
fastapi==0.115.6
|
||||
uvicorn[standard]==0.34.0
|
||||
requests==2.32.3
|
||||
redis>=5.0
|
||||
httpx>=0.27
|
||||
mutagen==1.47.0
|
||||
pytest>=8.0
|
||||
pytest-asyncio>=0.24
|
||||
respx>=0.21
|
||||
80
services/music-render/tests/test_nas_client.py
Normal file
80
services/music-render/tests/test_nas_client.py
Normal file
@@ -0,0 +1,80 @@
|
||||
"""nas_client — webhook adapter tests."""
|
||||
import os
|
||||
import pytest
|
||||
import respx
|
||||
import httpx
|
||||
|
||||
from nas_client import webhook_update_task, webhook_add_track
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _env(monkeypatch):
|
||||
monkeypatch.setenv("NAS_BASE_URL", "http://nas-test:18600")
|
||||
monkeypatch.setenv("INTERNAL_API_KEY", "test-key")
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_webhook_update_task_sends_x_internal_key():
|
||||
route = respx.post("http://nas-test:18600/api/internal/music/update").mock(
|
||||
return_value=httpx.Response(200, json={"ok": True})
|
||||
)
|
||||
webhook_update_task("task-1", "processing", 30, message="downloading")
|
||||
assert route.called
|
||||
req = route.calls[0].request
|
||||
assert req.headers["X-Internal-Key"] == "test-key"
|
||||
import json
|
||||
body = json.loads(req.content)
|
||||
assert body["task_id"] == "task-1"
|
||||
assert body["status"] == "processing"
|
||||
assert body["progress"] == 30
|
||||
assert body["message"] == "downloading"
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_webhook_update_task_with_audio_url():
|
||||
route = respx.post("http://nas-test:18600/api/internal/music/update").mock(
|
||||
return_value=httpx.Response(200, json={"ok": True})
|
||||
)
|
||||
webhook_update_task("task-2", "succeeded", 100, message="완료",
|
||||
audio_url="/media/music/task-2.mp3")
|
||||
import json
|
||||
payload = json.loads(route.calls[0].request.content)
|
||||
assert payload["audio_url"] == "/media/music/task-2.mp3"
|
||||
assert payload["status"] == "succeeded"
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_webhook_update_task_with_error():
|
||||
route = respx.post("http://nas-test:18600/api/internal/music/update").mock(
|
||||
return_value=httpx.Response(200, json={"ok": True})
|
||||
)
|
||||
webhook_update_task("task-3", "failed", 0, error="API rate limit")
|
||||
import json
|
||||
payload = json.loads(route.calls[0].request.content)
|
||||
assert payload["error"] == "API rate limit"
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_webhook_add_track_uses_track_field():
|
||||
"""add_track은 update와 동시에 (succeeded 시)."""
|
||||
route = respx.post("http://nas-test:18600/api/internal/music/update").mock(
|
||||
return_value=httpx.Response(200, json={"ok": True})
|
||||
)
|
||||
track = {"title": "x", "audio_url": "/media/music/t.mp3", "provider": "suno"}
|
||||
webhook_add_track("task-4", "succeeded", 100, message="ok",
|
||||
audio_url="/media/music/t.mp3", track=track)
|
||||
import json
|
||||
payload = json.loads(route.calls[0].request.content)
|
||||
assert payload["track"]["title"] == "x"
|
||||
assert payload["status"] == "succeeded"
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_webhook_swallows_network_error(caplog):
|
||||
"""webhook 실패해도 raise 안 함 (logger.error)."""
|
||||
respx.post("http://nas-test:18600/api/internal/music/update").mock(
|
||||
side_effect=httpx.ConnectError("no host")
|
||||
)
|
||||
# raise 안 하면 통과
|
||||
webhook_update_task("task-5", "processing", 10)
|
||||
assert "task-5" in caplog.text
|
||||
32
services/music-render/tests/test_suno_provider.py
Normal file
32
services/music-render/tests/test_suno_provider.py
Normal file
@@ -0,0 +1,32 @@
|
||||
"""providers/suno.py — _build_suno_payload 단위 테스트 + 1개 함수 mock 검증."""
|
||||
import pytest
|
||||
from providers.suno import _build_suno_payload
|
||||
|
||||
|
||||
def test_payload_custom_mode_with_lyrics():
|
||||
params = {"lyrics": "[Verse]\nhello", "genre": "lofi", "moods": ["chill"], "model": "V4"}
|
||||
p = _build_suno_payload(params)
|
||||
assert p["customMode"] is True
|
||||
assert p["prompt"] == "[Verse]\nhello"
|
||||
assert "lofi" in p["style"]
|
||||
assert "chill" in p["style"]
|
||||
|
||||
|
||||
def test_payload_simple_mode_no_lyrics_no_genre():
|
||||
params = {"prompt": "happy summer", "model": "V4"}
|
||||
p = _build_suno_payload(params)
|
||||
assert p["customMode"] is False
|
||||
assert "happy summer" in p["prompt"]
|
||||
|
||||
|
||||
def test_payload_instrumental_clears_prompt():
|
||||
params = {"genre": "ambient", "instrumental": True, "model": "V5"}
|
||||
p = _build_suno_payload(params)
|
||||
assert p["instrumental"] is True
|
||||
assert p["prompt"] == ""
|
||||
|
||||
|
||||
def test_payload_includes_optional_vocal_gender():
|
||||
params = {"genre": "pop", "vocal_gender": "f", "model": "V4"}
|
||||
p = _build_suno_payload(params)
|
||||
assert p["vocalGender"] == "f"
|
||||
169
services/music-render/tests/test_worker.py
Normal file
169
services/music-render/tests/test_worker.py
Normal file
@@ -0,0 +1,169 @@
|
||||
"""worker.py — job_type 디스패처 + paused 체크."""
|
||||
import json
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import worker
|
||||
|
||||
|
||||
def test_dispatch_suno_generation_calls_run_suno_generation():
|
||||
payload = {
|
||||
"task_id": "t1",
|
||||
"job_type": "suno_generation",
|
||||
"params": {"genre": "lofi", "title": "x"},
|
||||
}
|
||||
with patch("worker.run_suno_generation") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t1", {"genre": "lofi", "title": "x"})
|
||||
|
||||
|
||||
def test_dispatch_local_generation_calls_run_local_generation():
|
||||
payload = {
|
||||
"task_id": "t2",
|
||||
"job_type": "local_generation",
|
||||
"params": {"genre": "ambient"},
|
||||
}
|
||||
with patch("worker.run_local_generation") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t2", {"genre": "ambient"})
|
||||
|
||||
|
||||
def test_dispatch_unknown_job_type_logs_error():
|
||||
payload = {"task_id": "t3", "job_type": "weird_type", "params": {}}
|
||||
with patch("worker.webhook_update_task") as m:
|
||||
worker._dispatch(payload)
|
||||
# 알 수 없는 job_type은 failed로 보고
|
||||
m.assert_called_once()
|
||||
args = m.call_args[0]
|
||||
assert args[0] == "t3"
|
||||
assert args[1] == "failed"
|
||||
|
||||
|
||||
def test_dispatch_suno_extend_calls_run_suno_extend():
|
||||
payload = {"task_id": "t4", "job_type": "suno_extend", "params": {"suno_id": "abc"}}
|
||||
with patch("worker.run_suno_extend") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t4", {"suno_id": "abc"})
|
||||
|
||||
|
||||
def test_dispatch_vocal_removal_calls_run_vocal_removal():
|
||||
payload = {"task_id": "t5", "job_type": "vocal_removal", "params": {"suno_id": "abc"}}
|
||||
with patch("worker.run_vocal_removal") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t5", {"suno_id": "abc"})
|
||||
|
||||
|
||||
def test_dispatch_cover_image_calls_run_cover_image():
|
||||
payload = {"task_id": "t6", "job_type": "cover_image", "params": {"suno_task_id": "x"}}
|
||||
with patch("worker.run_cover_image") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t6", {"suno_task_id": "x"})
|
||||
|
||||
|
||||
def test_dispatch_wav_convert_calls_run_wav_convert():
|
||||
payload = {"task_id": "t7", "job_type": "wav_convert", "params": {"suno_task_id": "x", "suno_id": "y"}}
|
||||
with patch("worker.run_wav_convert") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t7", {"suno_task_id": "x", "suno_id": "y"})
|
||||
|
||||
|
||||
def test_dispatch_stem_split_calls_run_stem_split():
|
||||
payload = {"task_id": "t8", "job_type": "stem_split", "params": {"suno_task_id": "x", "suno_id": "y"}}
|
||||
with patch("worker.run_stem_split") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t8", {"suno_task_id": "x", "suno_id": "y"})
|
||||
|
||||
|
||||
def test_dispatch_video_generate_calls_run_video_generate():
|
||||
payload = {"task_id": "t9", "job_type": "video_generate", "params": {"suno_task_id": "x", "suno_id": "y"}}
|
||||
with patch("worker.run_video_generate") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t9", {"suno_task_id": "x", "suno_id": "y"})
|
||||
|
||||
|
||||
def test_dispatch_upload_cover_calls_run_upload_cover():
|
||||
payload = {"task_id": "t10", "job_type": "upload_cover", "params": {"upload_url": "u"}}
|
||||
with patch("worker.run_upload_cover") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t10", {"upload_url": "u"})
|
||||
|
||||
|
||||
def test_dispatch_upload_extend_calls_run_upload_extend():
|
||||
payload = {"task_id": "t11", "job_type": "upload_extend", "params": {"upload_url": "u"}}
|
||||
with patch("worker.run_upload_extend") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t11", {"upload_url": "u"})
|
||||
|
||||
|
||||
def test_dispatch_add_vocals_calls_run_add_vocals():
|
||||
payload = {"task_id": "t12", "job_type": "add_vocals", "params": {"upload_url": "u"}}
|
||||
with patch("worker.run_add_vocals") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t12", {"upload_url": "u"})
|
||||
|
||||
|
||||
def test_dispatch_add_instrumental_calls_run_add_instrumental():
|
||||
payload = {"task_id": "t13", "job_type": "add_instrumental", "params": {"upload_url": "u"}}
|
||||
with patch("worker.run_add_instrumental") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t13", {"upload_url": "u"})
|
||||
|
||||
|
||||
# ----- F6: ReliableQueue poll_once -----
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_acks_on_success(monkeypatch):
|
||||
"""F6 — _dispatch 정상 return → queue.ack(raw)."""
|
||||
payload = {"task_id": "t1", "job_type": "suno_generation", "params": {}}
|
||||
raw = json.dumps(payload).encode()
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=(payload, raw))
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
|
||||
monkeypatch.setattr(worker, "_dispatch", MagicMock())
|
||||
|
||||
handled = await worker.poll_once(fake_queue)
|
||||
assert handled is True
|
||||
fake_queue.ack.assert_awaited_once_with(raw)
|
||||
fake_queue.fail.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_calls_fail_on_dispatch_exception(monkeypatch):
|
||||
"""F6 — _dispatch unhandled exception → queue.fail(raw, payload)."""
|
||||
payload = {"task_id": "t2", "job_type": "suno_generation", "params": {}}
|
||||
raw = json.dumps(payload).encode()
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=(payload, raw))
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
|
||||
def _boom(p):
|
||||
raise RuntimeError("dispatch crash")
|
||||
|
||||
monkeypatch.setattr(worker, "_dispatch", _boom)
|
||||
|
||||
handled = await worker.poll_once(fake_queue)
|
||||
assert handled is True
|
||||
fake_queue.fail.assert_awaited_once_with(raw, payload)
|
||||
fake_queue.ack.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_returns_false_on_timeout(monkeypatch):
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=None)
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
dispatch_mock = MagicMock()
|
||||
monkeypatch.setattr(worker, "_dispatch", dispatch_mock)
|
||||
|
||||
handled = await worker.poll_once(fake_queue)
|
||||
assert handled is False
|
||||
dispatch_mock.assert_not_called()
|
||||
fake_queue.ack.assert_not_awaited()
|
||||
fake_queue.fail.assert_not_awaited()
|
||||
114
services/music-render/worker.py
Normal file
114
services/music-render/worker.py
Normal file
@@ -0,0 +1,114 @@
|
||||
"""Redis ReliableQueue worker — F6 신뢰성 패턴 (BLMOVE + ack/fail + recovery).
|
||||
|
||||
queue:paused 가 set이면 대기 (task-watcher가 박재오 활동 감지 시 set).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
from providers.suno import (
|
||||
run_suno_generation, run_suno_extend, run_vocal_removal,
|
||||
run_cover_image, run_wav_convert, run_stem_split,
|
||||
run_upload_cover, run_upload_extend, run_add_vocals,
|
||||
run_add_instrumental, run_video_generate,
|
||||
)
|
||||
from providers.local import run_local_generation
|
||||
from _shared.reliable_queue import ReliableQueue
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
REDIS_URL = os.getenv("REDIS_URL", "redis://192.168.45.54:6379")
|
||||
QUEUE_KEY = "queue:music-render"
|
||||
PAUSED_KEY = "queue:paused"
|
||||
|
||||
# Maps job_type → module-level function name (string).
|
||||
# _dispatch resolves the name via globals() at call time so unittest.mock.patch
|
||||
# on "worker.<name>" is correctly intercepted.
|
||||
_DISPATCH_TABLE: dict[str, str] = {
|
||||
"suno_generation": "run_suno_generation",
|
||||
"local_generation": "run_local_generation",
|
||||
"suno_extend": "run_suno_extend",
|
||||
"vocal_removal": "run_vocal_removal",
|
||||
"cover_image": "run_cover_image",
|
||||
"wav_convert": "run_wav_convert",
|
||||
"stem_split": "run_stem_split",
|
||||
"upload_cover": "run_upload_cover",
|
||||
"upload_extend": "run_upload_extend",
|
||||
"add_vocals": "run_add_vocals",
|
||||
"add_instrumental": "run_add_instrumental",
|
||||
"video_generate": "run_video_generate",
|
||||
}
|
||||
|
||||
|
||||
def _dispatch(payload: dict) -> None:
|
||||
"""payload[job_type] → provider 함수 호출 (sync, asyncio.to_thread로 래핑)."""
|
||||
import sys
|
||||
_self = sys.modules[__name__]
|
||||
job_type = payload.get("job_type", "")
|
||||
task_id = payload.get("task_id", "")
|
||||
params = payload.get("params", {})
|
||||
fn_name = _DISPATCH_TABLE.get(job_type)
|
||||
if fn_name is None:
|
||||
logger.error("unknown job_type=%s task=%s", job_type, task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"unknown job_type: {job_type}")
|
||||
return
|
||||
try:
|
||||
fn = getattr(_self, fn_name)
|
||||
except AttributeError:
|
||||
logger.error("dispatch table typo for job_type=%s name=%s task=%s", job_type, fn_name, task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"internal dispatch error: {fn_name}")
|
||||
return
|
||||
fn(task_id, params)
|
||||
|
||||
|
||||
async def poll_once(queue: ReliableQueue) -> bool:
|
||||
"""F6 — 1 cycle: dequeue → _dispatch → ack/fail. Returns True if a job handled."""
|
||||
result = await queue.dequeue(timeout=5)
|
||||
if result is None:
|
||||
return False
|
||||
payload, raw = result
|
||||
try:
|
||||
# sync provider 함수 — thread로 실행해서 이벤트 루프 블로킹 방지
|
||||
await asyncio.to_thread(_dispatch, payload)
|
||||
except Exception:
|
||||
logger.exception("dispatch unhandled exception task_id=%s",
|
||||
payload.get("task_id"))
|
||||
await queue.fail(raw, payload)
|
||||
return True
|
||||
await queue.ack(raw)
|
||||
return True
|
||||
|
||||
|
||||
async def worker_loop():
|
||||
redis = aioredis.from_url(REDIS_URL, decode_responses=False)
|
||||
queue = ReliableQueue(redis, queue_key=QUEUE_KEY)
|
||||
logger.info("music-render worker started worker_id=%s queue=%s",
|
||||
queue.worker_id, QUEUE_KEY)
|
||||
# F6: startup recovery
|
||||
try:
|
||||
recovered = await queue.recover()
|
||||
if recovered:
|
||||
logger.info("recovered %d orphaned items at startup", recovered)
|
||||
except Exception:
|
||||
logger.exception("startup recover failed")
|
||||
|
||||
while True:
|
||||
try:
|
||||
paused = await redis.get(PAUSED_KEY)
|
||||
if paused == b"1":
|
||||
await asyncio.sleep(10)
|
||||
continue
|
||||
await poll_once(queue)
|
||||
except asyncio.CancelledError:
|
||||
logger.info("worker_loop cancelled")
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("worker_loop iteration 실패, 5초 후 재시도")
|
||||
await asyncio.sleep(5)
|
||||
11
services/task-watcher/.env.example
Normal file
11
services/task-watcher/.env.example
Normal file
@@ -0,0 +1,11 @@
|
||||
# Plan-B-Infra — task-watcher
|
||||
|
||||
# NAS Redis
|
||||
REDIS_URL=redis://192.168.45.54:6379
|
||||
|
||||
# NAS stock holidays endpoint
|
||||
STOCK_BASE_URL=http://192.168.45.54:18500
|
||||
|
||||
# 트레이딩 윈도우 (KST, HH:MM) — 이 시간대에만 queue:paused
|
||||
TRADING_START=07:00
|
||||
TRADING_END=16:30
|
||||
16
services/task-watcher/Dockerfile
Normal file
16
services/task-watcher/Dockerfile
Normal file
@@ -0,0 +1,16 @@
|
||||
FROM python:3.12-slim-bookworm
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ca-certificates tzdata \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir --timeout 600 --retries 5 -r requirements.txt
|
||||
|
||||
COPY . .
|
||||
|
||||
EXPOSE 8000
|
||||
CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "1"]
|
||||
83
services/task-watcher/NSSM_SETUP.md
Normal file
83
services/task-watcher/NSSM_SETUP.md
Normal file
@@ -0,0 +1,83 @@
|
||||
# NSSM 자동 시작 설정 (SP-9)
|
||||
|
||||
Windows AI 머신 부팅 시 ai_trade(트레이딩) + WSL2 Docker(render workers + task-watcher) 자동 시작.
|
||||
|
||||
## 1. NSSM 다운로드
|
||||
|
||||
https://nssm.cc/download → nssm-2.24.zip → `C:\nssm\nssm.exe` 배치 (또는 PATH 등록).
|
||||
|
||||
## 2. ai_trade (Native Python, HIGH priority)
|
||||
|
||||
⚠️ spec의 signal_v2는 ai_trade로 rename됨. 경로/포트 확인.
|
||||
|
||||
```powershell
|
||||
# 관리자 PowerShell
|
||||
C:\nssm\nssm.exe install ai_trade "C:\Python312\python.exe" "-m uvicorn main:app --host 0.0.0.0 --port 8001"
|
||||
C:\nssm\nssm.exe set ai_trade AppDirectory "C:\Users\jaeoh\Desktop\workspace\web-ai\ai_trade"
|
||||
C:\nssm\nssm.exe set ai_trade Priority HIGH_PRIORITY_CLASS
|
||||
C:\nssm\nssm.exe set ai_trade Start SERVICE_AUTO_START
|
||||
C:\nssm\nssm.exe set ai_trade AppStdout "C:\Users\jaeoh\nssm-logs\ai_trade.log"
|
||||
C:\nssm\nssm.exe set ai_trade AppStderr "C:\Users\jaeoh\nssm-logs\ai_trade.log"
|
||||
```
|
||||
|
||||
(ai_trade의 실제 진입점이 main:app + port 8001인지 확인. 다르면 조정.)
|
||||
|
||||
## 3. WSL2 Docker (NORMAL priority — render workers + task-watcher)
|
||||
|
||||
```powershell
|
||||
C:\nssm\nssm.exe install wsl_docker "C:\Windows\System32\wsl.exe" "-d Ubuntu-24.04 -- sh -c 'sudo service docker start && cd /workspace/web-ai/services && docker compose up -d'"
|
||||
C:\nssm\nssm.exe set wsl_docker Priority NORMAL_PRIORITY_CLASS
|
||||
C:\nssm\nssm.exe set wsl_docker Start SERVICE_AUTO_START
|
||||
C:\nssm\nssm.exe set wsl_docker AppStdout "C:\Users\jaeoh\nssm-logs\wsl_docker.log"
|
||||
```
|
||||
|
||||
⚠️ 변경점: Ubuntu-22.04 → **Ubuntu-24.04**, web-ai-services → **web-ai/services**. WSL 경로는 박재오 WSL 마운트 기준 (`/workspace`가 web-ai에 매핑되어 있으면 그대로, 아니면 `/mnt/c/Users/jaeoh/Desktop/workspace/web-ai/services`).
|
||||
|
||||
`sudo service docker start`가 비밀번호 요구하면 sudoers에 NOPASSWD 추가:
|
||||
```bash
|
||||
# WSL2 안
|
||||
echo "$USER ALL=(ALL) NOPASSWD: /usr/sbin/service docker start" | sudo tee /etc/sudoers.d/docker-start
|
||||
```
|
||||
|
||||
## 4. 서비스 시작 + 확인
|
||||
|
||||
```powershell
|
||||
C:\nssm\nssm.exe start ai_trade
|
||||
C:\nssm\nssm.exe start wsl_docker
|
||||
|
||||
# 상태 확인
|
||||
C:\nssm\nssm.exe status ai_trade
|
||||
C:\nssm\nssm.exe status wsl_docker
|
||||
sc query ai_trade
|
||||
```
|
||||
|
||||
## 5. 검증
|
||||
|
||||
```powershell
|
||||
# ai_trade
|
||||
curl http://localhost:8001/health # 또는 ai_trade의 실제 health endpoint
|
||||
|
||||
# WSL2 docker 컨테이너 (재부팅 후 자동 시작 확인)
|
||||
wsl -d Ubuntu-24.04 -- docker ps
|
||||
# insta-render, music-render, video-render, task-watcher 4개 Up 확인
|
||||
```
|
||||
|
||||
## 6. 재부팅 테스트
|
||||
|
||||
Windows 재부팅 → 로그인 → 수동 조작 없이:
|
||||
- ai_trade 서비스 자동 시작 (HIGH priority)
|
||||
- WSL2 + Docker + 4 컨테이너 자동 시작 (NORMAL priority)
|
||||
- task-watcher가 trading window에 queue:paused 토글 시작
|
||||
|
||||
## task-watcher 동작 확인
|
||||
|
||||
```bash
|
||||
# WSL2
|
||||
docker logs task-watcher --tail 20
|
||||
# 기대: "task-watcher started" + mode 전환 로그 (trading/free)
|
||||
|
||||
# Redis 큐 상태 (NAS 또는 LAN)
|
||||
docker exec redis redis-cli GET queue:paused
|
||||
# 트레이딩 시간대(평일 07:00-16:30): "1"
|
||||
# 그 외: (nil)
|
||||
```
|
||||
36
services/task-watcher/main.py
Normal file
36
services/task-watcher/main.py
Normal file
@@ -0,0 +1,36 @@
|
||||
"""task-watcher FastAPI entry — health + lifespan (watcher loop spawn)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
import watcher
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
watcher_task = asyncio.create_task(watcher.watcher_loop())
|
||||
logger.info("task-watcher lifespan 시작")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
watcher_task.cancel()
|
||||
try:
|
||||
await watcher_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
logger.info("task-watcher lifespan 종료")
|
||||
|
||||
|
||||
app = FastAPI(lifespan=lifespan)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"ok": True, "service": "task-watcher"}
|
||||
57
services/task-watcher/mode.py
Normal file
57
services/task-watcher/mode.py
Normal file
@@ -0,0 +1,57 @@
|
||||
"""시간대 + 휴장일 기반 모드 판정 (idle 감지 생략 — 박재오 결정 2026-05-22).
|
||||
|
||||
trading: 비휴장 평일 07:00–16:30 (장중) → queue:paused SET
|
||||
free: 그 외 (장 전/후, 주말, 휴장) → queue:paused DEL
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
import logging
|
||||
import os
|
||||
from typing import Set
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
KST = ZoneInfo("Asia/Seoul")
|
||||
STOCK_BASE_URL = os.getenv("STOCK_BASE_URL", "http://192.168.45.54:18500")
|
||||
|
||||
# 트레이딩 윈도우 (HH:MM, KST). .env로 조정 가능.
|
||||
TRADING_START = os.getenv("TRADING_START", "07:00")
|
||||
TRADING_END = os.getenv("TRADING_END", "16:30")
|
||||
|
||||
|
||||
def _parse_hhmm(s: str) -> dt.time:
|
||||
hh, mm = s.split(":")
|
||||
return dt.time(int(hh), int(mm))
|
||||
|
||||
|
||||
def current_mode(now: dt.datetime, holidays: Set[str]) -> str:
|
||||
"""now(KST aware) + holidays(ISO date set) → 'trading' | 'free'."""
|
||||
# 주말 (토=5, 일=6)
|
||||
if now.weekday() >= 5:
|
||||
return "free"
|
||||
# 휴장일
|
||||
if now.date().isoformat() in holidays:
|
||||
return "free"
|
||||
# 트레이딩 윈도우 [start, end)
|
||||
start = _parse_hhmm(TRADING_START)
|
||||
end = _parse_hhmm(TRADING_END)
|
||||
t = now.timetz().replace(tzinfo=None)
|
||||
if start <= t < end:
|
||||
return "trading"
|
||||
return "free"
|
||||
|
||||
|
||||
def fetch_holidays() -> Set[str]:
|
||||
"""NAS stock /api/stock/holidays 조회. 실패 시 빈 set (안전 — free로 판정)."""
|
||||
try:
|
||||
r = httpx.get(f"{STOCK_BASE_URL}/api/stock/holidays", timeout=10.0)
|
||||
if r.status_code == 200:
|
||||
return set(r.json().get("holidays", []))
|
||||
logger.warning("holidays fetch returned %d", r.status_code)
|
||||
except Exception:
|
||||
logger.exception("holidays fetch 실패")
|
||||
return set()
|
||||
5
services/task-watcher/requirements.txt
Normal file
5
services/task-watcher/requirements.txt
Normal file
@@ -0,0 +1,5 @@
|
||||
fastapi==0.115.6
|
||||
uvicorn[standard]==0.34.0
|
||||
redis>=5.0
|
||||
httpx>=0.27
|
||||
pytest>=8.0
|
||||
0
services/task-watcher/tests/__init__.py
Normal file
0
services/task-watcher/tests/__init__.py
Normal file
44
services/task-watcher/tests/test_mode.py
Normal file
44
services/task-watcher/tests/test_mode.py
Normal file
@@ -0,0 +1,44 @@
|
||||
"""current_mode — 시간대 + 휴장일 판정 (순수 함수)."""
|
||||
import datetime as dt
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
from mode import current_mode
|
||||
|
||||
KST = ZoneInfo("Asia/Seoul")
|
||||
HOLIDAYS = {"2026-05-25"} # 가상 휴장일 (월요일)
|
||||
|
||||
|
||||
def _kst(y, m, d, hh, mm):
|
||||
return dt.datetime(y, m, d, hh, mm, tzinfo=KST)
|
||||
|
||||
|
||||
def test_weekday_trading_hours_is_trading():
|
||||
# 2026-05-22 금요일 10:00 — 트레이딩 시간대
|
||||
assert current_mode(_kst(2026, 5, 22, 10, 0), HOLIDAYS) == "trading"
|
||||
|
||||
|
||||
def test_weekday_before_open_is_free():
|
||||
# 평일 06:00 — 장 전
|
||||
assert current_mode(_kst(2026, 5, 22, 6, 0), HOLIDAYS) == "free"
|
||||
|
||||
|
||||
def test_weekday_after_close_is_free():
|
||||
# 평일 17:00 — 장 마감 후
|
||||
assert current_mode(_kst(2026, 5, 22, 17, 0), HOLIDAYS) == "free"
|
||||
|
||||
|
||||
def test_weekend_is_free():
|
||||
# 2026-05-23 토요일 10:00
|
||||
assert current_mode(_kst(2026, 5, 23, 10, 0), HOLIDAYS) == "free"
|
||||
|
||||
|
||||
def test_holiday_weekday_is_free():
|
||||
# 2026-05-25 월요일이지만 휴장일 → 트레이딩 시간대라도 free
|
||||
assert current_mode(_kst(2026, 5, 25, 10, 0), HOLIDAYS) == "free"
|
||||
|
||||
|
||||
def test_trading_boundary_inclusive_start_exclusive_end():
|
||||
# 07:00 정각 = 트레이딩 시작, 16:30 정각 = 마감 (16:30은 free)
|
||||
assert current_mode(_kst(2026, 5, 22, 7, 0), HOLIDAYS) == "trading"
|
||||
assert current_mode(_kst(2026, 5, 22, 16, 29), HOLIDAYS) == "trading"
|
||||
assert current_mode(_kst(2026, 5, 22, 16, 30), HOLIDAYS) == "free"
|
||||
59
services/task-watcher/watcher.py
Normal file
59
services/task-watcher/watcher.py
Normal file
@@ -0,0 +1,59 @@
|
||||
"""30초마다 current_mode 판정 → queue:paused 토글.
|
||||
|
||||
trading → SET queue:paused 1 EX 600 (10분 TTL — watcher 죽어도 자동 해제)
|
||||
free → DEL queue:paused
|
||||
holidays는 1시간마다 refresh (매 loop fetch 부하 회피).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import datetime as dt
|
||||
import logging
|
||||
import os
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
from mode import current_mode, fetch_holidays, KST
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
REDIS_URL = os.getenv("REDIS_URL", "redis://192.168.45.54:6379")
|
||||
PAUSED_KEY = "queue:paused"
|
||||
LOOP_INTERVAL = 30 # 초
|
||||
HOLIDAYS_REFRESH = 3600 # 1시간
|
||||
PAUSED_TTL = 600 # 10분 (watcher 죽어도 자동 해제)
|
||||
|
||||
|
||||
async def watcher_loop():
|
||||
redis = aioredis.from_url(REDIS_URL, decode_responses=False)
|
||||
holidays = fetch_holidays()
|
||||
last_holiday_refresh = dt.datetime.now(KST)
|
||||
last_mode = None
|
||||
logger.info("task-watcher started (trading window 토글)")
|
||||
|
||||
while True:
|
||||
try:
|
||||
now = dt.datetime.now(KST)
|
||||
# holidays 주기적 refresh
|
||||
if (now - last_holiday_refresh).total_seconds() >= HOLIDAYS_REFRESH:
|
||||
holidays = fetch_holidays()
|
||||
last_holiday_refresh = now
|
||||
|
||||
mode = current_mode(now, holidays)
|
||||
if mode == "trading":
|
||||
await redis.set(PAUSED_KEY, b"1", ex=PAUSED_TTL)
|
||||
else:
|
||||
await redis.delete(PAUSED_KEY)
|
||||
|
||||
if mode != last_mode:
|
||||
logger.info("mode 전환: %s → %s (paused=%s)", last_mode, mode, mode == "trading")
|
||||
last_mode = mode
|
||||
|
||||
await asyncio.sleep(LOOP_INTERVAL)
|
||||
except asyncio.CancelledError:
|
||||
logger.info("watcher_loop cancelled")
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("watcher_loop iteration 실패, 30초 후 재시도")
|
||||
await asyncio.sleep(LOOP_INTERVAL)
|
||||
27
services/video-render/.env.example
Normal file
27
services/video-render/.env.example
Normal file
@@ -0,0 +1,27 @@
|
||||
# Plan-B-Video — Windows video-render worker
|
||||
|
||||
# NAS Redis 큐
|
||||
REDIS_URL=redis://192.168.45.54:6379
|
||||
|
||||
# NAS internal webhook (video-lab port 18801)
|
||||
NAS_BASE_URL=http://192.168.45.54:18801
|
||||
INTERNAL_API_KEY=__copy_from_nas_dotenv__
|
||||
|
||||
# Sora 2 (OpenAI)
|
||||
OPENAI_API_KEY=__paste_openai_key__
|
||||
|
||||
# Veo (Google Gemini API — ai.google.dev. Vertex AI 경로 아님, GCS bucket 불필요)
|
||||
GEMINI_API_KEY=__paste_gemini_key__
|
||||
|
||||
# Kling (Native KlingAI — JWT auth with Access Key + Secret Key)
|
||||
KLING_ACCESS_KEY=__paste_kling_access_key__
|
||||
KLING_SECRET_KEY=__paste_kling_secret_key__
|
||||
|
||||
# Seedance 2.0 (BytePlus)
|
||||
SEEDANCE_API_KEY=__paste_seedance_key__
|
||||
|
||||
# NAS SMB mount 안의 video 디렉토리
|
||||
VIDEO_MEDIA_ROOT=/mnt/nas/webpage/data/video
|
||||
|
||||
# nginx 서빙 prefix (NAS webhook payload용)
|
||||
VIDEO_MEDIA_URL_PREFIX=/media/video
|
||||
19
services/video-render/Dockerfile
Normal file
19
services/video-render/Dockerfile
Normal file
@@ -0,0 +1,19 @@
|
||||
FROM python:3.12-slim-bookworm
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY video-render/requirements.txt /app/
|
||||
RUN pip install --no-cache-dir --timeout 600 --retries 5 -r requirements.txt
|
||||
|
||||
# F6: 공통 ReliableQueue 모듈 (services/_shared)
|
||||
COPY _shared /app/_shared
|
||||
COPY video-render/. /app/
|
||||
ENV PYTHONPATH=/app
|
||||
|
||||
EXPOSE 8000
|
||||
CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "1"]
|
||||
5
services/video-render/conftest.py
Normal file
5
services/video-render/conftest.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""Make services/ root importable so `from _shared.reliable_queue import ...` works during tests."""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
36
services/video-render/main.py
Normal file
36
services/video-render/main.py
Normal file
@@ -0,0 +1,36 @@
|
||||
"""video-render FastAPI entry — health + lifespan (worker loop spawn)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
import worker
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
worker_task = asyncio.create_task(worker.worker_loop())
|
||||
logger.info("video-render lifespan 시작")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
worker_task.cancel()
|
||||
try:
|
||||
await worker_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
logger.info("video-render lifespan 종료")
|
||||
|
||||
|
||||
app = FastAPI(lifespan=lifespan)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"ok": True, "service": "video-render"}
|
||||
54
services/video-render/nas_client.py
Normal file
54
services/video-render/nas_client.py
Normal file
@@ -0,0 +1,54 @@
|
||||
"""NAS webhook 어댑터 — Windows worker가 NAS DB 직접 접근 못하므로 HTTP로 위임.
|
||||
|
||||
Plan-B-Music nas_client와 동일 패턴 (call-time os.getenv으로 테스트 격리).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_TIMEOUT = 10.0
|
||||
|
||||
|
||||
def _post(payload: Dict[str, Any]) -> None:
|
||||
nas_base_url = os.getenv("NAS_BASE_URL", "http://192.168.45.54:18801")
|
||||
internal_api_key = os.getenv("INTERNAL_API_KEY", "")
|
||||
url = f"{nas_base_url}/api/internal/video/update"
|
||||
try:
|
||||
r = httpx.post(
|
||||
url,
|
||||
headers={"X-Internal-Key": internal_api_key},
|
||||
json=payload,
|
||||
timeout=_TIMEOUT,
|
||||
)
|
||||
if r.status_code != 200:
|
||||
logger.error("webhook %s returned %d: %s",
|
||||
payload.get("task_id"), r.status_code, r.text[:200])
|
||||
except Exception:
|
||||
logger.exception("webhook %s 호출 실패", payload.get("task_id"))
|
||||
|
||||
|
||||
def webhook_update_task(
|
||||
task_id: str,
|
||||
status: str,
|
||||
progress: int,
|
||||
message: str = "",
|
||||
video_url: Optional[str] = None,
|
||||
error: Optional[str] = None,
|
||||
) -> None:
|
||||
payload: Dict[str, Any] = {
|
||||
"task_id": task_id,
|
||||
"status": status,
|
||||
"progress": progress,
|
||||
"message": message,
|
||||
}
|
||||
if video_url is not None:
|
||||
payload["video_url"] = video_url
|
||||
if error is not None:
|
||||
payload["error"] = error
|
||||
_post(payload)
|
||||
0
services/video-render/providers/__init__.py
Normal file
0
services/video-render/providers/__init__.py
Normal file
153
services/video-render/providers/kling.py
Normal file
153
services/video-render/providers/kling.py
Normal file
@@ -0,0 +1,153 @@
|
||||
"""Kling AI video generation — Native KlingAI API (api.klingai.com).
|
||||
|
||||
JWT auth: HS256, payload {iss: ACCESS_KEY, exp: now+1800, nbf: now-5}.
|
||||
POST /v1/videos/text2video → GET /v1/videos/text2video/{task_id} → task_result.videos[0].url 다운로드.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import jwt as pyjwt
|
||||
import requests
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
KLING_BASE_URL = "https://api.klingai.com"
|
||||
VIDEO_MEDIA_ROOT = os.getenv("VIDEO_MEDIA_ROOT", "/mnt/nas/webpage/data/video")
|
||||
VIDEO_MEDIA_URL_PREFIX = os.getenv("VIDEO_MEDIA_URL_PREFIX", "/media/video")
|
||||
|
||||
POLL_INTERVAL = 10
|
||||
POLL_MAX_ATTEMPTS = 60 # 최대 ~10분
|
||||
|
||||
DEFAULT_MODEL = "kling-v1-6"
|
||||
|
||||
JWT_EXP_SECONDS = 1800 # 30분
|
||||
JWT_NBF_OFFSET = -5 # 5초 뒤로
|
||||
|
||||
|
||||
def _generate_jwt() -> Optional[str]:
|
||||
access_key = os.getenv("KLING_ACCESS_KEY", "")
|
||||
secret_key = os.getenv("KLING_SECRET_KEY", "")
|
||||
if not access_key or not secret_key:
|
||||
return None
|
||||
now = int(time.time())
|
||||
headers = {"alg": "HS256", "typ": "JWT"}
|
||||
payload = {"iss": access_key, "exp": now + JWT_EXP_SECONDS, "nbf": now + JWT_NBF_OFFSET}
|
||||
return pyjwt.encode(payload, secret_key, algorithm="HS256", headers=headers)
|
||||
|
||||
|
||||
def _headers() -> dict:
|
||||
token = _generate_jwt()
|
||||
return {
|
||||
"Authorization": f"Bearer {token}" if token else "",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
|
||||
def run_kling_generation(task_id: str, params: dict) -> None:
|
||||
"""Kling으로 영상 생성 → mp4 → NAS SMB → webhook."""
|
||||
try:
|
||||
if not os.getenv("KLING_ACCESS_KEY") or not os.getenv("KLING_SECRET_KEY"):
|
||||
webhook_update_task(task_id, "failed", 0, "",
|
||||
error="KLING_ACCESS_KEY 또는 KLING_SECRET_KEY 미설정")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 5, "Kling API 호출 중...")
|
||||
|
||||
# image_url 있으면 image2video, 없으면 text2video
|
||||
is_image2video = bool(params.get("image_url"))
|
||||
endpoint_path = "/v1/videos/image2video" if is_image2video else "/v1/videos/text2video"
|
||||
|
||||
body = {
|
||||
"model_name": params.get("model") or DEFAULT_MODEL,
|
||||
"prompt": params["prompt"][:2500],
|
||||
"duration": str(params.get("duration", 5)),
|
||||
"aspect_ratio": params.get("aspect_ratio", "16:9"),
|
||||
"mode": params.get("mode", "std"),
|
||||
}
|
||||
if params.get("negative_prompt"):
|
||||
body["negative_prompt"] = params["negative_prompt"][:2500]
|
||||
if params.get("cfg_scale") is not None:
|
||||
body["cfg_scale"] = float(params["cfg_scale"])
|
||||
if is_image2video:
|
||||
body["image"] = params["image_url"]
|
||||
|
||||
resp = requests.post(f"{KLING_BASE_URL}{endpoint_path}",
|
||||
headers=_headers(), json=body, timeout=30)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "",
|
||||
error=f"Kling API 오류: {resp.status_code} {resp.text[:300]}")
|
||||
return
|
||||
|
||||
body_json = resp.json()
|
||||
if body_json.get("code") != 0:
|
||||
webhook_update_task(task_id, "failed", 0, "",
|
||||
error=f"Kling API 거부: {body_json.get('message', '?')}")
|
||||
return
|
||||
|
||||
kling_task_id = (body_json.get("data") or {}).get("task_id", "")
|
||||
if not kling_task_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Kling 응답에 task_id 없음")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 15, "Kling 작업 등록됨")
|
||||
|
||||
# 폴링 — GET /v1/videos/{text2video|image2video}/{task_id}
|
||||
video_url = None
|
||||
for attempt in range(POLL_MAX_ATTEMPTS):
|
||||
time.sleep(POLL_INTERVAL)
|
||||
fetch = requests.get(f"{KLING_BASE_URL}{endpoint_path}/{kling_task_id}",
|
||||
headers=_headers(), timeout=30)
|
||||
if fetch.status_code != 200:
|
||||
continue
|
||||
fd = fetch.json()
|
||||
if fd.get("code") != 0:
|
||||
continue
|
||||
data = fd.get("data") or {}
|
||||
status = data.get("task_status", "")
|
||||
scaled = min(15 + int((attempt / POLL_MAX_ATTEMPTS) * 65), 79)
|
||||
webhook_update_task(task_id, "processing", scaled, f"Kling 생성 중... ({status})")
|
||||
|
||||
if status == "succeed":
|
||||
videos = ((data.get("task_result") or {}).get("videos") or [])
|
||||
if videos:
|
||||
video_url = videos[0].get("url", "")
|
||||
break
|
||||
elif status == "failed":
|
||||
err = data.get("task_status_msg") or "Kling 작업 실패"
|
||||
webhook_update_task(task_id, "failed", 0, "", error=err)
|
||||
return
|
||||
# submitted/processing → 계속 폴링
|
||||
else:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Kling 폴링 timeout (10분)")
|
||||
return
|
||||
|
||||
if not video_url:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Kling 완료했으나 video url 없음")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 85, "Kling 결과 다운로드 중...")
|
||||
filename = f"{task_id}.mp4"
|
||||
os.makedirs(VIDEO_MEDIA_ROOT, exist_ok=True)
|
||||
file_path = os.path.join(VIDEO_MEDIA_ROOT, filename)
|
||||
|
||||
# Kling 결과 url은 일반적으로 인증 불필요 (signed URL)
|
||||
dl = requests.get(video_url, stream=True, timeout=300)
|
||||
dl.raise_for_status()
|
||||
with open(file_path, "wb") as f:
|
||||
for chunk in dl.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
local_url = f"{VIDEO_MEDIA_URL_PREFIX}/{filename}"
|
||||
webhook_update_task(task_id, "succeeded", 100, "Kling 생성 완료", video_url=local_url)
|
||||
|
||||
except requests.Timeout:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Kling API 타임아웃")
|
||||
except Exception as e:
|
||||
logger.exception("Kling generation error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
121
services/video-render/providers/seedance.py
Normal file
121
services/video-render/providers/seedance.py
Normal file
@@ -0,0 +1,121 @@
|
||||
"""Seedance 2.0 video generation — ByteDance Volcano Engine (BytePlus 국제 endpoint).
|
||||
|
||||
POST https://api.byteplus.com/seedance/v1/videos → GET /videos/{id} 폴링 → output.video_url 다운로드.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SEEDANCE_BASE_URL = "https://api.byteplus.com/seedance/v1"
|
||||
VIDEO_MEDIA_ROOT = os.getenv("VIDEO_MEDIA_ROOT", "/mnt/nas/webpage/data/video")
|
||||
VIDEO_MEDIA_URL_PREFIX = os.getenv("VIDEO_MEDIA_URL_PREFIX", "/media/video")
|
||||
|
||||
POLL_INTERVAL = 8 # Seedance는 30~120초
|
||||
POLL_MAX_ATTEMPTS = 60
|
||||
|
||||
DEFAULT_MODEL = "seedance-2.0"
|
||||
|
||||
|
||||
def _headers() -> dict:
|
||||
api_key = os.getenv("SEEDANCE_API_KEY", "")
|
||||
return {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
|
||||
def run_seedance_generation(task_id: str, params: dict) -> None:
|
||||
"""Seedance로 영상 생성 → mp4 → NAS SMB → webhook."""
|
||||
try:
|
||||
if not os.getenv("SEEDANCE_API_KEY"):
|
||||
webhook_update_task(task_id, "failed", 0, "", error="SEEDANCE_API_KEY 미설정")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 5, "Seedance API 호출 중...")
|
||||
|
||||
body = {
|
||||
"model": params.get("model") or DEFAULT_MODEL,
|
||||
"prompt": params["prompt"][:2000],
|
||||
"resolution": params.get("resolution", "1080p"),
|
||||
"duration": params.get("duration", 5),
|
||||
"aspect_ratio": params.get("aspect_ratio", "16:9"),
|
||||
}
|
||||
if params.get("negative_prompt"):
|
||||
body["negative_prompt"] = params["negative_prompt"]
|
||||
if params.get("image_url"):
|
||||
body["references"] = [{"type": "image", "data": params["image_url"], "role": "subject"}]
|
||||
if params.get("audio") is not None:
|
||||
body["audio"] = bool(params["audio"])
|
||||
if params.get("seed") is not None:
|
||||
body["seed"] = int(params["seed"])
|
||||
|
||||
resp = requests.post(f"{SEEDANCE_BASE_URL}/videos", headers=_headers(), json=body, timeout=30)
|
||||
if resp.status_code not in (200, 201):
|
||||
webhook_update_task(task_id, "failed", 0, "",
|
||||
error=f"Seedance API 오류: {resp.status_code} {resp.text[:300]}")
|
||||
return
|
||||
|
||||
body_json = resp.json()
|
||||
job_id = body_json.get("id", "")
|
||||
if not job_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Seedance 응답에 id 없음")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 15, "Seedance 작업 등록됨")
|
||||
|
||||
# 폴링
|
||||
video_url = None
|
||||
for attempt in range(POLL_MAX_ATTEMPTS):
|
||||
time.sleep(POLL_INTERVAL)
|
||||
fetch = requests.get(f"{SEEDANCE_BASE_URL}/videos/{job_id}",
|
||||
headers=_headers(), timeout=30)
|
||||
if fetch.status_code != 200:
|
||||
continue
|
||||
fd = fetch.json()
|
||||
status = fd.get("status", "")
|
||||
scaled = min(15 + int((attempt / POLL_MAX_ATTEMPTS) * 65), 79)
|
||||
webhook_update_task(task_id, "processing", scaled, f"Seedance 생성 중... ({status})")
|
||||
|
||||
if status == "completed":
|
||||
video_url = (fd.get("output") or {}).get("video_url", "")
|
||||
break
|
||||
elif status == "failed":
|
||||
err = fd.get("error") or "Seedance 작업 실패"
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(err)[:300])
|
||||
return
|
||||
else:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Seedance 폴링 timeout (10분)")
|
||||
return
|
||||
|
||||
if not video_url:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Seedance 완료했으나 video_url 없음")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 85, "Seedance 결과 다운로드 중...")
|
||||
filename = f"{task_id}.mp4"
|
||||
os.makedirs(VIDEO_MEDIA_ROOT, exist_ok=True)
|
||||
file_path = os.path.join(VIDEO_MEDIA_ROOT, filename)
|
||||
|
||||
dl = requests.get(video_url, stream=True, timeout=300)
|
||||
dl.raise_for_status()
|
||||
with open(file_path, "wb") as f:
|
||||
for chunk in dl.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
local_url = f"{VIDEO_MEDIA_URL_PREFIX}/{filename}"
|
||||
webhook_update_task(task_id, "succeeded", 100, "Seedance 생성 완료", video_url=local_url)
|
||||
|
||||
except requests.Timeout:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Seedance API 타임아웃")
|
||||
except Exception as e:
|
||||
logger.exception("Seedance generation error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
119
services/video-render/providers/sora.py
Normal file
119
services/video-render/providers/sora.py
Normal file
@@ -0,0 +1,119 @@
|
||||
"""Sora 2 video generation — OpenAI Videos API.
|
||||
|
||||
POST /v1/videos → poll GET /v1/videos/{id} → GET /v1/videos/{id}/content download.
|
||||
⚠️ Deprecated, shutdown 2026-09-24. Spec 진행은 박재오 결정 따름.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SORA_BASE_URL = "https://api.openai.com/v1"
|
||||
VIDEO_MEDIA_ROOT = os.getenv("VIDEO_MEDIA_ROOT", "/mnt/nas/webpage/data/video")
|
||||
VIDEO_MEDIA_URL_PREFIX = os.getenv("VIDEO_MEDIA_URL_PREFIX", "/media/video")
|
||||
|
||||
POLL_INTERVAL = 15 # OpenAI 권장: 10~20초
|
||||
POLL_MAX_ATTEMPTS = 40 # 최대 ~10분
|
||||
|
||||
DEFAULT_MODEL = "sora-2"
|
||||
|
||||
|
||||
def _headers() -> dict:
|
||||
api_key = os.getenv("OPENAI_API_KEY", "")
|
||||
return {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
|
||||
def run_sora_generation(task_id: str, params: dict) -> None:
|
||||
"""Sora 2로 영상 생성 → mp4 → NAS SMB 저장 → webhook."""
|
||||
try:
|
||||
if not os.getenv("OPENAI_API_KEY"):
|
||||
webhook_update_task(task_id, "failed", 0, "", error="OPENAI_API_KEY 미설정 (Windows .env)")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 5, "Sora API 호출 중...")
|
||||
|
||||
payload = {
|
||||
"model": params.get("model") or DEFAULT_MODEL,
|
||||
"prompt": params["prompt"][:5000],
|
||||
}
|
||||
if params.get("duration"):
|
||||
payload["seconds"] = params["duration"]
|
||||
if params.get("size"):
|
||||
payload["size"] = params["size"]
|
||||
elif params.get("aspect_ratio") == "9:16":
|
||||
payload["size"] = "1080x1920"
|
||||
elif params.get("aspect_ratio") == "16:9":
|
||||
payload["size"] = "1920x1080"
|
||||
|
||||
resp = requests.post(f"{SORA_BASE_URL}/videos", headers=_headers(), json=payload, timeout=30)
|
||||
if resp.status_code not in (200, 201):
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"Sora API 오류: {resp.status_code} {resp.text[:300]}")
|
||||
return
|
||||
|
||||
body = resp.json()
|
||||
video_id = body.get("id", "")
|
||||
if not video_id:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Sora 응답에 video id 없음")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 15, f"Sora 작업 생성됨 (id={video_id[:16]})")
|
||||
|
||||
# 폴링
|
||||
for attempt in range(POLL_MAX_ATTEMPTS):
|
||||
time.sleep(POLL_INTERVAL)
|
||||
sr = requests.get(f"{SORA_BASE_URL}/videos/{video_id}", headers=_headers(), timeout=30)
|
||||
if sr.status_code != 200:
|
||||
continue
|
||||
sd = sr.json()
|
||||
status = sd.get("status", "")
|
||||
progress = sd.get("progress", 0)
|
||||
scaled = min(15 + int(progress * 0.65), 79)
|
||||
webhook_update_task(task_id, "processing", scaled, f"Sora 생성 중... {progress}%")
|
||||
|
||||
if status == "completed":
|
||||
break
|
||||
elif status == "failed":
|
||||
err = sd.get("error", {}).get("message", "Sora 작업 실패")
|
||||
webhook_update_task(task_id, "failed", 0, "", error=err)
|
||||
return
|
||||
else:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Sora 폴링 timeout (10분)")
|
||||
return
|
||||
|
||||
# 다운로드
|
||||
webhook_update_task(task_id, "processing", 80, "Sora 결과 다운로드 중...")
|
||||
filename = f"{task_id}.mp4"
|
||||
os.makedirs(VIDEO_MEDIA_ROOT, exist_ok=True)
|
||||
file_path = os.path.join(VIDEO_MEDIA_ROOT, filename)
|
||||
|
||||
dl = requests.get(
|
||||
f"{SORA_BASE_URL}/videos/{video_id}/content",
|
||||
headers=_headers(),
|
||||
params={"variant": "video"},
|
||||
stream=True,
|
||||
timeout=300,
|
||||
)
|
||||
dl.raise_for_status()
|
||||
with open(file_path, "wb") as f:
|
||||
for chunk in dl.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
local_url = f"{VIDEO_MEDIA_URL_PREFIX}/{filename}"
|
||||
webhook_update_task(task_id, "succeeded", 100, "Sora 생성 완료", video_url=local_url)
|
||||
|
||||
except requests.Timeout:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Sora API 타임아웃")
|
||||
except Exception as e:
|
||||
logger.exception("Sora generation error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
139
services/video-render/providers/veo.py
Normal file
139
services/video-render/providers/veo.py
Normal file
@@ -0,0 +1,139 @@
|
||||
"""Veo 3.1 video generation — Gemini API (ai.google.dev).
|
||||
|
||||
POST https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:predictLongRunning
|
||||
GET https://generativelanguage.googleapis.com/v1beta/{operation_name}
|
||||
→ done=true 시 response.generateVideoResponse.generatedSamples[0].video.uri 다운로드
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta"
|
||||
VIDEO_MEDIA_ROOT = os.getenv("VIDEO_MEDIA_ROOT", "/mnt/nas/webpage/data/video")
|
||||
VIDEO_MEDIA_URL_PREFIX = os.getenv("VIDEO_MEDIA_URL_PREFIX", "/media/video")
|
||||
|
||||
POLL_INTERVAL = 10 # Veo는 30~120초 소요
|
||||
POLL_MAX_ATTEMPTS = 60 # 최대 ~10분
|
||||
|
||||
DEFAULT_MODEL = "veo-3.1-fast-generate-preview"
|
||||
|
||||
|
||||
def _headers() -> dict:
|
||||
api_key = os.getenv("GEMINI_API_KEY", "")
|
||||
return {
|
||||
"x-goog-api-key": api_key,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
|
||||
def run_veo_generation(task_id: str, params: dict) -> None:
|
||||
"""Veo로 영상 생성 → mp4 → NAS SMB → webhook."""
|
||||
try:
|
||||
if not os.getenv("GEMINI_API_KEY"):
|
||||
webhook_update_task(task_id, "failed", 0, "", error="GEMINI_API_KEY 미설정 (Windows .env)")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 5, "Veo (Gemini API) 호출 중...")
|
||||
|
||||
model_id = params.get("model") or DEFAULT_MODEL
|
||||
body = {
|
||||
"instances": [{"prompt": params["prompt"]}],
|
||||
"parameters": {
|
||||
"aspectRatio": params.get("aspect_ratio") or "16:9",
|
||||
},
|
||||
}
|
||||
# numberOfVideos는 일부 모델(veo-3.0-fast 등) 미지원 — 호출자 명시 시에만 추가
|
||||
if params.get("number_of_videos"):
|
||||
body["parameters"]["numberOfVideos"] = int(params["number_of_videos"])
|
||||
if params.get("duration"):
|
||||
body["parameters"]["durationSeconds"] = int(params["duration"])
|
||||
if params.get("resolution"):
|
||||
body["parameters"]["resolution"] = params["resolution"]
|
||||
if params.get("negative_prompt"):
|
||||
body["parameters"]["negativePrompt"] = params["negative_prompt"]
|
||||
if params.get("person_generation"):
|
||||
body["parameters"]["personGeneration"] = params["person_generation"]
|
||||
|
||||
resp = requests.post(
|
||||
f"{GEMINI_BASE_URL}/models/{model_id}:predictLongRunning",
|
||||
headers=_headers(), json=body, timeout=30,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
webhook_update_task(task_id, "failed", 0, "",
|
||||
error=f"Veo Gemini API 오류: {resp.status_code} {resp.text[:300]}")
|
||||
return
|
||||
|
||||
op_name = resp.json().get("name", "")
|
||||
if not op_name:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Veo 응답에 operation name 없음")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 15, "Veo 작업 시작됨")
|
||||
|
||||
# 폴링 — GET /v1beta/{operation_name}
|
||||
video_uri = None
|
||||
for attempt in range(POLL_MAX_ATTEMPTS):
|
||||
time.sleep(POLL_INTERVAL)
|
||||
fetch = requests.get(
|
||||
f"{GEMINI_BASE_URL}/{op_name}",
|
||||
headers=_headers(),
|
||||
timeout=30,
|
||||
)
|
||||
if fetch.status_code != 200:
|
||||
continue
|
||||
fd = fetch.json()
|
||||
done = fd.get("done", False)
|
||||
scaled = min(15 + int((attempt / POLL_MAX_ATTEMPTS) * 65), 79)
|
||||
webhook_update_task(task_id, "processing", scaled, "Veo 생성 중...")
|
||||
|
||||
if done:
|
||||
if "error" in fd:
|
||||
webhook_update_task(task_id, "failed", 0, "",
|
||||
error=f"Veo 작업 실패: {fd['error'].get('message','?')}")
|
||||
return
|
||||
# response.generateVideoResponse.generatedSamples[0].video.uri
|
||||
response = fd.get("response") or {}
|
||||
gen = response.get("generateVideoResponse") or {}
|
||||
samples = gen.get("generatedSamples") or []
|
||||
if not samples:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Veo 완료했으나 generatedSamples 비어 있음")
|
||||
return
|
||||
video_uri = (samples[0].get("video") or {}).get("uri", "")
|
||||
break
|
||||
else:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Veo 폴링 timeout (10분)")
|
||||
return
|
||||
|
||||
if not video_uri:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Veo 응답에 video.uri 없음")
|
||||
return
|
||||
|
||||
webhook_update_task(task_id, "processing", 85, "Veo 결과 다운로드 중...")
|
||||
filename = f"{task_id}.mp4"
|
||||
os.makedirs(VIDEO_MEDIA_ROOT, exist_ok=True)
|
||||
file_path = os.path.join(VIDEO_MEDIA_ROOT, filename)
|
||||
|
||||
# 다운로드 — x-goog-api-key 헤더 그대로 사용 (Gemini API가 인증 처리)
|
||||
dl = requests.get(video_uri, headers=_headers(), stream=True, timeout=300)
|
||||
dl.raise_for_status()
|
||||
with open(file_path, "wb") as f:
|
||||
for chunk in dl.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
local_url = f"{VIDEO_MEDIA_URL_PREFIX}/{filename}"
|
||||
webhook_update_task(task_id, "succeeded", 100, "Veo 생성 완료", video_url=local_url)
|
||||
|
||||
except requests.Timeout:
|
||||
webhook_update_task(task_id, "failed", 0, "", error="Veo API 타임아웃")
|
||||
except Exception as e:
|
||||
logger.exception("Veo generation error task=%s", task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=str(e))
|
||||
10
services/video-render/requirements.txt
Normal file
10
services/video-render/requirements.txt
Normal file
@@ -0,0 +1,10 @@
|
||||
fastapi==0.115.6
|
||||
uvicorn[standard]==0.34.0
|
||||
requests==2.32.3
|
||||
redis>=5.0
|
||||
httpx>=0.27
|
||||
openai>=1.50.0
|
||||
PyJWT>=2.8.0
|
||||
pytest>=8.0
|
||||
pytest-asyncio>=0.24
|
||||
respx>=0.21
|
||||
0
services/video-render/tests/__init__.py
Normal file
0
services/video-render/tests/__init__.py
Normal file
70
services/video-render/tests/test_nas_client.py
Normal file
70
services/video-render/tests/test_nas_client.py
Normal file
@@ -0,0 +1,70 @@
|
||||
"""nas_client — webhook adapter for video-render."""
|
||||
import pytest
|
||||
import respx
|
||||
import httpx
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _env(monkeypatch):
|
||||
monkeypatch.setenv("NAS_BASE_URL", "http://nas-test:18801")
|
||||
monkeypatch.setenv("INTERNAL_API_KEY", "test-key")
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_webhook_update_task_sends_x_internal_key():
|
||||
route = respx.post("http://nas-test:18801/api/internal/video/update").mock(
|
||||
return_value=httpx.Response(200, json={"ok": True})
|
||||
)
|
||||
webhook_update_task("task-1", "processing", 30, message="downloading")
|
||||
assert route.called
|
||||
req = route.calls[0].request
|
||||
assert req.headers["X-Internal-Key"] == "test-key"
|
||||
import json
|
||||
body = json.loads(req.content)
|
||||
assert body["task_id"] == "task-1"
|
||||
assert body["status"] == "processing"
|
||||
assert body["progress"] == 30
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_webhook_update_task_with_video_url():
|
||||
route = respx.post("http://nas-test:18801/api/internal/video/update").mock(
|
||||
return_value=httpx.Response(200, json={"ok": True})
|
||||
)
|
||||
webhook_update_task("task-2", "succeeded", 100, message="완료",
|
||||
video_url="/media/video/task-2.mp4")
|
||||
import json
|
||||
payload = json.loads(route.calls[0].request.content)
|
||||
assert payload["video_url"] == "/media/video/task-2.mp4"
|
||||
assert payload["status"] == "succeeded"
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_webhook_update_task_with_error():
|
||||
route = respx.post("http://nas-test:18801/api/internal/video/update").mock(
|
||||
return_value=httpx.Response(200, json={"ok": True})
|
||||
)
|
||||
webhook_update_task("task-3", "failed", 0, error="Sora API rate limit")
|
||||
import json
|
||||
payload = json.loads(route.calls[0].request.content)
|
||||
assert payload["error"] == "Sora API rate limit"
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_webhook_swallows_network_error(caplog):
|
||||
respx.post("http://nas-test:18801/api/internal/video/update").mock(
|
||||
side_effect=httpx.ConnectError("no host")
|
||||
)
|
||||
webhook_update_task("task-5", "processing", 10)
|
||||
assert "task-5" in caplog.text
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_webhook_swallows_non_200(caplog):
|
||||
respx.post("http://nas-test:18801/api/internal/video/update").mock(
|
||||
return_value=httpx.Response(500, text="server error")
|
||||
)
|
||||
webhook_update_task("task-6", "processing", 50)
|
||||
assert "task-6" in caplog.text
|
||||
96
services/video-render/tests/test_worker.py
Normal file
96
services/video-render/tests/test_worker.py
Normal file
@@ -0,0 +1,96 @@
|
||||
"""worker.py — job_type 디스패처 (4 provider)."""
|
||||
import pytest
|
||||
from unittest.mock import patch
|
||||
|
||||
import worker
|
||||
|
||||
|
||||
def test_dispatch_sora_calls_run_sora_generation():
|
||||
payload = {"task_id": "t1", "job_type": "sora_generation", "params": {"prompt": "x"}}
|
||||
with patch("worker.run_sora_generation") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t1", {"prompt": "x"})
|
||||
|
||||
|
||||
def test_dispatch_veo_calls_run_veo_generation():
|
||||
payload = {"task_id": "t2", "job_type": "veo_generation", "params": {"prompt": "x"}}
|
||||
with patch("worker.run_veo_generation") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t2", {"prompt": "x"})
|
||||
|
||||
|
||||
def test_dispatch_kling_calls_run_kling_generation():
|
||||
payload = {"task_id": "t3", "job_type": "kling_generation", "params": {"prompt": "x"}}
|
||||
with patch("worker.run_kling_generation") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t3", {"prompt": "x"})
|
||||
|
||||
|
||||
def test_dispatch_seedance_calls_run_seedance_generation():
|
||||
payload = {"task_id": "t4", "job_type": "seedance_generation", "params": {"prompt": "x"}}
|
||||
with patch("worker.run_seedance_generation") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once_with("t4", {"prompt": "x"})
|
||||
|
||||
|
||||
def test_dispatch_unknown_job_type_logs_error():
|
||||
payload = {"task_id": "t5", "job_type": "weird_type", "params": {}}
|
||||
with patch("worker.webhook_update_task") as m:
|
||||
worker._dispatch(payload)
|
||||
m.assert_called_once()
|
||||
args = m.call_args[0]
|
||||
assert args[0] == "t5"
|
||||
assert args[1] == "failed"
|
||||
|
||||
|
||||
# ----- F6: ReliableQueue poll_once -----
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_acks_on_success(monkeypatch):
|
||||
payload = {"task_id": "t1", "job_type": "sora_generation", "params": {}}
|
||||
raw = json.dumps(payload).encode()
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=(payload, raw))
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
monkeypatch.setattr(worker, "_dispatch", MagicMock())
|
||||
handled = await worker.poll_once(fake_queue)
|
||||
assert handled is True
|
||||
fake_queue.ack.assert_awaited_once_with(raw)
|
||||
fake_queue.fail.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_calls_fail_on_dispatch_exception(monkeypatch):
|
||||
payload = {"task_id": "t2", "job_type": "sora_generation", "params": {}}
|
||||
raw = json.dumps(payload).encode()
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=(payload, raw))
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
|
||||
def _boom(p):
|
||||
raise RuntimeError("dispatch crash")
|
||||
|
||||
monkeypatch.setattr(worker, "_dispatch", _boom)
|
||||
handled = await worker.poll_once(fake_queue)
|
||||
assert handled is True
|
||||
fake_queue.fail.assert_awaited_once_with(raw, payload)
|
||||
fake_queue.ack.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_once_returns_false_on_timeout(monkeypatch):
|
||||
fake_queue = AsyncMock()
|
||||
fake_queue.dequeue = AsyncMock(return_value=None)
|
||||
fake_queue.ack = AsyncMock()
|
||||
fake_queue.fail = AsyncMock()
|
||||
monkeypatch.setattr(worker, "_dispatch", MagicMock())
|
||||
handled = await worker.poll_once(fake_queue)
|
||||
assert handled is False
|
||||
fake_queue.ack.assert_not_awaited()
|
||||
fake_queue.fail.assert_not_awaited()
|
||||
98
services/video-render/worker.py
Normal file
98
services/video-render/worker.py
Normal file
@@ -0,0 +1,98 @@
|
||||
"""Redis ReliableQueue worker — F6 신뢰성 패턴 (BLMOVE + ack/fail + recovery).
|
||||
|
||||
queue:paused 가 set이면 대기 (task-watcher가 박재오 활동 감지 시 set).
|
||||
string-based dispatch + getattr (테스트 patch 호환).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
from nas_client import webhook_update_task
|
||||
from providers.sora import run_sora_generation
|
||||
from providers.veo import run_veo_generation
|
||||
from providers.kling import run_kling_generation
|
||||
from providers.seedance import run_seedance_generation
|
||||
from _shared.reliable_queue import ReliableQueue
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
REDIS_URL = os.getenv("REDIS_URL", "redis://192.168.45.54:6379")
|
||||
QUEUE_KEY = "queue:video-render"
|
||||
PAUSED_KEY = "queue:paused"
|
||||
|
||||
# string names so `unittest.mock.patch` on `worker.<name>` is correctly intercepted
|
||||
_DISPATCH_TABLE = {
|
||||
"sora_generation": "run_sora_generation",
|
||||
"veo_generation": "run_veo_generation",
|
||||
"kling_generation": "run_kling_generation",
|
||||
"seedance_generation": "run_seedance_generation",
|
||||
}
|
||||
|
||||
|
||||
def _dispatch(payload: dict) -> None:
|
||||
"""payload[job_type] → provider 함수 호출 (sync, worker_loop에서 asyncio.to_thread로 wrap)."""
|
||||
job_type = payload.get("job_type", "")
|
||||
task_id = payload.get("task_id", "")
|
||||
params = payload.get("params", {})
|
||||
fn_name = _DISPATCH_TABLE.get(job_type)
|
||||
if fn_name is None:
|
||||
logger.error("unknown job_type=%s task=%s", job_type, task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"unknown job_type: {job_type}")
|
||||
return
|
||||
try:
|
||||
fn = getattr(sys.modules[__name__], fn_name)
|
||||
except AttributeError:
|
||||
logger.error("dispatch table typo for job_type=%s name=%s task=%s", job_type, fn_name, task_id)
|
||||
webhook_update_task(task_id, "failed", 0, "", error=f"internal dispatch error: {fn_name}")
|
||||
return
|
||||
fn(task_id, params)
|
||||
|
||||
|
||||
async def poll_once(queue: ReliableQueue) -> bool:
|
||||
"""F6 — 1 cycle: dequeue → _dispatch → ack/fail. Returns True if a job handled."""
|
||||
result = await queue.dequeue(timeout=5)
|
||||
if result is None:
|
||||
return False
|
||||
payload, raw = result
|
||||
try:
|
||||
await asyncio.to_thread(_dispatch, payload)
|
||||
except Exception:
|
||||
logger.exception("dispatch unhandled exception task_id=%s",
|
||||
payload.get("task_id"))
|
||||
await queue.fail(raw, payload)
|
||||
return True
|
||||
await queue.ack(raw)
|
||||
return True
|
||||
|
||||
|
||||
async def worker_loop():
|
||||
redis = aioredis.from_url(REDIS_URL, decode_responses=False)
|
||||
queue = ReliableQueue(redis, queue_key=QUEUE_KEY)
|
||||
logger.info("video-render worker started worker_id=%s queue=%s",
|
||||
queue.worker_id, QUEUE_KEY)
|
||||
try:
|
||||
recovered = await queue.recover()
|
||||
if recovered:
|
||||
logger.info("recovered %d orphaned items at startup", recovered)
|
||||
except Exception:
|
||||
logger.exception("startup recover failed")
|
||||
|
||||
while True:
|
||||
try:
|
||||
paused = await redis.get(PAUSED_KEY)
|
||||
if paused == b"1":
|
||||
await asyncio.sleep(10)
|
||||
continue
|
||||
await poll_once(queue)
|
||||
except asyncio.CancelledError:
|
||||
logger.info("worker_loop cancelled")
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("worker_loop iteration 실패, 5초 후 재시도")
|
||||
await asyncio.sleep(5)
|
||||
Reference in New Issue
Block a user