| |
| """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() |
|
|
