You’ve put in hours writing your script, downloading model weights, and testing inputs. Your AI agent works perfectly on your laptop. But as soon as you share the code with a teammate or try to run it on a cloud server, things fall apart. In AI, writing the model code is just the beginning. If you can’t package your application and deploy it reliably, your model stays stuck in the lab. That’s where Docker helps. In this article, I’ll show you how to Dockerize a simple AI agent from scratch.
Dockerize an AI Agent: Getting Started
Let’s dive in. We’ll build a simple conversational AI agent with the Hugging Face Transformers library and then package it using Docker.
Download Docker before you begin.
Create a new folder for your project. Inside, you’ll need three files: app.py, requirements.txt, and a Dockerfile.
Now, let’s start Dockerizing the AI agent step by step.
Step 1: The AI Agent
First, write the agent’s code. We’ll use distilgpt2, which is a smaller and faster version of GPT-2. It’s great for a lightweight tutorial.
Create a file named app.py and write the following code:
from transformers import pipeline
# Load a small text generation model
generator = pipeline("text-generation", model="distilgpt2")
def run_agent():
print("AI Agent is running. Type 'exit' to stop.\n")
while True:
user_input = input("You: ")
if user_input.lower() == "exit":
break
response = generator(user_input, max_length=50, num_return_sequences=1)
print("Agent:", response[0]['generated_text'], "\n")
if __name__ == "__main__":
run_agent()This code sets up a text-generation pipeline. The run_agent() function puts the user in a loop, reads their input, sends it to the distilgpt2 model, and prints the response until the user types “exit”.
Step 2: Defining Dependencies
Next, tell Docker which Python libraries your script needs.
Create a file named requirements.txt and add:
transformers
torch
These are the main libraries needed to run Hugging Face models. PyTorch (torch) is the engine for tensors, and transformers gives you the model interface.
Step 3: Writing the Dockerfile
The Dockerfile is like a blueprint for your container. It tells Docker how to build the environment your app needs.
Create a file named Dockerfile (no file extension) and add:
# Use official Python image
FROM python:3.10-slim
# Set working directory
WORKDIR /app
# Copy files into container
COPY . /app
# Install dependencies
RUN pip install --no-cache-dir -r requirements.txt
# Run the application
CMD ["python", "app.py"]Let’s break this down:
- FROM python:3.10-slim: We use a lightweight version of Python 3.10. Using an official slim image helps keep the final container size small.
- WORKDIR /app: This command creates a directory named /app in the container and sets it as the working directory.
- COPY . /app: This copies all files from your local project folder into the /app folder in the container.
- RUN pip install…: This installs the required libraries. The –no-cache-dir flag stops pip from saving cache files, which helps keep your Docker image smaller.
- CMD [“python”, “app.py”]: This sets the default command to run when the container starts.
Step 4: Building and Running the Container
Now let’s put everything together. Open your terminal, go to your project folder, and run the build command:
docker build -t ai-agent .
The dot at the end is important. It tells Docker to use the current directory to find the Dockerfile.
This process takes a few minutes while Docker downloads the Python image and installs PyTorch. When it’s finished, you’ll have a packaged image called ai-agent.
Next, run the container:
docker run -it ai-agent
Why use the -it flag? Our script needs user input from the terminal. The -it flags turn on interactive mode and a tty, so your terminal can connect to the container’s input.
When you run the container, you’ll see the Hugging Face pipeline download the distilgpt2 model weights the first time. After that, your interactive prompt will appear:
Downloading config.json: 100%|██████████| 762/762 [00:00<00:00, 2.51MB/s]
Downloading model.safetensors: 100%|████| 353M/353M [00:05<00:00, 68.4MB/s]
Downloading tokenizer_config.json: 100%|██| 26.0/26.0 [00:00<00:00, 89.2kB/s]
Downloading vocab.json: 100%|██████████| 1.04M/1.04M [00:00<00:00, 3.42MB/s]
Downloading merges.txt: 100%|██████████| 456k/456k [00:00<00:00, 1.63MB/s]
Downloading tokenizer.json: 100%|████████| 1.36M/1.36M [00:00<00:00, 4.31MB/s]
AI Agent is running. Type 'exit' to stop.
You: The future of artificial intelligence is
Agent: The future of artificial intelligence is not just about automation, but about augmenting human capabilities in ways we haven't fully imagined yet...
You: exit
Closing Thoughts
That’s how you Dockerize a simple AI agent from scratch.
Learning to use tools like Docker helps you move from running Jupyter notebooks locally to building scalable, production-ready systems as an AI Engineer. The environment you built can now be deployed to AWS, Google Cloud, or Azure, and it will work just like it did on your machine.
I hope you enjoyed the article! Follow me on Instagram for more AI and machine learning tips. You can also check out my book, Hands-On GenAI, LLMs & AI Agents, to get career-ready in AI.
