services: update comfy agent, senpai md consumer, and swapper deps

This commit is contained in:
Apple
2026-02-19 00:14:18 -08:00
parent c201d105f6
commit c57e6ed96b
15 changed files with 586 additions and 40 deletions

View File

@@ -1,5 +1,6 @@
# services/comfy-agent/Dockerfile # services/comfy-agent/Dockerfile
FROM python:3.11-slim ARG BASE_IMAGE=python:3.11-slim
FROM ${BASE_IMAGE}
WORKDIR /app WORKDIR /app
COPY requirements.txt /app/requirements.txt COPY requirements.txt /app/requirements.txt

View File

@@ -1,11 +1,64 @@
# services/comfy-agent/app/api.py # services/comfy-agent/app/api.py
from fastapi import APIRouter, HTTPException import hashlib
import json
from fastapi import APIRouter, Header, HTTPException
from .models import GenerateImageRequest, GenerateVideoRequest, JobStatus from .models import GenerateImageRequest, GenerateVideoRequest, JobStatus
from .jobs import JOB_STORE from .jobs import JOB_STORE
from .worker import enqueue from .worker import enqueue
from . import idempotency
from .config import settings
router = APIRouter() router = APIRouter()
def _req_hash(gen_type: str, payload: dict) -> str:
normalized = json.dumps({"type": gen_type, "payload": payload}, sort_keys=True, separators=(",", ":"))
return hashlib.sha256(normalized.encode("utf-8")).hexdigest()
def _resolve_idempotency_key(
*,
header_key: str | None,
body_key: str | None,
) -> str | None:
key = (header_key or body_key or "").strip()
return key or None
def _create_job_with_idempotency(
*,
gen_type: str,
idem_key: str | None,
req_hash: str,
) -> tuple[JobStatus, bool]:
"""
Returns:
(job_status, should_enqueue)
"""
if not idem_key or idempotency.IDEMPOTENCY_STORE is None:
return JOB_STORE.create(gen_type), True
candidate_job_id = JOB_STORE.new_job_id()
result = idempotency.IDEMPOTENCY_STORE.reserve(
idem_key=idem_key,
gen_type=gen_type,
req_hash=req_hash,
job_id=candidate_job_id,
)
if result.decision == "conflict":
raise HTTPException(status_code=409, detail="idempotency_key_reused_with_different_payload")
if result.decision == "exists":
existing = JOB_STORE.get(result.job_id)
if existing:
return existing, False
# If process was restarted and in-memory JOB_STORE was lost, return queued placeholder.
return JOB_STORE.create(gen_type, job_id=result.job_id), False
return JOB_STORE.create(gen_type, job_id=result.job_id), True
def _build_workflow_t2i(req: GenerateImageRequest) -> dict: def _build_workflow_t2i(req: GenerateImageRequest) -> dict:
# Basic SD 1.5 workflow # Basic SD 1.5 workflow
# Node structure: CheckpointLoader -> CLIP Encode -> KSampler -> VAE Decode -> SaveImage # Node structure: CheckpointLoader -> CLIP Encode -> KSampler -> VAE Decode -> SaveImage
@@ -70,24 +123,147 @@ def _build_workflow_t2i(req: GenerateImageRequest) -> dict:
} }
def _build_workflow_t2v(req: GenerateVideoRequest) -> dict: def _build_workflow_t2v(req: GenerateVideoRequest) -> dict:
# MVP placeholder for LTX-2 pipeline; replace with actual LTX-2 workflow. if not settings.LTX_TEXT_ENCODER:
raise HTTPException(status_code=503, detail="ltx_text_encoder_not_configured")
frame_rate = float(max(1, req.fps))
length = req.frames if req.frames and req.frames > 0 else (max(1, req.seconds) * max(1, req.fps) + 1)
neg = req.negative_prompt if req.negative_prompt else "low quality, worst quality, deformed, distorted, disfigured, motion artifacts"
# LTX-2 text-to-video pipeline with SaveVideo output node.
return { return {
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": req.prompt, "clip": ["2", 0]}}, "4": {
# TODO: Add complete workflow JSON for text-to-video with LTX-2 "class_type": "CheckpointLoaderSimple",
"inputs": {
"ckpt_name": settings.LTX_CKPT_NAME,
},
},
"5": {
"class_type": "LTXAVTextEncoderLoader",
"inputs": {
"text_encoder": settings.LTX_TEXT_ENCODER,
"ckpt_name": settings.LTX_CKPT_NAME,
"device": settings.LTX_DEVICE,
},
},
"6": {
"class_type": "CLIPTextEncode",
"inputs": {
"text": req.prompt,
"clip": ["5", 0],
},
},
"7": {
"class_type": "CLIPTextEncode",
"inputs": {
"text": neg,
"clip": ["5", 0],
},
},
"69": {
"class_type": "LTXVConditioning",
"inputs": {
"positive": ["6", 0],
"negative": ["7", 0],
"frame_rate": frame_rate,
},
},
"70": {
"class_type": "EmptyLTXVLatentVideo",
"inputs": {
"width": req.width,
"height": req.height,
"length": length,
"batch_size": 1,
},
},
"73": {
"class_type": "KSamplerSelect",
"inputs": {
"sampler_name": settings.LTX_SAMPLER,
},
},
"71": {
"class_type": "LTXVScheduler",
"inputs": {
"steps": req.steps,
"max_shift": settings.LTX_MAX_SHIFT,
"base_shift": settings.LTX_BASE_SHIFT,
"stretch": settings.LTX_STRETCH,
"terminal": settings.LTX_TERMINAL,
},
},
"72": {
"class_type": "SamplerCustom",
"inputs": {
"model": ["4", 0],
"add_noise": True,
"noise_seed": req.seed if req.seed is not None else 42,
"cfg": req.cfg,
"positive": ["69", 0],
"negative": ["69", 1],
"sampler": ["73", 0],
"sigmas": ["71", 0],
"latent_image": ["70", 0],
},
},
"8": {
"class_type": "VAEDecode",
"inputs": {
"samples": ["72", 0],
"vae": ["4", 2],
},
},
"78": {
"class_type": "CreateVideo",
"inputs": {
"images": ["8", 0],
"fps": frame_rate,
},
},
"79": {
"class_type": "SaveVideo",
"inputs": {
"video": ["78", 0],
"filename_prefix": "comfy-agent/video",
"format": req.format,
"codec": req.codec,
},
},
} }
@router.post("/generate/image", response_model=JobStatus) @router.post("/generate/image", response_model=JobStatus)
async def generate_image(req: GenerateImageRequest): async def generate_image(
job = JOB_STORE.create("text-to-image") req: GenerateImageRequest,
idempotency_key: str | None = Header(default=None, alias="Idempotency-Key"),
):
idem_key = _resolve_idempotency_key(header_key=idempotency_key, body_key=req.idempotency_key)
req_hash = _req_hash("text-to-image", req.model_dump(mode="json", exclude={"idempotency_key"}))
job, should_enqueue = _create_job_with_idempotency(
gen_type="text-to-image",
idem_key=idem_key,
req_hash=req_hash,
)
graph = _build_workflow_t2i(req) graph = _build_workflow_t2i(req)
enqueue(job.job_id, "text-to-image", graph) if should_enqueue:
enqueue(job.job_id, "text-to-image", graph)
return JOB_STORE.get(job.job_id) return JOB_STORE.get(job.job_id)
@router.post("/generate/video", response_model=JobStatus) @router.post("/generate/video", response_model=JobStatus)
async def generate_video(req: GenerateVideoRequest): async def generate_video(
job = JOB_STORE.create("text-to-video") req: GenerateVideoRequest,
idempotency_key: str | None = Header(default=None, alias="Idempotency-Key"),
):
idem_key = _resolve_idempotency_key(header_key=idempotency_key, body_key=req.idempotency_key)
req_hash = _req_hash("text-to-video", req.model_dump(mode="json", exclude={"idempotency_key"}))
job, should_enqueue = _create_job_with_idempotency(
gen_type="text-to-video",
idem_key=idem_key,
req_hash=req_hash,
)
graph = _build_workflow_t2v(req) graph = _build_workflow_t2v(req)
enqueue(job.job_id, "text-to-video", graph) if should_enqueue:
enqueue(job.job_id, "text-to-video", graph)
return JOB_STORE.get(job.job_id) return JOB_STORE.get(job.job_id)
@router.get("/status/{job_id}", response_model=JobStatus) @router.get("/status/{job_id}", response_model=JobStatus)

