Kubeflow TensorFlow Training Operator Add Evaluator

plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/__init__.py

@@ -10,4 +10,4 @@
    TfJob
 """
 
-from .task import PS, Chief, CleanPodPolicy, RestartPolicy, RunPolicy, TfJob, Worker
+from .task import PS, Chief, CleanPodPolicy, Evaluator, RestartPolicy, RunPolicy, TfJob, Worker

plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/task.py

@@ -85,6 +85,15 @@ class Worker:
     restart_policy: Optional[RestartPolicy] = None
 
 
+@dataclass
+class Evaluator:
+    image: Optional[str] = None
+    requests: Optional[Resources] = None
+    limits: Optional[Resources] = None
+    replicas: int = 0
+    restart_policy: Optional[RestartPolicy] = None
+
+
 @dataclass
 class TfJob:
     """
@@ -95,6 +104,7 @@ class TfJob:
         chief: Configuration for the chief replica group.
         ps: Configuration for the parameter server (PS) replica group.
         worker: Configuration for the worker replica group.
+        evaluator: Configuration for the evaluator replica group.
         run_policy: Configuration for the run policy.
         num_workers: [DEPRECATED] This argument is deprecated. Use `worker.replicas` instead.
         num_ps_replicas: [DEPRECATED] This argument is deprecated. Use `ps.replicas` instead.
@@ -104,11 +114,13 @@ class TfJob:
     chief: Chief = field(default_factory=lambda: Chief())
     ps: PS = field(default_factory=lambda: PS())
     worker: Worker = field(default_factory=lambda: Worker())
+    evaluator: Evaluator = field(default_factory=lambda: Evaluator())
     run_policy: Optional[RunPolicy] = field(default_factory=lambda: None)
     # Support v0 config for backwards compatibility
     num_workers: Optional[int] = None
     num_ps_replicas: Optional[int] = None
     num_chief_replicas: Optional[int] = None
+    num_evaluator_replicas: Optional[int] = None
 
 
 class TensorflowFunctionTask(PythonFunctionTask[TfJob]):
@@ -130,19 +142,23 @@ def __init__(self, task_config: TfJob, task_function: Callable, **kwargs):
             )
         if task_config.num_chief_replicas and task_config.chief.replicas:
             raise ValueError(
-                "Cannot specify both `num_workers` and `chief.replicas`. Please use `chief.replicas` as `num_chief_replicas` is depreacated."
+                "Cannot specify both `num_chief_replicas` and `chief.replicas`. Please use `chief.replicas` as `num_chief_replicas` is depreacated."
             )
         if task_config.num_chief_replicas is None and task_config.chief.replicas is None:
             raise ValueError(
-                "Must specify either `num_workers` or `chief.replicas`. Please use `chief.replicas` as `num_chief_replicas` is depreacated."
+                "Must specify either `num_chief_replicas` or `chief.replicas`. Please use `chief.replicas` as `num_chief_replicas` is depreacated."
             )
         if task_config.num_ps_replicas and task_config.ps.replicas:
             raise ValueError(
-                "Cannot specify both `num_workers` and `ps.replicas`. Please use `ps.replicas` as `num_ps_replicas` is depreacated."
+                "Cannot specify both `num_ps_replicas` and `ps.replicas`. Please use `ps.replicas` as `num_ps_replicas` is depreacated."
             )
         if task_config.num_ps_replicas is None and task_config.ps.replicas is None:
             raise ValueError(
-                "Must specify either `num_workers` or `ps.replicas`. Please use `ps.replicas` as `num_ps_replicas` is depreacated."
+                "Must specify either `num_ps_replicas` or `ps.replicas`. Please use `ps.replicas` as `num_ps_replicas` is depreacated."
+            )
+        if task_config.num_evaluator_replicas and task_config.evaluator.replicas > 0:
+            raise ValueError(
+                "Cannot specify both `num_evaluator_replicas` and `evaluator.replicas`. Please use `evaluator.replicas` as `num_evaluator_replicas` is depreacated."
             )
         super().__init__(
             task_type=self._TF_JOB_TASK_TYPE,
@@ -153,7 +169,7 @@ def __init__(self, task_config: TfJob, task_function: Callable, **kwargs):
         )
 
     def _convert_replica_spec(
-        self, replica_config: Union[Chief, PS, Worker]
+        self, replica_config: Union[Chief, PS, Worker, Evaluator]
     ) -> tensorflow_task.DistributedTensorflowTrainingReplicaSpec:
         resources = convert_resources_to_resource_model(requests=replica_config.requests, limits=replica_config.limits)
         return tensorflow_task.DistributedTensorflowTrainingReplicaSpec(
@@ -184,11 +200,16 @@ def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
         if self.task_config.num_ps_replicas:
             ps.replicas = self.task_config.num_ps_replicas
 
+        evaluator = self._convert_replica_spec(self.task_config.evaluator)
+        if self.task_config.num_evaluator_replicas:
+            evaluator.replicas = self.task_config.num_evaluator_replicas
+
         run_policy = self._convert_run_policy(self.task_config.run_policy) if self.task_config.run_policy else None
         training_task = tensorflow_task.DistributedTensorflowTrainingTask(
             chief_replicas=chief,
             worker_replicas=worker,
             ps_replicas=ps,
+            evaluator_replicas=evaluator,
             run_policy=run_policy,
         )
 

plugins/flytekit-kf-tensorflow/setup.py

@@ -4,7 +4,7 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit>=1.6.1"]
+plugin_requires = ["flyteidl>=1.10.0", "flytekit>=1.6.1"]
 
 __version__ = "0.0.0+develop"
 

plugins/flytekit-kf-tensorflow/tests/test_tensorflow_task.py

@@ -1,5 +1,5 @@
 import pytest
-from flytekitplugins.kftensorflow import PS, Chief, CleanPodPolicy, RestartPolicy, RunPolicy, TfJob, Worker
+from flytekitplugins.kftensorflow import PS, Chief, CleanPodPolicy, Evaluator, RestartPolicy, RunPolicy, TfJob, Worker
 
 from flytekit import Resources, task
 from flytekit.configuration import Image, ImageConfig, SerializationSettings
@@ -23,6 +23,7 @@ def test_tensorflow_task_with_default_config(serialization_settings: Serializati
         worker=Worker(replicas=1),
         chief=Chief(replicas=0),
         ps=PS(replicas=0),
+        evaluator=Evaluator(replicas=0),
     )
 
     @task(
@@ -52,6 +53,9 @@ def my_tensorflow_task(x: int, y: str) -> int:
         "psReplicas": {
             "resources": {},
         },
+        "evaluatorReplicas": {
+            "resources": {},
+        },
     }
     assert my_tensorflow_task.get_custom(serialization_settings) == expected_dict
 
@@ -75,6 +79,13 @@ def test_tensorflow_task_with_custom_config(serialization_settings: Serializatio
             replicas=2,
             restart_policy=RestartPolicy.ALWAYS,
         ),
+        evaluator=Evaluator(
+            replicas=5,
+            requests=Resources(cpu="2", mem="2Gi"),
+            limits=Resources(cpu="4", mem="2Gi"),
+            image="evaluator:latest",
+            restart_policy=RestartPolicy.FAILURE,
+        ),
     )
 
     @task(
@@ -122,7 +133,23 @@ def my_tensorflow_task(x: int, y: str) -> int:
             "replicas": 2,
             "restartPolicy": "RESTART_POLICY_ALWAYS",
         },
