Add support for submitting a prebuilt image

.pre-commit-config.yaml

@@ -82,7 +82,6 @@ repos:
     # files: "^notebooks/(api|tutorials|admin)"
     hooks:
       - id: nbqa-isort
-      - id: nbqa-black
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.

notebooks/admin/Custom API + Custom Worker.ipynb

@@ -114,8 +114,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "submit_result = domain_client.api.services.worker_image.submit_dockerfile(\n",
-    "    docker_config=docker_config\n",
+    "submit_result = domain_client.api.services.worker_image.submit_container_image(\n",
+    "    worker_config=docker_config\n",
     ")\n",
     "submit_result"
    ]

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

@@ -229,8 +229,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "submit_result = domain_client.api.services.worker_image.submit_dockerfile(\n",
-    "    docker_config=docker_config\n",
+    "submit_result = domain_client.api.services.worker_image.submit_container_image(\n",
+    "    worker_config=docker_config\n",
     ")"
    ]
   },
@@ -1097,8 +1097,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "submit_result = domain_client.api.services.worker_image.submit_dockerfile(\n",
-    "    docker_config=docker_config_2\n",
+    "submit_result = domain_client.api.services.worker_image.submit_container_image(\n",
+    "    worker_config=docker_config_2\n",
     ")\n",
     "submit_result"
    ]

notebooks/api/0.8/11-container-images-k8s.ipynb

@@ -265,8 +265,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "submit_result = domain_client.api.services.worker_image.submit_dockerfile(\n",
-    "    docker_config=docker_config\n",
+    "submit_result = domain_client.api.services.worker_image.submit_container_image(\n",
+    "    worker_config=docker_config\n",
     ")\n",
     "submit_result"
    ]
@@ -935,8 +935,8 @@
    "outputs": [],
    "source": [
     "submit_result = None\n",
-    "submit_result = domain_client.api.services.worker_image.submit_dockerfile(\n",
-    "    docker_config=docker_config_opendp\n",
+    "submit_result = domain_client.api.services.worker_image.submit_container_image(\n",
+    "    worker_config=docker_config_opendp\n",
     ")\n",
     "submit_result"
    ]

packages/syft/src/syft/__init__.py

@@ -26,6 +26,7 @@
 from .client.user_settings import UserSettings  # noqa: F401
 from .client.user_settings import settings  # noqa: F401
 from .custom_worker.config import DockerWorkerConfig  # noqa: F401
+from .custom_worker.config import PrebuiltWorkerConfig  # noqa: F401
 from .node.credentials import SyftSigningKey  # noqa: F401
 from .node.domain import Domain  # noqa: F401
 from .node.enclave import Enclave  # noqa: F401

packages/syft/src/syft/service/request/request.py

