Merge pull request #8802 from kiendang/prebuilt-image-worker

Add support for submitting a prebuilt image

.pre-commit-config.yaml

@@ -81,7 +81,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(\n",
+    "    worker_config=docker_config\n",
     ")\n",
     "submit_result"
    ]

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

@@ -227,8 +227,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(\n",
+    "    worker_config=docker_config\n",
     ")"
    ]
   },
@@ -1095,8 +1095,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(\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(\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(\n",
+    "    worker_config=docker_config_opendp\n",
     ")\n",
     "submit_result"
    ]

packages/syft/src/syft/__init__.py

@@ -25,6 +25,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/custom_worker/runner_k8s.py

@@ -13,7 +13,7 @@
 from .k8s import get_kr8s_client
 
 JSONPATH_AVAILABLE_REPLICAS = "{.status.availableReplicas}"
-CREATE_POOL_TIMEOUT_SEC = 60
+CREATE_POOL_TIMEOUT_SEC = 180
 SCALE_POOL_TIMEOUT_SEC = 60
 
 
@@ -60,8 +60,6 @@ def create_pool(
                 f"jsonpath='{JSONPATH_AVAILABLE_REPLICAS}'={replicas}",
                 timeout=CREATE_POOL_TIMEOUT_SEC,
             )
-        except Exception:
-            raise
         finally:
             if pull_secret:
                 pull_secret.delete(propagation_policy="Foreground")

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

@@ -6,6 +6,7 @@
 from typing import Any
 
 # third party
+from pydantic import model_validator
 from result import Err
 from result import Ok
 from result import Result
@@ -15,6 +16,7 @@
 from ...abstract_node import NodeSideType
 from ...client.api import APIRegistry
 from ...client.client import SyftClient
+from ...custom_worker.config import DockerWorkerConfig
 from ...custom_worker.config import WorkerConfig
 from ...custom_worker.k8s import IN_KUBERNETES
 from ...node.credentials import SyftVerifyKey
@@ -187,7 +189,7 @@ def __repr_syft_nested__(self) -> str:
 
 
 @serializable()
-class CreateCustomImageChange(Change):
+class CreateCustomImageChangeV2(Change):
     __canonical_name__ = "CreateCustomImageChange"
     __version__ = SYFT_OBJECT_VERSION_2
 
@@ -198,41 +200,70 @@ class CreateCustomImageChange(Change):
 
     __repr_attrs__ = ["config", "tag"]
 
+
+@serializable()
+class CreateCustomImageChange(Change):
+    __canonical_name__ = "CreateCustomImageChange"
+    __version__ = SYFT_OBJECT_VERSION_3
+
+    config: WorkerConfig
+    tag: str | None = None
+    registry_uid: UID | None = None
+    pull_image: bool = True
+
+    __repr_attrs__ = ["config", "tag"]
+
+    @model_validator(mode="after")
+    def _tag_required_for_dockerworkerconfig(self) -> Self:
+        if isinstance(self.config, DockerWorkerConfig) and self.tag is None:
+            raise ValueError("`tag` is required for `DockerWorkerConfig`.")
+        return self
+
     def _run(
         self, context: ChangeContext, apply: bool
     ) -> Result[SyftSuccess, SyftError]:
         try:
             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(
+                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:
+            if IN_KUBERNETES and not worker_image.is_prebuilt:
                 push_result = worker_image_service.push(
                     service_context,
                     image=worker_image.id,
@@ -245,11 +276,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 +322,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

@@ -331,7 +331,6 @@ def create_kubernetes_pool(
     **kwargs: Any,
 ) -> list[Pod] | SyftError:
     pool = None
-    error = False
 
     try:
         print(
@@ -366,11 +365,14 @@ def create_kubernetes_pool(
             reg_url=reg_url,
         )
     except Exception as e:
-        error = True
-        return SyftError(message=f"Failed to start workers {e}")
-    finally:
-        if error and pool:
+        if pool:
             pool.delete()
+        # stdlib
+        import traceback
+
+        return SyftError(
+            message=f"Failed to start workers {e} {e.__class__} {e.args} {traceback.format_exc()}."
+        )
 
     return runner.get_pool_pods(pool_name=pool_name)
 
@@ -564,7 +566,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,8 @@
 import pydantic
 
 # relative
-from ...custom_worker.config import DockerWorkerConfig
+from ...custom_worker.config import PrebuiltWorkerConfig
+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,16 +40,27 @@ 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",
+        name="submit",
         roles=DATA_OWNER_ROLE_LEVEL,
     )
-    def submit_dockerfile(
-        self, context: AuthedServiceContext, docker_config: DockerWorkerConfig
+    def submit(
+        self, context: AuthedServiceContext, worker_config: WorkerConfig
     ) -> SyftSuccess | SyftError:
+        image_identifier: SyftWorkerImageIdentifier | None = None
+        if isinstance(worker_config, PrebuiltWorkerConfig):
+            try:
+                image_identifier = SyftWorkerImageIdentifier.from_str(worker_config.tag)
+            except Exception:
+                return SyftError(
+                    f"Invalid Docker image name: {worker_config.tag}.\n"
+                    + "Please specify the image name in this format <registry>/<repo>:<tag>."
+                )
+
         worker_image = SyftWorkerImage(
-            config=docker_config,
+            config=worker_config,
             created_by=context.credentials,
+            image_identifier=image_identifier,
         )
         res = self.stash.set(context.credentials, worker_image)
 
@@ -278,14 +290,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)