Implement TaskQueueDB

run_local.sh

@@ -17,6 +17,7 @@ export DIRACX_CONFIG_BACKEND_URL="git+file://${tmp_dir}/cs_store/initialRepo"
 export DIRACX_DB_URL_AUTHDB="sqlite+aiosqlite:///:memory:"
 export DIRACX_DB_URL_JOBDB="sqlite+aiosqlite:///:memory:"
 export DIRACX_DB_URL_JOBLOGGINGDB="sqlite+aiosqlite:///:memory:"
+export DIRACX_DB_URL_TASKQUEUEDB="sqlite+aiosqlite:///:memory:"
 export DIRACX_SERVICE_AUTH_TOKEN_KEY="file://${tmp_dir}/signing-key/rs256.key"
 export DIRACX_SERVICE_AUTH_ALLOWED_REDIRECTS='["http://'$(hostname| tr -s '[:upper:]' '[:lower:]')':8000/docs/oauth2-redirect"]'
 

setup.cfg

@@ -40,6 +40,7 @@ diracx.dbs =
 	AuthDB = diracx.db:AuthDB
 	JobDB = diracx.db:JobDB
 	JobLoggingDB = diracx.db:JobLoggingDB
+	TaskQueueDB = diracx.db:TaskQueueDB
 	#DummyDB = diracx.db:DummyDB
 diracx.services =
 	jobs = diracx.routers.job_manager:router

src/diracx/cli/internal.py

@@ -48,9 +48,7 @@ def generate_cs(
         DefaultGroup=user_group,
         Users={},
         Groups={
-            user_group: GroupConfig(
-                JobShare=None, Properties=["NormalUser"], Quota=None, Users=[]
-            )
+            user_group: GroupConfig(Properties=["NormalUser"], Quota=None, Users=[])
         },
     )
     config = Config(

src/diracx/core/config/schema.py

@@ -5,7 +5,7 @@
 from typing import Any, Optional
 
 from pydantic import BaseModel as _BaseModel
-from pydantic import EmailStr, PrivateAttr, root_validator
+from pydantic import EmailStr, PrivateAttr, root_validator, validator
 
 from ..properties import SecurityProperty
 
@@ -49,7 +49,7 @@ class GroupConfig(BaseModel):
     AutoAddVOMS: bool = False
     AutoUploadPilotProxy: bool = False
     AutoUploadProxy: bool = False
-    JobShare: Optional[int]
+    JobShare: int = 1000
     Properties: list[SecurityProperty]
     Quota: Optional[int]
     Users: list[str]
@@ -86,9 +86,35 @@ class JobMonitoringConfig(BaseModel):
     useESForJobParametersFlag: bool = False
 
 
+class JobSchedulingConfig(BaseModel):
+    EnableSharesCorrection: bool = False
+    taskQueueCPUTimeIntervals: list[int] = [
+        6 * 60,
+        30 * 60,
+        1 * 3600,
+        6 * 3600,
+        12 * 3600,
+        1 * 86400,
+        2 * 86400,
+        3 * 86400,
+        4 * 86400,
+        6 * 86400,
+        8 * 86400,
+        10 * 86400,
+        int(12.5 * 86400),
+    ]
+
+    @validator("taskQueueCPUTimeIntervals")
+    def check_taskQueueCPUTimeIntervals(cls, v):
+        assert len(v) >= 1
+        assert all(x > 0 for x in v)
+        return sorted(v)
+
+
 class ServicesConfig(BaseModel):
     Catalogs: dict[str, Any] | None
     JobMonitoring: JobMonitoringConfig = JobMonitoringConfig()
+    JobScheduling: JobSchedulingConfig = JobSchedulingConfig()
 
 
 class OperationsConfig(BaseModel):

src/diracx/core/models.py

@@ -3,9 +3,11 @@
 from datetime import datetime
 from enum import Enum
 from typing import Literal, TypedDict
+from uuid import UUID
 
 from pydantic import BaseModel, Field
 
+from diracx.core.properties import SecurityProperty
 from diracx.core.utils import JobStatus
 
 
@@ -82,4 +84,36 @@ class SetJobStatusReturn(BaseModel):
     last_update_time: datetime | None = Field(alias="LastUpdateTime")
 
 
+class tqDefDict(BaseModel):
+    cpu_time: int = Field(alias="CPUTime")
+    owner: str = Field(alias="Owner")
+    owner_group: str = Field(alias="OwnerGroup")
+    sites: list[str] = Field(alias="Sites")
+    grid_ces: list[str] = Field(alias="GridCEs")
+    banned_sites: list[str] = Field(alias="BannedSites")
+    platforms: list[str] = Field(alias="Platforms")
+    job_types: list[str] = Field(alias="JobTypes")
+    tags: list[str] = Field(alias="Tags")
+
+
+class AuthInfo(BaseModel):
+    # raw token for propagation
+    bearer_token: str
+
+    # token ID in the DB for Component
+    # unique jwt identifier for user
+    token_id: UUID
+
+    # list of DIRAC properties
+    properties: list[SecurityProperty]
+
+
+class UserInfo(AuthInfo):
+    # dirac generated vo:sub
+    sub: str
+    preferred_username: str
+    dirac_group: str
+    vo: str
+
+
 SearchSpec = ScalarSearchSpec | VectorSearchSpec

src/diracx/db/__init__.py

@@ -1,6 +1,6 @@
-__all__ = ("AuthDB", "JobDB", "JobLoggingDB")
+__all__ = ("AuthDB", "JobDB", "JobLoggingDB", "TaskQueueDB")
 
 from .auth.db import AuthDB
-from .jobs.db import JobDB, JobLoggingDB
+from .jobs.db import JobDB, JobLoggingDB, TaskQueueDB
 
 # from .dummy.db import DummyDB