fix tests and lints

Signed-off-by: Yubo Wang <[email protected]>

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

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__(

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",