diff --git a/bazel/labgrid/protocol/util/target.py b/bazel/labgrid/protocol/util/target.py deleted file mode 100644 index 1ec7b4207f8ec58ac9644e0a83b1e3806e57d00e..0000000000000000000000000000000000000000 --- a/bazel/labgrid/protocol/util/target.py +++ /dev/null @@ -1,67 +0,0 @@ -import uuid -from contextlib import contextmanager -from pathlib import Path, PurePath -from shlex import join -from subprocess import CalledProcessError, CompletedProcess -from typing import Iterator, Mapping - -from labgrid.protocol.commandprotocol import CommandProtocol -from labgrid.protocol.filetransferprotocol import FileTransferProtocol - - -class Subprocess: - def __init__(self, shell: CommandProtocol): - self.__shell = shell - - @property - def shell(self) -> CommandProtocol: - return self.__shell - - def run( - self, args: tuple[str, ...], *, check: bool = False, text: bool = False - ) -> CompletedProcess: - assert text, "Non-text target execution is not supported" - out, err, code = self.shell.run(args) - if check and code: - raise CalledProcessError(cmd=args, returncode=code, output=out, stderr=err) - return CompletedProcess(returncode=code, stdout=out, stderr=err, args=args) - - -@contextmanager -def TemporaryDirectory( - shell: CommandProtocol, - transfer: FileTransferProtocol, - mktemp: Path, - rm: Path, - mkdir: Path, -) -> Iterator[PurePath]: - # Transfer mktemp - mktemp_remote = ( - f"/tmp/{uuid.getnode()}" - ) # avoid race-condition if there are parallel executions - transfer.put(mktemp, mktemp_remote) - target = Subprocess(shell) - result = target.run(f"{mktemp_remote} -d", check=True, text=True) - root = PurePath(result.stdout[0].rstrip()) - - tools = root / "tools" - execroot = root / "execroot" - # Transfer mkdir - transfer.put(mkdir, f"{root}") - # Create tools dir - target.run(f"{root / 'mkdir'} {tools}", check=True, text=True) - # Create execroot dir - target.run(f"{root / 'mkdir'} {execroot}", check=True, text=True) - # Transfer rm - transfer.put(rm, f"{tools}") - - ##cleanup - - # Delete mkdir - target.run(f"{tools / 'rm'} {root / 'mkdir'}", check=True, text=True) - # Delete mktemp - target.run(f"{tools / 'rm'} {mktemp_remote}", check=True, text=True) - - yield execroot - - target.run(f"{tools / 'rm'} -rf {root}", check=True, text=True) diff --git a/bazel/labgrid/protocol/util/BUILD.bazel b/bazel/labgrid/util/BUILD.bazel similarity index 100% rename from bazel/labgrid/protocol/util/BUILD.bazel rename to bazel/labgrid/util/BUILD.bazel diff --git a/bazel/labgrid/util/target.py b/bazel/labgrid/util/target.py new file mode 100644 index 0000000000000000000000000000000000000000..5ef2a148a64e0d1f8a3379b5f2f2134fdf6ac984 --- /dev/null +++ b/bazel/labgrid/util/target.py @@ -0,0 +1,75 @@ +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 + ) + + +@contextmanager +def TemporaryDirectory( + shell: CommandProtocol, + transfer: FileTransferProtocol, + mktemp: Path, + rm: Path, + mkdir: 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 some basic tools + transfer.put(f"{mkdir}", f"{root / 'mkdir'}") + transfer.put(f"{rm}", f"{root / 'rm'}") + + # Create the execution root + execroot = root / "execroot" + target.run((f"{root / 'mkdir'}", f"{execroot}"), check=True, text=True) + + # Remove unneeded tools + target.run((f"{root / 'rm'}", f"{root / 'mkdir'}"), check=True, text=True) + target.run((f"{root / 'rm'}", f"{unique}"), check=True, text=True) + + # Clean up the root directory when complete + try: + yield execroot + finally: + target.run((f"{root / 'rm'}", "-rf", f"{root}"), check=True, text=True) diff --git a/labgrid/run/BUILD.bazel b/labgrid/run/BUILD.bazel index 80f170ab2f8fa085bcaae1a5b0adbf97db807fd2..0a2da3e6d10b4c4a1cebb5c7462c216d56644c13 100644 --- a/labgrid/run/BUILD.bazel +++ b/labgrid/run/BUILD.bazel @@ -27,7 +27,7 @@ py_binary( tags = ["manual"], visibility = ["//visibility:public"], deps = [ - "//bazel/labgrid/protocol/util:target", + "//bazel/labgrid/util:target", "//labgrid:pkg", "//labgrid/config:deps", "@rules_python//python/runfiles", diff --git a/labgrid/run/run.py b/labgrid/run/run.py index 9027dce0572718fc90083f385bc3d49a38a36e07..c21ce9f84995c7b4a70f37c37aade447ee75f211 100644 --- a/labgrid/run/run.py +++ b/labgrid/run/run.py @@ -1,21 +1,30 @@ #! /usr/bin/env python3 +from __future__ import annotations + from argparse import ArgumentParser +from dataclasses import dataclass from os import environ, linesep from pathlib import Path, PurePath from shlex import join -from subprocess import PIPE, run +from subprocess import CalledProcessError from sys import argv, stderr, stdout -from typing import Collection, Protocol, Tuple from python.runfiles import Runfiles -from bazel.labgrid.protocol.util.target import TemporaryDirectory -from labgrid import Environment, Target +from bazel.labgrid.util.target import TemporaryDirectory +from labgrid import Environment + +class RunfileNotFoundError(FileNotFoundError): + pass -class Shell(Protocol): - def run(cmd: str) -> Tuple[Collection[str], Collection[str], int]: - ... + +def runfile(path: str) -> Path: + runfiles = Runfiles.Create() + resolved = Path(runfiles.Rlocation(path)) + if not resolved.exists(): + raise RunfileNotFoundError(f"runfile not found: {path}") + return resolved def arguments(prsr: ArgumentParser) -> None: @@ -41,58 +50,113 @@ def arguments(prsr: ArgumentParser) -> None: help="The state to transition the LabGrid strategy to at the end", default=environ["BZL_LG_FINAL_STATE"], ) + prsr.add_argument( + "--mktemp", + help="A `mktemp` binary for the target.", + type=runfile, + default=runfile("ape/ape/assimilate/mktemp.ape/mktemp"), + ) + prsr.add_argument( + "--rm", + help="A `rm` binary for the target.", + type=runfile, + default=runfile("ape/ape/assimilate/rm.ape/rm"), + ) + prsr.add_argument( + "--mkdir", + help="A `mkdir` binary for the target.", + type=runfile, + default=runfile("ape/ape/assimilate/mkdir.ape/mkdir"), + ) -def main(exe: Path, *args: str) -> int: - # Allow subprocesses to have their own runfiles - runfiles = Runfiles.Create() - - del environ["RUNFILES_DIR"] +@dataclass(frozen=True) +class State: + desired: str + final: str | None = None - prsr = ArgumentParser( - prog=str(exe), description="Runs a command over SSH to a LabGrid Docker device." - ) - arguments(prsr) - - parsed = prsr.parse_args(args) +# TODO: move this out to `bazel.labgrid.run` on day to allow downstream to re-use +def run( + program: Path, + *arguments: str, + config: Path, + state: State, + mktemp: Path, + mkdir: Path, + rm: Path, +) -> int: + # Allow any Bazel binaries used in the configuration to work + try: + del environ["RUNFILES_DIR"] + except KeyError: + pass # Start up Docker container with LabGrid - env = Environment(str(parsed.config)) + env = Environment(str(config)) target = env.get_target() strategy = target.get_driver("Strategy") try: - strategy.transition(parsed.state) + strategy.transition(state.desired) # Retrieve the communication protocols shell = target.get_driver("CommandProtocol") transfer = target.get_driver("FileTransferProtocol") - # utils - mktemp = runfiles.Rlocation("ape/ape/assimilate/mktemp.ape/mktemp") - mkdir = runfiles.Rlocation("ape/ape/assimilate/mkdir.ape/mkdir") - rm = runfiles.Rlocation("ape/ape/assimilate/rm.ape/rm") - with TemporaryDirectory(shell, transfer, mktemp, rm, mkdir) as temp_dir: - # Transfer the provided program over to the Docker image - program = "{}/{}".format(temp_dir, parsed.program.name) - transfer.put(parsed.program, program) + with TemporaryDirectory(shell, transfer, mktemp, rm, mkdir) as temp: + # Transfer the provided program over to the device + remote = temp / program.name + transfer.put(f"{program}", f"{remote}") # Transfer runfiles - src_runfiles = parsed.program.with_suffix(".runfiles") - dest_runfiles = Path(program).with_suffix(".runfiles") + # TODO: implement a `BazelFileTransfer` that transfers runfiles implicitly in `put` + # Make that the default file transfer implementation so _any_ file transfer automatically gets the runfiles as well + src_runfiles = program.with_suffix(".runfiles") + dest_runfiles = remote.with_suffix(".runfiles") if src_runfiles.exists(): - transfer.put(src_runfiles, dest_runfiles) + transfer.put(f"{src_runfiles}", f"{dest_runfiles}") # Run the transferred program - out, err, code = shell.run(join((program, *parsed.arguments))) + out, err, code = shell.run(join((f"{remote}", *arguments))) for line in out: stdout.write(f"{line}{linesep}") for line in err: stderr.write(f"{line}{linesep}") - + return code finally: - strategy.transition(parsed.final_state) + strategy.transition(state.final) - return code + +def main(exe: Path, *args: str) -> int: + prsr = ArgumentParser( + prog=str(exe), description="Runs a command on a LabGrid device." + ) + + arguments(prsr) + + parsed = prsr.parse_args(args) + + # TODO: implement an argument parser action to do this + state = State(final=parsed.final_state, desired=parsed.state) + + try: + return run( + parsed.program, + *parsed.arguments, + state=state, + config=parsed.config, + mktemp=parsed.mktemp, + rm=parsed.rm, + mkdir=parsed.mkdir, + ) + except CalledProcessError as e: + print(f"fatal: subprocess failed: {e}", file=stderr) + if e.stdout is not None: + print(e.stdout, file=stderr) + if e.stderr is not None: + print(e.stderr, file=stderr) + return e.returncode + except KeyboardInterrupt: + return 130 def entry():