fix: address review comments (iteration #1)

Author: aviruthen

v3-examples/ml-ops-examples/v3-mlflow-train-inference-e2e-example.ipynb

@@ -35,8 +35,8 @@
    "outputs": [],
    "source": [
     "# Install fix for MLflow path resolution issues\n",
-    "# Note: Version pins (mlflow==3.4.0) may need updating - check for latest compatible versions\n",
-    "%pip install mlflow==3.4.0"
+    "# Using minimum version constraint so the notebook stays compatible with future releases\n",
+    "%pip install 'mlflow>=3.4.0'"
    ]
   },
   {
@@ -63,7 +63,6 @@
    "source": [
     "import json\n",
     "import uuid\n",
-    "import boto3\n",
     "from sagemaker.core import image_uris\n",
     "from sagemaker.core.helper.session_helper import Session\n",
     "\n",
@@ -301,14 +300,14 @@
     "\n",
     "# Training on SageMaker managed infrastructure\n",
     "model_trainer = ModelTrainer(\n",
-    "    sagemaker_session=sagemaker_session,\n",
     "    training_image=PYTORCH_TRAINING_IMAGE,\n",
     "    source_code=SourceCode(\n",
     "        source_dir=training_code_dir,\n",
     "        entry_script=\"train.py\",\n",
     "        requirements=\"requirements.txt\",\n",
     "    ),\n",
     "    base_job_name=training_job_name,\n",
+    "    sagemaker_session=sagemaker_session,\n",
     ")\n",
     "\n",
     "# Start training job\n",
@@ -340,9 +339,11 @@
     "client = MlflowClient()\n",
     "\n",
     "# Use search_model_versions (compatible with MLflow 3.x)\n",
+    "# Note: order_by field name may vary across MLflow versions;\n",
+    "# 'creation_timestamp' is broadly supported.\n",
     "model_versions = client.search_model_versions(\n",
     "    filter_string=f\"name='{MLFLOW_REGISTERED_MODEL_NAME}'\",\n",
-    "    order_by=['version_number DESC'],\n",
+    "    order_by=['creation_timestamp DESC'],\n",
     "    max_results=1\n",
     ")\n",
     "latest_version = model_versions[0]\n",
@@ -413,7 +414,7 @@
     "        \"MLFLOW_MODEL_PATH\": mlflow_model_path,\n",
     "        \"MLFLOW_TRACKING_ARN\": MLFLOW_TRACKING_ARN\n",
     "    },\n",
-    "    dependencies={\"auto\": False, \"custom\": [\"mlflow==3.4.0\", \"sagemaker==3.3.1\", \"numpy==2.4.1\", \"cloudpickle==3.1.2\"]},\n",
+    "    dependencies={\"auto\": False, \"custom\": [\"mlflow>=3.4.0\", \"sagemaker>=3.3.1\", \"numpy>=2.4.1\", \"cloudpickle>=3.1.2\"]},\n",
     ")\n",
     "\n",
     "print(f\"ModelBuilder configured with MLflow model: {mlflow_model_path}\")"
@@ -464,11 +465,18 @@
     "test_data = [[0.1, 0.2, 0.3, 0.4]]\n",
     "\n",
     "result = core_endpoint.invoke(\n",
-    "    body=json.dumps(test_data),\n",
+    "    body=json.dumps(test_data).encode('utf-8'),\n",
     "    content_type='application/json'\n",
     ")\n",
     "\n",
-    "prediction = json.loads(result.body.read().decode('utf-8'))\n",
+    "# The invoke() response body may be a streaming object or bytes;\n",
+    "# handle both cases for robustness.\n",
+    "response_body = result.body\n",
+    "if hasattr(response_body, 'read'):\n",
+    "    response_body = response_body.read()\n",
+    "if isinstance(response_body, bytes):\n",
+    "    response_body = response_body.decode('utf-8')\n",
+    "prediction = json.loads(response_body)\n",
     "print(f\"Input: {test_data}\")\n",
     "print(f\"Prediction: {prediction}\")"
    ]