syntax = "proto3";

package flyteidl.plugins;

option go_package = "github.com/lyft/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, ps, chief replicas spawned in the cluster for this job
  int32 workers = 1;
  // PS -> Parameter server
  int32 ps_replicas = 2;
  int32 chief_replicas = 3;
}