+        "evaluatorReplicas": {
+            "replicas": 5,
+            "image": "evaluator:latest",
+            "resources": {
+                "requests": [
+                    {"name": "CPU", "value": "2"},
+                    {"name": "MEMORY", "value": "2Gi"},
+                ],
+                "limits": [
+                    {"name": "CPU", "value": "4"},
+                    {"name": "MEMORY", "value": "2Gi"},
+                ],
+            },
+            "restartPolicy": "RESTART_POLICY_ON_FAILURE",
+        },
     }
+
     assert my_tensorflow_task.get_custom(serialization_settings) == expected_custom_dict
 
 
@@ -131,6 +158,7 @@ def test_tensorflow_task_with_run_policy(serialization_settings: SerializationSe
         worker=Worker(replicas=1),
         ps=PS(replicas=0),
         chief=Chief(replicas=0),
+        evaluator=Evaluator(replicas=0),
         run_policy=RunPolicy(
             clean_pod_policy=CleanPodPolicy.RUNNING,
             backoff_limit=5,
@@ -166,19 +194,23 @@ def my_tensorflow_task(x: int, y: str) -> int:
         "psReplicas": {
             "resources": {},
         },
+        "evaluatorReplicas": {
+            "resources": {},
+        },
         "runPolicy": {
             "cleanPodPolicy": "CLEANPOD_POLICY_RUNNING",
             "backoffLimit": 5,
             "activeDeadlineSeconds": 100,
             "ttlSecondsAfterFinished": 100,
         },
     }
+
     assert my_tensorflow_task.get_custom(serialization_settings) == expected_dict
 
 
 def test_tensorflow_task():
     @task(
-        task_config=TfJob(num_workers=10, num_ps_replicas=1, num_chief_replicas=1),
+        task_config=TfJob(num_workers=10, num_ps_replicas=1, num_chief_replicas=1, num_evaluator_replicas=1),
         cache=True,
         requests=Resources(cpu="1"),
         cache_version="1",
@@ -212,7 +244,12 @@ def my_tensorflow_task(x: int, y: str) -> int:
             "replicas": 1,
             "resources": {},
         },
+        "evaluatorReplicas": {
+            "replicas": 1,
+            "resources": {},
+        },
     }
+
     assert my_tensorflow_task.get_custom(settings) == expected_dict
     assert my_tensorflow_task.resources.limits == Resources()
     assert my_tensorflow_task.resources.requests == Resources(cpu="1")