kafka-python 3.0.0__py3-none-any.whl

kafka/benchmarks/producer_encode_path.py

@@ -0,0 +1,201 @@
+#!/usr/bin/env python3
+"""Benchmarks for the producer encode hot path.
+
+Measures the cost of finalizing a record batch and encoding it into a
+ProduceRequest - the pipeline that runs on every send to the broker.
+
+To compare two implementations (e.g. before/after a change) run this
+script twice and diff the output:
+
+    # baseline
+    git stash
+    python -m kafka.benchmarks.producer_encode_path -o baseline.json
+    git stash pop
+
+    # new
+    python -m kafka.benchmarks.producer_encode_path -o new.json
+
+    pyperf compare_to baseline.json new.json
+
+For an allocation count (tracemalloc) run with --allocations, which is
+a separate mode (not a timing benchmark).
+"""
+import argparse
+import sys
+import tracemalloc
+
+import pyperf
+
+from kafka.producer.producer_batch import ProducerBatch
+from kafka.record.memory_records import MemoryRecordsBuilder
+from kafka.structs import TopicPartition
+from kafka.protocol.producer import ProduceRequest
+
+
+DEFAULT_RECORDS_PER_BATCH = 100
+DEFAULT_VALUE_SIZE = 128
+DEFAULT_BATCH_SIZE = 1 << 20  # 1 MiB - big enough to hold the largest case
+
+
+def _build_unclosed_batch(num_records, value_size):
+    tp = TopicPartition('bench', 0)
+    records = MemoryRecordsBuilder(
+        magic=2, compression_type=0, batch_size=DEFAULT_BATCH_SIZE)
+    batch = ProducerBatch(tp, records)
+    value = b'x' * value_size
+    for i in range(num_records):
+        future = batch.try_append(i, None, value, [])
+        assert future is not None, 'batch too small for %d records' % num_records
+    return batch
+
+
+def _build_closed_batch(num_records, value_size):
+    batch = _build_unclosed_batch(num_records, value_size)
+    batch.records.close()
+    return batch
+
+
+def bench_build_and_close(loops, num_records, value_size):
+    """Time build + close - isolates the batch-finalization path.
+
+    close() is fast enough that pyperf would need many thousands of loops
+    to measure it directly, so we include the build cost and rely on the
+    per-run diff to expose the close() delta.
+    """
+    value = b'x' * value_size
+    tp = TopicPartition('bench', 0)
+
+    t0 = pyperf.perf_counter()
+    for _ in range(loops):
+        records = MemoryRecordsBuilder(
+            magic=2, compression_type=0, batch_size=DEFAULT_BATCH_SIZE)
+        batch = ProducerBatch(tp, records)
+        for i in range(num_records):
+            batch.try_append(i, None, value, [])
+        batch.records.close()
+    return pyperf.perf_counter() - t0
+
+
+def bench_encode(loops, num_records, value_size):
+    """Time ProduceRequest.encode() with a pre-closed batch.
+
+    Isolates the protocol encoding cost from batch finalization.
+    """
+    batch = _build_closed_batch(num_records, value_size)
+    buf = batch.records.buffer()
+    topic_data = [('bench', [(0, buf)])]
+
+    t0 = pyperf.perf_counter()
+    for _ in range(loops):
+        req = ProduceRequest[8](
+            transactional_id=None,
+            acks=1,
+            timeout_ms=30000,
+            topic_data=topic_data,
+        )
+        req.with_header(correlation_id=1, client_id='bench')
+        req.encode(framed=True, header=True)
+    return pyperf.perf_counter() - t0
+
+
+def bench_full_pipeline(loops, num_records, value_size):
+    """Time the full producer hot path: append -> close -> encode.
+
+    Mirrors what the Sender thread does for each drained batch.
+    """
+    value = b'x' * value_size
+    tp = TopicPartition('bench', 0)
+
+    t0 = pyperf.perf_counter()
+    for _ in range(loops):
+        records = MemoryRecordsBuilder(
+            magic=2, compression_type=0, batch_size=DEFAULT_BATCH_SIZE)
+        batch = ProducerBatch(tp, records)
+        for i in range(num_records):
+            batch.try_append(i, None, value, [])
+        batch.records.close()
+        req = ProduceRequest[8](
+            transactional_id=None,
+            acks=1,
+            timeout_ms=30000,
+            topic_data=[('bench', [(0, batch.records.buffer())])],
+        )
+        req.with_header(correlation_id=1, client_id='bench')
+        req.encode(framed=True, header=True)
+    return pyperf.perf_counter() - t0
+
+
+def report_allocations(num_records, value_size):
+    """Use tracemalloc to count bytes allocated during a single encode."""
+    # Pre-build so we only measure the encode path.
+    batch = _build_closed_batch(num_records, value_size)
+    buf = batch.records.buffer()
+    topic_data = [('bench', [(0, buf)])]
+    req = ProduceRequest[8](
+        transactional_id=None,
+        acks=1,
+        timeout_ms=30000,
+        topic_data=topic_data,
+    )
+    req.with_header(correlation_id=1, client_id='bench')
+
+    # Warmup - populate the compiled-encoder cache so we don't count its
+    # one-time allocations.
+    req.encode(framed=True, header=True)
+
+    tracemalloc.start()
+    snap_before = tracemalloc.take_snapshot()
+    data = req.encode(framed=True, header=True)
+    snap_after = tracemalloc.take_snapshot()
+    tracemalloc.stop()
+
+    stats = snap_after.compare_to(snap_before, 'lineno')
+    total_bytes = sum(max(0, s.size_diff) for s in stats)
+    total_allocs = sum(max(0, s.count_diff) for s in stats)
+
+    print('Config        : %d records x %d bytes' % (num_records, value_size))
+    print('Wire bytes    : %d' % len(data))
+    print('Allocated     : %d bytes across %d allocations' % (total_bytes, total_allocs))
+    print()
+    print('Top allocation sites:')
+    for s in sorted(stats, key=lambda s: s.size_diff, reverse=True)[:10]:
+        if s.size_diff <= 0:
+            continue
+        frame = s.traceback[0]
+        loc = '%s:%d' % (frame.filename.split('/workspace/')[-1], frame.lineno)
+        print('  %8d bytes / %3d allocs  %s' % (s.size_diff, s.count_diff, loc))
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__,
+                                formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument('--records', type=int, default=DEFAULT_RECORDS_PER_BATCH,
+                   help='records per batch (default: %(default)d)')
+    p.add_argument('--value-size', type=int, default=DEFAULT_VALUE_SIZE,
+                   help='record value size in bytes (default: %(default)d)')
+    p.add_argument('--allocations', action='store_true',
+                   help='report tracemalloc allocations for one encode (not a timing run)')
+
+    # Let pyperf parse its own arguments (--output, --rigorous, etc).
+    args, pyperf_args = p.parse_known_args()
+
+    if args.allocations:
+        report_allocations(args.records, args.value_size)
+        return
+
+    # Rebuild sys.argv so pyperf sees only its own arguments.
+    sys.argv = [sys.argv[0]] + pyperf_args
+
+    runner = pyperf.Runner()
+    runner.metadata['records_per_batch'] = args.records
+    runner.metadata['value_size_bytes'] = args.value_size
+    runner.bench_time_func(
+        'build_and_close', bench_build_and_close, args.records, args.value_size)
+    runner.bench_time_func(
+        'encode', bench_encode, args.records, args.value_size)
+    runner.bench_time_func(
+        'full_pipeline', bench_full_pipeline, args.records, args.value_size)
+
+
+if __name__ == '__main__':
+    main()

