File size: 13,921 Bytes
750b88d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 | #!/usr/bin/env python3
"""Download, configure and supervise the MemOS backing services.
MEMEBENCH_MEMOS_WORK_ROOT sets the default --work-root; MEMEBENCH_MEMOS_RUN_ID names the run (default: pid<pid>).
"""
from __future__ import annotations
import argparse
from concurrent.futures import ThreadPoolExecutor
import fcntl
import json
import os
from pathlib import Path
import secrets
import shutil
import sys
import signal
import socket
import subprocess
import tarfile
import time
import urllib.request
DEFAULT_WORK = Path(os.environ.get("MEMEBENCH_MEMOS_WORK_ROOT", Path(__file__).resolve().parent / "work"))
ARTIFACTS = {
"java21": "https://corretto.aws/downloads/resources/21.0.12.9.1/amazon-corretto-21.0.12.9.1-linux-x64.tar.gz",
"neo4j": "https://dist.neo4j.org/neo4j-community-5.26.12-unix.tar.gz",
"qdrant": "https://github.com/qdrant/qdrant/releases/download/v1.16.3/qdrant-x86_64-unknown-linux-musl.tar.gz",
}
BGE_REVISION = "953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e"
LOCAL_FILESYSTEMS = {"xfs", "ext2/ext3", "ext4", "btrfs", "zfs", "tmpfs", "overlayfs"}
def atomic_json(path, value):
temporary = path.with_suffix(path.suffix + ".partial")
temporary.write_text(json.dumps(value, indent=2) + "\n")
temporary.replace(path)
def directories(work):
for name in ("runtime/downloads", "runtime/services", "runtime/logs", "runtime/run", "models", "config",
"validation", "caches/hf", "caches/nltk", "caches/tiktoken", "caches/torch", "caches/triton"):
(work / name).mkdir(parents=True, exist_ok=True)
def download_one(work, name):
url = ARTIFACTS[name]
archive = work / "runtime/downloads" / (name + ".tar.gz")
if not archive.exists():
temporary = archive.with_suffix(".partial")
with urllib.request.urlopen(url, timeout=60) as response, temporary.open("wb") as out:
shutil.copyfileobj(response, out, 1024 * 1024)
temporary.replace(archive)
target = work / "runtime/services" / name
if not target.exists():
staging = target.with_name(name + ".extracting")
staging.mkdir(exist_ok=False)
with tarfile.open(archive) as contents:
contents.extractall(staging, filter="data")
children = list(staging.iterdir())
if len(children) == 1 and children[0].is_dir():
children[0].rename(target)
staging.rmdir()
else:
staging.rename(target)
print(f"Prepared {name}", flush=True)
return name, {"url": url, "archive": str(archive), "path": str(target)}
def download(work):
with ThreadPoolExecutor(max_workers=3) as pool:
results = dict(pool.map(lambda name: download_one(work, name), ARTIFACTS))
atomic_json(work / "validation/service_artifacts.json", results)
def model(work):
os.environ["HF_HUB_DISABLE_XET"] = "1"
from huggingface_hub import snapshot_download
result = snapshot_download("BAAI/bge-reranker-v2-m3", revision=BGE_REVISION,
cache_dir=str(work / "models/hf"), token=False,
allow_patterns=["config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json",
"sentencepiece.bpe.model", "special_tokens_map.json"])
files = []
for path in sorted(Path(result).iterdir()):
if path.is_file():
files.append({"name": path.name, "bytes": path.stat().st_size})
config = {"repo_id": "BAAI/bge-reranker-v2-m3", "revision": BGE_REVISION,
"path": result, "score_semantics": "sigmoid(logit)", "device": "cpu", "files": files}
atomic_json(work / "config/bge_model.json", config)
print(json.dumps({"model": config["repo_id"], "revision": BGE_REVISION, "path": result}), flush=True)
def local_filesystem(path, allowed=()):
"""Return the filesystem type of `path`; raise unless it is a known local filesystem or allowed."""
existing = path
while not existing.exists():
if existing == existing.parent:
raise RuntimeError("Cannot locate active-state filesystem")
existing = existing.parent
result = subprocess.run(["stat", "-f", "-c", "%T", str(existing)],
text=True, capture_output=True, check=True)
filesystem = result.stdout.strip().lower()
if filesystem not in LOCAL_FILESYSTEMS | set(allowed):
raise RuntimeError(f"Active database state requires a local filesystem; found {filesystem!r}. "
f"Pass --local-state-root on a local disk, or --allow-filesystem {filesystem}")
return filesystem
def prepare_local_state(local_state_root=None, allowed=()):
run_id = os.environ.get("MEMEBENCH_MEMOS_RUN_ID") or f"pid{os.getpid()}"
root = Path(local_state_root or (Path(os.environ.get("TMPDIR", "/tmp")) / ("memos-" + run_id))).resolve()
filesystem = local_filesystem(root, allowed)
if root.exists() and root.stat().st_uid != os.getuid():
raise RuntimeError("Active database directory belongs to another user")
owner = {"schema": "memos-local-state", "run_id": run_id}
marker = root / "owner.json"
if marker.exists():
if json.loads(marker.read_text()) != owner:
raise RuntimeError("Active database state belongs to a different run; choose a fresh local root")
elif root.exists() and any(root.iterdir()):
raise RuntimeError("Refusing to adopt a nonempty local directory without the owner.json this script writes")
else:
root.mkdir(parents=True, mode=0o700, exist_ok=True)
atomic_json(marker, owner)
root.chmod(0o700)
for directory in ("neo4j/data", "neo4j/transactions", "neo4j/run", "qdrant/storage", "qdrant/snapshots"):
(root / directory).mkdir(parents=True, mode=0o700, exist_ok=True)
return root, filesystem, run_id
def configure(work, local_state_root=None, allowed=()):
runtime = work / "runtime"
neo = runtime / "services/neo4j"
local_state, filesystem, run_id = prepare_local_state(local_state_root, allowed)
password_file = work / "config/neo4j_password.txt"
if not password_file.exists():
with os.fdopen(os.open(password_file, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "w") as stream:
stream.write(secrets.token_urlsafe(32) + "\n")
(neo / "conf/neo4j.conf").write_text(f"""server.default_listen_address=127.0.0.1
server.directories.data={local_state}/neo4j/data
server.directories.transaction.logs.root={local_state}/neo4j/transactions
server.directories.run={local_state}/neo4j/run
server.directories.logs={runtime}/services/neo4j/logs
server.bolt.enabled=true
server.bolt.listen_address=127.0.0.1:17687
server.bolt.advertised_address=127.0.0.1:17687
server.http.enabled=true
server.http.listen_address=127.0.0.1:17474
server.http.advertised_address=127.0.0.1:17474
server.https.enabled=false
server.memory.heap.initial_size=1G
server.memory.heap.max_size=2G
server.memory.pagecache.size=1G
server.jvm.additional=-XX:ActiveProcessorCount=4
dbms.security.auth_enabled=true
""")
initialized = local_state / "neo4j/password-initialized"
if not initialized.exists():
password = password_file.read_text().strip()
result = subprocess.run([str(neo / "bin/neo4j-admin"), "dbms", "set-initial-password", password],
env=dict(os.environ, JAVA_HOME=str(runtime / "services/java21")), capture_output=True)
if result.returncode:
raise RuntimeError("Neo4j initialization failed; inspect Java/platform compatibility")
initialized.touch()
(work / "config/qdrant.yaml").write_text(f"""storage:
storage_path: {local_state}/qdrant/storage
snapshots_path: {local_state}/qdrant/snapshots
performance:
max_search_threads: 4
service:
host: 127.0.0.1
http_port: 16333
grpc_port: 16334
enable_cors: false
telemetry_disabled: true
""")
config = {"run_id": run_id, "local_state_root": str(local_state), "local_state_filesystem": filesystem,
"active_state_lifetime": "this service run",
"neo4j_uri": "bolt://127.0.0.1:17687", "neo4j_http": "http://127.0.0.1:17474",
"neo4j_user": "neo4j", "neo4j_db_name": "neo4j", "neo4j_password_file": str(password_file),
"qdrant_url": "http://127.0.0.1:16333", "reranker_url": "http://127.0.0.1:18086/rerank",
"reranker_health": "http://127.0.0.1:18086/health", "python": sys.executable,
"neo4j_version": "5.26.12", "qdrant_version": "1.16.3", "bge_revision": BGE_REVISION,
"bge_score": "sigmoid(logit)", "bge_device": "cpu", "bge_max_length": 1024}
atomic_json(work / "config/services.json", config)
return config
def occupied(port):
with socket.socket() as handle:
handle.settimeout(1)
return handle.connect_ex(("127.0.0.1", port)) == 0
def run(work, local_state_root=None, allowed=()):
runtime = work / "runtime"
with (runtime / "run/services.lock").open("a") as lock:
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
for port in (17687, 17474, 16333, 16334, 18086):
if occupied(port):
raise RuntimeError(f"Port {port} already occupied; refusing to interfere")
config = configure(work, local_state_root, allowed)
bge = json.loads((work / "config/bge_model.json").read_text())
if bge["revision"] != BGE_REVISION:
raise RuntimeError("BGE manifest revision mismatch")
common = dict(os.environ, PYTHONDONTWRITEBYTECODE="1", HF_HOME=str(work / "caches/hf"),
HF_HUB_OFFLINE="1", TOKENIZERS_PARALLELISM="false", NLTK_DATA=str(work / "caches/nltk"),
TIKTOKEN_CACHE_DIR=str(work / "caches/tiktoken"), TORCH_HOME=str(work / "caches/torch"))
commands = {
"neo4j": ([str(runtime / "services/neo4j/bin/neo4j"), "console"],
dict(common, JAVA_HOME=str(runtime / "services/java21"))),
"qdrant": ([str(runtime / "services/qdrant/qdrant"), "--config-path", str(work / "config/qdrant.yaml")],
common),
"bge": ([sys.executable, "-B", str(Path(__file__).resolve().parent / "memos_bge_server.py"),
"--model-path", bge["path"], "--threads", "8", "--batch-size", "8", "--max-length", "1024"],
dict(common, CUDA_VISIBLE_DEVICES="", OMP_NUM_THREADS="8", MKL_NUM_THREADS="8")),
}
children = {}
stopping = False
def stop(signum, frame):
nonlocal stopping
stopping = True
signal.signal(signal.SIGTERM, stop)
signal.signal(signal.SIGINT, stop)
ready_file = runtime / "run/services.ready.json"
ready_file.unlink(missing_ok=True)
try:
for name, (command, environment) in commands.items():
with (runtime / "logs" / (name + ".log")).open("ab") as log:
child = subprocess.Popen(command, cwd=runtime, env=environment,
stdin=subprocess.DEVNULL, stdout=log, stderr=subprocess.STDOUT, start_new_session=True)
children[name] = child
(runtime / "run" / (name + ".pid")).write_text(str(child.pid) + "\n")
print(f"Started owned {name} pid={child.pid}", flush=True)
deadline = time.monotonic() + 600
while not stopping:
for name, child in children.items():
if child.poll() is not None:
raise RuntimeError(f"Owned {name} exited ({child.returncode}); see runtime/logs/{name}.log")
if not ready_file.exists():
try:
for url in (config["neo4j_http"], config["qdrant_url"] + "/healthz", config["reranker_health"]):
with urllib.request.urlopen(url, timeout=3) as response:
if response.status != 200:
raise RuntimeError("Service not ready")
except Exception:
if time.monotonic() > deadline:
raise RuntimeError("Services did not become healthy within 600 seconds")
else:
atomic_json(ready_file, dict(config, ready_at=time.time(),
pids={k: v.pid for k, v in children.items()}))
print("All MemOS services healthy; configuration: " + str(work / "config/services.json"),
flush=True)
time.sleep(2)
finally:
ready_file.unlink(missing_ok=True)
for child in children.values():
if child.poll() is None:
os.killpg(child.pid, signal.SIGTERM)
deadline = time.monotonic() + 30
for child in children.values():
try:
child.wait(timeout=max(0.1, deadline - time.monotonic()))
except subprocess.TimeoutExpired:
os.killpg(child.pid, signal.SIGKILL)
child.wait()
def main():
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("action", choices=("download", "model", "run"))
parser.add_argument("--work-root", type=Path, default=DEFAULT_WORK)
parser.add_argument("--local-state-root", type=Path,
help="Run-only local disk directory; default: TMPDIR/memos-<run id>")
parser.add_argument("--allow-filesystem", action="append", default=[],
help="also accept this filesystem type (as printed by stat -f -c %%T) for the database state")
args = parser.parse_args()
work = args.work_root.resolve()
directories(work)
if args.action == "run":
run(work, args.local_state_root, args.allow_filesystem)
else:
{"download": download, "model": model}[args.action](work)
if __name__ == "__main__":
main()
|