Move annotation downloading after successful validation

Author: zhiltsov-max

packages/examples/cvat/exchange-oracle/src/core/annotation_meta.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from pydantic import BaseModel
 
 ANNOTATION_RESULTS_METAFILE_NAME = "annotation_meta.json"
@@ -9,7 +7,6 @@
 class JobMeta(BaseModel):
     job_id: int
     task_id: int
-    annotation_filename: Path
     annotator_wallet_address: str
     assignment_id: str
     start_frame: int

packages/examples/cvat/exchange-oracle/src/core/oracle_events.py

@@ -47,6 +47,10 @@ class ExchangeOracleEvent_JobFinished(OracleEvent):
     pass  # escrow is enough for now
 
 
+class ExchangeOracleEvent_EscrowRecorded(OracleEvent):
+    pass  # escrow is enough for now
+
+
 class ExchangeOracleEvent_EscrowCleaned(OracleEvent):
     pass
 
@@ -56,10 +60,12 @@ class ReputationOracleEvent_EscrowCompleted(OracleEvent):
 
 
 _event_type_map = {
+    # TODO: make sender-dependent
     JobLauncherEventTypes.escrow_created: JobLauncherEvent_EscrowCreated,
     JobLauncherEventTypes.escrow_canceled: JobLauncherEvent_EscrowCanceled,
     RecordingOracleEventTypes.job_completed: RecordingOracleEvent_JobCompleted,
     RecordingOracleEventTypes.submission_rejected: RecordingOracleEvent_SubmissionRejected,
+    ExchangeOracleEventTypes.escrow_recorded: ExchangeOracleEvent_EscrowRecorded,
     ExchangeOracleEventTypes.escrow_failed: ExchangeOracleEvent_EscrowFailed,
     ExchangeOracleEventTypes.job_finished: ExchangeOracleEvent_JobFinished,
     ExchangeOracleEventTypes.escrow_cleaned: ExchangeOracleEvent_EscrowCleaned,
@@ -68,6 +74,7 @@ class ReputationOracleEvent_EscrowCompleted(OracleEvent):
 
 
 def get_class_for_event_type(event_type: str) -> type[OracleEvent]:
+    # TODO: make sender-dependent
     event_class = next((v for k, v in _event_type_map.items() if k == event_type), None)
 
     if not event_class:

packages/examples/cvat/exchange-oracle/src/core/types.py

@@ -65,6 +65,7 @@ class ExchangeOracleEventTypes(str, Enum, metaclass=BetterEnumMeta):
     escrow_failed = "escrow_failed"
     job_finished = "job_finished"
     escrow_cleaned = "escrow_cleaned"
+    escrow_recorded = "escrow_recorded"
 
 
 class JobLauncherEventTypes(str, Enum, metaclass=BetterEnumMeta):
@@ -78,7 +79,6 @@ class RecordingOracleEventTypes(str, Enum, metaclass=BetterEnumMeta):
 
 
 class ReputationOracleEventTypes(str, Enum, metaclass=BetterEnumMeta):
-    # TODO: rename to ReputationOracleEventType
     escrow_completed = "escrow_completed"
 
 

packages/examples/cvat/exchange-oracle/src/crons/__init__.py

@@ -16,6 +16,7 @@
     process_outgoing_job_launcher_webhooks,
 )
 from src.crons.webhooks.recording_oracle import (
+    process_incoming_recording_oracle_webhook_job_completed,
     process_incoming_recording_oracle_webhooks,
     process_outgoing_recording_oracle_webhooks,
 )
@@ -41,6 +42,11 @@ def cron_record():
             "interval",
             seconds=Config.cron_config.process_recording_oracle_webhooks_int,
         )
