Spotify Staff AI Engineer Interview Preparation Guide

from typing import List

def pack_requests(requests: List[int], max_batch_tokens: int, latency_budget_ms: int) -> List[List[int]]:
    """
    Greedy pack in input order without considering latency budget in timing.
    Returns batches as lists of indices into `requests`.
    """
    batches = []
    cur_batch = []
    cur_tokens = 0

    for i, tokens in enumerate(requests):
        if tokens > max_batch_tokens:
            # cannot fit even alone: treat as its own batch (or raise)
            if cur_batch:
                batches.append(cur_batch)
                cur_batch = []
                cur_tokens = 0
            batches.append([i])
            continue

        if cur_tokens + tokens <= max_batch_tokens:
            cur_batch.append(i)
            cur_tokens += tokens
        else:
            batches.append(cur_batch)
            cur_batch = [i]
            cur_tokens = tokens

    if cur_batch:
        batches.append(cur_batch)

    return batches

import heapq
from typing import List, Optional, Tuple

def pack_with_latency(requests: List[int],
                      max_batch_tokens: int,
                      latency_budget_ms: int,
                      arrival_times: Optional[List[int]] = None,
                      base_ms: float = 10.0,
                      ms_per_token: float = 0.05) -> List[Tuple[List[int], float]]:
    """
    Packs requests respecting latency budgets when possible.
    Returns list of (batch_indices, completion_time).
    arrival_times defaults to all zeros.
    """
    n = len(requests)
    if arrival_times is None:
        arrival_times = [0]*n

    # pending heap by deadline: (deadline, idx)
    pending = []
    result = []
    i = 0
    current_time = 0.0

    while i < n or pending:
        # push arrivals up to current_time (or if no pending, fast-forward to next arrival)
        if not pending and i < n and arrival_times[i] > current_time:
            current_time = arrival_times[i]

        while i < n and arrival_times[i] <= current_time:
            deadline = arrival_times[i] + latency_budget_ms
            heapq.heappush(pending, (deadline, i))
            i += 1

        if not pending:
            continue

        # form batch greedily from earliest deadlines
        batch = []
        batch_tokens = 0
        # take a snapshot of pending as list sorted by deadline without popping permanently
        temp = []
        while pending and len(temp) < len(pending):
            # pop earliest deadline candidate
            deadline, idx = heapq.heappop(pending)
            temp.append((deadline, idx))

        # iterate through EDF list, try to add if fits tokens and deadlines
        for deadline, idx in temp:
            tokens = requests[idx]
            if tokens > max_batch_tokens:
                # cannot batch normally; decide policy: send alone
                if batch:
                    # push back unselected ones
                    for d, j in temp:
                        if j not in batch:
                            heapq.heappush(pending, (d,j))
                    break
                # singleton dispatch
                est = current_time + base_ms + ms_per_token * tokens
                result.append(([idx], est))
                current_time = est
                # push back remaining arrivals (they are still pending)
                for d, j in temp:
                    if j != idx:
                        heapq.heappush(pending, (d,j))
                break

            # tentative completion if we include this request now
            tentative_tokens = batch_tokens + tokens
            est = current_time + base_ms + ms_per_token * tentative_tokens
            if est <= deadline and tentative_tokens <= max_batch_tokens:
                batch.append(idx)
                batch_tokens = tentative_tokens
            else:
                # cannot include this request; keep it pending
                heapq.heappush(pending, (deadline, idx))

        else:
            # loop finished normally: dispatch formed batch if non-empty
            if batch:
                completion = current_time + base_ms + ms_per_token * batch_tokens
                result.append((batch, completion))
                current_time = completion

            # any temp items not selected were pushed back already

    return result

import math
from typing import Iterable, Optional

def per_token_perplexity(log_probs: Iterable[float], mask: Optional[Iterable[int]] = None) -> float:
    """
    Compute per-token perplexity from an iterable of log-probabilities.
    log_probs: natural log probabilities (ln p) for each token.
    mask: optional iterable of 0/1 to include/exclude tokens (e.g., ignore padding).
    Returns: perplexity (float)
    """
    total = 0.0
    count = 0
    if mask is None:
        for lp in log_probs:
            total += lp
            count += 1
    else:
        for lp, m in zip(log_probs, mask):
            if m:
                total += lp
                count += 1
    if count == 0:
        raise ValueError("No tokens to compute perplexity")
    avg_log_prob = total / count
    # perplexity = exp(-avg_log_prob) assuming natural logs
    return math.exp(-avg_log_prob)

import torch
import torch.nn.functional as F
from torch.cuda.amp import autocast, GradScaler

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
optimizer = ...            # assumed provided
dataloader = ...          # assumed provided
epochs = 3
grad_accum_steps = 4
pad_token_id = -100       # or model.config.pad_token_id
scaler = GradScaler()

model.train()
for epoch in range(epochs):
    optimizer.zero_grad()
    for step, batch in enumerate(dataloader):
        input_ids = batch['input_ids'].to(device)            # (B, L)
        attention_mask = batch['attention_mask'].to(device)  # (B, L)

        with autocast():  # mixed precision context
            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            logits = outputs.logits                          # (B, L, Vocab)
            # shift logits and labels for causal LM: predict token t from previous tokens
            shift_logits = logits[..., :-1, :].contiguous()  # (B, L-1, V)
            shift_labels = input_ids[..., 1:].contiguous()   # (B, L-1)
            # optionally shift attention_mask if needed
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=pad_token_id,
                reduction='mean'
            )
            loss = loss / grad_accum_steps  # scale for accumulation

        # scale + backward for mixed precision
        scaler.scale(loss).backward()

        # optimizer step when accumulated enough gradients
        if (step + 1) % grad_accum_steps == 0:
            scaler.step(optimizer)   # calls unscale internally
            scaler.update()
            optimizer.zero_grad()
            # optionally: scheduler.step()