Support building image from Dockerfile (#1571)

* 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 <[email protected]>
Co-authored-by: Gokul Anantha Narayanan <[email protected]>
Co-authored-by: Pravali Uppugunduri <[email protected]>

Author: 4 people

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.
 

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

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
+    )