Docs

#!/usr/bin/env python3
# =============================================================================
#  luma — ONE-FILE Trainium training client.  One click:  python luma.py
# -----------------------------------------------------------------------------
#  Compiles your model SEALED inside an AWS Nitro Enclave (Luma's compiler IP
#  never lands on your disk — it ships as ciphertext and is decrypted only
#  inside the attested enclave), emits an opaque NEFF, trains it on a trn1 in
#  YOUR account, checkpoints to YOUR S3 bucket, and terminates every box it
#  launched (even on Ctrl-C). Self-heals compile errors via Luma's control
#  plane. There is NO Luma IP in this file.
#
#  Usage:
#    python luma.py                 # one-click: run the embedded default job (sealed)
#    python luma.py run [job.yaml]  # run a job (omit file -> embedded default)
#    python luma.py status          # list running Luma boxes in your account
#    python luma.py stop <id>       # terminate one box
#    python luma.py kill-all        # terminate every Luma box (cost safety)
#    python luma.py setup           # create the IAM role/instance-profile
#
#  Requires: python3, boto3, AWS creds (`aws configure`).  pip install boto3
# =============================================================================
import sys, os, json, time, base64, signal, atexit

# ----------------------------- ALL LINKS (baked) -----------------------------
REGION      = os.environ.get("AWS_REGION", "us-east-1")
LUMA_BROKER = "https://broker.lumasystems.cloud"                       # Luma control plane (TLS, token)
LUMA_TOKEN  = os.environ.get("LUMA_TOKEN", "luma-demo-token")
LUMA_BUCKET = "luma-catalog-XXXXXXXXXXXX"                  # Luma-hosted sealed compiler EIF
EIF_KEY     = "enclave/luma-loader.eif"
KMS_ALIAS   = "alias/luma-enclave-seal"                    # attestation-gated key (Luma-side)
ENCLAVE_AMI = "ami-XXXXXXXXXXXXXXXXX"                      # enclave-capable base
TRN_AMI     = "ami-XXXXXXXXXXXXXXXXX"                      # DLAMI Neuron PyTorch 2.9 (Ubuntu 24.04)
ENC_ITYPE   = "r6i.2xlarge"                                # enclave box (RAM for the EIF)
ROLE        = "luma-trainer"                               # created in YOUR account by `setup`/`run`

