Loading, Unloading, and Scheduling Models

/** Lesson 38's Backend, plus alive() so the manager can notice a crashed engine. */
public interface Backend extends AutoCloseable {
    String name();
    String baseUrl();                      // ends with /v1
    void ensureReady() throws Exception;   // start if needed, then wait until healthy
    boolean alive();
    @Override void close();                // stop only what this backend started
}

@Override public void ensureReady() throws Exception {
    ProcessBuilder pb = new ProcessBuilder(command).redirectErrorStream(true)
            .redirectOutput(ProcessBuilder.Redirect.DISCARD);
    pb.environment().remove("JAVA_TOOL_OPTIONS");
    process = pb.start();
    long deadline = System.nanoTime() + Duration.ofSeconds(30).toNanos();
    long sleep = 25;
    while (System.nanoTime() < deadline) {
        if (!process.isAlive()) throw new IllegalStateException("engine exited with code " + process.exitValue());
        if (probe() == 200) return;
        Thread.sleep(sleep);
        sleep = Math.min(sleep * 2, 500);
    }
    throw new IllegalStateException("engine not ready after 30s");
}

@Override public void close() {
    if (process == null) return;
    process.descendants().forEach(ProcessHandle::destroy);
    process.destroy();
    try {
        if (!process.waitFor(stopTimeout.toMillis(), java.util.concurrent.TimeUnit.MILLISECONDS)) {
            process.descendants().forEach(ProcessHandle::destroyForcibly);
            process.destroyForcibly();
            process.waitFor();
        }
    } catch (InterruptedException e) {
        process.destroyForcibly();
        Thread.currentThread().interrupt();
    }
}

import demo.GatewayProperties.ModelSpec;
import java.io.IOException;
import java.net.ServerSocket;
import java.util.ArrayList;
import java.util.List;
import org.springframework.stereotype.Component;

@Component
class ProcessBackendFactory implements BackendFactory {
    private final GatewayProperties props;

    ProcessBackendFactory(GatewayProperties props) { this.props = props; }

    @Override public Backend create(String model, ModelSpec spec) {
        int port = freePort();
        List<String> command = new ArrayList<>();
        for (String part : props.engineCommand()) {
            command.add(part.replace("{port}", Integer.toString(port)).replace("{model}", model));
        }
        return new ProcessBackend(model, command, port, props.stopTimeout());
    }

    private static int freePort() {
        try (ServerSocket s = new ServerSocket(0)) {
            return s.getLocalPort();
        } catch (IOException e) {
            throw new IllegalStateException("no free port", e);
        }
    }
}

import com.sun.net.httpserver.HttpServer;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

/** A fake OpenAI-compatible engine: args are port, model name, load milliseconds, generation milliseconds. */
public final class FakeEngine {
    public static void main(String[] args) throws Exception {
        int port = Integer.parseInt(args[0]);
        String model = args[1];
        long loadMs = Long.parseLong(args[2]), genMs = Long.parseLong(args[3]);
        long readyAt = System.currentTimeMillis() + loadMs;
        AtomicInteger inFlight = new AtomicInteger();
        HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", port), 0);
        server.setExecutor(Executors.newVirtualThreadPerTaskExecutor());
        server.createContext("/health", ex -> {
            boolean ready = System.currentTimeMillis() >= readyAt;
            respond(ex, ready ? 200 : 503, ready ? "{\"status\":\"ok\"}" : "{\"error\":{\"message\":\"Loading model\"}}");
        });
        server.createContext("/v1/chat/completions", ex -> {
            ex.getRequestBody().readAllBytes();
            int concurrent = inFlight.incrementAndGet();
            try {
                Thread.sleep(genMs);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            } finally {
                inFlight.decrementAndGet();
            }
            respond(ex, 200, "{\"model\":\"" + model + "\",\"choices\":[{\"index\":0,\"message\":{\"role\":\"assistant\","
                    + "\"content\":\"reply from " + model + "\"},\"finish_reason\":\"stop\"}],\"in_flight_at_start\":" + concurrent + "}");
        });
        server.start();
    }

    private static void respond(com.sun.net.httpserver.HttpExchange ex, int status, String body) throws java.io.IOException {
        byte[] bytes = body.getBytes(StandardCharsets.UTF_8);
        ex.getResponseHeaders().add("Content-Type", "application/json");
        ex.sendResponseHeaders(status, bytes.length);
        ex.getResponseBody().write(bytes);
        ex.close();
    }
}

import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;

import org.junit.jupiter.api.Test;

