ttsd-colabcli 1.0.0

package/core/daemon.py

@@ -0,0 +1,184 @@
+import asyncio
+import os
+import sys
+import uuid
+import time
+import httpx
+import argparse
+import base64
+import logging
+import signal
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Add project root to python path so we can import app.colab_cli
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from app.colab_cli.auth import AuthProvider, get_credentials
+from app.colab_cli.client import Client, Prod, Variant, Accelerator, ColabRequestError, TooManyAssignmentsError
+from app.colab_cli.common import State
+from app.colab_cli.commands.session import spawn_keep_alive
+from app.colab_cli.state import SessionState
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
+logger = logging.getLogger("satellite")
+
+MASTER_URL = os.getenv("MASTER_URL", "http://localhost:8090")
+NODE_API_KEY = os.getenv("NODE_API_KEY", "satellite-secret-key-123")
+NODE_ID = os.getenv("NODE_ID", f"node-{uuid.uuid4().hex[:6]}")
+
+HEADERS = {"x-node-key": NODE_API_KEY}
+
+async def fetch_token(email: str):
+    token_dir = os.path.expanduser("~/.config/colab-cli")
+    os.makedirs(token_dir, exist_ok=True)
+    safe = email.replace("@", "_at_").replace(".", "_")
+    token_path = os.path.join(token_dir, f"token_{safe}.json")
+    
+    async with httpx.AsyncClient() as client:
+        res = await client.get(f"{MASTER_URL}/api/node/token/{email}", headers=HEADERS)
+        if res.status_code == 200:
+            with open(token_path, "w") as f:
+                f.write(res.text)
+            return token_path
+        raise Exception(f"Failed to fetch token: {res.text}")
+
+async def report_status(email: str, status: str, colab_pid: int = None, worker_session_id: str = None, error: str = None):
+    async with httpx.AsyncClient() as client:
+        payload = {
+            "email": email,
+            "status": status,
+            "colab_pid": colab_pid,
+            "worker_session_id": worker_session_id,
+            "error": error
+        }
+        await client.post(f"{MASTER_URL}/api/node/status", json=payload, headers=HEADERS)
+
+async def launch_worker(email: str, public_url: str):
+    logger.info("Launching worker for %s", email)
+    
+    try:
+        token_path = await fetch_token(email)
+    except Exception as e:
+        logger.error("Failed to fetch token: %s", e)
+        await report_status(email, "FAILED", error=str(e))
+        return
+        
+    _creds = get_credentials(provider=AuthProvider.OAUTH2, token_path=token_path)
+    _client = Client(Prod(), _creds)
+    
+    _state = State()
+    _state.config_path = os.path.expanduser(f"~/.config/colab-cli/store_{email.replace('@', '_at_').replace('.', '_')}.json")
+    _state.auth_provider = AuthProvider.OAUTH2
+    _state._client = _client
+    
+    name = f"worker-{NODE_ID}-{uuid.uuid4().hex[:4]}"
+    variant = Variant.GPU
+    accelerator = Accelerator.T4
+    
+    loop = asyncio.get_running_loop()
+    
+    try:
+        res = await loop.run_in_executor(None, _client.assign, uuid.uuid4(), variant, accelerator)
+    except Exception as e:
+        logger.error("Assign failed for %s: %s", email, e)
+        await report_status(email, "FAILED", error=str(e))
+        return
+
+    token = res.runtime_proxy_info.token if hasattr(res, "runtime_proxy_info") else getattr(res, "runtime_proxy_token", "")
+    url = res.runtime_proxy_info.url if hasattr(res, "runtime_proxy_info") else ""
+    endpoint = res.endpoint
+
+    s = SessionState(name=name, token=token, url=url, endpoint=endpoint, variant=variant.value, accelerator=accelerator.value)
+    
+    try:
+        await loop.run_in_executor(None, _client.keep_alive_assignment, endpoint)
+    except Exception as e:
+        logger.warning("Keep-alive failed for %s: %s", email, e)
+
+    _state.store.add(s)
+    pid = spawn_keep_alive(endpoint, name, auth_provider=_state.auth_provider, config_path=_state.config_path, token_path=token_path)
+    s.keep_alive_pid = pid
+    _state.store.add(s)
+
+    deploy_wsid = str(uuid.uuid4())
+    
+    try:
+        from app.colab_cli.runtime import ColabRuntime
+        _worker_bytes = open(os.path.join(os.path.dirname(__file__), "../colab/worker.py"), "rb").read()
+        _worker_b64 = base64.b64encode(_worker_bytes).decode()
+        
+        _rt = ColabRuntime(url, token, session_name=name)
+        _deploy_code = f"""
+import subprocess, sys, os, json, base64
+subprocess.run([sys.executable, "-m", "pip", "install", "-q", "websockets", "httpx", "soundfile", "pydub", "omnivoice", "--no-deps", "transformers>=4.44.0"], capture_output=True, timeout=120)
+print("DEPS_OK", flush=True)
+worker_path = "/content/worker.py"
+with open(worker_path, "wb") as f:
+    f.write(base64.b64decode("{_worker_b64}"))
+env = os.environ.copy()
+env["OMNIVOICE_NUM_STEP"] = "16"
+env["OMNIVOICE_GUIDANCE_SCALE"] = "4.0"
+proc = subprocess.Popen([sys.executable, worker_path,
+    "--server-url", "{public_url}",
+    "--email", "{email}",
+    "--worker-session-id", "{deploy_wsid}"], env=env, start_new_session=True,
+    stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+import time
+time.sleep(10)
+ret = proc.poll()
+print("DEPLOYED PID=" + str(proc.pid) + " ALIVE=" + str(ret is None), flush=True)
+"""
+        await loop.run_in_executor(None, lambda: _rt.execute_code(_deploy_code, timeout=300))
+        logger.info("Worker deployed on Colab runtime for %s", email)
+    except Exception as e:
+        logger.error("Worker deploy failed: %s", e)
+        await report_status(email, "FAILED", error=str(e))
+        return
+
+    await report_status(email, "READY", colab_pid=pid, worker_session_id=deploy_wsid)
+    running_workers[email] = {"pid": pid, "start_time": time.time()}
+
+running_workers = {}
+
+async def monitor_workers():
+    while True:
+        now = time.time()
+        for email, info in list(running_workers.items()):
+            if now - info["start_time"] > 9000:  # 2.5 hours
+                logger.info("Auto-rotating worker for %s (reached 2.5h limit)", email)
+                try:
+                    os.kill(info["pid"], signal.SIGTERM)
+                except Exception as e:
+                    logger.debug("Failed to kill PID %s: %s", info["pid"], e)
+                
+                await report_status(email, "FAILED", error="Auto-rotation limit reached")
+                del running_workers[email]
+                
+        await asyncio.sleep(60)
+
+async def main():
+    logger.info("Starting Satellite Daemon %s", NODE_ID)
+    logger.info("Master URL: %s", MASTER_URL)
+    
+    asyncio.create_task(monitor_workers())
+    
+    async with httpx.AsyncClient() as client:
+        while True:
+            try:
+                res = await client.get(f"{MASTER_URL}/api/node/jobs", headers=HEADERS)
+                if res.status_code == 200:
+                    job = res.json().get("job")
+                    if job:
+                        email = job["email"]
+                        public_url = job.get("public_url") or MASTER_URL
+                        # Launch in background so we don't block polling
+                        asyncio.create_task(launch_worker(email, public_url))
+            except Exception as e:
+                logger.debug("Polling error: %s", e)
+            
+            await asyncio.sleep(5)
+
+if __name__ == "__main__":
+    asyncio.run(main())

package/core/requirements.txt

@@ -0,0 +1,8 @@
+httpx==0.27.2
+websockets==12.0
+google-auth-oauthlib==1.2.0
+typer>=0.12.0
+nbformat>=5.9.0
+websocket-client>=1.7.0
+filelock>=3.13.0
+python-dotenv>=1.0.0

package/package.json

@@ -0,0 +1,22 @@
+{
+  "name": "ttsd-colabcli",
+  "version": "1.0.0",
+  "description": "Satellite Node CLI for TTS Dubbing Distributed Architecture",
+  "main": "cli.js",
+  "bin": {
+    "colabcli": "./cli.js"
+  },
+  "scripts": {
+    "start": "node cli.js"
+  },
+  "keywords": [
+    "cli",
+    "tts",
+    "worker"
+  ],
+  "author": "",
+  "license": "ISC",
+  "dependencies": {
+    "inquirer": "^8.2.5"
+  }
+}