feature: make estimator accept json and yaml file as modelparallel config

Author: Yongyan Rao

src/sagemaker/fw_utils.py

@@ -13,6 +13,7 @@
 """Utility methods used by framework classes"""
 from __future__ import absolute_import
 
+import json
 import logging
 import os
 import re
@@ -21,6 +22,7 @@
 import tempfile
 from collections import namedtuple
 from typing import Optional, Union, Dict
+import yaml
 
 import sagemaker.image_uris
 from sagemaker.session_settings import SessionSettings
@@ -234,6 +236,45 @@ def validate_source_code_input_against_pipeline_variables(
         )
 
 
+def parse_mp_parameters(params):
+    """Parse the model parallelism parameters provided by the user.
+
+    Args:
+        params: a string representing path to an existing config, or
+                a config dict.
+
+    Returns:
+        parsed: a dict of parsed config.
+
+    Raises:
+        ValueError: if params is not a string or a dict, or
+                    the config file cannot be parsed as json or yaml.
+    """
+    parsed = None
+    if isinstance(params, dict):
+        parsed = params
+    elif os.path.exists(params):
+        try:
+            with open(params, "r") as fp:
+                parsed = json.load(fp)
+        except json.decoder.JSONDecodeError:
+            try:
+                with open(params, "r") as fp:
+                    parsed = yaml.load(fp)
+            except yaml.YAMLError:
+                pass
+    else:
+        raise ValueError(
+            f"Expected a string path to an existing modelparallel config, or a dictionary. "
+            f"Received: {params}."
+        )
+
+    if parsed is None:
+        raise ValueError(f"Cannot parse {params} as a json or yaml file.")
+
+    return parsed
+
+
 def get_mp_parameters(distribution):
     """Get the model parallelism parameters provided by the user.
 
@@ -250,6 +291,7 @@ def get_mp_parameters(distribution):
         mp_dict = {}
     if mp_dict.get("enabled", False) is True:
         params = mp_dict.get("parameters", {})
+        params = parse_mp_parameters(params)
         validate_mp_config(params)
         return params
     return None

tests/unit/test_fw_utils.py

@@ -13,10 +13,12 @@
 from __future__ import absolute_import
 
 import inspect
+import json
 import os
 import tarfile
 from contextlib import contextmanager
 from itertools import product
+import yaml
 
 import pytest
 
@@ -192,6 +194,61 @@ def test_validate_source_dir_file_not_in_dir():
         fw_utils.validate_source_dir(script, directory)
 
 
+def test_parse_mp_parameters_input_dict():
+    mp_parameters = {
+        "partitions": 1,
+        "tensor_parallel_degree": 2,
+        "microbatches": 1,
+        "optimize": "speed",
+        "pipeline": "interleaved",
+        "ddp": 1,
+        "auto_partition": False,
+        "default_partition": 0,
+    }
+    assert mp_parameters == fw_utils.parse_mp_parameters(mp_parameters)
+
+
+def test_parse_mp_parameters_input_str_json():
+    mp_parameters = {
+        "partitions": 1,
+        "tensor_parallel_degree": 2,
+        "microbatches": 1,
+        "optimize": "speed",
+        "pipeline": "interleaved",
+        "ddp": 1,
+        "auto_partition": False,
+        "default_partition": 0,
+    }
+    json_file_path = "./params.json"
+    with open(json_file_path, "x") as fp:
+        json.dump(mp_parameters, fp)
+    assert mp_parameters == fw_utils.parse_mp_parameters(json_file_path)
+    os.remove(json_file_path)
+
+
+def test_parse_mp_parameters_input_str_yaml():
+    mp_parameters = {
+        "partitions": 1,
+        "tensor_parallel_degree": 2,
+        "microbatches": 1,
+        "optimize": "speed",
+        "pipeline": "interleaved",
+        "ddp": 1,
+        "auto_partition": False,
+        "default_partition": 0,
+    }
+    yaml_file_path = "./params.yaml"
+    with open(yaml_file_path, "x") as fp:
+        yaml.dump(mp_parameters, fp)
+    assert mp_parameters == fw_utils.parse_mp_parameters(yaml_file_path)
+    os.remove(yaml_file_path)
+
+
+def test_parse_mp_parameters_input_not_exit():
+    with pytest.raises(ValueError):
+        fw_utils.parse_mp_parameters(" !@#$%^&*()path probably in not there.!@#$%^&*()")
+
+
 def test_tar_and_upload_dir_not_s3(sagemaker_session):
     bucket = "mybucket"
     s3_key_prefix = "something/source"