From 1a656986c461c3b4a727fb75f5cb0cfbb0ce7599 Mon Sep 17 00:00:00 2001 From: Xuming Meng Date: Tue, 29 Oct 2024 15:36:08 +0100 Subject: [PATCH] feat(state): introduce labgrid_state --- bazel/labgrid/strategy/BUILD.bazel | 1 + bazel/labgrid/strategy/__init__.py | 1 + bazel/labgrid/strategy/transition.py | 12 +++++++++ e2e/docker/BUILD.bazel | 10 +++++-- labgrid/config/rule.bzl | 15 ++++++++++- labgrid/config/toolchain/macro.bzl | 5 ++-- labgrid/run/BUILD.bazel | 1 + labgrid/run/run.py | 39 ++++++++++++++++------------ labgrid/state/BUILD.bazel | 8 ++++++ labgrid/state/StateInfo.bzl | 28 ++++++++++++++++++++ labgrid/state/defs.bzl | 7 +++++ labgrid/state/rule.bzl | 32 +++++++++++++++++++++++ 12 files changed, 137 insertions(+), 22 deletions(-) create mode 100644 bazel/labgrid/strategy/transition.py create mode 100644 labgrid/state/BUILD.bazel create mode 100644 labgrid/state/StateInfo.bzl create mode 100644 labgrid/state/defs.bzl create mode 100644 labgrid/state/rule.bzl diff --git a/bazel/labgrid/strategy/BUILD.bazel b/bazel/labgrid/strategy/BUILD.bazel index 43ba4cb6..223b0ae7 100644 --- a/bazel/labgrid/strategy/BUILD.bazel +++ b/bazel/labgrid/strategy/BUILD.bazel @@ -5,6 +5,7 @@ py_library( srcs = [ "__init__.py", "qemu/strategy.py", + "transition.py" ], visibility = ["//visibility:public"], deps = [ diff --git a/bazel/labgrid/strategy/__init__.py b/bazel/labgrid/strategy/__init__.py index d1fc1fab..f5524725 100644 --- a/bazel/labgrid/strategy/__init__.py +++ b/bazel/labgrid/strategy/__init__.py @@ -1 +1,2 @@ from .qemu.strategy import QEMUStrategy +from .transition import transition diff --git a/bazel/labgrid/strategy/transition.py b/bazel/labgrid/strategy/transition.py new file mode 100644 index 00000000..afebcd30 --- /dev/null +++ b/bazel/labgrid/strategy/transition.py @@ -0,0 +1,12 @@ +from contextlib import contextmanager + +@contextmanager +def transition(strategy, initial, desired, final): + try: + if initial is not None: + strategy.force(initial) + strategy.transition(desired) + yield + finally: + if final is not None: + strategy.transition(final) diff --git a/e2e/docker/BUILD.bazel b/e2e/docker/BUILD.bazel index 596bd45e..21e9f3b9 100644 --- a/e2e/docker/BUILD.bazel +++ b/e2e/docker/BUILD.bazel @@ -2,6 +2,7 @@ load("@rules_diff//diff/file/test:defs.bzl", "diff_file_test") load("@rules_labgrid//labgrid/config/toolchain:defs.bzl", "labgrid_config_toolchain") load("@rules_labgrid//labgrid/run:defs.bzl", "labgrid_run") load("@rules_python//python:defs.bzl", "py_binary") +load("@rules_labgrid//labgrid/state:defs.bzl", "labgrid_state") # A constraint to register the toolchain to constraint_setting(name = "device") @@ -22,12 +23,17 @@ platform( ], ) +labgrid_state( + name = "docker-state", + desired = "accessible", + final = "gone", +) + # Provide the toolchain with Bazel labgrid_config_toolchain( name = "config", src = "local-ubuntu.16.04-gnu.yaml", - final_state = "gone", - state = "accessible", + state = ":docker-state", target_compatible_with = [ ":constraint", "@toolchain_utils//toolchain/constraint/cpu:amd64", diff --git a/labgrid/config/rule.bzl b/labgrid/config/rule.bzl index d02e5bfc..08de1589 100644 --- a/labgrid/config/rule.bzl +++ b/labgrid/config/rule.bzl @@ -1,4 +1,5 @@ load("@rules_python//python:defs.bzl", "PyInfo") +load("@rules_labgrid//labgrid/state:defs.bzl", "LabgridStateInfo") visibility("//...") @@ -63,6 +64,10 @@ ATTRS = { "env": attr.string_dict( doc = "The environment variables to set. Subject to `$(location)` and \"Make\" variable substitution.", ), + "state": attr.label( + doc = "The state of the configuration.", + providers = [LabgridStateInfo], + ), "data": attr.label_list( allow_files = True, doc = "The data files to use.", @@ -125,10 +130,18 @@ def _forward_runfiles(ctx, targets): runfiles = runfiles.merge_all([d.default_runfiles for d in targets]) return DefaultInfo(files = files, runfiles = runfiles) +def _state_env(state): + tuples = [ + ("LG_INITIAL_STATE", state.initial), + ("LG_STATE", state.desired), + ("BZL_LG_FINAL_STATE", state.final), + ] + return { k: v for k, v in tuples if v} + def implementation(ctx): data = ctx.attr.data + [ctx.attr.src] tools = ctx.attr.tools + ctx.attr.toolchains - env = {k: _expand(ctx, data, v) for k, v in ctx.attr.env.items()} + env = {k: _expand(ctx, data, v) for k, v in ctx.attr.env.items()} | _state_env(ctx.attr.state[LabgridStateInfo]) default = DefaultInfo(files = depset([ctx.file.src])) toolchain = platform_common.ToolchainInfo( src = ctx.file.src, diff --git a/labgrid/config/toolchain/macro.bzl b/labgrid/config/toolchain/macro.bzl index 82b5d0d3..6797546c 100644 --- a/labgrid/config/toolchain/macro.bzl +++ b/labgrid/config/toolchain/macro.bzl @@ -2,7 +2,7 @@ load("@rules_labgrid//labgrid/config:defs.bzl", "labgrid_config") visibility("//...") -def labgrid_config_toolchain(*, name, src, target_compatible_with, state = "shell", final_state = "off", managers = [], deps = [], env = {}, data = [], tools = [], toolchains = []): +def labgrid_config_toolchain(*, name, src, target_compatible_with, state = "@rules_labgrid//labgrid/state", managers = [], deps = [], env = {}, data = [], tools = [], toolchains = []): src = native.package_relative_label(src) labgrid_config( name = "{}-config".format(name), @@ -10,9 +10,8 @@ def labgrid_config_toolchain(*, name, src, target_compatible_with, state = "shel deps = [native.package_relative_label(d) for d in deps], env = env | { "LG_ENV": "$(location {})".format(src), - "LG_STATE": state, - "BZL_LG_FINAL_STATE": final_state, }, + state = state, data = data, tools = tools, toolchains = toolchains, diff --git a/labgrid/run/BUILD.bazel b/labgrid/run/BUILD.bazel index 0a2da3e6..f9e813d6 100644 --- a/labgrid/run/BUILD.bazel +++ b/labgrid/run/BUILD.bazel @@ -28,6 +28,7 @@ py_binary( visibility = ["//visibility:public"], deps = [ "//bazel/labgrid/util:target", + "//bazel/labgrid/strategy", "//labgrid:pkg", "//labgrid/config:deps", "@rules_python//python/runfiles", diff --git a/labgrid/run/run.py b/labgrid/run/run.py index c21ce9f8..a19e2971 100644 --- a/labgrid/run/run.py +++ b/labgrid/run/run.py @@ -2,6 +2,8 @@ from __future__ import annotations from argparse import ArgumentParser +import argparse +from contextlib import contextmanager from dataclasses import dataclass from os import environ, linesep from pathlib import Path, PurePath @@ -11,6 +13,7 @@ from sys import argv, stderr, stdout from python.runfiles import Runfiles +from bazel.labgrid.strategy import transition from bazel.labgrid.util.target import TemporaryDirectory from labgrid import Environment @@ -40,15 +43,23 @@ def arguments(prsr: ArgumentParser) -> None: default=PurePath(environ["LG_ENV"]), type=PurePath, ) + prsr.add_argument( + "--initial-state", + help="The state to transition the LabGrid strategy to initially", + dest="initial", + action=StateAction, + ) prsr.add_argument( "--state", help="The state to transition the LabGrid strategy to", - default=environ["LG_STATE"], + dest="desired", + action=StateAction, ) prsr.add_argument( "--final-state", help="The state to transition the LabGrid strategy to at the end", - default=environ["BZL_LG_FINAL_STATE"], + dest="final", + action=StateAction, ) prsr.add_argument( "--mktemp", @@ -70,11 +81,15 @@ def arguments(prsr: ArgumentParser) -> None: ) -@dataclass(frozen=True) +@dataclass class State: - desired: str - final: str | None = None + initial: str | None = environ.get("LG_INITIAL_STATE", None) + desired: str = environ["LG_STATE"] + final: str | None = environ.get("BZL_LG_FINAL_STATE", None) +class StateAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string = None): + setattr(namespace.state, self.dest, values) # TODO: move this out to `bazel.labgrid.run` on day to allow downstream to re-use def run( @@ -96,9 +111,7 @@ def run( env = Environment(str(config)) target = env.get_target() strategy = target.get_driver("Strategy") - try: - strategy.transition(state.desired) - + with transition(strategy, state.initial, state.desired, state.final): # Retrieve the communication protocols shell = target.get_driver("CommandProtocol") transfer = target.get_driver("FileTransferProtocol") @@ -114,7 +127,6 @@ def run( dest_runfiles = remote.with_suffix(".runfiles") if src_runfiles.exists(): transfer.put(f"{src_runfiles}", f"{dest_runfiles}") - # Run the transferred program out, err, code = shell.run(join((f"{remote}", *arguments))) for line in out: @@ -122,27 +134,23 @@ def run( for line in err: stderr.write(f"{line}{linesep}") return code - finally: - strategy.transition(state.final) def main(exe: Path, *args: str) -> int: prsr = ArgumentParser( prog=str(exe), description="Runs a command on a LabGrid device." ) + prsr.set_defaults(state=State()) 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, + state=parsed.state, config=parsed.config, mktemp=parsed.mktemp, rm=parsed.rm, @@ -158,7 +166,6 @@ def main(exe: Path, *args: str) -> int: except KeyboardInterrupt: return 130 - def entry(): exit(main(Path(argv[0]), *argv[1:])) diff --git a/labgrid/state/BUILD.bazel b/labgrid/state/BUILD.bazel new file mode 100644 index 00000000..e969b38f --- /dev/null +++ b/labgrid/state/BUILD.bazel @@ -0,0 +1,8 @@ +load(":defs.bzl", "labgrid_state") + +labgrid_state( + name = "state", + desired = "shell", + final = "off", + visibility = ["//visibility:public"], +) \ No newline at end of file diff --git a/labgrid/state/StateInfo.bzl b/labgrid/state/StateInfo.bzl new file mode 100644 index 00000000..e4e952f2 --- /dev/null +++ b/labgrid/state/StateInfo.bzl @@ -0,0 +1,28 @@ +load("@bazel_skylib//lib:types.bzl", "types") + +visibility("//...") + +def init(initial, desired, final): + if not (initial == None or types.is_string(initial)): + fail("`LabgridStateInfo.initial` must be `None` or a string: {}".format(initial)) + + if not types.is_string(desired): + fail("`LabgridStateInfo.desired` must be a string: {}".format(desired)) + + if not (final== None or types.is_string(final)): + fail("`LabgridStateInfo.final` must be `None` or a string: {}".format(final)) + + return { + "initial": initial, + "desired": desired, + "final": final, + } + +LabgridStateInfo, labgrid_state_info = provider( + doc = "A LabGrid configuration state.", + fields = ["initial","desired","final"], + init = init, +) + +StateInfo = LabgridStateInfo +state_info = labgrid_state_info \ No newline at end of file diff --git a/labgrid/state/defs.bzl b/labgrid/state/defs.bzl new file mode 100644 index 00000000..089023f7 --- /dev/null +++ b/labgrid/state/defs.bzl @@ -0,0 +1,7 @@ +load(":rule.bzl", _state = "state") +load(":StateInfo.bzl", "StateInfo") + +visibility("public") + +labgrid_state = _state +LabgridStateInfo = StateInfo \ No newline at end of file diff --git a/labgrid/state/rule.bzl b/labgrid/state/rule.bzl new file mode 100644 index 00000000..ae44b762 --- /dev/null +++ b/labgrid/state/rule.bzl @@ -0,0 +1,32 @@ +load("//labgrid/cfg:unstore.bzl", _cfg = "unstore") +load(":StateInfo.bzl", "StateInfo") + +visibility("//labgrid/state/...") + +DOC = "A rule to create a StateInfo object." + +ATTRS = { + "initial": attr.string( + doc = "The initial state of the configuration.", + ), + "desired": attr.string( + doc = "The desired state of the configuration.", + mandatory = True, + ), + "final": attr.string( + doc = "The final state of the configuration.", + ), +} + +def implementation(ctx): + return StateInfo(initial = ctx.attr.initial, desired = ctx.attr.desired, final = ctx.attr.final) + +labgrid_state = rule( + doc = DOC, + attrs = ATTRS, + implementation = implementation, + provides = [StateInfo], + cfg = _cfg, +) + +state = labgrid_state \ No newline at end of file -- GitLab