Merge pull request #203 from NERSC/improve-docker-bake

Improve local operator dev workflow with docker bake

Author: swelborn

backend/core/interactem/core/models/messages.py

@@ -31,7 +31,7 @@ class MessageSubject(str, Enum):
 
 class MessageHeader(BaseModel):
     subject: MessageSubject
-    meta: dict[str, Any] = {}
+    meta: bytes | dict[str, Any] = b"{}"
     tracking: (
         list[
             OperatorTrackingMetadata

backend/operators/interactem/operators/messengers/zmq.py

@@ -98,35 +98,42 @@ async def recv(self) -> BytesMessage | None:
         all_messages: list[BytesMessage] = []
 
         for id, msg_parts in id_msgs:
+            raw_meta = None
+
             if len(msg_parts) == 2:
                 _header, _data = msg_parts
             elif len(msg_parts) == 3:
-                _, _header, _data = msg_parts
+                _header, raw_meta, _data = msg_parts
             else:
-                logger.error(
-                    "Received an unexpected number of message parts: %s", len(msg_parts)
-                )
                 return None
 
+            # Decode header JSON
             if isinstance(_header, zmq.Message):
                 header = MessageHeader.model_validate_json(_header.bytes)
             elif isinstance(_header, bytes):
                 header = MessageHeader.model_validate_json(_header)
             else:
-                logger.error("Received an unexpected message type: %s", type(_header))
                 continue
 
+            # If meta was sent separately as bytes, restore it
+            if raw_meta is not None:
+                if isinstance(raw_meta, zmq.Message):
+                    header.meta = raw_meta.bytes
+                else:
+                    header.meta = raw_meta
+
             if header.subject != MessageSubject.BYTES:
                 logger.error(
                     "Received an unexpected message subject: %s", header.subject
                 )
                 continue
 
-            msg = (
-                BytesMessage(header=header, data=_data.bytes)
-                if isinstance(_data, zmq.Message)
-                else BytesMessage(header=header, data=_data)
-            )
+            if isinstance(_data, zmq.Message):
+                msg = BytesMessage(header=header, data=_data.bytes)
+            else:
+                msg = BytesMessage(header=header, data=_data)
+
+            # Tracking update
             if header.tracking is not None:
                 header.tracking.append(
                     InputPortTrackingMetadata(
@@ -157,21 +164,32 @@ async def send(self, message: BytesMessage):
         msg_futures = []
         data = message.data
         should_copy_header = len(self.output_ports) > 1
-        time_before_send = datetime.now()
         for socket in self.output_sockets.values():
-            # need to copy to send same message to multple sockets
-            # otherwise we will continue to append to the same header
             if should_copy_header:
-                header = message.header.model_copy(deep=True)
+                header = message.header.model_copy()
             else:
                 header = message.header
+
             if header.tracking is not None:
+                time_before_send = datetime.now()
                 meta = OutputPortTrackingMetadata(
                     id=socket.info.port_id, time_before_send=time_before_send
                 )
                 header.tracking.append(meta)
-            header = header.model_dump_json().encode()
-            things_to_send = [header, data]
+
+            # --- Separate meta if it's bytes ---
+            raw_meta = None
+            if isinstance(header.meta, bytes):
+                raw_meta = header.meta
+                header_json = header.model_dump_json(exclude={"meta"}).encode()
+            else:
+                header_json = header.model_dump_json().encode()
+
+            # --- Build multipart message ---
+            if raw_meta is not None:
+                things_to_send = [header_json, raw_meta, data]
+            else:
+                things_to_send = [header_json, data]
 
             msg_futures.append(self._send_and_update_metrics(socket, things_to_send))
         # TODO: look into creating tasks

operators/.env.example

@@ -0,0 +1,4 @@
+PODMAN_SERVICE_URI=unix:///<your-socket>
+BAKE_FILE=./docker-bake.hcl
+VERBOSE=true
+REGISTRY=host.containers.internal:5001/ghcr.io/nersc/interactem

operators/.gitignore

@@ -0,0 +1 @@
+.env

operators/README.md

@@ -0,0 +1,31 @@
+# Building operators
+
+## Building locally
+
+### MacOS
+
+1. Get `docker desktop` and `podman desktop`.
+1. Set up docker local docker registry by running:
+
+    ```sh
+    docker run -d -p 5001:5000 --restart always --name docker-registry registry:3
+    ```
+
+1. Use [bake.sh](./bake.sh) to build all containers with docker. This includes base image, operator, and distiller-streaming. You should do the following
+
+    ```sh
+    ./bake --push-local --build-base
+    ```
+
+    This will push everything to the local registry, instead of pushing up to GitHub packages. You can also omit `--build-base` to avoid building base images for faster iteration.
+
+1. Set your `.env` file in the operator directory to have the correct podman socket (see [`.env.example`](./.env.example))
+1. Use poetry environment from root directory [pyproject.toml](../pyproject.toml) and run the following:
+
+    ```sh
+    poetry run python pull_images_from_bake.py
+    ```
+
+    This will pull local registry images into podman and tag them appropriately with [`pull_images_from_bake.py`](./pull_images_from_bake.py). This also runs at the end of [`bake.sh`](./bake.sh) if `--pull-local` is given.
+
+The [build_all.sh](./build_all.sh) script was used before, but it is cumbersome so I am not updating it.

operators/bake.sh

@@ -1,60 +1,131 @@
 #!/bin/bash
+set -e
+
 SCRIPT_DIR=$(dirname "$0")
 REPO_ROOT_DIR=$(git rev-parse --show-toplevel)
 TAG=$(git rev-parse --short=6 HEAD)
+OPERATORS_DIR="$REPO_ROOT_DIR/operators"
+BAKE_FILE="$OPERATORS_DIR/docker-bake.hcl"
+OPERATOR_JSON="operator.json"
+VENV_DIR="$OPERATORS_DIR/distiller-streaming/.venv"
 
-cd $REPO_ROOT_DIR
+BUILD_BASE=false
+PUSH_LOCAL=false
+PUSH_REMOTE=false
+PULL_LOCAL=false
+TARGET=""
 
-# Build all images using Docker Bake
-docker buildx bake base --file $REPO_ROOT_DIR/operators/docker-bake.hcl
-docker buildx bake operator --file $REPO_ROOT_DIR/operators/docker-bake.hcl
-docker buildx bake distiller-streaming --file $REPO_ROOT_DIR/operators/docker-bake.hcl
-docker buildx bake operators --file $REPO_ROOT_DIR/operators/docker-bake.hcl --provenance=false
+# Parse args
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --build-base)
+            BUILD_BASE=true
+            shift
+            ;;
+        --push-local)
+            PUSH_LOCAL=true
+            VARS="$OPERATORS_DIR/vars-local.hcl"
+            shift
+            ;;
+        --pull-local)
+            PULL_LOCAL=true
+            shift
+            ;;
+        --push-remote)
+            PUSH_REMOTE=true
+            VARS="$OPERATORS_DIR/vars-prod.hcl"
+            shift
+            ;;
+        --target)
+            TARGET="$2"
+            shift 2
+            ;;
+        *)
+            shift
+            ;;
+    esac
+done
 
