Enhance dataset lookup and improve job message handling

pandaserver/api/v1/job_api.py

@@ -18,6 +18,7 @@
     request_validation,
 )
 from pandaserver.config import panda_config
+from pandaserver.dataservice.ddm import rucioAPI
 from pandaserver.jobdispatcher import Protocol
 from pandaserver.srvcore.panda_request import PandaRequest
 from pandaserver.taskbuffer import JobUtils
@@ -35,10 +36,14 @@
 # partition of jobs_StatusLog
 MAX_STATUS_HISTORY_JOBS = 1000
 
-# most names one get_datasets_by_name call may give: what Oracle takes in one IN list. The
-# db_read_probe_budget normally bounds it lower, since each name costs one index lookup per
+# most keys one get_datasets_by_name call may look up: what Oracle takes in one IN list. Each name
+# is looked up in both of its forms, so a name costs KEYS_PER_DATASET_NAME keys. The
+# db_read_probe_budget normally bounds it lower, since each key costs one index lookup per
 # partition of Datasets
-MAX_DATASET_NAMES = 1000
+MAX_DATASET_KEYS = 1000
+
+# keys one name given to get_datasets_by_name is looked up with: the name as given and its other form
+KEYS_PER_DATASET_NAME = 2
 
 # slack put on both ends of a get_status_history window before it is widened to whole days. A
 # creationTime is written with the clock of a server and jobs_StatusLog with the database clock, so
@@ -63,6 +68,28 @@ def _parse_utc_time(value: str) -> datetime.datetime:
     return parsed
 
 
+def _dataset_name_forms(name: str) -> list[str]:
+    """
+    The forms a dataset name is looked up in: the name as given, then its other form, which is the
+    name without its scope when it has one, and with a scope put on when it has none.
+
+    The Datasets table keeps one row per vuid, under the name its first writer used, and the vuid is
+    made from scope:name whichever form the name came in, so the row of a dataset may be under either
+    form. The scope put on is the one extract_scope gives, since the vuid is made with that scope.
+
+    Args:
+        name(str): the dataset name as given
+
+    Returns:
+        list: the name, followed by its other form unless that is empty or the name itself
+    """
+    scope, bare_name = rucioAPI.extract_scope(name)
+    other_form = bare_name if ":" in name else f"{scope}:{bare_name}"
+    if other_form and other_form != name:
+        return [name, other_form]
+    return [name]
+
+
 def init_task_buffer(task_buffer: TaskBuffer) -> None:
     """
     Initialize the task buffer. This method needs to be called before any other method in this module.
@@ -328,10 +355,15 @@ def get_datasets_by_name(req: PandaRequest, names: List[str]) -> Dict[str, Any]:
     Gets the rows of the PanDA Datasets table (the datasets jobs work with, such as the _sub and _dis ones made for
     them, not JEDI_Datasets) for the given names. Requires a secure connection.
 
-    Each name is looked up once in every partition of the table, which is partitioned on modificationdate, so the most
-    names one call may give is the db_read_probe_budget of the server divided by the number of partitions. A call giving
-    more is refused, not cut short. Both numbers come back with every answer, so a caller can split its names into
-    calls that fit.
+    The table keeps one row per dataset, under the form of the name its first writer used, scope:name or name. So each
+    name is looked up in both forms: a name with a scope also without it, and a name without one also with the scope
+    its first part gives (up to the first dot, or up to the second for user and group datasets), which is the scope
+    the vuid of the dataset is made with. Every row says which of the given names it answers.
+
+    Each form is a key, looked up once in every partition of the table, which is partitioned on modificationdate, so
+    the most names one call may give is the db_read_probe_budget of the server divided by keys_per_name times the
+    number of partitions. A call giving more is refused, not cut short. The three numbers come back with every answer,
+    so a caller can split its names into calls that fit.
 
     POST is accepted too, with the arguments in a JSON body (Content-Type: application/json). A GET puts every name in
     the URL, and a web server refuses a request line over its limit with an HTTP error before the call gets here. A
@@ -350,20 +382,26 @@ def get_datasets_by_name(req: PandaRequest, names: List[str]) -> Dict[str, Any]:
 
     Response data:
         dict: The rows of the datasets, and what tells how complete they are.
-            {"rows": [...], "as_of": datetime, "retained_from": datetime or None, "partitions": int, "budget": int}
+            {"rows": [...], "as_of": datetime, "retained_from": datetime or None, "partitions": int, "budget": int,
+             "keys_per_name": int}
             - rows: one dict per row, ordered by name and creationdate, with every column: vuid, name, version, type,
               status, numberfiles, currentfiles, creationdate, modificationdate, MoverID, transferStatus and subType,
-              with null for NULL. A name without rows has none
+              with null for NULL, and asked_as: the given names, in the order given, that the row answers in either
+              form. A name without rows has none
             - as_of: database time in UTC taken before the rows were read
             - retained_from: time from which the table holds every row, on modificationdate, or None when it cannot be
               told. Old partitions are dropped, so a dataset last modified before this time may have been in the table
-            - partitions, budget: the number of partitions, counted as 1 when there are none, and the
-              db_read_probe_budget; a call may give at most budget // partitions names
+            - partitions, budget, keys_per_name: the number of partitions, counted as 1 when there are none, the
+              db_read_probe_budget, and the keys a name is looked up with; a call may give at most
+              budget // (keys_per_name * partitions) names
         On failure, with the reason in the message, the reason as values:
             - {"rejected": "no_keys"}: no names
-            - {"rejected": "too_many_keys", "keys": int, "max_keys": int}: more names than one IN list takes
-            - {"rejected": "probe_budget", "keys": int, "partitions": int, "budget": int}: keys x partitions exceed the budget
+            - {"rejected": "too_many_keys", "keys": int, "max_keys": int, "keys_per_name": int}: more keys than one IN
+              list takes
+            - {"rejected": "probe_budget", "keys": int, "partitions": int, "budget": int, "keys_per_name": int}: keys x
+              partitions exceed the budget
             - {"error": str}: the database failed, with the error code (ORA-xxxxx, or the SQLSTATE on Postgres)
+        keys counts the forms looked up, and a form shared by two names, such as s:n given with n, counts once.
     """
     tmp_logger = LogWrapper(_logger, f"get_datasets_by_name n_names={len(names)}")
     tmp_logger.debug("Start")
