Mvs prepare (#138)

* bug-fix: reintroduce _modules.prepare for backwards compatibility

Author: mvsusp

CHANGELOG.rst

@@ -2,6 +2,11 @@
 CHANGELOG
 =========
 
+2.3.2
+=====
+
+* bug-fix: reintroduce _modules.prepare for backwards compatibility
+
 2.3.1
 =====
 

setup.py

@@ -34,7 +34,7 @@ def read(file_name):
 
 setup(
     name='sagemaker_containers',
-    version='2.3.1',
+    version='2.3.2',
     description='Open source library for creating containers to run on Amazon SageMaker.',
 
     packages=packages,

src/sagemaker_containers/_modules.py

@@ -16,6 +16,7 @@
 import os
 import shlex
 import sys
+import textwrap
 import warnings
 
 import six
@@ -48,6 +49,48 @@ def has_requirements(path):  # type: (str) -> None
     return os.path.exists(os.path.join(path, 'requirements.txt'))
 
 
+def prepare(path, name):  # type: (str, str) -> None
+    """Prepare a Python script (or module) to be imported as a module.
+    If the script does not contain a setup.py file, it creates a minimal setup.
+    Args:
+        path (str): path to directory with the script or module.
+        name (str): name of the script or module.
+    """
+    setup_path = os.path.join(path, 'setup.py')
+    if not os.path.exists(setup_path):
+        data = textwrap.dedent("""
+        from setuptools import setup
+        setup(packages=[''],
+              name="%s",
+              version='1.0.0',
+              include_package_data=True)
+        """ % name)
+
+        logger.info('Module %s does not provide a setup.py. \nGenerating setup.py' % name)
+
+        _files.write_file(setup_path, data)
+
+        data = textwrap.dedent("""
+        [wheel]
+        universal = 1
+        """)
+
+        logger.info('Generating setup.cfg')
+
+        _files.write_file(os.path.join(path, 'setup.cfg'), data)
+
+        data = textwrap.dedent("""
+        recursive-include . *
+        recursive-exclude . __pycache__*
+        recursive-exclude . *.pyc
+        recursive-exclude . *.pyo
+        """)
+
+        logger.info('Generating MANIFEST.in')
+
+        _files.write_file(os.path.join(path, 'MANIFEST.in'), data)
+
+
 def install(path):  # type: (str) -> None
     """Install a Python module in the executing Python environment.
     Args:
@@ -173,6 +216,7 @@ def run_module(uri, args, env_vars=None, name=DEFAULT_MODULE_NAME, cache=None, w
 
     _files.download_and_extract(uri, name, _env.code_dir)
 
+    prepare(_env.code_dir, name)
     install(_env.code_dir)
 
     _env.write_env_vars(env_vars)

test/functional/test_training_framework.py

@@ -241,6 +241,13 @@ def framework_training_with_script_mode_fn():
                     training_env.to_env_vars())
 
 
+def framework_training_with_run_modules_fn():
+    training_env = sagemaker_containers.training_env()
+
+    modules.run_module(training_env.module_dir, training_env.to_cmd_args(),
+                       training_env.to_env_vars(), training_env.module_name)
+
+
 def test_parameter_server():
     module = test.UserModule(test.File(name='user_script.py', data=PARAMETER_SERVER_SCRIPT))
     hyperparameters = dict(sagemaker_program='user_script.py')
@@ -254,8 +261,10 @@ def test_parameter_server():
     process.kill()
 
 
-@pytest.mark.parametrize('user_script', [USER_MODE_SCRIPT])
-def test_script_mode(user_script):
+@pytest.mark.parametrize('user_script, training_fn', [
+    [USER_MODE_SCRIPT, framework_training_with_script_mode_fn],
+    [USER_MODE_SCRIPT, framework_training_with_run_modules_fn]])
+def test_script_mode(user_script, training_fn):
     channel = test.Channel.create(name='training')
 
     features = [1, 2, 3, 4]
@@ -269,7 +278,7 @@ def test_script_mode(user_script):
 
     test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[channel])
 
-    assert execute_an_wrap_exit(framework_training_with_script_mode_fn) == trainer.SUCCESS_CODE
+    assert execute_an_wrap_exit(training_fn) == trainer.SUCCESS_CODE
 
     model_path = os.path.join(env.model_dir, 'saved_model')
 
@@ -281,8 +290,10 @@ def test_script_mode(user_script):
     assert model.optimizer == 'SGD'
 
 
-@pytest.mark.parametrize('user_script', [USER_MODE_SCRIPT])
-def test_script_mode_local_directory(user_script, tmpdir):
+@pytest.mark.parametrize('user_script, training_fn', [
+    [USER_MODE_SCRIPT, framework_training_with_script_mode_fn],
+    [USER_MODE_SCRIPT, framework_training_with_run_modules_fn]])
+def test_script_mode_local_directory(user_script, training_fn, tmpdir):
     channel = test.Channel.create(name='training')
 
     features = [1, 2, 3, 4]
@@ -300,7 +311,7 @@ def test_script_mode_local_directory(user_script, tmpdir):
 
     test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[channel], local=True)
 
-    assert execute_an_wrap_exit(framework_training_with_script_mode_fn) == trainer.SUCCESS_CODE
+    assert execute_an_wrap_exit(training_fn) == trainer.SUCCESS_CODE
 
     model_path = os.path.join(env.model_dir, 'saved_model')
 
@@ -318,7 +329,10 @@ def test_script_mode_local_directory(user_script, tmpdir):
 """
 
 
-def test_script_mode_client_error():
+@pytest.mark.parametrize('training_fn', [
+    framework_training_with_script_mode_fn,
+    framework_training_with_run_modules_fn])
+def test_script_mode_client_error(training_fn):
     channel = test.Channel.create(name='training')
 
     module = test.UserModule(test.File(name='user_script.py', data=USER_MODE_SCRIPT_WITH_ERROR))
@@ -328,13 +342,16 @@ def test_script_mode_client_error():
     test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[channel])
 
     with pytest.raises(errors.ExecuteUserScriptError) as e:
-        framework_training_with_script_mode_fn()
+        training_fn()
 
     message = str(e.value)
     assert 'ExecuteUserScriptError' in message
 
 
-def test_script_mode_client_import_error():
+@pytest.mark.parametrize('training_fn', [
+    framework_training_with_script_mode_fn,
+    framework_training_with_run_modules_fn])
+def test_script_mode_client_import_error(training_fn):
     channel = test.Channel.create(name='training')
 
     requirements_file = test.File('requirements.txt', '42/0')
@@ -347,7 +364,7 @@ def test_script_mode_client_import_error():
     test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[channel])
 
     with pytest.raises(errors.InstallModuleError) as e:
