Lab 5 — Watermarks

"""GatewayMeter stage 5."""

"""GatewayMeter stage 5. See this lab's README and common/flink_job.py."""
from __future__ import annotations

import sys
sys.path.insert(0, "/labs")

from common.flink_job import build_stage, run_stage  # noqa: E402

def build_job(env):
    return build_stage(env, 5)

if __name__ == "__main__":
    run_stage(5)

import pathlib
import sys

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2]))

from common.gatewaymeter import operator_watermark
from common.produce import kafka_partition_for_key  # noqa: E402

def test_one_active_partition_holds_the_minimum():
    marks = {0: 100_000, 1: 99_000, 2: 20_000}
    assert operator_watermark(marks) == 20_000
    assert operator_watermark(marks, idle={2}) == 99_000


def test_stall_partition_omits_only_that_kafka_key_route():
    import subprocess
    root = pathlib.Path(__file__).resolve().parents[2]
    out = subprocess.run(
        [sys.executable, str(root / "common/produce.py"), "--rate", "100",
         "--seconds", "1", "--seed", "7", "--stall-partition", "5",
         "--stall-after-seconds", "0"],
        capture_output=True, text=True, timeout=10,
    )
    assert out.returncode == 0, out.stderr
    keys = [line.partition("\t")[0] for line in out.stdout.splitlines() if line.strip()]
    assert keys
    assert 5 not in {kafka_partition_for_key(key) for key in keys}
    assert len({kafka_partition_for_key(key) for key in keys}) >= 5

"""The event envelope every lab shares.

The gateway emits one JSON object per thing that happens to a request. `ts` is **event time** —
the millisecond the gateway says the thing happened — and it is the only clock the job may bill by.
Kafka's append time is not it, and neither is the moment an operator sees the record.
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any, Dict

# The six things the gateway reports. Frozen: a lab that needs a new fact adds a payload field,
# never a seventh type.
REQUEST_STARTED = "RequestStarted"          # payload: prompt_tokens
TOKEN_GENERATED = "TokenGenerated"          # payload: n
TOOL_CALLED = "ToolCalled"                  # payload: tool, latency_ms
AGENT_STEP_COMPLETED = "AgentStepCompleted"  # payload: step, tokens_in, tokens_out
MODEL_ERROR = "ModelError"                  # payload: code
REQUEST_COMPLETED = "RequestCompleted"      # payload: status, completion_tokens, latency_ms

EVENT_TYPES = (
    REQUEST_STARTED, TOKEN_GENERATED, TOOL_CALLED,
    AGENT_STEP_COMPLETED, MODEL_ERROR, REQUEST_COMPLETED,
)


@dataclass
class Event:
    event_type: str
    request_id: str
    tenant_id: str
    model: str
    ts: int                                   # event time, ms since epoch
    payload: Dict[str, Any] = field(default_factory=dict)

    def to_json(self) -> str:
        return json.dumps({
            "event_type": self.event_type, "request_id": self.request_id,
            "tenant_id": self.tenant_id, "model": self.model, "ts": self.ts,
            "payload": self.payload,
        })


def parse_event(raw: str) -> Event:
    """JSON line -> Event. A plain map; no state, no side effects, no clock.

    Malformed records raise. Module 2 is where the job decides what to do about that; until then
    a bad record should stop you, not vanish.
    """
    d = json.loads(raw)
    return Event(
        event_type=d["event_type"], request_id=d["request_id"], tenant_id=d["tenant_id"],
        model=d["model"], ts=int(d["ts"]), payload=d.get("payload") or {},
    )

"""Pure GatewayMeter business semantics shared by the labs and their fast tests.

Nothing in this module imports Flink.  A lab can therefore distinguish two questions:

* is the invoice calculation correct for this event trace? and
* does Flink preserve that calculation under partitioning, time, and failure?

