Mercurial > code > home > repos > infra

annotate pi-setup/runner.py @ 332:d4893670f888 default tip

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
WIP: use watchdog reboot timer on pi
author drewp@bigasterisk.com
date Thu, 27 Feb 2025 11:09:29 -0800
parents 1cb4aeec8fc6
children
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
279
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
1 import asyncio
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
2 import itertools
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
3 import logging
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
4 import shlex
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
5 import time
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
6 from contextlib import asynccontextmanager, contextmanager
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
7 from pathlib import Path
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
8 from typing import Any, AsyncGenerator, Generator, Sequence, cast
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
9
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
10 import more_itertools
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
11 import psutil
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
12
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
13 log = logging.getLogger()
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
14
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
15 # This is the type of arg we pass to create_subprocess_exec.
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
16 RunArgType = str | Path
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
17
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
18 # This is what you can call run or get_output with, passing sublists of args for
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
19 # clarity.
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
20 ArgType = str | Path | Sequence[str | Path]
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
21
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
22
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
23 def _flatten_run_args(args: tuple[ArgType, ...]) -> tuple[RunArgType]:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
24 return tuple(more_itertools.collapse(args))
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
25
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
26
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
27 async def run(program: str, *args: ArgType, _stdin: str | None = None, _stdout: int | None = None):
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
28 run_args = _flatten_run_args(args)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
29 log.info(f'Running {program} {shlex.join(map(str,run_args))}')
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
30
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
31 proc = await asyncio.create_subprocess_exec(
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
32 program,
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
33 *run_args,
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
34 stdin=(asyncio.subprocess.PIPE if _stdin is not None else None),
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
35 stdout=_stdout,
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
36 )
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
37
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
38 async def log_busy_cpu(pid, secs=3):
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
39 pr = psutil.Process(pid)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
40 while True:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
41 pct = pr.cpu_percent()
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
42 if pct > 5:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
43 bar = '=' * int(pct / 400 * 50)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
44 log.info(f"{program} cpu {pct:5.1f} {bar}")
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
45 await asyncio.sleep(secs)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
46
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
47 busy = asyncio.create_task(log_busy_cpu(proc.pid))
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
48
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
49 out, _ = await proc.communicate(_stdin.encode() if _stdin is not None else None)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
50 busy.cancel()
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
51 if proc.returncode != 0:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
52 raise ValueError(f'{program} returned {proc.returncode}')
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
53 return out.decode() if _stdout is not None else None
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
54
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
55
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
56 async def get_output(program: str, *args: ArgType, stdin: None | str = None) -> str:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
57 out = await run(program, *args, _stdin=stdin, _stdout=asyncio.subprocess.PIPE)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
58 log.info(f" -> returned {out!r}")
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
59 return cast(str, out)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
60
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
61
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
62 _mount_count = itertools.count()
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
63
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
64
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
65 @contextmanager
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
66 def _new_mount_point(dir: Path) -> Generator[Path, Any, Any]:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
67 p = dir / f'mount{next(_mount_count)}'
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
68 p.mkdir()
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
69 try:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
70 yield p
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
71 finally:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
72 p.rmdir()
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
73
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
74
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
75 @asynccontextmanager
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
76 async def mount(work_dir: Path, src: Path, src_offset: int) -> AsyncGenerator[Path, Any]:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
77 with _new_mount_point(work_dir) as mount_point:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
78 args = []
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
79 if not str(src).startswith('/dev/'):
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
80 args = ['-o', f'loop,offset={src_offset}']
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
81 await run('mount', args, src, mount_point)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
82 try:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
83 yield mount_point
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
84 finally:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
85 try:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
86 await run('umount', mount_point)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
87 except Exception:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
88 time.sleep(.5)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
89 await run('umount', mount_point)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
90
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
91
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
92 @asynccontextmanager
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
93 async def sshfs(work_dir: Path, ssh_path: str) -> AsyncGenerator[Path, Any]:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
94 with _new_mount_point(work_dir) as mount_point:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
95 await run('sshfs', ssh_path, mount_point)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
96 try:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
97 yield mount_point
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
98 finally:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
99 await run('umount', mount_point)
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
100
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
101
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
102 @asynccontextmanager
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
103 async def iscsi_login(*args):
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
104 await run('iscsiadm', *(list(args) + ['--login']))
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
105 try:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
106 yield
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
107 finally:
1cb4aeec8fc6 pi_setup code to prepare a pi for netboot
drewp@bigasterisk.com
parents:
diff changeset
108 await run('iscsiadm', *(list(args) + ['--logout']))