View File

@@ -16,7 +16,30 @@ class Settings(BaseSettings):
STORAGE_PATH: str = "/data/comfy-results" STORAGE_PATH: str = "/data/comfy-results"
PUBLIC_BASE_URL: str = "http://212.8.58.133:8880/files" # NODE3 IP PUBLIC_BASE_URL: str = "http://212.8.58.133:8880/files" # NODE3 IP
S3_ENDPOINT: str = ""
S3_BUCKET: str = ""
S3_ACCESS_KEY: str = ""
S3_SECRET_KEY: str = ""
S3_REGION: str = "us-east-1"
S3_SECURE: bool = False
S3_URL_TTL_S: int = 900
S3_PREFIX: str = "comfy-results"
S3_FORCE_PATH_STYLE: bool = True
MINIO_ENDPOINT: str = ""
MINIO_BUCKET: str = ""
MINIO_ACCESS_KEY: str = ""
MINIO_SECRET_KEY: str = ""
LTX_CKPT_NAME: str = "ltx-2-19b-distilled-fp8.safetensors"
LTX_TEXT_ENCODER: str = ""
LTX_DEVICE: str = "default"
LTX_SAMPLER: str = "euler"
LTX_MAX_SHIFT: float = 2.05
LTX_BASE_SHIFT: float = 0.95
LTX_TERMINAL: float = 0.1
LTX_STRETCH: bool = True
MAX_CONCURRENCY: int = 1 # для LTX-2 стартово краще 1 MAX_CONCURRENCY: int = 1 # для LTX-2 стартово краще 1
IDEMPOTENCY_TTL_S: int = 24 * 60 * 60
IDEMPOTENCY_DB_PATH: str = ""
settings = Settings() settings = Settings()

