Typed task lock for mutually exclusive job types: Frontier Interview's solution

Question: Typed task lock for mutually exclusive job types

Companies: Microsoft AI, Microsoft, OpenAI

Difficulty: Medium

Problem statement

Some jobs on a shared accelerator pool must not run alongside jobs of a different kind. Each job runs on its own thread and has one of two types: "gen" or "score". Any number of jobs of the same type may run at once. Jobs of different types must never overlap.

Implement a thread-safe lock class MutexTaskLock:

class MutexTaskLock:
    def acquire(self, task_type: str) -> None: ...  # blocks until the job may run
    def release(self, task_type: str) -> None: ...  # called when the job finishes

Every job calls acquire(t) before it starts and release(t) when it ends.

Requirements

Two jobs of different types are never between acquire and release at the same time.
Jobs of the same type are not serialised. They can share the lock.
A blocked thread must sleep, not spin (no busy-waiting or sleep polling).
Many threads call both methods concurrently.
Be ready to discuss what happens when one type arrives continuously. Can the other type starve?

Example 1

Threads A (gen) and B (gen) call acquire at t=0 and both return immediately. Thread C (score) calls acquire at t=1 and blocks. A releases at t=2 and C is still blocked. B releases at t=3, the lock is empty, and C's acquire returns.

Example 2

score job S holds the lock. A gen job G arrives and waits. While S is still running, another score job S2 arrives. With a fair lock, S2 waits behind G: S finishes, G runs, then S2 runs. A lock without fairness would admit S2 at once, and G could starve.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

This is a group lock: a generalisation of a readers-writer lock where there
are two "reader" groups and each group excludes the other. Any number of
threads of the running type may share the lock; a thread of the other type
must wait until the lock is empty.

Naive version. Keep active_type and active_count behind one mutex and a
condition variable:

acquire(t): wait while active_count > 0 and active_type != t, then set
active_type = t and increment the count.
release(t): decrement; when the count reaches 0, clear active_type and
notify_all().

This is correct (never mixes types, never busy-waits, since Condition.wait
parks the thread), but it starves: if "gen" jobs keep arriving while
others of the same type are still running, the count never reaches zero and
a waiting "score" job waits forever. Interviewers almost always ask about
this next, so it is worth building fairness in from the start.

Fair version (phase switching). Also track how many threads of each type
are waiting, and which type ran last.

If the lock is held by type t and nobody of the other type is waiting, a
new t thread joins immediately (maximum sharing).
If the lock is held by type t and someone of the other type is waiting,
new t arrivals queue up instead of joining. The current batch drains.
When the lock becomes empty, the turn goes to the type that did not run
last, if it has waiters. All its waiters enter together as one batch.

So under contention the lock alternates gen batch -> score batch -> gen batch ..., and every waiting thread gets in within at most one opposing
batch. Under no contention it behaves like the naive lock.

Invariants (all checked under the mutex):

active_count > 0 implies every running thread has type active_type.
active_count == 0 iff active_type is None.
waiting[t] equals the number of threads blocked in acquire(t).
Waiters are woken only when the lock becomes empty. That is the only event
that can change any waiter's answer, because the "other type is waiting"
count only falls when that type enters, which needs an empty lock.

Always wait in a while loop (spurious and stolen wakeups), and use
notify_all, because a single notify could wake a thread of the wrong type
and lose the wakeup.

Code

import threading
from contextlib import contextmanager
from typing import Iterator, Optional

TYPES = ("gen", "score")


