Docker setup for AI/ML with GPU acceleration and image optimization

Heavy models and GPUs do not tolerate migration between environments well: builds fall apart, training sits idle. We set up the environment end to end — from image definition to running on accelerators — so the team gets reproducible results and builds speed up several times over.
Frequently Asked Questions

Why Docker for AI/ML breaks on GPU and large models

The most common failure is visible on the very first run at the customer's site: in our environment the container gets access to the GPU, on theirs it does not.

The reason is usually not in the code, but in the mismatch between the driver on the host and the libraries inside the image: the runtime does not pass through the device, and the model silently falls back to the CPU. Such a non-reproducible build costs days of correspondence and reinstallations.

The second line is GPU containerization issues due to version conflicts between drivers and compute libraries. An image built on top of a fresh CUDA branch does not start on a host with an old driver, and updating it breaks an already running container.

As long as dependencies are not pinned, each environment lives its own life, and “it works on my machine” stops meaning anything.

Heavy model images make the picture worse. The base layer with a full set of libraries, weights, package cache, and build tools remain in the final artifact — the image grows, takes a long time to deploy on a new node, and bloats the registry. Each experiment adds another such build, and deleting it without consequences already feels scary.

Next is manual assembly for each model. A separate Dockerfile for each architecture, dependencies installed by hand, steps reproduced from the engineer's memory.

Project handover turns into a retelling rather than a transfer of artifacts: the contractor and the customer end up with different environments for the same code.

In total, these are not minor inconveniences, but direct time losses on every run and update. Below we analyze how these failures are addressed in Docker for AI/ML — at the level of image building, dependency pinning, and accelerator access.

What image optimization and GPU acceleration provide in practice

An image with a model accumulates dependencies faster than code: weight files, acceleration libraries, system packages. As long as everything is in one layer, any edit rebuilds it entirely — image build time grows with each iteration, and the registry bloats.

Splitting into layers and caching rarely changing parts enable builds in minutes instead of tens of minutes.

GPU acceleration in containers rests on three things: device passthrough, a compatible runtime, and pinned driver versions.

We move them into the base layer and prevent arbitrary updates — the container sees the accelerator the same way on all environments, from development to production.

For the integration with NVIDIA, passthrough is placed in a separate layer: changing the driver does not break the entire build.

Reproducibility of runs is a derivative of the same order. Docker image optimization includes pinning library versions, controlling base images, and arranging dependencies so that heavy weights are cached separately from code. An image built yesterday starts today without “environment mismatch”.

The result is measurable: build time is reduced severalfold, model image size drops without loss of functionality, the share of successful runs grows, and the class of failures due to version mismatch disappears from support.

Building Docker for AI/ML: comparison of approaches and modes

Building an image for an AI workload differs from a regular service: GPU libraries, driver bindings, and model weights go inside, and the base and layer order determine how long a rebuild takes and how often startup on the accelerator fails. In practice, one chooses among four schemes — and each has a cost.

Option When suitable Limitations
Single base image Multiple models on a shared stack, space saving on nodes matters Bloats: dependencies of all tasks sit in one layer, any update rebuilds everything
Separate image per model Models live their own cycle and require different library versions Number of artifacts and maintenance time grows, a registry without a retention policy quickly fills up
Layered overlay Base stack is stable, weights and model output code change Layers must be ordered from rare to frequent, otherwise cache is not reused and the build pulls everything again
Ready third-party image Need to test a hypothesis in hours, not build an environment Opaque composition, extra dependencies in layers, harder to pass vulnerability scanning and reproduce the build

On top of this, the choice is constrained by the container runtime itself: access to the accelerator via device passthrough, driver version mismatch between host and image, the weight of model weights in a layer, and the prohibition of downloading from the network in a closed environment.

Therefore, the scheme is chosen not by the “purity” of the solution, but by how often models change and how strict the reproducibility requirements are. Describe your task — we will return with an estimate and a plan.

