Lab051: accelerate 500 ksymbol/s streaming receiver

experiments/lab051_receiver_performance.py

@@ -0,0 +1,320 @@
+"""Lab051: измеренный блочный RX и возраст JPEG в шестисекундном потоке.
+
+Источник синтетический, из реального фото. RF не выполняется. Виртуальное
+время работника использует фактическое время каждого push/сборки/JPEG decode.
+Это офлайн-воспроизведение; отставание явно означает требуемый внешний IQ backlog.
+"""
+
+from __future__ import annotations
+
+import argparse
+from collections import Counter
+from dataclasses import asdict
+import hashlib
+from io import BytesIO
+import json
+import math
+from pathlib import Path
+import statistics
+import sys
+import time
+
+if __package__ in (None, ""):
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+import numpy as np
+from PIL import Image
+
+from experiments import lab050_video_stream as video
+from experiments import lab051_streaming_receiver as fast
+from experiments.lab051_receiver_profile import OUTPUT, ROOT, StageProfiler, array_observation, process_memory
+from protocol.video_frame_scheduler import FramePolicy, schedule_video_frames
+
+
+def write_result(name: str, result: dict) -> None:
+    OUTPUT.mkdir(parents=True, exist_ok=True)
+    (OUTPUT / name).write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+
+
+class ImageSink:
+    """Только текущий объект; никаких ожидаемых ID, размеров JPEG или TX bytes."""
+
+    def __init__(self):
+        self.object_id = -1
+        self.last_completed_id = -1
+        self.assembler = None
+        self.first_rf_sample = None
+        self.first_rx_time = None
+        self.retained_payload_peak = 0
+        self.assembly_seconds = self.jpeg_decode_seconds = 0.0
+        self.rejected = Counter()
+
+    def offer(self, packet, received_time: float) -> dict | None:
+        started = time.perf_counter()
+        fragment = video.transport.recover_fragment_from_payload(packet.payload)
+        if fragment is None or fragment.original_size > 65536 or fragment.fragment_count > 269:
+            self.rejected["invalid_fragment"] += 1
+            self.assembly_seconds += time.perf_counter() - started
+            return None
+        if fragment.image_id < self.object_id or fragment.image_id <= self.last_completed_id:
+            self.rejected["older_object"] += 1
+            self.assembly_seconds += time.perf_counter() - started
+            return None
+        if fragment.image_id > self.object_id:
+            self.object_id = fragment.image_id
+            self.assembler = video.transport.ObjectAssembler()
+            self.first_rf_sample = packet.start_sample
+            self.first_rx_time = received_time
+        self.assembler.offer(fragment)
+        self.retained_payload_peak = max(self.retained_payload_peak, fragment.original_size)
+        jpeg = self.assembler.recovered_object()
+        self.assembly_seconds += time.perf_counter() - started
+        if jpeg is None:
+            return None
+        decode_started = time.perf_counter()
+        try:
+            with Image.open(BytesIO(jpeg)) as image:
+                image.load()
+                dimensions = image.size
+        except (OSError, ValueError):
+            self.rejected["jpeg_decode"] += 1
+            return None
+        finally:
+            self.jpeg_decode_seconds += time.perf_counter() - decode_started
+        self.last_completed_id = self.object_id
+        self.assembler = None
+        return {"image_id": self.object_id, "jpeg": jpeg, "dimensions": dimensions,
+                "first_rf_sample": self.first_rf_sample, "last_rf_sample": packet.end_sample,
+                "first_rx_time": self.first_rx_time}
+
+
+def run_stream(sequence: list[video.EncodedImage]) -> dict:
+    """Вход RX — только последовательные IQ-блоки; TX-диагностика после JPEG."""
+    mode = video.phy.PhyMode(500000)
+    receiver = fast.StreamingReceiver()
+    sink = ImageSink()
+    random_generator = np.random.default_rng(50050)
+    frames = video.make_frame_groups(sequence, 5)
+    packet_time = video.FRAME_INTERVALS / mode.symbol_rate
+    bitrate = (256 + video.HEADER_SIZE) * 8 / packet_time
+    # TX использует известную нижнюю задержку радио/буфера, без предсказания
+    # будущего CPU и без знания RX verdict. Перегрузка RX измеряется ниже.
+    buffer_delay = (receiver.lookahead + 2*receiver.core_samples) / mode.sample_rate_hz
+    schedule = schedule_video_frames(frames, (), FramePolicy.LATEST_ONLY, bitrate,
+                                     max_video_age_seconds=.2, completion_latency_seconds=buffer_delay)
+    starts = {}
+    for sent in schedule.transmitted:
+        starts.setdefault(sent.item.composite_frame_id, sent.start_seconds)
+    horizon = len(sequence) / 5
+    encoded_lookup = {image.image_id: image for image in sequence}
+    capture_lookup = {frame.composite_frame_id: frame.generation_time_seconds for frame in frames}
+    pending = np.empty(0, np.complex64)
+    total_samples = tx_packets = 0
+    worker_time = rx_processing = max_lag = 0.0
+    completed = []
+    block_metrics = []
+    rss_samples = []
+
+    def consume_block(block, *, final=False):
+        nonlocal total_samples, worker_time, rx_processing, max_lag
+        total_samples += block.size
+        arrival = total_samples / mode.sample_rate_hz
+        worker_start = max(worker_time, arrival)
+        started = time.perf_counter()
+        packets = receiver.finish() if final else receiver.push_iq_block(block)
+        cpu = time.perf_counter() - started
+        rx_processing += cpu
+        worker_time = worker_start + cpu
+        for packet in packets:
+            started = time.perf_counter()
+            item = sink.offer(packet, worker_time)
+            cost = time.perf_counter() - started
+            worker_time += cost
+            if item is not None:
+                # Приём/сборка/JPEG уже завершены. TX используется только теперь.
+                jpeg = item.pop("jpeg")
+                image_id = item["image_id"]
+                capture = capture_lookup.get(image_id)
+                age = None if capture is None else worker_time - capture
+                completed.append({**item, "decode_complete_timestamp": worker_time,
+                                  "capture_timestamp": capture, "image_age_seconds": age,
+                                  "fresh": age is not None and 0 <= age <= .2,
+                                  "sha256": hashlib.sha256(jpeg).hexdigest(),
+                                  "sha_match": image_id in encoded_lookup and jpeg == encoded_lookup[image_id].jpeg})
+        lag = max(0., worker_time - arrival)
+        max_lag = max(max_lag, lag)
+        block_metrics.append({"samples_received": total_samples, "arrival_time": arrival,
+                              "worker_start": worker_start, "worker_end": worker_time,
+                              "rx_cpu_seconds": cpu, "packets": len(packets), "lag_seconds": lag})
+        if len(block_metrics) % 64 == 0:
+            rss_samples.append({"samples": total_samples, **process_memory()})
+
+    def push(values):
+        nonlocal pending
+        offset = 0
+        if pending.size:
+            needed = receiver.block_samples - pending.size
+            joined = np.concatenate((pending, values[:needed]))
+            offset = min(needed, values.size)
+            if joined.size < receiver.block_samples:
+                pending = joined
+                return
+            consume_block(joined)
+            pending = np.empty(0, np.complex64)
+        while offset + receiver.block_samples <= values.size:
+            consume_block(values[offset:offset + receiver.block_samples])
+            offset += receiver.block_samples
+        pending = values[offset:].copy()
+
+    def background(count):
+        while count:
+            amount = min(count, receiver.block_samples)
+            values = .2 * (random_generator.normal(size=amount) + 1j*random_generator.normal(size=amount))
+            push(values.astype(np.complex64))
+            count -= amount
+
+    for image in sequence:
+        if image.image_id not in starts:
+            continue
+        target = math.ceil(starts[image.image_id] * mode.sample_rate_hz)
+        background(max(0, target - total_samples - pending.size))
+        for frame in video.transport.build_object_frames(image.jpeg, image.image_id):
+            raw = video.phy.build_parameterized_tx_cs8(frame.plan, mode)
+            components = np.frombuffer(raw, np.int8).astype(np.float32)
+            values = components[0::2] + 1j*components[1::2]
+            indexes = np.arange(values.size)
+            values = values * np.exp(1j*(.73 + 2*np.pi*10*indexes/mode.sample_rate_hz)) * (-1 if tx_packets % 2 else 1)
+            energy = float(np.sum(abs(values)**2)) / frame.plan.frame_symbol_count
+            sigma = math.sqrt(energy / (2*10**(.7)))
+            values += sigma*(random_generator.normal(size=values.size) + 1j*random_generator.normal(size=values.size))
+            push(values.astype(np.complex64))
+            tx_packets += 1
+    background(max(0, math.ceil(horizon * mode.sample_rate_hz) + receiver.lookahead + receiver.core_samples
+                   - total_samples - pending.size))
+    if pending.size:
+        consume_block(pending)
+    consume_block(np.empty(0,np.complex64), final=True)
+    fresh = [row for row in completed if row["fresh"]]
+    displayed_in_window = [row for row in fresh if row["decode_complete_timestamp"] <= horizon]
+    decoded_in_window = [row for row in completed if row["decode_complete_timestamp"] <= horizon]
+    ages = [row["image_age_seconds"] for row in completed if row["image_age_seconds"] is not None]
+    fresh_ages = [row["image_age_seconds"] for row in fresh]
+    processing = rx_processing + sink.assembly_seconds + sink.jpeg_decode_seconds
+    return {
+        "scope": "continuous synthetic IQ; measured per-block CPU event replay; camera/display latency UNKNOWN",
+        "symbol_rate":500000,"samples_per_symbol":16,"input_fps":5,"input_frames":len(sequence),
+        "observation_seconds":horizon,"iq_seconds":total_samples/mode.sample_rate_hz,
+        "tx_images":len(starts),"tx_physical_packets":tx_packets,
+        "dropped_before_tx":len(schedule.dropped_frames),"maximum_waiting_frames":schedule.maximum_waiting_frames,
+        "rx_physical_packets":receiver.packet_count,"received_images":len(completed),
+        "object_jpeg_integrity":len(completed)==len(starts) and all(row["sha_match"] for row in completed),
+        "dropped_stale_at_display":len(completed)-len(fresh),"displayed":len(fresh),"fresh_displayed":len(fresh),
+        "raw_received_fps":len(decoded_in_window)/horizon,
+        "raw_display_fps":len(displayed_in_window)/horizon,"fresh_display_fps":len(displayed_in_window)/horizon,
+        "drain_fresh_displayed":len(fresh)-len(displayed_in_window),
+        "image_age_min_seconds":min(ages) if ages else None,
+        "image_age_median_seconds":statistics.median(ages) if ages else None,
+        "image_age_max_seconds":max(ages) if ages else None,
+        "fresh_image_age_seconds":fresh_ages,"max_video_age_seconds":.2,
+        "processing_seconds":processing,"rx_processing_seconds":rx_processing,
+        "assembly_seconds":sink.assembly_seconds,"jpeg_decode_seconds":sink.jpeg_decode_seconds,
+        "processing_per_rf_second":processing/(total_samples/mode.sample_rate_hz),
+        "physical_packet_processing_seconds":processing/tx_packets if tx_packets else None,
+        "maximum_cpu_lag_seconds":max_lag,
+        "implied_external_iq_backlog_bytes":math.ceil(max_lag*mode.sample_rate_hz)*8,
+        "live_external_queue_implemented":False,
+        "retained_iq_peak_bytes":receiver.peak_buffer_samples*8,
+        "working_arrays_upper_estimate_bytes":receiver.memory_estimate_bytes + 20*receiver.buffer_limit_samples*8,
+        "object_payload_peak_bytes":sink.retained_payload_peak,"process_memory":process_memory(),
+        "rss_samples":rss_samples,"candidate_count":receiver.candidate_count,
+        "rejected_reasons":dict(receiver.rejected_reasons),"front_statistics":dict(receiver.front_statistics),
+        "decimated_samples_reused":receiver.decimator.reused_samples,
+        "actual_decimator_input_samples":receiver.decimator.filtered_input_samples,
+        "images":completed,"blocks":block_metrics,"rf_tx":0,
+    }
+
+
+def run_memory_test(block_count: int = 4096) -> dict:
+    """33,55 с без истории IQ/пакетов; нулевой фон отделён от noisy IQ опыта."""
+    receiver = fast.StreamingReceiver()
+    values = np.zeros(receiver.block_samples,np.complex64)
+    checkpoints=[]
+    for index in range(block_count):
+        assert receiver.push_iq_block(values) == []
+        assert receiver.buffer.size <= receiver.buffer_limit_samples
+        if (index+1) % 512 == 0:
+            checkpoints.append({"block":index+1,"iq_bytes":receiver.buffer.nbytes,
+                                "decimated_bytes":receiver.decimator.values.nbytes,**process_memory()})
+    assert receiver.finish()==[]
+    return {"blocks":block_count,"model_seconds":block_count*values.size/8000000,
+            "background":"zeros; noisy background covered by six-second integration",
+            "peak_iq_bytes":receiver.peak_buffer_samples*8,"iq_limit_bytes":receiver.buffer_limit_samples*8,
+            "retained_after_finish":receiver.buffer.nbytes+receiver.decimator.values.nbytes,
+            "checkpoints":checkpoints,"rf_tx":0}
+
+
+def run_candidate_replay() -> dict:
+    metadata=json.loads((OUTPUT/"optimized_examples.json").read_text(encoding="utf-8"))
+    result={}
+    for kind, entry in metadata.items():
+        raw=np.load(OUTPUT/f"optimized_{kind}_window.npy")
+        candidate=fast.offline.CoarseCandidate(entry["offset"],entry["score"])
+        def call(optimized=False):
+            with fast.previous.streaming_mode(video.phy.PhyMode(500000)):
+                def decode():
+                    return fast.offline.decode_candidate_from_memmap(raw,entry["sample_count"],candidate,1,
+                        read_iq_window=lambda data,start,end:data[2*start:2*end].view(np.complex64))
+                if not optimized:
+                    return decode()
+                front=fast.BlockFrontEnd(raw.view(np.complex64),65536,Counter())
+                with fast.optimized_context(front): return decode()
+        old=call(False)
+        new=call(True)
+        assert (old.accepted,old.reason,old.refined_start_sample)==(new.accepted,new.reason,new.refined_start_sample)
+        result[kind]={"reason":new.reason,"baseline_arrays":array_observation(lambda:call(False)),
+                      "optimized_arrays":array_observation(lambda:call(True))}
+    return result
+
+
+def run_benchmark(sequence) -> dict:
+    original=video.StreamingReceiver
+    try:
+        video.StreamingReceiver=fast.StreamingReceiver
+        repeats=[video.benchmark_receiver(sequence[:3],500000) for unused in range(3)]
+        profiler=StageProfiler()
+        with profiler.measure(optimized=True): profiled=video.benchmark_receiver(sequence[:3],500000)
+        metadata={}
+        for kind,(raw,count,candidate) in profiler.examples.items():
+            np.save(OUTPUT/f"optimized_{kind}_window.npy",raw)
+            metadata[kind]={"sample_count":count,"offset":candidate.sample_offset,"score":candidate.score}
+        write_result("optimized_examples.json",metadata)
+        return {"repeats":repeats,"median_processing_seconds":statistics.median(row["processing_seconds"] for row in repeats),
+                "profile":profiler.summary(),"profiled_benchmark":profiled,"rf_tx":0}
+    finally:
+        video.StreamingReceiver=original
+
+
+def run_functional_tests() -> None:
+    assert video.maximum_jpeg_bytes(500000,5)==4148
+    assert fast.lab045.conv.CONSTRAINT_LENGTH==7
+
+
+def main() -> None:
+    parser=argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--stage",choices=("benchmark","stream","memory","replay","all"),default="all")
+    args=parser.parse_args()
+    OUTPUT.mkdir(parents=True,exist_ok=True)
+    with Image.open(ROOT/"data/raw/lab011_scene.jpg") as source:
+        sequence=video.encode_sequence(source,video.PROFILES[2],30)
+    for name,function in (("benchmark",lambda:run_benchmark(sequence)),("stream",lambda:run_stream(sequence)),
+                          ("memory",run_memory_test),("replay",run_candidate_replay)):
+        if args.stage not in ("all",name): continue
+        print(f"Lab051 {name}: START; RF TX=0",flush=True)
+        result=function()
+        write_result(f"{name}.json",result)
+        print(json.dumps({key:value for key,value in result.items() if key not in (
+            "blocks","images","repeats","profiled_benchmark","rss_samples")},ensure_ascii=False),flush=True)
+
+
+if __name__=="__main__":
+    main()