From 4c98139592b67cbe898ea625cf091b5764b037e0 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 23 Apr 2024 18:06:06 -0700 Subject: [PATCH] Use default serialization settings if serialization_settings is None (#2367) Signed-off-by: Kevin Su Signed-off-by: Jan Fiedler --- flytekit/models/admin/task_execution.py | 2 +- flytekit/remote/remote.py | 24 +++++++++++++++--------- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/flytekit/models/admin/task_execution.py b/flytekit/models/admin/task_execution.py index d0a6d4ed2d..3eecad795e 100644 --- a/flytekit/models/admin/task_execution.py +++ b/flytekit/models/admin/task_execution.py @@ -42,7 +42,7 @@ def __init__( def phase(self): """ Enum value from flytekit.models.core.execution.TaskExecutionPhase - :rtype: int + :rtype: flytekit.models.core.execution.TaskExecutionPhase """ return self._phase diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 856563f0db..fc7711833b 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -759,7 +759,8 @@ async def _serialize_and_register( for entity, cp_entity in cp_task_entity_map.items(): tasks.append( loop.run_in_executor( - None, functools.partial(self.raw_register, cp_entity, settings, version, og_entity=entity) + None, + functools.partial(self.raw_register, cp_entity, serialization_settings, version, og_entity=entity), ) ) ident = [] @@ -767,7 +768,7 @@ async def _serialize_and_register( # serial register cp_other_entities = OrderedDict(filter(lambda x: not isinstance(x[1], task_models.TaskSpec), m.items())) for entity, cp_entity in cp_other_entities.items(): - ident.append(self.raw_register(cp_entity, settings, version, og_entity=entity)) + ident.append(self.raw_register(cp_entity, serialization_settings, version, og_entity=entity)) return ident[-1] def register_task( @@ -793,6 +794,8 @@ def register_task( serialization_settings = SerializationSettings( image_config=ImageConfig.auto_default_image(), source_root=project_root, + project=self.default_project, + domain=self.default_domain, ) ident = asyncio.run( @@ -825,13 +828,16 @@ def register_workflow( :param options: Additional execution options that can be configured for the default launchplan :return: """ - ident = self._resolve_identifier(ResourceType.WORKFLOW, entity.name, version, serialization_settings) - if serialization_settings: - b = serialization_settings.new_builder() - b.project = ident.project - b.domain = ident.domain - b.version = ident.version - serialization_settings = b.build() + if serialization_settings is None: + _, _, _, module_file = extract_task_module(entity) + project_root = _find_project_root(module_file) + serialization_settings = SerializationSettings( + image_config=ImageConfig.auto_default_image(), + source_root=project_root, + project=self.default_project, + domain=self.default_domain, + ) + self._resolve_identifier(ResourceType.WORKFLOW, entity.name, version, serialization_settings) ident = asyncio.run( self._serialize_and_register(entity, serialization_settings, version, options, default_launch_plan) )