Skip to content

proxystore.stream.shims.kafka

Kafka publisher and subscriber shims.

Shims to the confluent-kafka package.

KafkaPublisher 

KafkaPublisher(client: Producer)

Kafka publisher shim.

Messages are delivered in the background, and close() waits for every message to be delivered.

Parameters:

  • client (Producer) –

    Kafka producer client.

Source code in proxystore/stream/shims/kafka.py 
29
30
def __init__(self, client: confluent_kafka.Producer) -> None:
    self.client = client

close 

close() -> None

Close this publisher.

Source code in proxystore/stream/shims/kafka.py 
32
33
34
def close(self) -> None:
    """Close this publisher."""
    self.client.flush()

send_message 

send_message(topic: str, message: bytes) -> None

Publish a message to the stream.

Parameters:

  • topic (str) –

    Stream topic to publish message to.

  • message (bytes) –

    Message as bytes to publish to the stream.

Source code in proxystore/stream/shims/kafka.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def send_message(self, topic: str, message: bytes) -> None:
    """Publish a message to the stream.

    Args:
        topic: Stream topic to publish message to.
        message: Message as bytes to publish to the stream.
    """
    # Messages are sent in the background and flushed in close().
    # Flushing after each message would wait for every message to be
    # delivered before the next is sent.
    while True:
        try:
            self.client.produce(topic, message)
        except BufferError:
            # The local queue is full, so wait for some messages to be
            # delivered and try again.
            self.client.poll(_FULL_QUEUE_POLL_TIMEOUT)
        else:
            break
    # Serve delivery callbacks of earlier messages without blocking.
    self.client.poll(0)

KafkaSubscriber 

KafkaSubscriber(client: Consumer)

Kafka subscriber shim.

This shim is an iterable object which will yield bytes messages from the stream, blocking on the next message, until the stream is closed.

Parameters:

  • client (Consumer) –

    Kafka consumer client. The client must already be subscribed to the relevant topics.

Source code in proxystore/stream/shims/kafka.py 
71
72
def __init__(self, client: confluent_kafka.Consumer) -> None:
    self.client = client

next_message 

next_message() -> bytes

Get the next message.

Source code in proxystore/stream/shims/kafka.py 
80
81
82
83
84
85
86
87
def next_message(self) -> bytes:
    """Get the next message."""
    message = self.client.poll()
    # Should not be None because we do not specify a poll in timeout.
    assert message is not None
    if message.error() is not None:  # pragma: no cover
        raise confluent_kafka.KafkaException(message.error())
    return message.value()

close 

close() -> None

Close this subscriber.

Source code in proxystore/stream/shims/kafka.py 
89
90
91
def close(self) -> None:
    """Close this subscriber."""
    self.client.close()