diff --git a/bazel/labgrid/runner/BUILD.bazel b/bazel/labgrid/runner/BUILD.bazel index 5c0e712128bb35a0d1603121fd5bd95e49834190..98258a4065d57821ceeb5084ac0ac2f11709828b 100644 --- a/bazel/labgrid/runner/BUILD.bazel +++ b/bazel/labgrid/runner/BUILD.bazel @@ -22,12 +22,12 @@ py_library( srcs = [ "__init__.py", "runner.py", + "tools.py", ], data = [":{}".format(tool) for tool in TOOLS], visibility = ["//visibility:public"], deps = [ "//bazel/labgrid/strategy", - "//bazel/labgrid/util:target", "//bazel/python/runfiles", "//labgrid:pkg", "//labgrid/config:deps", diff --git a/bazel/labgrid/runner/runner.py b/bazel/labgrid/runner/runner.py index b92b93a1f34ae4c53965945d6fdefbc695a22e09..e124c0196c39ac45d11bcf09cfa3ae509ae17128 100644 --- a/bazel/labgrid/runner/runner.py +++ b/bazel/labgrid/runner/runner.py @@ -10,22 +10,16 @@ from typing import ( Iterator, Mapping, ) +from threading import get_native_id +from uuid import getnode +from subprocess import CalledProcessError from bazel.labgrid.strategy import State, transition -from bazel.labgrid.util.target import TemporaryDirectory -from bazel.python.runfiles import runfile from labgrid import Environment from labgrid.driver.exception import ExecutionError from labgrid.logging import StepLogger, basicConfig - -@dataclass -class _Tools: - mktemp: Path = runfile("ape/ape/assimilate/mktemp/mktemp") - rm: Path = runfile("ape/ape/assimilate/rm/rm") - mkdir: Path = runfile("ape/ape/assimilate/mkdir/mkdir") - env: Path = runfile("ape/ape/assimilate/env/env") - mv: Path = runfile("ape/ape/assimilate/mv/mv") +from .tools import Tools, Tool @dataclass(frozen=True) @@ -57,7 +51,6 @@ def _log_level(value): def runner(): config = PurePath(environ["LG_ENV"]) log_level = _log_level(environ.get("LOG_LEVEL", "ERROR")) - tools = _Tools() # Activating the StepLogger requires a DEBUG level of verbosity basicConfig(level=log_level) @@ -70,60 +63,50 @@ def runner(): # Retrieve the communication protocols shell = target.get_driver("CommandProtocol") transfer = target.get_driver("FileTransferProtocol") - with TemporaryDirectory( - shell, - transfer, - tools.mktemp, - tools.rm, - tools.mkdir, - tools.mv, - ) as temp: - r = Runner(temp, transfer, shell) - - # Transfer 'env' binary over to the device - r.put([FileTransfer(r.exec_path / tools.env.name, tools.env)]) - + with Runner(transfer, shell) as r: yield r class Runner: - def __init__(self, exec_path, transfer, shell): - self.exec_path = exec_path - self.transfer = transfer - self.shell = shell - self.runfiles_dir = environ.get("RUNFILES_DIR") + def __init__(self, transfer, shell): + self._transfer = transfer + self._shell = shell + self._exec_root = None + self._runfiles_dir = environ.get("RUNFILES_DIR") + self._paths_to_remove_on_exit = [] + self._tools = None + + @property + def exec_root(self): + return self._exec_root def put(self, uploads: Iterator[FileTransfer]): for upload in uploads: remote, local = self._resolve(upload) - self.exec_path.mkdir(remote.parent) - self.transfer.put(local, remote) + self._run(f"{self._tools.mkdir(remote.parent)}") + self._transfer.put(local, remote) if upload.include_runfiles: src_runfiles = local.with_suffix(".runfiles") dest_runfiles = remote.with_suffix(".runfiles") if Path(src_runfiles).exists(): - self.transfer.put(f"{src_runfiles}", f"{dest_runfiles}") - elif self.runfiles_dir is not None: - self.transfer.put(f"{self.runfiles_dir}", f"{dest_runfiles}") + self._transfer.put(f"{src_runfiles}", f"{dest_runfiles}") + elif self._runfiles_dir is not None: + self._transfer.put(f"{self._runfiles_dir}", f"{dest_runfiles}") def get(self, downloads: Iterator[FileTransfer], code: int): for download in downloads: remote, local = self._resolve(download) # FIXME: Check if file exists with `@ape//ape:test` instead of assuming any error is due to a missing file try: - self.transfer.get(remote, local) + self._transfer.get(remote, local) except ExecutionError as e: if not download.optional and code == 0: raise e def run(self, cmd: str, env: Mapping[str, str]) -> int: - # Construct the CLI using env, the env vars and then the program and args. - name_values = " ".join(f"{k}={v}" for k, v in env.items()) - cmd = f"cd {self.exec_path} && {self.exec_path / 'env'} {name_values} {cmd}" - - # Run the transferred program - out, err, code = self.shell.run(cmd) + cmd = f"cd {self._exec_root} && {self._tools.env(env, cmd)}" + out, err, code = self._shell.run(cmd) for line in out: stdout.write(f"{line}{linesep}") for line in err: @@ -131,8 +114,30 @@ class Runner: return code + def __enter__(self): + # Get a unique temporary directory, avoiding concurrent instances + mktemp = Tool("mktemp", PurePath(f".{getnode()}.{get_native_id()}")) + self._transfer.put(mktemp.local, mktemp.remote) + + tools_root = self._read_path(mktemp.cmd("-d", "--suffix=-TOOLS")) + self._tools = Tools(tools_root) + for tool in self._tools: + self._transfer.put(tool.local, tool.remote) + + self._exec_root = self._read_path(mktemp.cmd("-d", "--suffix=-EXECROOT")) + + self._paths_to_remove_on_exit.extend( + [self._exec_root, mktemp.remote, tools_root] + ) + + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + for path in self._paths_to_remove_on_exit: + self._run(f"{self._tools.rm(path)}") + def _resolve(self, transfer: FileTransfer) -> tuple[str, str]: - remote = self._join(self.exec_path, transfer.remote) + remote = self._join(self._exec_root, transfer.remote) local = self._join(Path.cwd(), transfer.local) return (remote, local) @@ -141,3 +146,16 @@ class Runner: if not path.is_absolute(): return base / path return path + + def _run(self, cmd, check=True): + out, err, code = self._shell.run(cmd) + if check and code: + raise CalledProcessError( + cmd=cmd, returncode=code, output=stdout, stderr=err + ) + + return out, err, code + + def _read_path(self, cmd): + out, _, _ = self._run(cmd) + return PurePath(linesep.join(out).rstrip()) diff --git a/bazel/labgrid/runner/tools.py b/bazel/labgrid/runner/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..21346bcef50a0582f08678e60d7b0b5b9e7a0247 --- /dev/null +++ b/bazel/labgrid/runner/tools.py @@ -0,0 +1,45 @@ +from pathlib import PurePath +from typing import Mapping, Iterable, Generator + +from bazel.python.runfiles import runfile + + +class Tool: + def __init__(self, name: str, remote: PurePath): + self._local = runfile(f"ape/ape/assimilate/{name}/{name}") + self._remote = remote + + @property + def local(self) -> PurePath: + return self._local + + @property + def remote(self) -> PurePath: + return self._remote + + def cmd(self, *args: Iterable[str]) -> str: + exec = self._remote + if not self._remote.is_absolute(): + exec = f"./{exec}" + return f"{exec} {' '.join(args)}" + + +class Tools: + _tool_names = ["rm", "mkdir", "env"] + + def __init__(self, tools_root: PurePath): + self._tools = {name: Tool(name, tools_root / name) for name in self._tool_names} + + def mkdir(self, path: PurePath) -> str: + return self._tools["mkdir"].cmd("-p", str(path)) + + def rm(self, path: PurePath) -> str: + return self._tools["rm"].cmd("-r", str(path)) + + def env(self, env: Mapping[str, str], cmd: str) -> str: + name_values = [f"{k}={v}" for k, v in env.items()] + return self._tools["env"].cmd(*name_values, cmd) + + def __iter__(self) -> Generator[Tool, None, None]: + for tool in self._tools.values(): + yield tool diff --git a/bazel/labgrid/util/BUILD.bazel b/bazel/labgrid/util/BUILD.bazel deleted file mode 100644 index 70cbb96fa2fc3653cba221dcaad2fa17f44e30d3..0000000000000000000000000000000000000000 --- a/bazel/labgrid/util/BUILD.bazel +++ /dev/null @@ -1,7 +0,0 @@ -load("@rules_python//python:defs.bzl", "py_library") - -py_library( - name = "target", - srcs = ["target.py"], - visibility = ["//visibility:public"], -) diff --git a/bazel/labgrid/util/target.py b/bazel/labgrid/util/target.py deleted file mode 100644 index 850efa25f6999bca3b639ba9876ef69e5dd04c18..0000000000000000000000000000000000000000 --- a/bazel/labgrid/util/target.py +++ /dev/null @@ -1,119 +0,0 @@ -from contextlib import contextmanager -from os import linesep -from pathlib import Path, PurePath -from shlex import join -from subprocess import CalledProcessError, CompletedProcess -from threading import get_native_id -from typing import Iterable, Iterator -from uuid import getnode - -from labgrid.protocol.commandprotocol import CommandProtocol -from labgrid.protocol.filetransferprotocol import FileTransferProtocol - - -class TargetCalledProcessError(CalledProcessError): - pass - - -class Subprocess: - def __init__(self, shell: CommandProtocol): - self.__shell = shell - - @property - def shell(self) -> CommandProtocol: - return self.__shell - - def run( - self, args: Iterable[str], *, check: bool = False, text: bool = False - ) -> CompletedProcess: - assert text, "Non-text target execution is not supported" - assert not isinstance(args, str), "`run` accepts iterables but not strings" - joined = join(args) - out, err, code = self.shell.run(joined) - stdout = linesep.join(out) - stderr = linesep.join(err) - if check and code: - raise TargetCalledProcessError( - cmd=joined, returncode=code, output=stdout, stderr=stderr - ) - return CompletedProcess( - returncode=code, stdout=stdout, stderr=stderr, args=args - ) - - -class _TemporaryDirectory: - """Acts like a PurePath with utility functions""" - - def __init__(self, target, root, execroot, tools): - self._target = target - self._root = root - self._execroot = execroot - self._tools = tools - - def mkdir(self, path): - self._target.run( - (f"{self._tools / 'mkdir'}", "-p", f"{path}"), check=True, text=True - ) - - def __getattr__(self, attr): - """Exposes attrs of execroot""" - try: - return getattr(self, attr) - except AttributeError: - pass - - try: - return getattr(self._execroot, attr) - except AttributeError: - raise AttributeError(f"'{type(self)}' object has no attribute '{attr}'") - - def __truediv__(self, path): - return self._execroot / path - - def __str__(self): - return str(self._execroot) - - def remove(self): - self._target.run( - (f"{self._tools / 'rm'}", "-rf", f"{self._root}"), check=True, text=True - ) - - -@contextmanager -def TemporaryDirectory( - shell: CommandProtocol, - transfer: FileTransferProtocol, - mktemp: Path, - rm: Path, - mkdir: Path, - mv: Path, -) -> Iterator[PurePath]: - # Get a unique temporary directory, avoiding concurrent `TemporaryDirectory` managers - unique = f".{getnode()}.{get_native_id()}" - transfer.put(f"{mktemp}", unique) - target = Subprocess(shell) - result = target.run((f"./{unique}", "-d"), check=True, text=True) - root = PurePath(result.stdout.rstrip()) - - # Transfer tools and make execroot and tools dir - transfer.put(f"{mkdir}", f"{root / 'mkdir'}") - - execroot = root / "execroot" - tools = root / "tools" - target.run((f"{root / 'mkdir'}", f"{execroot}"), check=True, text=True) - target.run((f"{root / 'mkdir'}", f"{tools}"), check=True, text=True) - - transfer.put(f"{mv}", f"{tools / 'mv'}") - transfer.put(f"{rm}", f"{tools / 'rm'}") - target.run( - (f"{tools / 'mv'}", f"{root / 'mkdir'}", f"{tools / 'mkdir'}"), - check=True, - text=True, - ) - - temp = _TemporaryDirectory(target, root, execroot, tools) - - try: - yield temp - finally: - temp.remove()