@flyteorg/flyteidl 1.5.0 → 1.5.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@flyteorg/flyteidl",
-  "version": "1.5.0",
+  "version": "1.5.1",
   "description": "Compiled protocol buffers and gRPC service clients/servers for Flyte IDLs",
   "repository": {
     "type": "git",

package/protos/flyteidl/plugins/kubeflow/common.proto

@@ -0,0 +1,33 @@
+syntax = "proto3";
+
+package flyteidl.plugins.kubeflow;
+
+option go_package = "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/plugins";
+
+
+enum RestartPolicy {
+    RESTART_POLICY_NEVER = 0;
+    RESTART_POLICY_ON_FAILURE = 1;
+    RESTART_POLICY_ALWAYS = 2;
+}
+
+enum CleanPodPolicy {
+    CLEANPOD_POLICY_NONE = 0;
+    CLEANPOD_POLICY_RUNNING = 1;
+    CLEANPOD_POLICY_ALL = 2;
+}
+
+message RunPolicy {
+    // Defines the policy to kill pods after the job completes. Default to None.
+    CleanPodPolicy clean_pod_policy = 1;
+
+    // TTL to clean up jobs. Default to infinite.
+    int32 ttl_seconds_after_finished = 2;
+
+    // Specifies the duration in seconds relative to the startTime that the job may be active
+    // before the system tries to terminate it; value must be positive integer.
+    int32 active_deadline_seconds = 3;
+
+    // Number of retries before marking this job failed.
+    int32 backoff_limit = 4;
+}

package/protos/flyteidl/plugins/kubeflow/mpi.proto

@@ -0,0 +1,43 @@
+syntax = "proto3";
+
+package flyteidl.plugins.kubeflow;
+
+option go_package = "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/plugins";
+
+import "flyteidl/core/tasks.proto";
+import "flyteidl/plugins/kubeflow/common.proto";
+
+// Proto for plugin that enables distributed training using https://github.com/kubeflow/mpi-operator
+message DistributedMPITrainingTask {
+  // Worker replicas spec
+  DistributedMPITrainingReplicaSpec worker_replicas = 1;
+
+  // Master replicas spec
+  DistributedMPITrainingReplicaSpec launcher_replicas = 2;
+  
+  // 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.
+  RunPolicy run_policy = 3;
+
+  // Number of slots per worker
+  int32 slots = 4; 
+}
+
+// Replica specification for distributed MPI training
+message DistributedMPITrainingReplicaSpec {
+  // Number of replicas
+  int32 replicas = 1;
+
+  // Image used for the replica group
+  string image = 2;
+
+  // Resources required for the replica group
+  core.Resources resources = 3;
+  
+  // Restart policy determines whether pods will be restarted when they exit
+  RestartPolicy restart_policy = 4;
+
+  // MPI sometimes requires different command set for different replica groups
+  repeated string command = 5;
+}

package/protos/flyteidl/plugins/kubeflow/pytorch.proto

@@ -0,0 +1,49 @@
+syntax = "proto3";
+
+package flyteidl.plugins.kubeflow;
+
+option go_package = "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/plugins";
+
+import "flyteidl/core/tasks.proto";
+import "flyteidl/plugins/kubeflow/common.proto";
+
+// Custom proto for torch elastic config for distributed training using 
+// https://github.com/kubeflow/training-operator/blob/master/pkg/apis/kubeflow.org/v1/pytorch_types.go
+message ElasticConfig {
+  string rdzv_backend = 1;
+  int32 min_replicas = 2;
+  int32 max_replicas = 3;
+  int32 nproc_per_node = 4;
+  int32 max_restarts = 5;
+}
+
+// Proto for plugin that enables distributed training using https://github.com/kubeflow/pytorch-operator
+message DistributedPyTorchTrainingTask {
+  // Worker replicas spec
+  DistributedPyTorchTrainingReplicaSpec worker_replicas = 1;
+
+  // Master replicas spec, master replicas can only have 1 replica
+  DistributedPyTorchTrainingReplicaSpec master_replicas = 2;
+
+  // 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.
+  RunPolicy run_policy = 3;
+
+  // config for an elastic pytorch job
+  ElasticConfig elastic_config = 4;
+}
+
+message DistributedPyTorchTrainingReplicaSpec {
+  // Number of replicas
+  int32 replicas = 1;
+
+  // Image used for the replica group
+  string image = 2;
+
+  // Resources required for the replica group
+  core.Resources resources = 3;
+  
+  // RestartPolicy determines whether pods will be restarted when they exit
+  RestartPolicy restart_policy = 4;
+}

package/protos/flyteidl/plugins/kubeflow/tensorflow.proto

@@ -0,0 +1,39 @@
+syntax = "proto3";
+
+package flyteidl.plugins.kubeflow;
+
+option go_package = "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/plugins";
+
+import "flyteidl/core/tasks.proto";
+import "flyteidl/plugins/kubeflow/common.proto";
+
+// Proto for plugin that enables distributed training using https://github.com/kubeflow/tf-operator
+message DistributedTensorflowTrainingTask {
+  // Worker replicas spec
+  DistributedTensorflowTrainingReplicaSpec worker_replicas = 1;
+
+  // Parameter server replicas spec
+  DistributedTensorflowTrainingReplicaSpec ps_replicas = 2;
+
+  // Chief replicas spec
+  DistributedTensorflowTrainingReplicaSpec chief_replicas = 3;
+
+  // 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.
+  RunPolicy run_policy = 4;
+}
+
+message DistributedTensorflowTrainingReplicaSpec {
+  // Number of replicas
+  int32 replicas = 1;
+
+  // Image used for the replica group
+  string image = 2;
+
+  // Resources required for the replica group
+  core.Resources resources = 3;
+
+  // RestartPolicy Determines whether pods will be restarted when they exit
+  RestartPolicy restart_policy = 4;
+}