Skip to content

Commit 5f2e01d

Browse files
author
Dewen Qi
committed
change: reorganize integ test files for workflow
1 parent 6670e30 commit 5f2e01d

19 files changed

+3358
-3045
lines changed

‎tests/integ/__init__.py

-6
Original file line numberDiff line numberDiff line change
@@ -148,12 +148,6 @@
148148
"eu-west-2",
149149
"us-east-1",
150150
]
151-
NO_SM_PIPELINE_MM_CLARIFY_CHECK_STEP_REGIONS = [
152-
"ap-northeast-3",
153-
"ap-south-1",
154-
"eu-north-1",
155-
"sa-east-1",
156-
]
157151
EDGE_PACKAGING_SUPPORTED_REGIONS = [
158152
"us-east-2",
159153
"us-west-2",

‎tests/integ/sagemaker/__init__.py

Whitespace-only changes.

‎tests/integ/sagemaker/lineage/conftest.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
artifact,
2727
)
2828
from sagemaker.model import ModelPackage
29-
from tests.integ.test_workflow import test_end_to_end_pipeline_successful_execution
29+
from tests.integ.sagemaker.workflow.test_workflow import test_end_to_end_pipeline_successful_execution
3030
from sagemaker.workflow.pipeline import _PipelineExecution
3131
from sagemaker.session import get_execution_role
3232
from smexperiments import trial_component, trial, experiment

‎tests/integ/sagemaker/workflow/__init__.py

Whitespace-only changes.

‎tests/integ/sagemaker/workflow/test_callback_steps.py

+118
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
from __future__ import absolute_import
14+
15+
import re
16+
17+
import pytest
18+
19+
from sagemaker import get_execution_role, utils
20+
from sagemaker.workflow.callback_step import CallbackOutput, CallbackStep, CallbackOutputTypeEnum
21+
from sagemaker.workflow.parameters import ParameterInteger
22+
from sagemaker.workflow.pipeline import Pipeline
23+
24+
25+
@pytest.fixture
26+
def role(sagemaker_session):
27+
return get_execution_role(sagemaker_session)
28+
29+
30+
@pytest.fixture
31+
def pipeline_name():
32+
return utils.unique_name_from_base("my-pipeline-callback")
33+
34+
35+
@pytest.fixture
36+
def region_name(sagemaker_session):
37+
return sagemaker_session.boto_session.region_name
38+
39+
40+
def test_one_step_callback_pipeline(sagemaker_session, role, pipeline_name, region_name):
41+
instance_count = ParameterInteger(name="InstanceCount", default_value=2)
42+
43+
outputParam1 = CallbackOutput(output_name="output1", output_type=CallbackOutputTypeEnum.String)
44+
step_callback = CallbackStep(
45+
name="callback-step",
46+
sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
47+
inputs={"arg1": "foo"},
48+
outputs=[outputParam1],
49+
)
50+
51+
pipeline = Pipeline(
52+
name=pipeline_name,
53+
parameters=[instance_count],
54+
steps=[step_callback],
55+
sagemaker_session=sagemaker_session,
56+
)
57+
58+
try:
59+
response = pipeline.create(role)
60+
create_arn = response["PipelineArn"]
61+
assert re.match(
62+
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
63+
create_arn,
64+
)
65+
66+
pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
67+
response = pipeline.update(role)
68+
update_arn = response["PipelineArn"]
69+
assert re.match(
70+
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
71+
update_arn,
72+
)
73+
finally:
74+
try:
75+
pipeline.delete()
76+
except Exception:
77+
pass
78+
79+
80+
def test_two_step_callback_pipeline_with_output_reference(
81+
sagemaker_session, role, pipeline_name, region_name
82+
):
83+
instance_count = ParameterInteger(name="InstanceCount", default_value=2)
84+
85+
outputParam1 = CallbackOutput(output_name="output1", output_type=CallbackOutputTypeEnum.String)
86+
step_callback1 = CallbackStep(
87+
name="callback-step1",
88+
sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
89+
inputs={"arg1": "foo"},
90+
outputs=[outputParam1],
91+
)
92+
93+
step_callback2 = CallbackStep(
94+
name="callback-step2",
95+
sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
96+
inputs={"arg1": outputParam1},
97+
outputs=[],
98+
)
99+
100+
pipeline = Pipeline(
101+
name=pipeline_name,
102+
parameters=[instance_count],
103+
steps=[step_callback1, step_callback2],
104+
sagemaker_session=sagemaker_session,
105+
)
106+
107+
try:
108+
response = pipeline.create(role)
109+
create_arn = response["PipelineArn"]
110+
assert re.match(
111+
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
112+
create_arn,
113+
)
114+
finally:
115+
try:
116+
pipeline.delete()
117+
except Exception:
118+
pass

