feat: Support selective pipeline execution between function step and regular step#4392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
knikure merged 1 commit into from
Feb 29, 2024
Merged

feat: Support selective pipeline execution between function step and regular step #4392

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 4 additions & 11 deletions src/sagemaker/workflow/function_step.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -33,12 +33,11 @@
PipelineVariable,
)

from sagemaker.workflow.execution_variables import ExecutionVariables
from sagemaker.workflow.properties import Properties
from sagemaker.workflow.retry import RetryPolicy
from sagemaker.workflow.steps import Step, ConfigurableRetryStep, StepTypeEnum
from sagemaker.workflow.step_collections import StepCollection
from sagemaker.workflow.step_outputs import StepOutput
from sagemaker.workflow.step_outputs import StepOutput, get_step
from sagemaker.workflow.utilities import trim_request_dict, load_step_compilation_context

from sagemaker.s3_utils import s3_path_join
Expand DownExpand Up@@ -277,14 +276,12 @@ def _to_json_get(self) -> JsonGet:
"""Expression structure for workflow service calls using JsonGet resolution."""
from sagemaker.remote_function.core.stored_function import (
JSON_SERIALIZED_RESULT_KEY,
RESULTS_FOLDER,
JSON_RESULTS_FILE,
)

if not self._step.name:
raise ValueError("Step name is not defined.")

s3_root_uri = self._step._job_settings.s3_root_uri
# Resolve json path --
# Deserializer will be able to resolve a JsonGet using path "Return[1]" to
# access value 10 from following serialized JSON:
Expand All@@ -308,13 +305,9 @@ def _to_json_get(self) -> JsonGet:

return JsonGet(
s3_uri=Join(
"/",
[
s3_root_uri,
ExecutionVariables.PIPELINE_NAME,
ExecutionVariables.PIPELINE_EXECUTION_ID,
self._step.name,
RESULTS_FOLDER,
on="/",
values=[
get_step(self)._properties.OutputDataConfig.S3OutputPath,
Comment thread
qidewenwhen marked this conversation as resolved.
JSON_RESULTS_FILE,
],
),
Expand Down
8 changes: 4 additions & 4 deletions src/sagemaker/workflow/functions.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -21,7 +21,7 @@
from sagemaker.workflow.entities import PipelineVariable
from sagemaker.workflow.execution_variables import ExecutionVariable
from sagemaker.workflow.parameters import Parameter
from sagemaker.workflow.properties import PropertyFile
from sagemaker.workflow.properties import PropertyFile, Properties

if TYPE_CHECKING:
from sagemaker.workflow.steps import Step
Expand DownExpand Up@@ -172,9 +172,9 @@ def _validate_json_get_s3_uri(self):
for join_arg in s3_uri.values:
if not is_pipeline_variable(join_arg):
continue
if not isinstance(join_arg, (Parameter, ExecutionVariable)):
if not isinstance(join_arg, (Parameter, ExecutionVariable, Properties)):
raise ValueError(
f"Invalid JsonGet function {self.expr}. "
f"The Join values in JsonGet's s3_uri can only be a primitive object "
f"or Parameteror ExecutionVariable."
f"The Join values in JsonGet's s3_uri can only be a primitive object, "
f"Parameter, ExecutionVariable or Properties."
)
27 changes: 1 addition & 26 deletions src/sagemaker/workflow/pipeline.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -41,7 +41,6 @@
RESOURCE_NOT_FOUND_EXCEPTION,
EXECUTION_TIME_PIPELINE_PARAMETER_FORMAT,
)
from sagemaker.workflow.function_step import DelayedReturn
from sagemaker.workflow.lambda_step import LambdaOutput, LambdaStep
from sagemaker.workflow.entities import (
Expression,
Expand DownExpand Up@@ -725,10 +724,7 @@ def _interpolate(
pipeline_name (str): The name of the pipeline to be interpolated.
"""
if isinstance(obj, (Expression, Parameter, Properties, StepOutput)):
updated_obj = _replace_pipeline_name_in_json_get_s3_uri(
obj=obj, pipeline_name=pipeline_name
)
return updated_obj.expr
return obj.expr

if isinstance(obj, CallbackOutput):
step_name = callback_output_to_step_map[obj.output_name]
Expand DownExpand Up@@ -760,27 +756,6 @@ def _interpolate(
return new


# TODO: we should remove this once the ExecutionVariables.PIPELINE_NAME is fixed in backend
def _replace_pipeline_name_in_json_get_s3_uri(obj: Union[RequestType, Any], pipeline_name: str):
"""Replace the ExecutionVariables.PIPELINE_NAME in DelayedReturn's JsonGet s3_uri

with the pipeline_name, because ExecutionVariables.PIPELINE_NAME
is parsed as all lower-cased str in backend.
"""
if not isinstance(obj, DelayedReturn):
return obj

json_get = obj._to_json_get()

if not json_get.s3_uri:
return obj
# the s3 uri has to be a Join, which has been validated in JsonGet init
for i in range(len(json_get.s3_uri.values)):
if json_get.s3_uri.values[i] == ExecutionVariables.PIPELINE_NAME:
json_get.s3_uri.values[i] = pipeline_name
return json_get


def _map_callback_outputs(steps: List[Step]):
"""Iterate over the provided steps, building a map of callback output parameters to step names.

Expand Down
4 changes: 2 additions & 2 deletions tests/integ/sagemaker/workflow/helpers.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -33,7 +33,7 @@ def create_and_execute_pipeline(
region_name,
role,
no_of_steps,
last_step_name,
last_step_name_prefix,
execution_parameters,
step_status,
step_result_type=None,
Expand DownExpand Up@@ -66,7 +66,7 @@ def create_and_execute_pipeline(
len(execution_steps) == no_of_steps
), f"Expected {no_of_steps}, instead found {len(execution_steps)}"

assert last_step_name in execution_steps[0]["StepName"]
assert last_step_name_prefix in execution_steps[0]["StepName"]
assert execution_steps[0]["StepStatus"] == step_status
if step_result_type:
result = execution.result(execution_steps[0]["StepName"])
Expand Down
110 changes: 105 additions & 5 deletions tests/integ/sagemaker/workflow/test_selective_execution.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@

import pytest

from sagemaker.processing import ProcessingInput
from tests.integ import DATA_DIR
from sagemaker.sklearn import SKLearnProcessor
from sagemaker.workflow.step_outputs import get_step
Expand DownExpand Up@@ -84,7 +85,7 @@ def sum(a, b):
region_name=region_name,
role=role,
no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
execution_parameters=dict(),
step_status="Succeeded",
step_result_type=int,
Expand All@@ -97,7 +98,7 @@ def sum(a, b):
region_name=region_name,
role=role,
no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
execution_parameters=dict(),
step_status="Succeeded",
step_result_type=int,
Expand All@@ -115,7 +116,7 @@ def sum(a, b):
pass


def test_selective_execution_of_regular_step_depended_by_function_step(
def test_selective_execution_of_regular_step_referenced_by_function_step(
sagemaker_session,
role,
pipeline_name,
Expand DownExpand Up@@ -168,7 +169,7 @@ def func_2(arg):
region_name=region_name,
role=role,
no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
execution_parameters=dict(),
step_status="Succeeded",
step_result_type=str,
Expand All@@ -182,7 +183,7 @@ def func_2(arg):
region_name=region_name,
role=role,
no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
execution_parameters=dict(),
step_status="Succeeded",
step_result_type=str,
Expand All@@ -199,3 +200,102 @@ def func_2(arg):
pipeline.delete()
except Exception:
pass


def test_selective_execution_of_function_step_referenced_by_regular_step(
pipeline_session,
role,
pipeline_name,
region_name,
dummy_container_without_error,
sklearn_latest_version,
):
# Test Selective Pipeline Execution on function step -> [select: regular step]
os.environ["AWS_DEFAULT_REGION"] = region_name
processing_job_instance_counts = 2

@step(
name="step1",
role=role,
image_uri=dummy_container_without_error,
instance_type=INSTANCE_TYPE,
keep_alive_period_in_seconds=60,
)
def func(var: int):
return 1, var

step_output = func(processing_job_instance_counts)

script_path = os.path.join(DATA_DIR, "dummy_script.py")
input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
inputs = [
ProcessingInput(source=input_file_path, destination="/opt/ml/processing/inputs/"),
]

sklearn_processor = SKLearnProcessor(
framework_version=sklearn_latest_version,
role=role,
instance_type=INSTANCE_TYPE,
instance_count=step_output[1],
command=["python3"],
sagemaker_session=pipeline_session,
base_job_name="test-sklearn",
)

step_args = sklearn_processor.run(
inputs=inputs,
code=script_path,
)
process_step = ProcessingStep(
name="MyProcessStep",
step_args=step_args,
)

pipeline = Pipeline(
name=pipeline_name,
steps=[process_step],
sagemaker_session=pipeline_session,
)

try:
execution, _ = create_and_execute_pipeline(
pipeline=pipeline,
pipeline_name=pipeline_name,
region_name=region_name,
role=role,
no_of_steps=2,
last_step_name_prefix=process_step.name,
execution_parameters=dict(),
step_status="Succeeded",
wait_duration=1000, # seconds
)

_, execution_steps2 = create_and_execute_pipeline(
pipeline=pipeline,
pipeline_name=pipeline_name,
region_name=region_name,
role=role,
no_of_steps=2,
last_step_name_prefix=process_step.name,
execution_parameters=dict(),
step_status="Succeeded",
wait_duration=1000, # seconds
selective_execution_config=SelectiveExecutionConfig(
source_pipeline_execution_arn=execution.arn,
selected_steps=[process_step.name],
),
)

execution_proc_job = pipeline_session.describe_processing_job(
execution_steps2[0]["Metadata"]["ProcessingJob"]["Arn"].split("/")[-1]
)
assert (
execution_proc_job["ProcessingResources"]["ClusterConfig"]["InstanceCount"]
== processing_job_instance_counts
)

finally:
try:
pipeline.delete()
except Exception:
pass
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
feat: Support selective pipeline execution between function step and regular step#4392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knikure
merged 1 commit into

from

Feb 29, 2024
 Merged


feat: Support selective pipeline execution between function step and regular step
#4392
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


15 changes: 4 additions & 11 deletions
src/sagemaker/workflow/function_step.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,12 +33,11 @@
 PipelineVariable,
)

from sagemaker.workflow.execution_variables import ExecutionVariables
from sagemaker.workflow.properties import Properties
from sagemaker.workflow.retry import RetryPolicy
from sagemaker.workflow.steps import Step, ConfigurableRetryStep, StepTypeEnum
from sagemaker.workflow.step_collections import StepCollection
from sagemaker.workflow.step_outputs import StepOutput
from sagemaker.workflow.step_outputs import StepOutput, get_step
from sagemaker.workflow.utilities import trim_request_dict, load_step_compilation_context

from sagemaker.s3_utils import s3_path_join
Expand DownExpand Up@@ -277,14 +276,12 @@ def _to_json_get(self) -> JsonGet:
 """Expression structure for workflow service calls using JsonGet resolution."""
 from sagemaker.remote_function.core.stored_function import (
 JSON_SERIALIZED_RESULT_KEY,
 RESULTS_FOLDER,
 JSON_RESULTS_FILE,
 )

 if not self._step.name:
 raise ValueError("Step name is not defined.")

 s3_root_uri = self._step._job_settings.s3_root_uri
 # Resolve json path --
 # Deserializer will be able to resolve a JsonGet using path "Return[1]" to
 # access value 10 from following serialized JSON:
Expand All@@ -308,13 +305,9 @@ def _to_json_get(self) -> JsonGet:

 return JsonGet(
 s3_uri=Join(
 "/",
 [
 s3_root_uri,
 ExecutionVariables.PIPELINE_NAME,
 ExecutionVariables.PIPELINE_EXECUTION_ID,
 self._step.name,
 RESULTS_FOLDER,
 on="/",
 values=[
 get_step(self)._properties.OutputDataConfig.S3OutputPath,
Comment thread
qidewenwhen marked this conversation as resolved.

 JSON_RESULTS_FILE,
 ],
 ),
Expand Down

8 changes: 4 additions & 4 deletions
src/sagemaker/workflow/functions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,7 +21,7 @@
from sagemaker.workflow.entities import PipelineVariable
from sagemaker.workflow.execution_variables import ExecutionVariable
from sagemaker.workflow.parameters import Parameter
from sagemaker.workflow.properties import PropertyFile
from sagemaker.workflow.properties import PropertyFile, Properties

if TYPE_CHECKING:
 from sagemaker.workflow.steps import Step
Expand DownExpand Up@@ -172,9 +172,9 @@ def _validate_json_get_s3_uri(self):
 for join_arg in s3_uri.values:
 if not is_pipeline_variable(join_arg):
 continue
 if not isinstance(join_arg, (Parameter, ExecutionVariable)):
 if not isinstance(join_arg, (Parameter, ExecutionVariable, Properties)):
 raise ValueError(
 f"Invalid JsonGet function {self.expr}. "
 f"The Join values in JsonGet's s3_uri can only be a primitive object "
 f"or Parameteror ExecutionVariable."
 f"The Join values in JsonGet's s3_uri can only be a primitive object, "
 f"Parameter, ExecutionVariable or Properties."
 )

27 changes: 1 addition & 26 deletions
src/sagemaker/workflow/pipeline.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,7 +41,6 @@
 RESOURCE_NOT_FOUND_EXCEPTION,
 EXECUTION_TIME_PIPELINE_PARAMETER_FORMAT,
)
from sagemaker.workflow.function_step import DelayedReturn
from sagemaker.workflow.lambda_step import LambdaOutput, LambdaStep
from sagemaker.workflow.entities import (
 Expression,
Expand DownExpand Up@@ -725,10 +724,7 @@ def _interpolate(
 pipeline_name (str): The name of the pipeline to be interpolated.
 """
 if isinstance(obj, (Expression, Parameter, Properties, StepOutput)):
 updated_obj = _replace_pipeline_name_in_json_get_s3_uri(
 obj=obj, pipeline_name=pipeline_name
 )
 return updated_obj.expr
 return obj.expr

 if isinstance(obj, CallbackOutput):
 step_name = callback_output_to_step_map[obj.output_name]
Expand DownExpand Up@@ -760,27 +756,6 @@ def _interpolate(
 return new


# TODO: we should remove this once the ExecutionVariables.PIPELINE_NAME is fixed in backend
def _replace_pipeline_name_in_json_get_s3_uri(obj: Union[RequestType, Any], pipeline_name: str):
 """Replace the ExecutionVariables.PIPELINE_NAME in DelayedReturn's JsonGet s3_uri

 with the pipeline_name, because ExecutionVariables.PIPELINE_NAME
 is parsed as all lower-cased str in backend.
 """
 if not isinstance(obj, DelayedReturn):
 return obj

 json_get = obj._to_json_get()

 if not json_get.s3_uri:
 return obj
 # the s3 uri has to be a Join, which has been validated in JsonGet init
 for i in range(len(json_get.s3_uri.values)):
 if json_get.s3_uri.values[i] == ExecutionVariables.PIPELINE_NAME:
 json_get.s3_uri.values[i] = pipeline_name
 return json_get


def _map_callback_outputs(steps: List[Step]):
 """Iterate over the provided steps, building a map of callback output parameters to step names.

Expand Down

4 changes: 2 additions & 2 deletions
tests/integ/sagemaker/workflow/helpers.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,7 +33,7 @@ def create_and_execute_pipeline(
 region_name,
 role,
 no_of_steps,
last_step_name,
last_step_name_prefix,
 execution_parameters,
 step_status,
 step_result_type=None,
Expand DownExpand Up@@ -66,7 +66,7 @@ def create_and_execute_pipeline(
 len(execution_steps) == no_of_steps
 ), f"Expected {no_of_steps}, instead found {len(execution_steps)}"

 assert last_step_name in execution_steps[0]["StepName"]
 assert last_step_name_prefix in execution_steps[0]["StepName"]
 assert execution_steps[0]["StepStatus"] == step_status
 if step_result_type:
 result = execution.result(execution_steps[0]["StepName"])
Expand Down

110 changes: 105 additions & 5 deletions
tests/integ/sagemaker/workflow/test_selective_execution.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@

import pytest

from sagemaker.processing import ProcessingInput
from tests.integ import DATA_DIR
from sagemaker.sklearn import SKLearnProcessor
from sagemaker.workflow.step_outputs import get_step
Expand DownExpand Up@@ -84,7 +85,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -97,7 +98,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -115,7 +116,7 @@ def sum(a, b):
 pass


def test_selective_execution_of_regular_step_depended_by_function_step(
def test_selective_execution_of_regular_step_referenced_by_function_step(
 sagemaker_session,
 role,
 pipeline_name,
Expand DownExpand Up@@ -168,7 +169,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -182,7 +183,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -199,3 +200,102 @@ def func_2(arg):
 pipeline.delete()
 except Exception:
 pass


def test_selective_execution_of_function_step_referenced_by_regular_step(
 pipeline_session,
 role,
 pipeline_name,
 region_name,
 dummy_container_without_error,
 sklearn_latest_version,
):
 # Test Selective Pipeline Execution on function step -> [select: regular step]
 os.environ["AWS_DEFAULT_REGION"] = region_name
 processing_job_instance_counts = 2

 @step(
 name="step1",
 role=role,
 image_uri=dummy_container_without_error,
 instance_type=INSTANCE_TYPE,
 keep_alive_period_in_seconds=60,
 )
 def func(var: int):
 return 1, var

 step_output = func(processing_job_instance_counts)

 script_path = os.path.join(DATA_DIR, "dummy_script.py")
 input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 inputs = [
 ProcessingInput(source=input_file_path, destination="/opt/ml/processing/inputs/"),
 ]

 sklearn_processor = SKLearnProcessor(
 framework_version=sklearn_latest_version,
 role=role,
 instance_type=INSTANCE_TYPE,
 instance_count=step_output[1],
 command=["python3"],
 sagemaker_session=pipeline_session,
 base_job_name="test-sklearn",
 )

 step_args = sklearn_processor.run(
 inputs=inputs,
 code=script_path,
 )
 process_step = ProcessingStep(
 name="MyProcessStep",
 step_args=step_args,
 )

 pipeline = Pipeline(
 name=pipeline_name,
 steps=[process_step],
 sagemaker_session=pipeline_session,
 )

 try:
 execution, _ = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 )

 _, execution_steps2 = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 selective_execution_config=SelectiveExecutionConfig(
 source_pipeline_execution_arn=execution.arn,
 selected_steps=[process_step.name],
 ),
 )

 execution_proc_job = pipeline_session.describe_processing_job(
 execution_steps2[0]["Metadata"]["ProcessingJob"]["Arn"].split("/")[-1]
 )
 assert (
 execution_proc_job["ProcessingResources"]["ClusterConfig"]["InstanceCount"]
 == processing_job_instance_counts
 )

 finally:
 try:
 pipeline.delete()
 except Exception:
 pass
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
feat: Support selective pipeline execution between function step and regular step#4392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knikure
merged 1 commit into

from

Feb 29, 2024
 Merged


feat: Support selective pipeline execution between function step and regular step
#4392
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


15 changes: 4 additions & 11 deletions
src/sagemaker/workflow/function_step.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,12 +33,11 @@
 PipelineVariable,
)

from sagemaker.workflow.execution_variables import ExecutionVariables
from sagemaker.workflow.properties import Properties
from sagemaker.workflow.retry import RetryPolicy
from sagemaker.workflow.steps import Step, ConfigurableRetryStep, StepTypeEnum
from sagemaker.workflow.step_collections import StepCollection
from sagemaker.workflow.step_outputs import StepOutput
from sagemaker.workflow.step_outputs import StepOutput, get_step
from sagemaker.workflow.utilities import trim_request_dict, load_step_compilation_context

from sagemaker.s3_utils import s3_path_join
Expand DownExpand Up@@ -277,14 +276,12 @@ def _to_json_get(self) -> JsonGet:
 """Expression structure for workflow service calls using JsonGet resolution."""
 from sagemaker.remote_function.core.stored_function import (
 JSON_SERIALIZED_RESULT_KEY,
 RESULTS_FOLDER,
 JSON_RESULTS_FILE,
 )

 if not self._step.name:
 raise ValueError("Step name is not defined.")

 s3_root_uri = self._step._job_settings.s3_root_uri
 # Resolve json path --
 # Deserializer will be able to resolve a JsonGet using path "Return[1]" to
 # access value 10 from following serialized JSON:
Expand All@@ -308,13 +305,9 @@ def _to_json_get(self) -> JsonGet:

 return JsonGet(
 s3_uri=Join(
 "/",
 [
 s3_root_uri,
 ExecutionVariables.PIPELINE_NAME,
 ExecutionVariables.PIPELINE_EXECUTION_ID,
 self._step.name,
 RESULTS_FOLDER,
 on="/",
 values=[
 get_step(self)._properties.OutputDataConfig.S3OutputPath,
Comment thread
qidewenwhen marked this conversation as resolved.

 JSON_RESULTS_FILE,
 ],
 ),
Expand Down

8 changes: 4 additions & 4 deletions
src/sagemaker/workflow/functions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,7 +21,7 @@
from sagemaker.workflow.entities import PipelineVariable
from sagemaker.workflow.execution_variables import ExecutionVariable
from sagemaker.workflow.parameters import Parameter
from sagemaker.workflow.properties import PropertyFile
from sagemaker.workflow.properties import PropertyFile, Properties

if TYPE_CHECKING:
 from sagemaker.workflow.steps import Step
Expand DownExpand Up@@ -172,9 +172,9 @@ def _validate_json_get_s3_uri(self):
 for join_arg in s3_uri.values:
 if not is_pipeline_variable(join_arg):
 continue
 if not isinstance(join_arg, (Parameter, ExecutionVariable)):
 if not isinstance(join_arg, (Parameter, ExecutionVariable, Properties)):
 raise ValueError(
 f"Invalid JsonGet function {self.expr}. "
 f"The Join values in JsonGet's s3_uri can only be a primitive object "
 f"or Parameteror ExecutionVariable."
 f"The Join values in JsonGet's s3_uri can only be a primitive object, "
 f"Parameter, ExecutionVariable or Properties."
 )

27 changes: 1 addition & 26 deletions
src/sagemaker/workflow/pipeline.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,7 +41,6 @@
 RESOURCE_NOT_FOUND_EXCEPTION,
 EXECUTION_TIME_PIPELINE_PARAMETER_FORMAT,
)
from sagemaker.workflow.function_step import DelayedReturn
from sagemaker.workflow.lambda_step import LambdaOutput, LambdaStep
from sagemaker.workflow.entities import (
 Expression,
Expand DownExpand Up@@ -725,10 +724,7 @@ def _interpolate(
 pipeline_name (str): The name of the pipeline to be interpolated.
 """
 if isinstance(obj, (Expression, Parameter, Properties, StepOutput)):
 updated_obj = _replace_pipeline_name_in_json_get_s3_uri(
 obj=obj, pipeline_name=pipeline_name
 )
 return updated_obj.expr
 return obj.expr

 if isinstance(obj, CallbackOutput):
 step_name = callback_output_to_step_map[obj.output_name]
Expand DownExpand Up@@ -760,27 +756,6 @@ def _interpolate(
 return new


# TODO: we should remove this once the ExecutionVariables.PIPELINE_NAME is fixed in backend
def _replace_pipeline_name_in_json_get_s3_uri(obj: Union[RequestType, Any], pipeline_name: str):
 """Replace the ExecutionVariables.PIPELINE_NAME in DelayedReturn's JsonGet s3_uri

 with the pipeline_name, because ExecutionVariables.PIPELINE_NAME
 is parsed as all lower-cased str in backend.
 """
 if not isinstance(obj, DelayedReturn):
 return obj

 json_get = obj._to_json_get()

 if not json_get.s3_uri:
 return obj
 # the s3 uri has to be a Join, which has been validated in JsonGet init
 for i in range(len(json_get.s3_uri.values)):
 if json_get.s3_uri.values[i] == ExecutionVariables.PIPELINE_NAME:
 json_get.s3_uri.values[i] = pipeline_name
 return json_get


def _map_callback_outputs(steps: List[Step]):
 """Iterate over the provided steps, building a map of callback output parameters to step names.

Expand Down

4 changes: 2 additions & 2 deletions
tests/integ/sagemaker/workflow/helpers.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,7 +33,7 @@ def create_and_execute_pipeline(
 region_name,
 role,
 no_of_steps,
last_step_name,
last_step_name_prefix,
 execution_parameters,
 step_status,
 step_result_type=None,
Expand DownExpand Up@@ -66,7 +66,7 @@ def create_and_execute_pipeline(
 len(execution_steps) == no_of_steps
 ), f"Expected {no_of_steps}, instead found {len(execution_steps)}"

 assert last_step_name in execution_steps[0]["StepName"]
 assert last_step_name_prefix in execution_steps[0]["StepName"]
 assert execution_steps[0]["StepStatus"] == step_status
 if step_result_type:
 result = execution.result(execution_steps[0]["StepName"])
Expand Down

110 changes: 105 additions & 5 deletions
tests/integ/sagemaker/workflow/test_selective_execution.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@

import pytest

from sagemaker.processing import ProcessingInput
from tests.integ import DATA_DIR
from sagemaker.sklearn import SKLearnProcessor
from sagemaker.workflow.step_outputs import get_step
Expand DownExpand Up@@ -84,7 +85,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -97,7 +98,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -115,7 +116,7 @@ def sum(a, b):
 pass


def test_selective_execution_of_regular_step_depended_by_function_step(
def test_selective_execution_of_regular_step_referenced_by_function_step(
 sagemaker_session,
 role,
 pipeline_name,
Expand DownExpand Up@@ -168,7 +169,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -182,7 +183,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -199,3 +200,102 @@ def func_2(arg):
 pipeline.delete()
 except Exception:
 pass


def test_selective_execution_of_function_step_referenced_by_regular_step(
 pipeline_session,
 role,
 pipeline_name,
 region_name,
 dummy_container_without_error,
 sklearn_latest_version,
):
 # Test Selective Pipeline Execution on function step -> [select: regular step]
 os.environ["AWS_DEFAULT_REGION"] = region_name
 processing_job_instance_counts = 2

 @step(
 name="step1",
 role=role,
 image_uri=dummy_container_without_error,
 instance_type=INSTANCE_TYPE,
 keep_alive_period_in_seconds=60,
 )
 def func(var: int):
 return 1, var

 step_output = func(processing_job_instance_counts)

 script_path = os.path.join(DATA_DIR, "dummy_script.py")
 input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 inputs = [
 ProcessingInput(source=input_file_path, destination="/opt/ml/processing/inputs/"),
 ]

 sklearn_processor = SKLearnProcessor(
 framework_version=sklearn_latest_version,
 role=role,
 instance_type=INSTANCE_TYPE,
 instance_count=step_output[1],
 command=["python3"],
 sagemaker_session=pipeline_session,
 base_job_name="test-sklearn",
 )

 step_args = sklearn_processor.run(
 inputs=inputs,
 code=script_path,
 )
 process_step = ProcessingStep(
 name="MyProcessStep",
 step_args=step_args,
 )

 pipeline = Pipeline(
 name=pipeline_name,
 steps=[process_step],
 sagemaker_session=pipeline_session,
 )

 try:
 execution, _ = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 )

 _, execution_steps2 = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 selective_execution_config=SelectiveExecutionConfig(
 source_pipeline_execution_arn=execution.arn,
 selected_steps=[process_step.name],
 ),
 )

 execution_proc_job = pipeline_session.describe_processing_job(
 execution_steps2[0]["Metadata"]["ProcessingJob"]["Arn"].split("/")[-1]
 )
 assert (
 execution_proc_job["ProcessingResources"]["ClusterConfig"]["InstanceCount"]
 == processing_job_instance_counts
 )

 finally:
 try:
 pipeline.delete()
 except Exception:
 pass
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
feat: Support selective pipeline execution between function step and regular step#4392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knikure
merged 1 commit into

from

Feb 29, 2024
 Merged


feat: Support selective pipeline execution between function step and regular step
#4392
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


15 changes: 4 additions & 11 deletions
src/sagemaker/workflow/function_step.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,12 +33,11 @@
 PipelineVariable,
)

from sagemaker.workflow.execution_variables import ExecutionVariables
from sagemaker.workflow.properties import Properties
from sagemaker.workflow.retry import RetryPolicy
from sagemaker.workflow.steps import Step, ConfigurableRetryStep, StepTypeEnum
from sagemaker.workflow.step_collections import StepCollection
from sagemaker.workflow.step_outputs import StepOutput
from sagemaker.workflow.step_outputs import StepOutput, get_step
from sagemaker.workflow.utilities import trim_request_dict, load_step_compilation_context

from sagemaker.s3_utils import s3_path_join
Expand DownExpand Up@@ -277,14 +276,12 @@ def _to_json_get(self) -> JsonGet:
 """Expression structure for workflow service calls using JsonGet resolution."""
 from sagemaker.remote_function.core.stored_function import (
 JSON_SERIALIZED_RESULT_KEY,
 RESULTS_FOLDER,
 JSON_RESULTS_FILE,
 )

 if not self._step.name:
 raise ValueError("Step name is not defined.")

 s3_root_uri = self._step._job_settings.s3_root_uri
 # Resolve json path --
 # Deserializer will be able to resolve a JsonGet using path "Return[1]" to
 # access value 10 from following serialized JSON:
Expand All@@ -308,13 +305,9 @@ def _to_json_get(self) -> JsonGet:

 return JsonGet(
 s3_uri=Join(
 "/",
 [
 s3_root_uri,
 ExecutionVariables.PIPELINE_NAME,
 ExecutionVariables.PIPELINE_EXECUTION_ID,
 self._step.name,
 RESULTS_FOLDER,
 on="/",
 values=[
 get_step(self)._properties.OutputDataConfig.S3OutputPath,
Comment thread
qidewenwhen marked this conversation as resolved.

 JSON_RESULTS_FILE,
 ],
 ),
Expand Down

8 changes: 4 additions & 4 deletions
src/sagemaker/workflow/functions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,7 +21,7 @@
from sagemaker.workflow.entities import PipelineVariable
from sagemaker.workflow.execution_variables import ExecutionVariable
from sagemaker.workflow.parameters import Parameter
from sagemaker.workflow.properties import PropertyFile
from sagemaker.workflow.properties import PropertyFile, Properties

if TYPE_CHECKING:
 from sagemaker.workflow.steps import Step
Expand DownExpand Up@@ -172,9 +172,9 @@ def _validate_json_get_s3_uri(self):
 for join_arg in s3_uri.values:
 if not is_pipeline_variable(join_arg):
 continue
 if not isinstance(join_arg, (Parameter, ExecutionVariable)):
 if not isinstance(join_arg, (Parameter, ExecutionVariable, Properties)):
 raise ValueError(
 f"Invalid JsonGet function {self.expr}. "
 f"The Join values in JsonGet's s3_uri can only be a primitive object "
 f"or Parameteror ExecutionVariable."
 f"The Join values in JsonGet's s3_uri can only be a primitive object, "
 f"Parameter, ExecutionVariable or Properties."
 )

27 changes: 1 addition & 26 deletions
src/sagemaker/workflow/pipeline.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,7 +41,6 @@
 RESOURCE_NOT_FOUND_EXCEPTION,
 EXECUTION_TIME_PIPELINE_PARAMETER_FORMAT,
)
from sagemaker.workflow.function_step import DelayedReturn
from sagemaker.workflow.lambda_step import LambdaOutput, LambdaStep
from sagemaker.workflow.entities import (
 Expression,
Expand DownExpand Up@@ -725,10 +724,7 @@ def _interpolate(
 pipeline_name (str): The name of the pipeline to be interpolated.
 """
 if isinstance(obj, (Expression, Parameter, Properties, StepOutput)):
 updated_obj = _replace_pipeline_name_in_json_get_s3_uri(
 obj=obj, pipeline_name=pipeline_name
 )
 return updated_obj.expr
 return obj.expr

 if isinstance(obj, CallbackOutput):
 step_name = callback_output_to_step_map[obj.output_name]
Expand DownExpand Up@@ -760,27 +756,6 @@ def _interpolate(
 return new


# TODO: we should remove this once the ExecutionVariables.PIPELINE_NAME is fixed in backend
def _replace_pipeline_name_in_json_get_s3_uri(obj: Union[RequestType, Any], pipeline_name: str):
 """Replace the ExecutionVariables.PIPELINE_NAME in DelayedReturn's JsonGet s3_uri

 with the pipeline_name, because ExecutionVariables.PIPELINE_NAME
 is parsed as all lower-cased str in backend.
 """
 if not isinstance(obj, DelayedReturn):
 return obj

 json_get = obj._to_json_get()

 if not json_get.s3_uri:
 return obj
 # the s3 uri has to be a Join, which has been validated in JsonGet init
 for i in range(len(json_get.s3_uri.values)):
 if json_get.s3_uri.values[i] == ExecutionVariables.PIPELINE_NAME:
 json_get.s3_uri.values[i] = pipeline_name
 return json_get


def _map_callback_outputs(steps: List[Step]):
 """Iterate over the provided steps, building a map of callback output parameters to step names.

Expand Down

4 changes: 2 additions & 2 deletions
tests/integ/sagemaker/workflow/helpers.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,7 +33,7 @@ def create_and_execute_pipeline(
 region_name,
 role,
 no_of_steps,
last_step_name,
last_step_name_prefix,
 execution_parameters,
 step_status,
 step_result_type=None,
Expand DownExpand Up@@ -66,7 +66,7 @@ def create_and_execute_pipeline(
 len(execution_steps) == no_of_steps
 ), f"Expected {no_of_steps}, instead found {len(execution_steps)}"

 assert last_step_name in execution_steps[0]["StepName"]
 assert last_step_name_prefix in execution_steps[0]["StepName"]
 assert execution_steps[0]["StepStatus"] == step_status
 if step_result_type:
 result = execution.result(execution_steps[0]["StepName"])
Expand Down

110 changes: 105 additions & 5 deletions
tests/integ/sagemaker/workflow/test_selective_execution.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@

import pytest

from sagemaker.processing import ProcessingInput
from tests.integ import DATA_DIR
from sagemaker.sklearn import SKLearnProcessor
from sagemaker.workflow.step_outputs import get_step
Expand DownExpand Up@@ -84,7 +85,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -97,7 +98,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -115,7 +116,7 @@ def sum(a, b):
 pass


def test_selective_execution_of_regular_step_depended_by_function_step(
def test_selective_execution_of_regular_step_referenced_by_function_step(
 sagemaker_session,
 role,
 pipeline_name,
Expand DownExpand Up@@ -168,7 +169,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -182,7 +183,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -199,3 +200,102 @@ def func_2(arg):
 pipeline.delete()
 except Exception:
 pass


def test_selective_execution_of_function_step_referenced_by_regular_step(
 pipeline_session,
 role,
 pipeline_name,
 region_name,
 dummy_container_without_error,
 sklearn_latest_version,
):
 # Test Selective Pipeline Execution on function step -> [select: regular step]
 os.environ["AWS_DEFAULT_REGION"] = region_name
 processing_job_instance_counts = 2

 @step(
 name="step1",
 role=role,
 image_uri=dummy_container_without_error,
 instance_type=INSTANCE_TYPE,
 keep_alive_period_in_seconds=60,
 )
 def func(var: int):
 return 1, var

 step_output = func(processing_job_instance_counts)

 script_path = os.path.join(DATA_DIR, "dummy_script.py")
 input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 inputs = [
 ProcessingInput(source=input_file_path, destination="/opt/ml/processing/inputs/"),
 ]

 sklearn_processor = SKLearnProcessor(
 framework_version=sklearn_latest_version,
 role=role,
 instance_type=INSTANCE_TYPE,
 instance_count=step_output[1],
 command=["python3"],
 sagemaker_session=pipeline_session,
 base_job_name="test-sklearn",
 )

 step_args = sklearn_processor.run(
 inputs=inputs,
 code=script_path,
 )
 process_step = ProcessingStep(
 name="MyProcessStep",
 step_args=step_args,
 )

 pipeline = Pipeline(
 name=pipeline_name,
 steps=[process_step],
 sagemaker_session=pipeline_session,
 )

 try:
 execution, _ = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 )

 _, execution_steps2 = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 selective_execution_config=SelectiveExecutionConfig(
 source_pipeline_execution_arn=execution.arn,
 selected_steps=[process_step.name],
 ),
 )

 execution_proc_job = pipeline_session.describe_processing_job(
 execution_steps2[0]["Metadata"]["ProcessingJob"]["Arn"].split("/")[-1]
 )
 assert (
 execution_proc_job["ProcessingResources"]["ClusterConfig"]["InstanceCount"]
 == processing_job_instance_counts
 )

 finally:
 try:
 pipeline.delete()
 except Exception:
 pass
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
feat: Support selective pipeline execution between function step and regular step#4392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knikure
merged 1 commit into

from

Feb 29, 2024
 Merged


feat: Support selective pipeline execution between function step and regular step
#4392
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


15 changes: 4 additions & 11 deletions
src/sagemaker/workflow/function_step.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,12 +33,11 @@
 PipelineVariable,
)

from sagemaker.workflow.execution_variables import ExecutionVariables
from sagemaker.workflow.properties import Properties
from sagemaker.workflow.retry import RetryPolicy
from sagemaker.workflow.steps import Step, ConfigurableRetryStep, StepTypeEnum
from sagemaker.workflow.step_collections import StepCollection
from sagemaker.workflow.step_outputs import StepOutput
from sagemaker.workflow.step_outputs import StepOutput, get_step
from sagemaker.workflow.utilities import trim_request_dict, load_step_compilation_context

from sagemaker.s3_utils import s3_path_join
Expand DownExpand Up@@ -277,14 +276,12 @@ def _to_json_get(self) -> JsonGet:
 """Expression structure for workflow service calls using JsonGet resolution."""
 from sagemaker.remote_function.core.stored_function import (
 JSON_SERIALIZED_RESULT_KEY,
 RESULTS_FOLDER,
 JSON_RESULTS_FILE,
 )

 if not self._step.name:
 raise ValueError("Step name is not defined.")

 s3_root_uri = self._step._job_settings.s3_root_uri
 # Resolve json path --
 # Deserializer will be able to resolve a JsonGet using path "Return[1]" to
 # access value 10 from following serialized JSON:
Expand All@@ -308,13 +305,9 @@ def _to_json_get(self) -> JsonGet:

 return JsonGet(
 s3_uri=Join(
 "/",
 [
 s3_root_uri,
 ExecutionVariables.PIPELINE_NAME,
 ExecutionVariables.PIPELINE_EXECUTION_ID,
 self._step.name,
 RESULTS_FOLDER,
 on="/",
 values=[
 get_step(self)._properties.OutputDataConfig.S3OutputPath,
Comment thread
qidewenwhen marked this conversation as resolved.

 JSON_RESULTS_FILE,
 ],
 ),
Expand Down

8 changes: 4 additions & 4 deletions
src/sagemaker/workflow/functions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,7 +21,7 @@
from sagemaker.workflow.entities import PipelineVariable
from sagemaker.workflow.execution_variables import ExecutionVariable
from sagemaker.workflow.parameters import Parameter
from sagemaker.workflow.properties import PropertyFile
from sagemaker.workflow.properties import PropertyFile, Properties

if TYPE_CHECKING:
 from sagemaker.workflow.steps import Step
Expand DownExpand Up@@ -172,9 +172,9 @@ def _validate_json_get_s3_uri(self):
 for join_arg in s3_uri.values:
 if not is_pipeline_variable(join_arg):
 continue
 if not isinstance(join_arg, (Parameter, ExecutionVariable)):
 if not isinstance(join_arg, (Parameter, ExecutionVariable, Properties)):
 raise ValueError(
 f"Invalid JsonGet function {self.expr}. "
 f"The Join values in JsonGet's s3_uri can only be a primitive object "
 f"or Parameteror ExecutionVariable."
 f"The Join values in JsonGet's s3_uri can only be a primitive object, "
 f"Parameter, ExecutionVariable or Properties."
 )

27 changes: 1 addition & 26 deletions
src/sagemaker/workflow/pipeline.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,7 +41,6 @@
 RESOURCE_NOT_FOUND_EXCEPTION,
 EXECUTION_TIME_PIPELINE_PARAMETER_FORMAT,
)
from sagemaker.workflow.function_step import DelayedReturn
from sagemaker.workflow.lambda_step import LambdaOutput, LambdaStep
from sagemaker.workflow.entities import (
 Expression,
Expand DownExpand Up@@ -725,10 +724,7 @@ def _interpolate(
 pipeline_name (str): The name of the pipeline to be interpolated.
 """
 if isinstance(obj, (Expression, Parameter, Properties, StepOutput)):
 updated_obj = _replace_pipeline_name_in_json_get_s3_uri(
 obj=obj, pipeline_name=pipeline_name
 )
 return updated_obj.expr
 return obj.expr

 if isinstance(obj, CallbackOutput):
 step_name = callback_output_to_step_map[obj.output_name]
Expand DownExpand Up@@ -760,27 +756,6 @@ def _interpolate(
 return new


# TODO: we should remove this once the ExecutionVariables.PIPELINE_NAME is fixed in backend
def _replace_pipeline_name_in_json_get_s3_uri(obj: Union[RequestType, Any], pipeline_name: str):
 """Replace the ExecutionVariables.PIPELINE_NAME in DelayedReturn's JsonGet s3_uri

 with the pipeline_name, because ExecutionVariables.PIPELINE_NAME
 is parsed as all lower-cased str in backend.
 """
 if not isinstance(obj, DelayedReturn):
 return obj

 json_get = obj._to_json_get()

 if not json_get.s3_uri:
 return obj
 # the s3 uri has to be a Join, which has been validated in JsonGet init
 for i in range(len(json_get.s3_uri.values)):
 if json_get.s3_uri.values[i] == ExecutionVariables.PIPELINE_NAME:
 json_get.s3_uri.values[i] = pipeline_name
 return json_get


def _map_callback_outputs(steps: List[Step]):
 """Iterate over the provided steps, building a map of callback output parameters to step names.

Expand Down

4 changes: 2 additions & 2 deletions
tests/integ/sagemaker/workflow/helpers.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,7 +33,7 @@ def create_and_execute_pipeline(
 region_name,
 role,
 no_of_steps,
last_step_name,
last_step_name_prefix,
 execution_parameters,
 step_status,
 step_result_type=None,
Expand DownExpand Up@@ -66,7 +66,7 @@ def create_and_execute_pipeline(
 len(execution_steps) == no_of_steps
 ), f"Expected {no_of_steps}, instead found {len(execution_steps)}"

 assert last_step_name in execution_steps[0]["StepName"]
 assert last_step_name_prefix in execution_steps[0]["StepName"]
 assert execution_steps[0]["StepStatus"] == step_status
 if step_result_type:
 result = execution.result(execution_steps[0]["StepName"])
Expand Down

110 changes: 105 additions & 5 deletions
tests/integ/sagemaker/workflow/test_selective_execution.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@

import pytest

from sagemaker.processing import ProcessingInput
from tests.integ import DATA_DIR
from sagemaker.sklearn import SKLearnProcessor
from sagemaker.workflow.step_outputs import get_step
Expand DownExpand Up@@ -84,7 +85,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -97,7 +98,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -115,7 +116,7 @@ def sum(a, b):
 pass


def test_selective_execution_of_regular_step_depended_by_function_step(
def test_selective_execution_of_regular_step_referenced_by_function_step(
 sagemaker_session,
 role,
 pipeline_name,
Expand DownExpand Up@@ -168,7 +169,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -182,7 +183,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -199,3 +200,102 @@ def func_2(arg):
 pipeline.delete()
 except Exception:
 pass


def test_selective_execution_of_function_step_referenced_by_regular_step(
 pipeline_session,
 role,
 pipeline_name,
 region_name,
 dummy_container_without_error,
 sklearn_latest_version,
):
 # Test Selective Pipeline Execution on function step -> [select: regular step]
 os.environ["AWS_DEFAULT_REGION"] = region_name
 processing_job_instance_counts = 2

 @step(
 name="step1",
 role=role,
 image_uri=dummy_container_without_error,
 instance_type=INSTANCE_TYPE,
 keep_alive_period_in_seconds=60,
 )
 def func(var: int):
 return 1, var

 step_output = func(processing_job_instance_counts)

 script_path = os.path.join(DATA_DIR, "dummy_script.py")
 input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 inputs = [
 ProcessingInput(source=input_file_path, destination="/opt/ml/processing/inputs/"),
 ]

 sklearn_processor = SKLearnProcessor(
 framework_version=sklearn_latest_version,
 role=role,
 instance_type=INSTANCE_TYPE,
 instance_count=step_output[1],
 command=["python3"],
 sagemaker_session=pipeline_session,
 base_job_name="test-sklearn",
 )

 step_args = sklearn_processor.run(
 inputs=inputs,
 code=script_path,
 )
 process_step = ProcessingStep(
 name="MyProcessStep",
 step_args=step_args,
 )

 pipeline = Pipeline(
 name=pipeline_name,
 steps=[process_step],
 sagemaker_session=pipeline_session,
 )

 try:
 execution, _ = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 )

 _, execution_steps2 = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 selective_execution_config=SelectiveExecutionConfig(
 source_pipeline_execution_arn=execution.arn,
 selected_steps=[process_step.name],
 ),
 )

 execution_proc_job = pipeline_session.describe_processing_job(
 execution_steps2[0]["Metadata"]["ProcessingJob"]["Arn"].split("/")[-1]
 )
 assert (
 execution_proc_job["ProcessingResources"]["ClusterConfig"]["InstanceCount"]
 == processing_job_instance_counts
 )

 finally:
 try:
 pipeline.delete()
 except Exception:
 pass
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
feat: Support selective pipeline execution between function step and regular step#4392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knikure
merged 1 commit into

from

Feb 29, 2024
 Merged


feat: Support selective pipeline execution between function step and regular step
#4392
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


15 changes: 4 additions & 11 deletions
src/sagemaker/workflow/function_step.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,12 +33,11 @@
 PipelineVariable,
)

from sagemaker.workflow.execution_variables import ExecutionVariables
from sagemaker.workflow.properties import Properties
from sagemaker.workflow.retry import RetryPolicy
from sagemaker.workflow.steps import Step, ConfigurableRetryStep, StepTypeEnum
from sagemaker.workflow.step_collections import StepCollection
from sagemaker.workflow.step_outputs import StepOutput
from sagemaker.workflow.step_outputs import StepOutput, get_step
from sagemaker.workflow.utilities import trim_request_dict, load_step_compilation_context

from sagemaker.s3_utils import s3_path_join
Expand DownExpand Up@@ -277,14 +276,12 @@ def _to_json_get(self) -> JsonGet:
 """Expression structure for workflow service calls using JsonGet resolution."""
 from sagemaker.remote_function.core.stored_function import (
 JSON_SERIALIZED_RESULT_KEY,
 RESULTS_FOLDER,
 JSON_RESULTS_FILE,
 )

 if not self._step.name:
 raise ValueError("Step name is not defined.")

 s3_root_uri = self._step._job_settings.s3_root_uri
 # Resolve json path --
 # Deserializer will be able to resolve a JsonGet using path "Return[1]" to
 # access value 10 from following serialized JSON:
Expand All@@ -308,13 +305,9 @@ def _to_json_get(self) -> JsonGet:

 return JsonGet(
 s3_uri=Join(
 "/",
 [
 s3_root_uri,
 ExecutionVariables.PIPELINE_NAME,
 ExecutionVariables.PIPELINE_EXECUTION_ID,
 self._step.name,
 RESULTS_FOLDER,
 on="/",
 values=[
 get_step(self)._properties.OutputDataConfig.S3OutputPath,
Comment thread
qidewenwhen marked this conversation as resolved.

 JSON_RESULTS_FILE,
 ],
 ),
Expand Down

8 changes: 4 additions & 4 deletions
src/sagemaker/workflow/functions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,7 +21,7 @@
from sagemaker.workflow.entities import PipelineVariable
from sagemaker.workflow.execution_variables import ExecutionVariable
from sagemaker.workflow.parameters import Parameter
from sagemaker.workflow.properties import PropertyFile
from sagemaker.workflow.properties import PropertyFile, Properties

if TYPE_CHECKING:
 from sagemaker.workflow.steps import Step
Expand DownExpand Up@@ -172,9 +172,9 @@ def _validate_json_get_s3_uri(self):
 for join_arg in s3_uri.values:
 if not is_pipeline_variable(join_arg):
 continue
 if not isinstance(join_arg, (Parameter, ExecutionVariable)):
 if not isinstance(join_arg, (Parameter, ExecutionVariable, Properties)):
 raise ValueError(
 f"Invalid JsonGet function {self.expr}. "
 f"The Join values in JsonGet's s3_uri can only be a primitive object "
 f"or Parameteror ExecutionVariable."
 f"The Join values in JsonGet's s3_uri can only be a primitive object, "
 f"Parameter, ExecutionVariable or Properties."
 )

27 changes: 1 addition & 26 deletions
src/sagemaker/workflow/pipeline.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,7 +41,6 @@
 RESOURCE_NOT_FOUND_EXCEPTION,
 EXECUTION_TIME_PIPELINE_PARAMETER_FORMAT,
)
from sagemaker.workflow.function_step import DelayedReturn
from sagemaker.workflow.lambda_step import LambdaOutput, LambdaStep
from sagemaker.workflow.entities import (
 Expression,
Expand DownExpand Up@@ -725,10 +724,7 @@ def _interpolate(
 pipeline_name (str): The name of the pipeline to be interpolated.
 """
 if isinstance(obj, (Expression, Parameter, Properties, StepOutput)):
 updated_obj = _replace_pipeline_name_in_json_get_s3_uri(
 obj=obj, pipeline_name=pipeline_name
 )
 return updated_obj.expr
 return obj.expr

 if isinstance(obj, CallbackOutput):
 step_name = callback_output_to_step_map[obj.output_name]
Expand DownExpand Up@@ -760,27 +756,6 @@ def _interpolate(
 return new


# TODO: we should remove this once the ExecutionVariables.PIPELINE_NAME is fixed in backend
def _replace_pipeline_name_in_json_get_s3_uri(obj: Union[RequestType, Any], pipeline_name: str):
 """Replace the ExecutionVariables.PIPELINE_NAME in DelayedReturn's JsonGet s3_uri

 with the pipeline_name, because ExecutionVariables.PIPELINE_NAME
 is parsed as all lower-cased str in backend.
 """
 if not isinstance(obj, DelayedReturn):
 return obj

 json_get = obj._to_json_get()

 if not json_get.s3_uri:
 return obj
 # the s3 uri has to be a Join, which has been validated in JsonGet init
 for i in range(len(json_get.s3_uri.values)):
 if json_get.s3_uri.values[i] == ExecutionVariables.PIPELINE_NAME:
 json_get.s3_uri.values[i] = pipeline_name
 return json_get


def _map_callback_outputs(steps: List[Step]):
 """Iterate over the provided steps, building a map of callback output parameters to step names.

Expand Down

4 changes: 2 additions & 2 deletions
tests/integ/sagemaker/workflow/helpers.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,7 +33,7 @@ def create_and_execute_pipeline(
 region_name,
 role,
 no_of_steps,
last_step_name,
last_step_name_prefix,
 execution_parameters,
 step_status,
 step_result_type=None,
Expand DownExpand Up@@ -66,7 +66,7 @@ def create_and_execute_pipeline(
 len(execution_steps) == no_of_steps
 ), f"Expected {no_of_steps}, instead found {len(execution_steps)}"

 assert last_step_name in execution_steps[0]["StepName"]
 assert last_step_name_prefix in execution_steps[0]["StepName"]
 assert execution_steps[0]["StepStatus"] == step_status
 if step_result_type:
 result = execution.result(execution_steps[0]["StepName"])
Expand Down

110 changes: 105 additions & 5 deletions
tests/integ/sagemaker/workflow/test_selective_execution.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@

import pytest

from sagemaker.processing import ProcessingInput
from tests.integ import DATA_DIR
from sagemaker.sklearn import SKLearnProcessor
from sagemaker.workflow.step_outputs import get_step
Expand DownExpand Up@@ -84,7 +85,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -97,7 +98,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -115,7 +116,7 @@ def sum(a, b):
 pass


def test_selective_execution_of_regular_step_depended_by_function_step(
def test_selective_execution_of_regular_step_referenced_by_function_step(
 sagemaker_session,
 role,
 pipeline_name,
Expand DownExpand Up@@ -168,7 +169,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -182,7 +183,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -199,3 +200,102 @@ def func_2(arg):
 pipeline.delete()
 except Exception:
 pass


def test_selective_execution_of_function_step_referenced_by_regular_step(
 pipeline_session,
 role,
 pipeline_name,
 region_name,
 dummy_container_without_error,
 sklearn_latest_version,
):
 # Test Selective Pipeline Execution on function step -> [select: regular step]
 os.environ["AWS_DEFAULT_REGION"] = region_name
 processing_job_instance_counts = 2

 @step(
 name="step1",
 role=role,
 image_uri=dummy_container_without_error,
 instance_type=INSTANCE_TYPE,
 keep_alive_period_in_seconds=60,
 )
 def func(var: int):
 return 1, var

 step_output = func(processing_job_instance_counts)

 script_path = os.path.join(DATA_DIR, "dummy_script.py")
 input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 inputs = [
 ProcessingInput(source=input_file_path, destination="/opt/ml/processing/inputs/"),
 ]

 sklearn_processor = SKLearnProcessor(
 framework_version=sklearn_latest_version,
 role=role,
 instance_type=INSTANCE_TYPE,
 instance_count=step_output[1],
 command=["python3"],
 sagemaker_session=pipeline_session,
 base_job_name="test-sklearn",
 )

 step_args = sklearn_processor.run(
 inputs=inputs,
 code=script_path,
 )
 process_step = ProcessingStep(
 name="MyProcessStep",
 step_args=step_args,
 )

 pipeline = Pipeline(
 name=pipeline_name,
 steps=[process_step],
 sagemaker_session=pipeline_session,
 )

 try:
 execution, _ = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 )

 _, execution_steps2 = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 selective_execution_config=SelectiveExecutionConfig(
 source_pipeline_execution_arn=execution.arn,
 selected_steps=[process_step.name],
 ),
 )

 execution_proc_job = pipeline_session.describe_processing_job(
 execution_steps2[0]["Metadata"]["ProcessingJob"]["Arn"].split("/")[-1]
 )
 assert (
 execution_proc_job["ProcessingResources"]["ClusterConfig"]["InstanceCount"]
 == processing_job_instance_counts
 )

 finally:
 try:
 pipeline.delete()
 except Exception:
 pass
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
feat: Support selective pipeline execution between function step and regular step#4392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knikure
merged 1 commit into

from

Feb 29, 2024
 Merged


feat: Support selective pipeline execution between function step and regular step
#4392
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


15 changes: 4 additions & 11 deletions
src/sagemaker/workflow/function_step.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,12 +33,11 @@
 PipelineVariable,
)

from sagemaker.workflow.execution_variables import ExecutionVariables
from sagemaker.workflow.properties import Properties
from sagemaker.workflow.retry import RetryPolicy
from sagemaker.workflow.steps import Step, ConfigurableRetryStep, StepTypeEnum
from sagemaker.workflow.step_collections import StepCollection
from sagemaker.workflow.step_outputs import StepOutput
from sagemaker.workflow.step_outputs import StepOutput, get_step
from sagemaker.workflow.utilities import trim_request_dict, load_step_compilation_context

from sagemaker.s3_utils import s3_path_join
Expand DownExpand Up@@ -277,14 +276,12 @@ def _to_json_get(self) -> JsonGet:
 """Expression structure for workflow service calls using JsonGet resolution."""
 from sagemaker.remote_function.core.stored_function import (
 JSON_SERIALIZED_RESULT_KEY,
 RESULTS_FOLDER,
 JSON_RESULTS_FILE,
 )

 if not self._step.name:
 raise ValueError("Step name is not defined.")

 s3_root_uri = self._step._job_settings.s3_root_uri
 # Resolve json path --
 # Deserializer will be able to resolve a JsonGet using path "Return[1]" to
 # access value 10 from following serialized JSON:
Expand All@@ -308,13 +305,9 @@ def _to_json_get(self) -> JsonGet:

 return JsonGet(
 s3_uri=Join(
 "/",
 [
 s3_root_uri,
 ExecutionVariables.PIPELINE_NAME,
 ExecutionVariables.PIPELINE_EXECUTION_ID,
 self._step.name,
 RESULTS_FOLDER,
 on="/",
 values=[
 get_step(self)._properties.OutputDataConfig.S3OutputPath,
Comment thread
qidewenwhen marked this conversation as resolved.

 JSON_RESULTS_FILE,
 ],
 ),
Expand Down

8 changes: 4 additions & 4 deletions
src/sagemaker/workflow/functions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,7 +21,7 @@
from sagemaker.workflow.entities import PipelineVariable
from sagemaker.workflow.execution_variables import ExecutionVariable
from sagemaker.workflow.parameters import Parameter
from sagemaker.workflow.properties import PropertyFile
from sagemaker.workflow.properties import PropertyFile, Properties

if TYPE_CHECKING:
 from sagemaker.workflow.steps import Step
Expand DownExpand Up@@ -172,9 +172,9 @@ def _validate_json_get_s3_uri(self):
 for join_arg in s3_uri.values:
 if not is_pipeline_variable(join_arg):
 continue
 if not isinstance(join_arg, (Parameter, ExecutionVariable)):
 if not isinstance(join_arg, (Parameter, ExecutionVariable, Properties)):
 raise ValueError(
 f"Invalid JsonGet function {self.expr}. "
 f"The Join values in JsonGet's s3_uri can only be a primitive object "
 f"or Parameteror ExecutionVariable."
 f"The Join values in JsonGet's s3_uri can only be a primitive object, "
 f"Parameter, ExecutionVariable or Properties."
 )

27 changes: 1 addition & 26 deletions
src/sagemaker/workflow/pipeline.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,7 +41,6 @@
 RESOURCE_NOT_FOUND_EXCEPTION,
 EXECUTION_TIME_PIPELINE_PARAMETER_FORMAT,
)
from sagemaker.workflow.function_step import DelayedReturn
from sagemaker.workflow.lambda_step import LambdaOutput, LambdaStep
from sagemaker.workflow.entities import (
 Expression,
Expand DownExpand Up@@ -725,10 +724,7 @@ def _interpolate(
 pipeline_name (str): The name of the pipeline to be interpolated.
 """
 if isinstance(obj, (Expression, Parameter, Properties, StepOutput)):
 updated_obj = _replace_pipeline_name_in_json_get_s3_uri(
 obj=obj, pipeline_name=pipeline_name
 )
 return updated_obj.expr
 return obj.expr

 if isinstance(obj, CallbackOutput):
 step_name = callback_output_to_step_map[obj.output_name]
Expand DownExpand Up@@ -760,27 +756,6 @@ def _interpolate(
 return new


# TODO: we should remove this once the ExecutionVariables.PIPELINE_NAME is fixed in backend
def _replace_pipeline_name_in_json_get_s3_uri(obj: Union[RequestType, Any], pipeline_name: str):
 """Replace the ExecutionVariables.PIPELINE_NAME in DelayedReturn's JsonGet s3_uri

 with the pipeline_name, because ExecutionVariables.PIPELINE_NAME
 is parsed as all lower-cased str in backend.
 """
 if not isinstance(obj, DelayedReturn):
 return obj

 json_get = obj._to_json_get()

 if not json_get.s3_uri:
 return obj
 # the s3 uri has to be a Join, which has been validated in JsonGet init
 for i in range(len(json_get.s3_uri.values)):
 if json_get.s3_uri.values[i] == ExecutionVariables.PIPELINE_NAME:
 json_get.s3_uri.values[i] = pipeline_name
 return json_get


def _map_callback_outputs(steps: List[Step]):
 """Iterate over the provided steps, building a map of callback output parameters to step names.

Expand Down

4 changes: 2 additions & 2 deletions
tests/integ/sagemaker/workflow/helpers.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,7 +33,7 @@ def create_and_execute_pipeline(
 region_name,
 role,
 no_of_steps,
last_step_name,
last_step_name_prefix,
 execution_parameters,
 step_status,
 step_result_type=None,
Expand DownExpand Up@@ -66,7 +66,7 @@ def create_and_execute_pipeline(
 len(execution_steps) == no_of_steps
 ), f"Expected {no_of_steps}, instead found {len(execution_steps)}"

 assert last_step_name in execution_steps[0]["StepName"]
 assert last_step_name_prefix in execution_steps[0]["StepName"]
 assert execution_steps[0]["StepStatus"] == step_status
 if step_result_type:
 result = execution.result(execution_steps[0]["StepName"])
Expand Down

110 changes: 105 additions & 5 deletions
tests/integ/sagemaker/workflow/test_selective_execution.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@

import pytest

from sagemaker.processing import ProcessingInput
from tests.integ import DATA_DIR
from sagemaker.sklearn import SKLearnProcessor
from sagemaker.workflow.step_outputs import get_step
Expand DownExpand Up@@ -84,7 +85,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -97,7 +98,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -115,7 +116,7 @@ def sum(a, b):
 pass


def test_selective_execution_of_regular_step_depended_by_function_step(
def test_selective_execution_of_regular_step_referenced_by_function_step(
 sagemaker_session,
 role,
 pipeline_name,
Expand DownExpand Up@@ -168,7 +169,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -182,7 +183,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -199,3 +200,102 @@ def func_2(arg):
 pipeline.delete()
 except Exception:
 pass


def test_selective_execution_of_function_step_referenced_by_regular_step(
 pipeline_session,
 role,
 pipeline_name,
 region_name,
 dummy_container_without_error,
 sklearn_latest_version,
):
 # Test Selective Pipeline Execution on function step -> [select: regular step]
 os.environ["AWS_DEFAULT_REGION"] = region_name
 processing_job_instance_counts = 2

 @step(
 name="step1",
 role=role,
 image_uri=dummy_container_without_error,
 instance_type=INSTANCE_TYPE,
 keep_alive_period_in_seconds=60,
 )
 def func(var: int):
 return 1, var

 step_output = func(processing_job_instance_counts)

 script_path = os.path.join(DATA_DIR, "dummy_script.py")
 input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 inputs = [
 ProcessingInput(source=input_file_path, destination="/opt/ml/processing/inputs/"),
 ]

 sklearn_processor = SKLearnProcessor(
 framework_version=sklearn_latest_version,
 role=role,
 instance_type=INSTANCE_TYPE,
 instance_count=step_output[1],
 command=["python3"],
 sagemaker_session=pipeline_session,
 base_job_name="test-sklearn",
 )

 step_args = sklearn_processor.run(
 inputs=inputs,
 code=script_path,
 )
 process_step = ProcessingStep(
 name="MyProcessStep",
 step_args=step_args,
 )

 pipeline = Pipeline(
 name=pipeline_name,
 steps=[process_step],
 sagemaker_session=pipeline_session,
 )

 try:
 execution, _ = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 )

 _, execution_steps2 = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 selective_execution_config=SelectiveExecutionConfig(
 source_pipeline_execution_arn=execution.arn,
 selected_steps=[process_step.name],
 ),
 )

 execution_proc_job = pipeline_session.describe_processing_job(
 execution_steps2[0]["Metadata"]["ProcessingJob"]["Arn"].split("/")[-1]
 )
 assert (
 execution_proc_job["ProcessingResources"]["ClusterConfig"]["InstanceCount"]
 == processing_job_instance_counts
 )

 finally:
 try:
 pipeline.delete()
 except Exception:
 pass
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
feat: Support selective pipeline execution between function step and regular step#4392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knikure
merged 1 commit into

from

Feb 29, 2024
 Merged


feat: Support selective pipeline execution between function step and regular step
#4392
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


15 changes: 4 additions & 11 deletions
src/sagemaker/workflow/function_step.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,12 +33,11 @@
 PipelineVariable,
)

from sagemaker.workflow.execution_variables import ExecutionVariables
from sagemaker.workflow.properties import Properties
from sagemaker.workflow.retry import RetryPolicy
from sagemaker.workflow.steps import Step, ConfigurableRetryStep, StepTypeEnum
from sagemaker.workflow.step_collections import StepCollection
from sagemaker.workflow.step_outputs import StepOutput
from sagemaker.workflow.step_outputs import StepOutput, get_step
from sagemaker.workflow.utilities import trim_request_dict, load_step_compilation_context

from sagemaker.s3_utils import s3_path_join
Expand DownExpand Up@@ -277,14 +276,12 @@ def _to_json_get(self) -> JsonGet:
 """Expression structure for workflow service calls using JsonGet resolution."""
 from sagemaker.remote_function.core.stored_function import (
 JSON_SERIALIZED_RESULT_KEY,
 RESULTS_FOLDER,
 JSON_RESULTS_FILE,
 )

 if not self._step.name:
 raise ValueError("Step name is not defined.")

 s3_root_uri = self._step._job_settings.s3_root_uri
 # Resolve json path --
 # Deserializer will be able to resolve a JsonGet using path "Return[1]" to
 # access value 10 from following serialized JSON:
Expand All@@ -308,13 +305,9 @@ def _to_json_get(self) -> JsonGet:

 return JsonGet(
 s3_uri=Join(
 "/",
 [
 s3_root_uri,
 ExecutionVariables.PIPELINE_NAME,
 ExecutionVariables.PIPELINE_EXECUTION_ID,
 self._step.name,
 RESULTS_FOLDER,
 on="/",
 values=[
 get_step(self)._properties.OutputDataConfig.S3OutputPath,
Comment thread
qidewenwhen marked this conversation as resolved.

 JSON_RESULTS_FILE,
 ],
 ),
Expand Down

8 changes: 4 additions & 4 deletions
src/sagemaker/workflow/functions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -21,7 +21,7 @@
from sagemaker.workflow.entities import PipelineVariable
from sagemaker.workflow.execution_variables import ExecutionVariable
from sagemaker.workflow.parameters import Parameter
from sagemaker.workflow.properties import PropertyFile
from sagemaker.workflow.properties import PropertyFile, Properties

if TYPE_CHECKING:
 from sagemaker.workflow.steps import Step
Expand DownExpand Up@@ -172,9 +172,9 @@ def _validate_json_get_s3_uri(self):
 for join_arg in s3_uri.values:
 if not is_pipeline_variable(join_arg):
 continue
 if not isinstance(join_arg, (Parameter, ExecutionVariable)):
 if not isinstance(join_arg, (Parameter, ExecutionVariable, Properties)):
 raise ValueError(
 f"Invalid JsonGet function {self.expr}. "
 f"The Join values in JsonGet's s3_uri can only be a primitive object "
 f"or Parameteror ExecutionVariable."
 f"The Join values in JsonGet's s3_uri can only be a primitive object, "
 f"Parameter, ExecutionVariable or Properties."
 )

27 changes: 1 addition & 26 deletions
src/sagemaker/workflow/pipeline.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -41,7 +41,6 @@
 RESOURCE_NOT_FOUND_EXCEPTION,
 EXECUTION_TIME_PIPELINE_PARAMETER_FORMAT,
)
from sagemaker.workflow.function_step import DelayedReturn
from sagemaker.workflow.lambda_step import LambdaOutput, LambdaStep
from sagemaker.workflow.entities import (
 Expression,
Expand DownExpand Up@@ -725,10 +724,7 @@ def _interpolate(
 pipeline_name (str): The name of the pipeline to be interpolated.
 """
 if isinstance(obj, (Expression, Parameter, Properties, StepOutput)):
 updated_obj = _replace_pipeline_name_in_json_get_s3_uri(
 obj=obj, pipeline_name=pipeline_name
 )
 return updated_obj.expr
 return obj.expr

 if isinstance(obj, CallbackOutput):
 step_name = callback_output_to_step_map[obj.output_name]
Expand DownExpand Up@@ -760,27 +756,6 @@ def _interpolate(
 return new


# TODO: we should remove this once the ExecutionVariables.PIPELINE_NAME is fixed in backend
def _replace_pipeline_name_in_json_get_s3_uri(obj: Union[RequestType, Any], pipeline_name: str):
 """Replace the ExecutionVariables.PIPELINE_NAME in DelayedReturn's JsonGet s3_uri

 with the pipeline_name, because ExecutionVariables.PIPELINE_NAME
 is parsed as all lower-cased str in backend.
 """
 if not isinstance(obj, DelayedReturn):
 return obj

 json_get = obj._to_json_get()

 if not json_get.s3_uri:
 return obj
 # the s3 uri has to be a Join, which has been validated in JsonGet init
 for i in range(len(json_get.s3_uri.values)):
 if json_get.s3_uri.values[i] == ExecutionVariables.PIPELINE_NAME:
 json_get.s3_uri.values[i] = pipeline_name
 return json_get


def _map_callback_outputs(steps: List[Step]):
 """Iterate over the provided steps, building a map of callback output parameters to step names.

Expand Down

4 changes: 2 additions & 2 deletions
tests/integ/sagemaker/workflow/helpers.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -33,7 +33,7 @@ def create_and_execute_pipeline(
 region_name,
 role,
 no_of_steps,
last_step_name,
last_step_name_prefix,
 execution_parameters,
 step_status,
 step_result_type=None,
Expand DownExpand Up@@ -66,7 +66,7 @@ def create_and_execute_pipeline(
 len(execution_steps) == no_of_steps
 ), f"Expected {no_of_steps}, instead found {len(execution_steps)}"

 assert last_step_name in execution_steps[0]["StepName"]
 assert last_step_name_prefix in execution_steps[0]["StepName"]
 assert execution_steps[0]["StepStatus"] == step_status
 if step_result_type:
 result = execution.result(execution_steps[0]["StepName"])
Expand Down

110 changes: 105 additions & 5 deletions
tests/integ/sagemaker/workflow/test_selective_execution.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@

import pytest

from sagemaker.processing import ProcessingInput
from tests.integ import DATA_DIR
from sagemaker.sklearn import SKLearnProcessor
from sagemaker.workflow.step_outputs import get_step
Expand DownExpand Up@@ -84,7 +85,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -97,7 +98,7 @@ def sum(a, b):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="sum",
last_step_name_prefix="sum",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=int,
Expand All@@ -115,7 +116,7 @@ def sum(a, b):
 pass


def test_selective_execution_of_regular_step_depended_by_function_step(
def test_selective_execution_of_regular_step_referenced_by_function_step(
 sagemaker_session,
 role,
 pipeline_name,
Expand DownExpand Up@@ -168,7 +169,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -182,7 +183,7 @@ def func_2(arg):
 region_name=region_name,
 role=role,
 no_of_steps=2,
last_step_name="func",
last_step_name_prefix="func",
 execution_parameters=dict(),
 step_status="Succeeded",
 step_result_type=str,
Expand All@@ -199,3 +200,102 @@ def func_2(arg):
 pipeline.delete()
 except Exception:
 pass


def test_selective_execution_of_function_step_referenced_by_regular_step(
 pipeline_session,
 role,
 pipeline_name,
 region_name,
 dummy_container_without_error,
 sklearn_latest_version,
):
 # Test Selective Pipeline Execution on function step -> [select: regular step]
 os.environ["AWS_DEFAULT_REGION"] = region_name
 processing_job_instance_counts = 2

 @step(
 name="step1",
 role=role,
 image_uri=dummy_container_without_error,
 instance_type=INSTANCE_TYPE,
 keep_alive_period_in_seconds=60,
 )
 def func(var: int):
 return 1, var

 step_output = func(processing_job_instance_counts)

 script_path = os.path.join(DATA_DIR, "dummy_script.py")
 input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 inputs = [
 ProcessingInput(source=input_file_path, destination="/opt/ml/processing/inputs/"),
 ]

 sklearn_processor = SKLearnProcessor(
 framework_version=sklearn_latest_version,
 role=role,
 instance_type=INSTANCE_TYPE,
 instance_count=step_output[1],
 command=["python3"],
 sagemaker_session=pipeline_session,
 base_job_name="test-sklearn",
 )

 step_args = sklearn_processor.run(
 inputs=inputs,
 code=script_path,
 )
 process_step = ProcessingStep(
 name="MyProcessStep",
 step_args=step_args,
 )

 pipeline = Pipeline(
 name=pipeline_name,
 steps=[process_step],
 sagemaker_session=pipeline_session,
 )

 try:
 execution, _ = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 )

 _, execution_steps2 = create_and_execute_pipeline(
 pipeline=pipeline,
 pipeline_name=pipeline_name,
 region_name=region_name,
 role=role,
 no_of_steps=2,
 last_step_name_prefix=process_step.name,
 execution_parameters=dict(),
 step_status="Succeeded",
 wait_duration=1000, # seconds
 selective_execution_config=SelectiveExecutionConfig(
 source_pipeline_execution_arn=execution.arn,
 selected_steps=[process_step.name],
 ),
 )

 execution_proc_job = pipeline_session.describe_processing_job(
 execution_steps2[0]["Metadata"]["ProcessingJob"]["Arn"].split("/")[-1]
 )
 assert (
 execution_proc_job["ProcessingResources"]["ClusterConfig"]["InstanceCount"]
 == processing_job_instance_counts
 )

 finally:
 try:
 pipeline.delete()
 except Exception:
 pass
Loading