#!/usr/bin/env python3 from __future__ import annotations import csv import hashlib import json import linecache import os import subprocess import sys import threading import time import traceback import urllib.parse import urllib.request from pathlib import Path from typing import Any Tensor: Any = None dtypes: Any = None model: Any = None tokenizer: Any = None role: dict[str, Any] = {} loaded: dict[str, Any] = {} _nvidia_smi_available: bool | None = None class CpuLineSampler: def __init__( self, *, phase: str, request_id: int | None, model_id: str | None, interval_secs: float, ) -> None: self.phase = phase self.request_id = request_id self.model_id = model_id self.interval_secs = interval_secs self.target_thread_id = threading.get_ident() self.samples: dict[tuple[str, int, str], int] = {} self.wall_start = time.perf_counter() self.process_cpu_start = time.process_time() self._running = True self._thread = threading.Thread(target=self._run, name="cpu-line-sampler", daemon=True) self._thread.start() def _run(self) -> None: while self._running: frame = sys._current_frames().get(self.target_thread_id) if frame is not None: code = frame.f_code key = (code.co_filename, frame.f_lineno, code.co_name) self.samples[key] = self.samples.get(key, 0) + 1 time.sleep(self.interval_secs) def stop(self) -> None: self._running = False self._thread.join(timeout=max(0.25, self.interval_secs * 4.0)) wall_elapsed_ms = (time.perf_counter() - self.wall_start) * 1000.0 process_cpu_elapsed_ms = (time.process_time() - self.process_cpu_start) * 1000.0 total_samples = sum(self.samples.values()) top = [] for (filename, line, function), count in sorted( self.samples.items(), key=lambda item: item[1], reverse=True )[:32]: top.append( { "file": filename, "line": line, "function": function, "source": linecache.getline(filename, line).strip(), "samples": count, "percent": round((count * 100.0 / total_samples), 2) if total_samples else 0.0, } ) control( type="CpuLineProfileSummary", phase=self.phase, request_id=self.request_id, model_id=self.model_id, interval_ms=round(self.interval_secs * 1000.0, 3), wall_elapsed_ms=round(wall_elapsed_ms, 3), process_cpu_elapsed_ms=round(process_cpu_elapsed_ms, 3), process_cpu_over_wall=round(process_cpu_elapsed_ms / wall_elapsed_ms, 4) if wall_elapsed_ms > 0.0 else 0.0, total_samples=total_samples, top=top, ) def start_cpu_line_sampler( *, phase: str, request_id: int | None, model_id: str | None, ) -> CpuLineSampler | None: raw = os.environ.get("MVP_CPU_LINE_PROFILE") if not env_flag("MVP_CPU_LINE_PROFILE", False): control( type="CpuLineProfileSkipped", phase=phase, request_id=request_id, model_id=model_id, env_value=raw, ) return None interval_ms = float(os.environ.get("MVP_CPU_LINE_PROFILE_INTERVAL_MS", "2")) interval_secs = max(0.0005, interval_ms / 1000.0) control( type="CpuLineProfileStarted", phase=phase, request_id=request_id, model_id=model_id, interval_ms=round(interval_secs * 1000.0, 3), ) return CpuLineSampler( phase=phase, request_id=request_id, model_id=model_id, interval_secs=interval_secs, ) def stop_cpu_line_sampler(sampler: CpuLineSampler | None) -> None: if sampler is not None: sampler.stop() def env_flag(name: str, default: bool = True) -> bool: raw = os.environ.get(name) if raw is None: return default return raw.strip().lower() not in {"0", "false", "no", "off"} def parse_int(value: str) -> int | None: stripped = value.strip() if not stripped or stripped == "[Not Supported]": return None try: return int(float(stripped)) except ValueError: return None def run_nvidia_smi(args: list[str]) -> tuple[bool, str, str]: global _nvidia_smi_available if _nvidia_smi_available is False: return False, "", "nvidia-smi unavailable" try: result = subprocess.run( ["nvidia-smi", *args], check=False, capture_output=True, text=True, timeout=float(os.environ.get("MVP_GPU_SAMPLE_TIMEOUT_SECS", "2")), ) except FileNotFoundError: _nvidia_smi_available = False return False, "", "nvidia-smi not found" except Exception as exc: return False, "", str(exc) _nvidia_smi_available = True return result.returncode == 0, result.stdout, result.stderr.strip() def parse_gpu_rows(raw: str) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] for row in csv.reader(raw.splitlines()): if len(row) < 5: continue memory_total = parse_int(row[2]) memory_used = parse_int(row[3]) utilization = parse_int(row[4]) rows.append( { "index": parse_int(row[0]), "name": row[1].strip(), "memory_total_mib": memory_total, "memory_used_mib": memory_used, "utilization_gpu_percent": utilization, } ) return rows def parse_process_rows(raw: str) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] for row in csv.reader(raw.splitlines()): if len(row) < 2: continue pid = parse_int(row[0]) memory_used = parse_int(row[1]) if pid is None: continue rows.append({"pid": pid, "used_memory_mib": memory_used}) return rows def gpu_sample(label: str, **fields: Any) -> None: if not env_flag("MVP_GPU_SAMPLE", True): control(type="GpuSample", label=label, pid=os.getpid(), enabled=False, **fields) return gpu_ok, gpu_stdout, gpu_error = run_nvidia_smi( [ "--query-gpu=index,name,memory.total,memory.used,utilization.gpu", "--format=csv,noheader,nounits", ] ) proc_ok, proc_stdout, proc_error = run_nvidia_smi( [ "--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits", ] ) pid = os.getpid() processes = parse_process_rows(proc_stdout) if proc_ok else [] worker_processes = [process for process in processes if process.get("pid") == pid] control( type="GpuSample", label=label, pid=pid, enabled=True, nvidia_smi_available=(_nvidia_smi_available is True), gpu_query_ok=gpu_ok, process_query_ok=proc_ok, gpu_query_error=None if gpu_ok else gpu_error, process_query_error=None if proc_ok else proc_error, gpus=parse_gpu_rows(gpu_stdout) if gpu_ok else [], processes=processes, worker_processes=worker_processes, **fields, ) def control(**event: Any) -> None: print(json.dumps(event, separators=(",", ":")), flush=True) def log(message: str) -> None: print(f"mvp_tinygrad_worker: {message}", file=sys.stderr, flush=True) def fatal(reason: str, **fields: Any) -> None: control(type="WorkerFatal", reason=reason, **fields) raise SystemExit(1) def test_mode() -> bool: return os.environ.get("MVP_TINYGRAD_TEST_MODE", "").strip().lower() in {"1", "true", "yes", "on"} def initialize(cmd: dict[str, Any]) -> None: global Tensor, dtypes if int(cmd.get("helper_abi_version", 1)) != 1: fatal("UnsupportedHelperAbi", helper_abi_version=cmd.get("helper_abi_version")) device = str(cmd.get("backend", {}).get("device") or os.environ.get("DEV") or "CUDA") os.environ["DEV"] = device started = time.monotonic() control(type="TinygradImportStarted", requested_device=device, env_DEV=os.environ.get("DEV")) from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes control(type="TinygradImportReady", requested_device=device, env_DEV=os.environ.get("DEV")) Tensor = TinyTensor dtypes = tiny_dtypes control(type="TinygradDeviceProbeStarted", requested_device=device) value = Tensor([1], dtype=dtypes.int32).realize().numpy().tolist() control(type="TinygradDeviceProbeReady", requested_device=device, probe_result=value) gpu_sample("after_worker_probe", requested_device=device) control( type="WorkerReady", pid=os.getpid(), backend={"requested_device": device, "env_DEV": os.environ.get("DEV"), "tinygrad_device": device}, cuda_probe=value, elapsed_ms=int((time.monotonic() - started) * 1000), ) def configure_role(cmd: dict[str, Any]) -> None: config = cmd.get("config", {}) role.clear() role.update( role_id=int(cmd.get("role_id", 1)), run_id=int(config.get("run_id", 1)), stage_index=int(config.get("stage_index", 0)), layer_start=int(config.get("layer_start", 0)), layer_end_exclusive=int(config.get("layer_end_exclusive", 0)), ) control(type="RoleConfigured", role_id=role["role_id"], stage_index=role["stage_index"]) def cache_root() -> Path: raw = os.environ.get("MVP_MODEL_CACHE_DIR", "").strip() root = Path(raw).expanduser() if raw else Path.home() / ".cache" / "mvp-node" root.mkdir(parents=True, exist_ok=True) return root def hf_url(repo: str, file: str, revision: str | None) -> str: encoded_file = "/".join(urllib.parse.quote(part) for part in file.split("/")) return f"https://huggingface.co/{repo}/resolve/{revision or 'main'}/{encoded_file}" def source_url(source: dict[str, Any]) -> str | None: if "HuggingFaceGguf" not in source: return None hf = source["HuggingFaceGguf"] return hf_url(str(hf["repo"]), str(hf["file"]), hf.get("revision")) def source_path(source: dict[str, Any]) -> Path | None: if "LocalPath" not in source: return None return Path(str(source["LocalPath"])).expanduser() def source_kind(source: dict[str, Any]) -> str: if "LocalPath" in source: return "LocalPath" if "HuggingFaceGguf" in source: return "HuggingFaceGguf" return "Unknown" def cache_path_for(url: str) -> Path: parsed = urllib.parse.urlparse(url) basename = Path(parsed.path).name or "model.gguf" digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:16] return cache_root() / f"{digest}-{basename}" def request_headers() -> dict[str, str]: headers = {"User-Agent": "swactor-mvp-node/0.1"} token = os.environ.get("HF_TOKEN", "").strip() if token: headers["Authorization"] = f"Bearer {token}" return headers def fetch_whole(source: dict[str, Any]) -> Path: local = source_path(source) if local is not None: control(type="GgufLocalPathStatStarted", path=str(local)) if not local.is_file(): fatal("GgufLocalPathMissing", path=str(local)) stat = local.stat() control(type="GgufCacheReady", path=str(local), bytes=stat.st_size, cache_hit=True, source="local") return local url = source_url(source) if not url: fatal("UnsupportedGgufSource", source=source) target = cache_path_for(url) if target.is_file() and target.stat().st_size > 0: control(type="GgufCacheReady", path=str(target), bytes=target.stat().st_size, cache_hit=True, url=url) return target partial = target.with_name(target.name + ".partial") started = time.monotonic() req = urllib.request.Request(url, headers=request_headers()) control(type="GgufDownloadStarted", url=url, path=str(target)) try: with urllib.request.urlopen(req, timeout=60) as response, partial.open("wb") as out: total = int(response.headers.get("Content-Length") or 0) done = 0 last_event = 0.0 while True: chunk = response.read(1024 * 1024) if not chunk: break out.write(chunk) done += len(chunk) now = time.monotonic() if now - last_event >= float(os.environ.get("MVP_DOWNLOAD_PROGRESS_SECS", "5")): control( type="GgufDownloadProgress", bytes_done=done, bytes_total=total, elapsed_ms=int((now - started) * 1000), ) last_event = now partial.replace(target) except Exception as exc: try: partial.unlink(missing_ok=True) except Exception: pass fatal("GgufDownloadFailed", url=url, error=str(exc)) control( type="GgufCacheReady", path=str(target), bytes=target.stat().st_size, cache_hit=False, elapsed_ms=int((time.monotonic() - started) * 1000), url=url, ) return target def require_tinygrad() -> Any: if Tensor is None: fatal("BackendNotInitialized") return Tensor def load_weights(cmd: dict[str, Any]) -> None: global model, tokenizer TensorCls = require_tinygrad() started = time.monotonic() model_id = str(cmd["model_id"]) source = cmd["gguf_source"] control( type="LoadWeightsStarted", model_id=model_id, source_kind=source_kind(source), layer_start=int(cmd.get("layer_start", 0)), layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)), ) if test_mode(): model = {"test_mode": True} tokenizer = {"test_mode": True} loaded.clear() loaded.update( model_id=model_id, path="mvp-tinygrad-test-mode", layer_start=int(cmd.get("layer_start", 0)), layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)), ) control( type="WeightsLoaded", model_id=model_id, path=loaded["path"], test_mode=True, elapsed_ms=int((time.monotonic() - started) * 1000), ) return control(type="GgufResolveStarted", model_id=model_id, source_kind=source_kind(source)) path = fetch_whole(source) model_bytes = path.stat().st_size control(type="GgufResolveReady", model_id=model_id, path=str(path), bytes=model_bytes) try: control(type="TinygradLlmImportStarted", model_id=model_id) from tinygrad.apps.llm import SimpleTokenizer, Transformer control(type="TinygradLlmImportReady", model_id=model_id) max_context_raw = os.environ.get("MVP_MAX_CONTEXT", "512") max_context = int(max_context_raw) if max_context_raw else 512 gpu_sample("before_model_load", model_id=model_id, path=str(path), bytes=model_bytes) control( type="TransformerFromGgufStarted", model_id=model_id, path=str(path), bytes=model_bytes, max_context=max_context, realize=True, requested_device=os.environ.get("DEV"), ) model, kv = Transformer.from_gguf(TensorCls(path), max_context=max_context, realize=True) control( type="TransformerFromGgufReady", model_id=model_id, path=str(path), bytes=model_bytes, max_context=max_context, realize=True, requested_device=os.environ.get("DEV"), ) gpu_sample("after_model_load", model_id=model_id, path=str(path), bytes=model_bytes) tok_src = cmd.get("tokenizer", {"EmbeddedGguf": None}) if "EmbeddedGguf" in tok_src: control(type="TokenizerBuildStarted", model_id=model_id, source="EmbeddedGguf") tokenizer = SimpleTokenizer.from_gguf_kv(kv) control(type="TokenizerBuildReady", model_id=model_id, source="EmbeddedGguf") else: fatal("UnsupportedTokenizerSource", tokenizer=tok_src) except SystemExit: raise except Exception as exc: tb = traceback.format_exc() print(tb, file=sys.stderr, flush=True) fatal("ModelLoadFailed", error=str(exc), traceback=tb) loaded.clear() loaded.update( model_id=model_id, path=str(path), layer_start=int(cmd.get("layer_start", 0)), layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)), ) control( type="WeightsLoaded", model_id=model_id, path=str(path), elapsed_ms=int((time.monotonic() - started) * 1000), ) def decode_greedy_device_resident( prompt_tokens: list[int], max_tokens: int, *, request_id: int | None, model_id: str | None, progress_every: int, ) -> list[int]: if max_tokens <= 0: return [] max_context = int(getattr(model, "max_context", len(prompt_tokens) + max_tokens)) generation_limit = min(max_tokens, max(0, max_context - len(prompt_tokens))) if generation_limit <= 0: control( type="DecodeContextFull", request_id=request_id, model_id=model_id, prompt_tokens=len(prompt_tokens), max_context=max_context, ) return [] if generation_limit < max_tokens: control( type="DecodeLimitedByContext", request_id=request_id, model_id=model_id, prompt_tokens=len(prompt_tokens), requested_tokens=max_tokens, generation_limit=generation_limit, max_context=max_context, ) TensorCls = require_tinygrad() from tinygrad.uop.ops import UOp if hasattr(model, "forward_jit"): model.forward_jit.reset() use_symbolic_pos = os.environ.get("SYM", "1").strip().lower() not in {"0", "false", "no", "off"} pos_upper_bound = max(1, max_context - 1) symbolic_start_pos = UOp.variable("start_pos", 1, pos_upper_bound) next_token = model(TensorCls([prompt_tokens], dtype="int32"), 0).realize() generated_tensors = [] for token_index in range(generation_limit): generated_tensors.append(next_token.clone().realize()) tokens_generated = token_index + 1 if tokens_generated == 1: control( type="FirstTokenReady", request_id=request_id, model_id=model_id, token_index=1, prompt_tokens=len(prompt_tokens), ) gpu_sample("after_first_token", request_id=request_id, model_id=model_id) elif progress_every > 0 and tokens_generated % progress_every == 0: control( type="TokenProgress", request_id=request_id, model_id=model_id, tokens_generated=tokens_generated, prompt_tokens=len(prompt_tokens), ) if tokens_generated >= generation_limit: break start_pos = len(prompt_tokens) + token_index pos = symbolic_start_pos.bind(start_pos) if use_symbolic_pos else start_pos next_token = model(next_token, pos).realize() generated_tensor = ( generated_tensors[0] if len(generated_tensors) == 1 else generated_tensors[0].cat(*generated_tensors[1:], dim=1) ) generated_array = generated_tensor.numpy().reshape(-1).tolist() return [int(token) for token in generated_array] def infer_prompt(cmd: dict[str, Any]) -> None: if model is None or tokenizer is None: fatal("WeightsNotLoaded") prompt = str(cmd.get("prompt", "")) max_tokens = int(cmd.get("max_tokens", 1)) request_id_raw = cmd.get("request_id") request_id = int(request_id_raw) if request_id_raw is not None else None started = time.monotonic() control( type="PromptStarted", request_id=request_id, model_id=loaded.get("model_id"), prompt_bytes=len(prompt.encode("utf-8")), prompt_chars=len(prompt), max_tokens=max_tokens, ) gpu_sample("before_prompt", request_id=request_id, model_id=loaded.get("model_id")) if test_mode(): text = f"mvp-test response: {prompt}" control( type="PromptCompleted", request_id=request_id, model_id=loaded.get("model_id"), prompt_tokens=[], generated_tokens=list(range(min(max_tokens, 3))), text=text, test_mode=True, elapsed_ms=int((time.monotonic() - started) * 1000), ) return control(type="PromptEncodeStarted", request_id=request_id, model_id=loaded.get("model_id")) prompt_tokens = tokenizer.encode(prompt) control( type="PromptEncodeReady", request_id=request_id, model_id=loaded.get("model_id"), prompt_bytes=len(prompt.encode("utf-8")), prompt_tokens=len(prompt_tokens), ) progress_every = int(os.environ.get("MVP_TOKEN_PROGRESS_EVERY", "16") or "16") control( type="DecodeStarted", request_id=request_id, model_id=loaded.get("model_id"), prompt_tokens=len(prompt_tokens), max_tokens=max_tokens, decode_impl="device_resident_greedy", ) cpu_sampler = start_cpu_line_sampler( phase="decode", request_id=request_id, model_id=loaded.get("model_id"), ) try: generated = decode_greedy_device_resident( prompt_tokens, max_tokens, request_id=request_id, model_id=loaded.get("model_id"), progress_every=progress_every, ) finally: stop_cpu_line_sampler(cpu_sampler) control( type="DecodeReady", request_id=request_id, model_id=loaded.get("model_id"), prompt_tokens=len(prompt_tokens), tokens_generated=len(generated), ) gpu_sample("after_decode", request_id=request_id, model_id=loaded.get("model_id")) control(type="TextDecodeStarted", request_id=request_id, model_id=loaded.get("model_id"), tokens_generated=len(generated)) text = tokenizer.decode(generated) if generated else "" control( type="TextDecodeReady", request_id=request_id, model_id=loaded.get("model_id"), tokens_generated=len(generated), text_bytes=len(text.encode("utf-8")), ) control( type="PromptCompleted", request_id=request_id, model_id=loaded.get("model_id"), prompt_tokens=prompt_tokens, generated_tokens=generated, text=text, elapsed_ms=int((time.monotonic() - started) * 1000), ) def shutdown_worker(_: dict[str, Any]) -> None: control(type="WorkerStopped", reason="Graceful") raise SystemExit(0) HANDLERS = { "InitializeWorker": initialize, "ConfigureRole": configure_role, "LoadWeights": load_weights, "InferPrompt": infer_prompt, "ShutdownWorker": shutdown_worker, } for raw in sys.stdin: if not raw.strip(): continue try: command = json.loads(raw) handler = HANDLERS.get(command.get("type")) if handler is None: fatal("UnknownCommand", command=command.get("type")) handler(command) except SystemExit: raise except Exception as exc: tb = traceback.format_exc() print(tb, file=sys.stderr, flush=True) fatal("UnhandledWorkerException", error=str(exc), traceback=tb)