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

Author: benieric

.gitignore

@@ -32,6 +32,7 @@ env/
 .python-version
 *.html
 **/_repack_script_launcher.sh
+src/sagemaker/modules/scripts/train.sh
 tests/data/**/_repack_model.py
 tests/data/experiment/sagemaker-dev-1.0.tar.gz
 src/sagemaker/serve/tmp_workspace

src/sagemaker/modules/configs.py

@@ -10,7 +10,15 @@
 # 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.
-"""Configuration classes."""
+"""This module provides the configuration classes used in `sagemaker.modules`.
+
+Some of these classes are re-exported from `sagemaker-core.shapes`. For convinence,
+users can import these classes directly from `sagemaker.modules.configs`.
+
+For more documentation on `sagemaker-core.shapes`, see:
+    - https://sagemaker-core.readthedocs.io/en/stable/#sagemaker-core-shapes
+"""
+
 from __future__ import absolute_import
 
 from typing import Optional
@@ -20,8 +28,8 @@
     ResourceConfig,
     StoppingCondition,
     OutputDataConfig,
-    AlgorithmSpecification,
     Channel,
+    DataSource,
     S3DataSource,
     FileSystemDataSource,
     TrainingImageConfig,
@@ -33,8 +41,8 @@
     "ResourceConfig",
     "StoppingCondition",
     "OutputDataConfig",
-    "AlgorithmSpecification",
     "Channel",
+    "DataSource",
     "S3DataSource",
     "FileSystemDataSource",
     "TrainingImageConfig",
@@ -62,7 +70,7 @@ class SourceCodeConfig(BaseModel):
             job container. If not specified, command must be provided.
     """
 
-    command: Optional[str]
-    source_dir: Optional[str]
-    requirements: Optional[str]
-    entry_script: Optional[str]
+    command: Optional[str] = None
+    source_dir: Optional[str] = None
+    requirements: Optional[str] = None
+    entry_script: Optional[str] = None

src/sagemaker/modules/constants.py

@@ -12,5 +12,19 @@
 # language governing permissions and limitations under the License.
 """Constants module."""
 from __future__ import absolute_import
+import os
 
 DEFAULT_INSTANCE_TYPE = "ml.m5.xlarge"
+
+SOURCE_CODE_CONTAINER_PATH = "/opt/ml/input/data/code"
+
+SM_CODE_CONTAINER_PATH = "/opt/ml/input/data/sm_code"
+SM_CODE_LOCAL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "scripts")
+TRAIN_SCRIPT = "train.sh"
+
+DEFAULT_CONTAINER_ENTRYPOINT = ["/bin/bash"]
+DEFAULT_CONTAINER_ARGUMENTS = [
+    "-c",
+    f"chmod +x {SM_CODE_CONTAINER_PATH}/{TRAIN_SCRIPT} "
+    + f"&& {SM_CODE_CONTAINER_PATH}/{TRAIN_SCRIPT}",
+]

src/sagemaker/modules/image_spec.py

@@ -69,8 +69,13 @@ def __init__(
         self.config_name = config_name
         self.sagemaker_session = sagemaker_session
 
-    def get_image_uri(self):
+    def get_image_uri(
+        self, image_scope: Optional[str] = None, instance_type: Optional[str] = None
+    ) -> str:
         """Get image URI for a specific framework version."""
+
+        self.image_scope = image_scope or self.image_scope
+        self.instance_type = instance_type or self.instance_type
         return image_uris.retrieve(
             framework=self.framework_name,
             image_scope=self.image_scope,

src/sagemaker/modules/scripts/train.sh

@@ -0,0 +1,51 @@
+# 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.
+"""Templates module."""
+from __future__ import absolute_import
+
+TRAIN_SCRIPT_TEMPLATE = """
+#!/bin/bash
+echo "Starting training script"
+
+echo "/opt/ml/input/config/resourceconfig.json:"
+cat /opt/ml/input/config/resourceconfig.json
+echo
+
+echo "/opt/ml/input/config/inputdataconfig.json:"
+cat /opt/ml/input/config/inputdataconfig.json
+echo
+
+echo "/opt/ml/input/config/hyperparameters.json:"
+cat /opt/ml/input/config/hyperparameters.json
+echo
+
+python --version
+{working_dir}
+{install_requirements}
+CMD="{command}"
+echo "Running command: $CMD"
+eval $CMD
+EXIT_STATUS=$?
+
+if [ $EXIT_STATUS -ne 0 ]; then
+    echo "Command failed with exit status $EXIT_STATUS"
+    if [ ! -s /opt/ml/output/failure ]; then
+        echo "Command failed with exit code $EXIT_STATUS.
+For more details, see CloudWatch logs at 'aws/sagemaker/TrainingJobs'.
+TrainingJob - $TRAINING_JOB_NAME" >> /opt/ml/output/failure
+    fi
+    exit $EXIT_STATUS
+else
+    echo "Command succeeded"
+fi
+"""

src/sagemaker/modules/templates.py

@@ -19,8 +19,32 @@
    "outputs": [],
    "source": [
     "from sagemaker.modules.train.model_trainer import ModelTrainer\n",
+    "from sagemaker.modules.image_spec import ImageSpec\n",
     "\n",
-    "model_trainer = ModelTrainer(training_image=\"python:3.10.15-slim\")\n"
+    "pytorch_image = ImageSpec(\n",
+    "    framework_name=\"pytorch\",\n",
+    "    version=\"1.13.1\",\n",
+    "    py_version=\"py39\"\n",
+    ")\n",
+    "\n",
+    "python_ecr_image = \"public.ecr.aws/docker/library/python:3.10.15-slim\"\n",
+    "python_docker_image = \"python:3.10.15-slim\"\n",
+    "\n",
+    "model_trainer = ModelTrainer(training_image=pytorch_image)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sagemaker.modules.configs import SourceCodeConfig\n",
+    "\n",
+    "source_code_config = SourceCodeConfig(\n",
+    "    command=\"echo 'Hello World' && env\",\n",
+    ")\n",
+    "model_trainer.train(source_code_config=source_code_config)"
    ]
   },
   {