Add 'cluster_type' parameter to make create.py common for aws and azu…

…re Databricks

Signed-off-by: Tim Liu <timl@nvidia.com>

jenkins/databricks/clusterutils.py

@@ -23,21 +23,22 @@ 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
         templ = {}
         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'

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 <workspace> -t <token> -k <sshkey> -n <clustername> -i <idletime> -r <runtime> -o <workernodetype> -d <drivernodetype> -e <numworkers>')
+          'create.py -w <workspace> -t <token> -k <sshkey> -n <clustername> -i <idletime> -r <runtime> -o <workernodetype> -d <drivernodetype> -e <numworkers> -s <clustertype>')
       sys.exit(2)
 
   for opt, arg in opts:
       if opt == '-h':
           print(
-              'create.py -w <workspace> -t <token> -k <sshkey> -n <clustername> -i <idletime> -r <runtime> -o <workernodetype> -d <drivernodetype> -e <numworkers>')
+              'create.py -w <workspace> -t <token> -k <sshkey> -n <clustername> -i <idletime> -r <runtime> -o <workernodetype> -d <drivernodetype> -e <numworkers> -s <clustertype>')
           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)