From 38f1ad54c325feeb46f69dcead9be63baf72a127 Mon Sep 17 00:00:00 2001 From: Callum Forrester Date: Wed, 10 Jul 2024 13:34:50 +0100 Subject: [PATCH] Automatically init scratch area and remove startup script (#520) Fixes #510 and #399 Write python code to do the job of `container-startup.sh` and add code to clone the required repos. Replace `container-startup.sh` with this code. Stop blueapi from initializing scratch automatically on container startup, replace this job with an [initContainer](https://kubernetes.io/docs/concepts/workloads/pods/init-containers/). Copy the blueapi venv to an [emptyDir](https://kubernetes.io/docs/concepts/storage/volumes/#emptydir) so the initContainers can manipulate it. --------- Co-authored-by: Joe Shannon --- Dockerfile | 5 +- container-startup.sh | 19 --- helm/blueapi/templates/deployment.yaml | 42 ++++- helm/blueapi/values.yaml | 20 ++- pyproject.toml | 5 +- src/blueapi/cli/cli.py | 11 ++ src/blueapi/cli/scratch.py | 101 ++++++++++++ src/blueapi/config.py | 11 ++ tests/cli/test_scratch.py | 214 +++++++++++++++++++++++++ tests/example_yaml/scratch.yaml | 5 + tests/test_cli.py | 23 +++ 11 files changed, 417 insertions(+), 39 deletions(-) delete mode 100755 container-startup.sh create mode 100644 src/blueapi/cli/scratch.py create mode 100644 tests/cli/test_scratch.py create mode 100644 tests/example_yaml/scratch.yaml diff --git a/Dockerfile b/Dockerfile index 6c6f3522c..dd1a2398f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -26,12 +26,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ git \ && rm -rf /var/lib/apt/lists/* COPY --from=build /venv/ /venv/ -COPY ./container-startup.sh /container-startup.sh ENV PATH=/venv/bin:$PATH - RUN mkdir -p /.cache/pip; chmod -R 777 /venv /.cache/pip -# change this entrypoint if it is not the same as the repo -ENTRYPOINT ["/container-startup.sh"] +ENTRYPOINT ["blueapi"] CMD ["serve"] diff --git a/container-startup.sh b/container-startup.sh deleted file mode 100755 index dc040ac3b..000000000 --- a/container-startup.sh +++ /dev/null @@ -1,19 +0,0 @@ -#!/usr/bin/env bash - -# Set umask to DLS standard -umask 0002 - -if [[ -z "${SCRATCH_AREA}" ]]; then - SCRATCH_AREA="/blueapi-plugins/scratch" -fi - -if [[ -d "${SCRATCH_AREA}" ]]; then - echo "Loading Python packages from ${SCRATCH_AREA}" - for DIR in ${SCRATCH_AREA}/*/; do # All directories - python -m pip install --no-deps -e "${DIR}" - done -else - echo "blueapi scratch area not present" -fi - -blueapi $@ diff --git a/helm/blueapi/templates/deployment.yaml b/helm/blueapi/templates/deployment.yaml index c211c26c6..142c5eb13 100644 --- a/helm/blueapi/templates/deployment.yaml +++ b/helm/blueapi/templates/deployment.yaml @@ -35,11 +35,39 @@ spec: sources: - configMap: name: {{ include "blueapi.fullname" . }}-config - {{- if .Values.scratch.hostPath }} + {{- if .Values.scratchHostPath }} - name: scratch-host hostPath: - path: {{ .Values.scratch.hostPath }} + path: {{ .Values.scratchHostPath }} type: Directory + - name: venv + emptyDir: + sizeLimit: 5Gi + {{- end }} + initContainers: + {{- if .Values.scratchHostPath }} + - name: setup-scratch + image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" + imagePullPolicy: {{ .Values.image.pullPolicy }} + resources: + {{- toYaml .Values.resources | nindent 12 }} + command: [/bin/sh, -c] + args: + - | + echo "Setting up scratch area" + blueapi -c /config/config.yaml setup-scratch + if [ $? -ne 0 ]; then echo 'Blueapi failed'; exit 1; fi; + echo "Exporting venv as artefact" + cp -r /venv/* /artefacts + volumeMounts: + - name: worker-config + mountPath: "/config" + readOnly: true + - name: scratch-host + mountPath: {{ .Values.worker.scratch.root }} + mountPropagation: HostToContainer + - name: venv + mountPath: /artefacts {{- end }} containers: - name: {{ .Chart.Name }} @@ -53,19 +81,21 @@ spec: - name: worker-config mountPath: "/config" readOnly: true - {{- if .Values.scratch.hostPath }} + {{- if .Values.scratchHostPath }} - name: scratch-host - mountPath: {{ .Values.scratch.containerPath }} + mountPath: {{ .Values.worker.scratch.root }} mountPropagation: HostToContainer + - name: venv + mountPath: /venv {{- end }} args: - "-c" - "/config/config.yaml" - "serve" + {{- with .Values.extraEnvVars }} env: - {{- if .Values.extraEnvVars }} {{- tpl .Values.extraEnvVars . | nindent 10 }} - {{- end }} + {{- end }} {{- with .Values.nodeSelector }} nodeSelector: {{- toYaml . | nindent 8 }} diff --git a/helm/blueapi/values.yaml b/helm/blueapi/values.yaml index 2e3788556..569d26076 100644 --- a/helm/blueapi/values.yaml +++ b/helm/blueapi/values.yaml @@ -74,20 +74,15 @@ listener: enabled: true resources: {} -scratch: - hostPath: "" # example: /usr/local/blueapi-software-scratch - containerPath: /blueapi-plugins/scratch - # Additional envVars to mount to the pod as a String -extraEnvVars: | - - name: SCRATCH_AREA - value: {{ .Values.scratch.containerPath }} +extraEnvVars: [] # - name: RABBITMQ_PASSWORD # valueFrom: # secretKeyRef: # name: rabbitmq-password # key: rabbitmq-password +# Config for the worker goes here, will be mounted into a config file worker: api: host: 0.0.0.0 # Allow non-loopback traffic @@ -108,4 +103,13 @@ worker: passcode: guest host: rabbitmq port: 61613 -# Config for the worker goes here, will be mounted into a config file + # Uncomment this to enable the scratch directory + # scratch: + # root: /blueapi-plugins/scratch + # repositories: [] + # - name: "dodal" + # remote_url: https://github.com/DiamondLightSource/dodal.git + +# Mount path for scratch area from host machine, setting +# this effectively enables scratch area management +scratchHostPath: "" # example: /usr/local/blueapi-software-scratch diff --git a/pyproject.toml b/pyproject.toml index 873896bd1..3b820255d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "nslsii", "pyepics", "aioca", - "pydantic<2.0", # Leave pinned until can check incompatibility + "pydantic<2.0", # Leave pinned until can check incompatibility "stomp-py", "aiohttp", "PyYAML", @@ -25,9 +25,10 @@ dependencies = [ "fastapi[all]", "uvicorn", "requests", - "dls-bluesky-core", #requires ophyd-async + "dls-bluesky-core", #requires ophyd-async "dls-dodal>=1.24.0", "super-state-machine", # See GH issue 553 + "GitPython", ] dynamic = ["version"] license.file = "LICENSE" diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index fdb910b29..bce515e11 100644 --- a/src/blueapi/cli/cli.py +++ b/src/blueapi/cli/cli.py @@ -28,6 +28,7 @@ from blueapi.worker import ProgressEvent, Task, WorkerEvent, WorkerState from .rest import BlueapiRestClient +from .scratch import setup_scratch @click.group(invoke_without_command=True) @@ -341,6 +342,16 @@ def env(obj: dict, reload: bool | None) -> None: print(client.get_environment()) +@main.command(name="setup-scratch") +@click.pass_obj +def scratch(obj: dict) -> None: + config: ApplicationConfig = obj["config"] + if config.scratch is not None: + setup_scratch(config.scratch) + else: + raise KeyError("No scratch config supplied") + + # helper function def process_event_after_finished(event: WorkerEvent, logger: logging.Logger): if event.is_error(): diff --git a/src/blueapi/cli/scratch.py b/src/blueapi/cli/scratch.py new file mode 100644 index 000000000..d731eeeb5 --- /dev/null +++ b/src/blueapi/cli/scratch.py @@ -0,0 +1,101 @@ +import logging +import os +import stat +from pathlib import Path +from subprocess import Popen + +from git import Repo + +from blueapi.config import ScratchConfig + +_DEFAULT_INSTALL_TIMEOUT: float = 300.0 + + +def setup_scratch( + config: ScratchConfig, install_timeout: float = _DEFAULT_INSTALL_TIMEOUT +) -> None: + """ + Set up the scratch area from the config. Clone all required repositories + if they are not cloned already. Install them into the scratch area. + + Args: + config: Configuration for the scratch directory + install_timeout: Timeout for installing packages + """ + + _validate_directory(config.root) + + logging.info(f"Setting up scratch area: {config.root}") + + for repo in config.repositories: + local_directory = config.root / repo.name + ensure_repo(repo.remote_url, local_directory) + scratch_install(local_directory, timeout=install_timeout) + + +def ensure_repo(remote_url: str, local_directory: Path) -> None: + """ + Ensure that a repository is checked out for use in the scratch area. + Clone it if it isn't. + + Args: + remote_url: Git remote URL + local_directory: Output path for cloning + """ + + # Set umask to DLS standard + os.umask(stat.S_IWOTH) + + if not local_directory.exists(): + logging.info(f"Cloning {remote_url}") + Repo.clone_from(remote_url, local_directory) + logging.info(f"Cloned {remote_url} -> {local_directory}") + elif local_directory.is_dir(): + Repo(local_directory) + logging.info(f"Found {local_directory}") + else: + raise KeyError( + f"Unable to open {local_directory} as a git repository because " + "it is a file" + ) + + +def scratch_install(path: Path, timeout: float = _DEFAULT_INSTALL_TIMEOUT) -> None: + """ + Install a scratch package. Make blueapi aware of a repository checked out in + the scratch area. Make it automatically follow code changes to that repository + (pending a restart). Do not install any of the package's dependencies as they + may conflict with each other. + + Args: + path: Path to the checked out repository + timeout: Time to wait for for installation subprocess + """ + + _validate_directory(path) + + # Set umask to DLS standard + os.umask(stat.S_IWOTH) + + logging.info(f"Installing {path}") + process = Popen( + [ + "python", + "-m", + "pip", + "install", + "--no-deps", + "-e", + str(path), + ] + ) + process.wait(timeout=timeout) + if process.returncode != 0: + raise RuntimeError(f"Failed to install {path}: Exit Code: {process.returncode}") + + +def _validate_directory(path: Path) -> None: + if not path.exists(): + raise KeyError(f"{path}: No such file or directory") + elif path.is_file(): + raise KeyError(f"{path}: Is a file, not a directory") diff --git a/src/blueapi/config.py b/src/blueapi/config.py index 4196b8123..60f929761 100644 --- a/src/blueapi/config.py +++ b/src/blueapi/config.py @@ -85,6 +85,16 @@ class RestConfig(BlueapiBaseModel): protocol: str = "http" +class ScratchRepository(BlueapiBaseModel): + name: str = "example" + remote_url: str = "https://github.com/example/example.git" + + +class ScratchConfig(BlueapiBaseModel): + root: Path = Path("/tmp/scratch/blueapi") + repositories: list[ScratchRepository] = Field(default_factory=list) + + class ApplicationConfig(BlueapiBaseModel): """ Config for the worker application as a whole. Root of @@ -95,6 +105,7 @@ class ApplicationConfig(BlueapiBaseModel): env: EnvironmentConfig = Field(default_factory=EnvironmentConfig) logging: LoggingConfig = Field(default_factory=LoggingConfig) api: RestConfig = Field(default_factory=RestConfig) + scratch: ScratchConfig | None = None def __eq__(self, other: object) -> bool: if isinstance(other, ApplicationConfig): diff --git a/tests/cli/test_scratch.py b/tests/cli/test_scratch.py new file mode 100644 index 000000000..7e4ca3b9e --- /dev/null +++ b/tests/cli/test_scratch.py @@ -0,0 +1,214 @@ +import os +import stat +import uuid +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import Mock, call, patch + +import pytest + +from blueapi.cli.scratch import ensure_repo, scratch_install, setup_scratch +from blueapi.config import ScratchConfig, ScratchRepository + + +@pytest.fixture +def directory_path() -> Path: # type: ignore + temporary_directory = TemporaryDirectory() + yield Path(temporary_directory.name) + temporary_directory.cleanup() + + +@pytest.fixture +def file_path(directory_path: Path) -> Path: # type: ignore + file_path = directory_path / str(uuid.uuid4()) + with file_path.open("w") as stream: + stream.write("foo") + yield file_path + os.remove(file_path) + + +@pytest.fixture +def nonexistant_path(directory_path: Path) -> Path: + file_path = directory_path / str(uuid.uuid4()) + assert not file_path.exists() + return file_path + + +@patch("blueapi.cli.scratch.Popen") +def test_scratch_install_installs_path( + mock_popen: Mock, + directory_path: Path, +): + mock_process = Mock() + mock_process.returncode = 0 + mock_popen.return_value = mock_process + + scratch_install(directory_path, timeout=1.0) + + mock_popen.assert_called_once_with( + [ + "python", + "-m", + "pip", + "install", + "--no-deps", + "-e", + str(directory_path), + ] + ) + + +def test_scratch_install_fails_on_file(file_path: Path): + with pytest.raises(KeyError): + scratch_install(file_path, timeout=1.0) + + +def test_scratch_install_fails_on_nonexistant_path(nonexistant_path: Path): + with pytest.raises(KeyError): + scratch_install(nonexistant_path, timeout=1.0) + + +@patch("blueapi.cli.scratch.Popen") +@pytest.mark.parametrize("code", [1, 2, 65536]) +def test_scratch_install_fails_on_non_zero_exit_code( + mock_popen: Mock, + directory_path: Path, + code: int, +): + mock_process = Mock() + mock_process.returncode = code + mock_popen.return_value = mock_process + + with pytest.raises(RuntimeError): + scratch_install(directory_path, timeout=1.0) + + +@patch("blueapi.cli.scratch.Repo") +def test_repo_not_cloned_and_validated_if_found_locally( + mock_repo: Mock, + directory_path: Path, +): + ensure_repo("http://example.com/foo.git", directory_path) + mock_repo.assert_called_once_with(directory_path) + mock_repo.clone_from.assert_not_called() + + +@patch("blueapi.cli.scratch.Repo") +def test_repo_cloned_if_not_found_locally( + mock_repo: Mock, + nonexistant_path: Path, +): + ensure_repo("http://example.com/foo.git", nonexistant_path) + mock_repo.assert_not_called() + mock_repo.clone_from.assert_called_once_with( + "http://example.com/foo.git", nonexistant_path + ) + + +@patch("blueapi.cli.scratch.Repo") +def test_repo_cloned_with_correct_umask( + mock_repo: Mock, + directory_path: Path, +): + repo_root = directory_path / "foo" + file_path = repo_root / "a" + + def write_repo_files(): + repo_root.mkdir() + with file_path.open("w") as stream: + stream.write("foo") + + mock_repo.clone_from.side_effect = lambda url, path: write_repo_files() + + ensure_repo("http://example.com/foo.git", repo_root) + assert file_path.exists() + assert file_path.is_file() + st = os.stat(file_path) + assert st.st_mode & stat.S_IWGRP + + +def test_repo_discovery_errors_if_file_found_with_repo_name(file_path: Path): + with pytest.raises(KeyError): + ensure_repo("http://example.com/foo.git", file_path) + + +def test_setup_scratch_fails_on_nonexistant_root( + nonexistant_path: Path, +): + config = ScratchConfig(root=nonexistant_path, repositories=[]) + with pytest.raises(KeyError): + setup_scratch(config) + + +def test_setup_scratch_fails_on_non_directory_root( + file_path: Path, +): + config = ScratchConfig(root=file_path, repositories=[]) + with pytest.raises(KeyError): + setup_scratch(config) + + +@patch("blueapi.cli.scratch.ensure_repo") +@patch("blueapi.cli.scratch.scratch_install") +def test_setup_scratch_iterates_repos( + mock_scratch_install: Mock, + mock_ensure_repo: Mock, + directory_path: Path, +): + config = ScratchConfig( + root=directory_path, + repositories=[ + ScratchRepository( + name="foo", + remote_url="http://example.com/foo.git", + ), + ScratchRepository( + name="bar", + remote_url="http://example.com/bar.git", + ), + ], + ) + setup_scratch(config, install_timeout=120.0) + + mock_ensure_repo.assert_has_calls( + [ + call("http://example.com/foo.git", directory_path / "foo"), + call("http://example.com/bar.git", directory_path / "bar"), + ] + ) + + mock_scratch_install.assert_has_calls( + [ + call(directory_path / "foo", timeout=120.0), + call(directory_path / "bar", timeout=120.0), + ] + ) + + +@patch("blueapi.cli.scratch.ensure_repo") +@patch("blueapi.cli.scratch.scratch_install") +def test_setup_scratch_continues_after_failure( + mock_scratch_install: Mock, + mock_ensure_repo: Mock, + directory_path: Path, +): + config = ScratchConfig( + root=directory_path, + repositories=[ + ScratchRepository( + name="foo", + remote_url="http://example.com/foo.git", + ), + ScratchRepository( + name="bar", + remote_url="http://example.com/bar.git", + ), + ScratchRepository( + name="baz", + remote_url="http://example.com/baz.git", + ), + ], + ) + mock_ensure_repo.side_effect = [None, RuntimeError("bar"), None] + with pytest.raises(RuntimeError, match="bar"): + setup_scratch(config) diff --git a/tests/example_yaml/scratch.yaml b/tests/example_yaml/scratch.yaml new file mode 100644 index 000000000..b64dc715f --- /dev/null +++ b/tests/example_yaml/scratch.yaml @@ -0,0 +1,5 @@ +scratch: + root: /tmp + repositories: + - name: dodal + remote_url: https://github.com/DiamondLightSource/dodal.git diff --git a/tests/test_cli.py b/tests/test_cli.py index 074efb732..31b6c0c74 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2,6 +2,7 @@ from collections.abc import Mapping from dataclasses import dataclass from io import StringIO +from pathlib import Path from textwrap import dedent from unittest.mock import MagicMock, Mock, patch @@ -16,6 +17,7 @@ from blueapi.cli.cli import main from blueapi.cli.event_bus_client import BlueskyRemoteError from blueapi.cli.format import OutputFormat +from blueapi.config import ScratchConfig, ScratchRepository from blueapi.core.bluesky_types import Plan from blueapi.service.handler import Handler, teardown_handler from blueapi.service.model import ( @@ -643,3 +645,24 @@ def test_generic_base_model_formatting(): """) OutputFormat.JSON.display(obj, output) assert exp == output.getvalue() + + +@patch("blueapi.cli.cli.setup_scratch") +def test_init_scratch_calls_setup_scratch(mock_setup_scratch: Mock, runner: CliRunner): + expected_config = ScratchConfig( + root=Path("/tmp"), + repositories=[ + ScratchRepository( + name="dodal", + remote_url="https://github.com/DiamondLightSource/dodal.git", + ) + ], + ) + + result = runner.invoke( + main, + ["-c", "tests/example_yaml/scratch.yaml", "setup-scratch"], + input="\n", + ) + assert result.exit_code == 0 + mock_setup_scratch.assert_called_once_with(expected_config)