From 6151086e145e036b78f89a29ad3932273d6064bb Mon Sep 17 00:00:00 2001
From: Zack Cerza <zack@redhat.com>
Date: Tue, 7 Nov 2017 16:05:21 -0700
Subject: [PATCH] Reimage machines in parallel

Signed-off-by: Zack Cerza <zack@redhat.com>
---
 teuthology/lock/cli.py | 29 +++++++++++++++++++++++------
 teuthology/lock/ops.py |  8 +++++---
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/teuthology/lock/cli.py b/teuthology/lock/cli.py
index a3c3002a2..4848e502e 100644
--- a/teuthology/lock/cli.py
+++ b/teuthology/lock/cli.py
@@ -7,6 +7,7 @@ import re
 import yaml
 
 import teuthology
+import teuthology.parallel
 import teuthology.provision
 from teuthology import misc
 from teuthology.config import set_config_attr
@@ -144,24 +145,40 @@ def main(ctx):
         return 0
 
     elif ctx.lock:
-        if not util.vps_version_or_type_valid(ctx.machine_type, ctx.os_type,
-                                         ctx.os_version):
+        if not util.vps_version_or_type_valid(
+                ctx.machine_type, ctx.os_type, ctx.os_version):
             log.error('Invalid os-type or version detected -- lock failed')
             return 1
+        reimage_types = teuthology.provision.fog.get_types()
+        reimage_machines = list()
+        updatekeys_machines = list()
         for machine in machines:
-            if not ops.lock_one(machine, user, ctx.desc):
+            resp = ops.lock_one(machine, user, ctx.desc)
+            if resp.ok:
+                machine_status = resp.json()
+                machine_type = machine_status['machine_type']
+            if not resp.ok:
                 ret = 1
                 if not ctx.f:
                     return ret
-            elif not query.is_vm(machine):
-                teuthology.provision.reimage(ctx, machine)
-                keys.do_update_keys([machine])
+            elif not query.is_vm(machine, machine_status):
+                if machine_type in reimage_types:
+                    # Reimage in parallel just below here
+                    reimage_machines.append(machine)
+                # Update keys last
+                updatekeys_machines = list()
             else:
                 machines_to_update.append(machine)
                 teuthology.provision.create_if_vm(
                     ctx,
                     misc.canonicalize_hostname(machine),
                 )
+        with teuthology.parallel.parallel() as p:
+            for machine in reimage_machines:
+                p.spawn(teuthology.provision.reimage, ctx, machine)
+        for machine in updatekeys_machines:
+            keys.do_update_keys([machine])
+
     elif ctx.unlock:
         if ctx.owner is None and user is None:
             user = misc.get_user()
diff --git a/teuthology/lock/ops.py b/teuthology/lock/ops.py
index 309ae044c..71581ef4d 100644
--- a/teuthology/lock/ops.py
+++ b/teuthology/lock/ops.py
@@ -4,6 +4,7 @@ import os
 
 import requests
 
+import teuthology.parallel
 import teuthology.provision
 from teuthology import misc
 from teuthology.config import config
@@ -109,9 +110,10 @@ def lock_many(ctx, num, machine_type, user=None, description=None,
                 return ok_machs
             elif machine_type in reimage_types:
                 reimaged = dict()
-                for machine in machines:
-                    teuthology.provision.reimage(ctx, machine)
-                    reimaged[machine] = machines[machine]
+                with teuthology.parallel.parallel() as p:
+                    for machine in machines:
+                        p.spawn(teuthology.provision.reimage, ctx, machine)
+                        reimaged[machine] = machines[machine]
                 reimaged = keys.do_update_keys(reimaged.keys())[1]
                 return reimaged
             return machines
-- 
2.47.3