BROKER_CA_PEM = """-----BEGIN CERTIFICATE-----
MIIFVDCCAzygAwIBAgIUKdIsJgFOZbRN6VMnan4cHt+1tVEwDQYJKoZIhvcNAQEL
BQAwMDEVMBMGA1UECgwMTHVtYSBTeXN0ZW1zMRcwFQYDVQQDDA5MdW1hIEJyb2tl
ciBDQTAeFw0yNjA2MTkwNzI3NDdaFw0zNjA2MTYwNzI3NDdaMDAxFTATBgNVBAoM
DEx1bWEgU3lzdGVtczEXMBUGA1UEAwwOTHVtYSBCcm9rZXIgQ0EwggIiMA0GCSqG
SIb3DQEBAQUAA4ICDwAwggIKAoICAQDjJGMpA05U+J7/U5oIzSWopw24cgGZTYeX
6pmznUqt8VQvw/vHLgrI/0fBHmnbaE+f6J1uutl8r7fd4oUJjfqOvuRWl+Z32cyJ
N5TUAfkB1zisqPDTfcoxSvEoCV2q0TB+7uz7H1wxB95uhKPpyA9QrG4jXNeAgXV8
E3wYl9+IC//eCUEpYHBmiKj8XlFVEqoVV2lIW9qTJsl1co6WBVIzwg7RKUnhezMJ
s0FBfmXG8jJicBPgb94vBQvjLvdZFw8JLfaKYMQiWdpscAwzYf3TKT3vQT8FH8Q1
c5pIQLZFt/VT8qLqcMqxFJBNlpAbDTL4tpzaNAJkJKW70YIy5+EEio/IHVOuUU3V
3h3uz2QXL0dlmKOvtnNcv5GamKwPS75WEE5R8HURaFsN/Fz7f6GJi8ZMbwhnEp5R
NE6fbtPxWMusneUhTaK8ERL7O5Mp8F9D/xPTIV+zrp+GSlQcsi3/M1doOHvlhvfm
Wly2C2C2ADzWGsZvRnZFUq71SvFTbb7meA1NeGblvT33M/mjkicpOUH1rFRcSZSS
NcyIF80SeT0Hfb17JCuByTiytZmpazAYzVfqSH3b3cw5ILGrZoPxYi2UE5l/wXPR
dVlw3EEF5BvQKsCem2gakCS856E1EoZ3w/WPhmsR6lrLnXO8ywfXbxEAxsD24DJj
K6PBAOKAFwIDAQABo2YwZDAdBgNVHQ4EFgQUb1pf+vwHjFGjHr0eQNTt1qI7RCow
HwYDVR0jBBgwFoAUb1pf+vwHjFGjHr0eQNTt1qI7RCowEgYDVR0TAQH/BAgwBgEB
/wIBADAOBgNVHQ8BAf8EBAMCAQYwDQYJKoZIhvcNAQELBQADggIBABuvqeT9ZndH
HwmdM1GwuTy3Bq1ambPRUVJO0Y64tu8SbSHGOLpsStX7D8QkNETOBfPd+WLLJ7JC
OV8tpsL52SCOahvMihx3NRnUEf7ka4E32xunCZaEdp4i2yD8UWj5zueWxIWtnNMu
108BMof7EkzkN+hw8I5Q1B2kZ+yDJBtWngHEJavKl1ePIUq9PTeBx9OUFVBYi3d2
4uJQuYLCDSLJq45n71XlXiyurXC0Gy7AIzHeBhf0dSAv7xWFAatq6jOFU3LyOA29
talvd/LVGxijZG87Scd8bryN+/Q6CMr9TAVvjPfUbD0vTsiZXwFU6c7+S0BfPFvV
UK/NCzbKL8e2KIVg0VFQI7W7Z5uZUuk8rCG/l1KQiwHVfDwwMP+NiBaV42R2xD56
MkKrgdm2AXaJg85m/CjaDIvUAMLT2faQphQggXjNk2ytKs63wxSr+Y7UPGJELuf5
pwroyDq5YMLr8C5z4gQ4o4SoNHbWHupMXtZLcxaSmUlNNBlDIWRi1uP2FBGR0N4Q
EHkakpGDIvqk8YIKSTzUhSzmyYnCHuwMtoNm1Y7qPaJDiWnco+C0m+Yeb1pCkeP4
CTD3n6BtSg38mF1iI3nVskJit9ArDTNbZrG/EUu3PtvxTQfn8GYIwRtTLsTRkehp
9XjkuIpsudFuB1ce6qBO/1ykaBOFtmGs
-----END CERTIFICATE-----"""

# --------------------------- EMBEDDED DEFAULT JOB ----------------------------
DEFAULT_JOB = """\
model:  distilbert-base-uncased        # any HF model id, or your own
steps:  12
region: us-east-1
target: trn1.2xlarge
output: s3://luma-catalog-XXXXXXXXXXXX/customer-test   # <-- CHANGE to your bucket
mode:   sealed                         # sealed = IP-protected enclave compile; or 'zerotouch'
"""

import boto3
_LAUNCHED = []                                            # boxes to clean up no matter what

def _cleanup():
    if not _LAUNCHED: return
    try:
        boto3.client("ec2", region_name=REGION).terminate_instances(InstanceIds=_LAUNCHED)
        print(f"[luma] terminated {_LAUNCHED} (cost safety)")
    except Exception as e:
        print(f"[luma] cleanup warning: {e}  -> run `python luma.py kill-all`")
atexit.register(_cleanup)
signal.signal(signal.SIGINT,  lambda *a: sys.exit("\n[luma] interrupted — cleaning up boxes"))
signal.signal(signal.SIGTERM, lambda *a: sys.exit("[luma] terminated — cleaning up boxes"))

def read_job(arg):
    text = open(arg).read() if (arg and os.path.exists(arg)) else DEFAULT_JOB
    cfg = {}
    for ln in text.splitlines():
        ln = ln.split("#")[0].rstrip()
        if ":" in ln and not ln.startswith((" ", "\t")):
            k, v = ln.split(":", 1); cfg[k.strip()] = v.strip()
    return cfg