Implementation order: from environment audit to handover

We structure the work order so that the working environment does not change blindly: first we record the current state, then we assemble the container scheme, and only after that we migrate the workload.

Each stage ends with an artifact that can be verified — without it, we do not proceed to the next step, so the discrepancy between “it works on my machine” and “it crashes on the server” surfaces during setup, not on launch day.

  1. Environment audit. We inventory servers, accelerator models, driver and library versions, and the dependency composition of each model. The output is a map: what is already deployed in containers, what lives on the host, and what conflicts with what.

  2. Container scheme. We determine the number of images, the base layer, the layer order taking cache into account, and tagging rules. We move heavy model weights outside the image — the image stops weighing gigabytes and rebuilds in minutes instead of tens of minutes.

  3. Access to accelerators. We configure device passthrough, user permissions inside the container, and isolation between services. We verify that when the host restarts, the container starts automatically and sees the accelerator without manual commands.

  4. Load testing. We run real inference and training scenarios: parallel requests, peak queues, behavior under memory shortage. We look at worst-case response times, not averages.

  5. Update policy. We fix the version policy, retention of recent builds, and the rollback procedure to the previous one. Updates stop being a risky event — they become a routine operation.

  6. Handover to operations. We provide build and run documentation, conduct team training, and go over typical failures. After that, support proceeds according to our policy or by your engineers.

Configurations, schemes, and policies you receive

The work ends not with a running container, but with a set of artifacts that your team can use to build, update, and roll back the environment without our involvement. Below is what exactly remains with you after handover.

  • Image map: which image is for which task — training, inference, data preparation — what lies in the shared base layer, where the application layer begins, and which dependencies are pinned strictly.
  • Accelerator placement scheme: which nodes provide accelerators, how they are shared between tasks, where memory limits are hit, and how device passthrough into the container is configured.
  • Registry access policies: roles and permissions, who publishes and who only reads, key lifetime, rules for retaining recent tags, and prohibition of overwriting already released ones.
  • Image update policy: frequency of rebuilding the base layer, verification procedure before publication, conditions and procedure for rolling back to the previous tag.
  • Build documentation: build files with comments, list of dependencies, rationale for multi-stage builds and for what is deliberately left outside the image.
  • Operations manual: start and stop, checking accelerator availability, log and metric collection, analysis of typical failures and the procedure for handling them.
  • Environment scheme: how the development, test, and production environments differ, what is transferred between them, and what is configured separately.

We accompany the set with a walkthrough with your team — so that decisions are further made based on documents, not on the memory of a specific engineer.

Case study: build time reduced, GPU failures eliminated

The project is training and deployment of models on several nodes with accelerators, multiple teams, and a shared server fleet. The environment for the task was assembled manually: dependency delivery, compatibility checks, launch.

Below is what it cost before implementation and what changed after switching to a unified image scheme.

What it was like

Out of roughly forty builds per month, half were reproduced with discrepancies: locally the model converged, on the server it crashed on library import.

Containers on Docker were built manually and stored outside a shared registry — each node lived with its own set of dependencies.

A separate category of losses was GPU failures: the driver on the node did not match the image's expectations, and the task died at startup after already occupying a training queue.

What was changed and what it became

We consolidated the environments into a single base image with a fixed set of libraries and layers: application images inherit from it and differ only in code and settings.

Driver compatibility check was moved into the build step — mismatch is detected before the task is queued, not in the middle of it.

Build takes minutes instead of hours, reproducibility of runs is close to complete, GPU failures at startup are reduced to isolated cases.

The registry does not grow: the retention policy keeps a limited number of layers per tag, and rolling back to the previous image is one command.

How does a shared base image differ from an image for a single model?

A shared base image is a single environment with a runtime, libraries for accelerator computing, and common utilities, on top of which each model is installed. An image for a model is a separate build with exactly the set of dependencies needed by one model.

The difference shows up where the team pays: with build time or with the risk of dependency conflict.

