Home | History | Annotate | Download | only in server 
      1 # Copyright 2010 Google Inc. All Rights Reserved.
      2 
      3 __author__ = 'asharif (at] google.com (Ahmad Sharif)'
      4 
      5 from operator import attrgetter
      6 import copy
      7 import csv
      8 import threading
      9 import os.path
     10 
     11 from automation.common import machine
     12 
     13 DEFAULT_MACHINES_FILE = os.path.join(os.path.dirname(__file__), 'test_pool.csv')
     14 
     15 
     16 class MachineManager(object):
     17   """Container for list of machines one can run jobs on."""
     18 
     19   @classmethod
     20   def FromMachineListFile(cls, filename):
     21     # Read the file and skip header
     22     csv_file = csv.reader(open(filename, 'rb'), delimiter=',', quotechar='"')
     23     csv_file.next()
     24 
     25     return cls([machine.Machine(hostname, label, cpu, int(cores), os, user)
     26                 for hostname, label, cpu, cores, os, user in csv_file])
     27 
     28   def __init__(self, machines):
     29     self._machine_pool = machines
     30     self._lock = threading.RLock()
     31 
     32   def _GetMachine(self, mach_spec):
     33     available_pool = [m for m in self._machine_pool if mach_spec.IsMatch(m)]
     34 
     35     if available_pool:
     36       # find a machine with minimum uses
     37       uses = attrgetter('uses')
     38 
     39       mach = min(available_pool, key=uses)
     40 
     41       if mach_spec.preferred_machines:
     42         preferred_pool = [m
     43                           for m in available_pool
     44                           if m.hostname in mach_spec.preferred_machines]
     45         if preferred_pool:
     46           mach = min(preferred_pool, key=uses)
     47 
     48       mach.Acquire(mach_spec.lock_required)
     49 
     50       return mach
     51 
     52   def GetMachines(self, required_machines):
     53     """Acquire machines for use by a job."""
     54 
     55     with self._lock:
     56       acquired_machines = [self._GetMachine(ms) for ms in required_machines]
     57 
     58       if not all(acquired_machines):
     59         # Roll back acquires
     60         while acquired_machines:
     61           mach = acquired_machines.pop()
     62           if mach:
     63             mach.Release()
     64 
     65       return acquired_machines
     66 
     67   def GetMachineList(self):
     68     with self._lock:
     69       return copy.deepcopy(self._machine_pool)
     70 
     71   def ReturnMachines(self, machines):
     72     with self._lock:
     73       for m in machines:
     74         m.Release()
     75 
     76   def __str__(self):
     77     return str(self._machine_pool)
     78