class MutexTaskLock:
    """Same-type jobs share the lock; different types never overlap.
    Fair: under contention the two types alternate in batches."""

    def __init__(self) -> None:
        self._cond = threading.Condition(threading.Lock())
        self._active_type: Optional[str] = None
        self._active_count = 0
        self._waiting = {t: 0 for t in TYPES}
        self._last_type: Optional[str] = None  # type of the last finished batch

    @staticmethod
    def _other(task_type: str) -> str:
        return "score" if task_type == "gen" else "gen"

    def _can_enter(self, task_type: str) -> bool:
        other_waiting = self._waiting[self._other(task_type)] > 0
        if self._active_count > 0:
            # Join a running batch of our type, unless the other side is queued.
            return self._active_type == task_type and not other_waiting
        # Lock is empty: give the turn to whichever type did not run last.
        return not other_waiting or self._last_type != task_type

    def acquire(self, task_type: str) -> None:
        if task_type not in TYPES:
            raise ValueError(f"unknown task type: {task_type!r}")
        with self._cond:
            self._waiting[task_type] += 1
            try:
                while not self._can_enter(task_type):
                    self._cond.wait()  # parks the thread; no busy-wait
            finally:
                self._waiting[task_type] -= 1
            self._active_type = task_type
            self._active_count += 1

    def release(self, task_type: str) -> None:
        with self._cond:
            if self._active_count == 0 or self._active_type != task_type:
                raise RuntimeError(f"release({task_type!r}) without matching acquire")
            self._active_count -= 1
            if self._active_count == 0:
                self._last_type = task_type
                self._active_type = None
                self._cond.notify_all()  # waiters of either type re-check

    @contextmanager
    def hold(self, task_type: str) -> Iterator[None]:
        self.acquire(task_type)
        try:
            yield
        finally:
            self.release(task_type)

Usage: with lock.hold("gen"): run_generation(), which releases the lock even
if the job raises.

Complexity

acquire: O(1) work per check under the mutex. A blocked thread uses no CPU.
It may wake once per time the lock empties.
release: O(1). notify_all wakes the W waiting threads, so one release that
empties the lock costs O(W) wakeups in total. This is the price of a single
condition variable (see follow-ups).
Space: O(1) bookkeeping, plus the condition variable's internal queue of
waiting threads, which is O(W).

Test cases

One "gen" acquire/release on an idle lock: returns at once; afterwards
active_count == 0.
5 "gen" threads that each sleep 0.1 s inside the lock: they run at the same
time (peak concurrency 5, total about 0.1 s, not 0.5 s).
1 "gen" holder, then a "score" thread calls acquire: the score thread
blocks until the gen thread releases.
Mixed stress test: 20 random-type threads, each recording the type while
inside. A shared counter checked under a separate mutex never sees both
types inside at once.
Starvation: a stream of overlapping "gen" jobs keeps the lock busy, then one
"score" job arrives. New gen arrivals queue, the running gens drain, and the
score job runs next.
release("score") while "gen" holds the lock, or release on an idle lock:
raises RuntimeError.
acquire("train"): raises ValueError.
A job that raises inside with lock.hold("gen") still releases the lock.

Follow-ups

N types instead of two. Replace the two-type logic with a FIFO queue of
pending types (or round-robin over types that have waiters). When the lock
empties, the next type in that order gets the turn. With a
Condition per type, you can notify_all only the chosen type and avoid
waking everyone.
Bounded batch size or capacity. Example: at most 4 "gen" jobs on one GPU at
once. Add a per-type limit to _can_enter (active_count < cap[t]), and
on every release notify one waiter of the same type, not only when the
lock empties.
Timeouts / cancellation. acquire(t, timeout) uses cond.wait(remaining)
with a deadline and returns False on expiry. The finally already undoes
waiting[t]. Because a leaving waiter can unblock the other side, also
notify_all when waiting[t] drops to 0.
Strict FIFO fairness. Give each arrival a ticket number. A thread may enter only when
all earlier tickets have entered and the lock is empty or holds its own type. This
gives less sharing (a gen behind a score cannot join an earlier gen batch)
but strict arrival order.
Distributed version. Jobs run on many machines. Hold the state in a
coordination service (ZooKeeper, etcd, or a DB row with
compare-and-set), use leases with TTL so crashed holders free the lock,
and use fencing tokens so a stale holder cannot act after its lease ends.

Interviewer expectations

Mid-level (L4)

Gets the naive version right with a mutex and condition variable: a while
loop around wait, and notify_all on the last release.
Explains why polling or sleep loops are wrong.
Pairs acquire and release with try/finally or a context manager.

Senior (L5)

Spots starvation without being prompted and adds batch alternation.
Explains why notify instead of notify_all can lose a wakeup with two waiter
kinds.
Validates misuse (wrong-type release), and writes a concurrency test that
checks the invariant, not only the happy path.

Staff (L6+)

Discusses the throughput-versus-fairness trade-off (batching vs strict FIFO)
and what the workload needs.
Generalises to N types and per-type capacity, and avoids the cost of waking
every waiter with per-type conditions.
Covers timeouts, crashed holders and the distributed version with leases and
fencing.