Skip to content

fix: container env generation for S3 URI and add test for the same #2971

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 4 commits into from
Mar 7, 2022
Merged

fix: container env generation for S3 URI and add test for the same #2971

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
48ca435
Fixes container env generation for S3 URI, update test
shreyapandit Mar 4, 2022
31528ac
move test to model dir, update test to test properties based model
shreyapandit Mar 4, 2022
845bda5
bring back script assignment, cleanup test
shreyapandit Mar 4, 2022
9967091
fix indents and cleanup
shreyapandit Mar 4, 2022
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
9 changes: 6 additions & 3 deletions src/sagemaker/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -466,7 +466,7 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
)

def _script_mode_env_vars(self):
"""Placeholder docstring"""
"""Returns a mapping of environment variables for script mode execution"""
script_name = None
dir_name = None
if self.uploaded_code:
Expand All @@ -478,8 +478,11 @@ def _script_mode_env_vars(self):
elif self.entry_point is not None:
script_name = self.entry_point
if self.source_dir is not None:
dir_name = "file://" + self.source_dir

dir_name = (
self.source_dir
if self.source_dir.startswith("s3://")
else "file://" + self.source_dir
)
return {
SCRIPT_PARAM_NAME.upper(): script_name or str(),
DIR_PARAM_NAME.upper(): dir_name or str(),
Expand Down
20 changes: 19 additions & 1 deletion tests/unit/sagemaker/model/test_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@
from sagemaker.sklearn.model import SKLearnModel
from sagemaker.tensorflow.model import TensorFlowModel
from sagemaker.xgboost.model import XGBoostModel
from sagemaker.workflow.properties import Properties


MODEL_DATA = "s3://bucket/model.tar.gz"
MODEL_IMAGE = "mi"
Expand All @@ -42,7 +44,6 @@
BRANCH = "test-branch-git-config"
COMMIT = "ae15c9d7d5b97ea95ea451e4662ee43da3401d73"
ENTRY_POINT_INFERENCE = "inference.py"

SCRIPT_URI = "s3://codebucket/someprefix/sourcedir.tar.gz"
IMAGE_URI = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.9.0-gpu-py38"

Expand Down Expand Up @@ -71,6 +72,23 @@ def sagemaker_session():
return sms


@patch("shutil.rmtree", MagicMock())
@patch("tarfile.open", MagicMock())
@patch("os.listdir", MagicMock(return_value=[ENTRY_POINT_INFERENCE]))
def test_prepare_container_def_with_model_src_s3_returns_correct_url(sagemaker_session):
model = Model(
entry_point=ENTRY_POINT_INFERENCE,
role=ROLE,
sagemaker_session=sagemaker_session,
source_dir=SCRIPT_URI,
image_uri=MODEL_IMAGE,
model_data=Properties("Steps.MyStep"),
)
container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")

assert container_def["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"] == SCRIPT_URI


def test_prepare_container_def_with_model_data():
model = Model(MODEL_IMAGE)
container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")
Expand Down