-        framework_training_with_script_mode_fn()
+        training_fn()
 
     message = str(e.value)
     assert 'InstallModuleError:' in message

test/unit/test_modules.py

@@ -15,8 +15,9 @@
 import contextlib
 import os
 import sys
+import textwrap
 
-from mock import call, patch
+from mock import call, mock_open, patch
 import pytest
 from six import PY2
 
@@ -77,6 +78,47 @@ def patch_tmpdir():
     yield '/tmp'
 
 
+@patch(builtins_open, mock_open())
+@patch('os.path.exists', lambda x: False)
+def test_prepare():
+    _modules.prepare('c:/path/to/', 'my-module')
+
+    open.assert_any_call('c:/path/to/setup.py', 'w')
+    open.assert_any_call('c:/path/to/setup.cfg', 'w')
+    open.assert_any_call('c:/path/to/MANIFEST.in', 'w')
+
+    data = textwrap.dedent("""
+    from setuptools import setup
+    setup(packages=[''],
+          name="my-module",
+          version='1.0.0',
+          include_package_data=True)
+    """)
+
+    open().write.assert_any_call(data)
+
+    data = textwrap.dedent("""
+    [wheel]
+    universal = 1
+    """)
+    open().write.assert_any_call(data)
+
+    data = textwrap.dedent("""
+    recursive-include . *
+    recursive-exclude . __pycache__*
+    recursive-exclude . *.pyc
+    recursive-exclude . *.pyo
+    """)
+    open().write.assert_any_call(data)
+
+
+@patch(builtins_open, mock_open())
+@patch('os.path.exists', lambda x: True)
+def test_prepare_already_prepared():
+    _modules.prepare('c:/path/to/', 'my-module')
+    open.assert_not_called()
+
+
 @patch('importlib.import_module')
 def test_exists(import_module):
     assert _modules.exists('my_module')