diracx-db 0.0.1a33__py3-none-any.whl → 0.0.1a35__py3-none-any.whl

diracx/db/os/job_parameters.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from DIRAC.Core.Utilities import TimeUtilities
+
 from diracx.db.os.utils import BaseOSDB
 
 
@@ -18,8 +20,17 @@ class JobParametersDB(BaseOSDB):
         "Status": {"type": "keyword"},
         "JobType": {"type": "keyword"},
     }
-    index_prefix = "mysetup_elasticjobparameters_index_"
+    # TODO: Does this need to be configurable?
+    index_prefix = "job_parameters"
+
+    def index_name(self, vo, doc_id: int) -> str:
+        split = int(int(doc_id) // 1e6)
+        return f"{self.index_prefix}_{vo}_{split}m"
 
-    def index_name(self, doc_id: int) -> str:
-        # TODO: Remove setup and replace "123.0m" with "120m"?
-        return f"{self.index_prefix}_{doc_id // 1e6:.1f}m"
+    def upsert(self, vo, doc_id, document):
+        document = {
+            "JobID": doc_id,
+            "timestamp": TimeUtilities.toEpochMilliSeconds(),
+            **document,
+        }
+        return super().upsert(vo, doc_id, document)

diracx/db/os/utils.py

@@ -77,7 +77,7 @@ class BaseOSDB(metaclass=ABCMeta):
     index_prefix: str
 
     @abstractmethod
-    def index_name(self, doc_id: int) -> str: ...
+    def index_name(self, vo: str, doc_id: int) -> str: ...
 
     def __init__(self, connection_kwargs: dict[str, Any]) -> None:
         self._client: AsyncOpenSearch | None = None
@@ -182,15 +182,20 @@ class BaseOSDB(metaclass=ABCMeta):
         )
         assert result["acknowledged"]
 
-    async def upsert(self, doc_id, document) -> None:
-        # TODO: Implement properly
+    async def upsert(self, vo: str, doc_id: int, document: Any) -> None:
+        index_name = self.index_name(vo, doc_id)
         response = await self.client.update(
-            index=self.index_name(doc_id),
+            index=index_name,
             id=doc_id,
             body={"doc": document, "doc_as_upsert": True},
             params=dict(retry_on_conflict=10),
         )
-        print(f"{response=}")
+        logger.debug(
+            "Upserted document %s in index %s with response: %s",
+            doc_id,
+            index_name,
+            response,
+        )
 
     async def search(
         self, parameters, search, sorts, *, per_page: int = 100, page: int | None = None

diracx/db/sql/job/db.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+__all__ = ["JobDB"]
+
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Iterable
 
 from sqlalchemy import bindparam, case, delete, func, insert, select, update
 
@@ -9,13 +11,12 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.elements import BindParameter
 
 from diracx.core.exceptions import InvalidQueryError
-from diracx.core.models import (
-    SearchSpec,
-    SortSpec,
-)
+from diracx.core.models import JobCommand, SearchSpec, SortSpec
 
 from ..utils import BaseSQLDB, apply_search_filters, apply_sort_constraints
+from ..utils.functions import utcnow
 from .schema import (
+    HeartBeatLoggingInfo,
     InputData,
     JobCommands,
     JobDBBase,
@@ -38,6 +39,16 @@ def _get_columns(table, parameters):
 class JobDB(BaseSQLDB):
     metadata = JobDBBase.metadata
 
+    # Field names which should be stored in the HeartBeatLoggingInfo table
+    heartbeat_fields = {
+        "LoadAverage",
+        "MemoryUsed",
+        "Vsize",
+        "AvailableDiskSpace",
+        "CPUConsumed",
+        "WallClockTime",
+    }
+
     # TODO: this is copied from the DIRAC JobDB
     # but is overwriten in LHCbDIRAC, so we need
     # to find a way to make it dynamic
@@ -156,6 +167,24 @@ class JobDB(BaseSQLDB):
             ],
         )
 
+    @staticmethod
+    def _set_job_attributes_fix_value(column, value):
+        """Apply corrections to the values before inserting them into the database.
+
+        TODO: Move this logic into the sqlalchemy model.
+        """
+        if column == "VerifiedFlag":
+            value_str = str(value)
+            if value_str in ("True", "False"):
+                return value_str
+        if column == "AccountedFlag":
+            value_str = str(value)
+            if value_str in ("True", "False", "Failed"):
+                return value_str
+        else:
+            return value
+        raise NotImplementedError(f"Unrecognized value for column {column}: {value}")
+
     async def set_job_attributes(self, job_data):
         """Update the parameters of the given jobs."""
         # TODO: add myDate and force parameters.