@@ -374,16 +412,31 @@ def get_datasets_by_name(req: PandaRequest, names: List[str]) -> Dict[str, Any]:
         message = "names must not be empty"
         tmp_logger.debug(message)
         return generate_response(False, message=message, data={"rejected": "no_keys"})
-    if len(names) > MAX_DATASET_NAMES:
-        message = f"at most {MAX_DATASET_NAMES} names are accepted, got {len(names)}"
+
+    # each name is looked up in both of its forms, and a form shared by two names, such as s:n given
+    # with n, is looked up once for both. The given names of each key, in the order given
+    asked_as: dict[str, list[str]] = {}
+    for name in names:
+        for key in _dataset_name_forms(name):
+            asked_as.setdefault(key, []).append(name)
+    keys = list(asked_as)
+    if len(keys) > MAX_DATASET_KEYS:
+        message = f"at most {MAX_DATASET_KEYS} keys are looked up, got {len(keys)} for {len(names)} names"
         tmp_logger.debug(message)
-        return generate_response(False, message=message, data={"rejected": "too_many_keys", "keys": len(names), "max_keys": MAX_DATASET_NAMES})
+        return generate_response(
+            False, message=message, data={"rejected": "too_many_keys", "keys": len(keys), "max_keys": MAX_DATASET_KEYS, "keys_per_name": KEYS_PER_DATASET_NAME}
+        )
 
-    success, message, data = global_task_buffer.get_datasets_by_name(names, panda_config.db_read_probe_budget)
+    success, message, data = global_task_buffer.get_datasets_by_name(keys, panda_config.db_read_probe_budget)
     if not success:
+        if data.get("rejected") == "probe_budget":
+            data["keys_per_name"] = KEYS_PER_DATASET_NAME
         tmp_logger.debug(f"Failed with {message}")
         return generate_response(False, message=message, data=data)
 
+    for row in data["rows"]:
+        row["asked_as"] = asked_as.get(row["name"], [])
+    data["keys_per_name"] = KEYS_PER_DATASET_NAME
     tmp_logger.debug(f"""Done with {len(data["rows"])} rows""")
     return generate_response(True, data=data)
 

