Skip to content

feature: MLFlow E2E Example Notebook (5513) #5699

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
aviruthen wants to merge 2 commits into from
Closed

feature: MLFlow E2E Example Notebook (5513) #5699

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e304702
feature: MLFlow E2E Example Notebook (5513)
aviruthen Mar 30, 2026
1aaab37
fix: address review comments (iteration #1)
aviruthen Mar 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 48 additions & 16 deletions v3-examples/ml-ops-examples/v3-mlflow-train-inference-e2e-example.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@
"outputs": [],
"source": [
"import uuid\n",
"import boto3\n",
Copy link
Copy Markdown
Collaborator

@sagemaker-bot sagemaker-bot Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The import boto3 was moved here from Step 7, which is good for consolidation. However, per V3 architecture tenets, subpackages should avoid calling boto3 directly — the preferred pattern is to let Session() handle the boto session internally. Consider simplifying to:

from sagemaker.core.helper.session_helper import Session

sagemaker_session = Session()
AWS_REGION = sagemaker_session.boto_region_name

This avoids the explicit boto3 import entirely and is more consistent with V3's "Core as the foundation" tenet. The boto3 import was previously used in Step 7 for raw sagemaker-runtime client calls, which this PR correctly removes — so there's no remaining need for it.

"from sagemaker.core import image_uris\n",
"from sagemaker.core.helper.session_helper import Session\n",
"\n",
Expand All @@ -71,7 +72,9 @@
"MLFLOW_TRACKING_ARN = \"XXXXX\"\n",
"\n",
"# AWS Configuration\n",
"AWS_REGION = Session.boto_region_name\n",
"boto_session = boto3.Session()\n",
"sagemaker_session = Session(boto_session=boto_session)\n",
"AWS_REGION = sagemaker_session.boto_region_name\n",
"\n",
"# Get PyTorch training image dynamically\n",
"PYTORCH_TRAINING_IMAGE = image_uris.retrieve(\n",
Expand Down Expand Up @@ -297,6 +300,7 @@
"\n",
"# Training on SageMaker managed infrastructure\n",
"model_trainer = ModelTrainer(\n",
" sagemaker_session=sagemaker_session,\n",
Comment thread
aviruthen marked this conversation as resolved.
" training_image=PYTORCH_TRAINING_IMAGE,\n",
" source_code=SourceCode(\n",
" source_dir=training_code_dir,\n",
Expand Down Expand Up @@ -333,22 +337,29 @@
"from mlflow import MlflowClient\n",
"\n",
"client = MlflowClient()\n",
"registered_model = client.get_registered_model(name=MLFLOW_REGISTERED_MODEL_NAME)\n",
"\n",
"latest_version = registered_model.latest_versions[0]\n",
"# Use search_model_versions (compatible with MLflow 3.x)\n",
"# Note: latest_versions attribute was removed in MLflow 3.x\n",
"model_versions = client.search_model_versions(\n",
" filter_string=f\"name='{MLFLOW_REGISTERED_MODEL_NAME}'\",\n",
" order_by=['version_number DESC'],\n",
Comment thread
aviruthen marked this conversation as resolved.
" max_results=1\n",
")\n",
"\n",
"if not model_versions:\n",
" raise ValueError(f\"No versions found for model '{MLFLOW_REGISTERED_MODEL_NAME}'\")\n",
"\n",
"latest_version = model_versions[0]\n",
"model_version = latest_version.version\n",
"model_source = latest_version.source\n",
"\n",
"# Get S3 URL of model files (for info only)\n",
"artifact_uri = client.get_model_version_download_uri(MLFLOW_REGISTERED_MODEL_NAME, model_version)\n",
"\n",
"# MLflow model registry path to use with ModelBuilder\n",
"mlflow_model_path = f\"models:/{MLFLOW_REGISTERED_MODEL_NAME}/{model_version}\"\n",
"\n",
"print(f\"Registered Model: {MLFLOW_REGISTERED_MODEL_NAME}\")\n",
"print(f\"Latest Version: {model_version}\")\n",
"print(f\"Source: {model_source}\")\n",
"print(f\"Model artifacts location: {artifact_uri}\")"
"print(f\"Source (artifact location): {model_source}\")\n",
"print(f\"MLflow model path for deployment: {mlflow_model_path}\")"
]
},
{
Expand Down Expand Up @@ -481,23 +492,44 @@
"metadata": {},
"outputs": [],
"source": [
"import boto3\n",
"\n",
"# Test with JSON input\n",
"test_data = [[0.1, 0.2, 0.3, 0.4]]\n",
"\n",
Comment thread
aviruthen marked this conversation as resolved.
"runtime_client = boto3.client('sagemaker-runtime')\n",
"response = runtime_client.invoke_endpoint(\n",
" EndpointName=core_endpoint.endpoint_name,\n",
" Body=json.dumps(test_data),\n",
" ContentType='application/json'\n",
"result = core_endpoint.invoke(\n",
" body=json.dumps(test_data),\n",
" content_type='application/json'\n",
")\n",
"\n",
"prediction = json.loads(response['Body'].read().decode('utf-8'))\n",
"# Decode and display the result\n",
"prediction = json.loads(result.body.read().decode('utf-8'))\n",
"print(f\"Input: {test_data}\")\n",
"print(f\"Prediction: {prediction}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Test with different tensor inputs\n",
"test_inputs = [\n",
" [[0.5, 0.3, 0.2, 0.1]],\n",
" [[0.9, 0.1, 0.8, 0.2]],\n",
" [[0.2, 0.7, 0.4, 0.6]]\n",
"]\n",
"\n",
"for i, test_input in enumerate(test_inputs, 1):\n",
" result = core_endpoint.invoke(\n",
" body=json.dumps(test_input),\n",
" content_type='application/json'\n",
" )\n",
" \n",
" prediction = json.loads(result.body.read().decode('utf-8'))\n",
" print(f\"Test {i} - Input {test_input}: {prediction}\")\n",
" print(\"-\" * 50)"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down
Loading