Merge remote-tracking branch 'origin/main'

broija · broija · commit 0b00cba0397a · 2025-11-07T16:34:58.000+01:00
diff --git a/dispatcher/api/constants/mvp.py b/dispatcher/api/constants/mvp.py
@@ -0,0 +1,11 @@
+# coding: utf-8
+
+from enum import IntEnum
+
+class AppRejectionCode(IntEnum):
+    FILE_TOO_BIG_FOR_SEEDING = 1
+    INSUFFICIENT_SPACE = 2
+    HTTP_404 = 3 # HTTP : Not found
+    HTTP_403 = 4 # HTTP : Forbidden
+    FETCH_FAILED = 5 # javascript fetch failure
+    NOT_SUITABLE_FOR_BETA = 100000
diff --git a/dispatcher/api/models/logic.py b/dispatcher/api/models/logic.py
@@ -1,14 +1,14 @@
 import json
-import os
 from pathlib import Path
 from typing import List, Dict, Optional
-from datetime import datetime  # Import spécifique de la classe datetime
 import logging
 import uuid
 
 from sqlalchemy.orm import Session
 
 from rescue_api.models import Asset, Rescue, Rescuer
+from rescue_api.models.mvp_downloader_library import MvpDownloaderLibrary
+
 from .payload import AssetModel
 from .priorizer_client import PriorizerClient
 
@@ -78,6 +78,14 @@ def _load_json(self, file: Path) -> List[Dict]:
     def _save_json(file: Path, data: List[Dict]):
         file.write_text(json.dumps(data, indent=2))
     