class BackendFactoryTest {
    @Test
    void factoryBuildsAndStopsAnEngine() throws Exception {
        GatewayProperties props = Lab.props(2, 1);
        BackendFactory factory = new ProcessBackendFactory(props);
        try (Backend backend = factory.create("coder:7b", props.models().get("coder:7b"))) {
            System.out.println("alive before start: " + backend.alive());
            backend.ensureReady();
            System.out.println("alive after ensureReady: " + backend.alive());
            System.out.println("base url shape: " + backend.baseUrl().replaceAll(":\\d+/", ":<port>/"));
            assertTrue(backend.alive());
        }
        long left = ProcessHandle.allProcesses().filter(p -> p.info().command().map(c -> c.endsWith("/java")).orElse(false))
                .filter(p -> p.info().commandLine().map(c -> c.contains("demo.dev.FakeEngine")).orElse(false)).count();
        System.out.println("engine processes after close: " + left);
        assertFalse(left > 0);
    }
}

import jakarta.validation.Valid;
import jakarta.validation.constraints.Min;
import jakarta.validation.constraints.NotEmpty;
import jakarta.validation.constraints.NotNull;
import java.nio.file.Path;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.context.properties.bind.DefaultValue;
import org.springframework.util.unit.DataSize;
import org.springframework.validation.annotation.Validated;

@Validated
@ConfigurationProperties("gateway")
public record GatewayProperties(
        @DefaultValue("5m") Duration keepAlive,
        @DefaultValue("3") @Min(1) int maxLoadedModels,
        @DefaultValue("1") @Min(1) int parallelPerModel,
        @DefaultValue("8") @Min(0) int maxQueue,
        @DefaultValue("30s") Duration queueTimeout,
        @DefaultValue("5s") Duration stopTimeout,
        @DefaultValue("5s") Duration reapInterval,
        @NotNull DataSize memoryBudget,
        @NotEmpty List<String> engineCommand,
        @NotEmpty Map<String, @Valid ModelSpec> models,
        @NotNull String hubUrl,
        @NotNull Path modelDir) {

    public record ModelSpec(
            @NotNull DataSize weights,
            @Min(1) long kvBytesPerToken,
            @DefaultValue("4096") @Min(1) int contextTokens) { }
}

import java.util.concurrent.atomic.AtomicBoolean;

/** One request's hold on a loaded model. Closing it frees the slot and starts the keep_alive timer when it was the last. */
public final class Lease implements AutoCloseable {
    private final String model;
    private final String baseUrl;
    private final Runnable release;
    private final AtomicBoolean closed = new AtomicBoolean();

    Lease(String model, String baseUrl, Runnable release) {
        this.model = model; this.baseUrl = baseUrl; this.release = release;
    }

    public String model() { return model; }
    public String baseUrl() { return baseUrl; }

    @Override public void close() {
        if (closed.compareAndSet(false, true)) release.run();
    }
}

private static final class Slot {
    final String name;
    final long bytes;
    final Semaphore permits;
    State state = State.LOADING;
    Backend backend;
    int reserved;          // requests that hold a permit or wait for one
    int waiting;           // requests waiting for a permit
    Duration keepAlive;
    Instant lastUsed;
    Instant expiresAt;     // null while busy or when keep_alive is forever

    Slot(String name, long bytes, int permits, Duration keepAlive, Instant now) {
        this.name = name; this.bytes = bytes; this.permits = new Semaphore(permits, true);
        this.keepAlive = keepAlive; this.lastUsed = now;
    }
}

private Slot reserve(String model, long bytes, Duration keepAlive) throws InterruptedException {
    Slot fresh;
    lock.lock();
    try {
        while (true) {
            if (closing) throw new ShuttingDownException();
            Slot s = slots.get(model);
            if (s != null && s.state == State.READY && !s.backend.alive()) {
                stop(s, "engine exited");
                continue;
            }
            if (s != null && s.state == State.READY) {
                s.reserved++;
                s.keepAlive = keepAlive;
                s.expiresAt = null;
                return s;
            }
            if (s != null) { changed.await(); continue; }            // LOADING or STOPPING: wait for the change
            if (slots.size() >= props.maxLoadedModels() || used() + bytes > props.memoryBudget().toBytes()) {
                Slot victim = victim();
                if (victim == null) { changed.await(); continue; }   // every loaded model is busy or stopping
                stop(victim, "evicted to make room for " + model);
                continue;
            }
            fresh = new Slot(model, bytes, props.parallelPerModel(), keepAlive, clock.instant());
            slots.put(model, fresh);                                 // memory is reserved before the engine starts
            break;
        }
    } finally {
        lock.unlock();
    }
    Backend backend = factory.create(model, props.models().get(model));
    try {
        backend.ensureReady();
    } catch (Exception e) {
        backend.close();
        lock.lock();
        try { slots.remove(model); changed.signalAll(); } finally { lock.unlock(); }
        if (e instanceof InterruptedException ie) throw ie;
        throw new LoadFailedException(model, e);
    }
    lock.lock();
    try {
        fresh.backend = backend;
        fresh.state = State.READY;
        fresh.reserved++;
        changed.signalAll();
    } finally {
        lock.unlock();
    }
    events.publishEvent(new ModelEvent.Loaded(model, bytes));
    return fresh;
}

