Feature/shepherd (#20)

* wip: data preparation scripts

* wip: update req builder

* update workflow

* add simulator

* minor

Author: xzyaoi

.gitignore

@@ -5,3 +5,4 @@ pyrightconfig.json
 .local
 .vscode
 .zed
+.data

meta/requirements.txt

@@ -18,3 +18,4 @@ prometheus_client
 tqdm
 einops
 pillow
+tenacity

scratchpad/utils/client/__init__.py

@@ -0,0 +1 @@
+from .req import LLM

scratchpad/utils/client/req.py

@@ -0,0 +1,51 @@
+import os
+from typing import Optional
+import requests
+from tenacity import retry, stop_after_attempt, wait_fixed
+
+
+class LLM:
+    def __init__(
+        self,
+        model: str,
+        endpoint: Optional[str] = None,
+        api_key: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+    ):
+        if not endpoint:
+            endpoint = os.environ.get("RC_API_BASE", None)
+        if not api_key:
+            api_key = os.environ.get("RC_API_KEY", None)
+        if not endpoint or not api_key:
+            raise ValueError("API key or endpoint not found")
+        if not system_prompt:
+            system_prompt = "You are a helpful assistant."
+        self.model = model
+        self.endpoint = endpoint + "/chat/completions"
+        self.api_key = api_key
+        self.system_prompt = system_prompt
+        self.headers = {"Authorization": f"Bearer {self.api_key}"}
+
+    @retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
+    def __call__(self, prompt: str):
+        data = {
+            "model": self.model,
+            "messages": [
+                {"role": "system", "content": self.system_prompt},
+                {
+                    "role": "user",
+                    "content": prompt,
+                },
+            ],
+        }
+        try:
+            res = requests.post(
+                self.endpoint,
+                headers=self.headers,
+                json=data,
+            )
+            result = res.json()
+        except Exception as e:
+            print(f"Error calling LLM: {res.text}")
+            return None
+        return result["choices"][0]["message"]["content"]

scratchpad/utils/ui/__init__.py

@@ -0,0 +1 @@
+from .ui import make_table

scratchpad/utils/ui/ui.py

@@ -0,0 +1,21 @@
+from rich.console import Console
+from rich.table import Table
+import numpy as np
+
+console = Console()
+
+
+def make_table(title, data):
+    keys = data[0].keys()
+    table = Table(title=title)
+    colors = ["cyan", "magenta", "green", "yellow", "blue", "red", "black"]
+
+    for idx, column in enumerate(keys):
+        table.add_column(column, justify="Right", style=colors[idx % len(colors)])
+
+    for row in data:
+        for key in keys:
+            if type(row[key]) == np.float64 or type(row[key]) == float:
+                row[key] = str(round(row[key], 2))
+        table.add_row(*[row[key] for key in keys])
+    return table

scripts/serve_llama_1b.sh

@@ -0,0 +1,59 @@
+import time
+import requests
+
+
+def update_peer(endpoint, model, ipv4, port):
+    peer = {
+        "service": [
+            {
+                "name": "llm",
+                "status": "online",
+                "hardware": [],
+                "host": ipv4,
+                "port": port,
+                "identity_group": [f"model={model}"],
+            }
+        ]
+    }
+    res = requests.post(endpoint + "/v1/dnt/_node", json=peer)
+    print(res.text)
+
+
+def health_check(args):
+    local_address = f"http://{args.local_ip}:{args.service_port}/health"
+    is_healthy = False
+    while not is_healthy:
+        try:
+            print(f"Checking health of service at {local_address}", flush=True)
+            res = requests.get(local_address, timeout=5)
+            print(f"Service health check response: {res.status_code}", flush=True)
+            if res.status_code == 200:
+                is_healthy = True
+            else:
+                print(f"Service not ready yet, waiting for 5 seconds", flush=True)
+                time.sleep(5)
+        except Exception as e:
+            print(f"Service not ready yet, waiting for 5 seconds", flush=True)
+            time.sleep(5)
+    return is_healthy
+
+
+def register(args):
+    print(f"Registering service with config: {args}")
+    if health_check(args):
+        update_peer(args.ocf_addr, args.model_name, args.local_ip, args.service_port)
+    else:
+        print("Service is not healthy")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Register local service to OCF network"
+    )
+    parser.add_argument("--model-name", type=str, help="Model Name")
+    parser.add_argument("--ocf-addr", type=str, default="http://localhost:8092")
+    parser.add_argument("--local-ip", type=str, default="localhost")
+    parser.add_argument("--service-port", type=str, default="8000")
+    register(parser.parse_args())

tools/client/__init__.py

@@ -1,6 +1,9 @@
+import os
+import anyio
+import requests
 import aiohttp
 import asyncio
-from typing import Dict
+from typing import Dict, Optional
 
 
 async def async_request(endpoint, req: Dict):