Offset management and multiprocessing embryo (#173)

* Update doc

* Add method to return Kafka offset. Doc missing.

* WIP multiprocessing

* Update method documentation

* PEP8

* merge multiprocessing

* Update requirements

* Reset offset properly

* Manage negative offsets

* Better handling of the verbosity

* Update README

* PEP8

* PEP8

* Bump version to 5.0

* Update README

Author: JulienPeloton

README.md

@@ -13,26 +13,28 @@
 
 `fink_client` requires a version of Python 3.9+.
 
-### install with pip
+### Install with pip
 
 ```bash
 pip install fink-client --upgrade
 ```
 
-### install within a conda environment
+### Use or develop in a controlled environment
 
 ```bash
 git clone https://github.com/astrolabsoftware/fink-client.git
 cd fink-client
 conda env create -f environment.yml
-# install the latest release
-pip install fink-client
+conda activate fink-client
+```
+
+and then install it locally:
+
+```bash
 # install for development
 pip install -e .
 ```
 
-Learn how to connect and use it by checking the [documentation](docs/).
-
 ## Registration
 
 In order to connect and poll alerts from Fink, you need to get your credentials:
@@ -83,16 +85,14 @@ optional arguments:
   -filename FILENAME  Path to an alert data file (avro format)
 ```
 
-More information at [docs/livestream](docs/livestream_manual.md).
+More information at [docs/livestream](https://fink-broker.readthedocs.io/en/latest/services/livestream).
 
 ## Data Transfer usage
 
 If you requested data using the [Data Transfer service](https://fink-portal.org/download), you can easily poll your stream using:
 
 ```bash
-fink_datatransfer -h
-usage: fink_datatransfer [-h] [-topic TOPIC] [-limit LIMIT] [-outdir OUTDIR] [-partitionby PARTITIONBY] [-batchsize BATCHSIZE] [--restart_from_beginning]
-                            [--verbose]
+usage: fink_datatransfer [-h] [-topic TOPIC] [-limit LIMIT] [-outdir OUTDIR] [-partitionby PARTITIONBY] [-batchsize BATCHSIZE] [-nconsumers NCONSUMERS] [-maxtimeout MAXTIMEOUT] [--restart_from_beginning] [--verbose]
 
 Kafka consumer to listen and archive Fink streams from the data transfer service
 
@@ -102,13 +102,16 @@ optional arguments:
   -limit LIMIT          If specified, download only `limit` alerts from the stream. Default is None, that is download all alerts.
   -outdir OUTDIR        Folder to store incoming alerts. It will be created if it does not exist.
   -partitionby PARTITIONBY
-                        Partition data by `time` (year=YYYY/month=MM/day=DD), or `finkclass` (finkclass=CLASS), or `tnsclass` (tnsclass=CLASS). Default is
-                        time.
+                        Partition data by `time` (year=YYYY/month=MM/day=DD), or `finkclass` (finkclass=CLASS), or `tnsclass` (tnsclass=CLASS). `classId` is also available for ELASTiCC data. Default is time.
   -batchsize BATCHSIZE  Maximum number of alert within the `maxtimeout` (see conf). Default is 1000 alerts.
+  -nconsumers NCONSUMERS
+                        Number of parallel consumer to use. Default is 1.
+  -maxtimeout MAXTIMEOUT
+                        Overwrite the default timeout (in seconds) from user configuration. Default is None.
   --restart_from_beginning
                         If specified, restart downloading from the 1st alert in the stream. Default is False.
   --verbose             If specified, print on screen information about the consuming.
 
 ```
 
-More information at [docs/datatransfer](docs/datatransfer.md).
+More information at [docs/datatransfer](https://fink-broker.readthedocs.io/en/latest/services/data_transfer/).

docs/datatransfer.md

@@ -2,7 +2,7 @@
 
 _date 26/01/2023_
 
-This manual has been tested for `fink-client` version 4.1. In case of trouble, send us an email ([email protected]) or open an issue (https://github.com/astrolabsoftware/fink-client).
+This manual has been tested for `fink-client` version 4.4. In case of trouble, send us an email ([email protected]) or open an issue (https://github.com/astrolabsoftware/fink-client).
 
 ## Installation of fink-client
 
@@ -90,6 +90,10 @@ fink_datatransfer \
     --restart_from_beginning
 ```
 
+## Re-using the same queue
+
+In case you
+
 ## Reading alerts
 
 Alerts are saved in the Apache Parquet format. Assuming you are using Python, you can easily read them using Pandas:
@@ -107,5 +111,7 @@ In case of trouble, send us an email ([email protected]) or open an issue
 
 ### Known bugs
 
-1. Data from 2019/2020/2021 and 2022/2023 are not compatible (different schemas). We will resolve the problem soon, but in the meantime, do not mix data from the two periods in a single query. 
+1. Data from 2019/2020/2021 and 2022/2023 are not compatible (different schemas). We will resolve the problem soon, but in the meantime, do not mix data from the two periods in a single query.
 2. With version 4.0, you wouldn't have the partitioning column when reading in a dataframe. This has been corrected in 4.1.
+3. If you have recurrent timeouts, try to increase the timeout in your configuration file `~/.finkclient/credentials.yml`.
+4.

environment.yml

@@ -15,4 +15,5 @@ dependencies:
   - pyyaml
   - tabulate
   - matplotlib
+  - tqdm
 

fink_client/__init__.py

@@ -12,5 +12,5 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-__version__ = "4.4"
+__version__ = "5.0"
 __schema_version__ = "distribution_schema_fink_ztf_{}.avsc"

fink_client/consumer.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 import io
 import json
+import time
 import fastavro
 import confluent_kafka
 
@@ -267,6 +268,87 @@ def close(self):
         self._consumer.close()
 
 
+def return_offsets(consumer, topic, waitfor=1, timeout=10, verbose=False):
+    """ Poll servers to get the total committed offsets, and remaining lag
+
+    Parameters
+    ----------
+    consumer: confluent_kafka.Consumer
+        Kafka consumer
+    topic: str
+        Topic name
+    waitfor: int, optional
+        Time in second to wait before polling. Default is 1 second.
+    timeout: int, optional
+        Timeout in second when polling the servers. Default is 10.
+    verbose: bool, optional
+        If True, prints useful table. Default is False.
+
+    Returns
+    ---------
+    total_offsets: int
+        Total number of messages committed across all partitions
+    total_lag: int
+        Remaining messages in the topic across all partitions.
+    """
+    time.sleep(waitfor)
+    # Get the topic's partitions
+    metadata = consumer.list_topics(topic, timeout=timeout)
+    if metadata.topics[topic].error is not None:
+        raise confluent_kafka.KafkaException(metadata.topics[topic].error)
+
+    # Construct TopicPartition list of partitions to query
+    partitions = [confluent_kafka.TopicPartition(topic, p) for p in metadata.topics[topic].partitions]
+
+    # Query committed offsets for this group and the given partitions
+    try:
+        committed = consumer.committed(partitions, timeout=timeout)
+    except confluent_kafka.KafkaException as exception:
+        kafka_error = exception.args[0]
+        if kafka_error.code() == confluent_kafka.KafkaError._TIMED_OUT:
+            return -1, -1
+        else:
+            return 0, 0
+
+    total_offsets = 0
+    total_lag = 0
+    if verbose:
+        print("%-50s  %9s  %9s" % ("Topic [Partition]", "Committed", "Lag"))
+        print("=" * 72)
+    for partition in committed:
+        # Get the partitions low and high watermark offsets.
+        (lo, hi) = consumer.get_watermark_offsets(partition, timeout=timeout, cached=False)
+
+        if partition.offset == confluent_kafka.OFFSET_INVALID:
+            offset = "-"
+        else:
+            offset = "%d" % (partition.offset)
+
+        if hi < 0:
+            lag = "no hwmark"  # Unlikely
+        elif partition.offset < 0:
+            # No committed offset, show total message count as lag.
+            # The actual message count may be lower due to compaction
+            # and record deletions.
+            lag = "%d" % (hi - lo)
+            partition.offset = 0
+        else:
+            lag = "%d" % (hi - partition.offset)
+
+        total_offsets = partition.offset
+        total_lag = int(lag)
+
+        if verbose:
+            print("%-50s  %9s  %9s" % (
+                "{} [{}]".format(partition.topic, partition.partition), offset, lag))
+    if verbose:
+        print("-" * 72)
+        print("%-50s  %9s  %9s" % (
+            "Total", total_offsets, total_lag))
+        print("-" * 72)
+
+    return total_offsets, total_lag
+
 def _get_kafka_config(config: dict) -> dict:
     """Returns configurations for a consumer instance