/** Returns a lease on a ready engine for the model, loading it and evicting others as needed. */
public Lease acquire(String model, Duration keepAlive) throws InterruptedException {
    ModelSpec spec = props.models().get(model);
    if (spec == null) throw new UnknownModelException(model);
    long bytes = MemoryPlan.estimate(spec, props.parallelPerModel());
    if (bytes > props.memoryBudget().toBytes()) throw new DoesNotFitException(model, bytes, props.memoryBudget().toBytes());

    Slot slot = reserve(model, bytes, keepAlive == null ? props.keepAlive() : keepAlive);
    try {
        awaitPermit(slot);
    } catch (RuntimeException | InterruptedException e) {
        finish(slot);
        throw e;
    }
    return new Lease(model, slot.backend.baseUrl(), () -> { slot.permits.release(); finish(slot); });
}

/** A manager over real fake-engine processes; every event is appended to log. */
static ModelManager manager(GatewayProperties props, Clock clock, List<String> log) {
    return new ModelManager(props, new ProcessBackendFactory(props), clock, event -> {
        switch ((ModelEvent) event) {
            case ModelEvent.Loaded l -> log.add("loaded   " + l.model());
            case ModelEvent.Unloaded u -> log.add("unloaded " + u.model() + " (" + u.reason() + ")");
        }
    });
}

static String view(ModelManager m) {
    List<String> parts = new ArrayList<>();
    for (var s : m.snapshot()) parts.add(s.model() + "=" + s.state().name().toLowerCase() + " in_flight=" + s.inFlight() + " queued=" + s.queued());
    return parts.isEmpty() ? "(nothing loaded)" : String.join(", ", parts);
}

import java.time.Duration;
import java.util.List;
import org.junit.jupiter.api.Test;

class FirstRequestTest {
    @Test
    void firstRequestLoadsTheModel() throws Exception {
        List<String> log = Lab.log();
        ModelManager manager = Lab.manager(Lab.props(2, 1), new Lab.MutableClock(), log);
        System.out.println("before: " + Lab.view(manager));
        try (Lease lease = manager.acquire("coder:7b", null)) {
            System.out.println("during: " + Lab.view(manager));
            System.out.println("lease model: " + lease.model());
        }
        System.out.println("after:  " + Lab.view(manager));
        try (Lease again = manager.acquire("coder:7b", null)) {
            System.out.println("second request reuses the engine");
        }
        log.forEach(System.out::println);
        manager.shutdown();
        log.stream().skip(1).forEach(System.out::println);
    }
}

import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import org.junit.jupiter.api.Test;

class ConcurrentFirstRequestsTest {
    @Test
    void eightFirstRequestsStartOneEngine() throws Exception {
        List<String> log = Lab.log();
        ModelManager manager = Lab.manager(Lab.props(2, 8), new Lab.MutableClock(), log);
        Set<String> urls = ConcurrentHashMap.newKeySet();
        List<Thread> threads = new ArrayList<>();
        for (int i = 0; i < 8; i++) {
            threads.add(Thread.ofVirtual().start(() -> {
                try (Lease lease = manager.acquire("coder:7b", null)) {
                    urls.add(lease.baseUrl());
                    Lab.sleep(50);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            }));
        }
        for (Thread t : threads) t.join();
        System.out.println("requests served: 8");
        System.out.println("distinct engine URLs: " + urls.size());
        System.out.println("engine starts: " + log.stream().filter(l -> l.startsWith("loaded")).count());
        manager.shutdown();
    }
}

import demo.GatewayExceptions.LoadFailedException;
import java.nio.file.Path;
import java.time.Duration;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.springframework.util.unit.DataSize;

class LoadFailureTest {
    @Test
    void failedLoadLeavesNoTrace() throws Exception {
        GatewayProperties ok = Lab.props(2, 1);
        GatewayProperties broken = new GatewayProperties(ok.keepAlive(), 2, 1, 8, Duration.ofSeconds(10), Duration.ofSeconds(5),
                Duration.ofSeconds(5), DataSize.ofGigabytes(64),
                List.of("java", "-cp", "target/classes", "demo.dev.NoSuchEngine", "{port}"), ok.models(), ok.hubUrl(), Path.of("target/models"));
        List<String> log = Lab.log();
        ModelManager manager = Lab.manager(broken, new Lab.MutableClock(), log);
        for (int attempt = 1; attempt <= 2; attempt++) {
            try {
                manager.acquire("coder:7b", null);
            } catch (LoadFailedException e) {
                System.out.println("attempt " + attempt + ": " + e.getMessage());
            }
            System.out.println("state: " + Lab.view(manager));
        }
        System.out.println("events: " + log);
    }
}

import java.time.Duration;

public class KeepAliveStage {
    /** Parses the keep_alive request field: a number of seconds, or a string such as "10m". Negative means forever. */
    static final class KeepAlive {
        public static final Duration FOREVER = Duration.ofSeconds(-1);