The first question belongs in a fast unit test.  The second belongs in the compose cluster.
Money is represented as integer micro-dollars so replay and partial aggregation are deterministic.
"""
from __future__ import annotations

from dataclasses import dataclass, field, replace
from typing import Iterable, Mapping, Sequence
import zlib

from common.events import REQUEST_COMPLETED, REQUEST_STARTED, TOKEN_GENERATED, Event

MINUTE_MS = 60_000
PRICE_PER_1K = {
    "gpt-5.5": (12_000, 36_000),
    "gpt-5.4-mini": (1_000, 4_000),
    "claude-opus-5": (15_000, 75_000),
    "gemini-3.6-flash": (500, 2_000),
}


@dataclass
class OpenRequest:
    request_id: str
    tenant_id: str
    model: str
    started_ts: int
    prompt_tokens: int
    completion_tokens: int = 0


@dataclass(frozen=True)
class RequestRecord:
    request_id: str
    tenant_id: str
    model: str
    started_ts: int
    completed_ts: int
    latency_ms: int
    prompt_tokens: int
    completion_tokens: int
    status: str
    price_multiplier_ppm: int = 1_000_000
    price_source: str = "list"


@dataclass(frozen=True)
class TenantCostMinute:
    tenant_id: str
    window_end: int
    micro_usd: int
    requests: int

    def key(self) -> tuple[str, int]:
        return self.tenant_id, self.window_end

    def to_dict(self) -> dict[str, int | str]:
        return {
            "tenant_id": self.tenant_id,
            "window_end": self.window_end,
            "micro_usd": self.micro_usd,
            "requests": self.requests,
        }


@dataclass(frozen=True)
class PartialCost:
    tenant_id: str
    salt: int
    window_end: int
    micro_usd: int
    requests: int


@dataclass
class RequestAssemblerModel:
    """A small deterministic oracle for the keyed RequestAssembler operator."""

    open_requests: dict[str, OpenRequest] = field(default_factory=dict)

    def apply(self, event: Event) -> RequestRecord | None:
        if event.event_type == REQUEST_STARTED:
            self.open_requests[event.request_id] = OpenRequest(
                event.request_id,
                event.tenant_id,
                event.model,
                event.ts,
                int(event.payload.get("prompt_tokens", 0)),
            )
            return None
        current = self.open_requests.get(event.request_id)
        if current is None:
            return None
        if event.event_type == TOKEN_GENERATED:
            current.completion_tokens += int(event.payload.get("n", 0))
            return None
        if event.event_type != REQUEST_COMPLETED:
            return None
        self.open_requests.pop(event.request_id, None)
        return RequestRecord(
            request_id=current.request_id,
            tenant_id=current.tenant_id,
            model=current.model,
            started_ts=current.started_ts,
            completed_ts=event.ts,
            latency_ms=int(event.payload.get("latency_ms", event.ts - current.started_ts)),
            prompt_tokens=current.prompt_tokens,
            completion_tokens=int(event.payload.get("completion_tokens", current.completion_tokens)),
            status=str(event.payload.get("status", "ok")),
        )


# From stage 6 the assembler sets an event-time timer at start + 5 minutes; a request still open
# when it fires is billed as abandoned, for the tokens seen so far.
ABANDON_AFTER_MS = 5 * MINUTE_MS


def abandoned_record(request: OpenRequest, fired_at: int) -> RequestRecord:
    """The record the assembler's timer emits for a request that never completed."""
    return RequestRecord(
        request_id=request.request_id,
        tenant_id=request.tenant_id,
        model=request.model,
        started_ts=request.started_ts,
        completed_ts=fired_at,
        latency_ms=fired_at - request.started_ts,
        prompt_tokens=request.prompt_tokens,
        completion_tokens=request.completion_tokens,
        status="abandoned",
    )


def assemble(events: Iterable[Event]) -> tuple[list[RequestRecord], dict[str, OpenRequest]]:
    model = RequestAssemblerModel()
    records: list[RequestRecord] = []
    for event in events:
        record = model.apply(event)
        if record is not None:
            records.append(record)
    return records, model.open_requests