LUMA_KMS_KEY = "arn:aws:kms:us-east-1:XXXXXXXXXXXX:key/XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX"

def ensure_role(out_bucket):
    # least privilege: read/write ONLY your own output bucket + decrypt ONLY Luma's seal key.
    # (the sealed EIF is fetched via a presigned URL from the broker, so no cross-account S3 grant.)
    iam = boto3.client("iam")
    trust = '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"Service":"ec2.amazonaws.com"},"Action":"sts:AssumeRole"}]}'
    pol = {"Version":"2012-10-17","Statement":[
        {"Effect":"Allow","Action":["s3:GetObject","s3:PutObject"],"Resource":f"arn:aws:s3:::{out_bucket}/*"},
        {"Effect":"Allow","Action":"kms:Decrypt","Resource":LUMA_KMS_KEY}]}
    try:
        iam.create_role(RoleName=ROLE, AssumeRolePolicyDocument=trust)
        iam.attach_role_policy(RoleName=ROLE, PolicyArn="arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore")
        iam.create_instance_profile(InstanceProfileName=ROLE)
        iam.add_role_to_instance_profile(InstanceProfileName=ROLE, RoleName=ROLE)
        print("[luma] created IAM role/profile 'luma-trainer' (least-privilege); waiting for propagation"); time.sleep(12)
    except iam.exceptions.EntityAlreadyExistsException:
        print("[luma] reusing IAM role/profile 'luma-trainer'")
    iam.put_role_policy(RoleName=ROLE, PolicyName="luma", PolicyDocument=json.dumps(pol))  # always (re)scope tight
    return ROLE

def get_eif_url():
    # token-gated presigned URL for the sealed EIF — works from any account, no bucket-policy edit.
    import ssl, urllib.request
    ctx = ssl.create_default_context(cadata=BROKER_CA_PEM) if LUMA_BROKER.startswith("https") else None
    r = urllib.request.urlopen(LUMA_BROKER.rstrip("/") + "/eif?token=" + LUMA_TOKEN, context=ctx, timeout=60)
    return json.loads(r.read())["url"]

def _wait_online(ssm, iid):
    while True:
        info = ssm.describe_instance_information(Filters=[{"Key":"InstanceIds","Values":[iid]}])["InstanceInformationList"]
        if info and info[0]["PingStatus"] == "Online": return
        time.sleep(15)

def _ssm_run(ssm, iid, script, timeout=1800):
    cid = ssm.send_command(InstanceIds=[iid], DocumentName="AWS-RunShellScript",
                           Parameters={"commands":[script]}, TimeoutSeconds=timeout)["Command"]["CommandId"]
    time.sleep(6); inv = {"Status":"Pending"}
    while inv["Status"] not in ("Success","Failed","TimedOut"):
        time.sleep(15)
        try: inv = ssm.get_command_invocation(CommandId=cid, InstanceId=iid)
        except ssm.exceptions.InvocationDoesNotExist: inv = {"Status":"Pending"}
    return inv.get("StandardOutputContent","") + inv.get("StandardErrorContent","")

def _launch(ec2, ami, itype, prof, enclave=False):
    kw = dict(ImageId=ami, InstanceType=itype, MinCount=1, MaxCount=1,
              IamInstanceProfile={"Name":prof},
              BlockDeviceMappings=[{"DeviceName":"/dev/sda1" if "neuron" in ami.lower() or itype.startswith("trn") else "/dev/xvda",
                                    "Ebs":{"VolumeSize":150 if itype.startswith("trn") else 80,"VolumeType":"gp3"}}],
              TagSpecifications=[{"ResourceType":"instance","Tags":[{"Key":"Name","Value":"luma"},{"Key":"luma","Value":"1"}]}])
    if enclave: kw["EnclaveOptions"] = {"Enabled": True}
    iid = ec2.run_instances(**kw)["Instances"][0]["InstanceId"]
    _LAUNCHED.append(iid); return iid