        private KeepAlive() { }

        public static Duration parse(Object value, Duration fallback) {
            if (value == null) return fallback;
            if (value instanceof Number n) {
                return n.doubleValue() < 0 ? FOREVER : Duration.ofMillis(Math.round(n.doubleValue() * 1000));
            }
            if (value instanceof String s) {
                Duration d = parseUnits(s.trim());
                return d.isNegative() ? FOREVER : d;
            }
            throw new IllegalArgumentException("keep_alive must be a number or a string, not " + value.getClass().getSimpleName());
        }

        private static Duration parseUnits(String s) {
            if (s.isEmpty()) throw new IllegalArgumentException("keep_alive is empty");
            if (s.equals("0")) return Duration.ZERO;           // Go accepts a bare 0 with no unit
            boolean negative = s.startsWith("-");
            String body = negative ? s.substring(1) : s;
            long millis = 0;
            int i = 0;
            while (i < body.length()) {
                int start = i;
                while (i < body.length() && (Character.isDigit(body.charAt(i)) || body.charAt(i) == '.')) i++;
                if (start == i) throw new IllegalArgumentException("invalid keep_alive: " + s);
                double amount = Double.parseDouble(body.substring(start, i));
                int unitStart = i;
                while (i < body.length() && Character.isLetter(body.charAt(i))) i++;
                double unit = switch (body.substring(unitStart, i)) {
                    case "ms" -> 1;
                    case "s" -> 1000;
                    case "m" -> 60_000;
                    case "h" -> 3_600_000;
                    default -> throw new IllegalArgumentException("invalid keep_alive unit in: " + s);
                };
                millis += Math.round(amount * unit);
            }
            return Duration.ofMillis(negative ? -millis : millis);
        }
    }

