From ae5d16f3942d5f365dc231bd027295555e4fd392 Mon Sep 17 00:00:00 2001
From: Yannic Kilcher <yk@users.noreply.github.com>
Date: Fri, 27 Jan 2023 15:40:21 +0100
Subject: [PATCH] added a tmux inference dev setup script

---
 inference/full-dev-setup.sh  | 19 +++++++++++++++++++
 inference/worker/__main__.py |  2 ++
 2 files changed, 21 insertions(+)
 create mode 100755 inference/full-dev-setup.sh

diff --git a/inference/full-dev-setup.sh b/inference/full-dev-setup.sh
new file mode 100755
index 00000000..98a5b173
--- /dev/null
+++ b/inference/full-dev-setup.sh
@@ -0,0 +1,19 @@
+#!/bin/bash
+
+# Creates a tmux window with splits for the individual services
+
+tmux new-session -d -s "inference-dev-setup"
+tmux send-keys "docker run --rm -it -p 6379:6379 redis" C-m
+tmux split-window -h
+tmux send-keys "docker run --rm -it -p 8001:80 -e MODEL_ID=distilgpt2 ykilcher/text-generation-inference" C-m
+tmux split-window -h
+tmux send-keys "cd server" C-m
+tmux send-keys "uvicorn main:app --reload" C-m
+tmux split-window -h
+tmux send-keys "cd worker" C-m
+tmux send-keys "python __main__.py" C-m
+tmux split-window -h
+tmux send-keys "cd text-client" C-m
+tmux send-keys "sleep 5" C-m
+tmux send-keys "python __main__.py" C-m
+tmux attach-session -t "inference-dev-setup"
diff --git a/inference/worker/__main__.py b/inference/worker/__main__.py
index 190ed788..96fe164a 100644
--- a/inference/worker/__main__.py
+++ b/inference/worker/__main__.py
@@ -18,8 +18,10 @@ def main(
     inference_server_url: str = "http://localhost:8001",
 ):
     def on_open(ws: websocket.WebSocket):
+        logger.info("Connected to backend, sending config...")
         worker_config = inference.WorkerConfig(model_name=model_name)
         ws.send(worker_config.json())
+        logger.info("Config sent, waiting for work...")
 
     def on_message(ws: websocket.WebSocket, message: str):
         # TODO: what if this comes in, but one is already in progress?