697 lines
23 KiB
Python
Executable file
697 lines
23 KiB
Python
Executable file
#!/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)
|