Skip to content

Commit 0dea793

Browse files
author
Dewen Qi
committed
fix doc8 and unit tests
1 parent 57f1370 commit 0dea793

File tree

3 files changed

+22
-10
lines changed

3 files changed

+22
-10
lines changed

‎tests/data/pipeline/model_step/pytorch_mnist/requirements.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
scipy>=1.8.1
1+
scipy>=1.8.1

‎tests/unit/sagemaker/workflow/test_model_step.py

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
from sagemaker.tensorflow import TensorFlowModel
3333
from sagemaker.transformer import Transformer
3434
from sagemaker.tuner import HyperparameterTuner
35+
from sagemaker.workflow._utils import REPACK_SCRIPT_LAUNCHER
3536
from sagemaker.workflow.condition_step import ConditionStep
3637
from sagemaker.workflow.conditions import ConditionGreaterThanOrEqualTo
3738
from sagemaker.workflow.model_step import (
@@ -188,7 +189,9 @@ def test_register_model_with_runtime_repack(pipeline_session, model_data_param,
188189
}
189190
assert arguments["HyperParameters"]["inference_script"] == '"dummy_script.py"'
190191
assert arguments["HyperParameters"]["model_archive"] == {"Get": "Parameters.ModelData"}
191-
assert arguments["HyperParameters"]["sagemaker_program"] == '"_repack_model.py"'
192+
assert (
193+
arguments["HyperParameters"]["sagemaker_program"] == f'"{REPACK_SCRIPT_LAUNCHER}"'
194+
)
192195
assert "s3://" in arguments["HyperParameters"]["sagemaker_submit_directory"]
193196
assert arguments["HyperParameters"]["dependencies"] == "null"
194197
assert step["RetryPolicies"] == [
@@ -269,7 +272,9 @@ def test_create_model_with_runtime_repack(pipeline_session, model_data_param, mo
269272
}
270273
assert arguments["HyperParameters"]["inference_script"] == '"dummy_script.py"'
271274
assert arguments["HyperParameters"]["model_archive"] == {"Get": "Parameters.ModelData"}
272-
assert arguments["HyperParameters"]["sagemaker_program"] == '"_repack_model.py"'
275+
assert (
276+
arguments["HyperParameters"]["sagemaker_program"] == f'"{REPACK_SCRIPT_LAUNCHER}"'
277+
)
273278
assert "s3://" in arguments["HyperParameters"]["sagemaker_submit_directory"]
274279
assert arguments["HyperParameters"]["dependencies"] == "null"
275280
assert "repack a model with customer scripts" in step["Description"]
@@ -360,7 +365,9 @@ def test_create_pipeline_model_with_runtime_repack(pipeline_session, model_data_
360365
}
361366
assert arguments["HyperParameters"]["inference_script"] == '"dummy_script.py"'
362367
assert arguments["HyperParameters"]["model_archive"] == {"Get": "Parameters.ModelData"}
363-
assert arguments["HyperParameters"]["sagemaker_program"] == '"_repack_model.py"'
368+
assert (
369+
arguments["HyperParameters"]["sagemaker_program"] == f'"{REPACK_SCRIPT_LAUNCHER}"'
370+
)
364371
assert "s3://" in arguments["HyperParameters"]["sagemaker_submit_directory"]
365372
assert arguments["HyperParameters"]["dependencies"] == "null"
366373
assert step["RetryPolicies"] == [
@@ -460,7 +467,9 @@ def test_register_pipeline_model_with_runtime_repack(pipeline_session, model_dat
460467
}
461468
assert arguments["HyperParameters"]["inference_script"] == '"dummy_script.py"'
462469
assert arguments["HyperParameters"]["model_archive"] == {"Get": "Parameters.ModelData"}
463-
assert arguments["HyperParameters"]["sagemaker_program"] == '"_repack_model.py"'
470+
assert (
471+
arguments["HyperParameters"]["sagemaker_program"] == f'"{REPACK_SCRIPT_LAUNCHER}"'
472+
)
464473
assert "s3://" in arguments["HyperParameters"]["sagemaker_submit_directory"]
465474
assert arguments["HyperParameters"]["dependencies"] == "null"
466475
elif step["Type"] == "RegisterModel":
@@ -641,7 +650,9 @@ def test_conditional_model_create_and_regis(
641650
}
642651
assert arguments["HyperParameters"]["inference_script"] == '"dummy_script.py"'
643652
assert arguments["HyperParameters"]["model_archive"] == {"Get": "Parameters.ModelData"}
644-
assert arguments["HyperParameters"]["sagemaker_program"] == '"_repack_model.py"'
653+
assert (
654+
arguments["HyperParameters"]["sagemaker_program"] == f'"{REPACK_SCRIPT_LAUNCHER}"'
655+
)
645656
assert "s3://" in arguments["HyperParameters"]["sagemaker_submit_directory"]
646657
assert arguments["HyperParameters"]["dependencies"] == "null"
647658
elif step["Type"] == "RegisterModel":

‎tests/unit/sagemaker/workflow/test_step_collections.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import pytest
2121

2222
from sagemaker.drift_check_baselines import DriftCheckBaselines
23+
from sagemaker.workflow._utils import REPACK_SCRIPT_LAUNCHER
2324
from sagemaker.workflow.condition_step import ConditionStep
2425
from sagemaker.workflow.conditions import ConditionEquals
2526
from sagemaker.workflow.model_step import (
@@ -799,7 +800,7 @@ def test_register_model_with_model_repack_with_estimator(
799800
"inference_script": '"dummy_script.py"',
800801
"dependencies": f'"{dummy_requirements}"',
801802
"model_archive": '"s3://my-bucket/model.tar.gz"',
802-
"sagemaker_program": '"_repack_model.py"',
803+
"sagemaker_program": f'"{REPACK_SCRIPT_LAUNCHER}"',
803804
"sagemaker_container_log_level": "20",
804805
"sagemaker_region": f'"{REGION}"',
805806
"source_dir": "null",
@@ -925,7 +926,7 @@ def test_register_model_with_model_repack_with_model(model, model_metrics, drift
925926
"HyperParameters": {
926927
"inference_script": '"dummy_script.py"',
927928
"model_archive": '"s3://my-bucket/model.tar.gz"',
928-
"sagemaker_program": '"_repack_model.py"',
929+
"sagemaker_program": f'"{REPACK_SCRIPT_LAUNCHER}"',
929930
"sagemaker_container_log_level": "20",
930931
"sagemaker_region": f'"{REGION}"',
931932
"dependencies": "null",
@@ -1055,7 +1056,7 @@ def test_register_model_with_model_repack_with_pipeline_model(
10551056
"dependencies": "null",
10561057
"inference_script": '"dummy_script.py"',
10571058
"model_archive": '"s3://my-bucket/model.tar.gz"',
1058-
"sagemaker_program": '"_repack_model.py"',
1059+
"sagemaker_program": f'"{REPACK_SCRIPT_LAUNCHER}"',
10591060
"sagemaker_container_log_level": "20",
10601061
"sagemaker_region": f'"{REGION}"',
10611062
"source_dir": "null",
@@ -1267,7 +1268,7 @@ def test_estimator_transformer_with_model_repack_with_estimator(estimator):
12671268
"model_archive": '"s3://my-bucket/model.tar.gz"',
12681269
"dependencies": "null",
12691270
"source_dir": "null",
1270-
"sagemaker_program": '"_repack_model.py"',
1271+
"sagemaker_program": f'"{REPACK_SCRIPT_LAUNCHER}"',
12711272
"sagemaker_container_log_level": "20",
12721273
"sagemaker_region": '"us-west-2"',
12731274
},

0 commit comments

Comments
 (0)