Fix memory format (#10)

KanchiShimono · Jul 9, 2021 · e206b60 · e206b60
1 parent ec36303
commit e206b60
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 36 deletions.
diff --git a/README.md b/README.md
@@ -29,11 +29,11 @@ sco = SparkConfOptimizer(executor_instance, num_nodes, deploy_mode)
 print(sco)
 
 # spark.driver.cores: 5
-# spark.driver.memory: 36
-# spark.driver.memoryOvearhead: 5
+# spark.driver.memory: 36g
+# spark.driver.memoryOvearhead: 5g
 # spark.executor.cores: 5
-# spark.executor.memory: 36
-# spark.executor.memoryOvearhead: 5
+# spark.executor.memory: 36g
+# spark.executor.memoryOvearhead: 5g
 # spark.executor.instances: 60
 # spark.default.parallelism: 600
 # spark.sql.shuffle.partitions: 600
@@ -48,11 +48,11 @@ sco = SparkConfOptimizer(executor_instance, num_nodes, deploy_mode)
 print(sco)
 
 # spark.driver.cores: 5
-# spark.driver.memory: 36
-# spark.driver.memoryOvearhead: 5
+# spark.driver.memory: 36g
+# spark.driver.memoryOvearhead: 5g
 # spark.executor.cores: 5
-# spark.executor.memory: 36
-# spark.executor.memoryOvearhead: 5
+# spark.executor.memory: 36g
+# spark.executor.memoryOvearhead: 5g
 # spark.executor.instances: 59
 # spark.default.parallelism: 590
 # spark.sql.shuffle.partitions: 590
@@ -75,11 +75,11 @@ sco = SparkConfOptimizer(
 print(sco)
 
 # spark.driver.cores: 3
-# spark.driver.memory: 26
-# spark.driver.memoryOvearhead: 3
+# spark.driver.memory: 26g
+# spark.driver.memoryOvearhead: 3g
 # spark.executor.cores: 5
-# spark.executor.memory: 36
-# spark.executor.memoryOvearhead: 5
+# spark.executor.memory: 36g
+# spark.executor.memoryOvearhead: 5g
 # spark.executor.instances: 60
 # spark.default.parallelism: 600
 # spark.sql.shuffle.partitions: 600
@@ -108,11 +108,11 @@ print(conf.getAll())
 conf.setAll(sco.as_list())
 # dict_items([
 #     ('spark.driver.cores', '5'),
-#     ('spark.driver.memory', '36'),
-#     ('spark.driver.memoryOvearhead', '5'),
+#     ('spark.driver.memory', '36g'),
+#     ('spark.driver.memoryOvearhead', '5g'),
 #     ('spark.executor.cores', '5'),
-#     ('spark.executor.memory', '36'),
-#     ('spark.executor.memoryOvearhead', '5'),
+#     ('spark.executor.memory', '36g'),
+#     ('spark.executor.memoryOvearhead', '5g'),
 #     ('spark.executor.instances', '60'),
 #     ('spark.default.parallelism', '600'),
 #     ('spark.sql.shuffle.partitions', '600')

diff --git a/src/scopt/optimizer.py b/src/scopt/optimizer.py
@@ -242,11 +242,11 @@ class SparkConfOptimizer:
     >>> print(sco)
 
     spark.driver.cores: 5
-    spark.driver.memory: 36
-    spark.driver.memoryOvearhead: 5
+    spark.driver.memory: 36g
+    spark.driver.memoryOvearhead: 5g
     spark.executor.cores: 5
-    spark.executor.memory: 36
-    spark.executor.memoryOvearhead: 5
+    spark.executor.memory: 36g
+    spark.executor.memoryOvearhead: 5g
     spark.executor.instances: 60
     spark.default.parallelism: 600
     spark.sql.shuffle.partitions: 600
@@ -255,11 +255,11 @@ class SparkConfOptimizer:
     >>> print(conf.getAll())
     dict_items([
         ('spark.driver.cores', '5'),
-        ('spark.driver.memory', '36'),
-        ('spark.driver.memoryOvearhead', '5'),
+        ('spark.driver.memory', '36g'),
+        ('spark.driver.memoryOvearhead', '5g'),
         ('spark.executor.cores', '5'),
-        ('spark.executor.memory', '36'),
-        ('spark.executor.memoryOvearhead', '5'),
+        ('spark.executor.memory', '36g'),
+        ('spark.executor.memoryOvearhead', '5g'),
         ('spark.executor.instances', '60'),
         ('spark.default.parallelism', '600'),
         ('spark.sql.shuffle.partitions', '600')
@@ -326,11 +326,11 @@ def _repr_html_(self) -> str:
     def as_dict(self) -> Dict[str, Union[int, str]]:
         return {
             'spark.driver.cores': self.optimizer.driver_cores,
-            'spark.driver.memory': self.optimizer.driver_memory,
-            'spark.driver.memoryOvearhead': self.optimizer.driver_memory_overhead,  # noqa: E501
+            'spark.driver.memory': f'{self.optimizer.driver_memory}g',
+            'spark.driver.memoryOvearhead': f'{self.optimizer.driver_memory_overhead}g',  # noqa: E501
             'spark.executor.cores': self.optimizer.executor_cores,
-            'spark.executor.memory': self.optimizer.executor_memory,
-            'spark.executor.memoryOvearhead': self.optimizer.executor_memory_overhead,  # noqa: E501
+            'spark.executor.memory': f'{self.optimizer.executor_memory}g',
+            'spark.executor.memoryOvearhead': f'{self.optimizer.executor_memory_overhead}g',  # noqa: E501
             'spark.executor.instances': self.optimizer.executor_instances,
             'spark.default.parallelism': self.optimizer.default_parallelism,
             'spark.sql.shuffle.partitions': self.optimizer.sql_shuffle_partitions,  # noqa: E501

diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
@@ -101,11 +101,11 @@ def test_as_dict(self) -> None:
         optimizer = SparkConfOptimizer(Instance(32, 248), 10, 'client')
         expected = {
             'spark.driver.cores': 5,
-            'spark.driver.memory': 36,
-            'spark.driver.memoryOvearhead': 5,
+            'spark.driver.memory': '36g',
+            'spark.driver.memoryOvearhead': '5g',
             'spark.executor.cores': 5,
-            'spark.executor.memory': 36,
-            'spark.executor.memoryOvearhead': 5,
+            'spark.executor.memory': '36g',
+            'spark.executor.memoryOvearhead': '5g',
             'spark.executor.instances': 60,
             'spark.default.parallelism': 600,
             'spark.sql.shuffle.partitions': 600,
@@ -116,11 +116,11 @@ def test_as_list(self) -> None:
         optimizer = SparkConfOptimizer(Instance(32, 248), 10, 'client')
         expected = [
             ('spark.driver.cores', 5),
-            ('spark.driver.memory', 36),
-            ('spark.driver.memoryOvearhead', 5),
+            ('spark.driver.memory', '36g'),
+            ('spark.driver.memoryOvearhead', '5g'),
             ('spark.executor.cores', 5),
-            ('spark.executor.memory', 36),
-            ('spark.executor.memoryOvearhead', 5),
+            ('spark.executor.memory', '36g'),
+            ('spark.executor.memoryOvearhead', '5g'),
             ('spark.executor.instances', 60),
             ('spark.default.parallelism', 600),
             ('spark.sql.shuffle.partitions', 600),