Skip to content

fix: change model_dir to training job name if it is for tuning. #179

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 5 commits into from
Apr 12, 2019
Merged

fix: change model_dir to training job name if it is for tuning. #179

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6401c18
Change model_dir to training job if the training job is for tuning.
Apr 12, 2019
d7fa2b6
Remove unused import.
Apr 12, 2019
9eac2d7
Update method comments and modify model_dir.
Apr 12, 2019
1d3d0d5
Update checkpoint_dir format.
Apr 12, 2019
0722739
Us os.path.join.
Apr 12, 2019
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
6 changes: 6 additions & 0 deletions src/sagemaker_tensorflow_container/training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,12 @@ def main():
"""
hyperparameters = framework.env.read_hyperparameters()
env = framework.training_env(hyperparameters=hyperparameters)

# If the training job is part of the multiple training jobs for tuning, we need to append the training job name to
# model_dir in case they read from/write to the same object
if '_tuning_objective_metric' in hyperparameters:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd add a comment similar to https://github.com/aws/sagemaker-tensorflow-container/blob/master/src/tf_container/train_entry_point.py#L118-L119 about why we have this logic

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added.

env.hyperparameters['model_dir'] = os.path.join(hyperparameters.get('model_dir'), env.job_name, 'checkpoints')

s3_utils.configure(env.hyperparameters.get('model_dir'), os.environ.get('SAGEMAKER_REGION'))
logger.setLevel(env.log_level)
train(env)
Expand Down
31 changes: 31 additions & 0 deletions test/unit/test_training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ def simple_training_env():
env.hosts = CURRENT_HOST
env.current_host = CURRENT_HOST
env.to_env_vars = lambda: {}
env.job_name = 'test-training-job'
return env


Expand Down Expand Up @@ -252,3 +253,33 @@ def test_main(configure_s3_env, read_hyperparameters, training_env,
training_env.assert_called_once_with(hyperparameters={})
train.assert_called_once_with(single_machine_training_env)
configure_s3_env.assert_called_once()


@patch('sagemaker_tensorflow_container.training.logger')
@patch('sagemaker_tensorflow_container.training.train')
@patch('logging.Logger.setLevel')
@patch('sagemaker_containers.beta.framework.training_env')
@patch('sagemaker_containers.beta.framework.env.read_hyperparameters', return_value={'model_dir': MODEL_DIR})
@patch('sagemaker_tensorflow_container.s3_utils.configure')
def test_main_simple_training_model_dir(configure_s3_env, read_hyperparameters, training_env,
set_level, train, logger, single_machine_training_env):
training_env.return_value = single_machine_training_env
os.environ['SAGEMAKER_REGION'] = REGION
training.main()
configure_s3_env.assert_called_once_with(MODEL_DIR, REGION)


@patch('sagemaker_tensorflow_container.training.logger')
@patch('sagemaker_tensorflow_container.training.train')
@patch('logging.Logger.setLevel')
@patch('sagemaker_containers.beta.framework.training_env')
@patch('sagemaker_containers.beta.framework.env.read_hyperparameters', return_value={'model_dir': MODEL_DIR,
'_tuning_objective_metric': 'auc'})
@patch('sagemaker_tensorflow_container.s3_utils.configure')
def test_main_tunning_model_dir(configure_s3_env, read_hyperparameters, training_env,
set_level, train, logger, single_machine_training_env):
training_env.return_value = single_machine_training_env
os.environ['SAGEMAKER_REGION'] = REGION
training.main()
expected_model_dir = os.path.join(MODEL_DIR, single_machine_training_env.job_name, 'checkpoints')
configure_s3_env.assert_called_once_with(expected_model_dir, REGION)