diff --git a/bazel/labgrid/client/common.py b/bazel/labgrid/client/common.py index e5e4b707ec62bbcc2a759deddd9509245072be64..f8d9c5027a329733d4b0725c7ff6055cb151e8be 100644 --- a/bazel/labgrid/client/common.py +++ b/bazel/labgrid/client/common.py @@ -1,29 +1,38 @@ import yaml -from subprocess import run, PIPE +from sys import stderr, stdout +from subprocess import run, TimeoutExpired, CalledProcessError from python.runfiles import Runfiles -def labgrid_client(*args, **kwargs): - env = None +def labgrid_client(*args, env=None, timeout=10, output_as_yaml=True, url=""): url_params = () - timeout = None runfiles = Runfiles.Create() labgrid_client_path = runfiles.Rlocation("rules_labgrid/labgrid/client/client") - if "env" in kwargs: - env = kwargs["env"] - if "url" in kwargs: - url_params = ("-x", kwargs["url"]) - if "timeout" in kwargs: - timeout = kwargs["timeout"] + # TODO: look at removing the url parameter and have the user specify it either + # via LG_COORDINATOR env var or by adding `-x ` to the call themselves. + if url: + url_params = ("-x", url) cmd = (labgrid_client_path,) + url_params + args - process = run( - cmd, check=True, stdout=PIPE, env=env, timeout=timeout, encoding="utf-8" - ) + try: + process = run( + cmd, + check=True, + capture_output=True, + env=env, + timeout=timeout, + encoding="utf-8", + ) + except (CalledProcessError, TimeoutExpired) as e: + stderr.write(e.output) + raise if process.stdout: - return yaml.safe_load(process.stdout) + stdout.write(process.stdout) + if output_as_yaml: + return yaml.safe_load(process.stdout) + return process.stdout else: - return [] + return {} diff --git a/bazel/labgrid/executor/BUILD.bazel b/bazel/labgrid/executor/BUILD.bazel deleted file mode 100644 index 5034f9800a10d052219c7b6e84067f40ab71f542..0000000000000000000000000000000000000000 --- a/bazel/labgrid/executor/BUILD.bazel +++ /dev/null @@ -1,7 +0,0 @@ -load("@rules_python//python:defs.bzl", "py_library") - -py_library( - name = "executor", - srcs = ["manager.py"], - visibility = ["//visibility:public"], -) diff --git a/bazel/labgrid/manager/BUILD.bazel b/bazel/labgrid/manager/BUILD.bazel new file mode 100644 index 0000000000000000000000000000000000000000..7b66fd3021d70473625ae0eada4a83fd1c12c979 --- /dev/null +++ b/bazel/labgrid/manager/BUILD.bazel @@ -0,0 +1,44 @@ +load("@rules_python//python:defs.bzl", "py_library") + +py_library( + name = "manager", + srcs = [ + "__init__.py", + "manager.py", + ], + visibility = ["//visibility:public"], +) + +py_library( + name = "acquire", + srcs = [ + "acquire.py", + ], + data = [ + "//labgrid/client", + ], + visibility = ["//visibility:public"], + deps = [ + ":manager", + "//bazel/labgrid/client", + "//labgrid:pkg", + "@rules_python//python/runfiles", + ], +) + +py_library( + name = "reserve", + srcs = [ + "reserve.py", + ], + data = [ + "//labgrid/client", + ], + visibility = ["//visibility:public"], + deps = [ + ":manager", + "//bazel/labgrid/client", + "//labgrid:pkg", + "@rules_python//python/runfiles", + ], +) diff --git a/bazel/labgrid/manager/__init__.py b/bazel/labgrid/manager/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3ada44112e9dfcd62dbefb3893a710d6cfe5b9df --- /dev/null +++ b/bazel/labgrid/manager/__init__.py @@ -0,0 +1,3 @@ +from .manager import Manager, Data + +__all__ = ["Manager", "Data"] diff --git a/bazel/labgrid/manager/acquire.py b/bazel/labgrid/manager/acquire.py new file mode 100644 index 0000000000000000000000000000000000000000..bb4f79aea2564b516f42d64575dcbf027aad4f36 --- /dev/null +++ b/bazel/labgrid/manager/acquire.py @@ -0,0 +1,17 @@ +from __future__ import annotations +from contextlib import contextmanager +from typing import Iterator + +from bazel.labgrid.manager.manager import Data +from bazel.labgrid.client import labgrid_client + + +@contextmanager +def manager(input_data: Data) -> Iterator[Data]: + env = input_data.env + labgrid_client("acquire", env=env) + + try: + yield Data(env=env) + finally: + labgrid_client("release", env=env) diff --git a/bazel/labgrid/executor/manager.py b/bazel/labgrid/manager/manager.py similarity index 100% rename from bazel/labgrid/executor/manager.py rename to bazel/labgrid/manager/manager.py diff --git a/bazel/labgrid/manager/reserve.py b/bazel/labgrid/manager/reserve.py new file mode 100644 index 0000000000000000000000000000000000000000..c35d7fbf2562870c60088aea666fa03c25d6dc93 --- /dev/null +++ b/bazel/labgrid/manager/reserve.py @@ -0,0 +1,36 @@ +from __future__ import annotations +from contextlib import contextmanager +from typing import Iterator + +from bazel.labgrid.manager.manager import Data +from bazel.labgrid.client import labgrid_client + + +@contextmanager +def manager(input_data: Data) -> Iterator[Data]: + env = input_data.env + tags = env.get("BZL_LG_TAGS", "") + token = "" + + if tags: + # Reserve and get token + result = labgrid_client("reserve", tags, env=env) + token = list(result.values())[0].get("token") + + # Wait for allocation + reserve_timeout = env.get("BZL_LG_RESERVE_TIMEOUT", 60) + result = labgrid_client("wait", token, env=env, timeout=reserve_timeout) + place = result.get("allocations", {}).get("main", "+") + + env.update( + { + "LG_TOKEN": token, + "LG_PLACE": place, + } + ) + + try: + yield Data(env=env) + finally: + if token: + labgrid_client("cancel-reservation", token, env=env) diff --git a/labgrid/executor/BUILD.bazel b/labgrid/executor/BUILD.bazel index 1effed64887fe57273b021f473a8df7cd873c248..ba4b44204509d5e27cb6f4ce833f241e3327593d 100644 --- a/labgrid/executor/BUILD.bazel +++ b/labgrid/executor/BUILD.bazel @@ -17,7 +17,7 @@ py_binary( tags = ["manual"], visibility = ["//:__subpackages__"], deps = [ - "//bazel/labgrid/executor", + "//bazel/labgrid/manager", "//bazel/python/runfiles", "//labgrid/config:managers", ], @@ -28,7 +28,7 @@ py_library( srcs = ["host.py"], data = ["hello-world.txt"], deps = [ - "//bazel/labgrid/executor", + "//bazel/labgrid/manager", "//bazel/python/runfiles", ], ) diff --git a/labgrid/executor/executor.py b/labgrid/executor/executor.py index 467d90276ba303fa00f60e063b8300f11ce4d7f8..a614e623748fe8a64c1bbce4523287350a82695a 100644 --- a/labgrid/executor/executor.py +++ b/labgrid/executor/executor.py @@ -12,7 +12,7 @@ from shutil import which from subprocess import CalledProcessError, run from sys import argv -from bazel.labgrid.executor.manager import Data, Manager +from bazel.labgrid.manager import Data, Manager from bazel.python.runfiles import runfile, RunfileNotFoundError diff --git a/labgrid/executor/host.py b/labgrid/executor/host.py index 0b3b913bb0bdda03303ffa8949c80162f79d9428..54e044cca2db8264fddf6fc66abbdc6ecef33584 100644 --- a/labgrid/executor/host.py +++ b/labgrid/executor/host.py @@ -3,8 +3,7 @@ from __future__ import annotations from contextlib import contextmanager from typing import Iterator - -from bazel.labgrid.executor.manager import Data +from bazel.labgrid.manager import Data from bazel.python.runfiles import runfile