Skip to content

Commit d1dfbbb

Browse files
committed
fix(processing): Fix V3 shape regressions in processors
PySparkProcessor still built ProcessingInput/ProcessingOutput with the V2 source/destination keyword signature, so spark_event_logs_s3_uri and the configuration channel failed on V3. Build them with the V3 shapes (ProcessingS3Input/ProcessingS3Output) instead. Validate that submit_py_files, submit_jars and submit_files are lists and raise a clear ValueError otherwise. FrameworkProcessor.run(requirements=...) now threads the file name into the generated runproc.sh instead of hardcoding requirements.txt. Processor.run defaults kms_key to the configured output_kms_key so uploaded code is encrypted with the same key as job outputs. The requirements file is also honored on the custom entry_point (S3 source_dir) path, and a relative path such as reqs/cpu.txt is kept as-is because the source bundle preserves the source_dir layout. Integration tests (real jobs, submit-then-stop where the fix is in the request path) cover the code-upload KMS fallback, the FrameworkProcessor requirements file, and a V3 PySparkProcessor job ported from V2 that asserts the event-log output and the conf/py-files input shape; V3 had no Spark integ coverage before. Fixes #3809 Fixes #6253 Fixes #6252 Fixes #5805 Fixes #4874 --- X-AI-Prompt: Fix S-effort PySDK V3 bugs, processing-spark theme X-AI-Tool: Kiro
1 parent 7ac2708 commit d1dfbbb

11 files changed

Lines changed: 1806 additions & 19 deletions

File tree

‎sagemaker-core/src/sagemaker/core/processing.py‎

