|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# -*- coding: utf-8 -*- |
| 3 | + |
| 4 | +import os |
| 5 | +import sys |
| 6 | +import shlex |
| 7 | +import shutil |
| 8 | +import signal |
| 9 | +import atexit |
| 10 | +import time |
| 11 | +import threading |
| 12 | +import subprocess |
| 13 | +import logging |
| 14 | +from glob import glob |
| 15 | +from pathlib import Path |
| 16 | +from multiprocessing import cpu_count |
| 17 | +import tempfile |
| 18 | +import argparse |
| 19 | +from typing import List, Optional, Any |
| 20 | +from types import FrameType |
| 21 | + |
| 22 | +# ---------- defaults from env ---------- |
| 23 | +SERVER_ARGS: str = os.environ.get("SERVER_ARGS", "") |
| 24 | +WORKER_ARGS: str = os.environ.get("WORKER_ARGS", "") |
| 25 | +SERVER_ADDR: str = os.environ.get("SERVER_ADDR", "127.0.0.1:27183") |
| 26 | +SPAWN_SERVER: bool = os.environ.get("SPAWN_SERVER", "true").lower() == "true" |
| 27 | +SPAWN_WORKERS: bool = os.environ.get("SPAWN_WORKERS", "true").lower() == "true" |
| 28 | +TM_LOGLEVEL: str = os.environ.get("TM_LOGLEVEL", "info").lower() |
| 29 | + |
| 30 | +# default: nproc - 1 (at least 1) |
| 31 | +default_nworkers: int = max(cpu_count() - 1, 1) |
| 32 | + |
| 33 | +# ---------- logging ---------- |
| 34 | +LOGGER_NAME = "tm_entrypoint" |
| 35 | +logger = logging.getLogger(LOGGER_NAME) |
| 36 | + |
| 37 | +def _map_tm_loglevel(level: str) -> int: |
| 38 | + table = { |
| 39 | + "error": logging.ERROR, |
| 40 | + "warn": logging.WARNING, |
| 41 | + "warning": logging.WARNING, |
| 42 | + "info": logging.INFO, |
| 43 | + "debug": logging.DEBUG, |
| 44 | + } |
| 45 | + return table.get(level, logging.ERROR) |
| 46 | + |
| 47 | +def configure_logging() -> None: |
| 48 | + logging.basicConfig( |
| 49 | + level=_map_tm_loglevel(TM_LOGLEVEL), |
| 50 | + format="[%(asctime)s %(levelname)s\t%(name)s::%(funcName)s] %(message)s", |
| 51 | + datefmt= "%Y-%m-%dT%H:%M:%S.xxxxxxxxxZ", |
| 52 | + stream=sys.stderr, |
| 53 | + ) |
| 54 | + eff = logging.getLevelName(logger.getEffectiveLevel()) |
| 55 | + logger.debug("Logger initialized at DEBUG (effective: %s).", eff) |
| 56 | + logger.info("Logging configured (level=%s).", eff) |
| 57 | + |
| 58 | +# ---------- CLI ---------- |
| 59 | +def cli() -> argparse.Namespace: |
| 60 | + parser = argparse.ArgumentParser( |
| 61 | + description="Entrypoint for task-maker-rust Docker image." |
| 62 | + ) |
| 63 | + parser.add_argument( |
| 64 | + "-j", "--jobs", |
| 65 | + type=int, |
| 66 | + default=default_nworkers, |
| 67 | + help=f"Number of workers to launch [default: <nproc>-1 = {default_nworkers}]" |
| 68 | + ) |
| 69 | + |
| 70 | + args = parser.parse_args() |
| 71 | + assert args.jobs > 0, f"Please specify a positive number of jobs, not {args.jobs}." |
| 72 | + |
| 73 | + return args |
| 74 | + |
| 75 | +# ---------- cleanup ---------- |
| 76 | +if SPAWN_SERVER or SPAWN_WORKERS: |
| 77 | + def cleanup() -> None: |
| 78 | + logger.debug("Cleaning up temporary stores") |
| 79 | + for p in glob("/tmp/tmserver.*"): |
| 80 | + logger.debug("Removing server store: %s", p) |
| 81 | + shutil.rmtree(p, ignore_errors=True) |
| 82 | + for p in glob("/tmp/tmworker.*"): |
| 83 | + logger.debug("Removing worker store: %s", p) |
| 84 | + shutil.rmtree(p, ignore_errors=True) |
| 85 | + logger.debug("Cleanup complete.") |
| 86 | + |
| 87 | + atexit.register(cleanup) |
| 88 | + |
| 89 | + def _signal_handler(signum: int, frame: Optional[FrameType]) -> None: |
| 90 | + name = signal.Signals(signum).name if hasattr(signal, "Signals") else str(signum) |
| 91 | + logger.warning("Received signal %s — shutting down.", name) |
| 92 | + # atexit will run cleanup |
| 93 | + code = 130 if signum == signal.SIGINT else 143 |
| 94 | + sys.exit(code) |
| 95 | + |
| 96 | + signal.signal(signal.SIGINT, _signal_handler) |
| 97 | + signal.signal(signal.SIGTERM, _signal_handler) |
| 98 | + |
| 99 | +# ---------- verbosity mapping for task-maker-tools ---------- |
| 100 | +def loglevel_verbosity_flag(level: str) -> str: |
| 101 | + flag = { |
| 102 | + "error": "", |
| 103 | + "warn": "-v", |
| 104 | + "warning": "-v", |
| 105 | + "info": "-vv", |
| 106 | + "debug": "-vvv", |
| 107 | + }.get(level, "") |
| 108 | + logger.debug(f"Mapped TM_LOGLEVEL={level} to " |
| 109 | + f"task-maker-tools flag '{flag or "<none>"}'") |
| 110 | + return flag |
| 111 | + |
| 112 | +# ---------- spawn helpers ---------- |
| 113 | +def make_worker_stores(nworkers: int) -> List[str]: |
| 114 | + base = tempfile.mktemp(prefix="tmworker.", dir="/tmp") |
| 115 | + stores: List[str] = [] |
| 116 | + for i in range(1, nworkers + 1): |
| 117 | + idx = f"{i:02d}" |
| 118 | + d = f"{base}-{idx}" |
| 119 | + Path(d).mkdir(parents=True, exist_ok=True) |
| 120 | + stores.append(d) |
| 121 | + logger.debug("Created %d worker stores (base=%s): %s", nworkers, base, stores) |
| 122 | + return stores |
| 123 | + |
| 124 | +def spawn_tmserver() -> int: |
| 125 | + verbosity_flag: str = loglevel_verbosity_flag(TM_LOGLEVEL) |
| 126 | + def make_server_store() -> Path: |
| 127 | + path = tempfile.mkdtemp(prefix="tmserver.", dir="/tmp") |
| 128 | + logger.debug("Created server store: %s", path) |
| 129 | + return path |
| 130 | + |
| 131 | + server_store: str = make_server_store() |
| 132 | + cmd: List[str] = ["task-maker-tools"] |
| 133 | + if verbosity_flag: |
| 134 | + cmd.append(verbosity_flag) |
| 135 | + cmd += ["server", "--store-dir", server_store] |
| 136 | + if SERVER_ARGS.strip(): |
| 137 | + cmd += shlex.split(SERVER_ARGS) |
| 138 | + |
| 139 | + logger.debug("Exec: %s", " ".join(shlex.quote(c) for c in cmd)) |
| 140 | + logger.info("Starting server (store=%s)…", server_store) |
| 141 | + rc = subprocess.call(cmd) |
| 142 | + logger.info("Server exited with rc=%d", rc) |
| 143 | + return rc |
| 144 | + |
| 145 | +def spawn_tmworker(store_dir: str) -> subprocess.Popen[Any]: |
| 146 | + verbosity_flag: str = loglevel_verbosity_flag(TM_LOGLEVEL) |
| 147 | + |
| 148 | + cmd: List[str] = ["task-maker-tools"] |
| 149 | + if verbosity_flag: |
| 150 | + cmd.append(verbosity_flag) |
| 151 | + cmd += ["worker", "--store-dir", store_dir] |
| 152 | + if WORKER_ARGS.strip(): |
| 153 | + cmd += shlex.split(WORKER_ARGS) |
| 154 | + cmd.append(SERVER_ADDR) |
| 155 | + |
| 156 | + logger.debug("Exec: %s", " ".join(shlex.quote(c) for c in cmd)) |
| 157 | + logger.info("Starting worker (store=%s) → %s", store_dir, SERVER_ADDR) |
| 158 | + proc = subprocess.Popen(cmd) |
| 159 | + logger.debug("Worker PID %s started for store %s", proc.pid, store_dir) |
| 160 | + return proc |
| 161 | + |
| 162 | +# ---------- main ---------- |
| 163 | +def main() -> int: |
| 164 | + configure_logging() |
| 165 | + args = cli() |
| 166 | + nworkers: int = int(args.jobs) |
| 167 | + |
| 168 | + logger.info( |
| 169 | + f"Config: jobs={nworkers}, server={SPAWN_SERVER}, worker={SPAWN_WORKERS}, " |
| 170 | + f"addr={SERVER_ADDR}, tm_loglevel={TM_LOGLEVEL}" |
| 171 | + ) |
| 172 | + if SERVER_ARGS.strip(): |
| 173 | + logger.debug("SERVER_ARGS=%r", SERVER_ARGS) |
| 174 | + if WORKER_ARGS.strip(): |
| 175 | + logger.debug("WORKER_ARGS=%r", WORKER_ARGS) |
| 176 | + |
| 177 | + if SPAWN_WORKERS: |
| 178 | + worker_stores: List[str] = make_worker_stores(nworkers) |
| 179 | + |
| 180 | + # worker only |
| 181 | + if (not SPAWN_SERVER) and SPAWN_WORKERS: |
| 182 | + logger.info("Mode: workers only.") |
| 183 | + procs: List[subprocess.Popen[Any]] = [spawn_tmworker(store_dir=w) |
| 184 | + for w in worker_stores] |
| 185 | + exit_codes: List[Optional[int]] = [p.wait() for p in procs] |
| 186 | + max_rc = max(code or 0 for code in exit_codes) |
| 187 | + logger.info("All workers exited. Max rc=%d", max_rc) |
| 188 | + return max_rc |
| 189 | + |
| 190 | + # server only |
| 191 | + elif SPAWN_SERVER and (not SPAWN_WORKERS): |
| 192 | + logger.info("Mode: server only.") |
| 193 | + return spawn_tmserver() |
| 194 | + |
| 195 | + # server + worker |
| 196 | + elif SPAWN_SERVER and SPAWN_WORKERS: |
| 197 | + logger.info("Mode: server + workers (delayed worker start).") |
| 198 | + procs: List[subprocess.Popen[Any]] = [] |
| 199 | + |
| 200 | + def launch_workers() -> None: |
| 201 | + logger.debug("Delaying worker launch by 2s to let server come up…") |
| 202 | + time.sleep(2.0) |
| 203 | + for w in worker_stores: |
| 204 | + procs.append(spawn_tmworker(store_dir=w)) |
| 205 | + for p in procs: |
| 206 | + rc = p.wait() |
| 207 | + logger.info("Worker PID %s exited with rc=%d", p.pid, rc) |
| 208 | + |
| 209 | + t = threading.Thread(target=launch_workers, daemon=True) |
| 210 | + t.start() |
| 211 | + server_rc: int = spawn_tmserver() |
| 212 | + |
| 213 | + logger.debug("Server finished (rc=%d). Joining worker launcher…", server_rc) |
| 214 | + t.join(timeout=1.0) |
| 215 | + # If any workers still running, terminate politely |
| 216 | + for p in procs: |
| 217 | + if p.poll() is None: |
| 218 | + logger.debug("Terminating lingering worker PID %s…", p.pid) |
| 219 | + try: |
| 220 | + p.terminate() |
| 221 | + except Exception as e: |
| 222 | + logger.warning("Failed to terminate worker PID %s: %s", p.pid, e) |
| 223 | + return server_rc |
| 224 | + |
| 225 | + # nothing to spawn -> shell |
| 226 | + else: |
| 227 | + logger.info("Mode: nothing to spawn — opening interactive shell.") |
| 228 | + shell: str = os.environ.get("SHELL", "/bin/bash") |
| 229 | + try: |
| 230 | + return subprocess.call([shell]) |
| 231 | + except FileNotFoundError: |
| 232 | + logger.error("Shell %r not found; exiting 0.", shell) |
| 233 | + return 0 |
| 234 | + |
| 235 | +if __name__ == "__main__": |
| 236 | + try: |
| 237 | + sys.exit(main()) |
| 238 | + except subprocess.CalledProcessError as e: |
| 239 | + logger.exception("Subprocess failed (rc=%s).", e.returncode) |
| 240 | + sys.exit(e.returncode) |
| 241 | + except Exception: |
| 242 | + logger.exception("Fatal error:") |
| 243 | + sys.exit(1) |
0 commit comments