celery 容灾升级

video-gen-api/app/tasks/generation_create_tasks.py

@@ -2,11 +2,13 @@ from __future__ import annotations
 
 import asyncio
 import json
+import uuid
 from datetime import datetime, timedelta, timezone
 from typing import Any, Optional
 
 from app.config import settings
-from app.enums.celery_queue import CeleryQueue
+from app.enums.celery_queue import CeleryQueue, CeleryTaskName
+from app.enums.celery_runtime import CeleryRuntimeDomain
 from app.enums.generation_status import GenerationRecordPipelineStage
 from app.enums.generation_task import (
     ALLOWED_GENERATION_MODES,
@@ -34,6 +36,7 @@ from app.services.generation.pipeline.owner_service import (
     owner_mode,
     owner_provider_task_id,
     set_owner_provider_task_id,
+    renew_generation_owner_claim_lease,
 )
 from app.services.generation.poll_schedule_service import ensure_video_poll_fields
 from app.services.generation.provider_service import create_provider_task
@@ -41,10 +44,8 @@ from app.services.media_token_usage_snapshot_service import (
     sync_chat_generation_task_media_token_snapshot,
     sync_generation_record_media_token_snapshot,
 )
-from app.services.redis_registry_service import (
-    RedisExecutionLockError,
-    RedisExecutionLockLease,
-)
+from app.services.redis_registry_service import RedisExecutionLockError
+from app.services.celery_runtime.runtime_service import CeleryRuntimeLease, RuntimeIdentity
 from app.tasks.async_runner import run_async
 from app.tasks.celery_app import celery_app
 
@@ -257,13 +258,36 @@ async def _run(
     if effective_attempt is None:
         return
 
-    lease = await RedisExecutionLockLease.acquire(
-        lock_key=_lock_key(normalized_owner_type, task_id, effective_attempt),
-        ttl_seconds=int(settings.GENERATION_CREATE_LOCK_TTL_SECONDS or 600),
-        log_context="generation_create",
-        renew_interval_seconds=int(
-            settings.REDIS_EXECUTION_LOCK_RENEW_INTERVAL_SECONDS or 30
+    lease_token = uuid.uuid4().hex
+
+    async def _renew_db_claim(token: str) -> bool:
+        return await renew_generation_owner_claim_lease(
+            owner_type=normalized_owner_type,
+            owner_id=task_id,
+            attempt_no=effective_attempt,
+            claim_field="provider_create_claim_token",
+            lease_field="provider_create_lease_until",
+            token=token,
+            lease_seconds=int(settings.GENERATION_CREATE_LOCK_TTL_SECONDS or 600),
+        )
+
+    lease = await CeleryRuntimeLease.acquire(
+        identity=RuntimeIdentity(
+            domain=CeleryRuntimeDomain.GENERATION_CREATE.value,
+            owner_type=normalized_owner_type,
+            owner_id=task_id,
+            attempt_no=effective_attempt,
+            task_name=CeleryTaskName.CHATAPI_CREATE.value,
+            queue=CeleryQueue.GEN_CHATAPI_CREATE.value,
         ),
+        lock_key=_lock_key(normalized_owner_type, task_id, effective_attempt),
+        hash_key=settings.GENERATION_CREATE_ACTIVE_REDIS_HASH_KEY,
+        zset_key=settings.GENERATION_CREATE_ACTIVE_REDIS_ZSET_KEY,
+        token=lease_token,
+        ttl_seconds=int(settings.GENERATION_CREATE_LOCK_TTL_SECONDS or 600),
+        heartbeat_interval_seconds=int(settings.REDIS_EXECUTION_LOCK_RENEW_INTERVAL_SECONDS or 30),
+        pipeline_stage=ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
+        db_heartbeat=_renew_db_claim,
     )
     if lease is None:
         return
@@ -320,6 +344,10 @@ async def _run(
                 _stage(owner, ChatGenerationPipelineStage.PREPARING),
                 _stage(owner, ChatGenerationPipelineStage.CREATING_PROVIDER_TASK),
             }
+            if is_image_main:
+                allowed_stages.add(
+                    _stage(owner, ChatGenerationPipelineStage.PROVIDER_RESULT_STAGED)
+                )
             if owner.pipeline_stage not in allowed_stages:
                 return