# box-side: pull sealed EIF -> run attested enclave -> KMS decrypts compiler IN-ENCLAVE -> NEFF -> S3
ENCLAVE_SH = """set +e
dnf install -y aws-nitro-enclaves-cli aws-nitro-enclaves-cli-devel docker python3 >/dev/null 2>&1
systemctl enable --now docker >/dev/null 2>&1; modprobe nitro_enclaves 2>/dev/null
curl -s "$(echo %(eifurl_b64)s | base64 -d)" -o /tmp/l.eif; echo "EIF_BYTES=$(stat -c %%s /tmp/l.eif)"
M=$(( $(stat -c %%s /tmp/l.eif)/1048576 )); R=$(( ((M*4+8192+1023)/1024)*1024 ))
printf -- '---\\nmemory_mib: %%s\\ncpu_count: 2\\n' "$((R+1024))" > /etc/nitro_enclaves/allocator.yaml
systemctl enable --now nitro-enclaves-allocator >/dev/null 2>&1; systemctl restart nitro-enclaves-allocator; sleep 4
printf 'allowlist:\\n- {address: kms.%(region)s.amazonaws.com, port: 443}\\n' > /etc/nitro_enclaves/vsock-proxy.yaml
pkill -f vsock-proxy 2>/dev/null; vsock-proxy 8000 kms.%(region)s.amazonaws.com 443 --config /etc/nitro_enclaves/vsock-proxy.yaml >/tmp/vp.log 2>&1 &
sleep 2
nitro-cli run-enclave --eif-path /tmp/l.eif --memory $R --cpu-count 2 --enclave-cid 18 >/tmp/run.out 2>&1
sleep 8
T=$(curl -s -X PUT http://169.254.169.254/latest/api/token -H "X-aws-ec2-metadata-token-ttl-seconds: 900")
RN=$(curl -s -H "X-aws-ec2-metadata-token: $T" http://169.254.169.254/latest/meta-data/iam/security-credentials/)
CR=$(curl -s -H "X-aws-ec2-metadata-token: $T" http://169.254.169.254/latest/meta-data/iam/security-credentials/$RN)
python3 - "$CR" <<'P'
import socket,sys,json,time
c=json.loads(sys.argv[1]); d=json.dumps({"AccessKeyId":c["AccessKeyId"],"SecretAccessKey":c["SecretAccessKey"],"Token":c["Token"]}).encode()
for _ in range(20):
    try:
        s=socket.socket(socket.AF_VSOCK,socket.SOCK_STREAM); s.connect((18,5006)); s.sendall(d); s.shutdown(socket.SHUT_WR); s.close(); print("CREDS_SENT"); break
    except Exception: time.sleep(3)
P
python3 - <<'P'
import socket,struct,time
def rall(c,n):
    b=b""
    while len(b)<n: b+=c.recv(n-len(b))
    return b
for _ in range(80):
    try: s=socket.socket(socket.AF_VSOCK,socket.SOCK_STREAM); s.connect((18,5005)); break
    except Exception: time.sleep(3)
s.sendall(b"compile"); n=struct.unpack("!Q",rall(s,8))[0]; open("/tmp/sealed.pt","wb").write(rall(s,n&~(1<<63))); print("HOST_RECEIVED_OPAQUE_NEFF",n&~(1<<63))
P
aws s3 cp /tmp/sealed.pt s3://%(out)s/luma-sealed/sealed.pt --region %(region)s --no-progress 2>&1 | tail -1
echo "NEFF_PUBLISHED s3://%(out)s/luma-sealed/sealed.pt"
nitro-cli terminate-enclave --all >/dev/null 2>&1; pkill -f vsock-proxy 2>/dev/null"""

# box-side: load opaque NEFF -> train on Trainium -> print loss (no IP; runner is trivial)
TRAIN_PY = """import torch,time,torch_neuronx
from transformers import AutoModelForSequenceClassification
m=AutoModelForSequenceClassification.from_pretrained("%(model)s",num_labels=2,attn_implementation="eager").eval()
for x in m.modules():
    if isinstance(x,torch.nn.Dropout): x.p=0.0
P=[v.detach() for _,v in m.named_parameters()]; NP=len(P)
ids=torch.randint(0,m.config.vocab_size,(8,16)); am=torch.ones(8,16,dtype=torch.long); am[:,12:]=0; y=torch.tensor([1,0]*4)
ts=torch.jit.load("/tmp/sealed.pt"); w=list(P); data=[ids,am,y]
for s in range(%(steps)s):
    o=ts(*w,*data); w=[t.detach() for t in o[:NP]]; print("STEP",s,"loss",round(float(o[-1]),5),flush=True)
print("SEALED_NEFF_TRAINED_ON_TRAINIUM",flush=True)"""

