syntax = "proto3";

package flyteidl.plugins;

option go_package = "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/plugins";

// Custom proto for plugin that enables distributed training using https://github.com/kubeflow/tf-operator
message DistributedTensorflowTrainingTask {
  // number of worker replicas spawned in the cluster for this job
  int32 workers = 1;
  // PS -> Parameter server
  // number of ps replicas spawned in the cluster for this job
  int32 ps_replicas = 2;
  // number of chief replicas spawned in the cluster for this job
  int32 chief_replicas = 3;
  // number of evaluator replicas spawned in the cluster for this job
  int32 evaluator_replicas = 4;
}