mirror of
https://github.com/MemeMeow-Studio/MemeMeow.git
synced 2026-10-16 20:31:33 +08:00
infstellar
5ecda39843
- Implement tests for the ImageProcessingWorker to validate task deduplication and policy interactions. - Create tests for operation policy to ensure proper grant handling and error responses. - Add tests for reverse image provider response structure to validate empty and malformed results. - Develop comprehensive tests for scope resolution and service validation to ensure correct behavior in various scenarios.
269 lines
12 KiB
Python
269 lines
12 KiB
Python
"""分页 seed 图片处理 job 并回填当前 scope 的单图文本向量。
|
|
|
|
该脚本位于部署迁移工具目录，读取 PostgreSQL 当前 Meme 记录作为唯一事实来源。
|
|
它使用 keyset 分页和持久化迁移 epoch，单页之外不保存全库图片 ID；外部 embedding
|
|
调用通过服务端配置创建，不接受命令行传入的 provider 密钥或普通 payload 授权字段。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import logging
|
|
import sys
|
|
from collections.abc import Iterator
|
|
from pathlib import Path
|
|
from uuid import UUID
|
|
|
|
from sqlalchemy import func, or_, select
|
|
|
|
if __package__ in {None, ""}:
|
|
# 允许维护者直接执行脚本，同时保持 ``python -m`` 的标准入口。
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from backend.config import Settings
|
|
from backend.database import DatabaseError, DatabaseResources, ImageProcessingStage, Meme, StorageOperation, Task, create_engine_for_settings
|
|
from backend.image_processing import ImageProcessingError, ImageProcessingRepository, SingleImageEmbeddingService
|
|
from backend.metadata import MetadataError
|
|
from backend.pg_services import PostgresMetadataService, PostgresSearchService
|
|
from backend.visual import identity_from_settings
|
|
|
|
|
|
LOGGER = logging.getLogger("mememeow.image_backfill")
|
|
DEFAULT_PAGE_SIZE = 100
|
|
MAX_PAGE_SIZE = 500
|
|
IMAGE_TASK_STAGES = {
|
|
"visual_embedding_generation": "visual",
|
|
"meme_context_generation": "agent",
|
|
"text_embedding_generation": "text_embedding",
|
|
}
|
|
|
|
|
|
def backfill_task_sources(resources: DatabaseResources, scope_id: str) -> dict[str, int]:
|
|
"""按可信阶段关系回填 pipeline 来源，并标记其余历史阶段为未归类。
|
|
|
|
只有 ``image_processing_stages.task_id`` 的复合 scope 关系能证明父 Job
|
|
归属；孤立任务不补写 standalone 或 pipeline，保留 NULL 来源供只读诊断。
|
|
"""
|
|
linked = 0
|
|
unclassified = 0
|
|
with resources.factory() as session:
|
|
linked_rows = list(
|
|
session.execute(
|
|
select(ImageProcessingStage, Task)
|
|
.join(
|
|
Task,
|
|
(Task.scope_id == ImageProcessingStage.scope_id)
|
|
& (Task.id == ImageProcessingStage.task_id),
|
|
)
|
|
.where(
|
|
ImageProcessingStage.scope_id == scope_id,
|
|
ImageProcessingStage.task_id.is_not(None),
|
|
)
|
|
)
|
|
)
|
|
linked_task_ids: set[str] = set()
|
|
for stage, task in linked_rows:
|
|
expected = IMAGE_TASK_STAGES.get(task.task_type)
|
|
if expected != stage.stage:
|
|
continue
|
|
linked_task_ids.add(task.id)
|
|
if task.submission_mode is None and task.processing_job_id is None:
|
|
task.submission_mode = "pipeline"
|
|
task.image_stage = stage.stage
|
|
task.processing_job_id = stage.job_id
|
|
payload = dict(task.payload or {})
|
|
payload.update({"submission_mode": "pipeline", "stage": stage.stage, "job_id": str(stage.job_id)})
|
|
task.payload = payload
|
|
linked += 1
|
|
|
|
historical = list(
|
|
session.scalars(
|
|
select(Task).where(
|
|
Task.scope_id == scope_id,
|
|
Task.task_type.in_(tuple(IMAGE_TASK_STAGES)),
|
|
Task.submission_mode.is_(None),
|
|
)
|
|
)
|
|
)
|
|
for task in historical:
|
|
if task.id in linked_task_ids:
|
|
continue
|
|
stage = IMAGE_TASK_STAGES[task.task_type]
|
|
if task.image_stage is None:
|
|
task.image_stage = stage
|
|
payload = dict(task.payload or {})
|
|
payload.setdefault("stage", stage)
|
|
task.payload = payload
|
|
unclassified += 1
|
|
session.commit()
|
|
return {"pipeline_linked": linked, "unclassified": unclassified}
|
|
|
|
|
|
def _visible_filter(scope_id: str):
|
|
"""返回排除正在跨存储提交的 Meme 的 scope 条件。"""
|
|
active_operation = select(StorageOperation.id).where(
|
|
StorageOperation.scope_id == scope_id,
|
|
StorageOperation.meme_id == Meme.id,
|
|
StorageOperation.status.in_(("prepared", "file_applied")),
|
|
).exists()
|
|
return (Meme.scope_id == scope_id, ~active_operation)
|
|
|
|
|
|
def count_current_memes(resources: DatabaseResources, scope_id: str) -> int:
|
|
"""统计当前 scope 可迁移图片数量，不读取全库主键列表。"""
|
|
with resources.factory() as session:
|
|
return int(session.scalar(select(func.count()).select_from(Meme).where(*_visible_filter(scope_id))) or 0)
|
|
|
|
|
|
def iter_meme_pages(resources: DatabaseResources, scope_id: str, *, page_size: int = DEFAULT_PAGE_SIZE) -> Iterator[list[Meme]]:
|
|
"""按 storage key 和 meme UUID 稳定 keyset 分页返回当前 scope 图片。"""
|
|
page_size = max(1, min(int(page_size), MAX_PAGE_SIZE))
|
|
last_key: str | None = None
|
|
last_id: UUID | None = None
|
|
while True:
|
|
with resources.factory() as session:
|
|
conditions = list(_visible_filter(scope_id))
|
|
if last_key is not None and last_id is not None:
|
|
conditions.append(or_(Meme.storage_key > last_key, (Meme.storage_key == last_key) & (Meme.id > last_id)))
|
|
rows = list(
|
|
session.scalars(
|
|
select(Meme).where(*conditions).order_by(Meme.storage_key.asc(), Meme.id.asc()).limit(page_size)
|
|
)
|
|
)
|
|
if not rows:
|
|
return
|
|
yield rows
|
|
last_key = rows[-1].storage_key
|
|
last_id = rows[-1].id
|
|
|
|
|
|
def _processing_config(settings: Settings) -> dict[str, object]:
|
|
"""构造迁移 job 使用的服务端配置指纹输入。"""
|
|
identity = identity_from_settings(settings)
|
|
return {
|
|
"embedding_model": settings.embedding_model,
|
|
"embedding_dimensions": 1024,
|
|
"visual_model": identity.model,
|
|
"visual_dimensions": identity.dimensions,
|
|
"preprocess_version": identity.preprocess_version,
|
|
"settings_version": settings.settings_version,
|
|
"migration_source": "image-backfill",
|
|
}
|
|
|
|
|
|
def run_backfill(
|
|
resources: DatabaseResources,
|
|
settings: Settings,
|
|
*,
|
|
scope_id: str = "local",
|
|
page_size: int = DEFAULT_PAGE_SIZE,
|
|
seed_only: bool = False,
|
|
switch: bool = False,
|
|
) -> dict[str, object]:
|
|
"""执行一轮分页 seed/向量回填，并按需原子切换迁移来源。"""
|
|
total = count_current_memes(resources, scope_id)
|
|
config = _processing_config(settings)
|
|
jobs = ImageProcessingRepository(resources, scope_id)
|
|
metadata = PostgresMetadataService(resources, scope_id=scope_id)
|
|
search = PostgresSearchService(settings, resources, metadata, scope_id=scope_id)
|
|
embedder = None if seed_only else search._embedding
|
|
embedding_service = None if embedder is None else SingleImageEmbeddingService(resources, scope_id=scope_id, model=settings.embedding_model, dimensions=1024, embedder=embedder)
|
|
|
|
with resources.environment(scope_id) as environment:
|
|
state = environment.search.begin_incremental_backfill(settings.embedding_model, total_count=total)
|
|
epoch = state.epoch
|
|
|
|
processed = 0
|
|
seeded = 0
|
|
embedded = 0
|
|
failures: list[dict[str, str]] = []
|
|
for page in iter_meme_pages(resources, scope_id, page_size=page_size):
|
|
for meme in page:
|
|
processed += 1
|
|
try:
|
|
image = metadata.blob_store.resolve(meme.storage_key)
|
|
record = metadata.embedding_record(image)
|
|
metadata_hash = record.get("metadata_hash") if isinstance(record.get("metadata_hash"), str) else None
|
|
jobs.create_or_reuse(
|
|
meme.id,
|
|
meme.sha256,
|
|
metadata_hash=metadata_hash,
|
|
config=config,
|
|
reverse_image_policy="forbid",
|
|
explicit_retry=False,
|
|
)
|
|
seeded += 1
|
|
if embedding_service is not None:
|
|
if not record.get("indexable") or not metadata_hash:
|
|
raise ImageProcessingError("query_embedding_not_ready")
|
|
embedding_service.upsert(
|
|
meme.id,
|
|
image_sha256=meme.sha256,
|
|
metadata_hash=metadata_hash,
|
|
semantic_document=str(record.get("text") or ""),
|
|
)
|
|
embedded += 1
|
|
except (DatabaseError, MetadataError, ImageProcessingError, OSError, ValueError, TypeError) as exc:
|
|
failures.append({"meme_id": str(meme.id), "error": getattr(exc, "code", type(exc).__name__)})
|
|
with resources.environment(scope_id) as environment:
|
|
if not environment.search.record_incremental_backfill(epoch=epoch, completed_count=processed, total_count=total, model=settings.embedding_model):
|
|
raise DatabaseError("migration_epoch_stale")
|
|
LOGGER.info("image_backfill_progress processed=%d total=%d seeded=%d embedded=%d failed=%d", processed, total, seeded, embedded, len(failures))
|
|
|
|
switched = False
|
|
if switch and not seed_only and processed == total and not failures:
|
|
with resources.environment(scope_id) as environment:
|
|
switched = environment.search.switch_incremental_only(epoch=epoch, model=settings.embedding_model)
|
|
source_backfill = backfill_task_sources(resources, scope_id)
|
|
return {
|
|
"scope_id": scope_id,
|
|
"model": settings.embedding_model,
|
|
"epoch": epoch,
|
|
"total": total,
|
|
"processed": processed,
|
|
"seeded": seeded,
|
|
"embedded": embedded,
|
|
"failed": len(failures),
|
|
"failures": failures[:100],
|
|
"switched": switched,
|
|
"mode": "incremental_only" if switched else "backfill",
|
|
**source_backfill,
|
|
}
|
|
|
|
|
|
def _parser() -> argparse.ArgumentParser:
|
|
"""构造命令行参数解析器。"""
|
|
parser = argparse.ArgumentParser(description="分页 seed 图片处理 job 并回填单图文本向量")
|
|
parser.add_argument("--scope-id", default="local")
|
|
parser.add_argument("--page-size", type=int, default=DEFAULT_PAGE_SIZE)
|
|
parser.add_argument("--seed-only", action="store_true", help="只创建图片 job，不调用 embedding provider")
|
|
parser.add_argument("--switch", action="store_true", help="全部图片成功回填后切换到 incremental_only")
|
|
parser.add_argument("--log-level", default="INFO", choices=("DEBUG", "INFO", "WARNING", "ERROR"))
|
|
return parser
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
"""执行迁移命令并返回进程退出码。"""
|
|
parser = _parser()
|
|
args = parser.parse_args(argv)
|
|
if args.page_size < 1 or args.page_size > MAX_PAGE_SIZE:
|
|
parser.error(f"--page-size 必须在 1 到 {MAX_PAGE_SIZE} 之间")
|
|
logging.basicConfig(level=getattr(logging, args.log_level), format="%(levelname)s %(name)s %(message)s")
|
|
settings = Settings(_env_file=None)
|
|
engine = create_engine_for_settings(settings)
|
|
resources = None
|
|
try:
|
|
resources = DatabaseResources(engine, image_root=settings.image_root, data_root=settings.data_root, settings=settings, require_local_scope=False)
|
|
result = run_backfill(resources, settings, scope_id=args.scope_id, page_size=args.page_size, seed_only=args.seed_only, switch=args.switch)
|
|
LOGGER.info("image_backfill_finished processed=%s embedded=%s failed=%s switched=%s", result["processed"], result["embedded"], result["failed"], result["switched"])
|
|
return 0 if not result["failed"] else 2
|
|
except (DatabaseError, OSError, ValueError) as exc:
|
|
LOGGER.error("image_backfill_failed error=%s", getattr(exc, "code", type(exc).__name__))
|
|
return 1
|
|
finally:
|
|
engine.dispose()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|