翻訳待ち：Bring your own model with Amazon SageMaker AI: Script mode in SDK v3

AI サービスが一時的に利用できないため、復旧後に翻訳を補完します。

In 2021, we published Bring your own model with Amazon SageMaker script mode. That post showed how to use script mode on managed framework containers from AWS to write custom training and inference code. Script mode was a leap forward: you didn’t need to build or maintain Docker images to run your own algorithm on Amazon SageMaker AI. The v3 SDK delivers a redesign from scratch that makes many workflows like the bring-your-own-model workflow even more streamlined. The new SDK replaces framework-specific estimator classes (SKLearn, PyTorch, XGBoost) with a single, unified ModelTrainer for training and ModelBuilder for deployment. In v3, the SDK syncs a local source code directory into the training job at runtime using the new SourceCode configuration object. You bring a container image from Amazon Elastic Container Registry (Amazon ECR): one you build, an AWS Deep Learning Container, or a third-party image. The SDK handles injecting your code at runtime. This means: Faster iterations: Change your training script, rerun. No container rebuild necessary. Full container control: Install system packages or CUDA libraries in your image. The SDK doesn’t assume what’s inside. One API for multiple frameworks: Whether you’re training with frameworks like scikit-learn, PyTorch, Stable Diffusion, or a custom C++ inference binary, the interface is identical. Solution overview In this post, we walk through two end-to-end examples that demonstrate how script mode works in the SageMaker Python SDK v3: Train and deploy a scikit-learn Random Forest – a classic tabular machine learning (ML) workflow that trains on the diabetes dataset and deploys to a real-time endpoint using Deep Java Library (DJL) Serving, a high-performance model server. Fine-tune Stable Diffusion 3.5 with LoRA – a generative AI workflow that uses Hugging Face Accelerate for multi-GPU distributed training. Both examples use the same two core classes: ModelTrainer replaces the v2 Estimator family. Configures and launches a SageMaker training job. ModelBuilder replaces the v2 Model/Predictor pattern. Packages your inference handler and deploys to an endpoint. A key concept is the SourceCode object. It accepts a source_dir (a path to your local code directory) and either a command string (for training) or an entry_script (for inference). At job launch, SageMaker syncs this directory into the container, and your code runs inside the container without being baked into the image. You can find the example code for this blog post in the GitHub repository. What changed from SDK v2 to v3? The following table summarizes the architectural shift: SDK v2 (Estimator pattern) SDK v3 (ModelTrainer pattern) Training class SKLearn, PyTorch, XGBoost, … ModelTrainer (one single class) Deployment class Model + Predictor ModelBuilder to deploy the endpoint, prediction handled as part of invoke() Container AWS managed framework image Any image: yours, AWS DLC, or third-party Code injection entry_point + source_dir, framework-specific SourceCode object with source_dir + command/entry_script Dependencies requirements.txt in source_dir requirements.txt in source_dir Prerequisites To follow along, you need: An AWS account with Amazon SageMaker AI access. An AWS Identity and Access Management (IAM) execution role with Amazon SageMaker AI and Amazon Simple Storage Service (Amazon S3) permissions. The SageMaker Python SDK v3 installed (pip install sagemaker>=3.0). A training container image pushed to Amazon ECR (this post shows an example of building and pushing a container to Amazon ECR). An Amazon S3 bucket for training data and model artifacts. (Optional) An MLflow app or tracking server on Amazon SageMaker AI for experiment tracking. (Optional) If you plan to build and run the example containers from a JupyterLab space in Amazon SageMaker Studio rather than a local machine, Docker access must be enabled at the domain level. For details, see Local mode support in Amazon SageMaker Studio. Example 1: Train and deploy a scikit-learn model Let’s start with a classic ML workflow. We train a Random Forest classifier on the diabetes dataset and deploy it to a real-time SageMaker endpoint. Step 1: Building the Docker container The training container is intentionally minimal. It contains only the runtime and framework libraries and no training code, so that we can reuse it for other scikit-learn models we might want to build. The complete Dockerfile for our scikit-learn container is: FROM python:3.13-slim RUN apt-get update && apt-get install -y \ build-essential jq git \ && rm -rf /var/lib/apt/lists/* COPY requirements.txt . RUN pip install -r requirements.txt --no-cache-dir The container is a stable, version-controlled runtime environment. The algorithm-specific code lives in your source_dir, and the SDK injects it at runtime. Build this container once, push it to Amazon ECR, and iterate on your training code as many times as you want without touching Docker again. The example notebook includes Docker build and push commands by using two shell scripts: ./build.sh --env .env.docker.sklearn ./push.sh --env .env.docker.sklearn Note that you need Docker installed on the environment you’re using to run the code samples. If you’re running this on a JupyterLab space within Amazon SageMaker AI, you need to enable Docker on the domain-level settings. Step 1a: Configuration First, we auto-detect the account-level configuration. Note that we omit the import statements required in the following code snippet for brevity, but the full code is available in the GitHub repository. ... boto_session = boto3.Session() sm_session = Session(boto_session=boto_session) AWS_REGION = boto_session.region_name AWS_ACCOUNT_ID = boto3.client("sts").get_caller_identity()["Account"] SAGEMAKER_EXECUTION_ROLE = sm_session.get_caller_identity_arn() # works for role or user S3_BUCKET = sm_session.default_bucket() In the following snippet, we point TRAINING_IMAGE_URI at the container we built ourselves in the previous step. This gives you full control over installed packages and runtime versions. For deployment, we show that you can also use a pre-existing managed DJL framework container if you don’t want to build your own. For more information about pre-built containers, see available Deep Learning Containers images. # Your custom training image pushed to ECR TRAINING_IMAGE_URI = ( f"{AWS_ACCOUNT_ID}.dkr.ecr.{AWS_REGION}.amazonaws.com/sklearn:latest" ) MODEL_OUTPUT_S3_PATH = f"s3://{S3_BUCKET}/random-forest/model-output" Optionally, if you’d like to track hyperparameters, metrics, and model artifacts across training runs, the example training script is already instrumented for fully managed MLflow on Amazon SageMaker AI. Set the MLFLOW_ARN and MLFLOW_EXPERIMENT_NAME variables in the following snippet to automatically enable logging. This is an optional step, and you can set the values to None to skip experiment tracking instead. MLFLOW_ARN = "arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_ID}:mlflow-app/{XYZ}" MLFLOW_EXPERIMENT_NAME = "random-forest-experiment" Step 2: Launch a training job with ModelTrainer The SourceCode object takes your local source_dir and a command string. At job launch, SageMaker syncs the entire source_dir into the container and runs your command. This decouples your code from your container image. Change the script, re-launch with no container rebuild needed. source_code = SourceCode( source_dir="./train/random_forest", command=( "python random_forest.py" " --n_jobs 4 --max_depth 10 --n_estimators 120" f" --mlflow_arn {MLFLOW_ARN}" f" --mlflow_experiment_name {MLFLOW_EXPERIMENT_NAME}" ), ) compute = Compute( instance_type="ml.m5.2xlarge", instance_count=1, volume_size_in_gb=30, keep_alive_period_in_seconds=3600, # Warm pool for faster re-runs ) stopping_condition = StoppingCondition(max_runtime_in_seconds=3600) output_config = OutputDataConfig(s3_output_path=MODEL_OUTPUT_S3_PATH) model_trainer = ModelTrainer( training_image=TRAINING_IMAGE_URI, source_code=source_code, compute=compute, output_data_config=output_config, stopping_condition=stopping_condition, role=SAGEMAKER_EXECUTION_ROLE, base_job_name="random-forest-training", sagemaker_session=Session(), ) model_trainer.train(wait=True) A few things to note: source_dir can contain your files such as utility modules, config files, and shell scripts, which are synced into the container. command is a shell command that runs inside the container. You can call a Python script, a bash script, or anything else your container supports. keep_alive_period_in_seconds turns on SageMaker warm pools. The instance stays warm for 1 hour, so iterative re-runs launch in seconds rather than minutes. OutputDataConfig sets the S3 destination where SageMaker uploads your training results when the job finishes. Anything your script saves to /opt/ml/model (the SM_MODEL_DIR environment variable) is packaged as model.tar.gz under this path, and that’s the model artifact we deploy in Step 3. For more information, see Using the SageMaker training and inference toolkits for the folder structure and the environment variables that SageMaker sets. Step 3: Deploy to a real-time endpoint with ModelBuilder After training completes, we deploy the model artifact to a SageMaker real-time endpoint. ModelBuilder packages your inference handler, repacks it with the model artifact, and creates the endpoint in a few lines. You can use the metadata from the training job to find the S3 path of the final model artifact, then supply that to the ModelBuilder object. For serving, we use a pre-built AWS Deep Learning Container rather than building a custom one, though you can bring your own if needed. For other pre-built containers, see available Deep Learning Containers images. # Locate the trained model artifact sm_session = Session() sm_client = sm_session.boto_session.client("sagemaker") training_job_desc = sm_client.describe_training_job( TrainingJobName=training_job_name ) model_artifact_s3_uri = training_job_desc["ModelArtifacts"]["S3ModelArtifacts"] # Define inference source code inference_source_code = SourceCode( source_dir="./deploy/random_forest", entry_script="inference.py", ) # Build and deploy model_builder = ModelBuilder( image_uri=INFERENCE_IMAGE_URI, model_server=ModelServer.DJL_SERVING, source_code=inference_source_code, s3_model_data_url=model_artifact_s3_uri, env_vars={"OPTION_ENTRYPOINT": "code/inference.py"}, sagemaker_session=Session(), role_arn=SAGEMAKER_EXECUTION_ROLE, ) model_builder.build(model_name="random-forest-endpoint", mode=Mode.SAGEMAKER_ENDPOINT) The build() step assembles a deployable model without launching any infrastructure. ModelBuilder takes your inference handler and model artifact and packages them together according to the conventions of your chosen model server (here, DJL Serving). It then registers a SageMaker model that points at your inference image and repacked artifact in Amazon S3. ModelBuilder can also do more than we illustrate here, such as auto-selecting a container, auto-capturing dependencies, and generating serialization code from a raw framework model. For more information, see Create a model in Amazon SageMaker AI with ModelBuilder. With the model built, we call deploy() to stand up the real-time endpoint, which returns an Endpoint interface: predictor = model_builder.deploy( endpoint_name="random-forest-endpoint", initial_instance_count=1, ) Notice the same SourceCode pattern for inference: point at a local directory containing your handler and specify the entry_script. The SDK repacks the handler into the model archive so DJL Serving can find it at runtime. A few notes on the preceding code snippets: The inference.py script implements a single handle(inputs) function per the DJL Python mode documentation, which SageMaker calls for every request. When the inference worker first starts up, an empt [truncated for AI cost control]