    public static void main(String[] args) {
        Duration fallback = Duration.ofMinutes(5);
        Object[] inputs = {null, 90, 0, -1, 0.5, "10m", "1h30m", "250ms", "0", "-1", "-1m", "5"};
        for (Object in : inputs) {
            String shown = in instanceof String ? "\"" + in + "\"" : String.valueOf(in);
            try {
                System.out.println(shown + " -> " + KeepAlive.parse(in, fallback));
            } catch (IllegalArgumentException e) {
                System.out.println(shown + " -> error: " + e.getMessage());
            }
        }
    }
}

/** Called when a request ends: the last one out starts the keep_alive timer, or unloads at once for keep_alive 0. */
private void finish(Slot s) {
    lock.lock();
    try {
        s.reserved--;
        if (s.reserved > 0) return;
        Instant now = clock.instant();
        s.lastUsed = now;
        if (s.keepAlive.isZero()) stop(s, "keep_alive 0");
        else s.expiresAt = s.keepAlive.isNegative() ? null : now.plus(s.keepAlive);
        changed.signalAll();
    } finally {
        lock.unlock();
    }
}

/** Unloads every idle model whose keep_alive has run out. */
@Scheduled(fixedDelayString = "${gateway.reap-interval:5s}")
public void reapIdle() {
    lock.lock();
    try {
        Instant now = clock.instant();
        for (Slot s : new ArrayList<>(slots.values())) {
            if (s.state == State.READY && s.reserved == 0 && s.expiresAt != null && !now.isBefore(s.expiresAt)) {
                stop(s, "keep_alive expired");
            }
        }
    } finally {
        lock.unlock();
    }
}

import java.time.Duration;
import java.util.List;
import org.junit.jupiter.api.Test;

class KeepAliveExpiryTest {
    @Test
    void idleModelsUnloadWhenKeepAliveEnds() throws Exception {
        List<String> log = Lab.log();
        Lab.MutableClock clock = new Lab.MutableClock();
        ModelManager manager = Lab.manager(Lab.props(3, 1), clock, log);

        manager.acquire("coder:7b", null).close();                        // default keep_alive: 5m
        manager.acquire("chat:8b", Duration.ofSeconds(30)).close();      // keep_alive 30s
        manager.acquire("tiny:1b", KeepAlive.FOREVER).close();           // keep_alive -1
        long[] stepsSeconds = {0, 29, 1, 270, 3_600};
        for (long step : stepsSeconds) {
            clock.advance(Duration.ofSeconds(step));
            manager.reapIdle();
            System.out.println(clock.time() + " loaded: " + manager.snapshot().stream().map(ModelManager.LoadedModel::model).toList());
        }
        log.stream().filter(l -> l.startsWith("unloaded")).forEach(System.out::println);

        manager.acquire("coder:7b", Duration.ZERO).close();               // keep_alive 0
        System.out.println("after keep_alive 0: " + manager.snapshot().stream().map(ModelManager.LoadedModel::model).toList());
        manager.shutdown();
    }
}

private void awaitPermit(Slot s) throws InterruptedException {
    if (s.permits.tryAcquire()) return;
    lock.lock();
    try {
        if (s.waiting >= props.maxQueue()) throw new QueueFullException(s.name, props.maxQueue());
        s.waiting++;
    } finally {
        lock.unlock();
    }
    boolean got = false;
    try {
        got = s.permits.tryAcquire(props.queueTimeout().toMillis(), TimeUnit.MILLISECONDS);
    } finally {
        lock.lock();
        try { s.waiting--; } finally { lock.unlock(); }
    }
    if (!got) throw new QueueTimeoutException(s.name);
}

import demo.GatewayExceptions.QueueFullException;
import demo.GatewayExceptions.QueueTimeoutException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import org.junit.jupiter.api.Test;
import org.springframework.util.unit.DataSize;

class QueueTest {
    @Test
    void oneSlotTwoWaitersThenReject() throws Exception {
        ModelManager manager = Lab.manager(Lab.props(2, 1, 2, Duration.ofSeconds(10), DataSize.ofGigabytes(64)),
                new Lab.MutableClock(), Lab.log());
        List<String> order = new CopyOnWriteArrayList<>();
        Lease first = manager.acquire("coder:7b", null);
        System.out.println("A running:  " + Lab.view(manager));

        List<Thread> waiters = new ArrayList<>();
        for (String name : new String[] {"B", "C"}) {
            waiters.add(Thread.ofVirtual().start(() -> {
                try (Lease l = manager.acquire("coder:7b", null)) {
                    order.add(name);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            }));
            Lab.sleep(100);                                           // B queues before C
        }
        System.out.println("B, C wait: " + Lab.view(manager));
        try {
            manager.acquire("coder:7b", null);
        } catch (QueueFullException e) {
            System.out.println("D rejected: " + e.getMessage());
        }
        System.out.println("after D:   " + Lab.view(manager));
        first.close();
        for (Thread t : waiters) t.join();
        System.out.println("service order after A: " + order);
        System.out.println("at the end: " + Lab.view(manager));
        manager.shutdown();
    }

    @Test
    void waitingHasATimeLimit() throws Exception {
        ModelManager manager = Lab.manager(Lab.props(2, 1, 4, Duration.ofMillis(300), DataSize.ofGigabytes(64)),
                new Lab.MutableClock(), Lab.log());
        try (Lease holder = manager.acquire("tiny:1b", null)) {
            try {
                manager.acquire("tiny:1b", null);
            } catch (QueueTimeoutException e) {
                System.out.println(e.getMessage());
            }
            System.out.println(Lab.view(manager));
        }
        manager.shutdown();
    }
}

/** The idle model used longest ago; models loaded with keep_alive -1 go last. */
private Slot victim() {
    return slots.values().stream()
            .filter(s -> s.state == State.READY && s.reserved == 0)
            .min(Comparator.comparing((Slot s) -> s.keepAlive.isNegative()).thenComparing(s -> s.lastUsed).thenComparing(s -> s.name))
            .orElse(null);
}

/** Stops the engine. The lock is released while the process exits; the slot keeps its memory reserved until then. */
private void stop(Slot s, String reason) {
    s.state = State.STOPPING;
    lock.unlock();
    try {
        if (s.backend != null) s.backend.close();
    } finally {
        lock.lock();
        slots.remove(s.name, s);
        changed.signalAll();
    }
    events.publishEvent(new ModelEvent.Unloaded(s.name, reason));
}

import java.util.List;
import org.junit.jupiter.api.Test;

class EvictionTest {
    @Test
    void leastRecentlyUsedIdleModelIsEvicted() throws Exception {
        List<String> log = Lab.log();
        Lab.MutableClock clock = new Lab.MutableClock();
        ModelManager manager = Lab.manager(Lab.props(2, 1), clock, log);

        manager.acquire("coder:7b", null).close();
        clock.advance(java.time.Duration.ofSeconds(10));
        manager.acquire("chat:8b", null).close();
        clock.advance(java.time.Duration.ofSeconds(10));
        manager.acquire("coder:7b", null).close();                 // coder is now newer than chat
        clock.advance(java.time.Duration.ofSeconds(10));
        System.out.println("loaded: " + Lab.view(manager));
        manager.acquire("tiny:1b", null).close();                  // limit is 2: one must go
        log.forEach(System.out::println);
        System.out.println("loaded: " + Lab.view(manager));
        manager.shutdown();
    }