def request_cost_micro_usd(
    record: RequestRecord,
    prices: Mapping[str, tuple[int, int]] = PRICE_PER_1K,
) -> int:
    """Return deterministic micro-dollars, rounded once at the request boundary."""
    input_per_1k, output_per_1k = prices[record.model]
    list_micro = (
        record.prompt_tokens * input_per_1k
        + record.completion_tokens * output_per_1k
        + 500
    ) // 1000
    return (list_micro * record.price_multiplier_ppm + 500_000) // 1_000_000


def window_end(timestamp_ms: int, size_ms: int = MINUTE_MS) -> int:
    return (timestamp_ms // size_ms + 1) * size_ms


def aggregate_costs(
    records: Iterable[RequestRecord],
    *,
    timestamps: Mapping[str, int] | None = None,
) -> list[TenantCostMinute]:
    """Aggregate by tenant-minute; timestamps can substitute processing/arrival time in Lab 4."""
    totals: dict[tuple[str, int], tuple[int, int]] = {}
    for record in records:
        ts = timestamps[record.request_id] if timestamps is not None else record.completed_ts
        key = record.tenant_id, window_end(ts)
        micro, count = totals.get(key, (0, 0))
        totals[key] = micro + request_cost_micro_usd(record), count + 1
    return [TenantCostMinute(t, end, micro, count)
            for (t, end), (micro, count) in sorted(totals.items())]


def salt_for(request_id: str, buckets: int = 8) -> int:
    return zlib.crc32(request_id.encode("utf-8")) % buckets


def partial_costs(records: Iterable[RequestRecord], buckets: int = 8) -> list[PartialCost]:
    totals: dict[tuple[str, int, int], tuple[int, int]] = {}
    for record in records:
        key = record.tenant_id, salt_for(record.request_id, buckets), window_end(record.completed_ts)
        micro, count = totals.get(key, (0, 0))
        totals[key] = micro + request_cost_micro_usd(record), count + 1
    return [PartialCost(t, salt, end, micro, count)
            for (t, salt, end), (micro, count) in sorted(totals.items())]


def merge_partials(updates: Sequence[PartialCost]) -> list[TenantCostMinute]:
    """Replace a salt's prior value; a late window re-fire is a correction, not another add."""
    latest: dict[tuple[str, int, int], PartialCost] = {}
    for update in updates:
        latest[update.tenant_id, update.window_end, update.salt] = update
    totals: dict[tuple[str, int], tuple[int, int]] = {}
    for update in latest.values():
        key = update.tenant_id, update.window_end
        micro, count = totals.get(key, (0, 0))
        totals[key] = micro + update.micro_usd, count + update.requests
    return [TenantCostMinute(t, end, micro, count)
            for (t, end), (micro, count) in sorted(totals.items())]


def upsert_minutes(emissions: Iterable[TenantCostMinute]) -> dict[tuple[str, int], TenantCostMinute]:
    """The billing-store view: recovery duplicates and late corrections share one natural key."""
    result: dict[tuple[str, int], TenantCostMinute] = {}
    for emission in emissions:
        result[emission.key()] = emission
    return result


def with_negotiated_price(record: RequestRecord, multiplier_ppm: int) -> RequestRecord:
    return replace(record, price_multiplier_ppm=multiplier_ppm, price_source="negotiated")


# The billing service's price book. The job's price lookup (stage 9 on) and the invoice oracle both
# read it, so they cannot disagree about what a tenant pays.
NEGOTIATED_PRICE_PPM = {"acme": 900_000}


def negotiated_price(record: RequestRecord) -> RequestRecord:
    """The price lookup's answer: the tenant's negotiated multiplier, or the list price."""
    multiplier = NEGOTIATED_PRICE_PPM.get(record.tenant_id)
    return record if multiplier is None else with_negotiated_price(record, multiplier)


def operator_watermark(partition_watermarks: Mapping[int, int], idle: Iterable[int] = ()) -> int:
    """The minimum active input watermark; useful as the deterministic oracle in Lab 5."""
    idle_set = set(idle)
    active = [value for partition, value in partition_watermarks.items() if partition not in idle_set]
    if not active:
        raise ValueError("an operator watermark needs at least one active input")
    return min(active)


def replay_from_checkpoint(
    records: Sequence[RequestRecord], checkpoint_index: int
) -> list[TenantCostMinute]:
    """Model restore-and-replay: the final materialized invoice must equal a failure-free run."""
    before = aggregate_costs(records[:checkpoint_index])
    replayed = aggregate_costs(records[checkpoint_index:])
    # In a real window operator the checkpoint also contains open accumulators.  Re-aggregating the
    # full logical input is the business oracle against which the cluster output is compared.
    del before, replayed
    return aggregate_costs(records)

"""The one PyFlink GatewayMeter job, parameterised by the course stage.

Each lab's ``gatewaymeter/job.py`` calls :func:`run_stage` with its module number.  Keeping the
operators here makes changes cumulative and keeps frozen UIDs spelled once.  The lab README tells
the learner which section to inspect and which environment switch deliberately breaks it.
"""
from __future__ import annotations

import asyncio
import json
import os

from common.events import REQUEST_COMPLETED, REQUEST_STARTED, TOKEN_GENERATED, parse_event
from common.gatewaymeter import (
    ABANDON_AFTER_MS,
    MINUTE_MS,
    OpenRequest,
    PartialCost,
    RequestRecord,
    TenantCostMinute,
    abandoned_record,
    negotiated_price,
    request_cost_micro_usd,
    salt_for,
)

BOOTSTRAP = "kafka:19092"
INPUT_TOPIC = "gateway-events"
INPUT_PARTITIONS = 8                 # gateway-events' partitions; from stage 4, one source subtask each
LATENESS_MS = 2 * MINUTE_MS


def _types():
    from pyflink.common import Types
    return Types.PICKLED_BYTE_ARRAY()


def _source(env, stage: int):
    from pyflink.common import Duration, WatermarkStrategy
    from pyflink.common.serialization import SimpleStringSchema
    from pyflink.datastream.connectors.kafka import KafkaOffsetsInitializer, KafkaSource
    from pyflink.common.watermark_strategy import TimestampAssigner

    source = (
        KafkaSource.builder()
        .set_bootstrap_servers(BOOTSTRAP)
        .set_topics(INPUT_TOPIC)
        .set_group_id(f"gatewaymeter-m{stage:02d}")
        .set_starting_offsets(KafkaOffsetsInitializer.earliest())
        .set_value_only_deserializer(SimpleStringSchema())
        .build()
    )
    if stage < 4:
        raw = env.from_source(source, WatermarkStrategy.no_watermarks(), "gateway-events")
        return raw.uid("gateway-events").map(parse_event, output_type=_types()).name("parse").uid("parse")

    class EventTime(TimestampAssigner):
        def extract_timestamp(self, value, record_timestamp):
            return value.ts

    # Event time is Event.ts. PyFlink's from_source hands the source only the Java half of a
    # WatermarkStrategy and drops a Python TimestampAssigner, which would leave event time at the
    # moment the producer sent each record. So the source carries no watermarks, and
    # assign_timestamps_and_watermarks, which does run a Python assigner, generates them right after
    # it. Watermarks are then per source subtask; with one subtask per partition that is per
    # partition, which is what a silent partition (module 5) needs. It takes all eight slots.
    strategy = WatermarkStrategy.for_bounded_out_of_orderness(Duration.of_seconds(20))
    if stage >= 5 and os.getenv("GM_DISABLE_IDLENESS") != "1":
        strategy = strategy.with_idleness(Duration.of_seconds(15))
    strategy = strategy.with_timestamp_assigner(EventTime())     # last: with_idleness drops it

    parallelism = env.get_parallelism()
    env.set_parallelism(int(os.getenv("GM_SOURCE_PARALLELISM", str(INPUT_PARTITIONS))))
    raw = env.from_source(source, WatermarkStrategy.no_watermarks(), "gateway-events")
    events = raw.uid("gateway-events").map(parse_event, output_type=_types()).name("parse").uid("parse")
    timed = events.assign_timestamps_and_watermarks(strategy).name("event-time").uid("event-time")
    env.set_parallelism(parallelism)
    return timed


def _records(events, stage: int):
    from pyflink.common import Time, Types
    from pyflink.datastream.functions import KeyedProcessFunction
    from pyflink.datastream.state import StateTtlConfig, ValueStateDescriptor

    class RequestAssembler(KeyedProcessFunction):
        def open(self, runtime_context):
            descriptor = ValueStateDescriptor("open", _types())
            if stage >= 3:
                ttl = (
                    StateTtlConfig.new_builder(Time.minutes(10))
                    .set_update_type(StateTtlConfig.UpdateType.OnCreateAndWrite)
                    .set_state_visibility(StateTtlConfig.StateVisibility.NeverReturnExpired)
                    .build()
                )
                descriptor.enable_time_to_live(ttl)
            self.open_request = runtime_context.get_state(descriptor)

        def process_element(self, event, ctx):
            current = self.open_request.value()
            if event.event_type == REQUEST_STARTED:
                current = OpenRequest(
                    event.request_id,
                    event.tenant_id,
                    event.model,
                    event.ts,
                    int(event.payload.get("prompt_tokens", 0)),
                )
                self.open_request.update(current)
                if stage >= 6:
                    ctx.timer_service().register_event_time_timer(event.ts + ABANDON_AFTER_MS)
                return
            if current is None:
                return
            if event.event_type == TOKEN_GENERATED:
                current.completion_tokens += int(event.payload.get("n", 0))
                self.open_request.update(current)
                return
            if event.event_type != REQUEST_COMPLETED:
                return
            self.open_request.clear()
            yield RequestRecord(
                current.request_id,
                current.tenant_id,
                current.model,
                current.started_ts,
                event.ts,
                int(event.payload.get("latency_ms", event.ts - current.started_ts)),
                current.prompt_tokens,
                int(event.payload.get("completion_tokens", current.completion_tokens)),
                str(event.payload.get("status", "ok")),
            )

        def on_timer(self, timestamp, ctx):
            current = self.open_request.value()
            if current is None:
                return
            self.open_request.clear()
            yield abandoned_record(current, timestamp)

    keyed = events.key_by(
        (lambda event: event.tenant_id)
        if os.getenv("GM_BREAK_KEY_BY") == "1"
        else (lambda event: event.request_id),
        key_type=Types.STRING(),
    )
    return keyed.process(RequestAssembler(), output_type=_types()).name("assemble").uid("assemble")


def _enrich(records, stage: int):
    if stage < 9:
        return records
    from pyflink.common import Time
    from pyflink.datastream.async_data_stream import AsyncDataStream
    from pyflink.datastream.functions import AsyncFunction

    class PriceLookup(AsyncFunction):
        def __init__(self, delay_seconds):
            self.delay_seconds = delay_seconds

        async def async_invoke(self, record):
            await asyncio.sleep(self.delay_seconds)
            # The same price book the invoice oracle uses: acme's negotiated 10% discount.
            return [negotiated_price(record)]

    return (
        AsyncDataStream.unordered_wait(
            records,
            PriceLookup(float(os.getenv("GM_PRICE_DELAY_SECONDS", "0.02"))),
            Time.seconds(5),
            capacity=int(os.getenv("GM_ASYNC_CAPACITY", "100")),
            output_type=_types(),
        )
        .name("enrich-price")
        .uid("enrich-price")
    )


def _cost_windows(records, stage: int):
    from pyflink.common import Time, Types
    from pyflink.datastream.functions import AggregateFunction, ProcessWindowFunction
    from pyflink.datastream.output_tag import OutputTag
    from pyflink.datastream.window import TumblingEventTimeWindows, TumblingProcessingTimeWindows

    class CostAggregate(AggregateFunction):
        def create_accumulator(self):
            return 0, 0

        def add(self, record, accumulator):
            return accumulator[0] + request_cost_micro_usd(record), accumulator[1] + 1

        def get_result(self, accumulator):
            return accumulator

        def merge(self, left, right):
            return left[0] + right[0], left[1] + right[1]

    class EmitCost(ProcessWindowFunction):
        def process(self, key, context, elements):
            micro, requests = next(iter(elements))
            yield TenantCostMinute(key, context.window().end, micro, requests)

    class EmitPartial(ProcessWindowFunction):
        def process(self, key, context, elements):
            tenant, salt = key
            micro, requests = next(iter(elements))
            yield PartialCost(tenant, salt, context.window().end, micro, requests)

    if stage >= 11:
        keyed = records.key_by(
            lambda record: (record.tenant_id, salt_for(record.request_id)),
            key_type=Types.TUPLE([Types.STRING(), Types.INT()]),
        )
        windowed = keyed.window(TumblingEventTimeWindows.of(Time.minutes(1)))
        if stage >= 6:
            windowed = windowed.allowed_lateness(LATENESS_MS)
        partials = (
            windowed.aggregate(CostAggregate(), EmitPartial(),
                               accumulator_type=Types.TUPLE([Types.LONG(), Types.LONG()]),
                               output_type=_types())
            .name("tenant-cost-partial")
            .uid("tenant-cost-partial")
        )
        return _merge_partial_stream(partials), None

    keyed = records.key_by(lambda record: record.tenant_id, key_type=Types.STRING())
    assigner = (
        TumblingProcessingTimeWindows.of(Time.minutes(1))
        if stage == 4 and os.getenv("GM_CLOCK", "event") == "processing"
        else TumblingEventTimeWindows.of(Time.minutes(1))
    )
    windowed = keyed.window(assigner)
    late_tag = OutputTag("tenant-cost-late", _types())
    if stage >= 6:
        windowed = windowed.allowed_lateness(LATENESS_MS).side_output_late_data(late_tag)
    result = (
        windowed.aggregate(CostAggregate(), EmitCost(),
                           accumulator_type=Types.TUPLE([Types.LONG(), Types.LONG()]),
                           output_type=_types())
        .name("tenant-cost")
        .uid("tenant-cost")
    )
    return result, result.get_side_output(late_tag) if stage >= 6 else None


def _merge_partial_stream(partials):
    from pyflink.common import Types
    from pyflink.datastream.functions import KeyedProcessFunction
    from pyflink.datastream.state import MapStateDescriptor

    class MergePartials(KeyedProcessFunction):
        def open(self, runtime_context):
            self.parts = runtime_context.get_map_state(
                MapStateDescriptor("partials", Types.STRING(), _types())
            )

        def _total(self, tenant_id, end):
            values = [value for value in self.parts.values() if value.window_end == end]
            return TenantCostMinute(
                tenant_id,
                end,
                sum(value.micro_usd for value in values),
                sum(value.requests for value in values),
            )

        def process_element(self, partial, ctx):
            self.parts.put(f"{partial.window_end}:{partial.salt}", partial)
            fire_at = partial.window_end - 1
            cleanup_at = partial.window_end + LATENESS_MS
            if ctx.timer_service().current_watermark() >= fire_at:
                yield self._total(partial.tenant_id, partial.window_end)
            else:
                ctx.timer_service().register_event_time_timer(fire_at)
            ctx.timer_service().register_event_time_timer(cleanup_at)

        def on_timer(self, timestamp, ctx):
            if timestamp % MINUTE_MS == MINUTE_MS - 1:
                yield self._total(ctx.get_current_key(), timestamp + 1)
                return
            end = timestamp - LATENESS_MS
            for key in list(self.parts.keys()):
                if key.startswith(f"{end}:"):
                    self.parts.remove(key)

    return (
        partials.key_by(lambda partial: partial.tenant_id, key_type=Types.STRING())
        .process(MergePartials(), output_type=_types())
        .name("tenant-cost")
        .uid("tenant-cost")
    )


def _configure_recovery(env, stage: int):
    if stage < 7:
        return
    from pyflink.common import Configuration
    from pyflink.datastream import CheckpointingMode, ExternalizedCheckpointRetention

    env.enable_checkpointing(10_000, CheckpointingMode.EXACTLY_ONCE)
    checkpoint = env.get_checkpoint_config()
    checkpoint.set_checkpoint_timeout(60_000)
    checkpoint.set_min_pause_between_checkpoints(2_000)
    checkpoint.set_max_concurrent_checkpoints(1)
    checkpoint.set_externalized_checkpoint_retention(
        ExternalizedCheckpointRetention.RETAIN_ON_CANCELLATION
    )
    if stage >= 9 and os.getenv("GM_UNALIGNED") == "1":
        checkpoint.enable_unaligned_checkpoints()
    config = Configuration()
    config.set_string("state.backend.type", os.getenv("GM_STATE_BACKEND", "rocksdb"))
    config.set_string("execution.checkpointing.storage", "filesystem")
    config.set_string("execution.checkpointing.dir", "file:///checkpoints")
    config.set_string("restart-strategy.type", "fixed-delay")
    config.set_string("restart-strategy.fixed-delay.attempts", os.getenv("GM_RESTART_ATTEMPTS", "3"))
    config.set_string("restart-strategy.fixed-delay.delay", "10 s")
    env.configure(config)


def _sink(costs, stage: int):
    from pyflink.common import Types

    if stage < 8:
        costs.map(lambda value: json.dumps(value.to_dict(), sort_keys=True),
                  output_type=Types.STRING()).print()
        return
    from pyflink.common.serialization import SimpleStringSchema
    from pyflink.datastream.connectors.base import DeliveryGuarantee
    from pyflink.datastream.connectors.kafka import KafkaRecordSerializationSchema, KafkaSink

    guarantee_name = os.getenv("GM_SINK_GUARANTEE", "exactly_once")
    guarantee = (
        DeliveryGuarantee.AT_LEAST_ONCE
        if guarantee_name == "at_least_once"
        else DeliveryGuarantee.EXACTLY_ONCE
    )
    serializer = (
        KafkaRecordSerializationSchema.builder()
        .set_topic("tenant-cost")
        .set_value_serialization_schema(SimpleStringSchema())
        .build()
    )
    builder = (
        KafkaSink.builder()
        .set_bootstrap_servers(BOOTSTRAP)
        .set_record_serializer(serializer)
        .set_delivery_guarantee(guarantee)
        .set_property("transaction.timeout.ms", "900000")
    )
    if guarantee == DeliveryGuarantee.EXACTLY_ONCE:
        builder = builder.set_transactional_id_prefix(
            os.getenv("GM_TRANSACTIONAL_ID_PREFIX", f"gatewaymeter-m{stage:02d}-tenant-cost")
        )
    strings = costs.map(lambda value: json.dumps(value.to_dict(), sort_keys=True),
                        output_type=Types.STRING())
    strings.sink_to(builder.build()).name("sink-tenant-cost").uid("sink-tenant-cost")


def build_stage(env, stage: int):
    from pyflink.common import Types

    _configure_recovery(env, stage)
    env.set_parallelism(int(os.getenv("GM_PARALLELISM", "4")))
    if stage >= 10:
        env.set_max_parallelism(int(os.getenv("GM_MAX_PARALLELISM", "128")))
    env.add_python_file("/labs/common")
    events = _source(env, stage)
    records = _records(events, stage)
    if stage <= 3:
        records.print()
        return records
    records = _enrich(records, stage)
    costs, late = _cost_windows(records, stage)
    if late is not None:
        late.map(lambda value: f"LATE {value.request_id}",
                 output_type=Types.STRING()).print()
    _sink(costs, stage)
    return costs


def run_stage(stage: int) -> None:
    from pyflink.datastream import StreamExecutionEnvironment

    env = StreamExecutionEnvironment.get_execution_environment()
    build_stage(env, stage)
    env.execute(f"GatewayMeter - module {stage}")