@@ -168,7 +197,10 @@ class JobDB(BaseSQLDB):
         case_expressions = {
             column: case(
                 *[
-                    (Jobs.__table__.c.JobID == job_id, attrs[column])
+                    (
+                        Jobs.__table__.c.JobID == job_id,
+                        self._set_job_attributes_fix_value(column, attrs[column]),
+                    )
                     for job_id, attrs in job_data.items()
                     if column in attrs
                 ],
@@ -197,7 +229,7 @@ class JobDB(BaseSQLDB):
 
         return {jobid: jdl for jobid, jdl in (await self.conn.execute(stmt)) if jdl}
 
-    async def set_job_commands(self, commands: list[tuple[int, str, str]]):
+    async def set_job_commands(self, commands: list[tuple[int, str, str]]) -> None:
         """Store a command to be passed to the job together with the next heart beat."""
         await self.conn.execute(
             insert(JobCommands),
@@ -246,3 +278,58 @@ class JobDB(BaseSQLDB):
         rows = await self.conn.execute(stmt, update_parameters)
 
         return rows.rowcount
+
+    async def add_heartbeat_data(
+        self, job_id: int, dynamic_data: dict[str, str]
+    ) -> None:
+        """Add the job's heartbeat data to the database.
+
+        NOTE: This does not update the HeartBeatTime column in the Jobs table.
+        This is instead handled by the `diracx.logic.jobs.status.set_job_statuses`
+        as it involves updating multiple databases.
+
+        :param job_id: the job id
+        :param dynamic_data: mapping of the dynamic data to store,
+            e.g. {"AvailableDiskSpace": 123}
+        """
+        if extra_fields := set(dynamic_data) - self.heartbeat_fields:
+            raise InvalidQueryError(
+                f"Not allowed to store heartbeat data for: {extra_fields}. "
+                f"Allowed keys are: {self.heartbeat_fields}"
+            )
+        values = [
+            {
+                "JobID": job_id,
+                "Name": key,
+                "Value": value,
+                "HeartBeatTime": utcnow(),
+            }
+            for key, value in dynamic_data.items()
+        ]
+        await self.conn.execute(HeartBeatLoggingInfo.__table__.insert().values(values))
+
+    async def get_job_commands(self, job_ids: Iterable[int]) -> list[JobCommand]:
+        """Get a command to be passed to the job together with the next heartbeat.
+
+        :param job_ids: the job ids
+        :return: mapping of job id to list of commands
+        """
+        # Get the commands
+        stmt = (
+            select(JobCommands.job_id, JobCommands.command, JobCommands.arguments)
+            .where(JobCommands.job_id.in_(job_ids), JobCommands.status == "Received")
+            .order_by(JobCommands.job_id)
+        )
+        commands = await self.conn.execute(stmt)
+        # Update the status of the commands
+        stmt = (
+            update(JobCommands)
+            .where(JobCommands.job_id.in_(job_ids))
+            .values(Status="Sent")
+        )
+        await self.conn.execute(stmt)
+        # Return the commands grouped by job id
+        return [
+            JobCommand(job_id=cmd.JobID, command=cmd.Command, arguments=cmd.Arguments)
+            for cmd in commands
+        ]

diracx/db/sql/job/schema.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
+import sqlalchemy.types as types
 from sqlalchemy import (
     DateTime,
-    Enum,
     ForeignKey,
     Index,
     Integer,
@@ -16,6 +16,36 @@ from ..utils import Column, EnumBackedBool, NullColumn
 JobDBBase = declarative_base()
 
 
+class AccountedFlagEnum(types.TypeDecorator):
+    """Maps a ``AccountedFlagEnum()`` column to True/False in Python."""
+
+    impl = types.Enum
+    cache_ok: bool = True
+
+    def __init__(self) -> None:
+        super().__init__("True", "False", "Failed")
+
+    def process_bind_param(self, value, dialect) -> str:
+        if value is True:
+            return "True"
+        elif value is False:
+            return "False"
+        elif value == "Failed":
+            return "Failed"
+        else:
+            raise NotImplementedError(value, dialect)
+
+    def process_result_value(self, value, dialect) -> bool | str:
+        if value == "True":
+            return True
+        elif value == "False":
+            return False
+        elif value == "Failed":
+            return "Failed"
+        else:
+            raise NotImplementedError(f"Unknown {value=}")
+
+
 class Jobs(JobDBBase):
     __tablename__ = "Jobs"
 
@@ -45,10 +75,7 @@ class Jobs(JobDBBase):
     user_priority = Column("UserPriority", Integer, default=0)
     reschedule_counter = Column("RescheduleCounter", Integer, default=0)
     verified_flag = Column("VerifiedFlag", EnumBackedBool(), default=False)
-    # TODO: Should this be True/False/"Failed"? Or True/False/Null?
-    accounted_flag = Column(
-        "AccountedFlag", Enum("True", "False", "Failed"), default="False"
-    )
+    accounted_flag = Column("AccountedFlag", AccountedFlagEnum(), default=False)
 
     __table_args__ = (
         Index("JobType", "JobType"),

{diracx_db-0.0.1a33.dist-info → diracx_db-0.0.1a35.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diracx-db
-Version: 0.0.1a33
+Version: 0.0.1a35
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research

{diracx_db-0.0.1a33.dist-info → diracx_db-0.0.1a35.dist-info}/RECORD

@@ -3,8 +3,8 @@ diracx/db/__main__.py,sha256=tU4tp3OAClYCiPMxlRj524sZGBx9oy4CoWHd8pMuEEs,1715
 diracx/db/exceptions.py,sha256=1nn-SZLG-nQwkxbvHjZqXhE5ouzWj1f3qhSda2B4ZEg,83
 diracx/db/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/os/__init__.py,sha256=IZr6z6SefrRvuC8sTC4RmB3_wwOyEt1GzpDuwSMH8O4,112
-diracx/db/os/job_parameters.py,sha256=Knca19uT2G-5FI7MOFlaOAXeHn4ecPVLIH30TiwhaTw,858
-diracx/db/os/utils.py,sha256=xQLhOgC9YidfeEIbmdigXg3_YZLKHwMcZzhHf4mRwyQ,11399
+diracx/db/os/job_parameters.py,sha256=W-2O0cnq91iExcQAEqnCKuA9Ko_6MQaTGALGG8ziQ-4,1149
+diracx/db/os/utils.py,sha256=V4T-taos64SFNcorfIr7mq5l5y88K6TzyCj1YqWk8VI,11562
 diracx/db/sql/__init__.py,sha256=JYu0b0IVhoXy3lX2m2r2dmAjsRS7IbECBUMEDvX0Te4,391
 diracx/db/sql/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/auth/db.py,sha256=O9sp-AKobQ4X4IF2WaJKig-8LU8Ra9Syn5jW8JNLRU4,9030
@@ -13,8 +13,8 @@ diracx/db/sql/dummy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 diracx/db/sql/dummy/db.py,sha256=whcF02IjFQTwe1l4iTHtCkg2XukEaLHpLeHYNGMuGGw,1645
 diracx/db/sql/dummy/schema.py,sha256=9zI53pKlzc6qBezsyjkatOQrNZdGCjwgjQ8Iz_pyAXs,789
 diracx/db/sql/job/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/job/db.py,sha256=qoK1toWekRP1BO5XbT9CaRuGKR9_oixMVfRZ9myZsjg,8644
-diracx/db/sql/job/schema.py,sha256=Rmt5ldVN6Pf4MJHMhiT_Zn-k9LVPBTVivmGlyPKFRZE,4647
+diracx/db/sql/job/db.py,sha256=PTivVrSl4hrOTgy77WECunPJUJiYzGhSlqO0wYwEuE8,11909
+diracx/db/sql/job/schema.py,sha256=eFgZshe6NEzOM2qI0HI9Y3abrqDMoQIwa9L0vZugHcU,5431
 diracx/db/sql/job_logging/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/job_logging/db.py,sha256=0WuRn3pKEeuV9IZ_3VDhlvqI9eM1W4C0dZwIVZKW9N8,5811
 diracx/db/sql/job_logging/schema.py,sha256=DyFr_4goxfC453xPJX4OkxTGHtlP_5q8PXpr8bhNQS0,960
@@ -31,8 +31,8 @@ diracx/db/sql/utils/__init__.py,sha256=QkvpqBuIAgkAOywAssYzdxSzUQVZlSUumK7mPxotX
 diracx/db/sql/utils/base.py,sha256=7UxHBNLOSjdrIdslMKW4C_c5H9-6Y1BEimxscri2poE,12367
 diracx/db/sql/utils/functions.py,sha256=iLqlUIQ6SrDUtDEnZ5szaFbdcINJW15KNbCdGXss6kc,3055
 diracx/db/sql/utils/types.py,sha256=yU-tXsu6hFGPsr9ba1n3ZjGPnHQI_06lbpkTeDCWJtg,1287
-diracx_db-0.0.1a33.dist-info/METADATA,sha256=gMmJC7ijyqAcDSjopiSPacDI4FXIFAs7xvELywUHGA0,644
-diracx_db-0.0.1a33.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-diracx_db-0.0.1a33.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
-diracx_db-0.0.1a33.dist-info/top_level.txt,sha256=vJx10tdRlBX3rF2Psgk5jlwVGZNcL3m_7iQWwgPXt-U,7
-diracx_db-0.0.1a33.dist-info/RECORD,,
+diracx_db-0.0.1a35.dist-info/METADATA,sha256=9KlPHZrT1qCFKxJMs6uagTjk8VHZBbKRls7sw2dHqog,644
+diracx_db-0.0.1a35.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+diracx_db-0.0.1a35.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
+diracx_db-0.0.1a35.dist-info/top_level.txt,sha256=vJx10tdRlBX3rF2Psgk5jlwVGZNcL3m_7iQWwgPXt-U,7
+diracx_db-0.0.1a35.dist-info/RECORD,,