Merge pull request #8254 from OpenMined/config-dockerfile

Build custom worker with config

.bumpversion.cfg

@@ -38,6 +38,8 @@ first_value = 1
 
 [bumpversion:file:packages/grid/VERSION]
 
+[bumpversion:file:packages/grid/backend/worker_cpu.dockerfile]
+
 [bumpversion:file:packages/grid/frontend/package.json]
 
 [bumpversion:file:packages/grid/helm/syft/Chart.yaml]

.github/workflows/cd-syft.yml

@@ -129,17 +129,17 @@ jobs:
           cache-from: type=registry,ref=openmined/grid-backend:cache-${{ steps.release_metadata.outputs.short_release_platform }}
           cache-to: type=registry,ref=openmined/grid-backend:cache-${{ steps.release_metadata.outputs.short_release_platform }},mode=max
 
-      - name: Build and push `syft-base-cpu` image to DockerHub
-        id: syft-base-cpu-build
-        uses: docker/build-push-action@v5
-        with:
-          context: ./packages
-          file: ./packages/grid/backend/syft_base_cpu.dockerfile
-          platforms: ${{ steps.release_metadata.outputs.release_platform }}
-          target: backend
-          outputs: type=image,name=openmined/syft-base-cpu,push-by-digest=true,name-canonical=true,push=true
-          cache-from: type=registry,ref=openmined/syft-base-cpu:cache-${{ steps.release_metadata.outputs.short_release_platform }}
-          cache-to: type=registry,ref=openmined/syft-base-cpu:cache-${{ steps.release_metadata.outputs.short_release_platform }},mode=max
+      # - name: Build and push `syft-base-cpu` image to DockerHub
+      #   id: syft-base-cpu-build
+      #   uses: docker/build-push-action@v5
+      #   with:
+      #     context: ./packages
+      #     file: ./packages/grid/backend/worker_cpu.dockerfile
+      #     platforms: ${{ steps.release_metadata.outputs.release_platform }}
+      #     target: backend
+      #     outputs: type=image,name=openmined/syft-base-cpu,push-by-digest=true,name-canonical=true,push=true
+      #     cache-from: type=registry,ref=openmined/syft-base-cpu:cache-${{ steps.release_metadata.outputs.short_release_platform }}
+      #     cache-to: type=registry,ref=openmined/syft-base-cpu:cache-${{ steps.release_metadata.outputs.short_release_platform }},mode=max
 
       - name: Export digest for grid-backend
         run: |

packages/grid/backend/backend.dockerfile

@@ -111,7 +111,7 @@ ENV PATH=$PATH:$HOME/.local/bin \
 COPY --chown=$USER_GRP --from=syft_deps $HOME/.local $HOME/.local
 
 # copy grid
-COPY --chown=$USER_GRP grid/backend/grid ./grid
+COPY --chown=$USER_GRP grid/backend/grid grid/backend/worker_cpu.dockerfile ./grid/
 
 # copy syft
 COPY --chown=$USER_GRP syft/ ./syft/

packages/grid/backend/grid/start.sh

@@ -17,17 +17,11 @@ ls -lisa $HOME/data
 ls -lisa $APPDIR/syft/
 ls -lisa $APPDIR/grid/
 
-MARKER="telemetry,data_science"
-if [[ ${SYFT_BASE_IMAGE} == "True" ]];
-then
-    MARKER="telemetry"
-fi
-
 if [[ ${DEV_MODE} == "True" ]];
 then
     echo "DEV_MODE Enabled"
     RELOAD="--reload"
-    pip install --user -e "$APPDIR/syft[${MARKER}]"
+    pip install --user -e "$APPDIR/syft[telemetry,data_science]"
 fi
 
 set +e

packages/grid/backend/worker_cpu.dockerfile

