autotest: Customize bot id in swarming.
This CL adds a feature to current swarming bot manager script:
1. Get the bot_code by bot_id.
2. Specify bot_id when starting bot.
BUG=chromium:828662
TEST=Verify it in cros-skylab-suite-server1.
Change-Id: I14401859447f652d2b6703807495d5069681083d
Reviewed-on: https://chromium-review.googlesource.com/1037903
Commit-Ready: Xixuan Wu <xixuan@chromium.org>
Tested-by: Xixuan Wu <xixuan@chromium.org>
Reviewed-by: Aviv Keshet <akeshet@chromium.org>
diff --git a/site_utils/chromeos_proxy/swarming_bot_manager.py b/site_utils/chromeos_proxy/swarming_bot_manager.py
index ec8b5c6..a34d14c 100755
--- a/site_utils/chromeos_proxy/swarming_bot_manager.py
+++ b/site_utils/chromeos_proxy/swarming_bot_manager.py
@@ -58,6 +58,9 @@
parser.add_argument(
'-v', '--verbose', dest='verbose', action='store_true',
help='Verbose mode')
+ parser.add_argument(
+ '--specify_bot_id', action='store_true',
+ help='Specify bot id in retrieving bot codes & staring bots')
return parser.parse_args(args)
@@ -104,10 +107,11 @@
def tick(afe, bot_manager):
"""One tick for swarming bot manager.
- @param afe: the afe to check server role.
+ @param afe: the afe to check server role. If afe is empty, skip checking.
@param bot_manager: a swarming_bots.BotManager instance.
"""
- if is_server_in_prod(socket.getfqdn(), afe):
+ if ((afe and is_server_in_prod(socket.getfqdn(), afe)) or
+ (not afe)):
bot_manager.check()
@@ -139,7 +143,8 @@
bot_manager = swarming_bots.BotManager(
swarming_bots.parse_range(args.id_range),
args.working_dir,
- args.swarming_proxy)
+ args.swarming_proxy,
+ specify_bot_id=args.specify_bot_id)
is_prod = False
retryable = True
with ts_mon_config.SetupTsMonGlobalState('swarming_bots', indirect=True):
diff --git a/site_utils/chromeos_proxy/swarming_bots.py b/site_utils/chromeos_proxy/swarming_bots.py
index 0ee7ab9..42c2281 100755
--- a/site_utils/chromeos_proxy/swarming_bots.py
+++ b/site_utils/chromeos_proxy/swarming_bots.py
@@ -43,6 +43,7 @@
import re
import shutil
import signal
+import socket
import subprocess
import sys
import threading
@@ -89,6 +90,10 @@
"""Raised when multiple processes are detected for the same bot id."""
+def get_hostname():
+ return socket.getfqdn().split(u'.', 1)[0]
+
+
class SwarmingBot(object):
"""Class represent a swarming bot."""
@@ -101,7 +106,8 @@
BOT_CMD_PATTERN = 'swarming_bot.*zip start_bot'
- def __init__(self, bot_id, parent_dir, swarming_proxy):
+ def __init__(self, bot_id, parent_dir, swarming_proxy,
+ specify_bot_id=False):
"""Initialize.
@param bot_id: An integer.
@@ -112,6 +118,10 @@
@param swarming_proxy: URL to the swarming instance.
"""
self.bot_id = bot_id
+ self.specify_bot_id = specify_bot_id
+ if specify_bot_id:
+ self.bot_id = '%s-%s' % (get_hostname(), str(self.bot_id))
+
self.swarming_proxy = swarming_proxy
self.parent_dir = os.path.abspath(os.path.expanduser(parent_dir))
self.bot_dir = os.path.join(self.parent_dir,
@@ -202,14 +212,22 @@
shutil.rmtree(self.bot_dir)
os.makedirs(self.bot_dir)
dest = os.path.join(self.bot_dir, self.BOT_FILENAME)
+ new_env = dict(os.environ)
logging.debug('[Bot %s] Getting bot code from: %s/bot_code',
self.bot_id, self.swarming_proxy)
- urllib.urlretrieve('%s/bot_code' % self.swarming_proxy, dest)
+ if self.specify_bot_id:
+ url = '%s/bot_code?bot_id=%s' % (self.swarming_proxy, self.bot_id)
+ new_env['SWARMING_BOT_ID'] = self.bot_id
+ else:
+ url = '%s/bot_code' % self.swarming_proxy
+
+ logging.info('Download bot code from %s', url)
+ urllib.urlretrieve(url, dest)
cmd = [sys.executable, self.BOT_FILENAME]
logging.debug('[Bot %s] Calling command: %s', self. bot_id, cmd)
process = subprocess.Popen(
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
- cwd=self.bot_dir)
+ cwd=self.bot_dir, env=new_env)
self.pid = process.pid
self._write_pid()
logging.info('[Bot %s] Created bot (pid: %d)', self.bot_id, self.pid)
@@ -253,7 +271,8 @@
CHECK_BOTS_PATTERN = '{executable} {working_dir}.*{bot_cmd_pattern}'
- def __init__(self, bot_ids, working_dir, swarming_proxy):
+ def __init__(self, bot_ids, working_dir, swarming_proxy,
+ specify_bot_id=False):
"""Initialize.
@param bot_ids: a set of integers.
@@ -263,7 +282,8 @@
"""
self.bot_ids = bot_ids
self.working_dir = os.path.abspath(os.path.expanduser(working_dir))
- self.bots = [SwarmingBot(bid, self.working_dir, swarming_proxy)
+ self.bots = [SwarmingBot(bid, self.working_dir, swarming_proxy,
+ specify_bot_id)
for bid in bot_ids]
def launch(self):
