fix(core): let remote Spark jobs use pre_execution_script

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,

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"):