Prashanth B | 489b91d | 2014-03-15 12:17:16 -0700 | [diff] [blame] | 1 | # Copyright (c) 2014 The Chromium OS Authors. All rights reserved. |
| 2 | # Use of this source code is governed by a BSD-style license that can be |
| 3 | # found in the LICENSE file. |
| 4 | |
| 5 | """This module manages translation between monitor_db and the rdb. """ |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 6 | |
| 7 | import common |
| 8 | from autotest_lib.scheduler import rdb |
Prashanth B | 489b91d | 2014-03-15 12:17:16 -0700 | [diff] [blame] | 9 | from autotest_lib.scheduler import rdb_hosts |
| 10 | from autotest_lib.scheduler import rdb_requests |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 11 | from autotest_lib.server.cros import provision |
| 12 | |
| 13 | |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 14 | # Adapters for scheduler specific objects: Convert job information to a |
| 15 | # format more ameanable to the rdb/rdb request managers. |
| 16 | class JobQueryManager(object): |
| 17 | """A caching query manager for all job related information.""" |
Fang Deng | 52a2393 | 2014-11-20 18:30:22 -0800 | [diff] [blame] | 18 | def __init__(self, queue_entries, suite_min_duts=None): |
| 19 | """Initialize. |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 20 | |
Fang Deng | 52a2393 | 2014-11-20 18:30:22 -0800 | [diff] [blame] | 21 | @param queue_entries: A list of HostQueueEntry objects. |
| 22 | @param suite_min_duts: A dictionary where the key is suite job id, |
| 23 | and the value is the value of 'suite_min_dut' in the suite's |
| 24 | job keyvals. It should cover all the suite jobs which |
| 25 | the jobs (associated with the queue_entries) belong to. |
| 26 | """ |
Prashanth B | f66d51b | 2014-05-06 12:42:25 -0700 | [diff] [blame] | 27 | # TODO(beeps): Break this dependency on the host_query_manager, |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 28 | # crbug.com/336934. |
Prashanth B | f66d51b | 2014-05-06 12:42:25 -0700 | [diff] [blame] | 29 | from autotest_lib.scheduler import query_managers |
| 30 | self.query_manager = query_managers.AFEHostQueryManager() |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 31 | jobs = [queue_entry.job_id for queue_entry in queue_entries] |
Prashanth B | f66d51b | 2014-05-06 12:42:25 -0700 | [diff] [blame] | 32 | self._job_acls = self.query_manager._get_job_acl_groups(jobs) |
| 33 | self._job_deps = self.query_manager._get_job_dependencies(jobs) |
| 34 | self._labels = self.query_manager._get_labels(self._job_deps) |
Fang Deng | 52a2393 | 2014-11-20 18:30:22 -0800 | [diff] [blame] | 35 | self._suite_min_duts = suite_min_duts or {} |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 36 | |
| 37 | |
| 38 | def get_job_info(self, queue_entry): |
| 39 | """Extract job information from a queue_entry/host-scheduler. |
| 40 | |
| 41 | @param queue_entry: The queue_entry for which we need job information. |
| 42 | |
| 43 | @return: A dictionary representing job related information. |
| 44 | """ |
| 45 | job_id = queue_entry.job_id |
Fang Deng | a9bc959 | 2015-01-27 17:09:57 -0800 | [diff] [blame] | 46 | job_deps, job_preferred_deps = [], [] |
| 47 | for dep in self._job_deps.get(job_id, []): |
Dan Shi | e44f9c0 | 2016-02-18 13:25:05 -0800 | [diff] [blame] | 48 | if not provision.is_for_special_action(self._labels[dep].name): |
| 49 | job_deps.append(dep) |
| 50 | elif provision.Provision.acts_on(self._labels[dep].name): |
| 51 | job_preferred_deps.append(dep) |
Fang Deng | a9bc959 | 2015-01-27 17:09:57 -0800 | [diff] [blame] | 52 | |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 53 | job_acls = self._job_acls.get(job_id, []) |
Fang Deng | 52a2393 | 2014-11-20 18:30:22 -0800 | [diff] [blame] | 54 | parent_id = queue_entry.job.parent_job_id |
| 55 | min_duts = self._suite_min_duts.get(parent_id, 0) if parent_id else 0 |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 56 | |
| 57 | return {'deps': job_deps, 'acls': job_acls, |
Fang Deng | a9bc959 | 2015-01-27 17:09:57 -0800 | [diff] [blame] | 58 | 'preferred_deps': job_preferred_deps, |
Prashanth B | 9bc32fa | 2014-02-20 12:58:40 -0800 | [diff] [blame] | 59 | 'host_id': queue_entry.host_id, |
Prashanth B | 2c1a22a | 2014-04-02 17:30:51 -0700 | [diff] [blame] | 60 | 'parent_job_id': queue_entry.job.parent_job_id, |
Fang Deng | 52a2393 | 2014-11-20 18:30:22 -0800 | [diff] [blame] | 61 | 'priority': queue_entry.job.priority, |
| 62 | 'suite_min_duts': min_duts} |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 63 | |
| 64 | |
Fang Deng | 52a2393 | 2014-11-20 18:30:22 -0800 | [diff] [blame] | 65 | def acquire_hosts(queue_entries, suite_min_duts=None): |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 66 | """Acquire hosts for the list of queue_entries. |
| 67 | |
Prashanth B | 489b91d | 2014-03-15 12:17:16 -0700 | [diff] [blame] | 68 | The act of acquisition involves leasing a host from the rdb. |
| 69 | |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 70 | @param queue_entries: A list of queue_entries that need hosts. |
Fang Deng | 52a2393 | 2014-11-20 18:30:22 -0800 | [diff] [blame] | 71 | @param suite_min_duts: A dictionary that maps suite job id to the minimum |
| 72 | number of duts required. |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 73 | |
Prashanth B | 489b91d | 2014-03-15 12:17:16 -0700 | [diff] [blame] | 74 | @yield: An rdb_hosts.RDBClientHostWrapper for each host acquired on behalf |
| 75 | of a queue_entry, or None if a host wasn't found. |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 76 | |
| 77 | @raises RDBException: If something goes wrong making the request. |
| 78 | """ |
Fang Deng | 52a2393 | 2014-11-20 18:30:22 -0800 | [diff] [blame] | 79 | job_query_manager = JobQueryManager(queue_entries, suite_min_duts) |
Prashanth B | 489b91d | 2014-03-15 12:17:16 -0700 | [diff] [blame] | 80 | request_manager = rdb_requests.BaseHostRequestManager( |
| 81 | rdb_requests.AcquireHostRequest, rdb.rdb_host_request_dispatcher) |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 82 | for entry in queue_entries: |
| 83 | request_manager.add_request(**job_query_manager.get_job_info(entry)) |
| 84 | |
| 85 | for host in request_manager.response(): |
Prashanth B | 489b91d | 2014-03-15 12:17:16 -0700 | [diff] [blame] | 86 | yield (rdb_hosts.RDBClientHostWrapper(**host) |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 87 | if host else None) |
| 88 | |
| 89 | |
| 90 | def get_hosts(host_ids): |
| 91 | """Get information about the hosts with ids in host_ids. |
| 92 | |
Prashanth B | 489b91d | 2014-03-15 12:17:16 -0700 | [diff] [blame] | 93 | get_hosts is different from acquire_hosts in that it is completely |
| 94 | oblivious to the leased state of a host. |
| 95 | |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 96 | @param host_ids: A list of host_ids. |
| 97 | |
Prashanth B | 489b91d | 2014-03-15 12:17:16 -0700 | [diff] [blame] | 98 | @return: A list of rdb_hosts.RDBClientHostWrapper objects. |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 99 | |
| 100 | @raises RDBException: If something goes wrong in making the request. |
| 101 | """ |
Prashanth B | 489b91d | 2014-03-15 12:17:16 -0700 | [diff] [blame] | 102 | request_manager = rdb_requests.BaseHostRequestManager( |
| 103 | rdb_requests.HostRequest, rdb.get_hosts) |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 104 | for host_id in host_ids: |
| 105 | request_manager.add_request(host_id=host_id) |
| 106 | |
| 107 | hosts = [] |
| 108 | for host in request_manager.response(): |
Prashanth B | 489b91d | 2014-03-15 12:17:16 -0700 | [diff] [blame] | 109 | hosts.append(rdb_hosts.RDBClientHostWrapper(**host) |
beeps | cc9fc70 | 2013-12-02 12:45:38 -0800 | [diff] [blame] | 110 | if host else None) |
| 111 | return hosts |