blob: 4cbc62629fafb9270e2d8b1cc84cd6d3a2ce0efa [file] [log] [blame]
#!/usr/bin/env python3
#
# Copyright 2018 The Bazel Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import getpass
import queue
import subprocess
import sys
import threading
DEBUG = True
LOCATION = 'europe-west1-d'
MACHINES = {
'buildkite-ubuntu1404': {
'count': 4,
'startup_script': 'startup-ubuntu.sh',
'machine_type': 'n1-standard-32',
'local_ssd': 'interface=nvme',
},
'buildkite-ubuntu1604': {
'count': 4,
'startup_script': 'startup-ubuntu.sh',
'machine_type': 'n1-standard-32',
'local_ssd': 'interface=nvme',
},
'buildkite-windows': {
'count': 4,
'startup_script': 'startup-windows.ps1',
'machine_type': 'n1-standard-32',
'local_ssd': 'interface=scsi',
},
'%s-ubuntu1604' % getpass.getuser(): {
'image': 'buildkite-ubuntu1604',
'count': 4,
'startup_script': 'startup-ubuntu.sh',
'machine_type': 'n1-standard-32',
'local_ssd': 'interface=nvme',
}
}
PRINT_LOCK = threading.Lock()
WORK_QUEUE = queue.Queue()
def debug(*args, **kwargs):
if DEBUG:
with PRINT_LOCK:
print(*args, **kwargs)
def run(args, **kwargs):
debug('Running: %s' % ' '.join(args))
return subprocess.run(args, **kwargs)
def delete_vm(vm):
return run(['gcloud', 'compute', 'instances', 'delete', '--quiet', vm])
def create_vm(vm, idx, params):
if idx > 0:
vm = '%s-%s' % (vm, idx)
image_family = params.get('image', vm)
if delete_vm(vm).returncode == 0:
with PRINT_LOCK:
print('Deleted existing VM: %s' % vm)
cmd = ['gcloud', 'compute', 'instances', 'create', vm]
cmd.extend(['--zone', LOCATION])
cmd.extend(['--machine-type', params['machine_type']])
cmd.extend(['--network', 'buildkite'])
if 'windows' in image_family:
cmd.extend(['--metadata-from-file', 'windows-startup-script-ps1=' + params['startup_script']])
else:
cmd.extend(['--metadata-from-file', 'startup-script=' + params['startup_script']])
cmd.extend(['--min-cpu-platform', 'Intel Skylake'])
cmd.extend(['--boot-disk-type', 'pd-ssd'])
cmd.extend(['--boot-disk-size', '50GB'])
if 'local_ssd' in params:
cmd.extend(['--local-ssd', params['local_ssd']])
cmd.extend(['--image-project', 'bazel-public'])
cmd.extend(['--image-family', image_family])
cmd.extend(['--service-account', 'remote-account@bazel-public.iam.gserviceaccount.com'])
cmd.extend(['--scopes', 'cloud-platform'])
run(cmd)
def worker():
while True:
item = WORK_QUEUE.get()
if not item:
break
try:
create_vm(**item)
finally:
WORK_QUEUE.task_done()
def main(argv=None):
if argv is None:
argv = sys.argv[1:]
# Put VM creation instructions into the work queue.
worker_count = 0
for vm, params in MACHINES.items():
if argv and vm not in argv:
continue
worker_count += params['count']
if params['count'] > 1:
for idx in range(1, params['count'] + 1):
WORK_QUEUE.put({
'vm': vm,
'idx': idx,
'params': params
})
else:
WORK_QUEUE.put({
'vm': vm,
'idx': 0,
'params': params
})
# Spawn worker threads that will create the VMs.
threads = []
for _ in range(worker_count):
t = threading.Thread(target=worker)
t.start()
threads.append(t)
# Wait for all VMs to be created.
WORK_QUEUE.join()
# Signal worker threads to exit.
for _ in range(worker_count):
WORK_QUEUE.put(None)
# Wait for worker threads to exit.
for t in threads:
t.join()
return 0
if __name__ == '__main__':
sys.exit(main())