chg

qubole · Jun 17, 2020 · 76fd60d · 76fd60d
1 parent 0e267aa
commit 76fd60d
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/qds_sdk/cluster_info_v22.py b/qds_sdk/cluster_info_v22.py
@@ -55,7 +55,8 @@ def set_cluster_info_from_arguments(self, arguments):
                               disable_autoscale_node_pause=arguments.disable_autoscale_node_pause,
                               paused_autoscale_node_timeout_mins=arguments.paused_autoscale_node_timeout_mins,
                               parent_cluster_id=arguments.parent_cluster_id,
-                              image_version=arguments.image_version)
+                              image_version=arguments.image_version,
+                              ebs_optimized=arguments.ebs_optimized)
 
         self.set_composition(master_type=arguments.master_type,
                              master_spot_block_duration=arguments.master_spot_block_duration,
@@ -113,7 +114,8 @@ def set_cluster_info(self,
                          disable_autoscale_node_pause=None,
                          paused_autoscale_node_timeout_mins=None,
                          parent_cluster_id=None,
-                         image_version=None):
+                         image_version=None,
+                         ebs_optimized=None):
         """
         Args:
 
@@ -191,6 +193,8 @@ def set_cluster_info(self,
 
             `image_version`: cluster image version
 
+            `ebs_optimized`: Specify EBS optimization to be disabled - default is enabled.
+
         Doc: For getting details about arguments
         http://docs.qubole.com/en/latest/rest-api/cluster_api/create-new-cluster.html#parameters
 
@@ -219,6 +223,7 @@ def set_cluster_info(self,
         self.cluster_info['rootdisk']['size'] = root_disk_size
         self.cluster_info['parent_cluster_id'] = parent_cluster_id
         self.cluster_info['cluster_image_version'] = image_version
+        self.cluster_info['ebs_optimized'] = ebs_optimized
         self.set_data_disk(disk_size, disk_count, disk_type,
                            upscaling_config, enable_encryption)
         self.set_monitoring(enable_ganglia_monitoring,
@@ -541,6 +546,9 @@ def cluster_info_parser(argparser, action):
         cluster_info.add_argument("--image-version",
                                   dest="image_version",
                                   help="cluster image version")
+        cluster_info.add_argument("--ebs-optimized",
+                                  dest="ebs_optimized",
+                                  help="Enable EBS optimization for the instance")
         termination = cluster_info.add_mutually_exclusive_group()
         termination.add_argument("--disallow-cluster-termination",
                                  dest="disallow_cluster_termination",

diff --git a/tests/test_clusterv22.py b/tests/test_clusterv22.py
@@ -212,6 +212,28 @@ def test_image_version_v22(self):
                                                 'cluster_image_version': '1.latest',
                                                 'composition': {'min_nodes': {'nodes': [{'percentage': 100, 'type': 'ondemand'}]}, 'master': {'nodes': [{'percentage': 100, 'type': 'ondemand'}]}, 'autoscaling_nodes': {'nodes': [{'percentage': 50, 'type': 'ondemand'}, {'timeout_for_request': 1, 'percentage': 50, 'type': 'spot', 'fallback': 'ondemand', 'maximum_bid_price_percentage': 100, 'allocation_strategy': None}]}}, 'label': ['test_label']}})
 
+    def test_ebs_optimized_v22(self):
+        sys.argv = ['qds.py', '--version', 'v2.2', 'cluster', 'create', '--label',
+                     'test_label', '--flavour', 'hive', '--slave-instance-type', 'c1.xlarge', '--min-nodes', '1', '--ebs-optimized', False]
+        Qubole.cloud = None
+        print_command()
+        Connection._api_call = Mock(return_value={})
+        qds.main()
+        Connection._api_call.assert_called_with('POST', 'clusters',
+                                                {'engine_config': {'flavour': 'hive'},
+                                                 'cluster_info': {'label': ['test_label'],
+                                                 'min_nodes': 1,
+                                                 'slave_instance_type': 'c1.xlarge',
+                                                 'ebs_optimized': False,
+                                                 'composition': {'min_nodes': {'nodes': [{'percentage': 100, 'type': 'ondemand'}]},
+                                                                'master': {'nodes': [{'percentage': 100, 'type': 'ondemand'}]},
+                                                                'autoscaling_nodes': {'nodes': [{'percentage': 50, 'type': 'ondemand'},
+                                                                                                {'timeout_for_request': 1, 'percentage': 50,
+                                                                                                 'type': 'spot', 'fallback': 'ondemand',
+                                                                                                 'maximum_bid_price_percentage': 100,
+                                                                                                 'allocation_strategy': None}]}},
+                                                 'label': ['test_label']}})
+
     def test_hive_settings(self):
         sys.argv = ['qds.py', '--version', 'v2.2', 'cluster', 'create', '--label', 'test_label',
                     '--is_hs2', 'true', '--hs2_thrift_port', '10001']