Implemented ollama container scaling

Author: gekinci

gen_compose.py

@@ -42,6 +42,19 @@ def expand_service(name, service, replicas):
     return services
 
 
+def expand_ollama_service(name, service, replicas):
+    services = expand_service(name, service, replicas)
+
+    if replicas == 1:
+        return services
+
+    for i in range(1, replicas + 1):
+        replica_name = f"{name}-{i}"
+        services[replica_name]["devices"] = [f"nvidia.com/gpu={i - 1}"]
+
+    return services
+
+
 def expand_depends_on(depends_on, replica_counts):
     if isinstance(depends_on, dict):
         expanded = {}
@@ -91,6 +104,7 @@ def haproxy_dependencies(
     services,
     backend_n,
     litellm_n,
+    ollama_n,
     available_mcp_servers,
     mcp_replica_n,
 ):
@@ -103,7 +117,8 @@ def haproxy_dependencies(
         )
 
     dependencies.extend(service_instance_names("litellm", litellm_n, services))
-    dependencies.extend(["mongodb", "ollama"])
+    dependencies.append("mongodb")
+    dependencies.extend(service_instance_names("ollama", ollama_n, services))
 
     return [dependency for dependency in dependencies if dependency in services]
 
@@ -122,7 +137,14 @@ def haproxy_backend(name, port, service_names, sticky_mode=None):
 
 
 def generate_haproxy(
-    services, backend_n, backend_port, litellm_n, server_list, replica_dict, port_dict
+    services,
+    backend_n,
+    backend_port,
+    litellm_n,
+    ollama_n,
+    server_list,
+    replica_dict,
+    port_dict,
 ):
     conf = []
 
@@ -151,6 +173,10 @@ def generate_haproxy(
         "frontend fe_litellm\n    bind *:4000\n    default_backend be_litellm\n\n"
     )
 
+    conf.append(
+        "frontend fe_ollama\n    bind *:11434\n    default_backend be_ollama\n\n"
+    )
+
     for s in server_list:
         conf.append(
             f"frontend fe_{s}\n"
@@ -176,6 +202,15 @@ def generate_haproxy(
         )
     )
 
+    conf.append(
+        haproxy_backend(
+            "ollama",
+            11434,
+            service_instance_names("ollama", ollama_n, services),
+            "leastconn",
+        )
+    )
+
     for s in server_list:
         conf.append(
             haproxy_backend(
@@ -201,6 +236,7 @@ def main():
     backend_target_port = os.environ.get("CLIMATECLAW_TARGET_PORT", "8502")
     backend_n = int(os.environ.get("CLIMATECLAW_BACKEND_REPLICAS", "1"))
     litellm_n = int(os.environ.get("CLIMATECLAW_LITELLM_REPLICAS", "1"))
+    ollama_n = int(os.environ.get("CLIMATECLAW_OLLAMA_REPLICAS", "1"))
 
     available_mcp_servers = [
         s
@@ -226,6 +262,7 @@ def main():
     replica_counts = {
         "climateclaw": backend_n,
         "litellm": litellm_n,
+        "ollama": ollama_n,
         **mcp_replica_n,
     }
 
@@ -234,6 +271,8 @@ def main():
             new_services.update(expand_service(name, svc, backend_n))
         elif name == "litellm":
             new_services.update(expand_service(name, svc, litellm_n))
+        elif name == "ollama":
+            new_services.update(expand_ollama_service(name, svc, ollama_n))
         elif name in MCP_SERVICES:
             if name in available_mcp_servers:
                 new_services.update(expand_service(name, svc, mcp_replica_n[name]))
@@ -257,16 +296,27 @@ def main():
 
     network_name = list(base["networks"].keys())[0]
 
+    haproxy_network = (
+        {
+            network_name: {
+                "aliases": ["ollama"],
+            }
+        }
+        if ollama_n > 1
+        else [network_name]
+    )
+
     new_services["haproxy"] = {
         "image": "haproxy:3.0-alpine",
         "user": "0:0",
         "ports": dev_ports if "dev" in compose_path else prod_ports,
         "volumes": ["./haproxy.cfg:/usr/local/etc/haproxy/haproxy.cfg:ro"],
-        "networks": [network_name],
+        "networks": haproxy_network,
         "depends_on": haproxy_dependencies(
             new_services,
             backend_n,
             litellm_n,
+            ollama_n,
             available_mcp_servers,
             mcp_replica_n,
         ),
@@ -285,6 +335,7 @@ def main():
         backend_n,
         backend_port,
         litellm_n,
+        ollama_n,
         available_mcp_servers,
         mcp_replica_n,
         port_dict,