add AddKeysToAgent for ssh config file and ssh cmd

examples/mpirun_test.yaml

@@ -0,0 +1,24 @@
+workdir: .
+
+resources:
+  cloud: aws
+
+num_nodes: 2  # Total number of nodes (1 head + 1 worker)
+
+setup: |
+  echo "Running setup on node ${SKYPILOT_NODE_RANK}."
+  # Install MPI if not already present. This will vary based on your OS/distro.
+  sudo apt update
+  sudo apt install -y openmpi-bin openmpi-common libopenmpi-dev
+
+run: |
+  if [ "${SKYPILOT_NODE_RANK}" == "0" ]; then
+    echo "head node"
+    num_nodes=`echo "$SKYPILOT_NODE_IPS" | wc -l`
+    mpi_nodes=$(echo "$SKYPILOT_NODE_IPS" | tr '\n' ',')
+    mpi_nodes=${mpi_nodes::-1}
+    echo "$mpi_nodes"
+    mpirun -np $num_nodes -H $mpi_nodes bash -c 'echo "mpirun hello from IP $(hostname -I)"'
+  else
+    echo "worker nodes"
+  fi

sky/backends/backend_utils.py

@@ -428,6 +428,7 @@ def _get_generated_config(cls, autogen_comment: str, host_name: str,
               HostName {ip}
               User {username}
               IdentityFile {ssh_key_path}
+              AddKeysToAgent yes
               IdentitiesOnly yes
               ForwardAgent yes
               StrictHostKeyChecking no

sky/provision/provisioner.py

@@ -259,6 +259,8 @@ def _ssh_probe_command(ip: str,
         '-o',
         'IdentitiesOnly=yes',
         '-o',
+        'AddKeysToAgent=yes',
+        '-o',
         'ExitOnForwardFailure=yes',
         '-o',
         'ServerAliveInterval=5',

sky/utils/command_runner.py

@@ -85,6 +85,10 @@ def ssh_options_list(
         'LogLevel': 'ERROR',
         # Try fewer extraneous key pairs.
         'IdentitiesOnly': 'yes',
+        # Add the current private key to the agent so that agent forwarding
+        # will forward the key used for authentication instead of the system
+        # default in the user's home directory.
+        'AddKeysToAgent': 'yes',
         # Abort if port forwarding fails (instead of just printing to
         # stderr).
         'ExitOnForwardFailure': 'yes',