-if [ $? -ne 0 ]; then
-    echo "Failed to build images"
+if [ -z "$VARS" ]; then
+    echo "Error: Must specify either --push-local or --push-remote"
     exit 1
 fi
 
-generate_label_args() {
-    local label_args=""
-    # Process each operator directory
-    for dir in "$REPO_ROOT_DIR/operators"/*; do
-        if [ -d "$dir" ] && [ -f "$dir/operator.json" ] && [ -f "$dir/Containerfile" ]; then
-            op_name=$(basename "$dir")
+BASE_VARS="$OPERATORS_DIR/vars-base.hcl"
+
+cd "$REPO_ROOT_DIR"
+
+
+if $BUILD_BASE; then
+    echo "=== Building base images ==="
+    docker buildx bake base --file "$BAKE_FILE" \
+        --file "$BASE_VARS"
+
+    docker buildx bake operator --file "$BAKE_FILE" \
+        --file "$BASE_VARS"
 
-            # Get json
-            operator_json=$(jq -c . "$dir/operator.json" 2>/dev/null)
-            if [ $? -ne 0 ] || [ -z "$operator_json" ]; then
-                continue
+    docker buildx bake distiller-streaming --file "$BAKE_FILE" \
+        --file "$BASE_VARS"
+fi
+
+#
+# === Build operators ===
+#
+build_operators() {
+    local cmd=(docker buildx bake)
+    local has_labels=false
+    
+    if [ -n "$TARGET" ]; then
+        cmd+=("$TARGET")
+        cmd+=(--provenance=false)
+    else
+        cmd+=(operators)
+    fi
+
+    cmd+=(--file "$BAKE_FILE")
+    cmd+=(--file "$VARS")
+    
+    # Add labels from operator.json files
+    for dir in "$OPERATORS_DIR"/*; do
+        if [ -d "$dir" ] && [ -f "$dir/$OPERATOR_JSON" ] && [ -f "$dir/Containerfile" ]; then
+            op_name=$(basename "$dir")
+            operator_json_file=$(jq -c . "$dir/$OPERATOR_JSON" 2>/dev/null)
+            if [ $? -eq 0 ] && [ -n "$operator_json_file" ]; then
+                cmd+=(--set "${op_name}.labels.interactem.operator.spec=${operator_json_file}")
+                has_labels=true
             fi
-            # Add set argument using file instead of inline string
-            label_args="${label_args} --set ${op_name}.labels.interactem.operator.spec='${operator_json}'"
         fi
     done
 
-    echo "$label_args"
+    # Check if we found any labels
+    if [ "$has_labels" = false ]; then
+        echo "No operator.json files found or all are empty."
+        return 1
+    fi
+    
+    cmd+=(--push)
+    
+    # Execute the command
+    "${cmd[@]}"
 }
 
-LABEL_ARGS=$(generate_label_args)
-if [ -z "$LABEL_ARGS" ]; then
-    echo "No operator.json files found or all are empty."
-    exit 1
-fi
-
-# Create a temporary script because escaping json is hard
-TMP_SCRIPT="/tmp/docker_buildx_script.sh"
-cat > "$TMP_SCRIPT" << EOL
-#!/bin/bash
-docker buildx bake operators --file $REPO_ROOT_DIR/operators/docker-bake.hcl --push --provenance=false $LABEL_ARGS
-EOL
-
-chmod +x "$TMP_SCRIPT"
-"$TMP_SCRIPT"
+echo "=== Building operators ==="
+build_operators
 build_status=$?
-rm -f "$TMP_SCRIPT"
 
 if [ $build_status -ne 0 ]; then
     echo "Failed to build/push images"
     exit 1
-fi
+fi
+
+#
+# === Pull locally built images to podman ===
+#
+if $PUSH_LOCAL && $PULL_LOCAL; then
+    echo "=== Pulling images back from local registry ==="
+    cd $OPERATORS_DIR
+    source $VENV_DIR/bin/activate
+    poetry run python pull_images_from_bake.py
+fi
+
+echo "=== Done ==="

operators/beam-compensation/Containerfile

@@ -1,6 +1,7 @@
-FROM ghcr.io/nersc/interactem/operator
+FROM ghcr.io/nersc/interactem/distiller-streaming
 
-RUN pip install scipy ncempy stempy
+# had to add cffi because ncempy messes with that package when adding, and nkeys acts up...
+RUN poetry add ncempy cffi
 
 COPY ./run.py /app/run.py