Merge pull request #8370 from OpenMined/yash/syft-registry

Syft Image Registry

docs/requirements.txt

@@ -4,9 +4,9 @@ markupsafe==2.0.1
 pydata-sphinx-theme==0.7.2
 pygments>=2.15.0 # not directly required, pinned by Snyk to avoid a vulnerability
 requests>=2.31.0 # not directly required, pinned by Snyk to avoid a vulnerability
+setuptools>=65.5.1 # not directly required, pinned by Snyk to avoid a vulnerability
 sphinx==4.3.0
 sphinx-autoapi==1.8.4
 sphinx-code-include==1.1.1
 sphinx-copybutton==0.4.0
 sphinx-panels==0.6.0
-setuptools>=65.5.1 # not directly required, pinned by Snyk to avoid a vulnerability

notebooks/api/0.8/10-container-images.ipynb

@@ -22,7 +22,7 @@
     "import os\n",
     "sy.requires(SYFT_VERSION)\n",
     "from syft.service.worker.worker_image import SyftWorkerImage\n",
-    "from syft.service.worker.worker_pool import WorkerStatus, SyftWorker\n",
+    "from syft.service.worker.image_registry import SyftImageRegistry\n",
     "from syft.custom_worker.config import DockerWorkerConfig"
    ]
   },
@@ -57,6 +57,14 @@
     "domain_client = domain.login(email=\"[email protected]\", password=\"changethis\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "3c7a124a",
