How to Run an AI Image Generation Worker on Linux with systemd, Retries and Cost Tracking

Teams that start using image generation APIs usually begin with a script on a laptop. It works for a demo, but as soon as other people depend on it, the usual operational questions appear. Where does the API key live? What happens when the provider times out? Who notices when half the jobs fail overnight? How much did last week’s batch cost?

This tutorial shows how to turn that script into a small, boring and reliable Linux service. We will build a worker that pulls prompts from a queue, calls an image generation API, saves every result locally, logs to journald, retries sensibly and keeps a running cost ledger. The examples use Ubuntu 24.04 and Python 3, but the approach works on any systemd-based distribution.

Step 1: Create a service user and directories

The worker should never run as root or as your personal account. Create a dedicated system user with no login shell and give it ownership of a working directory.

sudo useradd --system --home /srv/imagegen --shell /usr/sbin/nologin imagegen
sudo mkdir -p /srv/imagegen/{out,app}
sudo chown -R imagegen:imagegen /srv/imagegen
sudo apt install -y python3-venv sqlite3
sudo -u imagegen python3 -m venv /srv/imagegen/venv
sudo -u imagegen /srv/imagegen/venv/bin/pip install requests

Step 2: Keep the API key out of the code

Store credentials in an environment file that only root and the service user can read. The worker reads the key and the provider endpoint from the environment, so nothing sensitive ends up in version control.

sudo tee /etc/imagegen.env >/dev/null <<'EOF'
API_KEY=replace-with-your-key
API_URL=replace-with-your-provider-endpoint
DAILY_BUDGET_USD=5.00
EOF
sudo chown root:imagegen /etc/imagegen.env
sudo chmod 640 /etc/imagegen.env

Step 3: A queue that is easy to inspect

For a single host, SQLite is a perfectly good queue. It survives restarts, needs no extra daemon and can be inspected with one command. If you later run several workers on different machines, swap it for Redis or a managed queue without changing the rest of the design.

sudo -u imagegen sqlite3 /srv/imagegen/jobs.db "
CREATE TABLE IF NOT EXISTS jobs(
  id INTEGER PRIMARY KEY, prompt TEXT, size TEXT DEFAULT '1K',
  status TEXT DEFAULT 'queued', attempts INTEGER DEFAULT 0,
  output TEXT, error TEXT, created TEXT DEFAULT CURRENT_TIMESTAMP);"

Adding a job is a single insert, which makes it easy to feed the queue from a cron job, a web form or another service.

Step 4: Choose a model and know its costs

Before writing the worker, decide which model it will call and what each image will cost. That number drives the budget check later. As an example, Microsoft’s MAI-Image-2.6 supports text-to-image generation, editing of an existing image and composition from up to five reference images, with output at 1K or 2K and aspect ratios from 1:4 to 4:1. Through the MAI-Image-2.6 API on APIMart, it is billed by actual tokens, with listed planning estimates of roughly $0.031 per 1K image and $0.070 per 2K image. Each request returns a single PNG, and failed tasks are refunded, which keeps the ledger logic simple: one successful job equals one image and one estimated charge.

Step 5: The worker

The worker loops forever: it checks the daily budget, claims one queued job, calls the API, saves the image and records the cost. Provider details vary, so the request body below is deliberately generic; adapt the field names to your provider’s documentation. Many image APIs return a task ID first and the image later, which is handled by the polling helper.

#!/srv/imagegen/venv/bin/python
import os, time, sqlite3, csv, datetime, requests

DB, OUT = "/srv/imagegen/jobs.db", "/srv/imagegen/out"
KEY, URL = os.environ["API_KEY"], os.environ["API_URL"]
BUDGET = float(os.environ.get("DAILY_BUDGET_USD", "5"))
PRICE = {"1K": 0.031, "2K": 0.070}
H = {"Authorization": f"Bearer {KEY}"}

def spent_today():
    today = datetime.date.today().isoformat()
    try:
        with open("/srv/imagegen/cost.csv") as f:
            return sum(float(r[3]) for r in csv.reader(f) if r[0].startswith(today))
    except FileNotFoundError:
        return 0.0

