From 064dd9abf009eb112484f738801ea000dc72371e Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Fri, 12 May 2023 11:05:38 -0400 Subject: [PATCH] dont import ElasticConfig at toplevel Signed-off-by: Niels Bantilan --- plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py b/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py index d79fa785f6..aea2c9a2e6 100644 --- a/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py +++ b/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py @@ -7,7 +7,7 @@ from typing import Any, Callable, Dict, Optional, Union import cloudpickle -from flyteidl.plugins.pytorch_pb2 import DistributedPyTorchTrainingTask, ElasticConfig +from flyteidl.plugins.pytorch_pb2 import DistributedPyTorchTrainingTask from google.protobuf.json_format import MessageToDict import flytekit @@ -227,6 +227,8 @@ def get_custom(self, settings: SerializationSettings) -> Optional[Dict[str, Any] """ return super().get_custom(settings) else: + from flyteidl.plugins.pytorch_pb2 import ElasticConfig + elastic_config = ElasticConfig( rdzv_backend=self.rdzv_backend, min_replicas=self.min_nodes,