Add bugbash bootstrapping (#1598)

Author: benieric

src/sagemaker/modules/testing_notebooks/base_model_trainer.ipynb

@@ -326,7 +326,7 @@
     "    channel_name=\"dataset\",\n",
     "    data_source=training_input_path,\n",
     ")\n",
-    "model_trainer.train(input_data_config=[test_data])"
+    "model_trainer.train(input_data_config=[test_data], wait=False)"
    ]
   },
   {
@@ -417,7 +417,7 @@
     "    channel_name=\"dataset\",\n",
     "    data_source=training_input_path,\n",
     ")\n",
-    "model_trainer.train(input_data_config=[test_data])"
+    "model_trainer.train(input_data_config=[test_data], wait=False)"
    ]
   },
   {
@@ -487,10 +487,16 @@
     "    ]\n",
     ")\n",
     "\n",
+    "env = {\n",
+    "    \"NCCL_P2P_DISABLE\": \"1\",\n",
+    "    \"NCCL_DEBUG\": \"WARN\",\n",
+    "}\n",
+    "\n",
     "model_trainer = ModelTrainer(\n",
     "    training_image=hugging_face_image,\n",
     "    compute=compute,\n",
     "    hyperparameters=hyperparameters,\n",
+    "    environment=env,\n",
     "    source_code=source_code,\n",
     "    distributed_runner=mpi,\n",
     "    base_job_name=f\"{alias}-distributed-case-3\",\n",
@@ -507,7 +513,7 @@
     "    channel_name=\"dataset\",\n",
     "    data_source=training_input_path,\n",
     ")\n",
-    "model_trainer.train(input_data_config=[test_data])"
+    "model_trainer.train(input_data_config=[test_data], wait=False)"
    ]
   },
   {
@@ -570,7 +576,7 @@
     "    training_recipe=\"training/llama/hf_llama3_8b_seq8192_gpu\",\n",
     "    training_image=training_image,\n",
     "    recipe_overrides=recipe_overrides,\n",
-    "    compute=Compute(instance_type=\"ml.g5.48xlarge\"),\n",
+    "    compute=Compute(instance_type=\"ml.p4d.24xlarge\", keep_alive_period_in_seconds=3600),\n",
     "    base_job_name=f\"{alias}-recipe-case-1\",\n",
     ")"
    ]
@@ -581,7 +587,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_trainer.train()"
+    "model_trainer.train(wait=False)"
    ]
   },
   {
@@ -607,7 +613,7 @@
     "model_trainer = ModelTrainer.from_recipe(\n",
     "    training_recipe=\"recipes/custom-recipe.yaml\",\n",
     "    training_image=training_image,\n",
-    "    compute=Compute(instance_type=\"ml.g5.48xlarge\"),\n",
+    "    compute=Compute(instance_type=\"ml.p4d.24xlarge\", keep_alive_period_in_seconds=3600),\n",
     "    base_job_name=f\"{alias}-recipe-case-2\",\n",
     ")"
    ]
@@ -618,7 +624,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_trainer.train()"
+    "model_trainer.train(wait=False)"
    ]
   },
   {
@@ -678,11 +684,13 @@
     "    compute=Compute(\n",
     "        instance_type=\"ml.trn1.32xlarge\",\n",
     "        instance_count=2,\n",
+    "        keep_alive_period_in_seconds=3600\n",
     "    ),\n",
     "    stopping_condition=StoppingCondition(\n",
     "        max_runtime_in_seconds=3600\n",
     "    ),\n",
-    "    environment=env\n",
+    "    environment=env,\n",
+    "    base_job_name=f\"{alias}-recipe-case-3\",\n",
     ")"
    ]
   },
@@ -699,6 +707,13 @@
     "\n",
     "model_trainer.train(input_data_config=[train], wait=False)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

src/sagemaker/modules/testing_notebooks/bootstrap.sh renamed to src/sagemaker/modules/testing_notebooks/bootstrap/bootstrap.sh

@@ -1,5 +1,7 @@
 #!/bin/bash
 
+echo "Installing Docker"
+
 sudo apt-get -y install ca-certificates curl gnupg
 
 sudo install -m 0755 -d /etc/apt/keyrings
@@ -19,5 +21,27 @@ sudo apt-get install docker-ce-cli=$VERSION_STRING docker-compose-plugin -y
 # validate the Docker Client is able to access Docker Server at [unix:///docker/proxy.sock]
 docker version
 
+echo "Installing Local SageMaker Tarball"
 pip install "pydantic>=2.0.0"
-pip install sagemaker-2.232.4.dev0.tar.gz
+pip install sagemaker-2.232.4.dev0.tar.gz
+
+
+echo "Setting Up Read-Only SSH Access"
+eval "$(ssh-agent -s)"
+
+mkdir -p ~/.ssh/
+
+cp /home/sagemaker-user/bootstrap/adapter_deploy_key /home/sagemaker-user/.ssh/adapter_deploy_key
+chmod 600 ~/.ssh/adapter_deploy_key
+
+cp /home/sagemaker-user/bootstrap/launcher_deploy_key /home/sagemaker-user/.ssh/launcher_deploy_key 
+chmod 600 ~/.ssh/launcher_deploy_key
+
+cp /home/sagemaker-user/bootstrap/config /home/sagemaker-user/.ssh/config
+chmod 644 ~/.ssh/config
+
+ssh-add ~/.ssh/adapter_deploy_key
+ssh-add ~/.ssh/launcher_deploy_key
+
+
+

src/sagemaker/modules/testing_notebooks/bootstrap/config

@@ -0,0 +1,7 @@
+Host i.8713187.xyz-launcher
+        Hostname github.com
+        IdentityFile=~/.ssh/launcher_deploy_key
+
+Host i.8713187.xyz-adapter
+        Hostname github.com
+        IdentityFile=~/.ssh/adapter_deploy_key

src/sagemaker/modules/testing_notebooks/sagemaker-2.232.4.dev0-py3-none-any.whl

@@ -206,7 +206,6 @@ def get_mpirun_command(
         str(num_processes),
         "--allow-run-as-root",
         "--tag-output",
-        "--oversubscribe",
         "-mca",
         "btl_tcp_if_include",
         network_interface_name,
@@ -245,6 +244,7 @@ def get_mpirun_command(
     instance_type = os.environ["SM_CURRENT_INSTANCE_TYPE"]
     # EFA settings
     if instance_type in SM_EFA_NCCL_INSTANCES:
+        mpirun_command.extend(["-x", "FI_PROVIDER=efa"])
         # Use simple protocol to handle the out-of-order data delivery from EFA
         mpirun_command.extend(["-x", "NCCL_PROTO=simple"])
 

src/sagemaker/modules/testing_notebooks/sagemaker-2.232.4.dev0.tar.gz

@@ -1,6 +1,6 @@
 {
-    "adapter_repo": "[email protected]:aws/private-sagemaker-training-adapter-for-nemo-staging.git",
-    "launcher_repo": "[email protected]:aws/private-sagemaker-training-launcher-staging.git",
+    "adapter_repo": "[email protected]-adapter:benieric/private-sagemaker-hyperpod-training-adapter-for-nemo-staging.git",
+    "launcher_repo": "[email protected]-launcher:benieric/private-sagemaker-hyperpod-recipes-staging.git",
     "neuron_dist_repo": "https://github.com/aws-neuron/neuronx-distributed-training.git",
     "gpu_image" : {
         "framework": "pytorch-smp",