‎tests/integ/test_workflow_with_clarify_check_steps.py renamed to ‎tests/integ/sagemaker/workflow/test_clarify_check_steps.py

-9
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
import pytest
2020
from botocore.exceptions import WaiterError
2121

22-
import tests
2322
from sagemaker.clarify import (
2423
BiasConfig,
2524
DataConfig,
@@ -129,10 +128,6 @@ def data_bias_check_config(data_config, bias_config):
129128
)
130129

131130

132-
@pytest.mark.skipif(
133-
tests.integ.test_region() in tests.integ.NO_SM_PIPELINE_MM_CLARIFY_CHECK_STEP_REGIONS,
134-
reason=f"ClarifyCheckStep is not fully deployed in {tests.integ.test_region()}.",
135-
)
136131
def test_one_step_data_bias_pipeline_happycase(
137132
sagemaker_session,
138133
role,
@@ -220,10 +215,6 @@ def test_one_step_data_bias_pipeline_happycase(
220215
pass
221216

222217

223-
@pytest.mark.skipif(
224-
tests.integ.test_region() in tests.integ.NO_SM_PIPELINE_MM_CLARIFY_CHECK_STEP_REGIONS,
225-
reason=f"ClarifyCheckStep is not fully deployed in {tests.integ.test_region()}.",
226-
)
227218
def test_one_step_data_bias_pipeline_constraint_violation(
228219
sagemaker_session,
229220
role,

‎tests/integ/sagemaker/workflow/test_emr_steps.py

+82
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
from __future__ import absolute_import
14+
15+
import re
16+
17+
import pytest
18+
19+
from sagemaker import get_execution_role, utils
20+
from sagemaker.workflow.emr_step import EMRStep, EMRStepConfig
21+
from sagemaker.workflow.parameters import ParameterInteger
22+
from sagemaker.workflow.pipeline import Pipeline
23+
24+
25+
@pytest.fixture
26+
def role(sagemaker_session):
27+
return get_execution_role(sagemaker_session)
28+
29+
30+
@pytest.fixture
31+
def pipeline_name():
32+
return utils.unique_name_from_base("my-pipeline-emr")
33+
34+
35+
@pytest.fixture
36+
def region_name(sagemaker_session):
37+
return sagemaker_session.boto_session.region_name
38+
39+
40+
def test_two_steps_emr_pipeline(sagemaker_session, role, pipeline_name, region_name):
41+
instance_count = ParameterInteger(name="InstanceCount", default_value=2)
42+
43+
emr_step_config = EMRStepConfig(
44+
jar="s3://us-west-2.elasticmapreduce/libs/script-runner/script-runner.jar",
45+
args=["dummy_emr_script_path"],
46+
)
47+
48+
step_emr_1 = EMRStep(
49+
name="emr-step-1",
50+
cluster_id="j-1YONHTCP3YZKC",
51+
display_name="emr_step_1",
52+
description="MyEMRStepDescription",
53+
step_config=emr_step_config,
54+
)
55+
56+
step_emr_2 = EMRStep(
57+
name="emr-step-2",
58+
cluster_id=step_emr_1.properties.ClusterId,
59+
display_name="emr_step_2",
60+
description="MyEMRStepDescription",
61+
step_config=emr_step_config,
62+
)
63+
64+
pipeline = Pipeline(
65+
name=pipeline_name,
66+
parameters=[instance_count],
67+
steps=[step_emr_1, step_emr_2],
68+
sagemaker_session=sagemaker_session,
69+
)
70+
71+
try:
72+
response = pipeline.create(role)
73+
create_arn = response["PipelineArn"]
74+
assert re.match(
75+
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
76+
create_arn,
77+
)
78+
finally:
79+
try:
80+
pipeline.delete()
81+
except Exception:
82+
pass

‎tests/integ/test_workflow_experiment.py renamed to ‎tests/integ/sagemaker/workflow/test_experiment.py

File renamed without changes.

‎tests/integ/test_workflow_with_fail_steps.py renamed to ‎tests/integ/sagemaker/workflow/test_fail_steps.py

File renamed without changes.

‎tests/integ/sagemaker/workflow/test_lambda_steps.py

+132
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
from __future__ import absolute_import
14+
15+
import re
16+
17+
import pytest
18+
19+
from sagemaker import get_execution_role, utils
20+
from sagemaker.workflow.lambda_step import (
21+
LambdaStep,
22+
LambdaOutput,
23+
LambdaOutputTypeEnum,
24+
)
25+
from sagemaker.lambda_helper import Lambda
26+
from sagemaker.workflow.parameters import ParameterInteger
27+
from sagemaker.workflow.pipeline import Pipeline
28+
29+
30+
@pytest.fixture
31+
def role(sagemaker_session):
32+
return get_execution_role(sagemaker_session)
33+
34+
35+
@pytest.fixture
36+
def pipeline_name():
37+
return utils.unique_name_from_base("my-pipeline-lambda")
38+
39+
40+
@pytest.fixture
41+
def region_name(sagemaker_session):
42+
return sagemaker_session.boto_session.region_name
43+
44+
45+
def test_one_step_lambda_pipeline(sagemaker_session, role, pipeline_name, region_name):
46+
instance_count = ParameterInteger(name="InstanceCount", default_value=2)
47+
48+
outputParam1 = LambdaOutput(output_name="output1", output_type=LambdaOutputTypeEnum.String)
49+
step_lambda = LambdaStep(
50+
name="lambda-step",
51+
lambda_func=Lambda(
52+
function_arn=("arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda"),
53+
session=sagemaker_session,
54+
),
55+
inputs={"arg1": "foo"},
56+
outputs=[outputParam1],
57+
)
58+
59+
pipeline = Pipeline(
60+
name=pipeline_name,
61+
parameters=[instance_count],
62+
steps=[step_lambda],
63+
sagemaker_session=sagemaker_session,
64+
)
65+
66+
try:
67+
response = pipeline.create(role)
68+
create_arn = response["PipelineArn"]
69+
assert re.match(
70+
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
71+
create_arn,
72+
)
73+
74+
pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
75+
response = pipeline.update(role)
76+
update_arn = response["PipelineArn"]
77+
assert re.match(
78+
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
79+
update_arn,
80+
)
81+
finally:
82+
try:
83+
pipeline.delete()
84+
except Exception:
85+
pass
86+
87+
88+
def test_two_step_lambda_pipeline_with_output_reference(
89+
sagemaker_session, role, pipeline_name, region_name
90+
):
91+
instance_count = ParameterInteger(name="InstanceCount", default_value=2)
92+
93+
outputParam1 = LambdaOutput(output_name="output1", output_type=LambdaOutputTypeEnum.String)
94+
step_lambda1 = LambdaStep(
95+
name="lambda-step1",
96+
lambda_func=Lambda(
97+
function_arn=("arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda"),
98+
session=sagemaker_session,
99+
),
100+
inputs={"arg1": "foo"},
101+
outputs=[outputParam1],
102+
)
103+
104+
step_lambda2 = LambdaStep(
105+
name="lambda-step2",
106+
lambda_func=Lambda(
107+
function_arn=("arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda"),
108+
session=sagemaker_session,
109+
),
110+
inputs={"arg1": outputParam1},
111+
outputs=[],
112+
)
113+
114+
pipeline = Pipeline(
115+
name=pipeline_name,
116+
parameters=[instance_count],
117+
steps=[step_lambda1, step_lambda2],
118+
sagemaker_session=sagemaker_session,
119+
)
120+
121+
try:
122+
response = pipeline.create(role)
123+
create_arn = response["PipelineArn"]
124+
assert re.match(
125+
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
126+
create_arn,
127+
)
128+
finally:
129+
try:
130+
pipeline.delete()
131+
except Exception:
132+
pass

0 commit comments

Comments
 (0)