File: clobber_cache_utils.py

package info (click to toggle)
chromium 139.0.7258.127-1
  • links: PTS, VCS
  • area: main
  • in suites:
  • size: 6,122,068 kB
  • sloc: cpp: 35,100,771; ansic: 7,163,530; javascript: 4,103,002; python: 1,436,920; asm: 946,517; xml: 746,709; pascal: 187,653; perl: 88,691; sh: 88,436; objc: 79,953; sql: 51,488; cs: 44,583; fortran: 24,137; makefile: 22,147; tcl: 15,277; php: 13,980; yacc: 8,984; ruby: 7,485; awk: 3,720; lisp: 3,096; lex: 1,327; ada: 727; jsp: 228; sed: 36
file content (177 lines) | stat: -rw-r--r-- 5,763 bytes parent folder | download | duplicates (5)
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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
# Copyright 2019 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Utility functions for cache clobbering scripts."""

from __future__ import print_function

import json
import os
import subprocess
import textwrap

_SRC_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
_SWARMING_CLIENT = os.path.join(_SRC_ROOT, 'tools', 'luci-go', 'swarming')
_SWARMING_SERVER = 'chromium-swarm.appspot.com'


def get_xcode_caches_for_all_bots(bots_state):
  """Extracts Xcode cache names for each bot from a list of bot states.

  Args:
    bots_state: A list of bot state dictionaries.

  Returns:
    A dict where keys are bot IDs and values are lists of Xcode cache names.
  """
  xcode_caches_dict = {}
  for bot_state in bots_state:
    xcode_caches = []

    for dimension in bot_state['dimensions']:
      if dimension['key'] == 'caches':
        for cache in dimension['value']:
          if cache.startswith('xcode'):
            xcode_caches.append(cache)
    if xcode_caches:
      xcode_caches_dict[bot_state['bot_id']] = xcode_caches
  return xcode_caches_dict


def get_bots_state_by_dimensions(swarming_server, dimensions):
  """Gets the states of bots matching given dimensions."""
  cmd = [
      _SWARMING_CLIENT,
      'bots',
      '-S',
      swarming_server,
  ]
  for d in dimensions:
    cmd.extend(['-dimension', d])
  try:
    return json.loads(subprocess.check_output(cmd))
  except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
    print(f"Error getting bot states: {e}")
    return []


def trigger_clobber_cache(swarming_server, pool, realm, cache, bot_id,
                          mount_rel_path, dry_run):
  """Clobber a specific cache on a given bot.

  Args:
    swarming_server: The swarming_server instance to lookup bots to clobber
        caches on.
    pool: The pool of machines to lookup bots to clobber caches on.
    realm: The realm to trigger tasks into.
    cache: The name of the cache to clobber.
    bot_id: the id of the bot that you wish to clobber.
    mount_rel_path: The relative path to mount the cache to when clobbering.
    dry_run: Whether a dry-run should be performed where the commands that
        would be executed to trigger the clobber task are printed rather than
        actually triggering the clobber task.
  """
  cmd = [
      _SWARMING_CLIENT,
      'trigger',
      '-S',
      swarming_server,
      '-realm',
      realm,
      '-dimension',
      'pool=' + pool,
      '-dimension',
      'id=' + bot_id,
      '-cipd-package',
      'cpython3:infra/3pp/tools/cpython3/${platform}=latest',
      '-named-cache',
      cache + '=' + mount_rel_path,
      '-priority',
      '10',
      '--',
      'cpython3/bin/python3${EXECUTABLE_SUFFIX}',
      '-c',
      textwrap.dedent('''\
          import os, shutil, stat

          def remove_readonly(func, path, _):
              "Clear the readonly bit and reattempt the removal"
              os.chmod(path, stat.S_IWRITE)
              func(path)

          shutil.rmtree({mount_rel_path!r}, onerror=remove_readonly)
          '''.format(mount_rel_path=mount_rel_path)),
  ]
  if dry_run:
    print('Would run `%s`' % ' '.join(cmd))
  else:
    subprocess.check_call(cmd)


def add_common_args(argument_parser):
  """Add common arguments to the argument parser used for cache clobber scripts.

  The following arguments will be added to the argument parser:
    * swarming_server (-S/--swarming-server) - The swarming server instance to
      lookup bots to clobber caches on, with a default of the
      chromium-swarm.appspot.com.
    * dry_run (-n/--dry-run) - Whether a dry-run should be performed rather than
      actually clobbering caches, defaults to False.
  """
  argument_parser.add_argument(
      '-S', '--swarming-server', default=_SWARMING_SERVER)
  argument_parser.add_argument('-n', '--dry-run', action='store_true')


def confirm_and_trigger_clobber_bots(swarming_server,
                                     pool,
                                     realm,
                                     cache,
                                     mount_rel_path,
                                     dry_run,
                                     bot_id=None):
  """Gets bot IDs, confirms with the user, handles dry-run and removes caches.

  Args:
      swarming_server: The Swarming server URL.
      pool: The Swarming pool.
      realm: The Swarming realm.
      cache: The name of the cache to clobber.
      mount_rel_path - The relative path to mount the cache to when clobbering.
      dry_run:  If True, don't actually clobber or prompt.
      bot_id: Optional, the id of a specific bot that you wish to clobber.

  Returns:
      A list of botsto clobber, or an empty list if no bots
      should be clobbered (either none were found, or the user cancelled).
  """

  if bot_id:
    bot_ids = [bot_id]  # Use explicitly provided bot IDs.
    bots_state = get_bots_state_by_dimensions(swarming_server, ['id=' + bot_id])
  else:
    bots_state = get_bots_state_by_dimensions(
        swarming_server, ['pool=' + pool, 'caches=' + cache])
    bot_ids = [bot['bot_id'] for bot in bots_state]

  if not bot_ids:
    print(f"No bots found in pool {pool} with cache {cache}.")
    return []

  print(f"The following bots with {cache} will be clobbered:")
  for b_id in bot_ids:
    print(f"  {b_id}")
  print()

  if not dry_run:
    val = input('Proceed? [Y/n] ')
    if val and not val.lower().startswith('y'):
      print('Cancelled.')
      return []

  for b_id in bot_ids:
    trigger_clobber_cache(swarming_server, pool, realm, cache, b_id,
                          mount_rel_path, dry_run)

  return bots_state