Lines changed: 42 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -332,7 +332,8 @@ def run(
332332
is built with :class:`~sagemaker.workflow.pipeline_context.PipelineSession`.
333333
However, the value of `TrialComponentDisplayName` is honored for display in Studio.
334334
kms_key (str): The ARN of the KMS key that is used to encrypt the
335-
user code file (default: None).
335+
user code file. If not provided, the processor's configured
336+
``output_kms_key`` is used (default: None).
336337
Returns:
337338
None or pipeline step arguments in case the Processor instance is built with
338339
:class:`~sagemaker.workflow.pipeline_context.PipelineSession`
@@ -343,6 +344,12 @@ def run(
343344
raise ValueError("""Logs can only be shown if wait is set to True.
344345
Please either set wait to True or set logs to False.""")
345346

347+
# When no explicit code-encryption key is given, fall back to the
348+
# configured output KMS key so the uploaded code and job outputs are
349+
# encrypted with the same key.
350+
if kms_key is None:
351+
kms_key = self.output_kms_key
352+
346353
normalized_inputs, normalized_outputs = self._normalize_args(
347354
job_name=job_name,
348355
arguments=arguments,
@@ -903,7 +910,8 @@ def run(
903910
is built with :class:`~sagemaker.workflow.pipeline_context.PipelineSession`.
904911
However, the value of `TrialComponentDisplayName` is honored for display in Studio.
905912
kms_key (str): The ARN of the KMS key that is used to encrypt the
906-
user code file (default: None).
913+
user code file. If not provided, the processor's configured
914+
``output_kms_key`` is used (default: None).
907915
Returns:
908916
None or pipeline step arguments in case the Processor instance is built with
909917
:class:`~sagemaker.workflow.pipeline_context.PipelineSession`
@@ -914,7 +922,7 @@ def run(
914922
inputs=inputs,
915923
outputs=outputs,
916924
code=code,
917-
kms_key=kms_key,
925+
kms_key=kms_key if kms_key is not None else self.output_kms_key,
918926
)
919927

920928
experiment_config = check_and_get_run_experiment_config(experiment_config)
@@ -1481,6 +1489,7 @@ def _pack_and_upload_code(
14811489
entry_point,
14821490
source_dir,
14831491
install_requirements_dir,
1492+
requirements=requirements,
14841493
)
14851494

14861495
return s3_runproc_sh, inputs, job_name
@@ -1526,13 +1535,14 @@ def _create_and_upload_runproc(
15261535
entry_point=None,
15271536
source_dir=None,
15281537
install_requirements_dir=None,
1538+
requirements=None,
15291539
):
15301540
"""Create runproc shell script and upload to S3 bucket."""
15311541
from sagemaker.core.workflow.utilities import _pipeline_config, hash_object
15321542

15331543
if _pipeline_config and _pipeline_config.pipeline_name:
15341544
runproc_file_str = self._generate_framework_script(
1535-
user_script, entry_point, source_dir, install_requirements_dir
1545+
user_script, entry_point, source_dir, install_requirements_dir, requirements
15361546
)
15371547
runproc_file_hash = hash_object(runproc_file_str)
15381548
s3_uri = s3.s3_path_join(
@@ -1551,7 +1561,7 @@ def _create_and_upload_runproc(
15511561
else:
15521562
s3_runproc_sh = s3.S3Uploader.upload_string_as_file_body(
15531563
self._generate_framework_script(
1554-
user_script, entry_point, source_dir, install_requirements_dir
1564+
user_script, entry_point, source_dir, install_requirements_dir, requirements
15551565
),
15561566
desired_s3_uri=entrypoint_s3_uri,
15571567
kms_key=kms_key,
@@ -1560,20 +1570,36 @@ def _create_and_upload_runproc(
15601570

15611571
return s3_runproc_sh
15621572

1573+
@staticmethod
1574+
def _requirements_file_in_container(requirements: Optional[str]) -> str:
1575+
"""Return the requirements path as it appears inside the extracted source bundle.
1576+
1577+
``requirements`` is documented as relative to ``source_dir`` and ``_package_code``
1578+
preserves the directory layout, so a relative path (``reqs/cpu.txt``) is kept. An
1579+
absolute path cannot be located inside the bundle, so only its basename is used.
1580+
"""
1581+
if not requirements:
1582+
return "requirements.txt"
1583+
if os.path.isabs(requirements):
1584+
return os.path.basename(requirements)
1585+
return os.path.normpath(requirements).replace(os.sep, "/")
1586+
15631587
def _generate_framework_script(
15641588
self,
15651589
user_script: str,
15661590
entry_point: str = None,
15671591
source_dir: str = None,
15681592
install_requirements_dir: str = None,
1593+
requirements: str = None,
15691594
) -> str:
15701595
"""Generate the framework entrypoint file (as text) for a processing job."""
15711596
if entry_point:
15721597
return self._generate_custom_framework_script(
1573-
user_script, entry_point, source_dir, install_requirements_dir
1598+
user_script, entry_point, source_dir, install_requirements_dir, requirements
15741599
)
15751600

15761601
install_requirements_dir = install_requirements_dir or self._SOURCE_CODE_CONTAINER_DIR
1602+
requirements_file = self._requirements_file_in_container(requirements)
15771603

15781604
return dedent("""\
15791605
#!/bin/bash
@@ -1597,16 +1623,17 @@ def _generate_framework_script(
15971623
exit 1
15981624
fi
15991625
1600-
if [[ -f 'requirements.txt' ]]; then
1626+
if [[ -f '{requirements_file}' ]]; then
16011627
# Some py3 containers has typing, which may breaks pip install
16021628
pip uninstall --yes typing
16031629
1604-
python3 {install_requirements_dir}/install_requirements.py requirements.txt
1630+
python3 {install_requirements_dir}/install_requirements.py {requirements_file}
16051631
fi
16061632
16071633
{entry_point_command} {entry_point} "$@"
16081634
""").format(
16091635
install_requirements_dir=install_requirements_dir,
1636+
requirements_file=requirements_file,
16101637
entry_point_command=" ".join(self.command),
16111638
entry_point=user_script,
16121639
)
@@ -1617,6 +1644,7 @@ def _generate_custom_framework_script(
16171644
entry_point: str,
16181645
source_dir: str = None,
16191646
install_requirements_dir: str = None,
1647+
requirements: str = None,
16201648
) -> str:
16211649
"""Generate a custom framework script with a user-provided entrypoint embedded.
16221650
@@ -1630,6 +1658,8 @@ def _generate_custom_framework_script(
16301658
is relative, it will be combined with source_dir.
16311659
install_requirements_dir (str): Container directory that holds
16321660
``install_requirements.py`` (default: the extracted source code dir).
1661+
requirements (str): Path to the requirements file relative to source_dir
1662+
(default: ``requirements.txt``).
16331663
16341664
Returns:
16351665
str: The generated script content
@@ -1639,6 +1669,7 @@ def _generate_custom_framework_script(
16391669
# source bundle on the container.
16401670
if self._is_s3_uri(source_dir):
16411671
install_requirements_dir = install_requirements_dir or self._SOURCE_CODE_CONTAINER_DIR
1672+
requirements_file = self._requirements_file_in_container(requirements)
16421673
return dedent("""\
16431674
#!/bin/bash
16441675
@@ -1655,9 +1686,9 @@ def _generate_custom_framework_script(
16551686
exit 1
16561687
fi
16571688
1658-
if [[ -f 'requirements.txt' ]]; then
1689+
if [[ -f '{requirements_file}' ]]; then
16591690
pip uninstall --yes typing
1660-
python3 {install_requirements_dir}/install_requirements.py requirements.txt
1691+
python3 {install_requirements_dir}/install_requirements.py {requirements_file}
16611692
fi
16621693
16631694
# Execute custom entrypoint
@@ -1667,6 +1698,7 @@ def _generate_custom_framework_script(
16671698
{entry_point_command} {user_script} "$@"
16681699
""").format(
16691700
install_requirements_dir=install_requirements_dir,
1701+
requirements_file=requirements_file,
16701702
entry_point=entry_point,
16711703
entry_point_command=" ".join(self.command),
16721704
user_script=user_script,

‎sagemaker-core/src/sagemaker/core/spark/processing.py‎

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
from sagemaker.core import s3
4040
from sagemaker.core.local.image import _ecr_login_if_needed, _pull_image
4141
from sagemaker.core.processing import ProcessingInput, ProcessingOutput, ScriptProcessor
42+
from sagemaker.core.shapes import ProcessingS3Input, ProcessingS3Output
4243
from sagemaker.core.s3 import S3Uploader
4344
from sagemaker.core.helper.session_helper import Session
4445
from sagemaker.core.network import NetworkConfig
@@ -330,9 +331,12 @@ def _extend_processing_args(self, inputs, outputs, **kwargs):
330331
)
331332

332333
output = ProcessingOutput(
333-
source=_SparkProcessorBase._spark_event_log_default_local_path,
334-
destination=spark_event_logs_s3_uri,
335-
s3_upload_mode="Continuous",
334+
output_name="spark-event-logs",
335+
s3_output=ProcessingS3Output(
336+
s3_uri=spark_event_logs_s3_uri,
337+
local_path=_SparkProcessorBase._spark_event_log_default_local_path,
338+
s3_upload_mode="Continuous",
339+
),
336340
)
337341

338342
extended_outputs.append(output)
@@ -444,9 +448,13 @@ def _stage_configuration(self, configuration):
444448
)
445449

446450
conf_input = ProcessingInput(
447-
source=s3_uri,
448-
destination=f"{self._conf_container_base_path}{self._conf_container_input_name}",
449451
input_name=_SparkProcessorBase._conf_container_input_name,
452+
s3_input=ProcessingS3Input(
453+
s3_uri=s3_uri,
454+
local_path=f"{self._conf_container_base_path}{self._conf_container_input_name}",
455+
s3_data_type="S3Prefix",
456+
s3_input_mode="File",
457+
),
450458
)
451459
return conf_input
452460

@@ -473,6 +481,11 @@ def _stage_submit_deps(self, submit_deps, input_channel_name):
473481
)
474482
if not input_channel_name:
475483
raise ValueError("input_channel_name value may not be empty.")
484+
if not isinstance(submit_deps, (list, tuple)):
485+
raise ValueError(
486+
f"submit_deps must be a list of one or more paths, but got "
487+
f"{type(submit_deps).__name__}. {self._submit_deps_error_message}"
488+
)
476489

477490
use_input_channel = False
478491
spark_opt_s3_uris = []
@@ -544,15 +557,20 @@ def _stage_submit_deps(self, submit_deps, input_channel_name):
544557
# them to the Spark container and form the spark-submit option from a
545558
# combination of S3 URIs and container's local input path
546559
if use_input_channel:
560+
input_channel_local_path = f"{self._conf_container_base_path}{input_channel_name}"
547561
input_channel = ProcessingInput(
548-
source=input_channel_s3_uri,
549-
destination=f"{self._conf_container_base_path}{input_channel_name}",
550562
input_name=input_channel_name,
563+
s3_input=ProcessingS3Input(
564+
s3_uri=input_channel_s3_uri,
565+
local_path=input_channel_local_path,
566+
s3_data_type="S3Prefix",
567+
s3_input_mode="File",
568+
),
551569
)
552570
spark_opt = (
553-
Join(on=",", values=spark_opt_s3_uris + [input_channel.destination])
571+
Join(on=",", values=spark_opt_s3_uris + [input_channel_local_path])
554572
if spark_opt_s3_uris_has_pipeline_var
555-
else ",".join(spark_opt_s3_uris + [input_channel.destination])
573+
else ",".join(spark_opt_s3_uris + [input_channel_local_path])
556574
)
557575
# If no local files were uploaded, form the spark-submit option from a list of S3 URIs
558576
else:
Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
import argparse
14+
import time
15+
16+
from pyspark.sql import SparkSession, SQLContext
17+
from pyspark.sql.functions import udf
18+
from pyspark.sql.types import IntegerType
19+
20+
# Import local module to test spark-submit--py-files dependencies
21+
import hello_py_spark_udfs as udfs
22+
23+
if __name__ == "__main__":
24+
parser = argparse.ArgumentParser(description="inputs and outputs")
25+
parser.add_argument("--input", type=str, help="path to input data")
26+
parser.add_argument("--output", required=False, type=str, help="path to output data")
27+
args = parser.parse_args()
28+
spark = SparkSession.builder.appName("SparkContainerTestApp").getOrCreate()
29+
sqlContext = SQLContext(spark.sparkContext)
30+
31+
# Load test data set
32+
inputPath = args.input
33+
salesDF = spark.read.json(inputPath)
34+
salesDF.printSchema()
35+
36+
salesDF.createOrReplaceTempView("sales")
37+
topDF = spark.sql("SELECT date, sale FROM sales WHERE sale > 750")
38+
# Show the first 20 rows of the dataframe
39+
topDF.show()
40+
time.sleep(60)
41+
42+
# Calculate average sales by date
43+
averageSalesPerDay = salesDF.groupBy("date").avg().collect()
44+
print(averageSalesPerDay)
45+
46+
outputPath = args.output
47+
48+
# Define a UDF that doubles an integer column
49+
# The UDF function is imported from local module to test spark-submit--py-files dependencies
50+
double_udf_int = udf(udfs.double_x, IntegerType())
51+
52+
# Save transformed data set to disk
53+
salesDF.select("date", "sale", double_udf_int("sale").alias("sale_double")).write.json(
54+
outputPath
55+
)
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
def double_x(x):
2+
return x + x

0 commit comments

Comments
 (0)