ServiceX_DID_Finder_lib 2.0.2__tar.gz → 3.0.0a1__tar.gz

{servicex_did_finder_lib-2.0.2 → servicex_did_finder_lib-3.0.0a1}/PKG-INFO

@@ -1,16 +1,15 @@
 Metadata-Version: 2.1
 Name: ServiceX_DID_Finder_lib
-Version: 2.0.2
+Version: 3.0.0a1
 Summary: ServiceX DID Library Routines
 Author: Gordon Watts
 Author-email: gwatts@uw.edu
-Requires-Python: >=3.8,<4.0
+Requires-Python: >=3.8.1,<4.0.0
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: Celery (>=5.4,<6.0)
 Requires-Dist: make-it-sync (>=1.0.0,<2.0.0)
-Requires-Dist: pika (==1.1.0)
 Requires-Dist: requests (>=2.25.0,<3.0.0)

{servicex_did_finder_lib-2.0.2 → servicex_did_finder_lib-3.0.0a1}/pyproject.toml

@@ -1,21 +1,22 @@
 [tool.poetry]
 name = "ServiceX_DID_Finder_lib"
-version = "2.0.2"
+version = "3.0.0-alpha.1"
 description = "ServiceX DID Library Routines"
 authors = ["Gordon Watts <gwatts@uw.edu>"]
 
 [tool.poetry.dependencies]
-python = "^3.8"
-pika = "1.1.0"
+python = "^3.8.1"
 make-it-sync = "^1.0.0"
 requests = "^2.25.0"
+Celery= "^5.4"
+
 
 [tool.poetry.group.dev]
 optional = true
 
 [tool.poetry.group.dev.dependencies]
-pytest = "^7.4"
-flake8 = "^3.9.1"
+pytest = "^8.2"
+flake8 = "^7.1"
 pytest-mock = "^3.12.0"
 coverage = "^7.4.0"
 responses = "^0.14.0"

servicex_did_finder_lib-3.0.0a1/src/servicex_did_finder_lib/__init__.py

@@ -0,0 +1 @@
+from .did_finder_app import DIDFinderApp  # NOQA F401

servicex_did_finder_lib-3.0.0a1/src/servicex_did_finder_lib/accumulator.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2024, IRIS-HEP
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+from typing import List, Dict, Any, Union
+
+from servicex_did_finder_lib.did_summary import DIDSummary
+from servicex_did_finder_lib.servicex_adaptor import ServiceXAdapter
+
+
+class Accumulator:
+    """Track or cache files depending on the mode we are operating in"""
+
+    def __init__(self, sx: ServiceXAdapter, sum: DIDSummary):
+        self.servicex = sx
+        self.summary = sum
+        self.file_cache: List[Dict[str, Any]] = []
+
+    def add(self, file_info: Union[Dict[str, Any], List[Dict[str, Any]]]):
+        """
+        Track and inject the file back into the system
+        :param file_info: The file information to track can be a single record or a list
+        """
+        if isinstance(file_info, dict):
+            self.file_cache.append(file_info)
+        elif isinstance(file_info, list):
+            self.file_cache.extend(file_info)
+        else:
+            raise ValueError("Invalid input: expected a dictionary or a list of dictionaries")
+
+    @property
+    def cache_len(self) -> int:
+        return len(self.file_cache)
+
+    def send_on(self, count):
+        """
+        Send the accumulated files
+        :param count: The number of files to send. Set to -1 to send all
+        """
+
+        # Sort the list to insure reproducibility
+        files = sorted(self.file_cache, key=lambda x: x["paths"])
+        self.send_bulk(files[:count])
+        self.file_cache.clear()
+
+    def send_bulk(self, file_list: List[Dict[str, Any]]):
+        """
+        does a bulk put of files
+        :param file_list: The list of files to send
+        """
+        for ifl in file_list:
+            self.summary.add_file(ifl)
+        self.servicex.put_file_add_bulk(file_list)

servicex_did_finder_lib-3.0.0a1/src/servicex_did_finder_lib/did_finder_app.py

