Merge pull request #3741 from Yelp/MLCOMPUTE-1098_fix_parsing_error

MLCOMPUTE-1098 | fix error in parsing user spark args
Yelp · Nov 22, 2023 · a9ff4aa · a9ff4aa
2 parents fc39da2 + 2be37fd
commit a9ff4aa
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 13 deletions.
diff --git a/paasta_tools/cli/cmds/spark_run.py b/paasta_tools/cli/cmds/spark_run.py
@@ -698,21 +698,20 @@ def _parse_user_spark_args(
     enable_compact_bin_packing: bool = False,
     enable_spark_dra: bool = False,
 ) -> Dict[str, str]:
-    if not spark_args:
-        return {}
 
     user_spark_opts = {}
-    for spark_arg in spark_args.split():
-        fields = spark_arg.split("=", 1)
-        if len(fields) != 2:
-            print(
-                PaastaColors.red(
-                    "Spark option %s is not in format option=value." % spark_arg
-                ),
-                file=sys.stderr,
-            )
-            sys.exit(1)
-        user_spark_opts[fields[0]] = fields[1]
+    if spark_args:
+        for spark_arg in spark_args.split():
+            fields = spark_arg.split("=", 1)
+            if len(fields) != 2:
+                print(
+                    PaastaColors.red(
+                        "Spark option %s is not in format option=value." % spark_arg
+                    ),
+                    file=sys.stderr,
+                )
+                sys.exit(1)
+            user_spark_opts[fields[0]] = fields[1]
 
     if enable_compact_bin_packing:
         user_spark_opts["spark.kubernetes.executor.podTemplateFile"] = pod_template_path

diff --git a/tests/cli/test_cmds_spark_run.py b/tests/cli/test_cmds_spark_run.py
@@ -343,6 +343,7 @@ def test_get_spark_env(
         ),
         ("spark.cores.max", False, None),
         (None, False, {}),
+        (None, True, {"spark.dynamicAllocation.enabled": "true"}),
     ],
 )
 def test_parse_user_spark_args(spark_args, enable_spark_dra, expected, capsys):