    @Test
    void busyModelsAreNeverEvicted() throws Exception {
        List<String> log = Lab.log();
        ModelManager manager = Lab.manager(Lab.props(2, 1), new Lab.MutableClock(), log);
        Lease coder = manager.acquire("coder:7b", null);
        Lease chat = manager.acquire("chat:8b", null);
        Thread third = Thread.ofVirtual().start(() -> {
            try (Lease l = manager.acquire("tiny:1b", null)) {
                log.add("tiny:1b served");
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        });
        Lab.sleep(300);
        System.out.println("both busy, tiny waits: " + Lab.view(manager));
        coder.close();
        third.join();
        log.forEach(System.out::println);
        chat.close();
        manager.shutdown();
    }
}

import demo.GatewayProperties.ModelSpec;

/** The memory estimate behind `loop fit`: weights, KV cache for every slot, and a fixed reserve. */
public final class MemoryPlan {
    public static final long COMPUTE_RESERVE = 512L << 20;

    private MemoryPlan() { }

    public static long estimate(ModelSpec spec, int parallel) {
        return spec.weights().toBytes() + spec.kvBytesPerToken() * spec.contextTokens() * parallel + COMPUTE_RESERVE;
    }
}

import demo.GatewayExceptions.DoesNotFitException;
import java.time.Duration;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.springframework.util.unit.DataSize;

class MemoryCheckTest {
    @Test
    void memoryDecidesWhatStaysLoaded() throws Exception {
        List<String> log = Lab.log();
        GatewayProperties props = Lab.props(3, 2, 8, Duration.ofSeconds(10), DataSize.ofMegabytes(12_288));
        for (var e : props.models().entrySet()) {
            System.out.printf("%-9s needs %5d MiB%n", e.getKey(), MemoryPlan.estimate(e.getValue(), props.parallelPerModel()) >> 20);
        }
        System.out.println("budget    " + props.memoryBudget().toMegabytes() + " MiB");
        Lab.MutableClock clock = new Lab.MutableClock();
        ModelManager manager = Lab.manager(props, clock, log);
        manager.acquire("tiny:1b", null).close();
        clock.advance(Duration.ofSeconds(10));
        manager.acquire("coder:7b", null).close();
        clock.advance(Duration.ofSeconds(10));
        System.out.println("loaded: " + Lab.view(manager));
        manager.acquire("chat:8b", null).close();
        log.forEach(System.out::println);
        System.out.println("loaded: " + Lab.view(manager));
        try {
            manager.acquire("big:70b", null);
        } catch (DoesNotFitException e) {
            System.out.println(e.getMessage());
        }
        manager.shutdown();
    }
}

import java.util.Map;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.RestClient;

@RestController
class ChatController {
    private final ModelManager manager;
    private final RestClient http = RestClient.create();

    ChatController(ModelManager manager) { this.manager = manager; }

    @PostMapping("/v1/chat/completions")
    ResponseEntity<String> chat(@RequestBody Map<String, Object> body) throws InterruptedException {
        if (!(body.get("model") instanceof String model)) throw new IllegalArgumentException("model is required");
        try (Lease lease = manager.acquire(model, null)) {
            String upstream = http.post().uri(lease.baseUrl() + "/chat/completions")
                    .contentType(MediaType.APPLICATION_JSON).body(body).retrieve().body(String.class);
            return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(upstream);
        }
    }
}

import demo.GatewayExceptions.*;
import org.springframework.http.HttpStatus;
import org.springframework.http.ProblemDetail;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.ExceptionHandler;
import org.springframework.web.bind.annotation.RestControllerAdvice;

@RestControllerAdvice
class ApiErrors {
    @ExceptionHandler(UnknownModelException.class)
    ProblemDetail unknown(UnknownModelException e) { return problem(HttpStatus.NOT_FOUND, "Unknown model", e); }

    @ExceptionHandler(DoesNotFitException.class)
    ProblemDetail tooBig(DoesNotFitException e) { return problem(HttpStatus.UNPROCESSABLE_CONTENT, "Model does not fit", e); }

