Fix jupyter spark memory calculation with DRA (#156)

Author: chi-yelp

service_configuration_lib/spark_config.py

@@ -1235,12 +1235,28 @@ def parse_memory_string(memory_string: Optional[str]) -> int:
         )
 
 
-def compute_requested_memory_overhead(spark_opts: Mapping[str, str], executor_memory):
-    return max(
+def get_spark_executor_memory_overhead_mb(spark_opts: Mapping[str, str], executor_memory) -> float:
+    """Return memory overhead in MB."""
+    # By default, Spark adds an overhead of 10% of the executor memory, with a
+    # minimum of 384mb
+    min_mem_overhead = 384
+    default_overhead_factor = 0.1
+
+    memory_overhead = max(
         parse_memory_string(spark_opts.get('spark.executor.memoryOverhead')),
         parse_memory_string(spark_opts.get('spark.mesos.executor.memoryOverhead')),
-        float(spark_opts.get('spark.kubernetes.memoryOverheadFactor', 0)) * executor_memory,
     )
+    if memory_overhead:
+        return float(max(memory_overhead, min_mem_overhead))
+    else:
+        memory_overhead_factor = (
+            spark_opts.get('spark.executor.memoryOverheadFactor') or
+            spark_opts.get('spark.kubernetes.memoryOverheadFactor') or
+            spark_opts.get('spark.mesos.executor.memoryOverheadFactor') or
+            default_overhead_factor
+        )
+        calculated_overhead = float(memory_overhead_factor) * executor_memory
+        return float(max(calculated_overhead, min_mem_overhead))
 
 
 def get_grafana_url(spark_conf: Mapping[str, str]) -> str:
@@ -1253,32 +1269,21 @@ def get_grafana_url(spark_conf: Mapping[str, str]) -> str:
 
 
 def get_resources_requested(spark_opts: Mapping[str, str]) -> Mapping[str, int]:
+    dra_enabled = str(spark_opts.get('spark.dynamicAllocation.enabled')).lower() == 'true'
     num_executors = (
-        # spark on k8s directly configure num instances
-        int(spark_opts.get('spark.executor.instances', 0)) or
-        # spark on mesos use cores.max and executor.core to calculate number of
-        # executors.
-        int(spark_opts.get('spark.cores.max', 0)) // int(spark_opts.get('spark.executor.cores', 0))
-    )
-    num_cpus = (
-        # spark on k8s
-        int(spark_opts.get('spark.executor.instances', 0)) * int(spark_opts.get('spark.executor.cores', 0)) or
-        # spark on mesos
-        int(spark_opts.get('spark.cores.max', 0))
+        int(spark_opts.get('spark.dynamicAllocation.maxExecutors', 0)) if dra_enabled
+        else
+        int(spark_opts.get('spark.executor.instances', 0))
     )
+    num_cpus = num_executors * int(spark_opts.get('spark.executor.cores', 0))
     num_gpus = int(spark_opts.get('spark.mesos.gpus.max', 0))
 
     executor_memory = parse_memory_string(spark_opts.get('spark.executor.memory', ''))
-    requested_memory = compute_requested_memory_overhead(spark_opts, executor_memory)
-    # by default, spark adds an overhead of 10% of the executor memory, with a
-    # minimum of 384mb
-    memory_overhead: int = (
-        requested_memory
-        if requested_memory > 0
-        else max(384, int(0.1 * executor_memory))
-    )
-    total_memory = (executor_memory + memory_overhead) * num_executors
-    log.info(f'Requested total memory of {total_memory} MiB')
+    executor_memory_overhead = get_spark_executor_memory_overhead_mb(spark_opts, executor_memory)
+    total_memory = int((executor_memory + executor_memory_overhead) * num_executors)
+    dra_enabled_string = '(DRA enabled)' if dra_enabled else ''
+
+    log.info(f'Requested total memory of {total_memory} MiB {dra_enabled_string}')
     return {
         'cpus': num_cpus,
         'mem': total_memory,

service_configuration_lib/utils.py

@@ -163,6 +163,7 @@ def get_runtime_env() -> str:
         return 'unknown'
 
 
+# TODO: merge with spark_config.parse_memory_string
 def get_spark_memory_in_unit(mem: str, unit: Literal['k', 'm', 'g', 't']) -> float:
     """
     Converts Spark memory to the desired unit.
@@ -182,6 +183,7 @@ def get_spark_memory_in_unit(mem: str, unit: Literal['k', 'm', 'g', 't']) -> flo
     return round(memory_unit, 5)
 
 
+# TODO: use spark_config.parse_memory_string
 def get_spark_driver_memory_mb(spark_conf: Dict[str, str]) -> float:
     """
     Returns the Spark driver memory in MB.
@@ -194,6 +196,7 @@ def get_spark_driver_memory_mb(spark_conf: Dict[str, str]) -> float:
         return SPARK_DRIVER_MEM_DEFAULT_MB
 
 
+# TODO: merge with spark_config.compute_requested_memory_overhead
 def get_spark_driver_memory_overhead_mb(spark_conf: Dict[str, str]) -> float:
     """
     Returns the Spark driver memory overhead in bytes.

tests/spark_config_test.py

@@ -1528,6 +1528,48 @@ def test_parse_memory_string(memory_string, expected_output):
     assert spark_config.parse_memory_string(memory_string) == expected_output
 
 
+@pytest.mark.parametrize(
+    'spark_opts,executor_memory,expected_output',
+    [
+        # min_memory_overhead
+        ({}, 1024, 384),
+        # default_memory_overhead_factor
+        ({}, 4096, 409.6),
+        # executor_memoryOverhead_configured
+        ({'spark.executor.memoryOverhead': '1024'}, 4096, 1024),
+        # mesos_memoryOverhead_configured
+        ({'spark.mesos.executor.memoryOverhead': '2048'}, 4096, 2048),
+        # kubernetes_memoryOverheadFactor_configured
+        ({'spark.kubernetes.memoryOverheadFactor': '0.2'}, 4096, int(4096 * 0.2)),
+        # multiple_configs_highest_selected
+        (
+            {
+                'spark.executor.memoryOverhead': '1024',
+                'spark.mesos.executor.memoryOverhead': '2048',
+                'spark.kubernetes.memoryOverheadFactor': '0.2',
+            },
+            4096,
+            2048,
+        ),
+        # default_memory_overhead_small_executor
+        ({}, 1024, 384),
+    ],
+    ids=[
+        'min_memory_overhead',
+        'default_memory_overhead_factor',
+        'executor_memoryOverhead_configured',
+        'mesos_memoryOverhead_configured',
+        'kubernetes_memoryOverheadFactor_configured',
+        'multiple_configs_highest_selected',
+        'default_memory_overhead_small_executor',
+    ],
+)
+def test_compute_requested_memory_overhead(spark_opts, executor_memory, expected_output):
+    result = spark_config.get_spark_executor_memory_overhead_mb(spark_opts, executor_memory)
+    assert isinstance(result, float)
+    assert int(result) == int(expected_output)
+
+
 def test_get_grafana_url():
     spark_conf = {
         'spark.executorEnv.PAASTA_CLUSTER': 'test-cluster',
@@ -1544,38 +1586,10 @@ def test_get_grafana_url():
 
 @pytest.mark.parametrize(
     'spark_opts,expected_output', [
-        # mesos ( 2 instances, not configure memory overhead, default: 384m )
-        (
-            {
-                'spark.cores.max': '10',
-                'spark.executor.cores': '5',
-                'spark.executor.memory': '2g',
-            },
-            {
-                'cpus': 10,
-                'mem': (384 + 2048) * 2,
-                'disk': (384 + 2048) * 2,
-                'gpus': 0,
-            },
-        ),
-        # mesos ( 2 instances, not configure memory overhead, default: 409m )
-        (
-            {
-                'spark.cores.max': '10',
-                'spark.executor.cores': '5',
-                'spark.executor.memory': '4g',
-            },
-            {
-                'cpus': 10,
-                'mem': (409 + 4096) * 2,
-                'disk': (409 + 4096) * 2,
-                'gpus': 0,
-            },
-        ),
-        # mesos ( 2 instances, configure memory overhead)
+        # basic_config
         (
             {
-                'spark.cores.max': '10',
+                'spark.executor.instances': '2',
                 'spark.executor.cores': '5',
                 'spark.executor.memory': '4g',
                 'spark.executor.memoryOverhead': '3072',
@@ -1587,26 +1601,25 @@ def test_get_grafana_url():
                 'gpus': 0,
             },
         ),
-        # mesos ( 2 instances, Duplicate config, choose the higher memory overhead)
+        # kubernetes_memory_overhead
         (
             {
-                'spark.cores.max': '10',
+                'spark.executor.instances': '2',
                 'spark.executor.cores': '5',
                 'spark.executor.memory': '4g',
-                'spark.executor.memoryOverhead': '3072',
-                'spark.mesos.executor.memoryOverhead': '4096',
+                'spark.kubernetes.memoryOverheadFactor': '0.5',
             },
             {
                 'cpus': 10,
-                'mem': (4096 + 4096) * 2,
-                'disk': (4096 + 4096) * 2,
+                'mem': (4096 * 0.5 + 4096) * 2,
+                'disk': (4096 * 0.5 + 4096) * 2,
                 'gpus': 0,
             },
         ),
-        # mesos ( 2 instances, configure memory overhead)
+        # mesos_memory_overhead
         (
             {
-                'spark.cores.max': '10',
+                'spark.executor.instances': '2',
                 'spark.executor.cores': '5',
                 'spark.executor.memory': '4g',
                 'spark.mesos.executor.memoryOverhead': '3072',
@@ -1618,10 +1631,11 @@ def test_get_grafana_url():
                 'gpus': 0,
             },
         ),
-        # k8s
+        # gpu_enabled
         (
             {
                 'spark.executor.instances': '2',
+                'spark.mesos.gpus.max': '2',
                 'spark.executor.cores': '5',
                 'spark.executor.memory': '4g',
                 'spark.executor.memoryOverhead': '3072',
@@ -1630,41 +1644,33 @@ def test_get_grafana_url():
                 'cpus': 10,
                 'mem': (3072 + 4096) * 2,
                 'disk': (3072 + 4096) * 2,
-                'gpus': 0,
+                'gpus': 2,
             },
         ),
-        # k8s
+        # dynamic_allocation_enabled
         (
             {
-                'spark.executor.instances': '2',
+                'spark.executor.instances': '0',
+                'spark.dynamicAllocation.enabled': 'true',
+                'spark.dynamicAllocation.maxExecutors': '2',
                 'spark.executor.cores': '5',
                 'spark.executor.memory': '4g',
                 'spark.kubernetes.memoryOverheadFactor': '0.5',
             },
             {
                 'cpus': 10,
-                'mem': (4096 * 0.5 + 4096) * 2,
-                'disk': (4096 * 0.5 + 4096) * 2,
+                'mem': (4096 * 1.5) * 2,
+                'disk': (4096 * 1.5) * 2,
                 'gpus': 0,
             },
         ),
-        # gpu
-        (
-            {
-                'spark.cores.max': '10',
-                'spark.mesos.gpus.max': '2',
-                'spark.executor.cores': '5',
-                'spark.executor.memory': '4g',
-                'spark.executor.memoryOverhead': '3072',
-            },
-            {
-                'cpus': 10,
-                'mem': (3072 + 4096) * 2,
-                'disk': (3072 + 4096) * 2,
-                'gpus': 2,
-            },
-
-        ),
+    ],
+    ids=[
+        'basic_config',
+        'kubernetes_memory_overhead',
+        'mesos_memory_overhead',
+        'gpu_enabled',
+        'dynamic_allocation_enabled',
     ],
 )
 def test_get_resources_requested(spark_opts, expected_output):