#!/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). """ 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-") 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()