    @ExceptionHandler({QueueFullException.class, QueueTimeoutException.class, ShuttingDownException.class})
    ResponseEntity<ProblemDetail> busy(RuntimeException e) {
        return ResponseEntity.status(HttpStatus.SERVICE_UNAVAILABLE).header("Retry-After", "1")
                .body(problem(HttpStatus.SERVICE_UNAVAILABLE, "Server busy", e));
    }

    @ExceptionHandler(LoadFailedException.class)
    ProblemDetail load(LoadFailedException e) { return problem(HttpStatus.BAD_GATEWAY, "Engine failed to start", e); }

    @ExceptionHandler(IllegalArgumentException.class)
    ProblemDetail bad(IllegalArgumentException e) { return problem(HttpStatus.BAD_REQUEST, "Bad request", e); }

    private static ProblemDetail problem(HttpStatus status, String title, Exception e) {
        ProblemDetail p = ProblemDetail.forStatusAndDetail(status, e.getMessage());
        p.setTitle(title);
        return p;
    }
}

import com.fasterxml.jackson.annotation.JsonProperty;
import java.time.Duration;
import java.time.Instant;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;

@RestController
class LifecycleController {
    record PsEntry(String name, String model, long size, String state,
                   @JsonProperty("in_flight") int inFlight, int queued,
                   @JsonProperty("expires_at") Instant expiresAt) { }

    record PsResponse(List<PsEntry> models) { }

    private final ModelManager manager;
    private final GatewayProperties props;

    LifecycleController(ModelManager manager, GatewayProperties props) { this.manager = manager; this.props = props; }

    @GetMapping("/api/ps")
    PsResponse ps() {
        return new PsResponse(manager.snapshot().stream()
                .map(m -> new PsEntry(m.model(), m.model(), m.estimatedBytes(), m.state().name().toLowerCase(),
                        m.inFlight(), m.queued(), m.expiresAt())).toList());
    }

    /** Only the load and unload forms of Ollama's /api/generate: an empty prompt. */
    @PostMapping("/api/generate")
    Map<String, Object> generate(@RequestBody Map<String, Object> body) throws InterruptedException {
        if (!(body.get("model") instanceof String model)) throw new IllegalArgumentException("model is required");
        Object prompt = body.get("prompt");
        if (prompt != null && !prompt.toString().isEmpty()) {
            throw new IllegalArgumentException("only an empty prompt is handled here; generation goes through /v1/chat/completions");
        }
        Duration keepAlive = KeepAlive.parse(body.get("keep_alive"), props.keepAlive());
        Map<String, Object> out = new LinkedHashMap<>();
        out.put("model", model);
        out.put("response", "");
        out.put("done", true);
        if (keepAlive.isZero()) {
            manager.unload(model);
            out.put("done_reason", "unload");
        } else {
            manager.preload(model, keepAlive);
            out.put("done_reason", "load");
        }
        return out;
    }
}

public void pull(String model, Consumer<Progress> out) throws IOException, InterruptedException {
    out.accept(Progress.of("pulling manifest"));
    HttpResponse<String> m = http.send(HttpRequest.newBuilder(URI.create(props.hubUrl() + "/manifests/" + model)).build(),
            HttpResponse.BodyHandlers.ofString());
    if (m.statusCode() != 200) throw new IOException("registry answered " + m.statusCode() + " for " + model);
    Manifest manifest = json.readValue(m.body(), Manifest.class);

    Files.createDirectories(props.modelDir().resolve("blobs"));
    Path partial = props.modelDir().resolve("blobs").resolve(manifest.digest().replace(':', '-') + ".partial");
    String label = "pulling " + manifest.digest().substring(7, 19);
    out.accept(new Progress(label, manifest.digest(), manifest.size(), 0L));

    MessageDigest sha = sha256();
    long done = 0;
    HttpResponse<InputStream> blob = http.send(
            HttpRequest.newBuilder(URI.create(props.hubUrl() + "/blobs/" + manifest.digest())).build(),
            HttpResponse.BodyHandlers.ofInputStream());
    try (InputStream in = blob.body(); OutputStream file = Files.newOutputStream(partial)) {
        long nextReport = manifest.size() / 4;
        byte[] buf;
        while ((buf = in.readNBytes(CHUNK)).length > 0) {
            file.write(buf);
            sha.update(buf);
            done += buf.length;
            if (done >= nextReport) {
                out.accept(new Progress(label, manifest.digest(), manifest.size(), done));
                nextReport += manifest.size() / 4;
            }
        }
    }
    out.accept(Progress.of("verifying sha256 digest"));
    String actual = "sha256:" + HexFormat.of().formatHex(sha.digest());
    if (!actual.equals(manifest.digest())) {
        Files.deleteIfExists(partial);
        throw new IOException("digest mismatch: expected " + manifest.digest() + " but received " + actual);
    }
    Files.move(partial, props.modelDir().resolve("blobs").resolve(manifest.digest().replace(':', '-')),
            java.nio.file.StandardCopyOption.REPLACE_EXISTING);
    out.accept(Progress.of("writing manifest"));
    Files.createDirectories(props.modelDir().resolve("manifests"));
    Files.writeString(props.modelDir().resolve("manifests").resolve(model.replace(':', '_') + ".json"), m.body());
    out.accept(Progress.of("success"));
}

import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
import tools.jackson.databind.json.JsonMapper;

@RestController
class PullController {
    private final PullService pulls;
    private final JsonMapper json;