@@ -0,0 +1,197 @@
+# Copyright (c) 2022, IRIS-HEP
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+import argparse
+import logging
+from datetime import datetime
+from typing import Any, Generator, Callable, Dict, Optional
+
+from celery import Celery, Task
+
+from servicex_did_finder_lib.accumulator import Accumulator
+from servicex_did_finder_lib.did_logging import initialize_root_logger
+from servicex_did_finder_lib.did_summary import DIDSummary
+from servicex_did_finder_lib.servicex_adaptor import ServiceXAdapter
+from servicex_did_finder_lib.util_uri import parse_did_uri
+
+# The type for the callback method to handle DID's, supplied by the user.
+# Arguments are:
+#   - The DID to process
+#   - A dictionary of information about the DID request
+#   - A dictionary of arguments passed to the DID finder
+UserDIDHandler = Callable[
+    [str, Dict[str, Any], Dict[str, Any]],
+    Generator[Dict[str, Any], None, None]
+]
+
+
+__logging = logging.getLogger(__name__)
+__logging.addHandler(logging.NullHandler())
+
+
+class DIDFinderTask(Task):
+    """
+    A Celery task that will process a single DID request. This task will
+    call the user supplied DID finder to get the list of files associated
+    with the DID, and then send that list to ServiceX for processing.
+    """
+    def __init__(self):
+        super().__init__()
+        self.logger = logging.getLogger(__name__)
+        self.logger.addHandler(logging.NullHandler())
+
+    def do_lookup(self, did: str, dataset_id: int, endpoint: str, user_did_finder: UserDIDHandler):
+        """
+        Perform the DID lookup for the given DID. This will call the user supplied
+        DID finder to get the list of files associated with the DID, and then send
+        that list to ServiceX for processing.
+        After all of the files have been sent, send a message to ServiceX indicating
+        that the fileset is complete
+        Args:
+            did: The DID to process
+            dataset_id: The dataset ID for the request
+            endpoint: The ServiceX endpoint to send the request to
+            user_did_finder: The user supplied DID finder to call to get the list of files
+        """
+
+        self.logger.info(
+            f"Received DID request {did}",
+            extra={"dataset_id": dataset_id}
+        )
+
+        servicex = ServiceXAdapter(dataset_id=dataset_id, endpoint=endpoint)
+
+        info = {
+            "dataset-id": dataset_id,
+        }
+
+        start_time = datetime.now()
+
+        summary = DIDSummary(did)
+        did_info = parse_did_uri(did)
+        acc = Accumulator(servicex, summary)
+
+        try:
+            for file_info in user_did_finder(did_info.did, info, self.app.did_finder_args):
+                acc.add(file_info)
+
+        except Exception:
+            if did_info.get_mode == "all":
+                raise
+
+        acc.send_on(did_info.file_count)
+
+        elapsed_time = int((datetime.now() - start_time).total_seconds())
+        servicex.put_fileset_complete(
+            {
+                "files": summary.file_count,
+                "files-skipped": summary.files_skipped,
+                "total-events": summary.total_events,
+                "total-bytes": summary.total_bytes,
+                "elapsed-time": elapsed_time,
+            }
+        )
+
+
+class DIDFinderApp:
+    """
+    The main application for a DID finder. This will setup the Celery application
+    and start the worker to process the DID requests.
+    """
+    def __init__(self, did_finder_name: str,
+                 parsed_args: Optional[argparse.Namespace] = None):
+        """
+        Initialize the DID finder application
+        Args:
+            did_finder_name: The name of the DID finder
+            parsed_args: The parsed command line arguments. Leave as None to use the default parser
+        """
+
+        self.name = did_finder_name
+        self.parsed_args = vars(parsed_args) if parsed_args else None
+
+        # Setup command line parsing
+        if self.parsed_args is None:
+            parser = argparse.ArgumentParser()
+            self.add_did_finder_cnd_arguments(parser)
+            self.parsed_args = vars(parser.parse_args())
+
+        initialize_root_logger(self.name)
+
+        self.app = Celery(f"did_finder_{self.name}", broker_url=self.parsed_args['rabbit_uri'])
+
+        # Cache the args in the App so they are accessible to the tasks
+        self.app.did_finder_args = self.parsed_args
+
+    def did_lookup_task(self, name):
+        """
+        Decorator to create a new task to handle a DID lookup request wihout
+        needing to know about Celery tasks.
+        Usage:
+            @app.did_lookup_task(name="did_finder_cern_opendata.lookup_dataset")
+            def lookup_dataset(self, did: str, dataset_id: int, endpoint: str) -> None:
+                self.do_lookup(did=did, dataset_id=dataset_id,
+                               endpoint=endpoint, user_did_finder=find_files)
+
+        Args:
+            name: The name of the task
+        """
+        def decorator(func):
+            @self.app.task(base=DIDFinderTask, bind=True, name=name)
+            def wrapper(*args, **kwargs):
+                return func(*args, **kwargs)
+            return wrapper
+        return decorator
+
+    def start(self):
+        self.app.worker_main(argv=['worker', '--loglevel=INFO'])
+
+    @classmethod
+    def add_did_finder_cnd_arguments(cls, parser: argparse.ArgumentParser):
+        """add_did_finder_cnd_arguments Add required arguments to a parser
+
+        If you need to parse command line arguments for some special configuration, create your
+        own argument parser, and call this function to make sure the arguments needed
+        for running the back-end communication are filled in properly.
+
+        Then pass the results of the parsing to the DID Finder App's constructor method.
+
+        Args:
+            parser (argparse.ArgumentParser): The argument parser. Arguments needed for the
+                                              did finder/servicex communication will be added.
+        """
+        parser.add_argument(
+            "--rabbit-uri", dest="rabbit_uri", action="store", required=True
+        )
+        parser.add_argument(
+            "--prefix",
+            dest="prefix",
+            action="store",
+            required=False,
+            default="",
+            help="Prefix to add to use a caching proxy for URIs",
+        )

