From d4bb721ae7c95ab37626c91a98f2af1aebff3111 Mon Sep 17 00:00:00 2001 From: Kien Dang Date: Thu, 16 May 2024 21:46:39 +0800 Subject: [PATCH] Add tag for images from PrebuiltWorkerConfig --- .../syft/src/syft/service/worker/worker_image_service.py | 7 +++++++ tests/integration/container_workload/pool_image_test.py | 2 +- 2 files changed, 8 insertions(+), 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 index 1e2391132ab..0e700317717 100644 --- 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 index 95e8f8f7177..935e3ad16fd 100644 --- 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") )