View File

@@ -0,0 +1,111 @@
import sqlite3
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Literal, Optional
from .config import settings
Decision = Literal["created", "exists", "conflict"]
@dataclass
class IdempotencyResult:
decision: Decision
job_id: str
class IdempotencyStore:
def __init__(self, db_path: str, ttl_s: int) -> None:
self.db_path = db_path
self.ttl_s = max(60, int(ttl_s))
self._lock = threading.Lock()
self._init_db()
def _connect(self) -> sqlite3.Connection:
conn = sqlite3.connect(self.db_path, check_same_thread=False)
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA synchronous=NORMAL")
return conn
def _init_db(self) -> None:
Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
with self._connect() as conn:
conn.execute(
"""
CREATE TABLE IF NOT EXISTS idempotency_jobs (
idem_key TEXT NOT NULL,
gen_type TEXT NOT NULL,
req_hash TEXT NOT NULL,
job_id TEXT NOT NULL,
created_at INTEGER NOT NULL,
expires_at INTEGER NOT NULL,
PRIMARY KEY (idem_key, gen_type)
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_idem_expires_at ON idempotency_jobs(expires_at)"
)
conn.commit()
def reserve(
self,
*,
idem_key: str,
gen_type: str,
req_hash: str,
job_id: str,
) -> IdempotencyResult:
now = int(time.time())
expires = now + self.ttl_s
with self._lock, self._connect() as conn:
conn.execute("DELETE FROM idempotency_jobs WHERE expires_at < ?", (now,))
row = conn.execute(
"""
SELECT req_hash, job_id, expires_at
FROM idempotency_jobs
WHERE idem_key = ? AND gen_type = ?
""",
(idem_key, gen_type),
).fetchone()
if row:
existing_hash, existing_job_id, existing_exp = row
if existing_exp >= now:
if existing_hash == req_hash:
return IdempotencyResult(decision="exists", job_id=existing_job_id)
return IdempotencyResult(decision="conflict", job_id=existing_job_id)
conn.execute(
"DELETE FROM idempotency_jobs WHERE idem_key = ? AND gen_type = ?",
(idem_key, gen_type),
)
conn.execute(
"""
INSERT INTO idempotency_jobs
(idem_key, gen_type, req_hash, job_id, created_at, expires_at)
VALUES (?, ?, ?, ?, ?, ?)
""",
(idem_key, gen_type, req_hash, job_id, now, expires),
)
conn.commit()
return IdempotencyResult(decision="created", job_id=job_id)
IDEMPOTENCY_STORE: Optional[IdempotencyStore] = None
def init_idempotency_store() -> None:
global IDEMPOTENCY_STORE
if IDEMPOTENCY_STORE is not None:
return
db_path = settings.IDEMPOTENCY_DB_PATH
if not db_path:
db_path = str(Path(settings.STORAGE_PATH) / "idempotency.sqlite3")
IDEMPOTENCY_STORE = IdempotencyStore(db_path=db_path, ttl_s=settings.IDEMPOTENCY_TTL_S)

View File

@@ -7,8 +7,12 @@ class JobStore:
def __init__(self) -> None: def __init__(self) -> None:
self._jobs: Dict[str, JobStatus] = {} self._jobs: Dict[str, JobStatus] = {}
def create(self, gen_type: GenType) -> JobStatus: @staticmethod
job_id = f"job_{uuid.uuid4().hex}" def new_job_id() -> str:
return f"job_{uuid.uuid4().hex}"
def create(self, gen_type: GenType, job_id: Optional[str] = None) -> JobStatus:
job_id = job_id or self.new_job_id()
js = JobStatus(job_id=job_id, type=gen_type, status="queued", progress=0.0) js = JobStatus(job_id=job_id, type=gen_type, status="queued", progress=0.0)
self._jobs[job_id] = js self._jobs[job_id] = js
return js return js

View File

@@ -6,7 +6,8 @@ from .config import settings
from .api import router from .api import router
from .worker import worker_loop from .worker import worker_loop
from .nats_client import start_nats from .nats_client import start_nats
from .storage import ensure_storage from .storage import ensure_storage, init_object_storage
from .idempotency import init_idempotency_store
app = FastAPI(title="Comfy Agent Service", version="0.1.0") app = FastAPI(title="Comfy Agent Service", version="0.1.0")
app.include_router(router) app.include_router(router)
@@ -14,6 +15,8 @@ app.include_router(router)
@app.on_event("startup") @app.on_event("startup")
async def startup(): async def startup():
ensure_storage() ensure_storage()
init_object_storage()
init_idempotency_store()
# Static files for result URLs: /files/{job_id}/... # Static files for result URLs: /files/{job_id}/...
app.mount("/files", StaticFiles(directory=settings.STORAGE_PATH), name="files") app.mount("/files", StaticFiles(directory=settings.STORAGE_PATH), name="files")

View File

@@ -11,15 +11,24 @@ class GenerateImageRequest(BaseModel):
height: int = 1024 height: int = 1024
steps: int = 28 steps: int = 28
seed: Optional[int] = None seed: Optional[int] = None
idempotency_key: Optional[str] = None
workflow: Optional[str] = None workflow: Optional[str] = None
workflow_params: Dict[str, Any] = Field(default_factory=dict) workflow_params: Dict[str, Any] = Field(default_factory=dict)
class GenerateVideoRequest(BaseModel): class GenerateVideoRequest(BaseModel):
prompt: str = Field(min_length=1) prompt: str = Field(min_length=1)
negative_prompt: Optional[str] = None
width: int = 768
height: int = 512
frames: Optional[int] = None
seconds: int = 4 seconds: int = 4
fps: int = 24 fps: int = 24
steps: int = 30 steps: int = 30
cfg: float = 2.5
seed: Optional[int] = None seed: Optional[int] = None
format: str = "mp4"
codec: str = "h264"
idempotency_key: Optional[str] = None
workflow: Optional[str] = None workflow: Optional[str] = None
workflow_params: Dict[str, Any] = Field(default_factory=dict) workflow_params: Dict[str, Any] = Field(default_factory=dict)

View File

@@ -1,10 +1,21 @@
# services/comfy-agent/app/nats_client.py # services/comfy-agent/app/nats_client.py
import json import json
import asyncio import asyncio
import hashlib
from nats.aio.client import Client as NATS from nats.aio.client import Client as NATS
from .config import settings from .config import settings
from .jobs import JOB_STORE from .jobs import JOB_STORE
from .worker import enqueue from .worker import enqueue
from . import idempotency
def _hash_payload(gen_type: str, workflow: dict) -> str:
normalized = json.dumps(
{"type": gen_type, "workflow": workflow},
sort_keys=True,
separators=(",", ":"),
)
return hashlib.sha256(normalized.encode("utf-8")).hexdigest()
async def start_nats() -> NATS: async def start_nats() -> NATS:
nc = NATS() nc = NATS()
@@ -24,8 +35,24 @@ async def start_nats() -> NATS:
await nc.publish(reply, json.dumps({"error": "missing_workflow"}).encode()) await nc.publish(reply, json.dumps({"error": "missing_workflow"}).encode())
return return
job = JOB_STORE.create(gen_type) idem_key = (payload.get("idempotency_key") or "").strip()
enqueue(job.job_id, gen_type, workflow) if idem_key and idempotency.IDEMPOTENCY_STORE is not None:
result = idempotency.IDEMPOTENCY_STORE.reserve(
idem_key=idem_key,
gen_type=gen_type,
req_hash=_hash_payload(gen_type, workflow),
job_id=JOB_STORE.new_job_id(),
)
if result.decision == "conflict":
if reply:
await nc.publish(reply, json.dumps({"error": "idempotency_key_reused_with_different_payload"}).encode())
return
job = JOB_STORE.get(result.job_id) or JOB_STORE.create(gen_type, job_id=result.job_id)
if result.decision == "created":
enqueue(job.job_id, gen_type, workflow)
else:
job = JOB_STORE.create(gen_type)
enqueue(job.job_id, gen_type, workflow)
if reply: if reply:
await nc.publish(reply, json.dumps({"job_id": job.job_id}).encode()) await nc.publish(reply, json.dumps({"job_id": job.job_id}).encode())

View File

@@ -1,16 +1,117 @@
# services/comfy-agent/app/storage.py # services/comfy-agent/app/storage.py
import mimetypes
import os import os
from pathlib import Path from pathlib import Path
from .config import settings from .config import settings
try:
import boto3
from botocore.client import Config
except Exception: # pragma: no cover - optional runtime dependency fallback
boto3 = None
Config = None
_s3_client = None
_s3_bucket: str | None = None
_s3_prefix: str | None = None
def _s3_value(primary: str, fallback: str) -> str:
return (primary or fallback).strip()
def _build_endpoint_url(endpoint: str) -> str:
if endpoint.startswith("http://") or endpoint.startswith("https://"):
return endpoint
scheme = "https" if settings.S3_SECURE else "http"
return f"{scheme}://{endpoint}"
def _ensure_bucket(client, bucket: str) -> None:
try:
client.head_bucket(Bucket=bucket)
return
except Exception:
pass
client.create_bucket(Bucket=bucket)
def init_object_storage() -> None:
global _s3_client, _s3_bucket, _s3_prefix
endpoint = _s3_value(settings.S3_ENDPOINT, settings.MINIO_ENDPOINT)
bucket = _s3_value(settings.S3_BUCKET, settings.MINIO_BUCKET)
access_key = _s3_value(settings.S3_ACCESS_KEY, settings.MINIO_ACCESS_KEY)
secret_key = _s3_value(settings.S3_SECRET_KEY, settings.MINIO_SECRET_KEY)
if not endpoint or not bucket or not access_key or not secret_key:
_s3_client = None
_s3_bucket = None
_s3_prefix = None
return
if boto3 is None or Config is None:
print("S3 storage disabled: boto3 is not installed")
return
endpoint_url = _build_endpoint_url(endpoint)
s3_config = Config(
signature_version="s3v4",
s3={"addressing_style": "path" if settings.S3_FORCE_PATH_STYLE else "auto"},
)
client = boto3.client(
"s3",
endpoint_url=endpoint_url,
region_name=settings.S3_REGION,
aws_access_key_id=access_key,
aws_secret_access_key=secret_key,
use_ssl=settings.S3_SECURE,
config=s3_config,
)
_ensure_bucket(client, bucket)
_s3_client = client
_s3_bucket = bucket
_s3_prefix = settings.S3_PREFIX.strip("/")
print(f"S3 storage enabled: endpoint={endpoint_url} bucket={bucket}")
def ensure_storage() -> None: def ensure_storage() -> None:
Path(settings.STORAGE_PATH).mkdir(parents=True, exist_ok=True) Path(settings.STORAGE_PATH).mkdir(parents=True, exist_ok=True)
def make_job_dir(job_id: str) -> str: def make_job_dir(job_id: str) -> str:
ensure_storage() ensure_storage()
d = os.path.join(settings.STORAGE_PATH, job_id) d = os.path.join(settings.STORAGE_PATH, job_id)
Path(d).mkdir(parents=True, exist_ok=True) Path(d).mkdir(parents=True, exist_ok=True)
return d return d
def public_url(job_id: str, filename: str) -> str: def public_url(job_id: str, filename: str) -> str:
return f"{settings.PUBLIC_BASE_URL}/{job_id}/{filename}" return f"{settings.PUBLIC_BASE_URL}/{job_id}/{filename}"
def publish_result_url(job_id: str, filename: str, local_path: str) -> str:
if _s3_client is None or _s3_bucket is None:
return public_url(job_id, filename)
object_key = f"{job_id}/{filename}"
if _s3_prefix:
object_key = f"{_s3_prefix}/{object_key}"
content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
try:
_s3_client.upload_file(
local_path,
_s3_bucket,
object_key,
ExtraArgs={"ContentType": content_type},
)
return _s3_client.generate_presigned_url(
ClientMethod="get_object",
Params={"Bucket": _s3_bucket, "Key": object_key},
ExpiresIn=settings.S3_URL_TTL_S,
)
except Exception as e:
print(f"S3 upload/presign failed for job={job_id}: {e}")
return public_url(job_id, filename)

View File

@@ -5,7 +5,7 @@ import os
import json import json
from typing import Any, Dict, Optional, Tuple from typing import Any, Dict, Optional, Tuple
from .jobs import JOB_STORE from .jobs import JOB_STORE
from .storage import make_job_dir, public_url from .storage import make_job_dir, publish_result_url
from .comfyui_client import ComfyUIClient from .comfyui_client import ComfyUIClient
from .config import settings from .config import settings
@@ -14,15 +14,67 @@ _queue: "asyncio.Queue[Tuple[str, str, Dict[str, Any]]]" = asyncio.Queue()
def enqueue(job_id: str, gen_type: str, prompt_graph: Dict[str, Any]) -> None: def enqueue(job_id: str, gen_type: str, prompt_graph: Dict[str, Any]) -> None:
_queue.put_nowait((job_id, gen_type, prompt_graph)) _queue.put_nowait((job_id, gen_type, prompt_graph))
async def _extract_first_output(history: Dict[str, Any], job_dir: str) -> Optional[str]: async def _extract_first_output(history: Dict[str, Any], job_dir: str, client: ComfyUIClient) -> Optional[str]:
# ComfyUI history structure can vary; implement a conservative extraction: # Keep full history for debugging/reproducibility.
# try to find any "images" or "gifs"/"videos" outputs and download via /view
# For MVP: prefer /view?filename=...&type=output&subfolder=...
# Here we return a "manifest.json" to unblock integration even if file fetching needs refinement.
manifest_path = os.path.join(job_dir, "manifest.json") manifest_path = os.path.join(job_dir, "manifest.json")
with open(manifest_path, "w", encoding="utf-8") as f: with open(manifest_path, "w", encoding="utf-8") as f:
json.dump(history, f, ensure_ascii=False, indent=2) json.dump(history, f, ensure_ascii=False, indent=2)
return "manifest.json"
def _iter_assets() -> list[Dict[str, Any]]:
assets: list[Dict[str, Any]] = []
for prompt_data in history.values():
outputs = prompt_data.get("outputs", {}) if isinstance(prompt_data, dict) else {}
for node_out in outputs.values():
if not isinstance(node_out, dict):
continue
for key in ("images", "gifs", "videos"):
for item in node_out.get(key, []) or []:
if isinstance(item, dict) and item.get("filename"):
assets.append(item)
return assets
assets = _iter_assets()
if not assets:
return None
first = assets[0]
filename = os.path.basename(first.get("filename", "output.bin"))
params = {
"filename": first.get("filename"),
"subfolder": first.get("subfolder", ""),
"type": first.get("type", "output"),
}
try:
resp = await client.http.get("/view", params=params)
resp.raise_for_status()
out_path = os.path.join(job_dir, filename)
with open(out_path, "wb") as f:
f.write(resp.content)
return filename
except Exception:
# Fallback remains manifest-only if /view download fails.
return None
def _extract_history_error(history: Dict[str, Any]) -> Optional[str]:
for prompt_data in history.values():
if not isinstance(prompt_data, dict):
continue
status = prompt_data.get("status", {})
if isinstance(status, dict) and status.get("status_str") == "error":
messages = status.get("messages", [])
for item in messages:
if not (isinstance(item, list) and len(item) >= 2):
continue
if item[0] != "execution_error":
continue
payload = item[1] if isinstance(item[1], dict) else {}
msg = payload.get("exception_message") or payload.get("exception_type")
if msg:
return str(msg).strip()
return "comfy_execution_error"
return None
async def worker_loop() -> None: async def worker_loop() -> None:
client = ComfyUIClient() client = ComfyUIClient()
@@ -44,16 +96,22 @@ async def worker_loop() -> None:
hist = await client.get_history(prompt_id) hist = await client.get_history(prompt_id)
job_dir = make_job_dir(job_id) job_dir = make_job_dir(job_id)
fname = await _extract_first_output(hist, job_dir) hist_error = _extract_history_error(hist)
if hist_error:
await _extract_first_output(hist, job_dir, client)
JOB_STORE.update(job_id, status="failed", message="failed", error=hist_error)
return
fname = await _extract_first_output(hist, job_dir, client)
if not fname: if not fname:
JOB_STORE.update(job_id, status="failed", error="No outputs found in ComfyUI history") JOB_STORE.update(job_id, status="failed", message="failed", error="No outputs found in ComfyUI history")
return return
url = public_url(job_id, fname) local_path = os.path.join(job_dir, fname)
url = publish_result_url(job_id, fname, local_path)
JOB_STORE.update(job_id, status="succeeded", progress=1.0, result_url=url) JOB_STORE.update(job_id, status="succeeded", progress=1.0, result_url=url)
except Exception as e: except Exception as e:
JOB_STORE.update(job_id, status="failed", error=str(e)) JOB_STORE.update(job_id, status="failed", message="failed", error=str(e))
while True: while True:
job_id, gen_type, prompt_graph = await _queue.get() job_id, gen_type, prompt_graph = await _queue.get()

View File

@@ -7,3 +7,4 @@ websockets==12.0
nats-py==2.7.2 nats-py==2.7.2
python-multipart==0.0.9 python-multipart==0.0.9
orjson==3.10.7 orjson==3.10.7
boto3==1.35.36

View File

@@ -25,6 +25,7 @@ logger = logging.getLogger(__name__)
# These are set by main.py at startup # These are set by main.py at startup
_state: LatestState | None = None _state: LatestState | None = None
_stats_fn = None # callable → dict _stats_fn = None # callable → dict
_features_cache: dict[str, dict] = {} # symbol → last computed features
def set_state(state: LatestState) -> None: def set_state(state: LatestState) -> None:
@@ -37,6 +38,11 @@ def set_stats_fn(fn) -> None:
_stats_fn = fn _stats_fn = fn
def cache_features(symbol: str, features: dict) -> None:
"""Cache pre-computed features for fast API responses."""
_features_cache[symbol] = features
async def _handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter): async def _handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
"""Minimal HTTP request handler.""" """Minimal HTTP request handler."""
try: try:
@@ -117,15 +123,19 @@ async def _route(
return body, "application/json", "200 OK" return body, "application/json", "200 OK"
elif path == "/features/latest": elif path == "/features/latest":
symbol = params.get("symbol", "") symbol = params.get("symbol", "").upper()
if not symbol: if not symbol:
body = json.dumps({"error": "missing ?symbol=XXX"}).encode() body = json.dumps({"error": "missing ?symbol=XXX"}).encode()
return body, "application/json", "400 Bad Request" return body, "application/json", "400 Bad Request"
if not _state: cached = _features_cache.get(symbol)
if cached:
data = {"symbol": symbol, "features": cached}
elif _state:
# Fallback to live compute (slower)
data = {"symbol": symbol, "features": compute_features(_state, symbol)}
else:
body = json.dumps({"error": "not initialized"}).encode() body = json.dumps({"error": "not initialized"}).encode()
return body, "application/json", "503 Service Unavailable" return body, "application/json", "503 Service Unavailable"
features = compute_features(_state, symbol)
data = {"symbol": symbol.upper(), "features": features}
body = json.dumps(data, ensure_ascii=False).encode() body = json.dumps(data, ensure_ascii=False).encode()
return body, "application/json", "200 OK" return body, "application/json", "200 OK"

View File

@@ -79,12 +79,24 @@ async def process_events(
events_per_sec_count = 0 events_per_sec_count = 0
time.monotonic() time.monotonic()
feature_compute_interval = 1.0 / max(settings.features_pub_rate_hz, 1.0)
next_feature_compute: dict[str, float] = {}
next_signal_emit: dict[str, float] = {}
signal_cooldown_sec = 1.0
batch_counter = 0
while True: while True:
try: try:
event = await consumer.queue.get() event = await consumer.queue.get()
except asyncio.CancelledError: except asyncio.CancelledError:
break break
# Yield to event loop every N events so HTTP API stays responsive
batch_counter += 1
if batch_counter >= 5:
batch_counter = 0
await asyncio.sleep(0)
proc_start = time.monotonic() proc_start = time.monotonic()
try: try:
@@ -110,15 +122,24 @@ async def process_events(
else: else:
symbol = None symbol = None
# Compute features + publish (only for trade/quote events) # Compute features + publish with per-symbol throttling
if symbol and settings.features_enabled: if symbol and settings.features_enabled:
snapshot = make_feature_snapshot(state, symbol) now_mono = time.monotonic()
await publisher.publish_features(snapshot) due = next_feature_compute.get(symbol, 0.0)
if now_mono >= due:
snapshot = make_feature_snapshot(state, symbol)
# Cache for fast HTTP API responses
api.cache_features(symbol, snapshot.features)
await publisher.publish_features(snapshot)
# Check for trade signal # Check for trade signal with cooldown to avoid flood
sig = check_signal(snapshot.features, symbol) sig = check_signal(snapshot.features, symbol)
if sig: sig_due = next_signal_emit.get(symbol, 0.0)
await publisher.publish_signal(sig) if sig and now_mono >= sig_due:
await publisher.publish_signal(sig)
next_signal_emit[symbol] = now_mono + signal_cooldown_sec
next_feature_compute[symbol] = now_mono + feature_compute_interval
# Processing latency metric # Processing latency metric
proc_ms = (time.monotonic() - proc_start) * 1000 proc_ms = (time.monotonic() - proc_start) * 1000

View File

@@ -78,7 +78,7 @@ class Publisher:
symbol=signal.symbol, symbol=signal.symbol,
direction=signal.direction, direction=signal.direction,
).inc() ).inc()
logger.info( logger.debug(
"publisher.signal_emitted", "publisher.signal_emitted",
extra={ extra={
"symbol": signal.symbol, "symbol": signal.symbol,

View File

@@ -14,6 +14,7 @@ pillow>=10.0.0
tiktoken>=0.5.0 tiktoken>=0.5.0
sentencepiece>=0.1.99 sentencepiece>=0.1.99
einops>=0.7.0 einops>=0.7.0
verovio>=4.3.1
# STT (Speech-to-Text) dependencies # STT (Speech-to-Text) dependencies
faster-whisper>=1.0.0 faster-whisper>=1.0.0
@@ -41,4 +42,4 @@ googlesearch-python>=1.2.0
pdfplumber>=0.10.0 pdfplumber>=0.10.0
python-docx>=1.1.0 python-docx>=1.1.0
openpyxl>=3.1.2 openpyxl>=3.1.2
chardet>=5.2.0 chardet>=5.2.0