{servicex_did_finder_lib-2.0.2 → servicex_did_finder_lib-3.0.0a1}/src/servicex_did_finder_lib/servicex_adaptor.py

@@ -51,23 +51,6 @@ class ServiceXAdapter:
             'file_events': file_info['file_events']
         }
 
-    def put_file_add(self, file_info):
-        success = False
-        attempts = 0
-        while not success and attempts < MAX_RETRIES:
-            try:
-                mesg = self._create_json(file_info)
-                requests.put(f"{self.endpoint}{self.dataset_id}/files", json=mesg)
-                self.logger.info("adding file:", extra=file_info)
-                success = True
-            except requests.exceptions.ConnectionError:
-                self.logger.exception(f'Connection error to ServiceX App. Will retry '
-                                      f'(try {attempts} out of {MAX_RETRIES}')
-                attempts += 1
-        if not success:
-            self.logger.error(f'After {attempts} tries, failed to send ServiceX App a put_file '
-                              f'message: {str(file_info)} - Ignoring error.')
-
     def put_file_add_bulk(self, file_list, chunk_length=300):
         # we first chunk up file_list as it can be very large in
         # case there are a lot of replicas and a lot of files.

servicex_did_finder_lib-2.0.2/src/servicex_did_finder_lib/__init__.py

@@ -1,4 +0,0 @@
-__version__ = '1.0.0a1'
-
-from .communication import start_did_finder, \
-    add_did_finder_cnd_arguments  # NOQA

servicex_did_finder_lib-2.0.2/src/servicex_did_finder_lib/communication.py

