State

Most hard-to-follow code is hard because of state: who can change this value, and what else has to change with it? Make that answer short.

Parallel containers

A classic smell: several dicts keyed by the same id, which every code path must update together.

prompts: dict[str, list[int]] = {}
generated: dict[str, list[int]] = {}
budgets: dict[str, int] = {}
arrived_at: dict[str, float] = {}

def cancel(request_id):
    prompts.pop(request_id, None)
    generated.pop(request_id, None)
    budgets.pop(request_id, None)
    # arrived_at forgotten here: a slow leak nobody notices

@dataclass
class Sequence:
    request_id: str
    prompt: list[int]
    max_new_tokens: int
    arrived_at: float
    generated: list[int] = field(default_factory=list)

running: dict[str, Sequence] = {}

def cancel(request_id: str) -> None:
    running.pop(request_id, None)

Adding a field to Sequence can't create a leak, because removal is one line that removes everything.

Status as an enum

class Worker:
    def __init__(self):
        self.loaded = False
        self.busy = False
        self.draining = False
        self.failed = False

    def can_accept(self):
        return self.loaded and not self.busy and not self.draining and not self.failed

class WorkerState(Enum):
    LOADING = auto()
    IDLE = auto()
    BUSY = auto()
    DRAINING = auto()
    FAILED = auto()

class Worker:
    def __init__(self):
        self.state = WorkerState.LOADING

    def can_accept(self) -> bool:
        return self.state is WorkerState.IDLE

If transitions matter, put them in one method (def transition(self, to: WorkerState)) that rejects illegal moves. Then the full state machine is in one place.

Guidelines

Owners are objects, not modules

Module-level mutable globals have every function in the module as a potential writer. Put the state on an instance and pass the instance in.

Prefer immutable messages

Anything that crosses a boundary (a request on a queue, a batch handed to the model, a result coming back) should be a @dataclass(frozen=True) or a plain value. Then the receiver can't accidentally change the sender's copy, and you never need a defensive deep copy.

Derive, don't store

If a value can be computed from other state cheaply (len(self.running), "is the batch full?"), compute it. Every stored copy is another thing to keep in sync.

Make invalid states unrepresentable

If a sequence only has a finish_reason once it's finished, consider a separate FinishedSequence type, or an Optional that's checked in exactly one place, rather than a field that's "usually set".