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