diff --git a/jenkins/databricks/clusterutils.py b/jenkins/databricks/clusterutils.py index af085ee0f584..fbe7dfed0472 100644 --- a/jenkins/databricks/clusterutils.py +++ b/jenkins/databricks/clusterutils.py @@ -23,7 +23,7 @@ class ClusterUtils(object): @staticmethod def generate_create_templ(sshKey, cluster_name, runtime, idle_timeout, - num_workers, driver_node_type, worker_node_type, + num_workers, driver_node_type, worker_node_type, cluster_type, printLoc=sys.stdout): timeStr = str(int(time.time())) uniq_name = cluster_name + "-" + timeStr @@ -31,13 +31,15 @@ def generate_create_templ(sshKey, cluster_name, runtime, idle_timeout, templ['cluster_name'] = uniq_name print("cluster name is going to be %s" % uniq_name, file=printLoc) templ['spark_version'] = runtime - templ['aws_attributes'] = { - "zone_id": "us-west-2a", - "first_on_demand": 1, - "availability": "SPOT_WITH_FALLBACK", - "spot_bid_price_percent": 100, - "ebs_volume_count": 0 - } + if (cluster_type == 'aws'): { + templ['aws_attributes'] = { + "zone_id": "us-west-2a", + "first_on_demand": 1, + "availability": "SPOT_WITH_FALLBACK", + "spot_bid_price_percent": 100, + "ebs_volume_count": 0 + } + } templ['autotermination_minutes'] = idle_timeout templ['enable_elastic_disk'] = 'false' templ['enable_local_disk_encryption'] = 'false' diff --git a/jenkins/databricks/create.py b/jenkins/databricks/create.py index 10b1d99c2b41..4ae60638a388 100644 --- a/jenkins/databricks/create.py +++ b/jenkins/databricks/create.py @@ -33,20 +33,21 @@ def main(): num_workers = 1 worker_type = 'g4dn.xlarge' driver_type = 'g4dn.xlarge' + cluster_type = 'aws' try: - opts, args = getopt.getopt(sys.argv[1:], 'hw:t:k:n:i:r:o:d:e:', + opts, args = getopt.getopt(sys.argv[1:], 'hw:t:k:n:i:r:o:d:e:s:', ['workspace=', 'token=', 'sshkey=', 'clustername=', 'idletime=', - 'runtime=', 'workertype=', 'drivertype=', 'numworkers=']) + 'runtime=', 'workertype=', 'drivertype=', 'numworkers=', 'clustertype=']) except getopt.GetoptError: print( - 'create.py -w -t -k -n -i -r -o -d -e ') + 'create.py -w -t -k -n -i -r -o -d -e -s ') sys.exit(2) for opt, arg in opts: if opt == '-h': print( - 'create.py -w -t -k -n -i -r -o -d -e ') + 'create.py -w -t -k -n -i -r -o -d -e -s ') sys.exit() elif opt in ('-w', '--workspace'): workspace = arg @@ -66,6 +67,8 @@ def main(): driver_type = arg elif opt in ('-e', '--numworkers'): num_workers = arg + elif opt in ('-s', '--clustertype'): + cluster_type = arg print('-w is ' + workspace, file=sys.stderr) print('-k is ' + sshkey, file=sys.stderr) @@ -75,6 +78,7 @@ def main(): print('-o is ' + worker_type, file=sys.stderr) print('-d is ' + driver_type, file=sys.stderr) print('-e is ' + str(num_workers), file=sys.stderr) + print('-s is ' + cluster_type, file=sys.stderr) if not sshkey: print("You must specify an sshkey!", file=sys.stderr) @@ -85,7 +89,7 @@ def main(): sys.exit(2) templ = ClusterUtils.generate_create_templ(sshkey, cluster_name, runtime, idletime, - num_workers, driver_type, worker_type, printLoc=sys.stderr) + num_workers, driver_type, worker_type, cluster_type, printLoc=sys.stderr) clusterid = ClusterUtils.create_cluster(workspace, templ, token, printLoc=sys.stderr) ClusterUtils.wait_for_cluster_start(workspace, clusterid, token, printLoc=sys.stderr)