#!/usr/bin/env python3
"""Start an EMT experiment on its dedicated rented cluster. Python 3.9+, OpenSSH.

Download emt_compute.py alongside this file. EMT_API_KEY stays local; each remote
node receives only the experiment token. Existing SSH host keys are checked.
"""
import argparse
import json
import re
import shlex
import subprocess
from pathlib import Path
from urllib.parse import quote
from emt_compute import EMTCompute, parse_ssh_connection

# Runs inside the selected scientific image, retaining the cluster host network.
BOOTSTRAP = r'''
import hashlib, io, json, os, sys, tarfile, urllib.request
conf = json.load(open("/workspace/runtime.json"))
request = urllib.request.Request(conf["url"], headers={"User-Agent": "EMT-Compute-Cluster/1.0"})
data = urllib.request.urlopen(request, timeout=120).read(8 * 1024 * 1024 + 1)
if len(data) > 8 * 1024 * 1024 or hashlib.sha256(data).hexdigest() != conf["sha256"]:
    raise RuntimeError("EMT runtime checksum mismatch")
target = "/workspace/runtime"
os.makedirs(target, exist_ok=True)
with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as archive:
    for member in archive.getmembers():
        parts = member.name.split("/")
        if member.issym() or member.islnk() or member.name.startswith("/") or ".." in parts or not (member.isfile() or member.isdir()):
            raise RuntimeError("Unsafe runtime archive")
    archive.extractall(target, filter="data")
runner = os.path.join(target, os.listdir(target)[0], "emt_runner.py")
os.execv(sys.executable, [sys.executable, "-u", runner])
'''

# Only this fixed script is passed as a remote command; credentials travel on stdin.
REMOTE = r'''
import json, os, re, subprocess, sys
c = json.load(sys.stdin)
if not re.fullmatch(r"[0-9a-f-]{36}", c["job"]) or not isinstance(c["node"], int) or not 0 <= c["node"] < 4:
    raise ValueError("Invalid experiment identity")
name = "emt-exp-" + c["job"] + "-" + str(c["node"])
check = subprocess.run(["docker", "inspect", "--format", "{{.State.Running}}", name], capture_output=True, text=True)
if check.returncode == 0:
    if check.stdout.strip() == "true":
        print("already running", name)
        sys.exit(0)
    raise RuntimeError("This node's experiment container already ended. Check the EMT job; do not launch it again.")
base = "/root/emt-experiments/" + c["job"]
os.makedirs(base, mode=0o700, exist_ok=True)
os.chmod(base, 0o700)
def save(name, value):
    path = os.path.join(base, name)
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w") as f: f.write(value)
    return path
save("bootstrap.py", c["bootstrap"])
save("runtime.json", json.dumps(c["runtime"]))
environment = {
    "EMT_EXPERIMENT_JOB": c["job"], "EMT_EXPERIMENT_TOKEN": c["token"],
    "EMT_EXPERIMENT_API": c["api"], "EMT_EXPERIMENT_NODE": str(c["node"]),
    "PYTHONUNBUFFERED": "1",
}
if any("\n" in str(v) or "\r" in str(v) for v in environment.values()): raise ValueError("Invalid environment")
envfile = save("runner.env", "\n".join(k + "=" + v for k, v in environment.items()) + "\n")
try:
    subprocess.run(["docker", "run", "-d", "--name", name, "--gpus", "all", "--network", "host",
        "--env-file", envfile, "--mount", "type=bind,src=" + base + ",dst=/workspace",
        "--entrypoint", "python", c["image"], "-u", "/workspace/bootstrap.py"], check=True)
finally:
    os.unlink(envfile)
print("started", name)
'''


def launch(client, job, identity, run=subprocess.run):
    if not re.fullmatch(r"[0-9a-f-]{36}", job):
        raise ValueError("Use the experiment's UUID")
    conf = client.request("experiment-jobs/" + quote(job, safe="") + "/cluster-bootstrap", "POST", {})
    if conf.get("job") != job or conf.get("api") != client.base_url or conf.get("mode") != "independent-runs":
        raise ValueError("Unexpected cluster bootstrap response")
    if conf.get("image") not in ("pytorch/pytorch:2.8.0-cuda12.8-cudnn9-runtime", "pytorch/pytorch:2.8.0-cuda12.8-cudnn9-devel", "nvcr.io/nvidia/jax:25.04-py3"):
        raise ValueError("Unsupported scientific image")
    runtime = conf.get("runtime", {})
    if not re.fullmatch(r"/env-runtime/[0-9a-f]{16}/emt_env_runtime-1\.0\.0\.tar\.gz", runtime.get("path", "")) or not re.fullmatch(r"[0-9a-f]{64}", runtime.get("sha256", "")):
        raise ValueError("Invalid runtime manifest")
    members = conf.get("members", [])
    if len(members) not in (2, 3, 4) or sorted(m["index"] for m in members) != list(range(len(members))):
        raise ValueError("Invalid cluster members")
    # Parse every connection before starting any node.
    connections = [parse_ssh_connection(m["ssh"]) for m in members]
    results = []
    for member, (target, port) in zip(members, connections):
        payload = {"job": job, "node": member["index"], "api": conf["api"], "token": conf["token"],
                   "image": conf["image"], "bootstrap": BOOTSTRAP,
                   "runtime": {"url": conf["api"] + runtime["path"], "sha256": runtime["sha256"]}}
        args = ["ssh", "-p", str(port), "-i", str(Path(identity).expanduser().resolve()),
                "-o", "StrictHostKeyChecking=ask", "-o", "ConnectTimeout=20", "--", target,
                "python3 -c " + shlex.quote(REMOTE)]
        result = run(args, input=json.dumps(payload), text=True, timeout=900, check=False)
        results.append({"node": member["index"], "started": result.returncode == 0})
        if result.returncode:
            raise RuntimeError("Node startup failed. Check the experiment in EMT and stop it to release the entire cluster; do not prepare a second job.")
    return results


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--job", required=True)
    parser.add_argument("--identity", required=True, help="SSH private-key file; it is never uploaded")
    args = parser.parse_args()
    print(json.dumps(launch(EMTCompute(), args.job, args.identity), indent=2))


if __name__ == "__main__":
    main()