+   "metadata": {},
+   "source": [
+    "#### Submit Dockerfile"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -160,6 +168,76 @@
     "assert isinstance(workerimage, SyftWorkerImage)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "91a66871",
+   "metadata": {},
+   "source": [
+    "#### Create Registry"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cde8bfff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_add_res = domain_client.api.services.image_registry.add(\"localhost:5678\")\n",
+    "image_add_res"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "82321b35",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert isinstance(image_add_res, sy.SyftSuccess), str(image_add_res)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3d4a4c33",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "images = domain_client.api.services.image_registry.get_all()\n",
+    "assert len(images) == 1\n",
+    "images"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "22f6e2f6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "local_registry = images[0]\n",
+    "local_registry"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cb9664ca",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert isinstance(local_registry, SyftImageRegistry)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "637a9596",
+   "metadata": {},
+   "source": [
+    "#### Build Image"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -168,7 +246,7 @@
    "outputs": [],
    "source": [
     "docker_tag = \"openmined/test-nginx:0.7.8\"\n",
-    "docker_build_res = domain_client.api.services.worker_image.build(uid=workerimage.id, tag=docker_tag)"
+    "docker_build_res = domain_client.api.services.worker_image.build(image_uid=workerimage.id, tag=docker_tag)"
    ]
   },
   {
@@ -188,7 +266,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "assert isinstance(docker_build_res, sy.SyftSuccess)"
+    "assert isinstance(docker_build_res, sy.SyftSuccess), str(docker_build_res)"
    ]
   },
   {
@@ -292,6 +370,14 @@
     "        return None"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "f5007073",
+   "metadata": {},
+   "source": [
+    "#### Create Worker Pool From Image"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -542,7 +628,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.11.2"
   }
  },
  "nbformat": 4,

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

@@ -1,43 +1,82 @@
 # stdlib
 import contextlib
+import io
 import os.path
 from pathlib import Path
+from typing import Any
+from typing import Iterable
+from typing import Optional
 from typing import Tuple
 
 # third party
 import docker
+from docker.models.images import Image
 
 # relative
 from .config import CustomWorkerConfig
+from .config import DockerWorkerConfig
+from .config import WorkerConfig
 
 
 class CustomWorkerBuilder:
     TYPE_CPU = "cpu"
     TYPE_GPU = "gpu"
 
-    DOCKERFILE_PROD_PATH = os.path.expandvars("$APPDIR/grid/")
-    DOCKERFILE_DEV_PATH = "../../../../../grid/backend/"
+    TEMPLATE_DIR_PROD = os.path.expandvars("$APPDIR/grid/")
+    TEMPLATE_DIR_DEV = "../../../../../grid/backend/"
 
     CUSTOM_IMAGE_PREFIX = "custom-worker"
 
     BUILD_MAX_WAIT = 30 * 60
 
-    def build_image(self, config: CustomWorkerConfig) -> None:
+    def build_image(
+        self,
+        config: WorkerConfig,
+        tag: str = None,
+        **kwargs: Any,
+    ) -> Tuple[Image, Iterable[str]]:
         """
-        Builds a Docker image for the custom worker based on the provided configuration.
+        Builds a Docker image from the given configuration.
         Args:
-            config (CustomImageConfig): The configuration for building the Docker image.
-        Returns:
-            bool: True if the image was built successfully, raises Exception otherwise.
+            config (WorkerConfig): The configuration for building the Docker image.
+            tag (str): The tag to use for the image.
+        """
+
+        if isinstance(config, DockerWorkerConfig):
+            return self._build_dockerfile(config, tag, **kwargs)
+        elif isinstance(config, CustomWorkerConfig):
+            return self._build_template(config, **kwargs)
+        else:
+            raise TypeError("Unknown worker config type")
+
+    def push_image(self, tag: str, **kwargs: Any) -> str:
+        """
+        Pushes a Docker image to the given repo.
+        Args:
+            repo (str): The repo to push the image to.
+            tag (str): The tag to use for the image.
         """
 
+        return self._push_image(tag, **kwargs)
+
+    def _build_dockerfile(self, config: DockerWorkerConfig, tag: str, **kwargs):
+        print("Building with provided dockerfile")
+
+        # convert string to file-like object
+        file_obj = io.BytesIO(config.dockerfile.encode("utf-8"))
+        return self._build_image(fileobj=file_obj, tag=tag, **kwargs)
+
+    def _build_template(self, config: CustomWorkerConfig, **kwargs: Any):
+        # Builds a Docker pre-made CPU/GPU image template using a CustomWorkerConfig
+        print("Building with dockerfule template")
+
         # 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)
+        contextdir, dockerfile = self._find_template_dir(type)
 
         imgtag = config.get_signature()[:8]
 
@@ -54,27 +93,41 @@ def build_image(self, config: CustomWorkerConfig) -> None:
             f"with args={build_args}"
         )
 
-        try:
-            # TODO: Push logs to mongo/seaweed?
-            with contextlib.closing(docker.from_env()) as client:
-                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 self._build_image(
+            tag=f"{self.CUSTOM_IMAGE_PREFIX}-{type}:{imgtag}",
+            path=str(contextdir),
+            dockerfile=dockerfile,
+            buildargs=build_args,
+        )
+
+    def _build_image(self, tag: str, **build_opts) -> Tuple[Image, Iterable]:
+        # Core docker build call. Func signature should match with Docker SDK's BuildApiMixin
+        with contextlib.closing(docker.from_env()) as client:
+            image = client.images.build(
+                tag=tag,
+                pull=True,
+                timeout=self.BUILD_MAX_WAIT,
+                **build_opts,
+            )
+            return image
+
+    def _push_image(
+        self,
+        tag: str,
+        registry_url: Optional[str] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+    ) -> str:
+        with contextlib.closing(docker.from_env()) as client:
+            if registry_url and username and password:
+                client.login(
+                    username=username, password=password, registry=registry_url
                 )
 
-            return
-        except docker.errors.BuildError as e:
-            raise e
-        except docker.errors.APIError as e:
-            raise e
-        except Exception as e:
-            raise e
+            result = client.images.push(repository=tag)
+            return result
 
-    def find_worker_ctx(self, type: str) -> Tuple[Path, str]:
+    def _find_template_dir(self, type: str) -> Tuple[Path, str]:
         """
         Find the Worker Dockerfile and it's context path
         - PROD will be in `$APPDIR/grid/`
@@ -88,8 +141,8 @@ def find_worker_ctx(self, type: str) -> Tuple[Path, str]:
         """
         filename = f"worker_{type}.dockerfile"
         lookup_paths = [
-            Path(self.DOCKERFILE_PROD_PATH, filename).resolve(),
-            Path(__file__, self.DOCKERFILE_DEV_PATH, filename).resolve(),
+            Path(self.TEMPLATE_DIR_PROD, filename).resolve(),
+            Path(__file__, self.TEMPLATE_DIR_DEV, filename).resolve(),
         ]
         for path in lookup_paths:
             if path.exists():

packages/syft/src/syft/custom_worker/config.py

@@ -121,3 +121,6 @@ def __eq__(self, __value: object) -> bool:
 
     def __hash__(self) -> int:
         return hash(self.dockerfile)
+
+    def __str__(self) -> str:
+        return self.dockerfile

packages/syft/src/syft/node/node.py

@@ -95,6 +95,7 @@
 from ..service.user.user_roles import ServiceRole
 from ..service.user.user_service import UserService
 from ..service.user.user_stash import UserStash
+from ..service.worker.image_registry_service import SyftImageRegistryService
 from ..service.worker.utils import DEFAULT_WORKER_IMAGE_TAG
 from ..service.worker.utils import DEFAULT_WORKER_POOL_NAME
 from ..service.worker.utils import create_default_image
@@ -317,6 +318,7 @@ def __init__(
                 MigrateStateService,
                 SyftWorkerImageService,
                 SyftWorkerPoolService,
+                SyftImageRegistryService,
             ]
             if services is None
             else services
@@ -870,6 +872,7 @@ def _construct_services(self):
                 MigrateStateService,
                 SyftWorkerImageService,
                 SyftWorkerPoolService,
+                SyftImageRegistryService,
             ]
 
             if OBLV:
@@ -1434,7 +1437,7 @@ def create_default_worker_pool(node: Node) -> Optional[SyftError]:
 
     # Build the Image for given tag
     result = image_build_method(
-        context, uid=default_image.id, tag=DEFAULT_WORKER_IMAGE_TAG
+        context, image_uid=default_image.id, tag=DEFAULT_WORKER_IMAGE_TAG
     )
 
     if isinstance(result, SyftError):

packages/syft/src/syft/protocol/protocol_version.json

@@ -957,7 +957,7 @@
       "SyftWorkerImage": {
         "1": {
           "version": 1,
-          "hash": "abfca84237a47aab399970a0a1ad747bb7c0f325944b6647be38050e1e0ac697",
+          "hash": "2a9585b6a286e24f1a9f3f943d0128730cf853edc549184dc1809d19e1eec54b",
           "action": "add"
         }
       },
@@ -1061,6 +1061,13 @@
           "hash": "3644c2caf4c01499e23b4e94d6a41fed0e39ce63c7fc5318c4e205cae9555308",
           "action": "add"
         }
+      },
+      "SyftImageRegistry": {
+        "1": {
+          "version": 1,
+          "hash": "dc83910c91947e3d9eaa3e6f8592237448f0408668c7cca80450b5fcd54722e1",
+          "action": "add"
+        }
       }
     }
   }