fix: move sagemaker config loading to LocalSession since it is only used for local code support. (aws#1137)

* Move sagemaker config loading to local session since it is only used for local code support.

* Add a common error message function.

* FIx imports order.

* Improve function doc string.

Author: nadiaya

setup.py

@@ -39,14 +39,17 @@ def read_version():
     "protobuf>=3.1",
     "scipy>=0.19.0",
     "protobuf3-to-dict>=0.1.5",
-    "PyYAML>=3.10, <5",  # required for session
     "requests>=2.20.0, <2.21",
 ]
 
 # Specific use case dependencies
 extras = {
     "analytics": ["pandas"],
-    "local": ["urllib3>=1.21, <1.25", "docker-compose>=1.23.0"],
+    "local": [
+        "urllib3>=1.21, <1.25",
+        "docker-compose>=1.23.0",
+        "PyYAML>=3.10, <5",  # PyYAML version has to match docker-compose requirements
+    ],
     "tensorflow": ["tensorflow>=1.3.0"],
 }
 # Meta dependency groups

src/sagemaker/local/image.py

@@ -33,8 +33,6 @@
 from threading import Thread
 from six.moves.urllib.parse import urlparse
 
-import yaml
-
 import sagemaker
 import sagemaker.local.data
 import sagemaker.local.utils
@@ -461,6 +459,13 @@ def _generate_compose_file(self, command, additional_volumes=None, additional_en
         }
 
         docker_compose_path = os.path.join(self.container_root, DOCKER_COMPOSE_FILENAME)
+
+        try:
+            import yaml
+        except ImportError as e:
+            logging.error(sagemaker.utils._module_import_error("yaml", "Local mode", "local"))
+            raise e
+
         yaml_content = yaml.dump(content, default_flow_style=False)
         logger.info("docker compose file: \n%s", yaml_content)
         with open(docker_compose_path, "w") as f:

src/sagemaker/local/local_session.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 import logging
+import os
 import platform
 
 import boto3
@@ -28,14 +29,7 @@
     _LocalTransformJob,
 )
 from sagemaker.session import Session
-from sagemaker.utils import DeferredError, get_config_value
-
-try:
-    import urllib3
-except ImportError as e:
-    logging.warning("urllib3 failed to import. Local mode features will be impaired or broken.")
-    # Any subsequent attempt to use urllib3 will raise the ImportError
-    urllib3 = DeferredError(e)
+from sagemaker.utils import get_config_value, _module_import_error
 
 logger = logging.getLogger(__name__)
 
@@ -329,6 +323,12 @@ def __init__(self, config=None):
             config (dict): Optional configuration for this client. In particular only
                 the local port is read.
         """
+        try:
+            import urllib3
+        except ImportError as e:
+            logging.error(_module_import_error("urllib3", "Local mode", "local"))
+            raise e
+
         self.http = urllib3.PoolManager()
         self.serving_port = 8080
         self.config = config
@@ -403,6 +403,16 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
         self.sagemaker_runtime_client = LocalSagemakerRuntimeClient(self.config)
         self.local_mode = True
 
+        sagemaker_config_file = os.path.join(os.path.expanduser("~"), ".sagemaker", "config.yaml")
+        if os.path.exists(sagemaker_config_file):
+            try:
+                import yaml
+            except ImportError as e:
+                logging.error(_module_import_error("yaml", "Local mode", "local"))
+                raise e
+
+            self.config = yaml.load(open(sagemaker_config_file, "r"))
+
     def logs_for_job(self, job_name, wait=False, poll=5):
         """
 

src/sagemaker/session.py

@@ -25,7 +25,6 @@
 import boto3
 import botocore.config
 from botocore.exceptions import ClientError
-import yaml
 
 import sagemaker.logs
 from sagemaker import vpc_utils
@@ -95,11 +94,8 @@ def __init__(self, boto_session=None, sagemaker_client=None, sagemaker_runtime_c
         """
         self._default_bucket = None
 
-        sagemaker_config_file = os.path.join(os.path.expanduser("~"), ".sagemaker", "config.yaml")
-        if os.path.exists(sagemaker_config_file):
-            self.config = yaml.load(open(sagemaker_config_file, "r"))
-        else:
-            self.config = None
+        # currently is used for local_code in local mode
+        self.config = None
 
         self._initialize(boto_session, sagemaker_client, sagemaker_runtime_client)
 

src/sagemaker/utils.py

@@ -608,3 +608,22 @@ def __getattr__(self, name):
             name:
         """
         raise self.exc
+
+
+def _module_import_error(py_module, feature, extras):
+    """Return error message for module import errors, provide
+    installation details.
+
+    Args:
+        py_module (str): Module that failed to be imported
+        feature (str): Affected sagemaker feature
+        extras (str): Name of the extra_requirements to install all of the dependencies
+    Returns:
+        str: Error message with installation instructions.
+    """
+    error_msg = (
+        "Failed to import {}. {} features will be impaired or broken."
+        "Please run \"pip install 'sagemaker[{}]'\" "
+        "to install all required dependencies."
+    )
+    return error_msg.format(py_module, feature, extras)