+    @staticmethod
+    def _sort_helper(asset: Dict) -> tuple:
+        """Helper function to sort by priority (desc) and size (asc)"""
+        return (
+            asset['priority'],
+            asset.get('size', float('inf')) if asset.get('size') is not None else float('inf')
+            )
+    
     async def get_available_assets(self) -> List[Dict]:   
         """
         Retrieve all assets (without filtering by allocation).
@@ -109,12 +117,12 @@ async def allocate_assets(self, free_space_mb: float, node_id: str = None) -> Di
         remaining_space = free_space_mb
         
         # Tri par priorité (desc) puis taille (asc)
-        for asset in sorted(available, key=lambda x: (x['priority'], x['size_mb'])):
+        for asset in sorted(available, key=self._sort_helper):
             # Asset size may not be known yet
-            if asset['size_mb'] is None or asset['size_mb'] <= remaining_space:
+            if asset['size'] is None or asset['size'] <= remaining_space:
                 selected.append(asset)
-                if asset['size_mb'] is not None:
-                   remaining_space -= asset['size_mb']
+                if asset['size'] is not None:
+                   remaining_space -= asset['size']
         
         if not selected:
             return None
@@ -130,7 +138,7 @@ async def allocate_assets(self, free_space_mb: float, node_id: str = None) -> Di
             "res_id": a['res_id'],
             "asset_id": a['asset_id'],
             "name": a['name'],
-            "size_mb": float(a['size_mb']) if a['size_mb'] is not None else None,
+            "size": float(a['size']) if a['size'] is not None else None,
             "priority": int(a['priority']),
             "url": a['url']
         } for a in selected]
@@ -140,7 +148,7 @@ async def allocate_assets(self, free_space_mb: float, node_id: str = None) -> Di
 
         return {
             "node_id": node_id,
-            "allocated_size_mb": sum(a['size_mb'] for a in selected if a['size_mb'] is not None),
+            "allocated_size": sum(a['size'] for a in selected if a['size'] is not None),
             "assets": selected,
             "allocation_id": str(uuid.uuid4())
         }
@@ -180,6 +188,11 @@ def upsert_rescues_to_db(self, rescuer_id: int, assets: List[AssetModel], db: Se
                 rescue.magnet_link = asset.magnet_link
                 rescue.status = asset.status.value.lower()
 
+            # Updating MVP downloader library table with file size
+            db.query(MvpDownloaderLibrary) \
+                .filter(MvpDownloaderLibrary.resource_id == asset.res_id) \
+                .update({'deeplink_file_size': asset.size})
+
             try:
                 db.commit()
             except Exception as e:
@@ -212,6 +225,49 @@ def upsert_rescues_to_db(self, rescuer_id: int, assets: List[AssetModel], db: Se
         }
 
 
+    def update_mvp_resource_link_as_defective(self, rescuer_id: int, resource_id: int, db: Session) -> bool:
+        result = False
+
+        if not self._rescuer_exists(rescuer_id=rescuer_id, db=db):
+            logger.error(f"Rescuer with id={rescuer_id} doesn't exist in the database.")
+            return False
+
+        db.query(MvpDownloaderLibrary) \
+            .filter(MvpDownloaderLibrary.resource_id == resource_id) \
+            .update({'defective_link_flag': True})
+
+        # @todo Log update in some table : rescuer, resource, defective, created_at
+
+        try:
+            db.commit()
+            result = True
+        except Exception as e:
+            print(f'Error while updating res {resource_id} with defective link in MVP Downloader Library', e)
+
+        return result
+
+
+    def update_mvp_resource_link_size(self, rescuer_id: int, resource_id: int, size: int, db: Session) -> bool:
+        result = False
+
+        if not self._rescuer_exists(rescuer_id=rescuer_id, db=db):
+            logger.error(f"Rescuer with id={rescuer_id} doesn't exist in the database.")
+            return False
+
+        db.query(MvpDownloaderLibrary) \
+            .filter(MvpDownloaderLibrary.resource_id == resource_id) \
+            .update({'deeplink_file_size': size})
+
+        # @todo Log update in some table : rescuer, resource, size, created_at        
+
+        try:
+            db.commit()
+            result = True
+        except Exception as e:
+            print(f'Error while updating res {resource_id} with size {size} in MVP Downloader Library', e)
+
+        return result
+
     @staticmethod
     def _rescuer_exists(rescuer_id: int, db: Session) -> bool:
         rescuer = db.query(Rescuer).filter(Rescuer.id == rescuer_id).first()
diff --git a/dispatcher/api/models/payload.py b/dispatcher/api/models/payload.py
@@ -22,7 +22,7 @@ class BaseAssetModel(BaseModel):
     # @todo : remove name ?
     name: str = Field(..., description="Asset name")
     priority: int = Field(..., description="Asset priority according to ranker")
-    size_mb: Optional[float] = Field(..., description="Asset estimated size")
+    size: Optional[float] = Field(..., description="Asset size (may be estimated)")
     ds_id: int = Field(..., description="Dataset id")
     res_id: int = Field(..., description="Resource id")
     asset_id: Optional[int] = Field(None, description="Asset id")
@@ -75,3 +75,10 @@ class RescuesResponse(DispatchResponse):
         default=None,
         description="List of rescued assets for which an error happened when we tried to commit it to the database.",
     )
+
+class RejectRequest(BaseModel):
+    rescuer_id: int = Field(..., description="Rescuer id")
+    code: int = Field(..., description="Rejection code")
+    # Resource id for Beta. To be replaced by asset id.
+    res_id: int = Field(..., description="Resource id")
+    size: Optional[int] = Field(default=None, description="Asset size")
diff --git a/dispatcher/api/routers/dispatch.py b/dispatcher/api/routers/dispatch.py
@@ -1,7 +1,8 @@
 import json
 from os.path import join, dirname
 
-from models.payload import DispatchRequest, DispatchResponse, RescuesRequest, RescuesResponse
+from constants.mvp import AppRejectionCode as MVPAppRejectionCode
+from models.payload import DispatchRequest, DispatchResponse, RejectRequest, RescuesRequest, RescuesResponse
 from models.state import app_state
 from fastapi import HTTPException, APIRouter, Depends
 from fastapi.responses import JSONResponse
@@ -107,3 +108,21 @@ async def upsert_rescues(request: RescuesRequest, db: Session = Depends(get_db))
     print(f"Dispatch response: {response}")
 
     return response
+
+@router.post('/asset/reject')
+async def update_mvp_rejected_resource(request: RejectRequest, db: Session = Depends(get_db)):
+    result = False
+
+    app_state._logger.info(f"________Rejected MVP resource: {request.code}")
+
+    if request.code in [MVPAppRejectionCode.FILE_TOO_BIG_FOR_SEEDING, MVPAppRejectionCode.NOT_SUITABLE_FOR_BETA]:
+        result = app_state._dispatcher.update_mvp_resource_link_size(request.rescuer_id, request.res_id, request.size, db)
+    elif request.code in [MVPAppRejectionCode.HTTP_404, MVPAppRejectionCode.HTTP_403, MVPAppRejectionCode.FETCH_FAILED]:
+        result = app_state._dispatcher.update_mvp_resource_link_as_defective(request.rescuer_id, request.res_id, db)
+
+    if not result:
+        raise HTTPException(
+            status_code=422,
+            detail="The error can be the following: the rescuer doesn't exist in the database; or the code provided is not valid"
+            " or the resource provided does not exist",
+        )
diff --git a/priorizer/api/models/logic.py b/priorizer/api/models/logic.py
@@ -10,13 +10,15 @@
 from rescue_api.models.mvp_downloader_library import MvpDownloaderLibrary
 from rescue_api.models.rescues import Rescue
 from rescue_api.database import get_db
-from sqlalchemy import func, case, desc, and_
+from sqlalchemy import func, case, desc, and_, or_
 from datetime import datetime, timezone
 from typing import List
 
-_RANKING_LIMIT = 100
+_RANKING_LIMIT = 100
 _MVP_RANKING_ID = 8 # broija 2025-09-20 : MVP default ranking id
 
+_MVP_BETA_LIMIT_SIZE = 400 * 1024 * 1024
+
 logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -73,25 +75,34 @@ def get_rank(self) -> dict:
                 .join(Resource, Resource.id == MvpDownloaderLibrary.resource_id)
                 .join(asset_resource, asset_resource.c.resource_id == MvpDownloaderLibrary.resource_id)
                 .outerjoin(Rescue, Rescue.asset_id == asset_resource.c.asset_id)
-                .where(DatasetRank.ranking_id == last_ranking_id).where(Rescue.asset_id.is_(None))
+                .where(DatasetRank.ranking_id == last_ranking_id)
+                .where(Rescue.asset_id.is_(None))
+                .where(and_(
+                     MvpDownloaderLibrary.defective_link_flag.is_(False),
+                     or_(MvpDownloaderLibrary.deeplink_file_size.is_(None), MvpDownloaderLibrary.deeplink_file_size < _MVP_BETA_LIMIT_SIZE)
+                     ))
                 .order_by(DatasetRank.rank, Resource.id)
                 .limit(_RANKING_LIMIT)
         )
+        # self._write_query(no_magnet_ranks)
 
         results = [{
                 "path": "",
                 "name": "",
                 "priority": r.rank,
-                "size_mb": r.deeplink_file_size,
+                "size": r.deeplink_file_size,
                 "ds_id": r.dataset_id,
                 "res_id": r.resource_id,
                 "asset_id": r.asset_id,
                 "url": r.deeplink
                 }
                 for r in no_magnet_ranks]
 
+        # results = [] # Forcing torrent downloads
+
+        remaining_count = _RANKING_LIMIT - len(results)
         # If not enough results, complete with results with magnet link
-        if len(results) < _RANKING_LIMIT:
+        if remaining_count > 0:
                 magnet_ranks = (
                         session.query(
                                 MvpDownloaderLibrary.dataset_id,
@@ -106,22 +117,22 @@ def get_rank(self) -> dict:
                         .join(asset_resource, asset_resource.c.resource_id == MvpDownloaderLibrary.resource_id)
                         .outerjoin(Rescue, Rescue.asset_id == asset_resource.c.asset_id)
                         .where(DatasetRank.ranking_id == last_ranking_id)
-                        .where(Rescue.asset_id.is_not_(None))
+                        .where(Rescue.asset_id.is_not(None))
                         .order_by(DatasetRank.rank, Resource.id)
-                        .limit(_RANKING_LIMIT)
+                        .limit(remaining_count)
                 )
-                while len(results) < _RANKING_LIMIT:
-                        for r in magnet_ranks:
-                                results.append({
-                                        "path": "",
-                                        "name": "",
-                                        "priority": r.rank,
-                                        "size_mb": r.deeplink_file_size,
-                                        "ds_id": r.dataset_id,
-                                        "res_id": r.resource_id,
-                                        "asset_id": r.asset_id,
-                                        "url": r.magnet_link
-                                })
+                for r in magnet_ranks:
+                        results.append({
+                                "path": "",
+                                "name": "",
+                                "priority": r.rank,
+                                "size": r.deeplink_file_size,
+                                "ds_id": r.dataset_id,
+                                "res_id": r.resource_id,
+                                "asset_id": r.asset_id,
+                                "url": r.magnet_link
+                        })
+
         return {"assets": results}
         
     def compute_rank(self) -> List[dict]:
diff --git a/priorizer/api/models/priorizer.py b/priorizer/api/models/priorizer.py
@@ -8,7 +8,7 @@ class AssetModel(BaseModel):
     path: str = Field(..., description="Asset path")
     name: str = Field(..., description="Asset name")
     priority: int = Field(..., description="Asset priority according to ranker")
-    size_mb: Optional[float] = Field(..., description="Asset estimated size")
+    size: Optional[float] = Field(..., description="Asset size (may be estimated)")
     ds_id: int = Field(..., description="Dataset id")
     res_id: int = Field(..., description="Resource id")
     asset_id: Optional[int] = Field(None, description="Asset id")
diff --git a/priorizer/dev/pyproject.toml b/priorizer/dev/pyproject.toml
@@ -16,24 +16,10 @@ dependencies = [
     "uvicorn==0.35.0",
     "psycopg2-binary>=2.9.9",
     "sqlalchemy>=2.0.0",
+    "rescue-api",
 ]
 
-[dependency-groups]
-dev = ["rescue-api"]
-prod = ["rescue-api"]
-
 [tool.uv.sources]
 rescue-api = [
-    { git = "https://github.com/dataforgoodfr/offseason-shelter-for-science-rescue_db.git", rev = "v0.0.3", group = "prod" },
     { path = "/lib/rescue_db/", editable = true, group = "dev" },
-]
-
-[tool.uv]
-default-groups = ["prod"]
-
-conflicts = [
-    [
-        { group = "dev" },
-        { group = "prod" },
-    ]
-]
+]
diff --git a/priorizer/pyproject.toml b/priorizer/pyproject.toml
@@ -16,18 +16,10 @@ dependencies = [
     "uvicorn==0.35.0",
     "psycopg2-binary>=2.9.9",
     "sqlalchemy>=2.0.0",
+    "rescue-api"
 ]
 
-[dependency-groups]
-prod = ["rescue-api"]
-
 [tool.uv.sources]
 rescue-api = [
-    { git = "https://github.com/dataforgoodfr/offseason-shelter-for-science-rescue_db.git", rev = "v0.0.4", group = "prod" },
-]
-
-[tool.uv]
-default-groups = ["prod"]
-
-conflicts = [
-]
+    { git = "https://github.com/dataforgoodfr/offseason-shelter-for-science-rescue_db.git", rev = "v0.0.4" },
+]
