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 <[email protected]>

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"]

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 }}

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

pyproject.toml

@@ -17,17 +17,18 @@ 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",
     "click",
     "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"

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():

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")

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):