Add tests

Author: trungleduc

tests/unit/sagemaker/config/conftest.py

@@ -292,3 +292,15 @@ def s3_resource_mock():
 @pytest.fixture()
 def get_data_dir():
     return os.path.join(os.path.dirname(__file__), "..", "..", "..", "data", "config")
+
+
+@pytest.fixture()
+def base_local_mode_config():
+    return {
+        "local": {
+            "local_code": True,
+            "region_name": "",
+            "serving_port": 8080,
+            "container_config": {"shm_size": "128M"},
+        }
+    }

tests/unit/sagemaker/config/test_config.py

@@ -16,13 +16,15 @@
 import pytest
 import yaml
 import logging
-from mock import Mock, MagicMock
+from mock import Mock, MagicMock, patch
 
 from sagemaker.config.config import (
+    load_local_mode_config,
     load_sagemaker_config,
     logger,
     _DEFAULT_ADMIN_CONFIG_FILE_PATH,
     _DEFAULT_USER_CONFIG_FILE_PATH,
+    _DEFAULT_LOCAL_MODE_CONFIG_FILE_PATH,
 )
 from jsonschema import exceptions
 from yaml.constructor import ConstructorError
@@ -402,3 +404,13 @@ def test_logging_with_additional_configs_and_none_are_found(caplog):
         in caplog.text
     )
     logger.propagate = False
+
+
+@patch("sagemaker.config.config._load_config_from_file")
+def test_load_local_mode_config(mock_load_config):
+    load_local_mode_config()
+    mock_load_config.assert_called_with(_DEFAULT_LOCAL_MODE_CONFIG_FILE_PATH)
+
+
+def test_load_local_mode_config_when_config_file_is_not_found():
+    assert load_local_mode_config() is None

tests/unit/sagemaker/config/test_config_schema.py

@@ -13,7 +13,10 @@
 from __future__ import absolute_import
 from jsonschema import validate, exceptions
 import pytest
-from sagemaker.config.config_schema import SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA
+from sagemaker.config.config_schema import (
+    SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA,
+    SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA,
+)
 
 
 def _validate_config(base_config_with_schema, sagemaker_config):
@@ -291,3 +294,19 @@ def test_session_s3_object_key_prefix_schema(base_config_with_schema, prefix_nam
 def test_invalid_session_s3_object_key_prefix_schema(base_config_with_schema, invalid_prefix_name):
     with pytest.raises(exceptions.ValidationError):
         test_session_s3_object_key_prefix_schema(base_config_with_schema, invalid_prefix_name)
+
+
+def test_validate_local_mode_schema(base_local_mode_config):
+    validate(base_local_mode_config, SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA)
+
+
+def test_validate_local_mode_schema_with_additional_key(base_local_mode_config):
+    config = dict(**base_local_mode_config)
+    config["foo"] = "bar"
+    with pytest.raises(exceptions.ValidationError):
+        validate(config, SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA)
+
+    config2 = dict(**base_local_mode_config)
+    config2["local"]["foo"] = "bar"
+    with pytest.raises(exceptions.ValidationError):
+        validate(config2, SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA)

tests/unit/sagemaker/local/test_local_image.py

@@ -657,6 +657,23 @@ def test_container_does_not_enable_nvidia_docker_for_cpu_containers(sagemaker_se
     assert "runtime" not in docker_host
 
 
+def test_container_with_custom_config(sagemaker_session):
+    custom_config = {
+        "local": {
+            "container_config": {"shm_size": "128M"},
+        }
+    }
+    sagemaker_session.config = custom_config
+    instance_count = 1
+    image = "my-image"
+    sagemaker_container = _SageMakerContainer(
+        "local", instance_count, image, sagemaker_session=sagemaker_session
+    )
+
+    docker_host = sagemaker_container._create_docker_host("host-1", {}, set(), "train", [])
+    assert "shm_size" in docker_host
+
+
 @patch("sagemaker.local.image._HostingContainer.run", Mock())
 @patch("sagemaker.local.image._SageMakerContainer._prepare_serving_volumes", Mock(return_value=[]))
 @patch("shutil.copy", Mock())

tests/unit/sagemaker/local/test_local_session.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
+import jsonschema
 
 import pytest
 import urllib3
@@ -1031,3 +1032,52 @@ def test_default_bucket_prefix_with_sagemaker_config(boto_session, client):
         **session_kwargs,
     )
     assert session_with_no_prefix.default_bucket_prefix is None
+
+
+VALID_LOCAL_MODE_CONFIG = {
+    "local": {
+        "local_code": True,
+        "serving_port": 8888,
+        "container_config": {"shm_size": "128M"},
+    }
+}
+
+INVALID_LOCAL_MODE_CONFIG = {
+    "locals": {
+        "local_code": True,
+        "serving_port": 8888,
+        "container_config": {"shm_size": "128M"},
+    }
+}
+
+
+@patch("sagemaker.local.local_session.load_local_mode_config", return_value=VALID_LOCAL_MODE_CONFIG)
+def test_config_getter(load_config_mock):
+    boto_session = Mock(region_name="us-west-2")
+    session = LocalSession(boto_session=boto_session)
+    load_config_mock.assert_called()
+    assert session.config == VALID_LOCAL_MODE_CONFIG
+
+
+@patch(
+    "sagemaker.local.local_session.load_local_mode_config", return_value=INVALID_LOCAL_MODE_CONFIG
+)
+def test_config_validation(load_config_mock):
+    boto_session = Mock(region_name="us-west-2")
+
+    with pytest.raises(jsonschema.ValidationError):
+        LocalSession(boto_session=boto_session)
+
+
+def test_config_setter():
+    boto_session = Mock(region_name="us-west-2")
+
+    session = LocalSession(boto_session=boto_session)
+    session.config = VALID_LOCAL_MODE_CONFIG
+    assert (
+        session.sagemaker_runtime_client.serving_port
+        == VALID_LOCAL_MODE_CONFIG["local"]["serving_port"]
+    )
+
+    with pytest.raises(jsonschema.ValidationError):
+        session.config = INVALID_LOCAL_MODE_CONFIG