The first wins on layer reuse. The layer with system libraries and drivers is built once and then taken from cache for all models — a rebuild after code edits takes minutes instead of tens of minutes. One set of checks and one deployment scenario instead of several.

The downside is dependency conflict. When one model needs a newer library branch, version resolution breaks the neighboring environment: a model that responded correctly yesterday crashes at startup or changes its output, and the cause lies in the shared layer. Rollback is also shared — by reverting one model, you drag back the environment for everyone.

A separate image pays off when the model changes frequently and dependencies are strictly pinned. Code and weights lie in the top layer, the base remains from cache, rollback is pinpoint: swapping the tag returns the previous version without affecting neighboring services. The price is registry size and the number of build scenarios.

The choice of build strategy comes down to two questions: how often the model is updated and how much its requirements diverge from the rest.

With infrequent updates and a shared stack, one base image is sufficient; with frequent deployments and divergences, one takes a thin base layer with drivers and builds a model image on top of it — layer reuse is preserved, and dependency conflict remains within a single layer.

Limitations, compatibility, and common technical concerns

A container levels the application environment, but does not eliminate the layer beneath it: access to the accelerator, the driver version on the node, and kernel requirements remain outside the image.

This means the machine fleet must be kept at a consistent driver version — otherwise the same image on some nodes simply will not see the device. We fix the compatibility matrix in advance and verify it on a pilot node, not on deployment day.

Some things cannot be unified in principle: memory distribution between processes, real throughput under a specific model, behavior under memory shortage. Here, manual testing on a production load profile remains — synthetic tests show only the ceiling, not the operating mode.

And there is a boundary of applicability: where direct hardware access is needed or training on multiple nodes with a shared bus, part of it honestly remains on a bare-metal machine.

Our models already work outside containers — do we have to migrate everything from scratch?

No. We migrate step by step: first the environment — dependencies, weights, entry point, then verification of results on identical inputs, and only after that we switch traffic. The old launch remains available for rollback until the new one passes a full load cycle.

What to do if nodes are built on different drivers?

We build the image for the lowest agreed version — NVIDIA maintains backward compatibility within a major branch. Where a node lags, the driver is upgraded to the common branch; if that is impossible, we remove the node from the pool for that task rather than patching the image.

Let's discuss the task: environment, models, and launch requirements

Let's discuss the task concretely: so that the estimate is based on facts, not guesses, we will start with initial data on the environment and models. It is enough to send what you already have at hand.

From this data, it is clear where the bottleneck is — build, image size, driver compatibility, or reproducibility on another node.

What helps calculate the work accurately:

  • composition of models and launch mode: training, fine-tuning, or inference only;
  • environment parameters: type and number of accelerators, availability of drivers and runtime;
  • current Docker image build and complaints: time, size, differences between nodes;
  • image registry: where it is stored, who has access, how old versions are cleaned;
  • launch environment: single server, cloud, cluster, orchestrator;
  • data and volume handling: datasets, local weight copies, training artifacts;
  • limitations: closed environment, security policies, update policy.

The first conversation is led by the engineer who will then do the work — without retelling through intermediaries. We go over the launch scheme, contentious points, and decisions already made, after which we fix what we take first.

We will return with an estimate and an implementation plan: stages, timelines, responsible persons, checkpoints, and the procedure for rolling back to the previous image.

The next step is to describe the task and send initial data on the environment; we will clarify the rest in communication.

AI Development Areas

Latest works

  • B2B ADVANCE company website development
    B2B ADVANCE company website development
    1542
  • Development of a web application for FEEDME
    Development of a web application for FEEDME
    1393
  • Website development for BELFINGROUP
    Website development for BELFINGROUP
    1081
  • Development of an online store for the company FURNORO
    Development of an online store for the company FURNORO
    1368
  • B2B Advance company logo design
    B2B Advance company logo design
    815
  • Development of a web application for Enviok
    Development of a web application for Enviok
    1139