AWS sagemaker documentation change
Summary
Comprehensive rewrite of MLflow model registration documentation. Added sections on IAM permissions, inference specifications, evaluation metrics logging, lifecycle management using aliases, and resource tag-based access control. Enhanced UI integration details and Studio visualization.
Security assessment
The changes add significant documentation about security features including required IAM permissions for model registration, condition keys for lifecycle transitions, and resource tag-based policies to prevent updates. However, there's no evidence of addressing a specific vulnerability or incident - these are proactive security controls.
Diff
diff --git a/sagemaker/latest/dg/mlflow-track-experiments-model-registration.md b/sagemaker/latest/dg/mlflow-track-experiments-model-registration.md index ddee3817d..757c27db1 100644 --- a//sagemaker/latest/dg/mlflow-track-experiments-model-registration.md +++ b//sagemaker/latest/dg/mlflow-track-experiments-model-registration.md @@ -7 +7 @@ -Register models using the SageMaker Python SDKRegister models using the MLflow UIView registered models in Studio +Log and register modelsManage model lifecycle stagesPrevent updates to registered models using resource tags @@ -11 +11,14 @@ Register models using the SageMaker Python SDKRegister models using the MLflow U -You can log MLflow models and automatically register them with SageMaker Model Registry using either the Python SDK or directly through the MLflow UI. +When you register an MLflow model, SageMaker AI automatically creates a corresponding Model Package Group and Model Package version in the SageMaker Model Registry. This enables you to: + + * Register models from the MLflow SDK or the MLflow UI. + + * Attach an inference specification to enable direct deployment to SageMaker AI endpoints. + + * Log evaluation metrics that are included as a model card in the Model Package. + + * Manage model lifecycle stages (staging, production) and statuses (pending, active) through MLflow aliases. + + * Control lifecycle transitions and registration with IAM condition keys and resource tag-based policies. + + + @@ -15 +28 @@ You can log MLflow models and automatically register them with SageMaker Model R -Do not use spaces in a model name. While MLflow supports model names with spaces, SageMaker AI Model Package does not. The auto-registration process fails if you use spaces in your model name. +Model Registry sync is an opt-in feature. To enable it, call the [UpdateMlflowTrackingServer](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_UpdateMlflowTrackingServer.html) or [UpdateMlflowApp](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_UpdateMlflowApp.html) API with Model Registry sync enabled. @@ -17 +30 @@ Do not use spaces in a model name. While MLflow supports model names with spaces -## Register models using the SageMaker Python SDK +## Log and register models @@ -19 +32 @@ Do not use spaces in a model name. While MLflow supports model names with spaces -Use `create_registered_model` within your MLflow client to automatically create a model package group in SageMaker AI that corresponds to an existing MLflow model of your choice. +This section covers how to prepare a model for registration, register it with the SageMaker Model Registry, and view it in Studio. @@ -20,0 +34 @@ Use `create_registered_model` within your MLflow client to automatically create +### Required IAM permissions @@ -22,2 +36 @@ Use `create_registered_model` within your MLflow client to automatically create - import mlflow - from mlflow import MlflowClient +The MLflow tracking server or MLflow App IAM service role must have the following permissions to register models: @@ -25 +37,0 @@ Use `create_registered_model` within your MLflow client to automatically create - mlflow.set_tracking_uri(arn) @@ -27 +39,27 @@ Use `create_registered_model` within your MLflow client to automatically create - client = MlflowClient() + { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "sagemaker:CreateModelPackageGroup", + "sagemaker:DescribeModelPackageGroup", + "sagemaker:CreateModelPackage", + "sagemaker:UpdateModelPackage", + "sagemaker:AddTags", + "sagemaker:CreateAction", + "sagemaker:AddAssociation" + ], + "Resource": "*" + }, + { + "Effect": "Allow", + "Action": [ + "s3:GetObject", + "s3:PutObject", + "s3:ListBucket" + ], + "Resource": "*" + } + ] + } @@ -29,2 +67 @@ Use `create_registered_model` within your MLflow client to automatically create - mlflow_model_name = 'AutoRegisteredModel' - client.create_registered_model(mlflow_model_name, tags={"key1": "value1"}) +For more information about setting up IAM roles for MLflow Apps, see [Set up IAM permissions for MLflow Apps](./mlflow-app-setup-prerequisites-iam.html). @@ -32 +69 @@ Use `create_registered_model` within your MLflow client to automatically create -Use `mlflow.register_model()` to automatically register a model with the SageMaker Model Registry during model training. When registering the MLflow model, a corresponding model package group and model package version are created in SageMaker AI. +### Prepare a model for registration @@ -33,0 +71 @@ Use `mlflow.register_model()` to automatically register a model with the SageMak +Before registering a model, you log it during a training run. You can optionally attach an inference specification and evaluation metrics, which are carried forward to the SageMaker AI Model Package upon registration. @@ -35,2 +73,6 @@ Use `mlflow.register_model()` to automatically register a model with the SageMak - import mlflow.sklearn - from mlflow.models import infer_signature +#### Log a model + +Use the MLflow SDK to log a model during a training run: + + + import mlflow @@ -38,0 +81 @@ Use `mlflow.register_model()` to automatically register a model with the SageMak + from mlflow.models import infer_signature @@ -40,0 +84,2 @@ Use `mlflow.register_model()` to automatically register a model with the SageMak + mlflow.set_experiment("my-experiment") + @@ -44 +88,0 @@ Use `mlflow.register_model()` to automatically register a model with the SageMak - # Log MLflow entities @@ -46,2 +90,2 @@ Use `mlflow.register_model()` to automatically register a model with the SageMak - rfr = RandomForestRegressor(**params).fit(X, y) - signature = infer_signature(X, rfr.predict(X)) + model = RandomForestRegressor(**params).fit(X, y) + signature = infer_signature(X, model.predict(X)) @@ -49 +93,67 @@ Use `mlflow.register_model()` to automatically register a model with the SageMak - mlflow.sklearn.log_model(rfr, artifact_path="sklearn-model", signature=signature) + model_info = mlflow.sklearn.log_model( + model, name="sklearn-model", signature=signature, input_example=X[:3] + ) + +#### Log an inference specification (optional) + +An inference specification defines the container image and instance types required to deploy your model. When attached to a logged model before registration, the inference specification is automatically included in the SageMaker AI Model Package. This enables direct deployment to a SageMaker AI endpoint from the Model Registry without additional configuration. + +Use `sagemaker_mlflow.log_inference_specification()` to attach an inference specification: + + + import sagemaker_mlflow + + # Get the logged model's artifact location + logged_model = mlflow.MlflowClient().get_logged_model(model_info.model_id) + + inference_spec = { + "Containers": [{ + "Image": "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:2.0.0-cpu-py310", + "ModelDataSource": { + "S3DataSource": { + "S3Uri": logged_model.artifact_location + "/", + "S3DataType": "S3Prefix", + "CompressionType": "None", + } + }, + }], + "SupportedRealtimeInferenceInstanceTypes": ["ml.m5.xlarge"], + } + sagemaker_mlflow.log_inference_specification( + model_info.model_id, + inference_specification=inference_spec + ) + +###### Note + +The inference specification must follow the same schema as the `InferenceSpecification` parameter of the SageMaker AI `CreateModelPackage` API. For more information, see [InferenceSpecification](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_InferenceSpecification.html) in the SageMaker AI API Reference. + +After logging, the inference specification is stored as a `sagemaker_inference_specification.json` artifact alongside your model: + + + +#### Log evaluation metrics (optional) + +Evaluation metrics provide a standardized summary of your model's performance. When logged before registration, evaluation results are included in the SageMaker AI Model Package as a model card, making metrics visible in the SageMaker AI Model Registry and Studio. For more information about model cards, see [Model card schema](https://docs.aws.amazon.com/sagemaker/latest/dg/model-registry-details.html#model-card-schema). + +Use `sagemaker_mlflow.evaluate()` to log evaluation metrics: + + + import sagemaker_mlflow + import pandas as pd + + eval_df = pd.DataFrame(X, columns=["f1", "f2", "f3"]) + eval_df["target"] = y + dataset = mlflow.data.from_pandas( + eval_df, source="s3://my-bucket/eval.csv", name="eval_set", targets="target" + ) + sagemaker_mlflow.evaluate(model_info, data=dataset, model_type="regressor") + +After logging, the evaluation group is stored as an artifact alongside your model: + + + +### Register models using the SDK + +Use `mlflow.register_model()` to register a model with the SageMaker Model Registry. This automatically creates a Model Package Group (if one does not exist) and a Model Package version in SageMaker AI. + @@ -52 +162 @@ Use `mlflow.register_model()` to automatically register a model with the SageMak - mv = mlflow.register_model(model_uri, "RandomForestRegressionModel") + mv = mlflow.register_model(model_uri, "MyRegisteredModel") @@ -57 +167 @@ Use `mlflow.register_model()` to automatically register a model with the SageMak -## Register models using the MLflow UI +You can also use `create_registered_model` to create the Model Package Group before registering versions: @@ -59 +169,20 @@ Use `mlflow.register_model()` to automatically register a model with the SageMak -You can alternatively register a model with the SageMaker Model Registry directly in the MLflow UI. Within the **Models** menu in the MLflow UI, choose **Create Model**. Any models newly created in this way are added to the SageMaker Model Registry. + + from mlflow import MlflowClient + + client = MlflowClient() + client.create_registered_model("MyRegisteredModel", tags={"team": "ml-platform"}) + +After registration, the model version is tagged with the SageMaker AI Model Package ARN. You can retrieve it as follows: + + + mv = client.get_model_version("MyRegisteredModel", mv.version) + sm_arn = mv.tags["sagemaker.model_package_arn"] + print(f"SageMaker Model Package ARN: {sm_arn}") + +###### Note + +Do not use spaces in a model name. While MLflow supports model names with spaces, SageMaker AI Model Package does not. The registration process fails if you use spaces in your model name. + +### Register models using the MLflow UI + +You can register a model with the SageMaker Model Registry directly in the MLflow UI. Within the **Models** menu, choose **Create Model**. Models created this way are automatically added to the SageMaker Model Registry. @@ -63 +192,25 @@ You can alternatively register a model with the SageMaker Model Registry directl -After logging a model during experiment tracking, navigate to the run page in the MLflow UI. Choose the **Artifacts** pane and choose **Register model** in the upper right corner to register the model version in both MLflow and SageMaker Model Registry. +After logging a model during experiment tracking, navigate to the run page in the MLflow UI. Choose the **Artifacts** pane and choose **Register model** to register the model version in both MLflow and SageMaker Model Registry. +