AI models in Docker containers

Trying out AI models means installing a lot of things. Each one wants a particular version of Python, a particular PyTorch or TensorFlow, a pile of other packages and a few gigabytes of weights, and two models rarely agree on any of it. Installing them side by side on one machine doesn't go well for long.

So each model gets a container of its own. The AI branch of my Docker module has a service for each one, which builds an image with the model and everything it needs inside, and runs it behind a small WebSocket server. From Node, starting MusicGen is one call:

const Docker = require('Docker')

Docker.Service.AI.AUDIOCRAFT_MEDIUM('ai-musicgen-medium', 1340)

Then anything that can open a WebSocket can send it a prompt.

The models

They're grouped the way the module's tree is: a branch like Docker.Service.AI.Diffusions or Docker.Service.AI.Transformers holds the models that share a library, and each model is a function of its own, named in capitals after the model.

A model service

A model service looks just like any other service in the module: it writes a Dockerfile and passes it to Docker.Service.Start. Here's Stable Diffusion:

STABLE_DIFFUSION: (serviceName, socketPort, gpuAcceleration = false ) => {
    serviceName = Security.String(serviceName);
    socketPort = Security.Integer(socketPort);
    Security.Boolean(gpuAcceleration);

    const networkName = Docker.Config._networkName

    const dockerfileContent = [
        `FROM python:3.10`,
        `WORKDIR /app`,
        `RUN apt-get update && apt-get install -y nano`,
        `RUN pip install torch websockets diffusers transformers accelerate`,
        `COPY server.py ./`,
        `RUN python -c "from diffusers import StableDiffusionPipeline; StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4')"`,
        `CMD ["python", "server.py"]`
    ].join("\n");

    let containerArgs = [
        `-p${socketPort}:1337`,
        `--network`, networkName
    ];
    if (gpuAcceleration) {
        containerArgs = containerArgs.concat([`--gpus`, `all`])
    }

    const contextPath = __dirname+`/../../../../resources/Docker.Service.AI.Diffusions/Docker.Service.AI.Diffusions.STABLE_DIFFUSION`;
    return Docker.Service.Start(serviceName, dockerfileContent, contextPath, containerArgs);
}

Three things are the same in all of them:

• The weights are downloaded while the image builds. The RUN python -c "...from_pretrained(...)" line loads the model once during the build, which leaves the weights in the image. A container starts with the model already there, and never downloads anything while it's answering a request.
• The server is a Python file kept next to the module. Each service has a folder under the module's resources, named after the service, holding its server.py. That folder is the build context, so COPY server.py picks it up.
• The servers listen on port 1337 inside their containers. The port you pass in is mapped to it, so each model can be given its own port on the host without changing the servers. The few ready-made programs keep their own ports: the Stable Diffusion web interface is mapped to 7860 and Vosk to 2700.

The graphics card is optional. With gpuAcceleration set, the container is started with --gpus all; without it the model runs on the processor, which is much slower but works on any machine.

The WebSocket server

Each server.py is a small websockets server. The model is loaded once, when the server starts. A client connects, sends one JSON message with the prompt and its settings, and gets back progress messages, then either the result or an error, after which the server closes the connection:

model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
next_job_id = 1


class Messenger:
    def __init__(self, websocket):
        self.websocket = websocket

    async def progress(self, progress):
        payload = {
            "progress": progress
        }
        payload_serialized = json.dumps(payload)
        await self.websocket.send(payload_serialized)

    async def success(self, result):
        payload = {
            "type": "image",
            "success": result
        }
        payload_serialized = json.dumps(payload)
        await self.websocket.send(payload_serialized)
        await self.websocket.close()
...
async def handler(websocket):
    message = await websocket.recv()
    message_data = json.loads(message)
    messenger = Messenger(websocket)

    if message_data.get("started"):
...
        prompt: str = message_data.get("prompt")
        settings = message_data.get("settings")

        await messenger.progress(25)
        pipe = model.to(settings["device"])

        await messenger.progress(50)

        generated_images_arr: list = pipe(prompt).images

        encoded_images_arr: list = [to_base64(img) for _, img in enumerate(generated_images_arr)]

        print(f"Sequence complete for job id: {current_job_id}")
        await messenger.progress(100)

        # Send the result back to the client
        await messenger.success(encoded_images_arr)


if __name__ == "__main__":
    start_server = websockets.serve(handler, "0.0.0.0", 1337)

The messages are the same for every model:

Images come back as base64 data: URLs, so the result can go straight into a page or be written to a file. Since every model speaks the same protocol, the code that talks to them doesn't care whether it's asking for a picture, a paragraph or a piece of music.

Starting a model when it's needed

Nobody has to remember to start the containers. My Ai module, which gives each kind of task a function in Node with callbacks for the result and for errors, checks whether the model's container is running before it connects, and starts it if not:

//  Check if server is running & start docker container
const isStarted = Docker.Container.Exists(serviceName, true);
const connHandler = () => {
    console.log("Attempting to connect...")
    const ws = new WebSocket(serverUri);
    ws.on('open', function() {
        const payload = JSON.stringify({
            started: true,
            prompt,
            settings,
        })
        ws.send(payload);
    });
...
    ws.on('error', function (error) {
        setTimeout(connHandler, 5000)
    });
}
if (!isStarted) {
    console.log("Starting ai server...")
    let useGpu = settings.device.toLowerCase() === 'cuda' || settings.device.toLowerCase() === 'gpu';
    let containerId = null;
    if (settings.model === `DALLE`) {
        containerId = Docker.Service.AI.Diffusions.DALLE(serviceName, serverPort, useGpu);
    }
    else if (settings.model === `STABLE_DIFFUSION`) {
        containerId = Docker.Service.AI.Diffusions.STABLE_DIFFUSION(serviceName, serverPort, useGpu);
    }
    else if (settings.model === `LATENT_DIFFUSION`) {
        containerId = Docker.Service.AI.Diffusions.LATENT_DIFFUSION(serviceName, serverPort, useGpu);
    }

    console.log("Started diffusion server with ID: "+containerId)
    setTimeout(connHandler, 5000);
}
else {
    console.log("Ai server already running")
    connHandler()
}

A model takes a while to load once its container starts, so the client simply tries again every five seconds until the server answers. The first request after a restart is slow; the rest go straight through to a container that's already running with the model in memory.

The name of the container comes from the model, ai-diffusion-server-stable_diffusion here, so every script on the machine that asks for Stable Diffusion shares the one container.

Things to know

• The images are big, since each one holds its model's weights. Falcon's image downloads the whole 40-billion-parameter model while it builds.
• Starting a service with the Docker module rebuilds its image, which downloads the model again, so it's worth starting a model only when it isn't already running, as the Ai module does.
• Falcon's Dockerfile installs the processor-only build of PyTorch, so it runs on the processor even when --gpus all is given.
• The MusicGen servers load the model for the processor (device="cpu"). The image models take the device in each request's settings: "cpu" by default, or "cuda" for the graphics card.
• A server works on one request at a time: while a model is generating, the next connection waits.