kafka/benchmarks/producer_performance.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python
+# Adapted from https://github.com/mrafayaleem/kafka-jython
+
+import argparse
+import pprint
+import sys
+import threading
+import time
+import traceback
+
+from kafka import KafkaProducer
+
+
+class ProducerPerformance:
+    @staticmethod
+    def run(args):
+        try:
+            props = {}
+            for prop in args.producer_config:
+                k, v = prop.split('=')
+                try:
+                    v = int(v)
+                except ValueError:
+                    pass
+                if v == 'None':
+                    v = None
+                elif v == 'False':
+                    v = False
+                elif v == 'True':
+                    v = True
+                props[k] = v
+
+            print('Initializing producer...')
+            props['bootstrap_servers'] = args.bootstrap_servers
+            record = bytes(bytearray(args.record_size))
+            props['metrics_sample_window_ms'] = args.stats_interval * 1000
+
+            producer = KafkaProducer(**props)
+            for k, v in props.items():
+                print('---> {0}={1}'.format(k, v))
+            print('---> topic={0}'.format(args.topic))
+            print('---> send {0} byte records'.format(args.record_size))
+            print('---> report stats every {0} secs'.format(args.stats_interval))
+            print('---> raw metrics? {0}'.format(args.raw_metrics))
+            timer_stop = threading.Event()
+            timer = StatsReporter(args.stats_interval, producer,
+                                  event=timer_stop,
+                                  raw_metrics=args.raw_metrics)
+            timer.start()
+            print('-> OK!')
+            print()
+
+            def _benchmark():
+                # Hoist attribute lookups and use positional args to minimize
+                # per-iteration harness overhead. With a well-optimized
+                # producer.send(), kwargs dict packing and repeated attribute
+                # dereference on `args`/`producer` dominate main-thread CPU
+                # inside this loop and mask the library's true throughput.
+                results = []
+                send = producer.send
+                append = results.append
+                topic = args.topic
+                num_records = args.num_records
+                for _ in range(num_records):
+                    append(send(topic, record))
+                print("Send complete...")
+                producer.flush()
+                producer.close()
+                count_success, count_failure = 0, 0
+                for r in results:
+                    if r.succeeded():
+                        count_success += 1
+                    elif r.failed():
+                        count_failure += 1
+                    else:
+                        raise ValueError(r)
+                print("%d suceeded, %d failed" % (count_success, count_failure))
+
+            start_time = time.monotonic()
+            _benchmark()
+            end_time = time.monotonic()
+            timer_stop.set()
+            timer.join()
+            print('Execution time:', end_time - start_time, 'secs')
+
+        except Exception:
+            exc_info = sys.exc_info()
+            traceback.print_exception(*exc_info)
+            sys.exit(1)
+
+
+class StatsReporter(threading.Thread):
+    def __init__(self, interval, producer, event=None, raw_metrics=False):
+        super().__init__()
+        self.interval = interval
+        self.producer = producer
+        self.event = event
+        self.raw_metrics = raw_metrics
+
+    def print_stats(self):
+        metrics = self.producer.metrics()
+        if not metrics:
+            return
+        if self.raw_metrics:
+            pprint.pprint(metrics)
+        else:
+            print('{record-send-rate:.0f} records/sec ({byte-rate:.0f} B/sec),'
+                  ' {request-rate:.0f} avg requests/sec,'
+                  ' {request-latency-avg:.0f}ms avg latency,'
+                  ' {throttle-time-max:.0f}ms max throttle,'
+                  ' {record-size-avg:.0f} avg record size,'
+                  ' {batch-size-avg:.0f} avg batch size,'
+                  ' {records-per-request-avg:.0f} avg records/req'
+                  .format(**metrics['producer-metrics']))
+
+    def print_final(self):
+        self.print_stats()
+
+    def run(self):
+        while self.event and not self.event.wait(self.interval):
+            self.print_stats()
+        else:
+            self.print_final()
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser(
+        description='This tool is used to verify the producer performance.')
+
+    parser.add_argument(
+        '-b', '--bootstrap-servers', type=str, nargs='+', default=(),
+        help='host:port for cluster bootstrap server')
+    parser.add_argument(
+        '-t', '--topic', type=str,
+        help='Topic name for test (default: kafka-python-benchmark-test)',
+        default='kafka-python-benchmark-test')
+    parser.add_argument(
+        '--num-records', type=int,
+        help='number of messages to produce (default: 1000000)',
+        default=1000000)
+    parser.add_argument(
+        '--record-size', type=int,
+        help='message size in bytes (default: 100)',
+        default=100)
+    parser.add_argument(
+        '-c', '--producer-config', type=str, nargs='+', default=(),
+        help='kafka producer related configuaration properties like '
+             'bootstrap_servers,client_id etc..')
+    parser.add_argument(
+        '--stats-interval', type=int,
+        help='Interval in seconds for stats reporting to console (default: 5)',
+        default=5)
+    parser.add_argument(
+        '--raw-metrics', action='store_true',
+        help='Enable this flag to print full metrics dict on each interval')
+    return parser
+
+
+if __name__ == '__main__':
+    args = get_args_parser().parse_args()
+    ProducerPerformance.run(args)

