diff --git a/bazel/labgrid/runner/BUILD.bazel b/bazel/labgrid/runner/BUILD.bazel index 5c0e712128bb35a0d1603121fd5bd95e49834190..69d06f347f924174577b19bd4fb2970341486b7d 100644 --- a/bazel/labgrid/runner/BUILD.bazel +++ b/bazel/labgrid/runner/BUILD.bazel @@ -1,21 +1,30 @@ load("@rules_labgrid//labgrid/config:transition.bzl", "labgrid_config_transition") load("@rules_python//python:defs.bzl", "py_library") -TOOLS = [ - "mktemp", - "mkdir", - "rm", - "env", - "mv", -] +py_library( + name = "subprocess", + srcs = ["subprocess.py"], + deps = ["//labgrid:pkg"], + visibility = [ + "//bazel/labgrid/runner/tool_users:__subpackages__" + ], +) + +labgrid_config_transition( + name = "mktemp", + srcs = ["@ape//ape:mktemp".format("mktemp")], +) -[ - labgrid_config_transition( - name = tool, - srcs = ["@ape//ape:{}".format(tool)], - ) - for tool in TOOLS -] +py_library( + name = "temp_manager", + srcs = ["temp_manager.py"], + data = [":mktemp"], + deps = [ + ":subprocess", + "//bazel/labgrid/runner/tool_users:rm_tool", + "//bazel/python/runfiles", + ] +) py_library( name = "runner", @@ -23,11 +32,13 @@ py_library( "__init__.py", "runner.py", ], - data = [":{}".format(tool) for tool in TOOLS], visibility = ["//visibility:public"], deps = [ + ":temp_manager", + "//bazel/labgrid/runner/tool_users:put", + "//bazel/labgrid/runner/tool_users:get", + "//bazel/labgrid/runner/tool_users:exec", "//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 a2f88d36ce822cf23c97d61cc39a56c09041d4db..099f5f6276a79aa26b543e0d23eb60dd39a292e2 100644 --- a/bazel/labgrid/runner/runner.py +++ b/bazel/labgrid/runner/runner.py @@ -12,20 +12,14 @@ from typing import ( ) 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 .temp_manager import TempManager +from .tool_users.put import Put +from .tool_users.get import Get +from .tool_users.exec import Exec @dataclass(frozen=True) @@ -57,16 +51,6 @@ def _log_level(value): def runner(): config = PurePath(environ["LG_ENV"]) log_level = _log_level(environ.get("LOG_LEVEL", "ERROR")) - tools = _Tools() - - runfiles_dir = None - # Allow any Bazel binaries used in the configuration to work - try: - # FIXME: Convert into context manager - runfiles_dir = environ["RUNFILES_DIR"] - del environ["RUNFILES_DIR"] - except KeyError: - pass # Activating the StepLogger requires a DEBUG level of verbosity basicConfig(level=log_level) @@ -79,60 +63,47 @@ 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, runfiles_dir) - - # Transfer 'env' binary over to the device - r.put([FileTransfer(r.exec_path / tools.env.name, tools.env)]) + with TempManager(transfer, shell) as (temp_manager, tool_manager): + temp = temp_manager.mktemp() + r = Runner(temp, transfer, shell, tool_manager) yield r class Runner: - def __init__(self, exec_path, transfer, shell, runfiles_dir): + def __init__(self, exec_path, transfer, shell, tool_manager): self.exec_path = exec_path self.transfer = transfer self.shell = shell - self.runfiles_dir = runfiles_dir + + self._get = Get(transfer) + self._put = tool_manager.make(Put, (self.shell, transfer)) + self._exec = tool_manager.make(Exec, (exec_path, self.shell)) 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._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._put(f"{src_runfiles}", f"{dest_runfiles}") + elif get(environ["RUNFILES_DIR"], None) is not None: + self._put(environ["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._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) + out, err, code = self._exec(cmd, env) for line in out: stdout.write(f"{line}{linesep}") for line in err: diff --git a/bazel/labgrid/runner/subprocess.py b/bazel/labgrid/runner/subprocess.py new file mode 100644 index 0000000000000000000000000000000000000000..cdc395f2ab87b0a938509896fc5048c3664ab1a8 --- /dev/null +++ b/bazel/labgrid/runner/subprocess.py @@ -0,0 +1,36 @@ +from subprocess import CalledProcessError, CompletedProcess +from typing import Iterable +from shlex import join +from os import linesep + +from labgrid.protocol.commandprotocol import CommandProtocol + + +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 = True, text: bool = True + ) -> 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 + ) diff --git a/bazel/labgrid/runner/temp_manager.py b/bazel/labgrid/runner/temp_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..a8e50e551d19996d440fdcb82c5e920c87b50a6f --- /dev/null +++ b/bazel/labgrid/runner/temp_manager.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from threading import get_native_id +from uuid import getnode +from pathlib import PurePath +from typing import Tuple + +from labgrid.protocol.commandprotocol import CommandProtocol +from labgrid.protocol.filetransferprotocol import FileTransferProtocol + +from bazel.python.runfiles import runfile +from .tool_users.tool_user import ToolManager +from .subprocess import Subprocess +from .tool_users.rm import Rm + + +class TempManager: + """Creates temporary directories and a ToolManager. Cleans once finished.""" + + def __init__(self, transfer: FileTransferProtocol, shell: CommandProtocol): + self._transfer = transfer + self._shell = shell + + self._temps = set() + self._unique = None + + def __enter__(self) -> Tuple[TempManager, ToolManager]: + # Get a unique temporary directory, avoiding concurrent `TempManager` instances + self._unique = f".{getnode()}.{get_native_id()}" + self._transfer.put(str(runfile("ape/ape/assimilate/mktemp/mktemp")), self._unique) + + self._tools_dir = self._mktemp(auto_clean=False) + tool_manager = ToolManager(self._tools_dir, self._transfer) + self._rm = tool_manager.make(Rm, (self._shell,)) + + return self, tool_manager + + def __exit__(self, exc_type, exc_val, exc_tb): + for temp in self._temps: + self._rm(temp) + + self._rm(f"./{self._unique}") + self._rm(self._tools_dir) + + def mktemp(self) -> PurePath: + return self._mktemp(auto_clean=True) + + def _mktemp(self, auto_clean=True): + result = Subprocess(self._shell).run((f"./{self._unique}", "-d"), check=True, text=True) + temp = PurePath(result.stdout.rstrip()) + if auto_clean: + self._temps.add(temp) + return temp diff --git a/bazel/labgrid/runner/tool_users/BUILD.bazel b/bazel/labgrid/runner/tool_users/BUILD.bazel new file mode 100644 index 0000000000000000000000000000000000000000..0953996c410f88ddcc8794fd62f7dd35dcd2a6ca --- /dev/null +++ b/bazel/labgrid/runner/tool_users/BUILD.bazel @@ -0,0 +1,59 @@ +load("@rules_labgrid//labgrid/config:transition.bzl", "labgrid_config_transition") +load("@rules_python//python:defs.bzl", "py_library") + +TOOLS = [ + "mkdir", + "rm", + "env", +] + +[ + labgrid_config_transition( + name = tool, + srcs = ["@ape//ape:{}".format(tool)], + ) + for tool in TOOLS +] + +py_library( + name = "tool_user", + srcs = ["tool_user.py"], +) + +py_library( + name = "rm_tool", + srcs = ["rm.py"], + data = [":rm"], + deps = [ + ":tool_user", + "//bazel/labgrid/runner:subprocess", + ], + visibility = ["//bazel/labgrid/runner:__subpackages__"], +) + +py_library( + name = "put", + srcs = ["put.py"], + data = [":mkdir"], + deps = [ + ":tool_user", + "//bazel/labgrid/runner:subprocess", + ], + visibility = ["//bazel/labgrid/runner:__subpackages__"], +) + +py_library( + name = "get", + srcs = ["get.py"], + visibility = ["//bazel/labgrid/runner:__subpackages__"], +) + +py_library( + name = "exec", + srcs = ["exec.py"], + data = [":env"], + deps = [ + ":tool_user", + ], + visibility = ["//bazel/labgrid/runner:__subpackages__"], +) diff --git a/bazel/labgrid/runner/tool_users/exec.py b/bazel/labgrid/runner/tool_users/exec.py new file mode 100644 index 0000000000000000000000000000000000000000..2aa136d9fe37d648e74badee0aa55dd333595eea --- /dev/null +++ b/bazel/labgrid/runner/tool_users/exec.py @@ -0,0 +1,37 @@ +from shlex import join +from os import environ +from contextlib import contextmanager +from pathlib import PurePath +from typing import Tuple + +from labgrid.protocol.commandprotocol import CommandProtocol + +from bazel.python.runfiles import runfile +from bazel.labgrid.runner.subprocess import Subprocess +from .tool_user import ToolUser + + +class Exec(ToolUser): + """Executes a command with env vars on the remote""" + + def __init__(self, tools_root: PurePath, exec_root: PurePath, shell: CommandProtocol): + super().__init__(tools_root) + self._exec_root = exec_root + self._shell = shell + + @property + def required_tools(self): + return [("env", runfile("ape/ape/assimilate/env/env"))] + + @contextmanager + def _no_runfiles_env(self): + runfiles_dir = environ["RUNFILES_DIR"] + del environ["RUNFILES_DIR"] + yield None + environ["RUNFILES_DIR"] = runfiles_dir + + def __call__(self, cmd, env) -> Tuple[list[str], list[str], int]: + with self._no_runfiles_env(): + envs = " ".join(f"{k}={v}" for k, v in env.items()) + cmd = f"cd {self._exec_root} && {self._tool('env')} {envs} {cmd}" + return self._shell.run(cmd) diff --git a/bazel/labgrid/runner/tool_users/get.py b/bazel/labgrid/runner/tool_users/get.py new file mode 100644 index 0000000000000000000000000000000000000000..16e1359d52aa50898f989e0e5d9e87264821813e --- /dev/null +++ b/bazel/labgrid/runner/tool_users/get.py @@ -0,0 +1,11 @@ +from labgrid.protocol.filetransferprotocol import FileTransferProtocol + + +class Get: + """Gets files from remote""" + def __init__(self, transfer: FileTransferProtocol): + self._transfer = transfer + + def __call__(self, remote: str, local: str): + # FIXME: Check if file exists with `@ape//ape:test` instead of assuming any error is due to a missing file + self._transfer.get(remote, local) diff --git a/bazel/labgrid/runner/tool_users/put.py b/bazel/labgrid/runner/tool_users/put.py new file mode 100644 index 0000000000000000000000000000000000000000..4bf7cd584e3194311fc12eb58f07207c45c7876e --- /dev/null +++ b/bazel/labgrid/runner/tool_users/put.py @@ -0,0 +1,25 @@ +from pathlib import PurePath + +from labgrid.protocol.commandprotocol import CommandProtocol +from labgrid.protocol.filetransferprotocol import FileTransferProtocol + +from bazel.python.runfiles import runfile +from bazel.labgrid.runner.subprocess import Subprocess +from .tool_user import ToolUser + + +class Put(ToolUser): + """Transfers files to the remote and makes parent directories""" + + def __init__(self, tools_root: PurePath, shell: CommandProtocol, transfer: FileTransferProtocol): + super().__init__(tools_root) + self._shell = Subprocess(shell) + self._transfer = transfer + + @property + def required_tools(self): + return [("mkdir", runfile("ape/ape/assimilate/mkdir/mkdir"))] + + def __call__(self, local: str, remote: str): + self._shell.run((str(self._tool("mkdir")), "-p", str(PurePath(remote).parent))) + self._transfer.put(local, remote) diff --git a/bazel/labgrid/runner/tool_users/rm.py b/bazel/labgrid/runner/tool_users/rm.py new file mode 100644 index 0000000000000000000000000000000000000000..79b08335b5b72f3c861502d69bfc0eb3eaa69f9b --- /dev/null +++ b/bazel/labgrid/runner/tool_users/rm.py @@ -0,0 +1,23 @@ +from pathlib import PurePath + +from labgrid.protocol.commandprotocol import CommandProtocol + +from bazel.python.runfiles import runfile +from bazel.labgrid.runner.subprocess import Subprocess +from .tool_user import ToolUser + + + +class Rm(ToolUser): + """Recursively removes a directory""" + + def __init__(self, tools_root: PurePath, shell: CommandProtocol): + super().__init__(tools_root) + self._shell = Subprocess(shell) + + @property + def required_tools(self): + return [("rm", runfile("ape/ape/assimilate/rm/rm"))] + + def __call__(self, path: PurePath): + self._shell.run((str(self._tool("rm")), "-r", str(path))) diff --git a/bazel/labgrid/runner/tool_users/tool_user.py b/bazel/labgrid/runner/tool_users/tool_user.py new file mode 100644 index 0000000000000000000000000000000000000000..3c5a884c6aa9eb9711305ca94316fa9c00f945b3 --- /dev/null +++ b/bazel/labgrid/runner/tool_users/tool_user.py @@ -0,0 +1,47 @@ +from abc import ABC, abstractmethod +from pathlib import PurePath +from typing import Tuple, Mapping, Any, Iterable + +from labgrid.protocol.filetransferprotocol import FileTransferProtocol + + +class ToolUser(ABC): + """Abstracts functionality that requires tool(s) on the remote""" + def __init__(self, tools_root: PurePath): + self._tools_root = tools_root + + @property + @abstractmethod + def required_tools(self) -> Iterable[Iterable[Any]]: + ... + + def _tool(self, name: str) -> PurePath: + return self._tools_root / name + + +class ToolManager: + """Creates ToolUser instances""" + + def __init__(self, temp: PurePath, transfer: FileTransferProtocol): + self._temp = temp + self._transfer = transfer + self._transferred = set() + + def make( + self, + tool_type: type, + init_args:Tuple[Any]=(), + init_kwargs: Mapping[str, Any]={} + ) -> ToolUser: + if not issubclass(tool_type, ToolUser): + raise Exception(f"{tool_type} must be a {ToolUser.__class__} subclass") + + tool = tool_type(self._temp, *init_args, **init_kwargs) + self._transfer_tools(tool.required_tools) + + return tool + + def _transfer_tools(self, tools): + for tool in tools: + if tool not in self._transferred: + self._transfer.put(str(tool[1]), str( self._temp / tool[0])) 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()