Synchronize SSH keys

Bobber relies on SSH keys that are baked into the images to enable
multi-node communication. This forces users to build the image on one
machine, save the image locally, copy it to all remote nodes, and load
the copied image on those hosts. This process is long and tedious, but
by replacing it with a synchronization method, makes it possible to
run the build on each host and not need to copy images remotely.

Signed-Off-By: Robert Clark <roclark@nvidia.com>

Author: roclark

README.md

@@ -105,9 +105,8 @@ pip3 install https://github.com/NVIDIA/Bobber/releases/download/v6.1.1/nvidia_bo
 
 ## Build Bobber container (includes OSU Tests, NCCL Tests, fio, mdtest, DALI RN50 Pipeline, and the base NGC TensorFlow container)
 The Bobber application includes a built-in mechanism to build the Docker
-container where all tests will be run. This command should be run on a single
-system in the cluster as it will be copied in a future step. For single-node
-tests, run the command on the node to be tested.
+container where all tests will be run. Run the following command on all nodes
+that will be tested.
 
 ```bash
 $ bobber build
@@ -123,13 +122,13 @@ $ docker images | grep nvidia/bobber
 nvidia/bobber               6.1.1               c697a75ee482        36 minutes ago      12.4GB
 ```
 
-## Save container
+## Synchronize container keys
 Bobber relies on shared SSH keys to communicate between containers via MPI. This
-is done by generating an SSH key in the image during build time and using that
-same container on all hosts. This requires saving the image to a local tarball
-and transferring the image to all other nodes. The `export` command saves the
-image as a local tarball. Run the command on the node from the previous step
-where the Docker image is located.
+is done by generating an SSH key on a single node and copying that key to all
+other containers in the cluster. Bobber includes a tool to automatically
+synchronize keys amongst all containers on all nodes. To make the process more
+seamless, it is highly recommended to create passwordless SSH keys to
+communicate between the hosts in order to automatically copy the keys.
 
 If running on a single node, this step is not required.
 

