Skip to content

Commit 2922621

Browse files
chuyang-denglaurenyu
authored andcommitted
fix: use name_from_base in auto_ml.py but unique_name_from_base in tests. (#1182)
1 parent 42d9c1e commit 2922621

File tree

2 files changed

+27
-21
lines changed

2 files changed

+27
-21
lines changed

src/sagemaker/automl/automl.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
from sagemaker.automl.candidate_estimator import CandidateEstimator
2020
from sagemaker.job import _Job
2121
from sagemaker.session import Session
22-
from sagemaker.utils import unique_name_from_base
22+
from sagemaker.utils import name_from_base
2323

2424

2525
class AutoML(object):
@@ -381,7 +381,7 @@ def _prepare_for_auto_ml_job(self, job_name=None):
381381
else:
382382
base_name = "automl"
383383
# CreateAutoMLJob API validates that member length less than or equal to 32
384-
self.current_job_name = unique_name_from_base(base_name, max_length=32)
384+
self.current_job_name = name_from_base(base_name, max_length=32)
385385

386386
if self.output_path is None:
387387
self.output_path = "s3://{}/".format(self.sagemaker_session.default_bucket())

tests/unit/sagemaker/automl/test_auto_ml.py

Lines changed: 25 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,14 @@
1313
from __future__ import absolute_import
1414

1515
import pytest
16-
from mock import Mock
16+
from mock import Mock, patch
1717
from sagemaker import AutoML, AutoMLJob, AutoMLInput, CandidateEstimator
1818

1919
MODEL_DATA = "s3://bucket/model.tar.gz"
2020
MODEL_IMAGE = "mi"
2121
ENTRY_POINT = "blah.py"
2222

23+
TIMESTAMP = "2017-11-06-14:14:15.671"
2324
BUCKET_NAME = "mybucket"
2425
INSTANCE_COUNT = 1
2526
INSTANCE_TYPE = "ml.c5.2xlarge"
@@ -31,6 +32,7 @@
3132
DEFAULT_OUTPUT_PATH = "s3://{}/".format(BUCKET_NAME)
3233
LOCAL_DATA_PATH = "file://data"
3334
DEFAULT_MAX_CANDIDATES = 500
35+
DEFAULT_JOB_NAME = "automl-{}".format(TIMESTAMP)
3436

3537
JOB_NAME = "default-job-name"
3638
JOB_NAME_2 = "banana-auto-ml-job"
@@ -281,34 +283,38 @@ def test_auto_ml_additional_optional_params(sagemaker_session):
281283
}
282284

283285

284-
def test_auto_ml_default_fit(sagemaker_session):
286+
@patch("time.strftime", return_value=TIMESTAMP)
287+
def test_auto_ml_default_fit(strftime, sagemaker_session):
285288
auto_ml = AutoML(
286289
role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
287290
)
288291
inputs = DEFAULT_S3_INPUT_DATA
289292
auto_ml.fit(inputs)
290293
sagemaker_session.auto_ml.assert_called_once()
291294
_, args = sagemaker_session.auto_ml.call_args
292-
assert args["input_config"] == [
293-
{
294-
"DataSource": {
295-
"S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": DEFAULT_S3_INPUT_DATA}
295+
assert args == {
296+
"input_config": [
297+
{
298+
"DataSource": {
299+
"S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": DEFAULT_S3_INPUT_DATA}
300+
},
301+
"TargetAttributeName": TARGET_ATTRIBUTE_NAME,
302+
}
303+
],
304+
"output_config": {"S3OutputPath": DEFAULT_OUTPUT_PATH},
305+
"auto_ml_job_config": {
306+
"CompletionCriteria": {"MaxCandidates": DEFAULT_MAX_CANDIDATES},
307+
"SecurityConfig": {
308+
"EnableInterContainerTrafficEncryption": ENCRYPT_INTER_CONTAINER_TRAFFIC
296309
},
297-
"TargetAttributeName": TARGET_ATTRIBUTE_NAME,
298-
}
299-
]
300-
assert args["output_config"] == {"S3OutputPath": DEFAULT_OUTPUT_PATH}
301-
assert args["auto_ml_job_config"] == {
302-
"CompletionCriteria": {"MaxCandidates": DEFAULT_MAX_CANDIDATES},
303-
"SecurityConfig": {
304-
"EnableInterContainerTrafficEncryption": ENCRYPT_INTER_CONTAINER_TRAFFIC
305310
},
311+
"role": ROLE,
312+
"job_name": DEFAULT_JOB_NAME,
313+
"problem_type": None,
314+
"job_objective": None,
315+
"generate_candidate_definitions_only": GENERATE_CANDIDATE_DEFINITIONS_ONLY,
316+
"tags": None,
306317
}
307-
assert args["role"] == ROLE
308-
assert args["problem_type"] is None
309-
assert args["job_objective"] is None
310-
assert args["generate_candidate_definitions_only"] == GENERATE_CANDIDATE_DEFINITIONS_ONLY
311-
assert args["tags"] is None
312318

313319

314320
def test_auto_ml_local_input(sagemaker_session):

0 commit comments

Comments
 (0)