    PullController(PullService pulls, JsonMapper json) { this.pulls = pulls; this.json = json; }

    @PostMapping("/api/pull")
    ResponseEntity<StreamingResponseBody> pull(@RequestBody Map<String, Object> body) throws Exception {
        if (!(body.get("model") instanceof String model)) throw new IllegalArgumentException("model is required");
        if (Boolean.FALSE.equals(body.get("stream"))) {
            AtomicReference<PullService.Progress> last = new AtomicReference<>();
            pulls.pull(model, last::set);                         // runs to the end before any byte is sent
            StreamingResponseBody single = out -> write(out, json.writeValueAsString(last.get()));
            return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(single);
        }
        StreamingResponseBody stream = (OutputStream out) -> {
            try {
                pulls.pull(model, p -> write(out, json.writeValueAsString(p)));
            } catch (IOException e) {
                write(out, json.writeValueAsString(Map.of("error", e.getMessage())));
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        };
        return ResponseEntity.ok().contentType(MediaType.parseMediaType("application/x-ndjson")).body(stream);
    }

    private static void write(OutputStream out, String line) {
        try {
            out.write((line + "\n").getBytes(StandardCharsets.UTF_8));
            out.flush();
        } catch (IOException e) {
            throw new java.io.UncheckedIOException(e);       // the client left: stop downloading
        }
    }
}

@PreDestroy
void shutdown() {
    lock.lock();
    try {
        closing = true;
        changed.signalAll();
        for (Slot s : new ArrayList<>(slots.values())) {
            if (s.state == State.READY) stop(s, "gateway shutdown");
        }
    } finally {
        lock.unlock();
    }
}

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;

public class VictimPolicies {
    // keepAliveSeconds: Long.MAX_VALUE stands for "forever", as in Ollama (negative keep_alive becomes MaxInt64)
    record Runner(String name, long keepAliveSeconds, long lastUsed, int refCount) { }

    /** Ollama v0.35.1 findRunnerToUnload: sort by keep_alive, then name; first idle runner, else the first one. */
    static Runner ollama(List<Runner> loaded) {
        List<Runner> sorted = new ArrayList<>(loaded);
        sorted.sort(Comparator.comparingLong(Runner::keepAliveSeconds).thenComparing(Runner::name));
        for (Runner r : sorted) if (r.refCount() == 0) return r;
        return sorted.get(0);                          // the caller then waits until it goes idle
    }

    /** The gateway: the idle runner used longest ago, keep_alive forever last; none when every runner is busy. */
    static Runner lru(List<Runner> loaded) {
        return loaded.stream().filter(r -> r.refCount() == 0)
                .min(Comparator.comparing((Runner r) -> r.keepAliveSeconds() == Long.MAX_VALUE)
                        .thenComparingLong(Runner::lastUsed).thenComparing(Runner::name)).orElse(null);
    }

    static String name(Runner r) { return r == null ? "(wait for one to go idle)" : r.name(); }

    public static void main(String[] args) {
        List<Runner> mixed = List.of(
                new Runner("coder", 300, 100, 0),
                new Runner("chat", 1800, 10, 0),
                new Runner("embed", 60, 150, 1));
        System.out.println("mixed   ollama: " + name(ollama(mixed)) + ", gateway: " + name(lru(mixed)));

        List<Runner> allBusy = List.of(
                new Runner("coder", 300, 100, 1),
                new Runner("chat", 1800, 10, 2));
        System.out.println("busy    ollama: " + name(ollama(allBusy)) + ", gateway: " + name(lru(allBusy)));

        List<Runner> sameKeepAlive = List.of(
                new Runner("b-model", 300, 5, 0),
                new Runner("a-model", 300, 90, 0));
        System.out.println("equal   ollama: " + name(ollama(sameKeepAlive)) + ", gateway: " + name(lru(sameKeepAlive)));

        List<Runner> pinned = List.of(
                new Runner("pinned", Long.MAX_VALUE, 1, 0),
                new Runner("recent", 300, 500, 0));
        System.out.println("pinned  ollama: " + name(ollama(pinned)) + ", gateway: " + name(lru(pinned)));
    }
}