Support intelligent parameters (#1540)

pintaoz-aws · pintaoz-aws · commit 89ec0275eba1 · 2024-12-04T04:38:29.000-08:00
* Support intelligent parameters

* fix codestyle
diff --git a/src/sagemaker/modules/scripts/intelligent_params.py b/src/sagemaker/modules/scripts/intelligent_params.py
@@ -0,0 +1,101 @@
+import argparse
+import json
+import os
+import re
+
+HYPERPARAMETERS_FILE_PATH = "/opt/ml/input/config/hyperparameters.json"
+
+
+def set_intelligent_params(path: str) -> None:
+    """
+    Set intelligent parameters for all python files under the given path.
+    For python code with comment sm_hyper_param or sm_hp_{variable_name}, the value will be found in
+        /opt/ml/input/config/hyperparameters.json, and this function will rewrite lines with these comments.
+
+    Args:
+        path (str): The folder path to set intellingent parameters
+    """
+    with open(HYPERPARAMETERS_FILE_PATH, "r") as f:
+        hyperparameters = json.load(f)
+    for root, dirs, files in os.walk(path):
+        for file in files:
+            if file.endswith(".py"):
+                file_path = os.path.join(root, file)
+                rewrite_file(file_path, hyperparameters)
+
+
+def rewrite_file(file_path: str, hyperparameters: dict) -> None:
+    """
+    Rewrite a single python file with intelligent parameters.
+
+    Args:
+        file_path (str): The file path to rewrite
+        hyperparameters (dict): The hyperparameter names and values
+    """
+    with open(file_path, "r", encoding="utf-8") as f:
+        lines = f.readlines()
+    for i in range(len(lines)):
+        lines[i] = rewrite_line(lines[i], hyperparameters)
+    with open(file_path, "w", encoding="utf-8") as f:
+        f.writelines(lines)
+
+
+def rewrite_line(line: str, hyperparameters: dict) -> None:
+    """
+    Rewrite a single line of python code with intelligent parameters.
+
+    Args:
+        line (str): The python code to rewrite
+        hyperparameters (dict): The hyperparameter names and values
+    """
+    # Remove strings from the line to avoid = and # in strings
+    line_without_strings = re.sub(r'".*?"', '""', line.strip())
+    line_without_strings = re.sub(r"'.*?'", '""', line_without_strings)
+
+    # Match lines with format "a = 1 # comment"
+    assignment_pattern = r"^[a-zA-Z_][a-zA-Z0-9_]*\s*=.*#.*"
+    if re.match(assignment_pattern, line_without_strings):
+        indent = (len(line) - len(line.lstrip())) * " "
+        variable = line_without_strings.split("=")[0].strip()
+        comment = line_without_strings.split("#")[-1].strip()
+        value = get_parameter_value(variable, comment, hyperparameters)
+        if value is None:
+            return line
+        if isinstance(value, str):
+            new_line = f'{indent}{variable} = "{value}" # set by intelligent parameters\n'
+        else:
+            new_line = f"{indent}{variable} = {str(value)} # set by intelligent parameters\n"
+        return new_line
+    return line
+
+
+def get_parameter_value(variable: str, comment: str, hyperparameters: dict) -> None:
+    """
+    Get the parameter value by the variable name and comment.
+
+    Args:
+        variable (str): The variable name
+        comment (str): The comment string in the python code
+        hyperparameters (dict): The hyperparameter names and values
+    """
+    if comment == "sm_hyper_param":
+        # Get the hyperparameter value by the variable name
+        return hyperparameters.get(variable, None)
+    if comment.startswith("sm_hp_"):
+        # Get the hyperparameter value by the suffix of comment
+        return hyperparameters.get(comment[6:], None)
+    # Get the hyperparameter value from environment variables
+    if comment.startswith("sm_"):
+        return os.environ.get(comment.upper(), None)
+    return None
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Intelligent parameters")
+    parser.add_argument(
+        "-p", "--path", help="The folder path to set intellingent parameters", required=True
+    )
+
+    args = parser.parse_args()
+
+    set_intelligent_params(args.path)
diff --git a/src/sagemaker/modules/train/local_snapshot.py b/src/sagemaker/modules/train/local_snapshot.py
@@ -1,9 +1,23 @@
-import boto3
-import docker
+# Copyright 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.
+"""Utility function to capture local environment"""
 import logging
 import subprocess
 import sys
 from typing import Optional
+
+import boto3
+import docker
 import yaml
 
 logger = logging.getLogger(__name__)
@@ -66,38 +80,41 @@ def capture_local_environment(
     """
     Capture all dependency packages installed in the local environment and build a docker image.
     When using this utility method, the docker daemon must be active in the environment.
-    Please note that this is an experimental feature. This utility function is not be able to detect the package
-    compatability between platforms. It is also not able to detect dependency conflicts between the local environment
-    and the additional dependencies.
+    Please note that this is an experimental feature. This utility function is not be able to
+    detect the package compatability between platforms. It is also not able to detect dependency
+    conflicts between the local environment and the additional dependencies.
 
     Args:
         image_name (str): The name of the docker image.
-        env_name (str): The name of the virtual environment to be activated in the image, defaults to "saved_local_env".
+        env_name (str): The name of the virtual environment to be activated in the image,
+            defaults to "saved_local_env".
         package_manager (str): The package manager, must be one of "conda" or "pip".
-        deploy_to_ecr (bool): Whether to deploy the docker image to AWS ECR, defaults to False. If set to True, the AWS
-        credentials must be configured in the environment.
-        base_image_name (Optional[str]): If provided will be used as the base image, else the utility will evaluate
-            from local environment in following manner:
+        deploy_to_ecr (bool): Whether to deploy the docker image to AWS ECR, defaults to False.
+            If set to True, the AWS credentials must be configured in the environment.
+        base_image_name (Optional[str]): If provided will be used as the base image, else the
+            utility will evaluate from local environment in following manner:
                 1. If package manager is conda, it will use ubuntu:latest.
-                2. If package manager is pip, it is resolved to base python image with the same python version
-                    as the environment running the local code.
-        job_conda_env (Optional[str]): If set, the dependencies will be captured from this specific conda Env,
-            otherwise the dependencies will be the installed packages in the current active environment. This parameter
-            is only valid when the package manager is conda.
-        additional_dependencies (Optional[str]): Either the path to a dependencies file (conda environment.yml OR pip
-            requirements.txt file). Regardless of this setting utility will automatically generate the dependencies
-            file corresponding to the current active environment’s snapshot. In addition to this, additional dependencies
-            is configurable.
-        ecr_repo_name (Optional[str]): The AWS ECR repo to push the docker image. If not specified, it will use image_name as
-            the ECR repo name. This parameter is only valid when deploy_to_ecr is True.
-        boto_session (Optional[boto3.Session]): The boto3 session with AWS account info. If not provided, a new boto session
-            will be created.
+                2. If package manager is pip, it is resolved to base python image with the same
+                    python version as the environment running the local code.
+        job_conda_env (Optional[str]): If set, the dependencies will be captured from this specific
+            conda Env, otherwise the dependencies will be the installed packages in the current
+            active environment. This parameter is only valid when the package manager is conda.
+        additional_dependencies (Optional[str]): Either the path to a dependencies file (conda
+            environment.yml OR pip requirements.txt file). Regardless of this setting utility will
+            automatically generate the dependencies file corresponding to the current active
+            environment’s snapshot. In addition to this, additional dependencies is configurable.
+        ecr_repo_name (Optional[str]): The AWS ECR repo to push the docker image. If not specified,
+            it will use image_name as the ECR repo name. This parameter is only valid when
+            deploy_to_ecr is True.
+        boto_session (Optional[boto3.Session]): The boto3 session with AWS account info. If not
+            provided, a new boto session will be created.
 
     Exceptions:
         docker.errors.DockerException: Error while fetching server API version:
             The docker engine is not running in your environment.
-        docker.errors.BuildError: The docker failed to build the image. The most likely reason is: 1) Some packages are not
-            supported in the base image. 2) There are dependency conflicts between your local environment and additional dependencies.
+        docker.errors.BuildError: The docker failed to build the image. The most likely reason is:
+            1) Some packages are not supported in the base image. 2) There are dependency conflicts
+            between your local environment and additional dependencies.
         botocore.exceptions.ClientError: AWS credentials are not configured.
     """
 
@@ -118,7 +135,8 @@ def capture_local_environment(
                 ".yml"
             ) and not additional_dependencies.endswith(".txt"):
                 raise ValueError(
-                    "When package manager is conda, additional dependencies file must be a yml file or a txt file."
+                    "When package manager is conda, additional dependencies "
+                    "file must be a yml file or a txt file."
                 )
             if additional_dependencies.endswith(".yml"):
                 _merge_environment_ymls(
@@ -153,7 +171,7 @@ def capture_local_environment(
                 additional_requirements = f.read()
             with open(REQUIREMENT_TXT_PATH, "a") as f:
                 f.write(additional_requirements)
-                logger.info(f"Merged requirements file saved to {REQUIREMENT_TXT_PATH}")
+                logger.info("Merged requirements file saved to %s", REQUIREMENT_TXT_PATH)
 
             if not base_image_name:
                 version = sys.version_info
@@ -165,23 +183,24 @@ def capture_local_environment(
 
     else:
         raise ValueError(
-            "The provided package manager is not supported. Use conda or pip as the package manager."
+            "The provided package manager is not supported. "
+            "Use conda or pip as the package manager."
         )
 
     # Create the Dockerfile
     with open(DOCKERFILE_PATH, "w") as f:
         f.write(dockerfile_contents)
 
     client = docker.from_env()
-    image, logs = client.images.build(
+    _, logs = client.images.build(
         path="/tmp",
         dockerfile=DOCKERFILE_PATH,
         rm=True,
         tag=image_name,
     )
     for log in logs:
         logger.info(log.get("stream", "").strip())
-    logger.info(f"Docker image {image_name} built successfully")
+    logger.info("Docker image %s built successfully", image_name)
 
     if deploy_to_ecr:
         if boto_session is None:
@@ -232,14 +251,15 @@ def _merge_environment_ymls(env_name: str, env_file1: str, env_file2: str, outpu
     with open(output_file, "w") as f:
         yaml.dump(merged_env, f, sort_keys=False)
 
-    logger.info(f"Merged environment file saved to '{output_file}'")
+    logger.info("Merged environment file saved to '%s'", output_file)
 
 
 def _merge_environment_yml_with_requirement_txt(
     env_name: str, env_file: str, req_txt: str, output_file: str
 ):
     """
-    Merge an environment.yml file with a requirements.txt file and save to a new environment.yml file.
+    Merge an environment.yml file with a requirements.txt file and save to a new
+        environment.yml file.
 
     Args:
         env_name (str): The name of the virtual environment to be activated in the image.
@@ -278,7 +298,7 @@ def _merge_environment_yml_with_requirement_txt(
     with open(output_file, "w") as f:
         yaml.dump(merged_env, f, sort_keys=False)
 
-    logger.info(f"Merged environment file saved to '{output_file}'")
+    logger.info("Merged environment file saved to '%s'", output_file)
 
 
 def _push_image_to_ecr(image_name: str, ecr_repo_name: str, boto_session: Optional[boto3.Session]):
@@ -317,4 +337,4 @@ def _push_image_to_ecr(image_name: str, ecr_repo_name: str, boto_session: Option
     docker_push_cmd = f"docker push {ecr_image_uri}"
     subprocess.run(docker_push_cmd, shell=True, check=True)
 
-    logger.info(f"Image {image_name} pushed to {ecr_image_uri}")
+    logger.info("Image %s pushed to %s", image_name, ecr_image_uri)
diff --git a/tests/unit/sagemaker/modules/scripts/test_intelligent_params.py b/tests/unit/sagemaker/modules/scripts/test_intelligent_params.py
@@ -0,0 +1,109 @@
+import os
+import pytest
+from sagemaker.modules.scripts.intelligent_params import (
+    rewrite_file,
+    rewrite_line,
+)
+
+
+@pytest.fixture()
+def hyperparameters():
+    return {"n_estimators": 3, "epochs": 4, "state": "In Progress", "list_value": [1, 2, 3]}
+
+
+@pytest.fixture()
+def python_code():
+    return """
+import sagemaker
+
+def main(args):
+    n_estimators = 10 # sm_hp_n_estimators
+    state = "Not Started" # sm_hyper_param
+    random_seed = 0.1 # sm_hyper_param
+    
+    output_dir = "local/dir/" # sm_model_dir
+    epochs = 5 # sm_hp_epochs
+    input_data = [0, 0, 0] # sm_hp_list_value
+
+    # Load the Iris dataset
+    iris = load_iris()
+    y = iris.target
+
+    # Make predictions on the test set
+    y_pred = clf.predict(input_data)
+
+    accuracy = accuracy_score(y, y_pred) # calculate the accuracy
+    print(f"# Model accuracy: {accuracy:.2f}")
+"""
+
+
+@pytest.fixture()
+def expected_output_code():
+    return """
+import sagemaker
+
+def main(args):
+    n_estimators = 3 # set by intelligent parameters
+    state = "In Progress" # set by intelligent parameters
+    random_seed = 0.1 # sm_hyper_param
+    
+    output_dir = "/opt/ml/input" # set by intelligent parameters
+    epochs = 4 # set by intelligent parameters
+    input_data = [1, 2, 3] # set by intelligent parameters
+
+    # Load the Iris dataset
+    iris = load_iris()
+    y = iris.target
+
+    # Make predictions on the test set
+    y_pred = clf.predict(input_data)
+
+    accuracy = accuracy_score(y, y_pred) # calculate the accuracy
+    print(f"# Model accuracy: {accuracy:.2f}")
+"""
+
+
+def test_rewrite_line(hyperparameters):
+    line = "n_estimators = 4 # sm_hyper_param"
+    new_line = rewrite_line(line, hyperparameters)
+    assert new_line == "n_estimators = 3 # set by intelligent parameters\n"
+
+    line = "    epochs = 5 # sm_hp_epochs"
+    new_line = rewrite_line(line, hyperparameters)
+    assert new_line == "    epochs = 4 # set by intelligent parameters\n"
+
+    os.environ["SM_MODEL_DIR"] = "/opt/ml/input"
+    line = 'output_dir = "local/dir/" # sm_model_dir    '
+    new_line = rewrite_line(line, hyperparameters)
+    assert new_line == 'output_dir = "/opt/ml/input" # set by intelligent parameters\n'
+
+    line = "    random_state = 1 # not an intelligent parameter comment    \n"
+    new_line = rewrite_line(line, hyperparameters)
+    assert new_line == "    random_state = 1 # not an intelligent parameter comment    \n"
+
+    line = "not_an_intelligent_parameter = 4 # sm_hyper_param\n"
+    new_line = rewrite_line(line, hyperparameters)
+    assert new_line == "not_an_intelligent_parameter = 4 # sm_hyper_param\n"
+
+    line = "not_found_in_hyper_params = 4 # sm_hp_not_found_in_hyper_params\n"
+    new_line = rewrite_line(line, hyperparameters)
+    assert new_line == "not_found_in_hyper_params = 4 # sm_hp_not_found_in_hyper_params\n"
+
+    line = "list_value = [4, 5, 6] # sm_hyper_param"
+    new_line = rewrite_line(line, hyperparameters)
+    assert new_line == "list_value = [1, 2, 3] # set by intelligent parameters\n"
+
+
+def test_rewrite_file(hyperparameters, python_code, expected_output_code):
+    test_file_path = "temp_test.py"
+
+    os.environ["SM_MODEL_DIR"] = "/opt/ml/input"
+    with open(test_file_path, "w") as f:
+        f.write(python_code)
+    rewrite_file(test_file_path, hyperparameters)
+
+    with open(test_file_path, "r") as f:
+        new_python_code = f.read()
+    assert new_python_code == expected_output_code
+
+    os.remove(test_file_path)