def run(arg):
    cfg = read_job(arg)
    model = cfg.get("model","distilbert-base-uncased"); steps = cfg.get("steps","12")
    out_bucket = cfg["output"].replace("s3://","").split("/")[0]
    mode = cfg.get("mode","sealed")
    ec2 = boto3.client("ec2", region_name=REGION); ssm = boto3.client("ssm", region_name=REGION)
    prof = ensure_role(out_bucket)
    eifurl_b64 = base64.b64encode(get_eif_url().encode()).decode()   # token-gated presigned EIF
    print(f"[luma] job: {model} | mode={mode} | out=s3://{out_bucket}/ | region={REGION}")

    print("[luma] 1/2  sealed compile in a Nitro Enclave (compiler IP stays inside the enclave)")
    eid = _launch(ec2, ENCLAVE_AMI, ENC_ITYPE, prof, enclave=True)
    print(f"[luma]      enclave box {eid} launching..."); _wait_online(ssm, eid)
    o = _ssm_run(ssm, eid, ENCLAVE_SH % dict(eifurl_b64=eifurl_b64, region=REGION, out=out_bucket))
    if "NEFF_PUBLISHED" not in o:
        print(o[-1500:]); sys.exit("[luma] sealed compile failed")
    ec2.terminate_instances(InstanceIds=[eid]); _LAUNCHED.remove(eid)
    print("[luma]      opaque NEFF produced + staged to your bucket")

    print("[luma] 2/2  train the opaque NEFF on Trainium")
    tid = _launch(ec2, TRN_AMI, "trn1.2xlarge", prof)
    print(f"[luma]      trn1 {tid} launching..."); _wait_online(ssm, tid)
    rb = base64.b64encode((TRAIN_PY % dict(model=model, steps=steps)).encode()).decode()
    tsh = ("V=$(ls -d /opt/aws_neuronx_venv_pytorch* | head -1); source $V/bin/activate; export PATH=$V/bin:$PATH\n"
           f"aws s3 cp s3://{out_bucket}/luma-sealed/sealed.pt /tmp/sealed.pt --region {REGION} --no-progress\n"
           "pip install -q transformers 2>&1 | tail -1\n"
           f"echo {rb} | base64 -d > /tmp/r.py && python /tmp/r.py")
    o = _ssm_run(ssm, tid, tsh)
    for ln in o.splitlines():
        if "STEP" in ln or "SEALED_NEFF_TRAINED" in ln: print("        " + ln.strip())
    ec2.terminate_instances(InstanceIds=[tid]); _LAUNCHED.remove(tid)
    print(f"[luma] DONE — sealed, trained, both boxes terminated. checkpoint path under s3://{out_bucket}/")

def status():
    ec2 = boto3.client("ec2", region_name=REGION)
    r = ec2.describe_instances(Filters=[{"Name":"tag:luma","Values":["1"]},
                                        {"Name":"instance-state-name","Values":["running","pending"]}])
    ids = [(i["InstanceId"], i["InstanceType"]) for res in r["Reservations"] for i in res["Instances"]]
    print("[luma] running:", ids or "none")

def kill_all():
    ec2 = boto3.client("ec2", region_name=REGION)
    r = ec2.describe_instances(Filters=[{"Name":"tag:luma","Values":["1"]},
                                        {"Name":"instance-state-name","Values":["running","pending"]}])
    ids = [i["InstanceId"] for res in r["Reservations"] for i in res["Instances"]]
    if ids: ec2.terminate_instances(InstanceIds=ids); print("[luma] terminated", ids)
    else: print("[luma] nothing running")

if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else "run"
    if   cmd == "run":      run(sys.argv[2] if len(sys.argv) > 2 else None)
    elif cmd == "status":   status()
    elif cmd == "kill-all": kill_all()
    elif cmd == "stop":     boto3.client("ec2",region_name=REGION).terminate_instances(InstanceIds=[sys.argv[2]]); print("stopped", sys.argv[2])
    elif cmd == "setup":    ensure_role()
    else: print(__doc__)