bobber/bin/sync-keys.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+set -e
+# Pass the list of hosts in as a string
+hosts=$1
+# Optionally pass a username to login to remote nodes
+user=$2
+
+# Generate a new RSA key locally for SSH to share across the cluster
+mkdir -p /tmp/bobber
+rm -f /tmp/bobber/*
+ssh-keygen -t rsa -b 4096 -f /tmp/bobber/id_rsa -N ""
+
+echo "Copying keys to containers on all hosts"
+echo "For remote hosts, if passwordless-ssh is not configured, you will be prompted for the password for all nodes"
+
+if [[ $hosts=="localhost" || -z "$hosts" ]]; then
+  docker cp /tmp/bobber/id_rsa bobber:/root/.ssh/id_rsa
+  docker cp /tmp/bobber/id_rsa.pub bobber:/root/.ssh/authorized_keys
+fi
+
+# Copy the key to the container
+for host in ${hosts//,/ }; do
+  if [ ! -z "$user" ]; then
+    scp -r /tmp/bobber $user@$host:/tmp/
+    ssh $user@$host 'docker cp /tmp/bobber/id_rsa bobber:/root/.ssh/id_rsa && docker cp /tmp/bobber/id_rsa.pub bobber:/root/.ssh/authorized_keys && rm /tmp/bobber/id_rsa*'
+  else
+    scp -r /tmp/bobber $host:/tmp/
+    ssh $host 'docker cp /tmp/bobber/id_rsa bobber:/root/.ssh/id_rsa && docker cp /tmp/bobber/id_rsa.pub bobber:/root/.ssh/authorized_keys && rm /tmp/bobber/id_rsa*'
+  fi
+done
+
+# Cleanup the local key
+rm -f /tmp/bobber/id_rsa*

bobber/bobber.py

@@ -20,10 +20,12 @@
     RUN_STG_BW,
     RUN_STG_IOPS,
     RUN_STG_META,
+    SYNC,
     SYSTEMS
 )
 from bobber.lib.analysis import parse_results
 from bobber.lib.system.file_handler import create_directory
+from bobber.lib.system.shell import copy_keys
 from bobber.lib.tests import run_tests
 from typing import NoReturn
 
@@ -227,6 +229,21 @@ def parse_args(version: str) -> Namespace:
                                'binary')
     load.add_argument('filename', help='Filename of local *.tar file of '
                       'the image to load')
+
+    # Options specific to synchronizing SSH keys in containers
+    sync = commands.add_parser(SYNC, help='Create SSH keys and add them to all'
+                               ' Bobber containers in a cluster. Requires the '
+                               'container to be running on all nodes using '
+                               '"bobber cast".')
+    sync.add_argument('--hosts', help='A comma-separated list of hostnames or '
+                      'IP address of the nodes to add SSH keys to. Required '
+                      'for multi-node tests. If left empty, it is assumed '
+                      'that keys should only be copied to the container on '
+                      'the local node.', type=str, default='')
+    sync.add_argument('--user', help='Optionally specify a user to use to '
+                      'login to remote hosts to copy keys to containers. If '
+                      'left blank, will use the currently logged-in user.',
+                      type=str, default='')
     return parser.parse_args()
 
 
@@ -360,6 +377,8 @@ def execute_command(args: Namespace, version: str) -> NoReturn:
         bobber.lib.docker.cast(args.storage_path, args.ignore_gpu, version)
     elif args.command == LOAD:
         bobber.lib.docker.load(args.filename)
+    elif args.command == SYNC:
+        copy_keys(args.hosts, args.user)
     else:
         # Update the version to be used in filenames
         version_underscore = version.replace('.', '_')

bobber/lib/constants.py

@@ -10,6 +10,7 @@
 RUN_STG_BW = 'run-stg-bw'
 RUN_STG_IOPS = 'run-stg-iops'
 RUN_STG_META = 'run-stg-meta'
+SYNC = 'sync'
 
 DGX_A100_SINGLE = {
     'gpus': 8,

bobber/lib/docker/management.py

@@ -234,3 +234,24 @@ def execute(self, command: str, environment: Optional[dict] = None,
                 print(result.output)
             except StopIteration:
                 break
+
+    def running(self):
+        """
+        Determine if the Bobber container is running on the system.
+
+        Check to see if the current version of the Bobber container is running
+        on the local machine and return the status. This method can be used to
+        determine whether or not to run a command that depends on the container
+        being launched.
+
+        Returns
+        -------
+        boolean
+            Returns `True` when the container is running and `False` when not.
+        """
+        try:
+            bobber = self.client.containers.get('bobber')
+        except docker.errors.NotFound:
+            return False
+        else:
+            return True

bobber/lib/system/shell.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: MIT
+import subprocess
+import sys
+from bobber.lib.docker import manager
+from typing import NoReturn
+
+
+def copy_keys(hosts: str, user: str) -> NoReturn:
+    """
+    Generate and copy SSH keys to all hosts.
+
+    Launch a shell script included with the package which generates a local SSH
+    key that is copied to all Bobber containers on all nodes to allow
+    passwordless communication for MPI.
+
+    Parameters
+    ----------
+    hosts : string
+        A comma-separated list as a ``string`` representing all hosts, such as
+        'host1,host2,host3,...'.
+    user : string
+        A ``string`` of the user to use to login to remote hosts as, if
+        necessary.
+    """
+    if not manager.running():
+        print('Bobber container is not running. Please ensure Bobber is '
+              'running on all nodes using the "bobber cast" command before '
+              'running "bobber sync".')
+        sys.exit(-1)
+    try:
+        subprocess.run(['bobber/bin/sync-keys.sh', hosts, user], check=True)
+    except subprocess.CalledProcessError:
+        print('Error synchronizing keys. See output from the sync script '
+              'above.')
+        sys.exit(-1)