pandaserver/api/v1/tests/job_api_tests.py

@@ -162,14 +162,14 @@ def test_get_datasets_by_name(self) -> None:
         self.assertEqual(status, 0)
         self.assertTrue(output["success"], output["message"])
         self.assertEqual(output["data"]["rows"], [])
-        self.assertLessEqual(2 * output["data"]["partitions"], output["data"]["budget"])
+        self.assertLessEqual(2 * output["data"]["keys_per_name"] * output["data"]["partitions"], output["data"]["budget"])
 
     def test_get_datasets_by_name_over_budget(self) -> None:
         # the most names a call may give, plus one
         url = f"{api_url_ssl}/job/get_datasets_by_name"
         status, output = self.http_client.post(url, {"names": ["no.such.dataset"]})
         self.assertTrue(output["success"], output["message"])
-        n_names = output["data"]["budget"] // output["data"]["partitions"] + 1
+        n_names = output["data"]["budget"] // (output["data"]["keys_per_name"] * output["data"]["partitions"]) + 1
         status, output = self.http_client.post(url, {"names": [f"no.such.dataset_sub{i}" for i in range(n_names)]})
         print(output)
         self.assertEqual(status, 0)
@@ -308,27 +308,72 @@ def _call(self, names: Any, method: str = "POST", answer: tuple[bool, str, dict[
             out = job_api.get_datasets_by_name(self._request(method), names=names)
         return out, task_buffer
 
+    # the same dataset in both forms of its name
+    SCOPED = "mc23_13p6TeV:mc23_13p6TeV.601229.PhPy8EG.evgen.EVNT.e8514_tid41234567_00"
+    BARE = "mc23_13p6TeV.601229.PhPy8EG.evgen.EVNT.e8514_tid41234567_00"
+
+    def _answer(self, *row_names: str) -> tuple[bool, str, dict[str, Any]]:
+        """A successful lookup with one row for each of the given names."""
+        return True, "", dict(self.ANSWER, rows=[{"vuid": f"vuid-{i}", "name": name} for i, name in enumerate(row_names)])
+
     def test_names_looked_up_once(self) -> None:
-        out, task_buffer = self._call(["a_sub1", "b_sub2", "a_sub1"])
+        out, task_buffer = self._call(["a.x_sub1", "b.y_sub2", "a.x_sub1"])
+        self.assertTrue(out["success"], out["message"])
+        self.assertEqual(task_buffer.get_datasets_by_name.call_args[0], (["a.x_sub1", "a:a.x_sub1", "b.y_sub2", "b:b.y_sub2"], 1000))
+
+    def test_scoped_name_looked_up_without_scope(self) -> None:
+        out, task_buffer = self._call([self.SCOPED])
+        self.assertEqual(task_buffer.get_datasets_by_name.call_args[0][0], [self.SCOPED, self.BARE])
+
+    def test_unscoped_name_looked_up_with_scope(self) -> None:
+        out, task_buffer = self._call([self.BARE])
+        self.assertEqual(task_buffer.get_datasets_by_name.call_args[0][0], [self.BARE, self.SCOPED])
+        # the scope of a user dataset is its first two parts
+        out, task_buffer = self._call(["user.alice.test.v1_sub123"])
+        self.assertEqual(task_buffer.get_datasets_by_name.call_args[0][0], ["user.alice.test.v1_sub123", "user.alice:user.alice.test.v1_sub123"])
+
+    def test_no_empty_form(self) -> None:
+        out, task_buffer = self._call(["s:"])
+        self.assertEqual(task_buffer.get_datasets_by_name.call_args[0][0], ["s:"])
+
+    def test_both_forms_given(self) -> None:
+        out, task_buffer = self._call([self.SCOPED, self.BARE], answer=self._answer(self.BARE, self.SCOPED))
         self.assertTrue(out["success"], out["message"])
-        self.assertEqual(task_buffer.get_datasets_by_name.call_args[0], (["a_sub1", "b_sub2"], 1000))
+        self.assertEqual(task_buffer.get_datasets_by_name.call_args[0][0], [self.SCOPED, self.BARE])
+        self.assertEqual([row["asked_as"] for row in out["data"]["rows"]], [[self.SCOPED, self.BARE], [self.SCOPED, self.BARE]])
+
+    def test_rows_say_which_names_they_answer(self) -> None:
+        out, task_buffer = self._call([self.SCOPED, "other.x_sub9"], answer=self._answer(self.BARE, "other:other.x_sub9"))
+        self.assertTrue(out["success"], out["message"])
+        self.assertEqual(
+            [(row["name"], row["asked_as"]) for row in out["data"]["rows"]], [(self.BARE, [self.SCOPED]), ("other:other.x_sub9", ["other.x_sub9"])]
+        )
+        self.assertEqual(out["data"]["keys_per_name"], 2)
+        self.assertEqual(out["data"]["partitions"], 5)
 
     def test_get(self) -> None:
-        out, task_buffer = self._call(["a_sub1"], method="GET")
+        out, task_buffer = self._call(["a.x_sub1"], method="GET")
         self.assertTrue(out["success"], out["message"])
-        self.assertEqual(task_buffer.get_datasets_by_name.call_args[0][0], ["a_sub1"])
+        self.assertEqual(task_buffer.get_datasets_by_name.call_args[0][0], ["a.x_sub1", "a:a.x_sub1"])
 
     def test_refused(self) -> None:
         out, task_buffer = self._call([])
         self.assertEqual((out["success"], out["data"]), (False, {"rejected": "no_keys"}))
-        out, task_buffer = self._call([f"a_sub{i}" for i in range(1001)])
-        self.assertEqual((out["success"], out["data"]), (False, {"rejected": "too_many_keys", "keys": 1001, "max_keys": 1000}))
+        # the cap counts keys, two a name
+        out, task_buffer = self._call([f"a.x_sub{i}" for i in range(501)])
+        self.assertEqual((out["success"], out["data"]), (False, {"rejected": "too_many_keys", "keys": 1002, "max_keys": 1000, "keys_per_name": 2}))
         task_buffer.get_datasets_by_name.assert_not_called()
+        out, task_buffer = self._call([f"a.x_sub{i}" for i in range(500)])
+        self.assertEqual(len(task_buffer.get_datasets_by_name.call_args[0][0]), 1000)
 
     def test_refusal_of_the_lookup_passed_on(self) -> None:
-        refusal = {"rejected": "probe_budget", "keys": 201, "partitions": 5, "budget": 1000}
-        out, task_buffer = self._call([f"a_sub{i}" for i in range(201)], answer=(False, "over", refusal))
-        self.assertEqual((out["success"], out["message"], out["data"]), (False, "over", refusal))
+        refusal = {"rejected": "probe_budget", "keys": 402, "partitions": 5, "budget": 1000}
+        out, task_buffer = self._call([f"a.x_sub{i}" for i in range(201)], answer=(False, "over", dict(refusal)))
+        self.assertEqual((out["success"], out["message"], out["data"]), (False, "over", dict(refusal, keys_per_name=2)))
+
+    def test_error_passed_on(self) -> None:
+        out, task_buffer = self._call(["a.x_sub1"], answer=(False, "database error ORA-00942", {"error": "ORA-00942"}))
+        self.assertEqual((out["success"], out["message"], out["data"]), (False, "database error ORA-00942", {"error": "ORA-00942"}))
 
 
 if __name__ == "__main__":

pandaserver/daemons/scripts/cache_schedconfig.py

@@ -35,7 +35,7 @@ class cacheSchedConfig:
     def __init__(self, tbuf: Any) -> None:
         self.tbuf = tbuf
         # the queue rows, read by getQueueData() before anything dumps them
-        self.queueData: Any = None
+        self.queueData: list[dict[str, Any]] = []
         self.cloudStatus = None
         # Define this here, but could be more flexible...
         self.queueDataFields = {
@@ -114,9 +114,9 @@ def __init__(self, tbuf: Any) -> None:
             "all": None,
         }
 
-    def query_column_sql(self, sql: str, varMap: dict[str, Any] | None = None, arraySize: int = 100) -> list[Any]:
+    def query_column_sql(self, sql: str, varMap: dict[str, Any] | None = None, arraySize: int = 100) -> list[dict[str, Any]]:
         res = self.tbuf.querySQL(sql, varMap, arraySize=arraySize)
-        retList = []
+        retList: list[dict[str, Any]] = []
         for (
             panda_queue,
             data,

pandaserver/daemons/utils.py

@@ -15,14 +15,18 @@
 from collections.abc import Callable
 from multiprocessing.connection import Connection
 from types import FrameType
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import psutil
 from pandacommon.pandalogger import logger_utils
 from pandacommon.pandautils.thread_utils import GenericThread, LockPool
 
 from pandaserver.config import daemon_config, panda_config
 
+if TYPE_CHECKING:
+    # imported for annotations only; _make_tbif() imports it at runtime, where the order matters
+    from pandaserver.taskbuffer.TaskBufferInterface import TaskBufferInterface
+
 # list of signals accepted to end the main process
 END_SIGNALS = [
     signal.SIGINT,
@@ -462,9 +466,8 @@ def __init__(self, logger: logging.Logger, n_workers: int = 1, n_dbconn: int = 1
         self._make_dem_run_map()
         # map to store global states
         self.global_state_map: dict[str, Any] = {}
-        # shared taskBufferIF, built by _make_tbif() only when use_tbif is on. The class it
-        # holds is imported inside that method on purpose, so it cannot be named here
-        self.tbif: Any = None
+        # shared taskBufferIF, built by _make_tbif() only when use_tbif is on
+        self.tbif: "TaskBufferInterface | None" = None
         self._make_tbif()
         # shared lock pool
         self.lock_pool = LockPool()

pandaserver/dataservice/adder_gen.py

@@ -19,6 +19,7 @@
 from pandaserver.brokerage.SiteMapper import SiteMapper
 from pandaserver.config import panda_config
 from pandaserver.dataservice import DataServiceUtils, closer
+from pandaserver.dataservice.adder_result import AdderResult
 from pandaserver.srvcore.CoreUtils import normalize_cpu_model
 
 if TYPE_CHECKING:
@@ -89,7 +90,7 @@ def __init__(
         self.lock_pool = lock_pool
         self.report_dict: dict[str, Any] = {}
         self.adder_plugin: Any = None
-        self.add_result: Any = None
+        self.add_result: AdderResult | None = None
         self.adder_plugin_class: Any = None
         # logger
         self.logger = LogWrapper(_logger, str(self.job_id))

pandaserver/srvcore/srv_msg_utils.py

@@ -1,8 +1,13 @@
 import json
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from pandacommon.pandautils.PandaUtils import naive_utcnow
 
+if TYPE_CHECKING:
+    # imported for annotations only: msg_bkr_utils sets up its loggers and imports stomp
+    # when imported, and DBProxy.get_mb_proxy defers that until a daemon has opened its log
+    from pandacommon.pandamsgbkr.msg_bkr_utils import MBSenderProxy
+
 
 # make message
 def make_message(msg_type: str, **kwargs: Any) -> str:
@@ -13,7 +18,7 @@ def make_message(msg_type: str, **kwargs: Any) -> str:
 
 
 # send a job message
-def send_job_message(msg_queue: Any, msg_topic: Any, task_id: int, job_id: int) -> None:
+def send_job_message(msg_queue: "MBSenderProxy", msg_topic: "MBSenderProxy", task_id: int, job_id: int) -> None:
     # make message
     msg = make_message("get_job", taskid=task_id, jobid=job_id)
     # use job ID for selector
@@ -25,7 +30,7 @@ def send_job_message(msg_queue: Any, msg_topic: Any, task_id: int, job_id: int)
 
 
 # delete a job message
-def delete_job_message(msg_queue: Any, job_id: int, time_out: int = 10) -> None:
+def delete_job_message(msg_queue: "MBSenderProxy", job_id: int, time_out: int = 10) -> None:
     # job ID for selector
     headers = {"selector": "type='{0}' OR JMSType='{0}'".format(job_id)}
     # subscribe to remove job messages
@@ -35,7 +40,7 @@ def delete_job_message(msg_queue: Any, job_id: int, time_out: int = 10) -> None:
 
 
 # send a task message
-def send_task_message(msg_topic: Any, command_str: str, task_id: int) -> None:
+def send_task_message(msg_topic: "MBSenderProxy", command_str: str, task_id: int) -> None:
     # make message
     msg = make_message(f"{command_str}_task", taskid=task_id)
     # send message to topic