Updating Trainer wait loop refresh time (#5361)

Author: jam-jee

sagemaker-train/src/sagemaker/train/common_utils/trainer_wait.py

@@ -178,10 +178,8 @@ def wait(
             iteration = 0
             while True:
                 iteration += 1
-                time.sleep(1)
-                if iteration == poll:
-                    iteration = 0
-                    training_job.refresh()
+                time.sleep(poll)
+                training_job.refresh()
                 clear_output(wait=False)
 
                 status = training_job.training_job_status

v3-examples/model-customization-examples/rlaif_finetuning_example_notebook_v3_prod.ipynb

@@ -70,15 +70,15 @@
     "from rich import print as rprint\n",
     "from rich.pretty import pprint\n",
     "from sagemaker.core.resources import ModelPackage\n",
-    "#import os\n",
+    "import os\n",
     "#os.environ['SAGEMAKER_REGION'] = 'us-east-1'\n",
     "#os.environ['SAGEMAKER_STAGE'] = 'prod'\n",
     "\n",
     "import boto3\n",
     "from sagemaker.core.helper.session_helper import Session\n",
     "\n",
     "# For MLFlow native metrics in Trainer wait, run below line with approriate region\n",
-    "# os.environ[\"SAGEMAKER_MLFLOW_CUSTOM_ENDPOINT\"] = \"https://mlflow.sagemaker.us-east-1.app.aws\"\n",
+    "os.environ[\"SAGEMAKER_MLFLOW_CUSTOM_ENDPOINT\"] = \"https://mlflow.sagemaker.us-west-2.app.aws\"\n",
     "\n"
    ]
   },

v3-examples/model-customization-examples/rlvr_finetuning_example_notebook_v3-prod.ipynb

@@ -23,45 +23,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
    "id": "10c2ef37-2425-4676-bc80-6d278d4e609a",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[12/02/25 17:54:14] </span><span style=\"color: #0069ff; text-decoration-color: #0069ff; font-weight: bold\">INFO    </span> Found credentials in shared credentials file: ~<span style=\"color: #e100e1; text-decoration-color: #e100e1\">/.aws/credentials</span>   <a href=\"file:///Users/rsareddy/workplace/virtual_envs/sagemaker-v3/lib/python3.12/site-packages/botocore/credentials.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">credentials.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/rsareddy/workplace/virtual_envs/sagemaker-v3/lib/python3.12/site-packages/botocore/credentials.py#1364\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1364</span></a>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "\u001B[2;36m[12/02/25 17:54:14]\u001B[0m\u001B[2;36m \u001B[0m\u001B[1;38;2;0;105;255mINFO    \u001B[0m Found credentials in shared credentials file: ~\u001B[38;2;225;0;225m/.aws/\u001B[0m\u001B[38;2;225;0;225mcredentials\u001B[0m   \u001B]8;id=600716;file:///Users/rsareddy/workplace/virtual_envs/sagemaker-v3/lib/python3.12/site-packages/botocore/credentials.py\u001B\\\u001B[2mcredentials.py\u001B[0m\u001B]8;;\u001B\\\u001B[2m:\u001B[0m\u001B]8;id=966746;file:///Users/rsareddy/workplace/virtual_envs/sagemaker-v3/lib/python3.12/site-packages/botocore/credentials.py#1364\u001B\\\u001B[2m1364\u001B[0m\u001B]8;;\u001B\\\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "sagemaker.config INFO - Not applying SDK defaults from location: /Library/Application Support/sagemaker/config.yaml\n",
-      "sagemaker.config INFO - Not applying SDK defaults from location: /Users/rsareddy/Library/Application Support/sagemaker/config.yaml\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #0069ff; text-decoration-color: #0069ff; font-weight: bold\">INFO    </span> Found credentials in shared credentials file: ~<span style=\"color: #e100e1; text-decoration-color: #e100e1\">/.aws/credentials</span>   <a href=\"file:///Users/rsareddy/workplace/virtual_envs/sagemaker-v3/lib/python3.12/site-packages/botocore/credentials.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">credentials.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/rsareddy/workplace/virtual_envs/sagemaker-v3/lib/python3.12/site-packages/botocore/credentials.py#1364\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1364</span></a>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "\u001B[2;36m                   \u001B[0m\u001B[2;36m \u001B[0m\u001B[1;38;2;0;105;255mINFO    \u001B[0m Found credentials in shared credentials file: ~\u001B[38;2;225;0;225m/.aws/\u001B[0m\u001B[38;2;225;0;225mcredentials\u001B[0m   \u001B]8;id=25130;file:///Users/rsareddy/workplace/virtual_envs/sagemaker-v3/lib/python3.12/site-packages/botocore/credentials.py\u001B\\\u001B[2mcredentials.py\u001B[0m\u001B]8;;\u001B\\\u001B[2m:\u001B[0m\u001B]8;id=835202;file:///Users/rsareddy/workplace/virtual_envs/sagemaker-v3/lib/python3.12/site-packages/botocore/credentials.py#1364\u001B\\\u001B[2m1364\u001B[0m\u001B]8;;\u001B\\\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
    "source": [
     "\n",
     "from sagemaker.train.rlvr_trainer import RLVRTrainer\n",
@@ -73,12 +36,15 @@
     "\n",
     "\n",
     "import boto3\n",
+    "import os\n",
     "from sagemaker.core.helper.session_helper import Session\n",
     "\n",
     "# For MLFlow native metrics in Trainer wait, run below line with approriate region\n",
-    "# os.environ[\"SAGEMAKER_MLFLOW_CUSTOM_ENDPOINT\"] = \"https://mlflow.sagemaker.us-east-1.app.aws\"\n",
+    "os.environ[\"SAGEMAKER_MLFLOW_CUSTOM_ENDPOINT\"] = \"https://mlflow.sagemaker.us-west-2.app.aws\"\n",
     "\n"
-   ]
+   ],
+   "outputs": [],
+   "execution_count": null
   },
   {
    "cell_type": "markdown",