testing/legion/run_task.py

   1 #!/usr/bin/env python
   2 # Copyright 2015 The Chromium Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 """The main task entrypoint."""
   7
   8 import argparse
   9 import logging
  10 import socket
  11 import sys
  12 import time
  13
  14 #pylint: disable=relative-import
  15 import common_lib
  16 import rpc_server
  17
  18
  19 def main():
  20   print ' '.join(sys.argv)
  21   common_lib.InitLogging()
  22   logging.info('Task starting')
  23
  24   parser = argparse.ArgumentParser()
  25   parser.add_argument('--otp',
  26                       help='One time token used to authenticate with the host')
  27   parser.add_argument('--controller',
  28                       help='The ip address of the controller machine')
  29   parser.add_argument('--idle-timeout', type=int,
  30                       default=common_lib.DEFAULT_TIMEOUT_SECS,
  31                       help='The idle timeout for the rpc server in seconds')
  32   args, _ = parser.parse_known_args()
  33
  34   logging.info(
  35       'Registering with registration server at %s using OTP "%s"',
  36       args.controller, args.otp)
  37   rpc_server.RpcServer.Connect(args.controller).RegisterTask(args.otp,
  38                                                              common_lib.MY_IP)
  39
  40   server = rpc_server.RpcServer(args.controller, args.idle_timeout)
  41
  42   server.serve_forever()
  43   logging.info('Server shutdown complete')
  44   return 0
  45
  46
  47 if __name__ == '__main__':
  48   sys.exit(main())