Notebook usability update (#5522)

Author: mollyheamazon

v3-examples/inference-examples/train-inference-e2e-example.ipynb

@@ -62,7 +62,7 @@
     "\n",
     "# AWS Configuration\n",
     "AWS_REGION = \"us-west-2\"\n",
-    "PYTORCH_TRAINING_IMAGE = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.13.1-cpu-py39\"\n",
+    "PYTORCH_TRAINING_IMAGE = f\"763104351884.dkr.ecr.{AWS_REGION}.amazonaws.com/pytorch-training:1.13.1-cpu-py39\"\n",
     "\n",
     "# Generate unique identifiers\n",
     "unique_id = str(uuid.uuid4())[:8]\n",

v3-examples/training-examples/custom-distributed-training-example.ipynb

@@ -22,7 +22,8 @@
     "from sagemaker.train.model_trainer import ModelTrainer\n",
     "from sagemaker.train.configs import SourceCode\n",
     "from sagemaker.train.distributed import DistributedConfig\n",
-    "from sagemaker.core.helper.session_helper import Session, get_execution_role"
+    "from sagemaker.core.helper.session_helper import Session, get_execution_role\n",
+    "from sagemaker.core import image_uris"
    ]
   },
   {
@@ -42,8 +43,16 @@
    "source": [
     "sagemaker_session = Session()\n",
     "role = get_execution_role()\n",
-    "\n",
-    "DEFAULT_CPU_IMAGE = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310\"\n",
+    "region = sagemaker_session.boto_region_name\n",
+    "\n",
+    "DEFAULT_CPU_IMAGE = image_uris.retrieve(\n",
+    "    framework=\"pytorch\",\n",
+    "    region=region,\n",
+    "    version=\"2.0.0\",\n",
+    "    py_version=\"py310\",\n",
+    "    instance_type=\"ml.m5.xlarge\",\n",
+    "    image_scope=\"training\"\n",
+    ")\n",
     "\n",
     "# Create temporary directories\n",
     "temp_dir = tempfile.mkdtemp()\n",

v3-examples/training-examples/distributed-local-training-example.ipynb

@@ -24,7 +24,8 @@
     "from sagemaker.train.model_trainer import ModelTrainer, Mode\n",
     "from sagemaker.train.configs import SourceCode, Compute, InputData\n",
     "from sagemaker.train.distributed import Torchrun\n",
-    "from sagemaker.core.helper.session_helper import Session"
+    "from sagemaker.core.helper.session_helper import Session\n",
+    "from sagemaker.core import image_uris"
    ]
   },
   {
@@ -54,7 +55,16 @@
    "outputs": [],
    "source": [
     "sagemaker_session = Session()\n",
-    "DEFAULT_CPU_IMAGE = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310\"\n",
+    "region = sagemaker_session.boto_region_name\n",
+    "\n",
+    "DEFAULT_CPU_IMAGE = image_uris.retrieve(\n",
+    "    framework=\"pytorch\",\n",
+    "    region=region,\n",
+    "    version=\"2.0.0\",\n",
+    "    py_version=\"py310\",\n",
+    "    instance_type=\"ml.m5.xlarge\",\n",
+    "    image_scope=\"training\"\n",
+    ")\n",
     "\n",
     "# Create temporary directories\n",
     "temp_dir = tempfile.mkdtemp()\n",

v3-examples/training-examples/hyperparameter-training-example.ipynb

@@ -23,7 +23,8 @@
     "\n",
     "from sagemaker.train.model_trainer import ModelTrainer\n",
     "from sagemaker.train.configs import SourceCode\n",
-    "from sagemaker.core.helper.session_helper import Session, get_execution_role"
+    "from sagemaker.core.helper.session_helper import Session, get_execution_role\n",
+    "from sagemaker.core import image_uris"
    ]
   },
   {
@@ -43,6 +44,7 @@
    "source": [
     "sagemaker_session = Session()\n",
     "role = get_execution_role()\n",
+    "region = sagemaker_session.boto_region_name\n",
     "\n",
     "# Expected hyperparameters\n",
     "EXPECTED_HYPERPARAMETERS = {\n",
@@ -61,7 +63,14 @@
     "    },\n",
     "}\n",
     "\n",
-    "DEFAULT_CPU_IMAGE = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310\"\n",
+    "DEFAULT_CPU_IMAGE = image_uris.retrieve(\n",
+    "    framework=\"pytorch\",\n",
+    "    region=region,\n",
+    "    version=\"2.0.0\",\n",
+    "    py_version=\"py310\",\n",
+    "    instance_type=\"ml.m5.xlarge\",\n",
+    "    image_scope=\"training\"\n",
+    ")\n",
     "\n",
     "# Create temporary directory\n",
     "temp_dir = tempfile.mkdtemp()\n",

v3-examples/training-examples/jumpstart-training-example.ipynb

@@ -276,6 +276,71 @@
     "    print(f\"{model_name}: trainer.latest_training_job.describe()\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 7: Discovery other models available in Jumpstart Hub (Optional)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sagemaker.core.jumpstart.notebook_utils import list_jumpstart_models\n",
+    "from sagemaker.core.jumpstart.search import search_public_hub_models\n",
+    "\n",
+    "# List all available JumpStart models\n",
+    "models = list_jumpstart_models()\n",
+    "\n",
+    "# Filter by framework (e.g., HuggingFace)\n",
+    "huggingface_models = list_jumpstart_models(filter=\"framework == huggingface\")\n",
+    "\n",
+    "print(huggingface_models)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Search for specific models\n",
+    "results = search_public_hub_models(query=\"bert\")\n",
+    "\n",
+    "print(results)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Search with logical expressions\n",
+    "text_gen_models = search_public_hub_models(\n",
+    "    query=\"@task:text-generation\"\n",
+    ")\n",
+    "\n",
+    "print(text_gen_models)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Complex queries\n",
+    "hf_bert = search_public_hub_models(\n",
+    "    query=\"@framework:huggingface AND bert\"\n",
+    ")\n",
+    "\n",
+    "print(hf_bert)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},