Generic OpenMPI support (#157)

* Generic MPI - Create library to changehostname when mpi starts (#153)

Author: mvsusp

.flake8

@@ -1,3 +1,3 @@
 [flake8]
-application_import_names = sagemaker_containers, test
-import-order-style = google
+application_import_names = sagemaker_containers, test, libchangehostname
+import-order-style = google

README.md

@@ -485,4 +485,4 @@ SM_TRAINING_ENV='
 Provides the entire training information as a JSON encoded dictionary.
 ## License
 
-This library is licensed under the Apache 2.0 License. 
+This library is licensed under the Apache 2.0 License.

setup.py

@@ -16,25 +16,31 @@
 import os
 import sys
 
-from setuptools import find_packages, setup
+import setuptools
 
 
 def read(file_name):
     return open(os.path.join(os.path.dirname(__file__), file_name)).read()
 
 
-packages = find_packages(where='src', exclude=('test',))
+packages = setuptools.find_packages(where='src', exclude=('test',))
 packages.append('sagemaker_containers.etc')
 
 required_packages = [
-  'boto3', 'six', 'pip', 'flask', 'gunicorn', 'gevent', 'inotify_simple', 'werkzeug'
+    'numpy', 'boto3', 'six', 'pip', 'flask', 'gunicorn', 'typing',
+    'gevent', 'inotify_simple', 'werkzeug', 'paramiko==2.4.2', 'psutil==5.4.8'
 ]
 
 # enum is introduced in Python 3.4. Installing enum back port
 if sys.version_info < (3, 4):
     required_packages.append('enum34 >= 1.1.6')
 
-setup(
+gethostname = setuptools.Extension('libchangehostname',
+                                   sources=['src/sagemaker_containers/c/libchangehostname.c'],
+                                   extra_compile_args=['-Wall', '-shared', '-export-dynamic',
+                                                       '-ldl'])
+
+setuptools.setup(
     name='sagemaker_containers',
     version='2.3.5',
     description='Open source library for creating containers to run on Amazon SageMaker.',
@@ -46,6 +52,7 @@ def read(file_name):
     },
     package_data={'sagemaker_containers.etc': ['*']},
     py_modules=[os.path.splitext(os.path.basename(path))[0] for path in glob('src/*.py')],
+    ext_modules=[gethostname],
     long_description=read('README.md'),
     author='Amazon Web Services',
     url='https://github.com/aws/sagemaker-containers/',
@@ -64,11 +71,11 @@ def read(file_name):
     install_requires=required_packages,
 
     extras_require={
-        'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'mock', 'sagemaker', 'numpy']
+        'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'mock', 'sagemaker==1.16.2']
     },
 
     entry_points={
-          'console_scripts': ['serve=sagemaker_containers.cli.serve:main',
-                              'train=sagemaker_containers.cli.train:main'],
+        'console_scripts': ['serve=sagemaker_containers.cli.serve:main',
+                            'train=sagemaker_containers.cli.train:main'],
     }
 )

src/sagemaker_containers/_entry_point_type.py

@@ -0,0 +1,42 @@
+# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License'). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the 'license' file accompanying this file. This file is
+# distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+import enum
+import os
+
+
+class _EntryPointType(enum.Enum):
+    PYTHON_PACKAGE = 'PYTHON_PACKAGE'
+    PYTHON_PROGRAM = 'PYTHON_PROGRAM'
+    COMMAND = 'COMMAND'
+
+
+PYTHON_PACKAGE = _EntryPointType.PYTHON_PACKAGE
+PYTHON_PROGRAM = _EntryPointType.PYTHON_PROGRAM
+COMMAND = _EntryPointType.COMMAND
+
+
+def get(path, name):  # type: (str, str) -> _EntryPointType
+    """
+    Args:
+        path (string): Directory where the entry point is located
+        name (string): Name of the entry point file
+
+    Returns:
+        (_EntryPointType): The type of the entry point
+    """
+    if 'setup.py' in os.listdir(path):
+        return _EntryPointType.PYTHON_PACKAGE
+    elif name.endswith('.py'):
+        return _EntryPointType.PYTHON_PROGRAM
+    else:
+        return _EntryPointType.COMMAND

src/sagemaker_containers/_env.py

@@ -138,7 +138,6 @@ def _set_base_path_env():  # type: () -> None
     str: the path to the intermediate output directory, e.g. /opt/ml/output/intermediate.
 """
 
-
 HYPERPARAMETERS_FILE = 'hyperparameters.json'  # type: str
 RESOURCE_CONFIG_FILE = 'resourceconfig.json'  # type: str
 INPUT_DATA_CONFIG_FILE = 'inputdataconfig.json'  # type: str
@@ -164,7 +163,7 @@ def _create_training_directories():
 
     resources_dict = {
         "current_host": host_name,
-        "hosts":        [host_name]
+        "hosts": [host_name]
     }
     _write_json(resources_dict, resource_config_file_dir)
 
@@ -534,11 +533,11 @@ def __init__(self, resource_config=None, input_data_config=None, hyperparameters
         resource_config = resource_config or read_resource_config()
         current_host = resource_config['current_host']
         hosts = resource_config['hosts']
-        network_interface_name = resource_config.get('network_interface_name', 'ethwe')
         input_data_config = input_data_config or read_input_data_config()
 
         all_hyperparameters = hyperparameters or read_hyperparameters()
-        split_result = _mapping.split_by_criteria(all_hyperparameters, keys=_params.SAGEMAKER_HYPERPARAMETERS,
+        split_result = _mapping.split_by_criteria(all_hyperparameters,
+                                                  keys=_params.SAGEMAKER_HYPERPARAMETERS,
                                                   prefix=_params.SAGEMAKER_PREFIX)
 
         sagemaker_hyperparameters = split_result.included
@@ -547,14 +546,17 @@ def __init__(self, resource_config=None, input_data_config=None, hyperparameters
             if k not in _params.SAGEMAKER_HYPERPARAMETERS
         }
 
-        sagemaker_region = sagemaker_hyperparameters.get(_params.REGION_NAME_PARAM, boto3.session.Session().region_name)
+        sagemaker_region = sagemaker_hyperparameters.get(_params.REGION_NAME_PARAM,
+                                                         boto3.session.Session().region_name)
 
         os.environ[_params.JOB_NAME_ENV] = sagemaker_hyperparameters.get(_params.JOB_NAME_PARAM, '')
         os.environ[_params.CURRENT_HOST_ENV] = current_host
         os.environ[_params.REGION_NAME_ENV] = sagemaker_region or ''
 
         self._hosts = hosts
-        self._network_interface_name = network_interface_name
+
+        self._network_interface_name = resource_config.get('network_interface_name', 'eth0')
+
         self._hyperparameters = split_result.excluded
         self._additional_framework_parameters = additional_framework_parameters
         self._resource_config = resource_config
@@ -567,7 +569,8 @@ def __init__(self, resource_config=None, input_data_config=None, hyperparameters
         # override base class attributes
         if self._module_name is None:
             self._module_name = str(sagemaker_hyperparameters.get(_params.USER_PROGRAM_PARAM, None))
-        self._user_entry_point = self._user_entry_point or sagemaker_hyperparameters.get(_params.USER_PROGRAM_PARAM)
+        self._user_entry_point = self._user_entry_point or sagemaker_hyperparameters.get(
+            _params.USER_PROGRAM_PARAM)
 
         self._module_dir = str(sagemaker_hyperparameters.get(_params.SUBMIT_DIR_PARAM, code_dir))
         self._log_level = sagemaker_hyperparameters.get(_params.LOG_LEVEL_PARAM, logging.INFO)
@@ -580,6 +583,21 @@ def __init__(self, resource_config=None, input_data_config=None, hyperparameters
         self._output_dir = output_dir
         self._job_name = os.environ.get(_params.TRAINING_JOB_ENV.upper(), None)
 
+        self._master_hostname = list(hosts)[0]
+        self._is_master = current_host == self._master_hostname
+
+    @property
+    def is_master(self):  # type: () -> bool
+        """Returns True if host is master
+        """
+        return self._is_master
+
+    @property
+    def master_hostname(self):  # type: () -> str
+        """Returns the hostname of the master node
+        """
+        return self._master_hostname
+
     @property
     def job_name(self):  # type: () -> str
         """The name of the current training job.
@@ -625,16 +643,19 @@ def to_env_vars(self):
         """
 
         env = {
-            'hosts':            self.hosts, 'network_interface_name': self.network_interface_name,
-            'hps':              self.hyperparameters, 'user_entry_point': self.user_entry_point,
+            'hosts': self.hosts, 'network_interface_name': self.network_interface_name,
+            'hps': self.hyperparameters, 'user_entry_point': self.user_entry_point,
             'framework_params': self.additional_framework_parameters,
-            'resource_config':  self.resource_config, 'input_data_config': self.input_data_config,
-            'output_data_dir':  self.output_data_dir, 'channels': sorted(self.channel_input_dirs.keys()),
-            'current_host':     self.current_host, 'module_name': self.module_name, 'log_level': self.log_level,
+            'resource_config': self.resource_config, 'input_data_config': self.input_data_config,
+            'output_data_dir': self.output_data_dir,
+            'channels': sorted(self.channel_input_dirs.keys()),
+            'current_host': self.current_host, 'module_name': self.module_name,
+            'log_level': self.log_level,
             'framework_module': self.framework_module, 'input_dir': self.input_dir,
-            'input_config_dir': self.input_config_dir, 'output_dir': self.output_dir, 'num_cpus': self.num_cpus,
-            'num_gpus':         self.num_gpus, 'model_dir': self.model_dir, 'module_dir': self.module_dir,
-            'training_env':     dict(self), 'user_args': self.to_cmd_args(),
+            'input_config_dir': self.input_config_dir, 'output_dir': self.output_dir,
+            'num_cpus': self.num_cpus,
+            'num_gpus': self.num_gpus, 'model_dir': self.model_dir, 'module_dir': self.module_dir,
+            'training_env': dict(self), 'user_args': self.to_cmd_args(),
             'output_intermediate_dir': self.output_intermediate_dir
         }