aboutsummaryrefslogtreecommitdiff
path: root/automation/server/machine_manager.py
blob: b7186077ad0c519678f56bfc986e5ad300ae736b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
# Copyright 2010 Google Inc. All Rights Reserved.

__author__ = 'asharif@google.com (Ahmad Sharif)'

from operator import attrgetter
import copy
import csv
import threading
import os.path

from automation.common import machine

DEFAULT_MACHINES_FILE = os.path.join(os.path.dirname(__file__), 'test_pool.csv')


class MachineManager(object):
  """Container for list of machines one can run jobs on."""

  @classmethod
  def FromMachineListFile(cls, filename):
    # Read the file and skip header
    csv_file = csv.reader(open(filename, 'rb'), delimiter=',', quotechar='"')
    csv_file.next()

    return cls([machine.Machine(hostname, label, cpu, int(cores), os, user)
                for hostname, label, cpu, cores, os, user in csv_file])

  def __init__(self, machines):
    self._machine_pool = machines
    self._lock = threading.RLock()

  def _GetMachine(self, mach_spec):
    available_pool = [m for m in self._machine_pool if mach_spec.IsMatch(m)]

    if available_pool:
      # find a machine with minimum uses
      uses = attrgetter('uses')

      mach = min(available_pool, key=uses)

      if mach_spec.preferred_machines:
        preferred_pool = [m
                          for m in available_pool
                          if m.hostname in mach_spec.preferred_machines]
        if preferred_pool:
          mach = min(preferred_pool, key=uses)

      mach.Acquire(mach_spec.lock_required)

      return mach

  def GetMachines(self, required_machines):
    """Acquire machines for use by a job."""

    with self._lock:
      acquired_machines = [self._GetMachine(ms) for ms in required_machines]

      if not all(acquired_machines):
        # Roll back acquires
        while acquired_machines:
          mach = acquired_machines.pop()
          if mach:
            mach.Release()

      return acquired_machines

  def GetMachineList(self):
    with self._lock:
      return copy.deepcopy(self._machine_pool)

  def ReturnMachines(self, machines):
    with self._lock:
      for m in machines:
        m.Release()

  def __str__(self):
    return str(self._machine_pool)