+        scheduler.add_job(
+            process_incoming_recording_oracle_webhook_job_completed,
+            "interval",
+            seconds=Config.cron_config.process_recording_oracle_webhooks_int,
+        )
         scheduler.add_job(
             process_outgoing_recording_oracle_webhooks,
             "interval",

packages/examples/cvat/exchange-oracle/src/crons/webhooks/recording_oracle.py

@@ -19,6 +19,7 @@
 from src.crons._cron_job import cron_job
 from src.crons.webhooks._common import handle_webhook, process_outgoing_webhooks
 from src.db.utils import ForUpdateParams
+from src.handlers.completed_escrows import handle_escrow_export
 from src.models.webhook import Webhook
 
 
@@ -30,6 +31,29 @@ def process_incoming_recording_oracle_webhooks(logger: logging.Logger, session:
     webhooks = oracle_db_service.inbox.get_pending_webhooks(
         session,
         OracleWebhookTypes.recording_oracle,
+        event_type_not_in=[RecordingOracleEventTypes.job_completed],
+        limit=CronConfig.process_recording_oracle_webhooks_chunk_size,
+        for_update=ForUpdateParams(skip_locked=True),
+    )
+
+    for webhook in webhooks:
+        with handle_webhook(logger, session, webhook, queue=oracle_db_service.inbox):
+            handle_recording_oracle_event(webhook, db_session=session, logger=logger)
+
+
+@cron_job
+def process_incoming_recording_oracle_webhook_job_completed(
+    logger: logging.Logger, session: Session
+):
+    """
+    Process incoming oracle webhooks of type job_completed
+    We do it in a separate job as this is a long operation that should not block
+    other message handling.
+    """
+    webhooks = oracle_db_service.inbox.get_pending_webhooks(
+        session,
+        OracleWebhookTypes.recording_oracle,
+        event_type_in=[RecordingOracleEventTypes.job_completed],
         limit=CronConfig.process_recording_oracle_webhooks_chunk_size,
         for_update=ForUpdateParams(skip_locked=True),
     )
@@ -70,6 +94,7 @@ def handle_recording_oracle_event(webhook: Webhook, *, db_session: Session, logg
                 )
                 return
 
+            recorded_project_cvat_ids: set[int] = set()
             chunk_size = CronConfig.process_accepted_projects_chunk_size
             for ids_chunk in take_by(project_ids, chunk_size):
                 projects_chunk = cvat_db_service.get_projects_by_cvat_ids(
@@ -89,15 +114,11 @@ def handle_recording_oracle_event(webhook: Webhook, *, db_session: Session, logg
                         )
                         return
 
-                    new_status = ProjectStatuses.recorded
-                    logger.info(
-                        "Changing project status to {} (escrow_address={}, project={})".format(
-                            new_status, webhook.escrow_address, project.cvat_id
-                        )
+                    recorded_project_cvat_ids.add(project.cvat_id)
+                    cvat_db_service.update_project_status(
+                        db_session, project.id, ProjectStatuses.recorded
                     )
 
-                    cvat_db_service.update_project_status(db_session, project.id, new_status)
-
                 cvat_db_service.touch_final_assignments(
                     db_session,
                     cvat_project_ids=[p.cvat_id for p in projects_chunk],
@@ -111,6 +132,27 @@ def handle_recording_oracle_event(webhook: Webhook, *, db_session: Session, logg
                 status=EscrowValidationStatuses.completed,
             )
 
+            logger.info(
+                f"Escrow '{webhook.escrow_address}' is accepted, trying to export annotations..."
+            )
+            handle_escrow_export(
+                logger=logger,
+                session=db_session,
+                escrow_address=webhook.escrow_address,
+                chain_id=webhook.chain_id,
+            )
+
+            if recorded_project_cvat_ids:
+                # Print it after successful export so that the logs were consistent
+                # in the case of export error
+                logger.info(
+                    "Changing project statuses to {} (escrow_address={}, project={})".format(
+                        ProjectStatuses.recorded,
+                        webhook.escrow_address,
+                        sorted(recorded_project_cvat_ids),
+                    )
+                )
+
         case RecordingOracleEventTypes.submission_rejected:
             event = RecordingOracleEvent_SubmissionRejected.model_validate(webhook.event_data)
 

packages/examples/cvat/exchange-oracle/src/handlers/completed_escrows.py

@@ -14,7 +14,10 @@
 from src.chain.escrow import get_escrow_manifest, validate_escrow
 from src.core.annotation_meta import ANNOTATION_RESULTS_METAFILE_NAME, RESULTING_ANNOTATIONS_FILE
 from src.core.config import CronConfig, StorageConfig
-from src.core.oracle_events import ExchangeOracleEvent_JobFinished
+from src.core.oracle_events import (
+    ExchangeOracleEvent_EscrowRecorded,
+    ExchangeOracleEvent_JobFinished,
+)
 from src.core.storage import compose_results_bucket_filename
 from src.core.types import EscrowValidationStatuses, OracleWebhookTypes, TaskTypes
 from src.db import SessionLocal
@@ -68,9 +71,23 @@ def _export_escrow_annotations(
 ) -> None:
     manifest = parse_manifest(get_escrow_manifest(chain_id, escrow_address))
 
-    logger.debug(f"Downloading results for the escrow ({escrow_address=})")
+    escrow_creation = cvat_service.get_escrow_creation_by_escrow_address(
+        session,
+        escrow_address,
+        chain_id,
+        active=False,
+    )
+    if not escrow_creation:
+        raise AssertionError(f"Can't find escrow creation for escrow '{escrow_address}'")
 
     jobs = cvat_service.get_jobs_by_escrow_address(session, escrow_address, chain_id)
+    if len(jobs) != escrow_creation.total_jobs:
+        raise AssertionError(
+            f"Unexpected number of jobs fetched for escrow "
+            f"'{escrow_address}': {len(jobs)}, expected {escrow_creation.total_jobs}"
+        )
+
+    logger.debug(f"Downloading results for the escrow ({escrow_address=})")
 
     annotation_format = CVAT_EXPORT_FORMAT_MAPPING[manifest.annotation.type]
     # FUTURE-TODO: probably can be removed in the future since
@@ -79,7 +96,7 @@ def _export_escrow_annotations(
 
     if manifest.annotation.type == TaskTypes.image_skeletons_from_boxes.value:
         # we'll have to merge annotations ourselves for skeletons
-        # might want to make this the only behaviour in the future
+        # might want to make this the only behavior in the future
         project_annotations_file = None
         project_images = None
     else:
@@ -116,13 +133,13 @@ def _export_escrow_annotations(
         manifest=manifest,
         project_images=project_images,
     )
-    logger.debug(f"Uploading results for the escrow ({escrow_address=})")
+    logger.debug(f"Uploading annotations for the escrow ({escrow_address=})")
 
     _upload_annotations(
         annotation_files=(
             resulting_annotations_file_desc,
             *job_annotations.values(),
-            prepare_annotation_metafile(jobs=jobs, job_annotations=job_annotations),
+            prepare_annotation_metafile(jobs=jobs),
         ),
         chain_id=chain_id,
         escrow_address=escrow_address,
@@ -133,7 +150,7 @@ def _export_escrow_annotations(
         escrow_address=escrow_address,
         chain_id=chain_id,
         type=OracleWebhookTypes.recording_oracle,
-        event=ExchangeOracleEvent_JobFinished(),
+        event=ExchangeOracleEvent_EscrowRecorded(),
     )
 
     logger.info(
@@ -142,6 +159,38 @@ def _export_escrow_annotations(
     )
 
 
+def _request_escrow_validation(
+    logger: logging.Logger,
+    chain_id: int,
+    escrow_address: str,
+    escrow_projects: Sequence[Project],
+    session: Session,
+) -> None:
+    # TODO: lock escrow once there is such a DB object
+    assert escrow_projects  # unused, but must hold a lock
+
+    # TODO: maybe upload only current iteration jobs
+    jobs = cvat_service.get_jobs_by_escrow_address(session, escrow_address, chain_id)
+
+    logger.debug(f"Uploading assignment info for the escrow ({escrow_address=})")
+
+    _upload_annotations(
+        annotation_files=(prepare_annotation_metafile(jobs=jobs),),
+        chain_id=chain_id,
+        escrow_address=escrow_address,
+    )
+
+    oracle_db_service.outbox.create_webhook(
+        session,
+        escrow_address=escrow_address,
+        chain_id=chain_id,
+        type=OracleWebhookTypes.recording_oracle,
+        event=ExchangeOracleEvent_JobFinished(),
+    )
+
+    logger.info(f"The escrow ({escrow_address=}) annotation is finished, " f"requesting validation")
+
+
 def _upload_annotations(
     annotation_files: Sequence[FileDescriptor], chain_id: int, escrow_address: str
 ) -> None:
@@ -249,7 +298,7 @@ def _handle_escrow_validation(
     escrow_projects = cvat_service.get_projects_by_escrow_address(
         session, escrow_address, limit=None, for_update=ForUpdateParams(nowait=True)
     )
-    _export_escrow_annotations(logger, chain_id, escrow_address, escrow_projects, session)
+    _request_escrow_validation(logger, chain_id, escrow_address, escrow_projects, session)
 
 
 def handle_escrows_validations(logger: logging.Logger) -> None:
@@ -281,3 +330,17 @@ def handle_escrows_validations(logger: logging.Logger) -> None:
                 increase_attempts=True,  # increase attempts always to allow escrow rotation
                 **update_kwargs,
             )
+
+
+def handle_escrow_export(
+    logger: logging.Logger,
+    session: Session,
+    escrow_address: str,
+    chain_id: int,
+):
+    validate_escrow(chain_id, escrow_address)
+
+    escrow_projects = cvat_service.get_projects_by_escrow_address(
+        session, escrow_address, limit=None, for_update=ForUpdateParams(nowait=True)
+    )
+    _export_escrow_annotations(logger, chain_id, escrow_address, escrow_projects, session)

packages/examples/cvat/exchange-oracle/src/handlers/cvat_events.py

@@ -142,6 +142,7 @@ def handle_create_job_event(payload: dict) -> None:
                 session,
                 escrow_address=project.escrow_address,
                 chain_id=project.chain_id,
+                active=True,
                 for_update=True,
             )
 

packages/examples/cvat/exchange-oracle/src/handlers/job_export.py

@@ -43,9 +43,7 @@ class FileDescriptor:
     file: io.RawIOBase | None
 
 
-def prepare_annotation_metafile(
-    jobs: list[Job], job_annotations: dict[int, FileDescriptor]
-) -> FileDescriptor:
+def prepare_annotation_metafile(jobs: list[Job]) -> FileDescriptor:
     """
     Prepares a task/project annotation descriptor file with annotator mapping.
     """
@@ -54,7 +52,6 @@ def prepare_annotation_metafile(
         jobs=[
             JobMeta(
                 job_id=job.cvat_id,
-                annotation_filename=job_annotations[job.cvat_id].filename,
                 annotator_wallet_address=job.latest_assignment.user_wallet_address,
                 assignment_id=job.latest_assignment.id,
                 task_id=job.cvat_task_id,

packages/examples/cvat/exchange-oracle/src/services/cvat.py

@@ -439,32 +439,26 @@ def create_escrow_creation(
     return escrow_creation_id
 
 
-def get_escrow_creation_by_id(
-    session: Session,
-    escrow_creation_id: str,
-    *,
-    for_update: bool | ForUpdateParams = False,
-) -> EscrowCreation | None:
-    return (
-        _maybe_for_update(session.query(EscrowCreation), enable=for_update)
-        .where(EscrowCreation.id == escrow_creation_id, EscrowCreation.finished_at.is_(None))
-        .first()
-    )
-
-
 def get_escrow_creation_by_escrow_address(
     session: Session,
     escrow_address: str,
     chain_id: int,
     *,
+    active: bool | None,
     for_update: bool | ForUpdateParams = False,
 ) -> EscrowCreation | None:
+    is_active_filter = []
+    if active is True:
+        is_active_filter = [EscrowCreation.finished_at.is_(None)]
+    elif active is False:
+        is_active_filter = [EscrowCreation.finished_at.is_not(None)]
+
     return (
         _maybe_for_update(session.query(EscrowCreation), enable=for_update)
         .where(
             EscrowCreation.escrow_address == escrow_address,
             EscrowCreation.chain_id == chain_id,
-            EscrowCreation.finished_at.is_(None),
+            *is_active_filter,
         )
         .first()
     )