Two GPUs, DataParallel, and Why There Is Still No Broker
Part two of the YOLO vision platform: why two GPUs are not two lanes, why DataParallel and not DDP, why a turn per GPU before Celery, how the AI service went from script to container, and four decisions I would make differently today.
Related project: YOLO Training & Inference Orchestration Architecture
This article is the second part of No queue until it is needed, about an internal computer vision platform for training and running YOLO models. The first part covered the boundary between the web and the GPU and how the platform grew by evidence. This one goes one floor down: what happens close to the GPU, how the AI service became a service, and a few decisions I would make differently today. Like the architecture case study, it names no organization, application field or measured result, and includes no code.
Two GPUs are not two lanes
The platform ran on a single machine with two NVIDIA GPUs. It is easy to read that as capacity for two jobs at once, but it was not, and understanding why shapes everything else.
PyTorch offers two ways to train on both GPUs. DistributedDataParallel (DDP) scales better: it starts one process per GPU, each with its own copy of the model, and synchronizes gradients between them. I tried it and could not make it stable. DDP needs a launcher and processes that find each other, and that way of working fits poorly inside a web service, which has its own process model. DataParallel is more modest: a single process splits each batch across both GPUs and gathers the results. It scales less, because one GPU acts as coordinator, but it trained on both without changing the architecture of the service.
The consequence is that, while training, the two GPUs were a single lane. Inference, for its part, ran on one selected GPU. A queue in front of that lane would not have added capacity: it would only have ordered the wait.
Admission before a broker
The natural next step is to spread inference across both GPUs. Then there really are two lanes, and a new problem appears: two jobs must not land on the same GPU and fight over its memory. That is admission control, and on a single machine it can be as simple as one turn per GPU: a semaphore a job takes before touching the device and releases when it is done.
Celery does not solve this on its own. The usual way to run it with GPUs is one worker per GPU, pinned with CUDA_VISIBLE_DEVICES and a concurrency of one: the same two lanes, only outside the service process. What a broker does buy is something else: restarting the API does not kill a running job, retries, cancellation, and the option to add a second machine. As long as none of those needs is real, the semaphore is enough.
Here I correct my own roadmap. The documentation of the revision said that a second GPU would be a reason for a broker. It is not: a second GPU on the same machine is a second admission lane. A broker is justified by a second machine or by the need to retry and cancel. The revision left admission uncontrolled and wrote it down as the open risk of the design: it is what comes next.
From script to service
The AI service was not born as a service. It started as Python scripts run from the terminal. When the web application needed to request work from it, I wrapped it in FastAPI, still on the host, and the Django application, running in its Docker container, called it over the network.
Keeping the GPU service on the host was deliberate: it used the NVIDIA driver, CUDA and PyTorch directly, without fighting GPU access from inside a container. The cost arrived over time, as the host and the containers drifted apart in versions. In the next version everything moved into one project and one Docker Compose file: the web application and the GPU service, each in its own container, with GPU access declared explicitly and both images built and tested in continuous integration. Those tests run without a GPU, against a simulated runtime, and the documentation says so: the real GPU path is not validated in CI.
In that same version, training left the platform. Its environment needed newer versions than inference, and keeping them together would have held back every inference release for a training upgrade. Training now happens separately, and the weights come in through an import step that fingerprints them.
Linux, drivers, and why it was not a whim
On a GPU platform, the most fragile part is not the code: it is getting the host's NVIDIA driver, the CUDA version in the container and the PyTorch version to agree. Ubuntu was the baseline with the least friction for that combination. On Windows the simple cases worked, but multi-GPU training ran into process handling: Python creates processes differently there, and PyTorch does not offer NCCL, the GPU communication backend that distributed training uses on Linux.
Cloud or local: two questions, not one
The cloud question is often answered with "it is expensive" or "it is simpler". The criterion that holds up is a different one: how many GPU hours are used per month, and where the data is born. Long training runs over large image sets that already live locally pay for GPU by the hour and for uploading the data on top. That is why the analysis pointed to a hybrid model: the GPU next to the data, and the application and the selected models in the cloud. What that analysis still lacks, and I say so, is the break-even point in numbers: without it, it is a direction, not a conclusion.
What I would do differently today
Separate submission from execution from day one. In the first version, the HTTP request stayed open for the whole job. Submitting, receiving an identifier and polling for status costs a table and an endpoint, and removes the whole problem. The next version did it, with job records that survive a restart.
Choose FastAPI for the contracts, not for async. Code that trains or infers on the GPU blocks either way. An async endpoint that calls that code without handing it to a thread freezes the event loop, health check included. What FastAPI does bring between two services is typed contracts and automatic API documentation.
Name the reason for separating correctly. Splitting the web from the GPU was not about scaling each part on its own: on a single machine that is theory. The real reasons were that a crashed training run does not take the site down, and that the CUDA and PyTorch environment gets its own upgrade pace. That second reason is the same one that later moved training out of the platform.
Put admission in place before enabling a second job type on the GPU. It is cheaper than the memory incident that would make it necessary.
I designed and documented this architecture. The decisions (ADRs), the limits, the roadmap and the later revision are in the case study repository.
Comments
Comments live on GitHub Discussions: sign in with GitHub in the box below to reply, or open the thread on GitHub.
No GitHub account? Leave a comment here
Comments are reviewed before they appear. Your email is optional and is never published.