From 44db4560a8657176f70efd911ab721d8c3f05eed Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Mon, 12 Oct 2026 12:48:29 -0500 Subject: [PATCH] fix(core): let remote Spark jobs use pre_execution_script #5816 made _JobSettings add the sagemaker-feature-store-pyspark install and JAR-copy commands to pre_execution_commands for Spark jobs, before the check that only one of pre_execution_commands and pre_execution_script is set. So @remote(spark_config=SparkConfig(), pre_execution_script=...) always failed with 'Only one of pre_execution_commands or pre_execution_script can be specified!' although the user set only the script (this worked in v2). The commands were also appended to the caller's own list in place. Check the user's configuration before adding the Spark commands; when a script is set, leave the setup to it (with a warning), and append the commands to a copy of the user's list. --- .../src/sagemaker/core/remote_function/job.py | 40 +++++++++------ .../tests/unit/remote_function/test_job.py | 50 +++++++++++++++++++ 2 files changed, 76 insertions(+), 14 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/remote_function/job.py b/sagemaker-core/src/sagemaker/core/remote_function/job.py index efdcf7db4f..90d05403c9 100644 --- a/sagemaker-core/src/sagemaker/core/remote_function/job.py +++ b/sagemaker-core/src/sagemaker/core/remote_function/job.py @@ -671,6 +671,18 @@ def __init__( sagemaker_session=self.sagemaker_session, ) + self.pre_execution_script = resolve_value_from_config( + direct_input=pre_execution_script, + config_path=REMOTE_FUNCTION_PRE_EXECUTION_SCRIPT, + sagemaker_session=self.sagemaker_session, + ) + + # Check what the user configured, before the Spark setup commands below are added. + if self.pre_execution_commands is not None and self.pre_execution_script is not None: + raise ValueError( + "Only one of pre_execution_commands or pre_execution_script can be specified!" + ) + # When using Spark, ensure sagemaker-feature-store-pyspark is installed # and its version-matched JAR is on Spark's classpath before spark-submit if spark_config: @@ -689,22 +701,22 @@ def __init__( "[shutil.copy(j, '/usr/lib/spark/jars/') " "for j in glob.glob(os.path.join(jars_dir, '*' + spark_ver + '*.jar'))]\"" ) - if self.pre_execution_commands is None: + if self.pre_execution_script is not None: + # Only one of commands or a script can run before the job, so leave the + # user's script in charge of the setup. + logger.warning( + "pre_execution_script is set, so sagemaker-feature-store-pyspark is not " + "installed automatically for this Spark job; install it in the script if " + "it is needed." + ) + elif self.pre_execution_commands is None: self.pre_execution_commands = [install_cmd, copy_jar_cmd] elif install_cmd not in self.pre_execution_commands: - self.pre_execution_commands.append(install_cmd) - self.pre_execution_commands.append(copy_jar_cmd) - - self.pre_execution_script = resolve_value_from_config( - direct_input=pre_execution_script, - config_path=REMOTE_FUNCTION_PRE_EXECUTION_SCRIPT, - sagemaker_session=self.sagemaker_session, - ) - - if self.pre_execution_commands is not None and self.pre_execution_script is not None: - raise ValueError( - "Only one of pre_execution_commands or pre_execution_script can be specified!" - ) + # Copy, so the caller's list is not modified. + self.pre_execution_commands = list(self.pre_execution_commands) + [ + install_cmd, + copy_jar_cmd, + ] self.include_local_workdir = resolve_value_from_config( direct_input=include_local_workdir, diff --git a/sagemaker-core/tests/unit/remote_function/test_job.py b/sagemaker-core/tests/unit/remote_function/test_job.py index 39e2d56d28..7823f6c323 100644 --- a/sagemaker-core/tests/unit/remote_function/test_job.py +++ b/sagemaker-core/tests/unit/remote_function/test_job.py @@ -106,6 +106,56 @@ def test_init_with_pre_execution_commands_and_script_raises_error(self, mock_ses image_uri="test-image", ) + @patch.object(_JobSettings, "_get_default_spark_image", return_value="spark-image") + def test_init_with_spark_and_pre_execution_script(self, mock_spark_image, mock_session): + """The SDK's own Spark setup commands must not make a user's script conflict.""" + settings = _JobSettings( + sagemaker_session=mock_session, + role="arn:aws:iam::123456789012:role/role", + instance_type="ml.m5.xlarge", + spark_config=SparkConfig(), + pre_execution_script="/path/to/script.sh", + ) + + assert settings.pre_execution_script == "/path/to/script.sh" + assert settings.pre_execution_commands is None + + @patch.object(_JobSettings, "_get_default_spark_image", return_value="spark-image") + def test_init_with_spark_does_not_mutate_user_pre_execution_commands( + self, mock_spark_image, mock_session + ): + user_commands = ["echo test"] + + settings = _JobSettings( + sagemaker_session=mock_session, + role="arn:aws:iam::123456789012:role/role", + instance_type="ml.m5.xlarge", + spark_config=SparkConfig(), + pre_execution_commands=user_commands, + ) + + assert user_commands == ["echo test"] + assert settings.pre_execution_commands[0] == "echo test" + assert any( + "sagemaker-feature-store-pyspark" in cmd for cmd in settings.pre_execution_commands + ) + + @patch.object(_JobSettings, "_get_default_spark_image", return_value="spark-image") + def test_init_with_spark_and_user_commands_and_script_raises_error( + self, mock_spark_image, mock_session + ): + with pytest.raises( + ValueError, match="Only one of pre_execution_commands or pre_execution_script" + ): + _JobSettings( + sagemaker_session=mock_session, + role="arn:aws:iam::123456789012:role/role", + instance_type="ml.m5.xlarge", + spark_config=SparkConfig(), + pre_execution_commands=["echo test"], + pre_execution_script="/path/to/script.sh", + ) + def test_init_without_instance_type_raises_error(self, mock_session): """Test that instance_type is required.""" with pytest.raises(ValueError, match="instance_type is a required parameter"):