xref: /aosp_15_r20/external/toolchain-utils/crosperf/machine_image_manager.py (revision 760c253c1ed00ce9abd48f8546f08516e57485fe)
1# -*- coding: utf-8 -*-
2# Copyright 2015 The ChromiumOS Authors
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""MachineImageManager allocates images to duts."""
7
8
9import functools
10
11
12class MachineImageManager(object):
13    """Management of allocating images to duts.
14
15    * Data structure we have -
16
17      duts_ - list of duts, for each duts, we assume the following 2 properties
18      exist - label_ (the current label the duts_ carries or None, if it has an
19      alien image) and name (a string)
20
21      labels_ - a list of labels, for each label, we assume these properties
22      exist - remote (a set/vector/list of dut names (not dut object), to each
23      of which this image is compatible), remote could be none, which means
24      universal compatible.
25
26      label_duts_ - for each label, we maintain a list of duts, onto which the
27      label is imaged. Note this is an array of lists. Each element of each list
28      is an integer which is dut oridnal. We access this array using label
29      ordinal.
30
31      allocate_log_ - a list of allocation record. For example, if we allocate
32      l1 to d1, then l2 to d2, then allocate_log_ would be [(1, 1), (2, 2)].
33      This is used for debug/log, etc. All tuples in the list are integer pairs
34      (label_ordinal, dut_ordinal).
35
36      n_duts_ - number of duts.
37
38      n_labels_ - number of labels.
39
40      dut_name_ordinal_ - mapping from dut name (a string) to an integer,
41      starting from 0. So that duts_[dut_name_ordinal_[a_dut.name]]= a_dut.
42
43    * Problem abstraction -
44
45      Assume we have the following matrix - label X machine (row X col). A 'X'
46      in (i, j) in the matrix means machine and lable is not compatible, or that
47      we cannot image li to Mj.
48
49           M1   M2   M3
50      L1    X
51
52      L2             X
53
54      L3    X    X
55
56      Now that we'll try to find a way to fill Ys in the matrix so that -
57
58        a) - each row at least get a Y, this ensures that each label get imaged
59        at least once, an apparent prerequiste.
60
61        b) - each column get at most N Ys. This make sure we can successfully
62        finish all tests by re-image each machine at most N times. That being
63        said, we could *OPTIONALLY* reimage some machines more than N times to
64        *accelerate* the test speed.
65
66      How to choose initial N for b) -
67      If number of duts (nd) is equal to or more than that of labels (nl), we
68      start from N == 1. Else we start from N = nl - nd + 1.
69
70      We will begin the search with pre-defined N, if we fail to find such a
71      solution for such N, we increase N by 1 and continue the search till we
72      get N == nl, at this case we fails.
73
74      Such a solution ensures minimal number of reimages.
75
76    * Solution representation
77
78      The solution will be placed inside the matrix, like below
79
80           M1   M2   M3   M4
81      L1    X    X    Y
82
83      L2    Y         X
84
85      L3    X    Y    X
86
87    * Allocation algorithm
88
89      When Mj asks for a image, we check column j, pick the first cell that
90      contains a 'Y', and mark the cell '_'. If no such 'Y' exists (like M4 in
91      the above solution matrix), we just pick an image that the minimal reimage
92      number.
93
94      After allocate for M3
95           M1   M2   M3   M4
96      L1    X    X    _
97
98      L2    Y         X
99
100      L3    X    Y    X
101
102      After allocate for M4
103           M1   M2   M3   M4
104      L1    X    X    _
105
106      L2    Y         X    _
107
108      L3    X    Y    X
109
110      After allocate for M2
111           M1   M2   M3   M4
112      L1    X    X    _
113
114      L2    Y         X    _
115
116      L3    X    _    X
117
118      After allocate for M1
119           M1   M2   M3   M4
120      L1    X    X    _
121
122      L2    _         X    _
123
124      L3    X    _    X
125
126      After allocate for M2
127           M1   M2   M3   M4
128      L1    X    X    _
129
130      L2    _    _    X    _
131
132      L3    X    _    X
133
134      If we try to allocate for M1 or M2 or M3 again, we get None.
135
136    * Special / common case to handle seperately
137
138      We have only 1 dut or if we have only 1 label, that's simple enough.
139    """
140
141    def __init__(self, labels, duts):
142        self.labels_ = labels
143        self.duts_ = duts
144        self.n_labels_ = len(labels)
145        self.n_duts_ = len(duts)
146        self.dut_name_ordinal_ = dict()
147        for idx, dut in enumerate(self.duts_):
148            self.dut_name_ordinal_[dut.name] = idx
149
150        # Generate initial matrix containg 'X' or ' '.
151        self.matrix_ = [
152            ["X" if l.remote else " " for _ in range(self.n_duts_)]
153            for l in self.labels_
154        ]
155        for ol, l in enumerate(self.labels_):
156            if l.remote:
157                for r in l.remote:
158                    self.matrix_[ol][self.dut_name_ordinal_[r]] = " "
159
160        self.label_duts_ = [[] for _ in range(self.n_labels_)]
161        self.allocate_log_ = []
162
163    def compute_initial_allocation(self):
164        """Compute the initial label-dut allocation.
165
166        This method finds the most efficient way that every label gets imaged at
167        least once.
168
169        Returns:
170          False, only if not all labels could be imaged to a certain machine,
171          otherwise True.
172        """
173
174        if self.n_duts_ == 1:
175            for i, v in self.matrix_vertical_generator(0):
176                if v != "X":
177                    self.matrix_[i][0] = "Y"
178            return
179
180        if self.n_labels_ == 1:
181            for j, v in self.matrix_horizontal_generator(0):
182                if v != "X":
183                    self.matrix_[0][j] = "Y"
184            return
185
186        if self.n_duts_ >= self.n_labels_:
187            n = 1
188        else:
189            n = self.n_labels_ - self.n_duts_ + 1
190        while n <= self.n_labels_:
191            if self._compute_initial_allocation_internal(0, n):
192                break
193            n += 1
194
195        return n <= self.n_labels_
196
197    def _record_allocate_log(self, label_i, dut_j):
198        self.allocate_log_.append((label_i, dut_j))
199        self.label_duts_[label_i].append(dut_j)
200
201    def allocate(self, dut, schedv2=None):
202        """Allocate a label for dut.
203
204        Args:
205          dut: the dut that asks for a new image.
206          schedv2: the scheduling instance, we need the benchmark run
207                   information with schedv2 for a better allocation.
208
209        Returns:
210          a label to image onto the dut or None if no more available images for
211          the dut.
212        """
213        j = self.dut_name_ordinal_[dut.name]
214        # 'can_' prefix means candidate label's.
215        can_reimage_number = 999
216        can_i = 999
217        can_label = None
218        can_pending_br_num = 0
219        for i, v in self.matrix_vertical_generator(j):
220            label = self.labels_[i]
221
222            # 2 optimizations here regarding allocating label to dut.
223            # Note schedv2 might be None in case we do not need this
224            # optimization or we are in testing mode.
225            if schedv2 is not None:
226                pending_br_num = len(schedv2.get_label_map()[label])
227                if pending_br_num == 0:
228                    # (A) - we have finished all br of this label,
229                    # apparently, we do not want to reimaeg dut to
230                    # this label.
231                    continue
232            else:
233                # In case we do not have a schedv2 instance, mark
234                # pending_br_num as 0, so pending_br_num >=
235                # can_pending_br_num is always True.
236                pending_br_num = 0
237
238            # For this time being, I just comment this out until we have a
239            # better estimation how long each benchmarkrun takes.
240            # if (pending_br_num <= 5 and
241            #     len(self.label_duts_[i]) >= 1):
242            #     # (B) this is heuristic - if there are just a few test cases
243            #     # (say <5) left undone for this label, and there is at least
244            #     # 1 other machine working on this lable, we probably not want
245            #     # to bother to reimage this dut to help with these 5 test
246            #     # cases
247            #     continue
248
249            if v == "Y":
250                self.matrix_[i][j] = "_"
251                self._record_allocate_log(i, j)
252                return label
253            if v == " ":
254                label_reimage_number = len(self.label_duts_[i])
255                if (can_label is None) or (
256                    label_reimage_number < can_reimage_number
257                    or (
258                        label_reimage_number == can_reimage_number
259                        and pending_br_num >= can_pending_br_num
260                    )
261                ):
262                    can_reimage_number = label_reimage_number
263                    can_i = i
264                    can_label = label
265                    can_pending_br_num = pending_br_num
266
267        # All labels are marked either '_' (already taken) or 'X' (not
268        # compatible), so return None to notify machine thread to quit.
269        if can_label is None:
270            return None
271
272        # At this point, we don't find any 'Y' for the machine, so we go the
273        # 'min' approach.
274        self.matrix_[can_i][j] = "_"
275        self._record_allocate_log(can_i, j)
276        return can_label
277
278    def matrix_vertical_generator(self, col):
279        """Iterate matrix vertically at column 'col'.
280
281        Yield row number i and value at matrix_[i][col].
282        """
283        for i, _ in enumerate(self.labels_):
284            yield i, self.matrix_[i][col]
285
286    def matrix_horizontal_generator(self, row):
287        """Iterate matrix horizontally at row 'row'.
288
289        Yield col number j and value at matrix_[row][j].
290        """
291        for j, _ in enumerate(self.duts_):
292            yield j, self.matrix_[row][j]
293
294    def _compute_initial_allocation_internal(self, level, N):
295        """Search matrix for d with N."""
296
297        if level == self.n_labels_:
298            return True
299
300        for j, v in self.matrix_horizontal_generator(level):
301            if v == " ":
302                # Before we put a 'Y', we check how many Y column 'j' has.
303                # Note y[0] is row idx, y[1] is the cell value.
304                ny = functools.reduce(
305                    lambda x, y: x + 1 if (y[1] == "Y") else x,
306                    self.matrix_vertical_generator(j),
307                    0,
308                )
309                if ny < N:
310                    self.matrix_[level][j] = "Y"
311                    if self._compute_initial_allocation_internal(level + 1, N):
312                        return True
313                    self.matrix_[level][j] = " "
314
315        return False
316