kafka/benchmarks/profile_protocol.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""Per-function and per-line profiling of new protocol encode/decode hot methods.
+
+Uses cProfile for call-level breakdown and a targeted manual timer
+for line-level insight into the key hot methods.
+
+Usage:
+    python kafka/benchmarks/profile_protocol.py [encode|decode|create] [N]
+
+    encode  - profile encode path (default)
+    decode  - profile decode path
+    create  - profile object creation path
+    N       - number of iterations (default: 1000)
+"""
+
+import cProfile
+import io
+import pstats
+import sys
+import time
+
+from kafka.protocol.metadata.metadata import MetadataResponse as NewMetadataResp
+from kafka.protocol.consumer.fetch import FetchResponse as NewFetchResp
+
+# === Test data (same as benchmark script) ===
+_BROKERS = [(i, 'broker-%d.example.com' % i, 9092 + i) for i in range(3)]
+_PARTITIONS = [(0, p, p % 3, [0, 1, 2], [0, 1]) for p in range(3)]
+_TOPICS = [(0, 'test-topic-%d' % t, _PARTITIONS) for t in range(5)]
+
+_RECORDS_BYTES = b'\x00' * 1024
+_FETCH_TOPICS = [
+    ('test-topic-%d' % t, [
+        (p, 0, 1000 + p, 999 + p, None, _RECORDS_BYTES)
+        for p in range(3)
+    ])
+    for t in range(2)
+]
+
+
+def profile_encode(n):
+    """Profile the encode path for MetadataResponse and FetchResponse."""
+    obj_meta = NewMetadataResp(version=0, brokers=_BROKERS, topics=_TOPICS)
+    obj_fetch = NewFetchResp(version=4, throttle_time_ms=0, responses=_FETCH_TOPICS)
+
+    # Warmup
+    obj_meta.encode(version=0)
+    obj_fetch.encode(version=4)
+
+    pr = cProfile.Profile()
+    pr.enable()
+    for _ in range(n):
+        obj_meta.encode(version=0)
+        obj_fetch.encode(version=4)
+    pr.disable()
+
+    print("=" * 80)
+    print("ENCODE PROFILE (%d iterations: MetadataResponse v0 + FetchResponse v4)" % n)
+    print("=" * 80)
+    stats = pstats.Stats(pr)
+    stats.sort_stats('cumulative')
+    stats.print_stats(40)
+    print()
+    stats.sort_stats('tottime')
+    stats.print_stats(40)
+
+
+def profile_decode(n):
+    """Profile the decode path for MetadataResponse and FetchResponse."""
+    obj_meta = NewMetadataResp(version=0, brokers=_BROKERS, topics=_TOPICS)
+    obj_fetch = NewFetchResp(version=4, throttle_time_ms=0, responses=_FETCH_TOPICS)
+    encoded_meta = obj_meta.encode(version=0)
+    encoded_fetch = obj_fetch.encode(version=4)
+
+    # Warmup
+    NewMetadataResp.decode(io.BytesIO(encoded_meta), version=0)
+    NewFetchResp.decode(io.BytesIO(encoded_fetch), version=4)
+
+    pr = cProfile.Profile()
+    pr.enable()
+    for _ in range(n):
+        NewMetadataResp.decode(io.BytesIO(encoded_meta), version=0)
+        NewFetchResp.decode(io.BytesIO(encoded_fetch), version=4)
+    pr.disable()
+
+    print("=" * 80)
+    print("DECODE PROFILE (%d iterations: MetadataResponse v0 + FetchResponse v4)" % n)
+    print("=" * 80)
+    stats = pstats.Stats(pr)
+    stats.sort_stats('cumulative')
+    stats.print_stats(40)
+    print()
+    stats.sort_stats('tottime')
+    stats.print_stats(40)
+
+
+def profile_create(n):
+    """Profile the object creation path."""
+    # Warmup
+    NewMetadataResp(version=0, brokers=_BROKERS, topics=_TOPICS)
+    NewFetchResp(version=4, throttle_time_ms=0, responses=_FETCH_TOPICS)
+
+    pr = cProfile.Profile()
+    pr.enable()
+    for _ in range(n):
+        NewMetadataResp(version=0, brokers=_BROKERS, topics=_TOPICS)
+        NewFetchResp(version=4, throttle_time_ms=0, responses=_FETCH_TOPICS)
+    pr.disable()
+
+    print("=" * 80)
+    print("CREATE PROFILE (%d iterations: MetadataResponse v0 + FetchResponse v4)" % n)
+    print("=" * 80)
+    stats = pstats.Stats(pr)
+    stats.sort_stats('cumulative')
+    stats.print_stats(40)
+    print()
+    stats.sort_stats('tottime')
+    stats.print_stats(40)
+
+
+if __name__ == '__main__':
+    mode = sys.argv[1] if len(sys.argv) > 1 else 'encode'
+    iterations = int(sys.argv[2]) if len(sys.argv) > 2 else 1000
+
+    if mode == 'encode':
+        profile_encode(iterations)
+    elif mode == 'decode':
+        profile_decode(iterations)
+    elif mode == 'create':
+        profile_create(iterations)
+    elif mode == 'all':
+        profile_encode(iterations)
+        print("\n\n")
+        profile_decode(iterations)
+        print("\n\n")
+        profile_create(iterations)
+    else:
+        print("Usage: python profile_protocol.py [encode|decode|create|all] [N]")
+        sys.exit(1)