Changes for temp location and out_dir with Sagemaker in mind (aws#154)

* Make outdir optional arg, use default path in sagemaker environment, also change temp location when writing local files

* remove is_s3 import

* add tests and fix case when / is at the front of filepath

* add comments

* change to .tmp suffix

* update testing script to take a tag

Author: rahul003

bin/sagemaker-containers/tensorflow/run_sagemaker.py

@@ -4,8 +4,10 @@
 dir_path = os.path.dirname(os.path.realpath(__file__))
 train_script_path = os.path.join(dir_path, 'tf-train.py')
 
+tag = os.environ.get('SM_TESTING_TAG', 'DEFAULTTAGWHICHWILLFAIL')
+
 estimator = TensorFlow(entry_point=train_script_path,
-                       image_name='072677473360.dkr.ecr.us-east-1.amazonaws.com/tornasole-preprod-tf-1.13.1-cpu:latest',
+                       image_name='072677473360.dkr.ecr.us-east-1.amazonaws.com/tornasole-preprod-tf-1.13.1-cpu:' + tag,
                        role='AmazonSageMaker-ExecutionRole-20190614T145575', # hardcode role name
                        base_job_name='tornasole', #there are some restrictions on base job name so keep it simple
                        train_instance_count=1,
@@ -15,7 +17,7 @@
 estimator.fit()
 
 estimator = TensorFlow(entry_point=train_script_path,
-                       image_name='072677473360.dkr.ecr.us-east-1.amazonaws.com/tornasole-preprod-tf-1.13.1-gpu:latest',
+                       image_name='072677473360.dkr.ecr.us-east-1.amazonaws.com/tornasole-preprod-tf-1.13.1-gpu:' + tag,
                        role='AmazonSageMaker-ExecutionRole-20190614T145575', # hardcode role name
                        base_job_name='tornasole', #there are some restrictions on base job name so keep it simple
                        train_instance_count=1,

tests/core/test_paths.py

@@ -0,0 +1,45 @@
+from tornasole.core.hook_utils import verify_and_get_out_dir
+from tornasole.core.json_config import DEFAULT_SAGEMAKER_TORNASOLE_PATH
+from tornasole.core.access_layer.file import get_temp_path, SAGEMAKER_TEMP_PATH_SUFFIX, NON_SAGEMAKER_TEMP_PATH_PREFIX
+import os
+import uuid
+
+
+def test_outdir_non_sagemaker():
+    id = str(uuid.uuid4())
+    path = '/tmp/tests/' + id
+    out_dir = verify_and_get_out_dir(path)
+    assert out_dir == path
+    os.makedirs(path)
+    try:
+        verify_and_get_out_dir(path)
+        # should raise exception as dir present
+        assert False
+    except RuntimeError as e:
+        pass
+
+
+def test_outdir_sagemaker():
+    os.environ['TRAINING_JOB_NAME'] = 'a'
+    id = str(uuid.uuid4())
+    paths = ['/tmp/tests/' + id, 's3://tmp/tests/' + id]
+    for path in paths:
+        out_dir = verify_and_get_out_dir(path)
+        assert out_dir == DEFAULT_SAGEMAKER_TORNASOLE_PATH
+    del os.environ['TRAINING_JOB_NAME']
+
+
+def test_temp_paths():
+    for path in ['/opt/ml/output/tensors/events/a',
+                 '/opt/ml/output/tensors/a',
+                 '/opt/ml/output/tensors/events/a/b']:
+        tp = get_temp_path(path)
+        assert tp.endswith(SAGEMAKER_TEMP_PATH_SUFFIX)
+        assert not tp.startswith(NON_SAGEMAKER_TEMP_PATH_PREFIX)
+
+    for path in ['/a/b/c', '/opt/ml/output/a', 'a/b/c']:
+        tp = get_temp_path(path)
+        assert not SAGEMAKER_TEMP_PATH_SUFFIX in tp
+        assert tp.startswith(NON_SAGEMAKER_TEMP_PATH_PREFIX)
+
+

tornasole/core/access_layer/file.py

@@ -1,8 +1,13 @@
+from tornasole.core.json_config import DEFAULT_SAGEMAKER_TORNASOLE_PATH
 from .base import TSAccessBase
 import os
 import shutil
 
 
+NON_SAGEMAKER_TEMP_PATH_PREFIX = '/tmp'
+SAGEMAKER_TEMP_PATH_SUFFIX = '.tmp'
+
+
 def ensure_dir(file_path, is_file=True):
     if is_file:
         directory = os.path.dirname(file_path)
@@ -13,6 +18,18 @@ def ensure_dir(file_path, is_file=True):
         os.makedirs(directory)
 
 
+def get_temp_path(file_path):
+    directory = os.path.dirname(file_path)
+    basename = os.path.basename(file_path)
+    if directory.startswith(DEFAULT_SAGEMAKER_TORNASOLE_PATH):
+        temp_path = file_path + SAGEMAKER_TEMP_PATH_SUFFIX
+    else:
+        if len(file_path) > 0 and file_path[0] == '/':
+            file_path = file_path[1:]
+        temp_path = os.path.join(NON_SAGEMAKER_TEMP_PATH_PREFIX, file_path)
+    return temp_path
+
+
 WRITE_MODES = ['w', 'w+', 'wb', 'wb+', 'a', 'a+', 'ab', 'ab+']
 
 
@@ -24,7 +41,7 @@ def __init__(self, path, mode):
         ensure_dir(path)
 
         if mode in WRITE_MODES:
-            self.temp_path = os.path.join('/tmp', self.path)
+            self.temp_path = get_temp_path(self.path)
             ensure_dir(self.temp_path)
             self.open(self.temp_path, mode)
         else:

tornasole/core/access_layer/utils.py

@@ -26,6 +26,7 @@ def training_has_ended(trial_prefix):
     else:
         writer = TSAccessFile(file_path, 'a+')
     writer.write("end of training job")
+    writer.flush()
     writer.close()
 
 

tornasole/core/hook_utils.py

@@ -0,0 +1,36 @@
+import os
+
+from .json_config import DEFAULT_SAGEMAKER_TORNASOLE_PATH
+from .sagemaker_utils import is_sagemaker_job, get_sagemaker_out_dir
+from .utils import is_s3, check_dir_exists, get_logger
+
+logger = get_logger()
+
+
+def verify_and_get_out_dir(out_dir):
+    if is_sagemaker_job():
+        if out_dir is not None:
+            logger.warning(
+                    'The out_dir parameter was set but job is running in a '
+                    'SageMaker environment, hence it is being ignored. '
+                    'Writing tensors to '
+                    '{}'.format(DEFAULT_SAGEMAKER_TORNASOLE_PATH))
+        out_dir = get_sagemaker_out_dir()
+        # here we don't check whether the directory exists because
+        # sagemaker creates the directory and it will exist.
+        # sagemaker assures us that this directory is unique for each job
+        # so there is no chance of tensors from two jobs being merged
+    else:
+        if out_dir is None:
+            raise RuntimeError(
+                'out_dir is a required argument when '
+                'running outside of SageMaker environments')
+        is_s3_path, _, _ = is_s3(out_dir)
+        if not is_s3_path:
+            out_dir = os.path.expanduser(out_dir)
+        # we check and raise error if directory already exists because
+        # we don't want to merge tensors from current job with
+        # tensors from previous job
+        check_dir_exists(out_dir)
+
+    return out_dir

tornasole/core/sagemaker_utils.py

@@ -1,4 +1,21 @@
 import boto3
+import os
+from .json_config import DEFAULT_SAGEMAKER_TORNASOLE_PATH
+
+
+def is_sagemaker_job():
+    """
+    If this variable is defined we are assuming that this is
+    a Sagemaker job. This is guaranteed to be defined
+    for all Sagemaker jobs.
+    https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo-running-container.html#your-algorithms-training-algo-running-container-environment-variables
+    :return: True or False
+    """
+    return 'TRAINING_JOB_NAME' in os.environ
+
+
+def get_sagemaker_out_dir():
+    return DEFAULT_SAGEMAKER_TORNASOLE_PATH
 
 
 class SageMakerUtils:

tornasole/mxnet/hook.py

@@ -3,10 +3,11 @@
 from tornasole.core.save_config import SaveConfig
 from tornasole.core.save_manager import SaveManager
 from tornasole.core.modes import ModeKeys, ALLOWED_MODES
-from tornasole.core.utils import get_logger, is_s3
+from tornasole.core.utils import get_logger
 from tornasole.core.reductions import get_reduction_tensor_name
 from tornasole.core.json_config import TORNASOLE_CONFIG_DEFAULT_WORKER_NAME, create_hook_from_json_config
 from tornasole.core.access_layer.utils import training_has_ended
+from tornasole.core.hook_utils import verify_and_get_out_dir
 from .mxnet_collection import get_collection_manager, get_collection
 from .util import get_aggregated_data, make_numpy_array
 import re as _re
@@ -27,20 +28,15 @@
 
 class TornasoleHook:
     def __init__(self,
-                 out_dir,
+                 out_dir=None,
                  dry_run=False,
                  worker=TORNASOLE_CONFIG_DEFAULT_WORKER_NAME,
                  reduction_config=None,
                  save_config=None,
                  include_regex=None,
                  include_collections=DEFAULT_INCLUDE_COLLECTIONS,
                  save_all=False):
-        if not is_s3(out_dir)[0]:
-            out_dir = os.path.expanduser(out_dir)
-        ## This is commented becuase Sagemaker creates out_dir
-        # This was created because we don't want user to overwrite their existing data
-        # check_dir_exists(out_dir)
-        self.out_dir = out_dir
+        self.out_dir = verify_and_get_out_dir(out_dir)
         self.out_base_dir = os.path.dirname(out_dir)
         self.run_id = os.path.basename(out_dir)
         self.include_collections = include_collections

tornasole/pytorch/hook.py

@@ -3,7 +3,8 @@
 from tornasole.core.save_config import SaveConfig
 from tornasole.core.save_manager import SaveManager
 from tornasole.core.modes import ModeKeys, ALLOWED_MODES
-from tornasole.core.utils import get_logger, is_s3
+from tornasole.core.utils import get_logger
+from tornasole.core.hook_utils import verify_and_get_out_dir
 from tornasole.core.reductions import get_reduction_tensor_name
 from tornasole.core.json_config import create_hook_from_json_config
 from tornasole.pytorch.torch_collection import get_collection_manager, get_collection
@@ -27,20 +28,15 @@
 
 class TornasoleHook:
     def __init__(self,
-                 out_dir,
+                 out_dir=None,
                  dry_run=False,
                  worker=DEFAULT_WORKER_NAME,
                  reduction_config=None,
                  save_config=None,
                  include_regex=None,
                  include_collections=DEFAULT_INCLUDE_COLLECTIONS,
                  save_all=False):
-        if not is_s3(out_dir)[0]:
-            out_dir = os.path.expanduser(out_dir)
-        # this is commented because SM creates dir. 
-        # This was created because we don't want user to overwrite their existing data
-        #check_dir_exists(out_dir)
-        self.out_dir = out_dir
+        self.out_dir = verify_and_get_out_dir(out_dir)
         self.out_base_dir = os.path.dirname(out_dir)
         self.run_id = os.path.basename(out_dir)
         self.include_collections = include_collections

tornasole/tensorflow/hook.py

@@ -6,7 +6,8 @@
 from .reductions import get_tensorflow_reduction
 from .collection import *
 from tornasole.core.writer import FileWriter
-from tornasole.core.utils import get_logger, flatten, is_s3, match_inc
+from tornasole.core.utils import get_logger, flatten, match_inc
+from tornasole.core.hook_utils import verify_and_get_out_dir
 from tornasole.core.reductions import get_reduction_tensor_name
 from tornasole.core.json_config import TORNASOLE_CONFIG_DEFAULT_WORKER_NAME, create_hook_from_json_config
 from tornasole.core.modes import ModeKeys, ALLOWED_MODES
@@ -17,7 +18,7 @@
 DEFAULT_INCLUDE_COLLECTIONS = ['weights', 'gradients', 'default']
 
 class TornasoleHook(tf.train.SessionRunHook):
-    def __init__(self, out_dir,
+    def __init__(self, out_dir=None,
                  dry_run=False,
                  worker=TORNASOLE_CONFIG_DEFAULT_WORKER_NAME,
                  reduction_config=None,
@@ -68,12 +69,7 @@ def __init__(self, out_dir,
             a shortcut for saving all tensors in the model.
             they are all saved in the collection `all`
         """
-        if not is_s3(out_dir)[0]:
-            out_dir = os.path.expanduser(out_dir)
-        # this is commented because SM creates dir. 
-        # This was created because we don't want user to overwrite their existing data
-        #check_dir_exists(out_dir)
-        self.out_dir = out_dir
+        self.out_dir = verify_and_get_out_dir(out_dir)
         self.out_base_dir = os.path.dirname(out_dir)
         self.run_id = os.path.basename(out_dir)