fix gpu assignment in local-cluster mode and bash cli's (#973)

Fixes gpu assignment in local-cluster mode (used in test of no-code
change cli's) to match what is done local mode.

This gpu assignment error was silently failing in the cli runs in
run_test.sh in CI runs because the cli's weren't propagating Spark
errors to their exit codes. This PR also fixes this issue in the clis by
setting `check=True` in the Python `subprocess.run` invocations and adds
a test to check the Spark errors are propagated correctly.

Signed-off-by: Erik Ordentlich <[email protected]>

Author: eordentlich

python/run_test.sh

@@ -28,10 +28,20 @@ fi
 python -m spark_rapids_ml tests_no_import_change/test_no_import_change.py 0.2
 # runs on cpu
 python tests_no_import_change/test_no_import_change.py 0.2
+
 # runs on gpu with spark-submit (note: local[1] and pyspark<3.5.6 for spark-rapids-submit hangs probably due to barrier rdd timer threads. TBD root cause)
 pip install pyspark==3.5.6
 spark-rapids-submit --master local-cluster[1,1,1024] tests_no_import_change/test_no_import_change.py 0.2
+# test that failure mode returns non-zero exit code
+set +e
+spark-rapids-submit --master local-cluster[1,1,1024] tests_no_import_change/test_no_import_change.py -0.2
+if [ $? -eq 0 ]; then
+    echo "test should have returned non-zero exit code"
+    exit 1
+fi
+set -e
 pip install -r requirements_dev.txt
+
 # runs on cpu with spark-submit
 spark-submit --master local-cluster[1,1,1024] tests_no_import_change/test_no_import_change.py 0.2
 

python/src/spark_rapids_ml/pyspark_rapids.py

@@ -41,4 +41,4 @@ def main_cli() -> None:
     command_line = "pyspark " + " ".join(sys.argv[1:])
     env = dict(os.environ)
     env["PYTHONSTARTUP"] = f"{spark_rapids_ml.__path__[0]}/install.py"
-    subprocess.run(command_line, shell=True, env=env)
+    subprocess.run(command_line, shell=True, env=env, check=True)

python/src/spark_rapids_ml/spark_rapids_submit.py

@@ -46,4 +46,4 @@ def main_cli() -> None:
         + " ".join(sys.argv[i:])
     )
 
-    subprocess.run(command_line, shell=True)
+    subprocess.run(command_line, shell=True, check=True)

python/src/spark_rapids_ml/utils.py

@@ -96,7 +96,7 @@ def _get_spark_session() -> SparkSession:
 
 def _is_local(sc: SparkContext) -> bool:
     """Whether it is Spark local mode"""
-    return sc._jsc.sc().isLocal()  # type: ignore
+    return sc._jsc.sc().isLocal() or sc.getConf().get("spark.master").startswith("local-cluster")  # type: ignore
 
 
 def _is_standalone_or_localcluster(conf: SparkConf) -> bool: