sdg/launch_mongodb.py

"""
Standalone script that starts a MongoDB server locally and prints the mongo-uri.

Uses ``mongod`` with a configurable port and dbpath. Runs indefinitely
(foreground process) so that worker shards can connect to it.

Usage:
    python -m sdg.launch_mongodb
    python -m sdg.launch_mongodb --port 27018 --dbpath /tmp/sdg_mongo

nlprun command (john queue, no GPU, 16 GB RAM):
    nlprun -a tonyreasoningtraces -q john --exclude john17 -c 4 -g 0 --memory 64g \
      -w /nlp/scr4/nlp/crfm/text2image/text2image-rlhf/reasoning/virtual-world-data \
      --job-name sdg-mongodb "python -m sdg.launch_mongodb 2>&1 | tee mongo.log"
"""

import argparse
import shutil
import socket
import subprocess
import sys
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser(description="Start a MongoDB server for SDG cache")
    parser.add_argument("--port", type=int, default=27017, help="MongoDB port (default: 27017)")
    parser.add_argument("--dbpath", default="/nlp/scr4/nlp/crfm/text2image/text2image-rlhf/reasoning/virtual-world-data/cache", help="Data directory")
    args = parser.parse_args()

    # Ensure mongod is available
    mongod = shutil.which("mongod")
    if mongod is None:
        print("ERROR: mongod not found on PATH. See README.md for install instructions.", file=sys.stderr)
        sys.exit(1)

    # Create dbpath if needed
    Path(args.dbpath).mkdir(parents=True, exist_ok=True)

    hostname = socket.gethostname()
    uri = f"mongodb://{hostname}:{args.port}"
    print(f"mongo-uri: {uri}", flush=True)

    # Launch mongod in foreground (replaces this process on Unix)
    cmd = [mongod, "--port", str(args.port), "--dbpath", args.dbpath, "--bind_ip_all"]
    print(f"Running: {' '.join(cmd)}", flush=True)
    proc = subprocess.run(cmd)
    sys.exit(proc.returncode)


if __name__ == "__main__":
    main()