@@ -0,0 +1,27 @@
+# Syft Worker
+# Build as-is to create a base worker image
+# Build with args to create a custom worker image
+
+# NOTE: This dockerfile will be built inside a grid-backend container in PROD
+# Hence COPY will not work the same way in DEV vs. PROD
+
+# FIXME: Due to dependency on grid-backend base, python can only be changed from 3.11 to 3.11-dev
+# Later we'd want to uninstall old python, and then install a new python runtime...
+# ... but pre-built syft deps may break!
+
+FROM openmined/grid-backend:0.8.3-beta.2
+
+ARG PYTHON_VERSION="3.11"
+ARG SYSTEM_PACKAGES=""
+ARG PIP_PACKAGES="pip --dry-run"
+ARG CUSTOM_CMD='echo "No custom commands passed"'
+
+# Worker specific environment variables go here
+ENV SYFT_WORKER="true"
+
+RUN --mount=type=cache,target=/var/cache/apk,sharing=locked \
+    --mount=type=cache,target=$HOME/.cache/pip,sharing=locked \
+    apk update && \
+    apk add ${SYSTEM_PACKAGES} && \
+    pip install --user ${PIP_PACKAGES} && \
+    bash -c "$CUSTOM_CMD"

packages/syft/setup.cfg

@@ -61,6 +61,8 @@ syft =
     # numpy and pandas are ML packages but are needed in many places througout the codebase
     numpy>=1.23.5,<=1.24.4
     pandas==1.5.3
+    docker==6.1.3
+    PyYAML>=6.0,<7
 
 
 install_requires =

packages/syft/src/syft/custom_worker/builder.py

@@ -0,0 +1,98 @@
+# stdlib
+import os.path
+from pathlib import Path
+from typing import Tuple
+
+# third party
+import docker
+
+# relative
+from .config import CustomWorkerConfig
+
+
+class CustomWorkerBuilder:
+    TYPE_CPU = "cpu"
+    TYPE_GPU = "gpu"
+
+    DOCKERFILE_PROD_PATH = os.path.expandvars("$APPDIR/grid/")
+    DOCKERFILE_DEV_PATH = "../../../../../grid/backend/"
+
+    CUSTOM_IMAGE_PREFIX = "custom-worker"
+
+    BUILD_MAX_WAIT = 30 * 60
+
+    def build_image(self, config: CustomWorkerConfig) -> None:
+        """
+        Builds a Docker image for the custom worker based on the provided configuration.
+        Args:
+            config (CustomImageConfig): The configuration for building the Docker image.
+        Returns:
+            bool: True if the image was built successfully, raises Exception otherwise.
+        """
+
+        # remove once GPU is supported
+        if config.build.gpu:
+            raise Exception("GPU custom worker is not supported yet")
+
+        type = self.TYPE_GPU if config.build.gpu else self.TYPE_CPU
+
+        contextdir, dockerfile = self.find_worker_ctx(type)
+
+        imgtag = config.get_signature()[:8]
+
+        build_args = {
+            "PYTHON_VERSION": config.build.python_version,
+            "SYSTEM_PACKAGES": config.build.merged_system_pkgs(),
+            "PIP_PACKAGES": config.build.merged_python_pkgs(),
+            "CUSTOM_CMD": config.build.merged_custom_cmds(),
+        }
+
+        print(
+            f"Building dockerfile={dockerfile} "
+            f"in context={contextdir} "
+            f"with args={build_args}"
+        )
+
+        try:
+            client = docker.from_env()
+
+            # TODO: Push logs to mongo/seaweed?
+            client.images.build(
+                path=str(contextdir),
+                dockerfile=dockerfile,
+                pull=True,
+                tag=f"{self.CUSTOM_IMAGE_PREFIX}-{type}:{imgtag}",
+                timeout=self.BUILD_MAX_WAIT,
+                buildargs=build_args,
+            )
+
+            return
+        except docker.errors.BuildError as e:
+            raise e
+        except docker.errors.APIError as e:
+            raise e
+        except Exception as e:
+            raise e
+
+    def find_worker_ctx(self, type: str) -> Tuple[Path, str]:
+        """
+        Find the Worker Dockerfile and it's context path
+        - PROD will be in `$APPDIR/grid/`
+        - DEV will be in `packages/grid/backend`
+        - In both the cases context dir does not matter (unless we're calling COPY)
+
+        Args:
+            type (str): The type of worker.
+        Returns:
+            Path: The path to the Dockerfile.
+        """
+        filename = f"worker_{type}.dockerfile"
+        lookup_paths = [
+            Path(self.DOCKERFILE_PROD_PATH, filename).resolve(),
+            Path(__file__, self.DOCKERFILE_DEV_PATH, filename).resolve(),
+        ]
+        for path in lookup_paths:
+            if path.exists():
+                return path.parent, filename
+
+        raise FileNotFoundError(f"Cannot find the {filename}")