def poll(task_url):
    delay = 3
    for _ in range(40):
        r = requests.get(task_url, headers=H, timeout=30).json()
        if r.get("status") in ("succeeded", "failed"):
            return r
        time.sleep(delay); delay = min(delay * 2, 30)
    raise TimeoutError("task did not finish")

The main loop claims work with an update statement, so two workers on the same database never pick the same job.

while True:
    if spent_today() >= BUDGET:
        print("daily budget reached, sleeping"); time.sleep(600); continue
    db = sqlite3.connect(DB)
    row = db.execute("SELECT id,prompt,size,attempts FROM jobs "
                     "WHERE status='queued' ORDER BY id LIMIT 1").fetchone()
    if not row:
        time.sleep(10); continue
    jid, prompt, size, attempts = row
    db.execute("UPDATE jobs SET status='running', attempts=attempts+1 WHERE id=?", (jid,)); db.commit()
    try:
        r = requests.post(URL, headers=H, json={"prompt": prompt, "size": size}, timeout=60)
        r.raise_for_status()
        result = poll(r.json()["task_url"])
        img = requests.get(result["image_url"], timeout=120).content
        path = f"{OUT}/{jid}.png"; open(path, "wb").write(img)
        with open("/srv/imagegen/cost.csv", "a") as f:
            csv.writer(f).writerow([datetime.datetime.now().isoformat(), jid, size, PRICE[size]])
        db.execute("UPDATE jobs SET status='done', output=? WHERE id=?", (path, jid))
        print(f"job {jid} done")
    except Exception as e:
        state = "queued" if attempts + 1 < 3 else "failed"
        db.execute("UPDATE jobs SET status=?, error=? WHERE id=?", (state, str(e)[:500], jid))
        print(f"job {jid} error: {e}"); time.sleep(5 * (attempts + 1))
    db.commit(); db.close()

Note that the image is downloaded immediately. Result links from image APIs are usually temporary, so a worker that only stores the provider URL will eventually serve broken images.

Step 6: Run it under systemd

A systemd unit gives you automatic restarts, logging and basic sandboxing for free.

sudo tee /etc/systemd/system/imagegen.service >/dev/null <<'EOF'
[Unit]
Description=Image generation worker
After=network-online.target

[Service]
User=imagegen
EnvironmentFile=/etc/imagegen.env
ExecStart=/srv/imagegen/app/worker.py
Restart=on-failure
RestartSec=10
ProtectSystem=strict
ReadWritePaths=/srv/imagegen
NoNewPrivileges=true

[Install]
WantedBy=multi-user.target
EOF
sudo systemctl daemon-reload
sudo systemctl enable --now imagegen
journalctl -u imagegen -f

Step 7: Watch failures and spending

Two quick queries cover most day-to-day questions: how many jobs failed recently and how much has been spent today.

sqlite3 /srv/imagegen/jobs.db "SELECT status, COUNT(*) FROM jobs GROUP BY status;"
awk -F, -v d="$(date +%F)" '$1 ~ d {s+=$4} END {printf "%.2f USD\n", s}' /srv/imagegen/cost.csv

Wire the failure count into whatever monitoring you already use. A simple rule, such as alerting when more than ten percent of jobs fail within an hour, catches expired keys, provider incidents and malformed prompts quickly.

Troubleshooting common problems

Every job fails immediately. Check the journal first. An authentication error usually means the environment file was not loaded; run systemctl show imagegen -p EnvironmentFiles and confirm the file permissions allow the service user to read it.

Jobs stay in the running state. This happens when the worker is killed mid-job. Add a small startup step that resets jobs stuck in running for more than fifteen minutes back to queued.

Images look wrong after a model update. Keep a fixed set of ten test prompts and run them through the queue after any configuration change. Compare the new outputs with the previous set before switching production traffic.

The disk fills up. Add a systemd timer that moves files older than a week to object storage, and monitor free space on the output volume.

Where to go next

Once the basic worker is stable, the natural extensions are syncing the output directory to object storage with rclone, adding a small HTTP endpoint for submitting jobs, and keeping the model name and price table in the environment file so a model change does not require a code edit. None of these change the core idea: treat image generation like any other external dependency, with a queue, bounded retries, local copies and a ledger. That is what turns a clever script into a service your team can rely on.