update warning message

Author: Chuyang Deng

src/sagemaker/fw_utils.py

@@ -49,8 +49,8 @@
     "Please set the argument \"py_version='py3'\" to use the Python 3 {framework} image."
 )
 PARAMETER_SERVER_MULTI_GPU_WARNING = (
-    "You have selected a multi-GPU training instance type. "
-    "You have also enabled parameter server for distributed training. "
+    "If you have selected a multi-GPU training instance type, "
+    "and have also enabled parameter server for distributed training. "
     "Distributed training with the default parameter server configuration will not "
     "fully leverage all GPU cores; the parameter server will be configured to run "
     "only one worker per host regardless of the number of GPUs."
@@ -617,9 +617,9 @@ def warn_if_parameter_server_with_multi_gpu(training_instance_type, distribution
         return
 
     is_multi_gpu_instance = (
-        training_instance_type.split(".")[1].startswith("p")
-        and training_instance_type not in SINGLE_GPU_INSTANCE_TYPES
-    )
+        training_instance_type == "local_gpu"
+        or training_instance_type.split(".")[1].startswith("p")
+    ) and training_instance_type not in SINGLE_GPU_INSTANCE_TYPES
 
     ps_enabled = "parameter_server" in distributions and distributions["parameter_server"].get(
         "enabled", False

tests/unit/test_fw_utils.py

@@ -1272,3 +1272,13 @@ def test_warn_if_parameter_server_with_multi_gpu(caplog):
         training_instance_type=train_instance_type, distributions=distributions
     )
     assert fw_utils.PARAMETER_SERVER_MULTI_GPU_WARNING in caplog.text
+
+
+def test_war_if_parameter_server_with_multi_gpu(caplog):
+    train_instance_type = "local_gpu"
+    distributions = {"parameter_server": {"enabled": True}}
+
+    fw_utils.warn_if_parameter_server_with_multi_gpu(
+        training_instance_type=train_instance_type, distributions=distributions
+    )
+    assert fw_utils.PARAMETER_SERVER_MULTI_GPU_WARNING in caplog.text