balance_pool: Add sku argument

BUG=chromium:789766
TEST=dry_run the tool

balance_pool -v -n --sku reef_intel_celeron_n3350_4Gb -a 1 -s crosperf
suites reef

atest label remove -m chromeos6-row3-rack12-host9 pool:crosperf
atest label add -m chromeos6-row3-rack12-host9 pool:suites

ERROR:root:Caught exception while flushing: Metrics were sent before
initializing the global Monitor.

Change-Id: Ic47f32d1345b295db1ce0c5ea67d26fc1c71e674
Signed-off-by: Chung-yih Wang <cywang@google.com>
Reviewed-on: https://chromium-review.googlesource.com/799652
Commit-Ready: Chung-yih Wang <cywang@chromium.org>
Tested-by: Chung-yih Wang <cywang@chromium.org>
Reviewed-by: Prathmesh Prabhu <pprabhu@chromium.org>
diff --git a/site_utils/balance_pools.py b/site_utils/balance_pools.py
index bf2787b..709d248 100755
--- a/site_utils/balance_pools.py
+++ b/site_utils/balance_pools.py
@@ -31,6 +31,7 @@
   -s POOL, --spare POOL
                         Pool from which to draw replacement spares (default:
                         pool:suites)
+  --sku SKU             The specific SKU we intend to swap with
   -n, --dry-run         Report actions to take in the form of shell commands
 
 
@@ -406,8 +407,9 @@
                   target_total, add_msg)
 
         _log_info(dry_run,
-                  '%s %s pool has %d spares available.',
-                  labels, main_pool.pool, len(spare_pool.working_hosts))
+                  '%s %s pool has %d spares available for balancing pool %s',
+                  labels, spare_pool.pool, len(spare_pool.working_hosts),
+                  main_pool.pool)
 
         if spares_needed > len(spare_duts):
             _log_error('Not enough spares: need %d, only have %d.',
@@ -604,6 +606,9 @@
         help='Names of boards or models to balance. (See also: --as-model)',
     )
 
+    parser.add_argument('--sku', type=str,
+                        help='Optional name of sku to restrict to.')
+
     arguments = parser.parse_args(argv[1:])
 
     # Error-check arguments.
@@ -662,6 +667,8 @@
                     labels['model'] = board_or_model
                 else:
                     labels['board'] = board_or_model
+            if arguments.sku:
+                labels['sku'] = arguments.sku
             balancer_targets.append((pool, labels.getlabels()))
     return balancer_targets