Add tag for images from PrebuiltWorkerConfig

OpenMined · May 16, 2024 · d4bb721 · d4bb721
1 parent 749bc1e
commit d4bb721
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/packages/syft/src/syft/service/worker/worker_image_service.py b/packages/syft/src/syft/service/worker/worker_image_service.py
@@ -6,8 +6,10 @@
 import pydantic
 
 # relative
+from ...custom_worker.config import PrebuiltWorkerConfig
 from ...custom_worker.config import WorkerConfig
 from ...custom_worker.k8s import IN_KUBERNETES
+from ...custom_worker.utils import ImageUtils
 from ...serde.serializable import serializable
 from ...store.document_store import DocumentStore
 from ...types.datetime import DateTime
@@ -49,6 +51,11 @@ def submit_container_image(
         worker_image = SyftWorkerImage(
             config=worker_config,
             created_by=context.credentials,
+            image_identifier=(
+                ImageUtils.parse_tag(worker_config.tag)
+                if isinstance(worker_config, PrebuiltWorkerConfig)
+                else None
+            ),
         )
         res = self.stash.set(context.credentials, worker_image)
 

diff --git a/tests/integration/container_workload/pool_image_test.py b/tests/integration/container_workload/pool_image_test.py
@@ -121,7 +121,7 @@ def test_pool_launch(
 
     # Submit Worker Image
     worker_config, docker_tag = (
-        (PrebuiltWorkerConfig(tag=(_tag := "docker.io/python:3-slim")), _tag)
+        (PrebuiltWorkerConfig(tag=(_tag := "docker.io/library/python:3-slim")), _tag)
         if prebuilt
         else make_docker_config_test_case("opendp")
     )