@@ -1,272 +0,0 @@
-import argparse
-from datetime import datetime
-import json
-import logging
-import time
-from typing import Any, AsyncGenerator, Callable, Dict, List, Optional
-import sys
-
-import pika
-from make_it_sync import make_sync
-
-from servicex_did_finder_lib.did_summary import DIDSummary
-from servicex_did_finder_lib.did_logging import initialize_root_logger
-from servicex_did_finder_lib.util_uri import parse_did_uri
-from .servicex_adaptor import ServiceXAdapter
-
-# The type for the callback method to handle DID's, supplied by the user.
-UserDIDHandler = Callable[[str, Dict[str, Any]], AsyncGenerator[Dict[str, Any], None]]
-
-# Given name, build the RabbitMQ queue name by appending this.
-# This is backed into how ServiceX works - do not change unless it
-# is also changed in the ServiceX_App
-QUEUE_NAME_POSTFIX = "_did_requests"
-
-# Easy to use local logger
-__logging = logging.getLogger(__name__)
-__logging.addHandler(logging.NullHandler())
-
-# TODO: get rid of different modes. It should always be batch.
-
-
-class _accumulator:
-    "Track or cache files depending on the mode we are operating in"
-
-    def __init__(self, sx: ServiceXAdapter, sum: DIDSummary, hold_till_end: bool):
-        self._servicex = sx
-        self._summary = sum
-        self._hold_till_end = hold_till_end
-        self._file_cache: List[Dict[str, Any]] = []
-
-    def add(self, file_info: Dict[str, Any]):
-        "Track and inject the file back into the system"
-        if self._hold_till_end:
-            self._file_cache.append(file_info)
-        else:
-            self._summary.add_file(file_info)
-            self._servicex.put_file_add(file_info)
-
-    def send_on(self, count):
-        "Send the accumulated files"
-        if self._hold_till_end:
-            self._hold_till_end = False
-            files = sorted(self._file_cache, key=lambda x: x["paths"])
-            self.send_bulk(files[:count])
-
-    def send_bulk(self, file_list: List[Dict[str, Any]]):
-        "does a bulk put of files"
-        if self._hold_till_end:
-            for f in file_list:
-                self.add(f)
-        else:
-            for ifl in file_list:
-                self._summary.add_file(ifl)
-            self._servicex.put_file_add_bulk(file_list)
-
-
-async def run_file_fetch_loop(
-    did: str,
-    servicex: ServiceXAdapter,
-    info: Dict[str, Any],
-    user_callback: UserDIDHandler,
-):
-    start_time = datetime.now()
-
-    summary = DIDSummary(did)
-    did_info = parse_did_uri(did)
-    hold_till_end = did_info.file_count != -1
-    acc = _accumulator(servicex, summary, hold_till_end)
-
-    try:
-        async for file_info in user_callback(did_info.did, info):
-            if type(file_info) is dict:
-                acc.add(file_info)
-            else:
-                acc.send_bulk(file_info)
-
-    except Exception:
-        if did_info.get_mode == "all":
-            raise
-
-    # If we've been holding onto any files, we need to send them now.
-    acc.send_on(did_info.file_count)
-
-    elapsed_time = int((datetime.now() - start_time).total_seconds())
-    servicex.put_fileset_complete(
-        {
-            "files": summary.file_count,
-            "files-skipped": summary.files_skipped,
-            "total-events": summary.total_events,
-            "total-bytes": summary.total_bytes,
-            "elapsed-time": elapsed_time,
-        }
-    )
-
-
-def rabbit_mq_callback(
-    user_callback: UserDIDHandler, channel, method, properties, body
-):
-    """rabbit_mq_callback Respond to RabbitMQ Message
-
-    When a request to resolve a DID comes into the DID finder, we
-    respond with this callback. This callback will remain active
-    until the request has been completed satisfied (so for some
-    DID finders, this could be a fairly long time.)
-
-    Args:
-        channel ([type]): RabbitMQ channel
-        method ([type]): Delivery method
-        properties ([type]): Properties of the message
-        body ([type]): The body (json for us) of the message
-    """
-    dataset_id = None  # set this in case we get an exception while loading request
-    try:
-        # Unpack the message. Really bad if we fail up here!
-        did_request = json.loads(body)
-        did = did_request["did"]
-        dataset_id = did_request["dataset_id"]
-        endpoint = did_request["endpoint"]
-        __logging.info(
-            f"Received DID request {did_request}",
-            extra={"dataset_id": dataset_id}
-        )
-        servicex = ServiceXAdapter(dataset_id=dataset_id,
-                                   endpoint=endpoint)
-
-        info = {
-            "dataset-id": dataset_id,
-        }
-
-        # Process the request and resolve the DID
-        try:
-            make_sync(run_file_fetch_loop)(did, servicex, info, user_callback)
-
-        except Exception as e:
-            _, exec_value, _ = sys.exc_info()
-            __logging.exception(f"DID Request Failed {str(e)}", extra={"dataset_id": dataset_id})
-            raise
-
-    except Exception as e:
-        __logging.exception(
-            f"DID request failed {str(e)}", extra={"dataset_id": dataset_id}
-        )
-
-    finally:
-        channel.basic_ack(delivery_tag=method.delivery_tag)
-
-
-def init_rabbit_mq(
-    user_callback: UserDIDHandler,
-    rabbitmq_url: str,
-    queue_name: str,
-    retries: int,
-    retry_interval: float,
-):  # type: ignore
-    rabbitmq = None
-    retry_count = 0
-
-    while not rabbitmq:
-        try:
-            rabbitmq = pika.BlockingConnection(pika.URLParameters(rabbitmq_url))
-            _channel = rabbitmq.channel()
-            _channel.queue_declare(queue=queue_name)
-
-            __logging.info("Connected to RabbitMQ. Ready to start consuming requests")
-
-            _channel.basic_consume(
-                queue=queue_name,
-                auto_ack=False,
-                on_message_callback=lambda c, m, p, b: rabbit_mq_callback(
-                    user_callback, c, m, p, b
-                ),
-            )
-            _channel.start_consuming()
-
-        except pika.exceptions.AMQPConnectionError:  # type: ignore
-            rabbitmq = None
-            retry_count += 1
-            if retry_count <= retries:
-                __logging.exception(
-                    f"Failed to connect to RabbitMQ at {rabbitmq_url} "
-                    f"(try #{retry_count}). Waiting {retry_interval} seconds "
-                    "before trying again"
-                )
-                time.sleep(retry_interval)
-            else:
-                __logging.exception(
-                    f"Failed to connect to RabbitMQ. Giving Up after {retry_count}"
-                    " tries"
-                )
-                raise
-
-
-def add_did_finder_cnd_arguments(parser: argparse.ArgumentParser):
-    """add_did_finder_cnd_arguments Add required arguments to a parser
-
-    If you need to parse command line arguments for some special configuration, create your
-    own argument parser, and call this function to make sure the arguments needed
-    for running the back-end communication are filled in properly.
-
-    Then pass the results of the parsing to the `start_did_finder` method.
-
-    Args:
-        parser (argparse.ArgumentParser): The argument parser. Arguments needed for the
-                                          did finder/servicex communication will be added.
-    """
-    parser.add_argument(
-        "--rabbit-uri", dest="rabbit_uri", action="store", required=True
-    )
-    parser.add_argument(
-        "--prefix",
-        dest="prefix",
-        action="store",
-        required=False,
-        default="",
-        help="Prefix to add to use a caching proxy for URIs",
-    )
-
-
-def start_did_finder(
-    did_finder_name: str,
-    callback: UserDIDHandler,
-    parsed_args: Optional[argparse.Namespace] = None,
-):
-    """start_did_finder Start the DID finder
-
-    Top level method that starts the DID finder, hooking it up to rabbitmq queues, etc.,
-    and sets up the callback to be called each time ServiceX wants to render a DID into
-    files.
-
-    Once called this will not return unless it totally fails to connect to the rabbit
-    mq server (which must have been specified on the command line that started this).
-    If it can't return, then it will raise the appropriate exception.
-
-    Args:
-        did_name (str): Name of the DID finder (baked into the rabbit mq name)
-        callback (UserDIDHandler): Callback to handle the DID rendering requests
-        parser (Optional[argparse.Namespace], optional): If you need to parse your own
-                                                         an arg parser, make sure to call
-                                                         add_did_finder_cnd_arguments, run the
-                                                         parser and pass in the resulting
-                                                         parsed arguments.
-                                                         Defaults to None (automatically
-                                                         parses)
-                                                         command line arguments, create
-    """
-    # Setup command line parsing
-    if parsed_args is None:
-        parser = argparse.ArgumentParser()
-        add_did_finder_cnd_arguments(parser)
-        parsed_args = parser.parse_args()
-
-    # Initialize the root logger
-    initialize_root_logger(did_finder_name)
-
-    # Start up rabbit mq and also callbacks
-    init_rabbit_mq(
-        callback,
-        parsed_args.rabbit_uri,
-        f"{did_finder_name}{QUEUE_NAME_POSTFIX}",
-        retries=12,
-        retry_interval=10,
-    )