From 8d8d9494e0a05ee90a02ebcd0364ba184ae62d17 Mon Sep 17 00:00:00 2001
From: Yubo Wang <yubwang@linkedin.com>
Date: Tue, 9 May 2023 13:25:23 -0700
Subject: [PATCH] fix tests and lints

Signed-off-by: Yubo Wang <yubwang@linkedin.com>
---
 .../flytekitplugins/kftensorflow/models/__init__.py        | 4 ++--
 .../flytekitplugins/kftensorflow/models/run_policy.py      | 7 ++-----
 setup.py                                                   | 2 +-
 3 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/__init__.py b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/__init__.py
index 2a25e199967..0b1e528b5b3 100644
--- a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/__init__.py
+++ b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/__init__.py
@@ -1,6 +1,6 @@
-from .restart_policy import RestartPolicy
 from .chief import Chief
 from .ps import PS
-from .worker import Worker
+from .restart_policy import RestartPolicy
 from .run_policy import CleanPodPolicy, RunPolicy
 from .tensorflow_job import TensorFlowJob
+from .worker import Worker
diff --git a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/run_policy.py b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/run_policy.py
index e9f44ad00b2..fd2aab80315 100644
--- a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/run_policy.py
+++ b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/run_policy.py
@@ -21,11 +21,8 @@ class CleanPodPolicy(Enum):
 
 class RunPolicy(common.FlyteIdlEntity):
     """
-    Configuration for a dask worker group
-
-    :param replicas: Number of workers in the group, minimum is 1
-    :param image: Optional image to use for the pods of the worker group
-    :param resources: Optional resources to use for the pods of the worker group
+    RunPolicy encapsulates various runtime policies of the distributed training job,
+    for example how to clean up resources and how long the job can stay active.
     """
 
     def __init__(
diff --git a/setup.py b/setup.py
index cc8114713fc..c435b41619b 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@
     },
     install_requires=[
         "googleapis-common-protos>=1.57",
-        "flyteidl>=1.3.16,<1.4.0",
+        "flyteidl>=1.5.2",
         "wheel>=0.30.0,<1.0.0",
         "pandas>=1.0.0,<2.0.0",
         "pyarrow>=4.0.0,<11.0.0",