@@ -205,32 +205,42 @@ def _run(
             worker_image_service = context.node.get_service("SyftWorkerImageService")
 
             service_context = context.to_service_ctx()
-            result = worker_image_service.submit_dockerfile(
-                service_context, docker_config=self.config
+            result = worker_image_service.submit_container_image(
+                service_context, worker_config=self.config
             )
 
             if isinstance(result, SyftError):
                 return Err(result)
 
-            result = worker_image_service.stash.get_by_docker_config(
+            result = worker_image_service.stash.get_by_worker_config(
                 service_context.credentials, config=self.config
             )
 
             if result.is_err():
                 return Err(SyftError(message=f"{result.err()}"))
 
-            worker_image = result.ok()
+            if (worker_image := result.ok()) is None:
+                return Err(SyftError(message="The worker image does not exist."))
 
-            build_result = worker_image_service.build(
-                service_context,
-                image_uid=worker_image.id,
-                tag=self.tag,
-                registry_uid=self.registry_uid,
-                pull=self.pull_image,
-            )
+            build_success_message = "Image was pre-built."
+
+            if not worker_image.is_prebuilt:
+                build_result = worker_image_service.build(
+                    service_context,
+                    image_uid=worker_image.id,
+                    tag=self.tag,
+                    registry_uid=self.registry_uid,
+                    pull=self.pull_image,
+                )
 
-            if isinstance(build_result, SyftError):
-                return Err(build_result)
+                if isinstance(build_result, SyftError):
+                    return Err(build_result)
+
+                build_success_message = build_result.message
+
+            build_success = SyftSuccess(
+                message=f"Build result: {build_success_message}"
+            )
 
             if IN_KUBERNETES:
                 push_result = worker_image_service.push(
@@ -245,11 +255,11 @@ def _run(
 
                 return Ok(
                     SyftSuccess(
-                        message=f"Build Result: {build_result.message} \n Push Result: {push_result.message}"
+                        message=f"{build_success}\nPush result: {push_result.message}"
                     )
                 )
 
-            return Ok(build_result)
+            return Ok(build_success)
 
         except Exception as e:
             return Err(SyftError(message=f"Failed to create/build image: {e}"))
@@ -291,7 +301,7 @@ def _run(
             service_context: AuthedServiceContext = context.to_service_ctx()
 
             if self.config is not None:
-                result = worker_pool_service.image_stash.get_by_docker_config(
+                result = worker_pool_service.image_stash.get_by_worker_config(
                     service_context.credentials, self.config
                 )
                 if result.is_err():

packages/syft/src/syft/service/worker/utils.py

@@ -564,7 +564,7 @@ def create_default_image(
         image_identifier=SyftWorkerImageIdentifier.from_str(tag),
     )
 
-    result = image_stash.get_by_docker_config(
+    result = image_stash.get_by_worker_config(
         credentials=credentials,
         config=worker_config,
     )

packages/syft/src/syft/service/worker/worker_image_service.py

@@ -6,7 +6,7 @@
 import pydantic
 
 # relative
-from ...custom_worker.config import DockerWorkerConfig
+from ...custom_worker.config import WorkerConfig
 from ...custom_worker.k8s import IN_KUBERNETES
 from ...serde.serializable import serializable
 from ...store.document_store import DocumentStore
@@ -39,15 +39,15 @@ def __init__(self, store: DocumentStore) -> None:
         self.stash = SyftWorkerImageStash(store=store)
 
     @service_method(
-        path="worker_image.submit_dockerfile",
-        name="submit_dockerfile",
+        path="worker_image.submit_container_image",
+        name="submit_container_image",
         roles=DATA_OWNER_ROLE_LEVEL,
     )
-    def submit_dockerfile(
-        self, context: AuthedServiceContext, docker_config: DockerWorkerConfig
+    def submit_container_image(
+        self, context: AuthedServiceContext, worker_config: WorkerConfig
     ) -> SyftSuccess | SyftError:
         worker_image = SyftWorkerImage(
-            config=docker_config,
+            config=worker_config,
             created_by=context.credentials,
         )
         res = self.stash.set(context.credentials, worker_image)
@@ -278,14 +278,14 @@ def get_by_uid(
         roles=DATA_SCIENTIST_ROLE_LEVEL,
     )
     def get_by_config(
-        self, context: AuthedServiceContext, docker_config: DockerWorkerConfig
+        self, context: AuthedServiceContext, worker_config: WorkerConfig
     ) -> SyftWorkerImage | SyftError:
-        res = self.stash.get_by_docker_config(
-            credentials=context.credentials, config=docker_config
+        res = self.stash.get_by_worker_config(
+            credentials=context.credentials, config=worker_config
         )
         if res.is_err():
             return SyftError(
-                message=f"Failed to get image with docker config {docker_config}. Error: {res.err()}"
+                message=f"Failed to get image with docker config {worker_config}. Error: {res.err()}"
             )
         image: SyftWorkerImage = res.ok()
         return image

packages/syft/src/syft/service/worker/worker_image_stash.py

@@ -48,7 +48,7 @@ def set(
         )
 
         if isinstance(obj.config, DockerWorkerConfig):
-            result = self.get_by_docker_config(
+            result = self.get_by_worker_config(
                 credentials=credentials, config=obj.config
             )
             if result.is_ok() and result.ok() is not None:
@@ -62,8 +62,8 @@ def set(
             ignore_duplicates=ignore_duplicates,
         )
 
-    def get_by_docker_config(
-        self, credentials: SyftVerifyKey, config: DockerWorkerConfig
+    def get_by_worker_config(
+        self, credentials: SyftVerifyKey, config: WorkerConfig
     ) -> Result[SyftWorkerImage | None, str]:
         qks = QueryKeys(qks=[WorkerConfigPK.with_obj(config)])
         return self.query_one(credentials=credentials, qks=qks)

packages/syft/src/syft/service/worker/worker_pool_service.py

@@ -257,7 +257,7 @@ def create_image_and_pool_request(
             return SyftError(message="Registry UID is required in Kubernetes mode.")
 
         # Check if an image already exists for given docker config
-        search_result = self.image_stash.get_by_docker_config(
+        search_result = self.image_stash.get_by_worker_config(
             credentials=context.credentials, config=config
         )
 

packages/syft/tests/syft/worker_pool/worker_pool_service_test.py

@@ -1,9 +1,12 @@
 # third party
 from faker import Faker
+import pytest
 
 # syft absolute
 import syft as sy
 from syft.custom_worker.config import DockerWorkerConfig
+from syft.custom_worker.config import PrebuiltWorkerConfig
+from syft.custom_worker.config import WorkerConfig
 from syft.node.worker import Worker
 from syft.service.request.request import CreateCustomWorkerPoolChange
 from syft.service.response import SyftSuccess
@@ -13,8 +16,35 @@
 # relative
 from ..request.request_code_accept_deny_test import get_ds_client
 
+PREBUILT_IMAGE_TAG = f"openmined/grid-backend:{sy.__version__}"
 
-def test_create_image_and_pool_request_accept(faker: Faker, worker: Worker):
+CUSTOM_DOCKERFILE = f"""
+FROM {PREBUILT_IMAGE_TAG}
+
+RUN pip install recordlinkage
+"""
+
+CUSTOM_IMAGE_TAG = "openmined/custom-worker-recordlinkage:latest"
+
+WORKER_CONFIG_TEST_CASES_WITH_N_IMAGES = [
+    (
+        CUSTOM_IMAGE_TAG,
+        DockerWorkerConfig(dockerfile=CUSTOM_DOCKERFILE),
+        2,  # total number of images.
+        # 2 since we pull a pre-built image (1) as the base image to build a custom image (2)
+    ),
+    (PREBUILT_IMAGE_TAG, PrebuiltWorkerConfig(tag=PREBUILT_IMAGE_TAG), 1),
+]
+
+WORKER_CONFIG_TEST_CASES = [
+    test_case[:2] for test_case in WORKER_CONFIG_TEST_CASES_WITH_N_IMAGES
+]
+
+
+@pytest.mark.parametrize("docker_tag,worker_config", WORKER_CONFIG_TEST_CASES)
+def test_create_image_and_pool_request_accept(
+    faker: Faker, worker: Worker, docker_tag: str, worker_config: WorkerConfig
+) -> None:
     """
     Test the functionality of `SyftWorkerPoolService.create_image_and_pool_request`
     when the request is accepted
@@ -25,22 +55,15 @@ def test_create_image_and_pool_request_accept(faker: Faker, worker: Worker):
     assert root_client.credentials != ds_client.credentials
 
     # the DS makes a request to create an image and a pool based on the image
-    custom_dockerfile = f"""
-        FROM openmined/grid-backend:{sy.__version__}
-
-        RUN pip install recordlinkage
-    """
-    docker_config = DockerWorkerConfig(dockerfile=custom_dockerfile)
-    docker_tag = "openmined/custom-worker-recordlinkage:latest"
     request = ds_client.api.services.worker_pool.create_image_and_pool_request(
         pool_name="recordlinkage-pool",
         num_workers=2,
         tag=docker_tag,
-        config=docker_config,
+        config=worker_config,
         reason="I want to do some more cool data science with PySyft and Recordlinkage",
     )
     assert len(request.changes) == 2
-    assert request.changes[0].config == docker_config
+    assert request.changes[0].config == worker_config
     assert request.changes[1].num_workers == 2
     assert request.changes[1].pool_name == "recordlinkage-pool"
 
@@ -61,7 +84,17 @@ def test_create_image_and_pool_request_accept(faker: Faker, worker: Worker):
     assert len(launched_pool.worker_list) == 2
 
 
-def test_create_pool_request_accept(faker: Faker, worker: Worker):
+@pytest.mark.parametrize(
+    "docker_tag,worker_config,n_images",
+    WORKER_CONFIG_TEST_CASES_WITH_N_IMAGES,
+)
+def test_create_pool_request_accept(
+    faker: Faker,
+    worker: Worker,
+    docker_tag: str,
+    worker_config: WorkerConfig,
+    n_images: int,
+) -> None:
     """
     Test the functionality of `SyftWorkerPoolService.create_pool_request`
     when the request is accepted
@@ -72,29 +105,27 @@ def test_create_pool_request_accept(faker: Faker, worker: Worker):
     assert root_client.credentials != ds_client.credentials
 
     # the DO submits the docker config to build an image
-    custom_dockerfile_str = f"""
-        FROM openmined/grid-backend:{sy.__version__}
-
-        RUN pip install opendp
-    """
-    docker_config = DockerWorkerConfig(dockerfile=custom_dockerfile_str)
-    submit_result = root_client.api.services.worker_image.submit_dockerfile(
-        docker_config=docker_config
+    submit_result = root_client.api.services.worker_image.submit_container_image(
+        worker_config=worker_config
     )
     assert isinstance(submit_result, SyftSuccess)
-    assert len(root_client.images.get_all()) == 2
+    assert len(root_client.images.get_all()) == n_images
 
     # The root client builds the image
-    worker_image: SyftWorkerImage = root_client.images[1]
-    docker_tag = "openmined/custom-worker-opendp:latest"
-    docker_build_result = root_client.api.services.worker_image.build(
-        image_uid=worker_image.id,
-        tag=docker_tag,
+    worker_image: SyftWorkerImage = root_client.api.services.worker_image.get_by_config(
+        worker_config
     )
-    # update the worker image variable after the image was built
-    worker_image: SyftWorkerImage = root_client.images[1]
-    assert isinstance(docker_build_result, SyftSuccess)
-    assert worker_image.image_identifier.repo_with_tag == docker_tag
+    if not worker_image.is_prebuilt:
+        docker_build_result = root_client.api.services.worker_image.build(
+            image_uid=worker_image.id,
+            tag=docker_tag,
+        )
+        # update the worker image variable after the image was built
+        worker_image: SyftWorkerImage = (
+            root_client.api.services.worker_image.get_by_config(worker_config)
+        )
+        assert isinstance(docker_build_result, SyftSuccess)
+        assert worker_image.image_identifier.repo_with_tag == docker_tag
 
     # The DS client submits a request to create a pool from an existing image
     request = ds_client.api.services.worker_pool.pool_creation_request(
@@ -113,3 +144,21 @@ def test_create_pool_request_accept(faker: Faker, worker: Worker):
     launched_pool = root_client.worker_pools["opendp-pool"]
     assert isinstance(launched_pool, WorkerPool)
     assert len(launched_pool.worker_list) == 3
+
+
+WORKER_CONFIGS = [test_case[1] for test_case in WORKER_CONFIG_TEST_CASES]
+
+
+@pytest.mark.parametrize("worker_config", WORKER_CONFIGS)
+def test_get_by_worker_config(
+    worker: Worker,
+    worker_config: WorkerConfig,
+) -> None:
+    root_client = worker.root_client
+    for config in WORKER_CONFIGS:
+        root_client.api.services.worker_image.submit_container_image(
+            worker_config=config
+        )
+
+    worker_image = root_client.api.services.worker_image.get_by_config(worker_config)
+    assert worker_image.config == worker_config