Support building image from Dockerfile (#1571)

pintaoz-aws · benieric · nargokul · pintaoz-aws · commit 8cf963107164 · 2024-12-04T01:41:19.000-08:00
* Base model trainer (#1521) * Base model trainer * flake8 * add testing notebook * add param validation & set defaults * Implement simple train method * feature: support script mode with local train.sh (#1523) * feature: support script mode with local train.sh * Stop tracking train.sh and add it to .gitignore * update message * make dir if not exist * fix docs * fix: docstyle * Address comments * fix hyperparams * Revert pydantic custom error * pylint * Image Spec refactoring and updates (#1525) * Image Spec refactoring and updates * Unit tests and update function for Image Spec * Fix hugging face test * Fix Tests * Add unit tests for ModelTrainer (#1527) * Add unit tests for ModelTrainer * Flake8 * format * Add example notebook (#1528) * Add testing notebook * format * use smaller data * remove large dataset * update * pylint * flake8 * ignore docstyle in directories with test * format * format * Add enviornment variable bootstrapping script (#1530) * Add enviornment variables scripts * format * fix comment * add docstrings * fix comment * feature: add utility function to capture local snapshot (#1524) * local snapshot * Update pip list command * Remove function calls * Address comments * Address comments * Support intelligent parameters (#1540) * Support intelligent parameters * fix codestyle * Revert Image Spec (#1541) * Cleanup ModelTrainer (#1542) * General image builder (#1546) * General image builder * General image builder * Fix codestyle * Fix codestyle * Move location * Add warnings * Add integ tests * Fix integ test * Fix integ test * Fix region error * Add region * Latest Container Image (#1545) * Latest Container Image * Test Fixes * Parameterized tests and some logic updates * Test fixes * Move to Image URI * Fixes for unit test * Fixes for unit test * Fix codestyle error checks * Cleanup ModelTrainer code (#1552) * feat: add pre-processing and post-processing logic to inference_spec (#1560) * add pre-processing and post-processing logic to inference_spec * fix format * make accept_type and content_type optional * remove accept_type and content_type from pre/post processing * correct typo * Add Distributed Training Support Model Trainer (#1536) * Add path to set Additional Settings in ModelTrainer (#1555) * Support building image from Dockerfile * Fix test * Fix test * Rename functions --------- Co-authored-by: Erick Benitez-Ramos <141277478+benieric@users.noreply.github.com> Co-authored-by: Gokul Anantha Narayanan <166456257+nargokul@users.noreply.github.com> Co-authored-by: Pravali Uppugunduri <46845440+pravali96@users.noreply.github.com>
diff --git a/src/sagemaker/modules/image_builder.py b/src/sagemaker/modules/image_builder.py
@@ -81,7 +81,7 @@
 """
 
 
-def build_image(
+def build_image_from_base(
     image_name: str = "sm-custom-image",
     env_name: str = "sm_custom_env",
     deploy_to_ecr: bool = False,
@@ -162,7 +162,7 @@ def build_image(
     return None
 
 
-def capture_local_environment(
+def build_image_from_local(
     image_name: str = "sm-local-capture",
     env_name: str = "saved_local_env",
     package_manager: str = "pip",
@@ -285,13 +285,13 @@ def capture_local_environment(
                 f.write(additional_requirements)
                 logger.info("Merged requirements file saved to %s", requirement_txt_path)
 
-            if not base_image:
-                version = sys.version_info
-                base_image = f"python:{version.major}.{version.minor}.{version.micro}"
-            dockerfile_contents = PIP_DOCKERFILE_TEMPLATE.format(
-                base_image=base_image,
-                env_name=env_name,
-            )
+        if not base_image:
+            version = sys.version_info
+            base_image = f"python:{version.major}.{version.minor}.{version.micro}"
+        dockerfile_contents = PIP_DOCKERFILE_TEMPLATE.format(
+            base_image=base_image,
+            env_name=env_name,
+        )
 
     else:
         raise ValueError(
@@ -306,6 +306,55 @@ def capture_local_environment(
     return None
 
 
+def build_image_from_dockerfile(
+    image_name: str,
+    dockerfile: str,
+    deploy_to_ecr: bool = False,
+    ecr_repo_name: Optional[str] = None,
+    boto_session: Optional[boto3.Session] = None,
+    region: Optional[str] = None,
+) -> Optional[str]:
+    """Build a Docker image with Dockerfile.
+
+    Args:
+        image_name (str): The name of the docker image.
+        dockerfile (str): The file path to the Dockerfile.
+        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.
+        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.
+        region (Optional[str]): The AWS region.
+
+    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.
+        botocore.exceptions.ClientError: AWS credentials are not configured.
+    """
+    absolute_path = os.path.abspath(dockerfile)
+    directory = os.path.dirname(absolute_path)
+
+    client = docker.from_env()
+    _, logs = client.images.build(
+        path=directory,
+        dockerfile=absolute_path,
+        rm=True,
+        tag=image_name,
+    )
+    for log in logs:
+        logger.info(log.get("stream", "").strip())
+    logger.info("Docker image %s built successfully", image_name)
+
+    if deploy_to_ecr:
+        return _push_image_to_ecr(image_name, ecr_repo_name, boto_session, region)
+    return None
+
+
 def _merge_environment_ymls(env_name: str, env_file1: str, env_file2: str, output_file: str):
     """Merge two environment.yml files and save to a new environment.yml file.
 
diff --git a/tests/integ/sagemaker/modules/Dockerfile b/tests/integ/sagemaker/modules/Dockerfile
@@ -0,0 +1,5 @@
+FROM python:3.12.4
+
+WORKDIR /app
+
+RUN pip install --no-cache-dir numpy scipy sagemaker
diff --git a/tests/integ/sagemaker/modules/test_image_builder.py b/tests/integ/sagemaker/modules/test_image_builder.py
@@ -15,24 +15,33 @@
 
 import os
 
-from sagemaker.modules.image_builder import build_image
+from sagemaker.modules.image_builder import build_image_from_base, build_image_from_dockerfile
 
 
 def test_build_public_image_locally():
-    build_image(image_name="python_310", base_image="python:3.10")
+    build_image_from_base(image_name="python_310", base_image="python:3.10")
 
 
 def test_build_with_dependency_file():
     dependency_file_path = os.getcwd() + "/tests/integ/sagemaker/modules/requirements.txt"
-    build_image(image_name="ubuntu_with_dependencies", dependency_file=dependency_file_path)
+    build_image_from_base(
+        image_name="ubuntu_with_dependencies", dependency_file=dependency_file_path
+    )
 
 
 def test_build_image_and_push_to_ecr():
     dependency_file_path = os.getcwd() + "/tests/integ/sagemaker/modules/environment.yml"
-    build_image(
+    build_image_from_base(
         image_name="ecr_test_image",
         dependency_file=dependency_file_path,
         base_image="debian",
         deploy_to_ecr=True,
         ecr_repo_name="image_builder_integ_test",
     )
+
+
+def test_build_image_from_dockerfile():
+    dockerfile_path = os.getcwd() + "/tests/integ/sagemaker/modules/Dockerfile"
+    build_image_from_dockerfile(
+        image_name="image_from_dockerfile", dockerfile=dockerfile_path, deploy_to_ecr=True
+    )
