Ornith experiment - Load testing

I did a fun local AI experiment with the Ornith 9b open-weights model. I was about to take a flight without internet access, so I fired up Ollama and pulled the model to see if it could yield any meaningful output. The computer I use when I am on the go is an M1 16GB MacBook Air.

TL;DR it failed, not miserably, but promisingly.

The task was to create a load testing script with Python, to test a new internal project that needs to be remarkably resilient under heavy load. I prompted the whole thing through the Ollama UI, as I had not installed a coding agent on this computer that worked with local models.

A few minutes after my first prompt and a significant amount of battery, it yielded a complete Python CLI that included typings, configurable CLI arguments with appropriate defaults and the server's semantics built into the client. There was also an impressive amount of load-testing primitives included, like capacity token buckets, threadpools, etc. It required though the requests package, which was not available in this environment. I sent a second prompt to rewrite it using Python's built-in urllib, which it did after a few minutes.

The final script looked like this, after truncating much of the business logic for the sake of brevity:

# Rate limiter — one shared bucket across all worker threads
@dataclass
class TokenBucket:
    rate: float
    capacity: float = 1000.0
    tokens: float = field(default=1000.0)

    def consume(self, n: int = 1) -> bool:
        self._refill_tokens()
        
        if self._has_sufficient_tokens(n):
            return self._deduct_tokens(n)
        
        return self._throttle_and_deplete(n)


# Payload
@dataclass
class Config:
    url: str
    rate: float = 1000.0  # total requests/sec across workers
    concurrency: int = 10  # parallel senders
    duration: float = 60.0  # seconds; 0 => run forever
    bucket_capacity: float = 1000.0


class PayloadBuilder:
    def __init__(self, cfg: Config):
        self.cfg = cfg

    def build(self, n: int) -> dict:
        return {...}

def run(cfg: Config) -> dict:
    builder = PayloadBuilder(cfg)
    bucket = TokenBucket(cfg.rate, cfg.bucket_capacity)
    errors = 0
    sent = 0
    lock = threading.Lock()
    start = time.monotonic()

    with ThreadPoolExecutor(max_workers=cfg.concurrency) as ex:
        while True:
            elapsed = _log_progress(start, sent, errors, lock)

            if _should_stop(cfg, elapsed):
                break
            
            _dispatch_workers(ex, cfg, bucket, builder, lock, sent, errors)

    return {"sent": sent, "errors": errors, "elapsed": time.monotonic() - start}


def print_summary(stats: dict, cfg: Config) -> None:
    sent, errors, elapsed = stats["sent"], stats["errors"], stats["elapsed"]
    print("\n=== Summary ===\n")
    print(f"  Duration:        {elapsed}s")
    print(f"  Requests sent:   {sent}")
    print(f"  Errors:          {errors}")
    print(f"  Throughput:      {sent / elapsed:.0f} req/s")


def main() -> None:
    p = argparse.ArgumentParser(description="Load-test a URL endpoint.")
    p.add_argument("--url", default=None, help="The URL endpoint to load test")

    p.add_argument(
        "--rate",
        type=float,
        default=1000.0,
        help="target requests/sec (total across all workers)",
    )
    p.add_argument(
        "--concurrency", type=int, default=10, help="number of parallel senders"
    )
    p.add_argument(
        "--duration", type=float, default=60.0, help="seconds to run (0 = forever)"
    )
    p.add_argument(
        "--bucket-capacity", type=float, default=1000.0, help="burst allowance"
    )
    args = p.parse_args()

    cfg = Config(
        url=args.url,
        rate=args.rate,
        concurrency=args.concurrency,
        duration=args.duration,
        bucket_capacity=args.bucket_capacity,
    )

    stats = run(cfg)
    print_summary(stats, cfg)


if __name__ == "__main__":
    main()

It still ran into trouble working with some internal libraries needed, which were not easily available to Ollama UI and had these redundant from __future__ import annotation import that all LLMs love this days, but that's OK. 90% of the work done on a 6 year old CPU on a fanless laptop without internet. This is even more impressive for a coding model that is just 9B (~6GB) in size.

Kudos to the Ornith team and looking forward to the future!