sagemaker-core 1.0.62__py3-none-any.whl → 2.3.1__py3-none-any.whl

sagemaker/core/local/image.py

@@ -0,0 +1,1243 @@
+# 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.
+"""Placeholder docstring"""
+from __future__ import absolute_import, annotations
+
+import base64
+import copy
+import errno
+import json
+import logging
+import os
+import platform
+import random
+import re
+import shlex
+import shutil
+import string
+import subprocess
+import sys
+import tarfile
+import tempfile
+
+from threading import Thread
+from typing import Dict, List
+from six.moves.urllib.parse import urlparse
+
+import sagemaker
+from sagemaker.core.config.config_schema import CONTAINER_CONFIG, LOCAL
+import sagemaker.core
+from sagemaker.core.common_utils import custom_extractall_tarfile
+
+CONTAINER_PREFIX = "algo"
+STUDIO_HOST_NAME = "sagemaker-local"
+DOCKER_COMPOSE_FILENAME = "docker-compose.yaml"
+DOCKER_COMPOSE_HTTP_TIMEOUT_ENV = "COMPOSE_HTTP_TIMEOUT"
+DOCKER_COMPOSE_HTTP_TIMEOUT = "120"
+
+# Environment variables to be set during training
+REGION_ENV_NAME = "AWS_REGION"
+TRAINING_JOB_NAME_ENV_NAME = "TRAINING_JOB_NAME"
+S3_ENDPOINT_URL_ENV_NAME = "S3_ENDPOINT_URL"
+SM_STUDIO_LOCAL_MODE = "SM_STUDIO_LOCAL_MODE"
+
+# SELinux Enabled
+SELINUX_ENABLED = os.environ.get("SAGEMAKER_LOCAL_SELINUX_ENABLED", "False").lower() in [
+    "1",
+    "true",
+    "yes",
+]
+
+logger = logging.getLogger(__name__)
+
+
+class _SageMakerContainer(object):
+    """Handle the lifecycle and configuration of a local container execution.
+
+    This class is responsible for creating the directories and configuration
+    files that the docker containers will use for either training or serving.
+    """
+
+    def __init__(
+        self,
+        instance_type,
+        instance_count,
+        image,
+        sagemaker_session=None,
+        container_entrypoint=None,
+        container_arguments=None,
+    ):
+        """Initialize a SageMakerContainer instance
+
+        It uses a :class:`sagemaker.core.helper.session.Session` for general interaction
+        with user configuration such as getting the default sagemaker S3 bucket.
+        However this class does not call any of the SageMaker APIs.
+
+        Args:
+            instance_type (str): The instance type to use. Either 'local' or
+                'local_gpu'
+            instance_count (int): The number of instances to create.
+            image (str): docker image to use.
+            sagemaker_session (sagemaker.core.helper.session.Session): a sagemaker session
+                to use when interacting with SageMaker.
+            container_entrypoint (str): the container entrypoint to execute
+            container_arguments (str): the container entrypoint arguments
+        """
+        from sagemaker.core.local.local_session import LocalSession
+
+        # check if docker-compose is installed
+
+        self.compose_cmd_prefix = _SageMakerContainer._get_compose_cmd_prefix()
+        self.sagemaker_session = sagemaker_session or LocalSession()
+        self.instance_type = instance_type
+        self.instance_count = instance_count
+        self.image = image
+        self.container_entrypoint = container_entrypoint
+        self.container_arguments = container_arguments
+        # Since we are using a single docker network, Generate a random suffix to attach to the
+        # container names. This way multiple jobs can run in parallel.
+        suffix = "".join(random.choice(string.ascii_lowercase + string.digits) for _ in range(5))
+        self.is_studio = sagemaker.core.local.utils.check_for_studio()
+        if self.is_studio:
+            if self.instance_count > 1:
+                raise NotImplementedError(
+                    "Multi instance Local Mode execution is "
+                    "currently not supported in SageMaker Studio."
+                )
+            # For studio use-case, directories need to be created in `~/tmp`, rather than /tmp
+            home = os.path.expanduser("~")
+            root_dir = os.path.join(home, "tmp")
+            if not os.path.isdir(root_dir):
+                os.mkdir(root_dir)
+            if self.sagemaker_session.config:
+                self.sagemaker_session.config["local"]["container_root"] = root_dir
+            else:
+                self.sagemaker_session.config = {"local": {"container_root": root_dir}}
+            # Studio only supports single instance run
+            self.hosts = [STUDIO_HOST_NAME]
+        else:
+            self.hosts = [
+                "{}-{}-{}".format(CONTAINER_PREFIX, i, suffix)
+                for i in range(1, self.instance_count + 1)
+            ]
+
+        self.container_root = None
+        self.container = None
+
+    @staticmethod
+    def _get_compose_cmd_prefix():
+        """Gets the Docker Compose command.
+
+        The method initially looks for 'docker compose' v2
+        executable, if not found looks for 'docker-compose' executable.
+
+        Returns:
+            Docker Compose executable split into list.
+
+        Raises:
+            ImportError: If Docker Compose executable was not found.
+        """
+        compose_cmd_prefix = []
+
+        output = None
+        try:
+            output = subprocess.check_output(
+                ["docker", "compose", "version"],
+                stderr=subprocess.DEVNULL,
+                encoding="UTF-8",
+            )
+        except subprocess.CalledProcessError:
+            logger.info(
+                "'Docker Compose' is not installed. "
+                "Proceeding to check for 'docker-compose' CLI."
+            )
+
+        if output and "v2" in output.strip():
+            logger.info("'Docker Compose' found using Docker CLI.")
+            compose_cmd_prefix.extend(["docker", "compose"])
+            return compose_cmd_prefix
+
+        if shutil.which("docker-compose") is not None:
+            logger.info("'Docker Compose' found using Docker Compose CLI.")
+            compose_cmd_prefix.extend(["docker-compose"])
+            return compose_cmd_prefix
+
+        raise ImportError(
+            "Docker Compose is not installed. "
+            "Local Mode features will not work without docker compose. "
+            "For more information on how to install 'docker compose', please, see "
+            "https://docs.docker.com/compose/install/"
+        )
+
+    def process(
+        self,
+        processing_inputs,
+        processing_output_config,
+        environment,
+        processing_job_name,
+    ):
+        """Run a processing job locally using docker-compose.
+
+        Args:
+            processing_inputs (dict): The processing input specification.
+            processing_output_config (dict): The processing output configuration specification.
+            environment (dict): The environment collection for the processing job.
+            processing_job_name (str): Name of the local processing job being run.
+        """
+
+        self.container_root = self._create_tmp_folder()
+
+        # A shared directory for all the containers;
+        # it is only mounted if the processing script is Local.
+        shared_dir = os.path.join(self.container_root, "shared")
+        os.mkdir(shared_dir)
+
+        data_dir = self._create_tmp_folder()
+        volumes = self._prepare_processing_volumes(
+            data_dir, processing_inputs, processing_output_config
+        )
+
+        # Create the configuration files for each container that we will create.
+        for host in self.hosts:
+            _create_processing_config_file_directories(self.container_root, host)
+            self.write_processing_config_files(
+                host,
+                environment,
+                processing_inputs,
+                processing_output_config,
+                processing_job_name,
+            )
+
+        self._generate_compose_file(
+            "process", additional_volumes=volumes, additional_env_vars=environment
+        )
+
+        if _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image):
+            _pull_image(self.image)
+
+        compose_command = self._compose()
+        process = subprocess.Popen(
+            compose_command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+        )
+
+        try:
+            _stream_output(process)
+        finally:
+            # Uploading processing outputs back to Amazon S3.
+            self._upload_processing_outputs(data_dir, processing_output_config)
+
+            try:
+                # Deleting temporary directories.
+                dirs_to_delete = [shared_dir, data_dir]
+                self._cleanup(dirs_to_delete)
+            except OSError:
+                pass
+
+        # Print our Job Complete line to have a similar experience to training on SageMaker where
+        # you see this line at the end.
+        logger.info("===== Job Complete =====")
+
+    def train(self, input_data_config, output_data_config, hyperparameters, environment, job_name):
+        """Run a training job locally using docker-compose.
+
+        Args:
+            input_data_config (dict): The Input Data Configuration, this contains data such as the
+                channels to be used for training.
+            output_data_config: The configuration of the output data.
+            hyperparameters (dict): The HyperParameters for the training job.
+            environment (dict): The environment collection for the training job.
+            job_name (str): Name of the local training job being run.
+
+        Returns (str): Location of the trained model.
+        """
+        self.container_root = self._create_tmp_folder()
+        os.mkdir(os.path.join(self.container_root, "output"))
+        # create output/data folder since sagemaker-containers 2.0 expects it
+        os.mkdir(os.path.join(self.container_root, "output", "data"))
+        # A shared directory for all the containers. It is only mounted if the training script is
+        # Local.
+        shared_dir = os.path.join(self.container_root, "shared")
+        os.mkdir(shared_dir)
+
+        data_dir = self._create_tmp_folder()
+        volumes = self._prepare_training_volumes(
+            data_dir, input_data_config, output_data_config, hyperparameters
+        )
+        # If local, source directory needs to be updated to mounted /opt/ml/code path
+        hyperparameters = self._update_local_src_path(
+            hyperparameters, key=sagemaker.serve.model_builder.DIR_PARAM_NAME
+        )
+
+        # Create the configuration files for each container that we will create
+        # Each container will map the additional local volumes (if any).
+        for host in self.hosts:
+            _create_config_file_directories(self.container_root, host)
+            self.write_config_files(host, hyperparameters, input_data_config)
+            shutil.copytree(data_dir, os.path.join(self.container_root, host, "input", "data"))
+
+        training_env_vars = {
+            REGION_ENV_NAME: self.sagemaker_session.boto_region_name,
+            TRAINING_JOB_NAME_ENV_NAME: job_name,
+        }
+        training_env_vars.update(environment)
+        if self.sagemaker_session.s3_resource is not None:
+            training_env_vars[S3_ENDPOINT_URL_ENV_NAME] = (
+                self.sagemaker_session.s3_resource.meta.client._endpoint.host
+            )
+
+        compose_data = self._generate_compose_file(
+            "train", additional_volumes=volumes, additional_env_vars=training_env_vars
+        )
+
+        if _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image):
+            _pull_image(self.image)
+
+        compose_command = self._compose()
+        process = subprocess.Popen(
+            compose_command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+        )
+
+        try:
+            _stream_output(process)
+        finally:
+            artifacts = self.retrieve_artifacts(compose_data, output_data_config, job_name)
+
+            # free up the training data directory as it may contain
+            # lots of data downloaded from S3. This doesn't delete any local
+            # data that was just mounted to the container.
+            dirs_to_delete = [data_dir, shared_dir]
+            self._cleanup(dirs_to_delete)
+
+        # Print our Job Complete line to have a similar experience to training on SageMaker where
+        # you see this line at the end.
+        logger.info("===== Job Complete =====")
+        return artifacts
+
+    def serve(self, model_dir, environment):
+        """Host a local endpoint using docker-compose.
+
+        Args:
+            primary_container (dict): dictionary containing the container runtime settings
+                for serving. Expected keys:
+                - 'ModelDataUrl' pointing to a file or s3:// location.
+                - 'Environment' a dictionary of environment variables to be passed to the
+                    hosting container.
+        """
+        logger.info("serving")
+
+        self.container_root = self._create_tmp_folder()
+        logger.info("creating hosting dir in %s", self.container_root)
+
+        volumes = self._prepare_serving_volumes(model_dir)
+
+        # If the user script was passed as a file:// mount it to the container.
+        if sagemaker.serve.model_builder.DIR_PARAM_NAME.upper() in environment:
+            script_dir = environment[sagemaker.serve.model_builder.DIR_PARAM_NAME.upper()]
+            parsed_uri = urlparse(script_dir)
+            if parsed_uri.scheme == "file":
+                host_dir = os.path.abspath(parsed_uri.netloc + parsed_uri.path)
+                volumes.append(_Volume(host_dir, "/opt/ml/code"))
+                # Update path to mount location
+                environment = environment.copy()
+                environment[sagemaker.serve.model_builder.DIR_PARAM_NAME.upper()] = "/opt/ml/code"
+
+        if _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image):
+            _pull_image(self.image)
+
+        self._generate_compose_file(
+            "serve", additional_env_vars=environment, additional_volumes=volumes
+        )
+
+        compose_command = self._compose()
+
+        self.container = _HostingContainer(compose_command)
+        self.container.start()
+
+    def stop_serving(self):
+        """Stop the serving container.
+
+        The serving container runs in async mode to allow the SDK to do other
+        tasks.
+        """
+        if self.container:
+            self.container.down()
+            self.container.join()
+            self._cleanup()
+        # for serving we can delete everything in the container root.
+        _delete_tree(self.container_root)
+
+    def retrieve_artifacts(self, compose_data, output_data_config, job_name):
+        """Get the model artifacts from all the container nodes.
+
+        Used after training completes to gather the data from all the
+        individual containers. As the official SageMaker Training Service, it
+        will override duplicate files if multiple containers have the same file
+        names.
+
+        Args:
+            compose_data (dict): Docker-Compose configuration in dictionary
+                format.
+            output_data_config: The configuration of the output data.
+            job_name: The name of the job.
+
+        Returns: Local path to the collected model artifacts.
+        """
+        # We need a directory to store the artfiacts from all the nodes
+        # and another one to contained the compressed final artifacts
+        artifacts = os.path.join(self.container_root, "artifacts")
+        compressed_artifacts = os.path.join(self.container_root, "compressed_artifacts")
+        os.mkdir(artifacts)
+
+        model_artifacts = os.path.join(artifacts, "model")
+        output_artifacts = os.path.join(artifacts, "output")
+
+        artifact_dirs = [model_artifacts, output_artifacts, compressed_artifacts]
+        for d in artifact_dirs:
+            os.mkdir(d)
+
+        # Gather the artifacts from all nodes into artifacts/model and artifacts/output
+        for host in self.hosts:
+            volumes = compose_data["services"][str(host)]["volumes"]
+            volumes = [v[:-2] if v.endswith(":z") else v for v in volumes]
+            for volume in volumes:
+                if re.search(r"^[A-Za-z]:", volume):
+                    unit, host_dir, container_dir = volume.split(":")
+                    host_dir = unit + ":" + host_dir
+                else:
+                    host_dir, container_dir = volume.split(":")
+                if container_dir == "/opt/ml/model":
+                    sagemaker.core.local.utils.recursive_copy(host_dir, model_artifacts)
+                elif container_dir == "/opt/ml/output":
+                    sagemaker.core.local.utils.recursive_copy(host_dir, output_artifacts)
+
+        # Tar Artifacts -> model.tar.gz and output.tar.gz
+        model_files = [os.path.join(model_artifacts, name) for name in os.listdir(model_artifacts)]
+        output_files = [
+            os.path.join(output_artifacts, name) for name in os.listdir(output_artifacts)
+        ]
+        sagemaker.core.common_utils.create_tar_file(
+            model_files, os.path.join(compressed_artifacts, "model.tar.gz")
+        )
+        sagemaker.core.common_utils.create_tar_file(
+            output_files, os.path.join(compressed_artifacts, "output.tar.gz")
+        )
+
+        if output_data_config["S3OutputPath"] == "":
+            output_data = "file://%s" % compressed_artifacts
+        else:
+            # Now we just need to move the compressed artifacts to wherever they are required
+            output_data = sagemaker.core.local.utils.move_to_destination(
+                compressed_artifacts,
+                output_data_config["S3OutputPath"],
+                job_name,
+                self.sagemaker_session,
+                prefix="output",
+            )
+
+        _delete_tree(model_artifacts)
+        _delete_tree(output_artifacts)
+
+        return os.path.join(output_data, "model.tar.gz")
+
+    def write_processing_config_files(
+        self,
+        host,
+        environment,
+        processing_inputs,
+        processing_output_config,
+        processing_job_name,
+    ):
+        """Write the config files for the processing containers.
+
+        This method writes the hyperparameters, resources and input data
+        configuration files.
+
+        Args:
+            host (str): Host to write the configuration for
+            environment (dict): Environment variable collection.
+            processing_inputs (dict): Processing inputs.
+            processing_output_config (dict): Processing output configuration.
+            processing_job_name (str): Processing job name.
+        """
+        config_path = os.path.join(self.container_root, host, "config")
+
+        resource_config = {
+            "current_host": host,
+            "hosts": self.hosts,
+            "network_interface_name": "eth0",
+            "current_instance_type": self.instance_type,
+        }
+        _write_json_file(os.path.join(config_path, "resourceconfig.json"), resource_config)
+
+        processing_job_config = {
+            "ProcessingJobArn": processing_job_name,
+            "ProcessingJobName": processing_job_name,
+            "AppSpecification": {
+                "ImageUri": self.image,
+                "ContainerEntrypoint": self.container_entrypoint,
+                "ContainerArguments": self.container_arguments,
+            },
+            "Environment": environment,
+            "ProcessingInputs": processing_inputs,
+            "ProcessingOutputConfig": processing_output_config,
+            "ProcessingResources": {
+                "ClusterConfig": {
+                    "InstanceCount": self.instance_count,
+                    "InstanceType": self.instance_type,
+                    "VolumeSizeInGB": 30,
+                    "VolumeKmsKeyId": None,
+                }
+            },
+            "RoleArn": "<no_role>",
+            "StoppingCondition": {"MaxRuntimeInSeconds": 86400},
+        }
+
+        _write_json_file(
+            os.path.join(config_path, "processingjobconfig.json"), processing_job_config
+        )
+
+    def write_config_files(self, host, hyperparameters, input_data_config):
+        """Write the config files for the training containers.
+
+        This method writes the hyperparameters, resources and input data
+        configuration files.
+
+        Returns: None
+
+        Args:
+            host (str): Host to write the configuration for
+            hyperparameters (dict): Hyperparameters for training.
+            input_data_config (dict): Training input channels to be used for
+                training.
+        """
+        config_path = os.path.join(self.container_root, host, "input", "config")
+
+        resource_config = {
+            "current_host": host,
+            "hosts": self.hosts,
+            "network_interface_name": "eth0",
+            "current_instance_type": self.instance_type,
+        }
+
+        json_input_data_config = {}
+        for c in input_data_config:
+            channel_name = c["ChannelName"]
+            json_input_data_config[channel_name] = {"TrainingInputMode": "File"}
+            if "ContentType" in c:
+                json_input_data_config[channel_name]["ContentType"] = c["ContentType"]
+
+        _write_json_file(os.path.join(config_path, "hyperparameters.json"), hyperparameters)
+        _write_json_file(os.path.join(config_path, "resourceconfig.json"), resource_config)
+        _write_json_file(os.path.join(config_path, "inputdataconfig.json"), json_input_data_config)
+
+    def _prepare_training_volumes(
+        self, data_dir, input_data_config, output_data_config, hyperparameters
+    ):
+        """Prepares the training volumes based on input and output data configs.
+
+        Args:
+            data_dir:
+            input_data_config:
+            output_data_config:
+            hyperparameters:
+        """
+        shared_dir = os.path.join(self.container_root, "shared")
+        model_dir = os.path.join(self.container_root, "model")
+        volumes = []
+
+        volumes.append(_Volume(model_dir, "/opt/ml/model"))
+
+        # Mount the metadata directory if present.
+        # Only expected to be present on SM notebook instances.
+        # This is used by some DeepEngine libraries
+        metadata_dir = "/opt/ml/metadata"
+        if os.path.isdir(metadata_dir):
+            volumes.append(_Volume(metadata_dir, metadata_dir))
+
+        # Set up the channels for the containers. For local data we will
+        # mount the local directory to the container. For S3 Data we will download the S3 data
+        # first.
+        for channel in input_data_config:
+            uri = channel["DataUri"]
+            channel_name = channel["ChannelName"]
+            channel_dir = os.path.join(data_dir, channel_name)
+            os.mkdir(channel_dir)
+
+            data_source = sagemaker.core.local.data.get_data_source_instance(
+                uri, self.sagemaker_session
+            )
+            volumes.append(_Volume(data_source.get_root_dir(), channel=channel_name))
+
+        # If there is a training script directory and it is a local directory,
+        # mount it to the container.
+        if sagemaker.serve.model_builder.DIR_PARAM_NAME in hyperparameters:
+            training_dir = json.loads(hyperparameters[sagemaker.serve.model_builder.DIR_PARAM_NAME])
+            parsed_uri = urlparse(training_dir)
+            if parsed_uri.scheme == "file":
+                host_dir = os.path.abspath(parsed_uri.netloc + parsed_uri.path)
+                volumes.append(_Volume(host_dir, "/opt/ml/code"))
+                # Also mount a directory that all the containers can access.
+                volumes.append(_Volume(shared_dir, "/opt/ml/shared"))
+
+        parsed_uri = urlparse(output_data_config["S3OutputPath"])
+        if (
+            parsed_uri.scheme == "file"
+            and sagemaker.serve.model_builder.SAGEMAKER_OUTPUT_LOCATION in hyperparameters
+        ):
+            dir_path = os.path.abspath(parsed_uri.netloc + parsed_uri.path)
+            intermediate_dir = os.path.join(dir_path, "output", "intermediate")
+            if not os.path.exists(intermediate_dir):
+                os.makedirs(intermediate_dir)
+            volumes.append(_Volume(intermediate_dir, "/opt/ml/output/intermediate"))
+
+        return volumes
+
+    def _prepare_processing_volumes(self, data_dir, processing_inputs, processing_output_config):
+        """Prepares local container volumes for the processing job.
+
+        Args:
+            data_dir: The local data directory.
+            processing_inputs: The configuration of processing inputs.
+            processing_output_config: The configuration of processing outputs.
+
+        Returns:
+            The volumes configuration.
+        """
+        shared_dir = os.path.join(self.container_root, "shared")
+        volumes = []
+
+        # Set up the input/outputs for the container.
+
+        for item in processing_inputs:
+            uri = item["DataUri"]
+            input_container_dir = item["S3Input"]["LocalPath"]
+
+            data_source = sagemaker.core.local.data.get_data_source_instance(
+                uri, self.sagemaker_session
+            )
+            volumes.append(_Volume(data_source.get_root_dir(), input_container_dir))
+
+        if processing_output_config and "Outputs" in processing_output_config:
+            for item in processing_output_config["Outputs"]:
+                output_name = item["OutputName"]
+                output_container_dir = item["S3Output"]["LocalPath"]
+
+                output_dir = os.path.join(data_dir, "output", output_name)
+                os.makedirs(output_dir)
+
+                volumes.append(_Volume(output_dir, output_container_dir))
+
+        volumes.append(_Volume(shared_dir, "/opt/ml/shared"))
+
+        return volumes
+
+    def _upload_processing_outputs(self, data_dir, processing_output_config):
+        """Uploads processing outputs to Amazon S3.
+
+        Args:
+            data_dir: The local data directory.
+            processing_output_config: The processing output configuration.
+        """
+        if processing_output_config and "Outputs" in processing_output_config:
+            for item in processing_output_config["Outputs"]:
+                output_name = item["OutputName"]
+                output_s3_uri = item["S3Output"]["S3Uri"]
+                output_dir = os.path.join(data_dir, "output", output_name)
+
+                sagemaker.core.local.utils.move_to_destination(
+                    output_dir, output_s3_uri, "", self.sagemaker_session
+                )
+
+    def _update_local_src_path(self, params, key):
+        """Updates the local path of source code.
+
+        Args:
+            params: Existing configuration parameters.
+            key: Lookup key for the path of the source code in the configuration parameters.
+
+        Returns:
+            The updated parameters.
+        """
+        if key in params:
+            src_dir = json.loads(params[key])
+            parsed_uri = urlparse(src_dir)
+            if parsed_uri.scheme == "file":
+                new_params = params.copy()
+                new_params[key] = json.dumps("/opt/ml/code")
+                return new_params
+        return params
+
+    def _prepare_serving_volumes(self, model_location):
+        """Prepares the serving volumes.
+
+        Args:
+            model_location: Location of the models.
+        """
+        volumes = []
+        host = self.hosts[0]
+        # Make the model available to the container. If this is a local file just mount it to
+        # the container as a volume. If it is an S3 location, the DataSource will download it, we
+        # just need to extract the tar file.
+        host_dir = os.path.join(self.container_root, host)
+        os.makedirs(host_dir)
+
+        model_data_source = sagemaker.core.local.data.get_data_source_instance(
+            model_location, self.sagemaker_session
+        )
+
+        for filename in model_data_source.get_file_list():
+            if tarfile.is_tarfile(filename):
+                with tarfile.open(filename) as tar:
+                    custom_extractall_tarfile(tar, model_data_source.get_root_dir())
+
+        volumes.append(_Volume(model_data_source.get_root_dir(), "/opt/ml/model"))
+
+        return volumes
+
+    def _generate_compose_file(self, command, additional_volumes=None, additional_env_vars=None):
+        """Writes a config file describing a training/hosting environment.
+
+        This method generates a docker compose configuration file, it has an
+        entry for each container that will be created (based on self.hosts). it
+        calls
+        :meth:~sagemaker.local_session.SageMakerContainer._create_docker_host to
+        generate the config for each individual container.
+
+        Args:
+            command (str): either 'train' or 'serve'
+            additional_volumes (list): a list of volumes that will be mapped to
+                the containers
+            additional_env_vars (dict): a dictionary with additional environment
+                variables to be passed on to the containers.
+
+        Returns: (dict) A dictionary representation of the configuration that was written.
+        """
+        boto_session = self.sagemaker_session.boto_session
+        additional_volumes = additional_volumes or []
+        additional_env_vars = additional_env_vars or {}
+        environment = []
+        optml_dirs = set()
+
+        aws_creds = _aws_credentials(boto_session)
+        if aws_creds is not None:
+            environment.extend(aws_creds)
+
+        additional_env_var_list = ["{}={}".format(k, v) for k, v in additional_env_vars.items()]
+        environment.extend(additional_env_var_list)
+
+        if self.is_studio:
+            environment.extend([f"{SM_STUDIO_LOCAL_MODE}=True"])
+
+        if os.environ.get(DOCKER_COMPOSE_HTTP_TIMEOUT_ENV) is None:
+            os.environ[DOCKER_COMPOSE_HTTP_TIMEOUT_ENV] = DOCKER_COMPOSE_HTTP_TIMEOUT
+
+        if command == "train":
+            optml_dirs = {"output", "output/data", "input"}
+        elif command == "process":
+            optml_dirs = {"output", "config"}
+
+        services = {
+            h: self._create_docker_host(h, environment, optml_dirs, command, additional_volumes)
+            for h in self.hosts
+        }
+
+        if self.is_studio:
+            content = {
+                # Use version 2.3 as a minimum so that we can specify the runtime
+                "version": "2.3",
+                "services": services,
+            }
+        else:
+            content = {
+                # Use version 2.3 as a minimum so that we can specify the runtime
+                "version": "2.3",
+                "services": services,
+                "networks": {"sagemaker-local": {"name": "sagemaker-local"}},
+            }
+
+        docker_compose_path = os.path.join(self.container_root, DOCKER_COMPOSE_FILENAME)
+
+        try:
+            import yaml
+        except ImportError as e:
+            logger.error(
+                sagemaker.core.common_utils._module_import_error("yaml", "Local mode", "local")
+            )
+            raise e
+
+        yaml_content = yaml.dump(content, default_flow_style=False)
+        # Mask all environment vars for logging, could contain secrects.
+        masked_content = copy.deepcopy(content)
+        for _, service_data in masked_content["services"].items():
+            service_data["environment"] = ["[Masked]" for _ in service_data["environment"]]
+
+        masked_content_for_logging = yaml.dump(masked_content, default_flow_style=False)
+        logger.info("docker compose file: \n%s", masked_content_for_logging)
+        with open(docker_compose_path, "w") as f:
+            f.write(yaml_content)
+
+        return content
+
+    def _compose(self, detached=False):
+        """Invokes the docker compose command.
+
+        Args:
+            detached:
+        """
+        compose_cmd = self.compose_cmd_prefix
+
+        command = [
+            "-f",
+            os.path.join(self.container_root, DOCKER_COMPOSE_FILENAME),
+            "up",
+            "--build",
+            "--abort-on-container-exit" if not detached else "--detach",  # mutually exclusive
+        ]
+
+        compose_cmd.extend(command)
+
+        logger.info("docker command: %s", " ".join(compose_cmd))
+        return compose_cmd
+
+    def _create_docker_host(
+        self,
+        host: str,
+        environment: List[str],
+        optml_subdirs: set[str],
+        command: str,
+        volumes: List,
+    ) -> Dict:
+        """Creates the docker host configuration.
+
+        Args:
+            host (str): The host address
+            environment (List[str]): List of environment variables
+            optml_subdirs (Set[str]): Set of subdirs
+            command (str): Either 'train' or 'serve'
+            volumes (list): List of volumes that will be mapped to the containers
+        """
+        optml_volumes = self._build_optml_volumes(host, optml_subdirs)
+        optml_volumes.extend(volumes)
+
+        container_name_prefix = "".join(
+            random.choice(string.ascii_lowercase + string.digits) for _ in range(10)
+        )
+        container_default_config = (
+            sagemaker.core.config.config_manager.SageMakerConfig.get_config_value(
+                f"{LOCAL}.{CONTAINER_CONFIG}", self.sagemaker_session.config
+            )
+            or {}
+        )
+
+        host_config = {
+            **container_default_config,
+            "image": self.image,
+            "container_name": f"{container_name_prefix}-{host}",
+            "stdin_open": True,
+            "tty": True,
+            "volumes": [v.map for v in optml_volumes],
+            "environment": environment,
+        }
+
+        is_train_with_entrypoint = False
+        if command == "train" and self.container_entrypoint:
+            # Remote function or Pipeline function step is translated into a training job
+            # with container_entrypoint configured
+            is_train_with_entrypoint = True
+
+        if command != "process" and not is_train_with_entrypoint:
+            host_config["command"] = command
+        else:
+            if self.container_entrypoint:
+                host_config["entrypoint"] = self.container_entrypoint
+            if self.container_arguments:
+                host_config["entrypoint"] = host_config["entrypoint"] + self.container_arguments
+
+        if self.is_studio:
+            host_config["network_mode"] = "sagemaker"
+        else:
+            host_config["networks"] = {"sagemaker-local": {"aliases": [host]}}
+
+        # for GPU support pass in nvidia as the runtime, this is equivalent
+        # to setting --runtime=nvidia in the docker commandline.
+        if self.instance_type == "local_gpu":
+            host_config["deploy"] = {
+                "resources": {
+                    "reservations": {"devices": [{"count": "all", "capabilities": ["gpu"]}]}
+                }
+            }
+
+        if not self.is_studio and command == "serve":
+            serving_port = (
+                sagemaker.core.config.config_manager.SageMakerConfig.get_config_value(
+                    "local.serving_port", self.sagemaker_session.config
+                )
+                or 8080
+            )
+            host_config.update({"ports": ["%s:8080" % serving_port]})
+
+        return host_config
+
+    def _create_tmp_folder(self):
+        """Placeholder docstring"""
+        root_dir = sagemaker.core.config.config_manager.SageMakerConfig.get_config_value(
+            "local.container_root", self.sagemaker_session.config
+        )
+        if root_dir:
+            root_dir = os.path.abspath(root_dir)
+
+        working_dir = tempfile.mkdtemp(dir=root_dir)
+
+        # Docker cannot mount Mac OS /var folder properly see
+        # https://forums.docker.com/t/var-folders-isnt-mounted-properly/9600
+        # Only apply this workaround if the user didn't provide an alternate storage root dir.
+        if root_dir is None and platform.system() == "Darwin":
+            working_dir = "/private{}".format(working_dir)
+
+        return os.path.abspath(working_dir)
+
+    def _build_optml_volumes(self, host, subdirs):
+        """Generate a list of :class:`~sagemaker.local_session.Volume`.
+
+        These are required for the container to start. It takes a folder with
+        the necessary files for training and creates a list of opt volumes
+        that the Container needs to start.
+
+        Args:
+            host (str): container for which the volumes will be generated.
+            subdirs (list): list of subdirectories that will be mapped. For
+                example: ['input', 'output', 'model']
+
+        Returns: (list) List of :class:`~sagemaker.local_session.Volume`
+        """
+        volumes = []
+
+        for subdir in subdirs:
+            host_dir = os.path.join(self.container_root, host, subdir)
+            container_dir = "/opt/ml/{}".format(subdir)
+            volume = _Volume(host_dir, container_dir)
+            volumes.append(volume)
+
+        return volumes
+
+    def _cleanup(self, dirs_to_delete=None):
+        """Cleans up directories and the like.
+
+        Args:
+            dirs_to_delete:
+        """
+        if dirs_to_delete:
+            for d in dirs_to_delete:
+                _delete_tree(d)
+
+        # Free the container config files.
+        for host in self.hosts:
+            container_config_path = os.path.join(self.container_root, host)
+            _delete_tree(container_config_path)
+
+
+class _HostingContainer(Thread):
+    """Placeholder docstring."""
+
+    def __init__(self, command):
+        """Creates a new threaded hosting container.
+
+        Args:
+            command (dict): docker compose command
+        """
+        Thread.__init__(self)
+        self.command = command
+        self.process = None
+
+    def run(self):
+        """Placeholder docstring"""
+        self.process = subprocess.Popen(
+            self.command, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        )
+        try:
+            _stream_output(self.process)
+        except RuntimeError as e:
+            # _stream_output() doesn't have the command line. We will handle the exception
+            # which contains the exit code and append the command line to it.
+            msg = "Failed to run: %s, %s" % (self.command, str(e))
+            raise RuntimeError(msg)
+
+    def down(self):
+        """Placeholder docstring"""
+        if os.name != "nt":
+            sagemaker.core.local.utils.kill_child_processes(self.process.pid)
+        self.process.terminate()
+
+
+class _Volume(object):
+    """Represent a Volume that will be mapped to a container."""
+
+    def __init__(self, host_dir, container_dir=None, channel=None):
+        """Create a Volume instance.
+
+        The container path can be provided as a container_dir or as a channel name but not both.
+
+        Args:
+            host_dir (str): path to the volume data in the host
+            container_dir (str): path inside the container that host_dir will be mapped to
+            channel (str): channel name that the host_dir represents. It will be mapped as
+                /opt/ml/input/data/<channel> in the container.
+        """
+        if not container_dir and not channel:
+            raise ValueError("Either container_dir or channel must be declared.")
+
+        if container_dir and channel:
+            raise ValueError("container_dir and channel cannot be declared together.")
+
+        self.container_dir = container_dir if container_dir else "/opt/ml/input/data/" + channel
+        self.host_dir = host_dir
+        map_format = "{}:{}"
+        if platform.system() == "Linux" and SELINUX_ENABLED:
+            # Support mounting shared volumes in SELinux enabled hosts
+            map_format += ":z"
+        if platform.system() == "Darwin" and host_dir.startswith("/var"):
+            self.host_dir = os.path.join("/private", host_dir)
+
+        self.map = map_format.format(self.host_dir, self.container_dir)
+
+
+def _stream_output(process):
+    """Stream the output of a process to stdout
+
+    This function takes an existing process that will be polled for output.
+    Only stdout will be polled and sent to sys.stdout.
+
+    Args:
+        process (subprocess.Popen): a process that has been started with
+            stdout=PIPE and stderr=STDOUT
+
+    Returns (int): process exit code
+    """
+    exit_code = None
+
+    while exit_code is None:
+        stdout = process.stdout.readline().decode("utf-8")
+        sys.stdout.write(stdout)
+        exit_code = process.poll()
+
+    if exit_code not in [0, 130]:
+        raise RuntimeError(f"Failed to run: {process.args}. Process exited with code: {exit_code}")
+
+    return exit_code
+
+
+def _check_output(cmd, *popenargs, **kwargs):
+    """Makes a call to `subprocess.check_output` for the given command and args.
+
+    Args:
+        cmd:
+        *popenargs:
+        **kwargs:
+    """
+    if isinstance(cmd, str):
+        cmd = shlex.split(cmd)
+
+    success = True
+    try:
+        output = subprocess.check_output(cmd, *popenargs, **kwargs)
+    except subprocess.CalledProcessError as e:
+        output = e.output
+        success = False
+
+    output = output.decode("utf-8")
+    if not success:
+        logger.error("Command output: %s", output)
+        raise Exception("Failed to run %s" % ",".join(cmd))
+
+    return output
+
+
+def _create_processing_config_file_directories(root, host):
+    """Creates the directory for the processing config files.
+
+    Args:
+        root: The root path.
+        host: The current host.
+    """
+    for d in ["config"]:
+        os.makedirs(os.path.join(root, host, d))
+
+
+def _create_config_file_directories(root, host):
+    """Creates the directories for the config files.
+
+    Args:
+        root:
+        host:
+    """
+    for d in ["input", "input/config", "output", "model"]:
+        os.makedirs(os.path.join(root, host, d))
+
+
+def _delete_tree(path):
+    """Makes a call to `shutil.rmtree` for the given path.
+
+    Args:
+        path:
+    """
+    try:
+        shutil.rmtree(path)
+    except OSError as exc:
+        # on Linux, when docker writes to any mounted volume, it uses the container's user. In most
+        # cases this is root. When the container exits and we try to delete them we can't because
+        # root owns those files. We expect this to happen, so we handle EACCESS. Any other error
+        # we will raise the exception up.
+        if exc.errno == errno.EACCES:
+            logger.warning("Failed to delete: %s Please remove it manually.", path)
+        else:
+            logger.error("Failed to delete: %s", path)
+            raise
+
+
+def _aws_credentials(session):
+    """Provides the AWS credentials of the session as a paired list of strings.
+
+    These can be used to set environment variables on command execution.
+
+    Args:
+        session:
+    """
+    try:
+        creds = session.get_credentials()
+        access_key = creds.access_key
+        secret_key = creds.secret_key
+        token = creds.token
+
+        # The presence of a token indicates the credentials are short-lived and as such are risky
+        # to be used as they might expire while running.
+        # Long-lived credentials are available either through
+        # 1. boto session
+        # 2. EC2 Metadata Service (SageMaker Notebook instances or EC2 instances with roles
+        #       attached them)
+        # Short-lived credentials available via boto session are permitted to support running on
+        # machines with no EC2 Metadata Service but a warning is provided about their danger
+        if token is None:
+            logger.info("Using the long-lived AWS credentials found in session")
+            return [
+                "AWS_ACCESS_KEY_ID=%s" % (str(access_key)),
+                "AWS_SECRET_ACCESS_KEY=%s" % (str(secret_key)),
+            ]
+        if _use_short_lived_credentials() or not _aws_credentials_available_in_metadata_service():
+            logger.warning(
+                "Using the short-lived AWS credentials found in session. They might expire while "
+                "running."
+            )
+            return [
+                "AWS_ACCESS_KEY_ID=%s" % (str(access_key)),
+                "AWS_SECRET_ACCESS_KEY=%s" % (str(secret_key)),
+                "AWS_SESSION_TOKEN=%s" % (str(token)),
+            ]
+        logger.info(
+            "No AWS credentials found in session but credentials from EC2 Metadata Service are "
+            "available."
+        )
+        return None
+    except Exception as e:  # pylint: disable=broad-except
+        logger.info("Could not get AWS credentials: %s", e)
+
+    return None
+
+
+def _aws_credentials_available_in_metadata_service():
+    """Placeholder docstring"""
+    import botocore
+    from botocore.credentials import InstanceMetadataProvider
+    from botocore.utils import InstanceMetadataFetcher
+
+    session = botocore.session.Session()
+    instance_metadata_provider = InstanceMetadataProvider(
+        iam_role_fetcher=InstanceMetadataFetcher(
+            timeout=session.get_config_variable("metadata_service_timeout"),
+            num_attempts=session.get_config_variable("metadata_service_num_attempts"),
+            user_agent=session.user_agent(),
+        )
+    )
+    return not instance_metadata_provider.load() is None
+
+
+def _use_short_lived_credentials():
+    """Use short-lived AWS credentials found in session."""
+    return os.environ.get("USE_SHORT_LIVED_CREDENTIALS") == "1"
+
+
+def _write_json_file(filename, content):
+    """Write the contents dict as json to the file.
+
+    Args:
+        filename:
+        content:
+    """
+    with open(filename, "w") as f:
+        json.dump(content, f)
+
+
+def _ecr_login_if_needed(boto_session, image):
+    """Log into ECR, if needed.
+
+    Of note, only ECR images need login.
+
+    Args:
+        boto_session:
+        image:
+    """
+    sagemaker_pattern = re.compile(sagemaker.core.common_utils.ECR_URI_PATTERN)
+    sagemaker_match = sagemaker_pattern.match(image)
+    if not sagemaker_match:
+        return False
+
+    # do we have the image?
+    if _check_output("docker images -q %s" % image).strip():
+        return False
+
+    if not boto_session:
+        raise RuntimeError(
+            "A boto session is required to login to ECR."
+            "Please pull the image: %s manually." % image
+        )
+
+    ecr = boto_session.client("ecr")
+    auth = ecr.get_authorization_token(registryIds=[image.split(".")[0]])
+    authorization_data = auth["authorizationData"][0]
+
+    raw_token = base64.b64decode(authorization_data["authorizationToken"])
+    token = raw_token.decode("utf-8").strip("AWS:")
+    ecr_url = auth["authorizationData"][0]["proxyEndpoint"]
+
+    # Log in to ecr, but use communicate to not print creds to the console
+    cmd = f"docker login {ecr_url} -u AWS --password-stdin".split()
+    proc = subprocess.Popen(
+        cmd,
+        stdin=subprocess.PIPE,
+    )
+
+    proc.communicate(input=token.encode())
+
+    return True
+
+
+def _pull_image(image):
+    """Invokes the docker pull command for the given image.
+
+    Args:
+        image:
+    """
+    pull_image_command = ("docker pull %s" % image).strip()
+    logger.info("docker command: %s", pull_image_command)
+
+    subprocess.check_output(pull_image_command.split())
+    logger.info("image pulled: %s", image)