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

* Change model_dir to training job if the training job is for tuning.

Author: chuyang-deng

src/sagemaker_tensorflow_container/training.py

@@ -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:
+        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)

test/unit/test_training.py

@@ -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
 
 
@@ -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)