minerva/cluster/run_cluster.py

59 lines
1.7 KiB
Python

import sys
import minerva
from minerva.pier import Pier
########### PREP ############################
DASK_BASE = "ami-0399a4f70ca684620" # dask on ubuntu 22.04 x86
NUM_WORK = int(sys.argv[1])
if sys.argv[2] == "large":
WORKER_TYPE = "r5.xlarge"
else:
WORKER_TYPE = "m5.large"
def worker(pier, n):
mach = pier.machine(ami = DASK_BASE,
instance_type = WORKER_TYPE,
username = "ubuntu",
name = f"dask-worker-{n}",
variables = {"type": "worker",
"number": n},
disk_size = 512)
return mach
def scheduler(pier):
mach = pier.machine(ami = DASK_BASE,
instance_type = "m5.large",
username = "ubuntu",
name = f"dask-scheduler",
variables = {"type": "scheduler"},
disk_size = 32)
return mach
########## CLUSTER ##########################
m = minerva.Minerva("hay")
pier = m.pier(subnet_id = "subnet-05eb26d8649a093e1", # project-subnet-public1-us-east-1a
sg_groups = ["sg-0f9e555954e863954", # ssh
"sg-0b34a3f7398076545", # default
"sg-04cd2626d91ac093c"], # dask (8786, 8787)
key_pair = ("Ari-Brown-HAY", "~/.ssh/Ari-Brown-HAY.pem"),
iam = "Minerva")
cluster = pier.cluster(scheduler, worker, num_workers=NUM_WORK)
cluster.start()
print()
print(f"dashboard: http://{cluster.scheduler.public_ip}:8787/")
print(f"cluster: {cluster.public_location}")
print()
print("type `exit()` to terminate the cluster")
print()
import IPython
IPython.embed()
cluster.terminate()