Handle BufferErrors by backing off and retrying (#673)

Author: moritzmeister

python/hsfs/engine/python.py

@@ -724,6 +724,10 @@ def _write_dataframe_kafka(
         # setup row writer function
         writer = self._get_encoder_func(feature_group._get_encoded_avro_schema())
 
+        def acked(err, msg):
+            if err is not None:
+                print("Failed to deliver message: %s: %s" % (str(msg), str(err)))
+
         # loop over rows
         for r in dataframe.itertuples(index=False):
             # itertuples returns Python NamedTyple, to be able to serialize it using
@@ -752,13 +756,27 @@ def _write_dataframe_kafka(
             # assemble key
             key = "".join([str(row[pk]) for pk in sorted(feature_group.primary_key)])
 
-            # produce
-            producer.produce(
-                topic=feature_group._online_topic_name, key=key, value=encoded_row
-            )
+            while True:
+                # if BufferError is thrown, we can be sure, message hasn't been send so we retry
+                try:
+                    # produce
+                    producer.produce(
+                        topic=feature_group._online_topic_name,
+                        key=key,
+                        value=encoded_row,
+                        callback=acked
+                        if offline_write_options.get("debug_kafka", False)
+                        else None,
+                    )
 
-            # Trigger internal callbacks to empty op queue
-            producer.poll(0)
+                    # Trigger internal callbacks to empty op queue
+                    producer.poll(0)
+                    break
+                except BufferError as e:
+                    if offline_write_options.get("debug_kafka", False):
+                        print("Caught: {}".format(e))
+                    # backoff for 1 second
+                    producer.poll(1)
 
         # make sure producer blocks and everything is delivered
         producer.flush()
@@ -803,12 +821,15 @@ def _get_encoder_func(self, writer_schema: str) -> callable:
         return lambda record, outf: writer.write(record, avro.io.BinaryEncoder(outf))
 
     def _get_kafka_config(self, write_options: dict = {}) -> dict:
+        # producer configuration properties
+        # https://docs.confluent.io/platform/current/clients/librdkafka/html/md_CONFIGURATION.html
         config = {
             "security.protocol": "SSL",
             "ssl.ca.location": client.get_instance()._get_ca_chain_path(),
             "ssl.certificate.location": client.get_instance()._get_client_cert_path(),
             "ssl.key.location": client.get_instance()._get_client_key_path(),
             "client.id": socket.gethostname(),
+            **write_options.get("kafka_producer_config", {}),
         }
 
         if isinstance(client.get_instance(), hopsworks.Client) or write_options.get(
@@ -831,5 +852,4 @@ def _get_kafka_config(self, write_options: dict = {}) -> dict:
                     )
                 ]
             )
-
         return config