1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
|
#!/usr/bin/env python3
"""Provision a large virtual machine on Digital Ocean and run a nix remote
build, then destroy the VM.
I tested this like so:
$ ./nix-build-omatic "'<nixpkgs>'" -A ocaml
The special quoting is necessary until I setup an argparse CLI, currently it
just passes all args through to the nix-build call. But it works!
"""
import json
import os
import sys
import subprocess
import time
def noti(msg):
print(f"nix-build-omatic: {msg}")
class Builder:
def __init__(self):
# TODO: generate a temporary privkey, and use that instead
self.keys = [k["id"] for k in json.loads(subprocess.check_output(["doctl", "compute", "ssh-key", "list", "--output", "json"]).decode('utf-8'))]
self.privkey_path = os.path.expanduser("~/.ssh/id_rsa")
self.max_jobs = 32
self.ip = None
def __enter__(self):
"Startup a droplet."
noti("preparing droplet...")
res = self.droplet(
"create",
# this is my custom nixos image, I had to download from nixos.org and then
# upload it to DO. There's gotta be a better way to do this..
"--image", "121034270",
"--size", "c2-32vcpu-64gb", # TODO: set this in argparse, match '32vcpu' with 'max_jobs'
"--region", "nyc3",
"--tag-name", "nix-build-omatic",
"--ssh-keys", ",".join([str(k) for k in self.keys]),
"nix-build-omatic"
)
self.id = str(res["id"])
self.active = False
while not self.active:
res = self.droplet("get", self.id)
if res["status"] == "active":
self.active = True
noti("droplet ready!")
else:
continue
while self.ip is None:
res = self.droplet("get", self.id)
self.ip = [n["ip_address"] for n in res["networks"]["v4"] if n["type"] == "public" ][0]
self.url = f"ssh://root@{self.ip}"
self.do_stupid_handshake()
return self
def __exit__(self, type, value, traceback):
"Destroy the droplet."
self.droplet("delete", self.id, "--force", parse=False)
noti("droplet destroyed")
def droplet(self, *args, parse=True):
"Run a `doctl compute droplet *args` command"
proc = subprocess.run(
["doctl", "compute", "droplet", *args, "--output", "json"],
capture_output=True
)
stderr = proc.stderr.decode('utf-8')
stdout = proc.stdout.decode('utf-8')
if proc.returncode > 0:
noti("stderr:", stderr)
noti("stdout:", stdout)
if parse:
return json.loads(stdout)[0]
else:
return stdout
def do_stupid_handshake(self):
"""
this whole thing is doing a key handshake, which shouldn't be necessary!
i should just generate the keys on my end and set them via the API! or
retrieve them via the API! i'm already authenticated! wtf digital ocean
"""
noti("exchanging keys...")
keyscan = None
while keyscan is None:
cmd = f"ssh-keyscan -t ed25519 {self.ip}"
proc = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stderr = proc.stderr.decode('utf-8')
stdout = proc.stdout.decode('utf-8')
if proc.returncode > 0:
continue
elif stdout == "":
continue
else:
keyscan = stdout
with open(os.path.expanduser("~/.ssh/known_hosts"), 'a') as known_hosts:
known_hosts.write(keyscan)
self.host_pubkey = subprocess.check_output(f"ssh root@{self.ip} cat /etc/ssh/ssh_host_ed25519_key.pub | base64 -w0", shell=True).decode('utf-8')
if __name__ == '__main__':
# TODO: use argparse, properly quote '<nixpkgs>' and so on
args = sys.argv[1:]
start_time = time.time()
with Builder() as b:
cmd = " ".join([
"nix-build",
*args,
# TODO: obv get rid of this when I'm not testing...
"--option", "substitute", "false",
# TODO: make this --builders line structured data somehow
"--builders", f"'{b.url} x86_64-linux {b.privkey_path} {b.max_jobs} - - - {b.host_pubkey}'"
])
noti(f"running: {cmd}")
subprocess.check_call(cmd, shell=True)
duration = time.time() - start_time
duration_human = time.strftime('%H:%M:%S', duration)
# the 32 vCPU machine costs $1/hr, so convert that to cost per second
cost = (1/60/60)*duration
print(f"duration: {duration_human}")
print(f"estimated cost: {cost:0.2f} USD")
|
