From 6d63d79e4690f5278745868ad22d53be2667d3c1 Mon Sep 17 00:00:00 2001 From: Timma Reddy Kunduru Date: Tue, 4 Mar 2025 03:20:55 -0800 Subject: [PATCH 1/4] Fixed the timeout regex of spark-submit command to match the options also --- paasta_tools/spark_tools.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/paasta_tools/spark_tools.py b/paasta_tools/spark_tools.py index fddb42db7a..72fee24adc 100644 --- a/paasta_tools/spark_tools.py +++ b/paasta_tools/spark_tools.py @@ -184,8 +184,10 @@ def auto_add_timeout_for_spark_job( if "spark-submit" not in cmd: return cmd try: + options_regex = "((-[a-zA-Z]+|\\-\\-[a-zA-Z\\-]+)(\\s+[a-zA-Z0-9\\-]+)?\\s*)*" + duration_regex = r"[\d]+[\.]?[\d]*[m|h]" timeout_present = re.match( - r"^.*timeout[\s]+[\d]+[\.]?[\d]*[m|h][\s]+spark-submit .*$", cmd + rf"^.*timeout[\s]+{options_regex}{duration_regex}[\s]+spark-submit .*$", cmd ) if not timeout_present: split_cmd = cmd.split("spark-submit") From 1e1f5a5a7a6669be54b2b021dca1f82dc7d4c5cf Mon Sep 17 00:00:00 2001 From: Timma Reddy Kunduru Date: Wed, 5 Mar 2025 05:56:48 -0800 Subject: [PATCH 2/4] Added unit tests for auto_add_timeout_for_spark_job function --- tests/test_spark_tools.py | 40 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/test_spark_tools.py b/tests/test_spark_tools.py index 7cbce1d99c..0d84d4a659 100644 --- a/tests/test_spark_tools.py +++ b/tests/test_spark_tools.py @@ -4,6 +4,7 @@ import pytest from paasta_tools import spark_tools +from paasta_tools.spark_tools import auto_add_timeout_for_spark_job def test_get_webui_url(): @@ -133,3 +134,42 @@ def test_get_volumes_from_spark_k8s_configs(mock_sys, spark_conf, expected): def test_get_spark_driver_monitoring_annotations(spark_config, expected): result = spark_tools.get_spark_driver_monitoring_annotations(spark_config) assert result == expected + + +@pytest.mark.parametrize( + argnames=[ + "cmd", + "timeout_duration", + "expected", + ], + argvalues=[ + pytest.param( + "spark-submit abc.py", + "4h", + "timeout 4h spark-submit abc.py", + id="No timeout", + ), + pytest.param( + "timeout 2h spark-submit abc.py", + "12h", + "timeout 2h spark-submit abc.py", + id="Timeout without options", + ), + pytest.param( + "timeout -v 2h spark-submit abc.py", + "12h", + "timeout -v 2h spark-submit abc.py", + id="Timeout with options", + ), + pytest.param( + "timeout -v -s 1 2h spark-submit abc.py", + "12h", + "timeout -v -s 1 2h spark-submit abc.py", + id="Timeout with multiple options", + ), + ], +) +def test_auto_add_timeout_for_spark_job(cmd, timeout_duration, expected): + result = auto_add_timeout_for_spark_job(cmd, timeout_duration) + + assert result == expected From 5c3c692a6a25df1a41fe06350908d20065785532 Mon Sep 17 00:00:00 2001 From: Timma Reddy Kunduru Date: Wed, 5 Mar 2025 09:39:57 -0800 Subject: [PATCH 3/4] Added the comment to explain why we are not supporting `=` Added unit test for `--` option --- paasta_tools/spark_tools.py | 3 +++ tests/test_spark_tools.py | 6 ++++++ 2 files changed, 9 insertions(+) diff --git a/paasta_tools/spark_tools.py b/paasta_tools/spark_tools.py index 72fee24adc..ec0f27c872 100644 --- a/paasta_tools/spark_tools.py +++ b/paasta_tools/spark_tools.py @@ -184,6 +184,9 @@ def auto_add_timeout_for_spark_job( if "spark-submit" not in cmd: return cmd try: + # Not supporting `=` for now, to support `=` we need to do lot of error checking, + # regex becomes a bit complex, use cases are very less, same command can be + # achieved without `=` also options_regex = "((-[a-zA-Z]+|\\-\\-[a-zA-Z\\-]+)(\\s+[a-zA-Z0-9\\-]+)?\\s*)*" duration_regex = r"[\d]+[\.]?[\d]*[m|h]" timeout_present = re.match( diff --git a/tests/test_spark_tools.py b/tests/test_spark_tools.py index 0d84d4a659..bc1730cf3f 100644 --- a/tests/test_spark_tools.py +++ b/tests/test_spark_tools.py @@ -167,6 +167,12 @@ def test_get_spark_driver_monitoring_annotations(spark_config, expected): "timeout -v -s 1 2h spark-submit abc.py", id="Timeout with multiple options", ), + pytest.param( + "timeout --signal SIGKILL 2h spark-submit abc.py", + "12h", + "timeout --signal SIGKILL 2h spark-submit abc.py", + id="Timeout with double dash option", + ), ], ) def test_auto_add_timeout_for_spark_job(cmd, timeout_duration, expected): From 1a32394d79dbb255a4f4813845bc26cacfdf5922 Mon Sep 17 00:00:00 2001 From: Timma Reddy Kunduru Date: Thu, 6 Mar 2025 06:17:28 -0800 Subject: [PATCH 4/4] Simplified the regex and added one more test case --- paasta_tools/spark_tools.py | 12 ++++++------ tests/test_spark_tools.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/paasta_tools/spark_tools.py b/paasta_tools/spark_tools.py index ec0f27c872..580291a88c 100644 --- a/paasta_tools/spark_tools.py +++ b/paasta_tools/spark_tools.py @@ -184,13 +184,13 @@ def auto_add_timeout_for_spark_job( if "spark-submit" not in cmd: return cmd try: - # Not supporting `=` for now, to support `=` we need to do lot of error checking, - # regex becomes a bit complex, use cases are very less, same command can be - # achieved without `=` also - options_regex = "((-[a-zA-Z]+|\\-\\-[a-zA-Z\\-]+)(\\s+[a-zA-Z0-9\\-]+)?\\s*)*" - duration_regex = r"[\d]+[\.]?[\d]*[m|h]" + # This is not an exhaustive regex, matches the invalid ones also, where as the invalid + # timeout command will fail during execution + options_regex = r"(--?[a-z][a-z-]*((\s+|=)[\w\d-]+)?\s+)*" + duration_regex = r"\d+\.?\d*[smhd]?" + timeout_present = re.match( - rf"^.*timeout[\s]+{options_regex}{duration_regex}[\s]+spark-submit .*$", cmd + rf"^.*timeout\s+{options_regex}{duration_regex}\s+spark-submit .*$", cmd ) if not timeout_present: split_cmd = cmd.split("spark-submit") diff --git a/tests/test_spark_tools.py b/tests/test_spark_tools.py index bc1730cf3f..784c6cdb9a 100644 --- a/tests/test_spark_tools.py +++ b/tests/test_spark_tools.py @@ -168,9 +168,9 @@ def test_get_spark_driver_monitoring_annotations(spark_config, expected): id="Timeout with multiple options", ), pytest.param( - "timeout --signal SIGKILL 2h spark-submit abc.py", + "timeout -k 10m --signal=SIGKILL 2h spark-submit abc.py", "12h", - "timeout --signal SIGKILL 2h spark-submit abc.py", + "timeout -k 10m --signal=SIGKILL 2h spark-submit abc.py", id="Timeout with double dash option", ), ],