Skip to content

change: use sagemaker_session in workflow tests #2152

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
merged 6 commits into from
Feb 17, 2021
Merged

change: use sagemaker_session in workflow tests #2152

Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
47cef38
change: use sagemaker_session in workflow tests
Feb 16, 2021
8a5dbfa
remove unused imports
Feb 16, 2021
c9c30bd
Merge branch 'master' into master
metrizable Feb 17, 2021
f2d9f45
Merge branch 'master' into master
metrizable Feb 17, 2021
a50f41c
Merge branch 'master' into master
metrizable Feb 17, 2021
6f68232
Merge branch 'master' into master
metrizable Feb 17, 2021
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
28 changes: 4 additions & 24 deletions tests/integ/test_workflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,8 @@
import time
import uuid

import boto3
import pytest

from botocore.config import Config
from botocore.exceptions import WaiterError
from sagemaker.debugger import (
DebuggerHookConfig,
Expand All @@ -32,7 +30,7 @@
from sagemaker.model import Model
from sagemaker.processing import ProcessingInput, ProcessingOutput
from sagemaker.pytorch.estimator import PyTorch
from sagemaker.session import get_execution_role, Session
from sagemaker.session import get_execution_role
from sagemaker.sklearn.estimator import SKLearn
from sagemaker.sklearn.processing import SKLearnProcessor
from sagemaker.workflow.conditions import ConditionGreaterThanOrEqualTo
Expand Down Expand Up @@ -74,21 +72,6 @@ def role(sagemaker_session):
return get_execution_role(sagemaker_session)


@pytest.fixture(scope="module")
def workflow_session(region_name):
boto_session = boto3.Session(region_name=region_name)

sagemaker_client_config = dict()
sagemaker_client_config.setdefault("config", Config(retries=dict(max_attempts=2)))
sagemaker_client = boto_session.client("sagemaker", **sagemaker_client_config)

return Session(
boto_session=boto_session,
sagemaker_client=sagemaker_client,
sagemaker_runtime_client=None,
)


@pytest.fixture(scope="module")
def script_dir():
return os.path.join(DATA_DIR, "sklearn_processing")
Expand Down Expand Up @@ -119,7 +102,6 @@ def athena_dataset_definition(sagemaker_session):

def test_three_step_definition(
sagemaker_session,
workflow_session,
region_name,
role,
script_dir,
Expand Down Expand Up @@ -205,7 +187,7 @@ def test_three_step_definition(
name=pipeline_name,
parameters=[instance_type, instance_count, output_prefix],
steps=[step_process, step_train, step_model],
sagemaker_session=workflow_session,
sagemaker_session=sagemaker_session,
)

definition = json.loads(pipeline.definition())
Expand Down Expand Up @@ -277,7 +259,6 @@ def test_three_step_definition(

def test_one_step_sklearn_processing_pipeline(
sagemaker_session,
workflow_session,
role,
sklearn_latest_version,
cpu_instance_type,
Expand Down Expand Up @@ -313,7 +294,7 @@ def test_one_step_sklearn_processing_pipeline(
name=pipeline_name,
parameters=[instance_count],
steps=[step_sklearn],
sagemaker_session=workflow_session,
sagemaker_session=sagemaker_session,
)

try:
Expand Down Expand Up @@ -363,7 +344,6 @@ def test_one_step_sklearn_processing_pipeline(

def test_conditional_pytorch_training_model_registration(
sagemaker_session,
workflow_session,
role,
cpu_instance_type,
pipeline_name,
Expand Down Expand Up @@ -433,7 +413,7 @@ def test_conditional_pytorch_training_model_registration(
name=pipeline_name,
parameters=[good_enough_input, instance_count, instance_type],
steps=[step_cond],
sagemaker_session=workflow_session,
sagemaker_session=sagemaker_session,
)

try:
Expand Down