// @generated by protoc-gen-es v1.7.2 with parameter "target=ts"
// @generated from file flyteidl/plugins/kubeflow/pytorch.proto (package flyteidl.plugins.kubeflow, syntax proto3)
/* eslint-disable */
// @ts-nocheck

import type { BinaryReadOptions, FieldList, JsonReadOptions, JsonValue, PartialMessage, PlainMessage } from "@bufbuild/protobuf";
import { Message, proto3 } from "@bufbuild/protobuf";
import { RestartPolicy, RunPolicy } from "./common_pb.js";
import { Resources } from "../../core/tasks_pb.js";

/**
 * 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
 *
 * @generated from message flyteidl.plugins.kubeflow.ElasticConfig
 */
export class ElasticConfig extends Message<ElasticConfig> {
  /**
   * @generated from field: string rdzv_backend = 1;
   */
  rdzvBackend = "";

  /**
   * @generated from field: int32 min_replicas = 2;
   */
  minReplicas = 0;

  /**
   * @generated from field: int32 max_replicas = 3;
   */
  maxReplicas = 0;

  /**
   * @generated from field: int32 nproc_per_node = 4;
   */
  nprocPerNode = 0;

  /**
   * @generated from field: int32 max_restarts = 5;
   */
  maxRestarts = 0;

  constructor(data?: PartialMessage<ElasticConfig>) {
    super();
    proto3.util.initPartial(data, this);
  }

  static readonly runtime: typeof proto3 = proto3;
  static readonly typeName = "flyteidl.plugins.kubeflow.ElasticConfig";
  static readonly fields: FieldList = proto3.util.newFieldList(() => [
    { no: 1, name: "rdzv_backend", kind: "scalar", T: 9 /* ScalarType.STRING */ },
    { no: 2, name: "min_replicas", kind: "scalar", T: 5 /* ScalarType.INT32 */ },
    { no: 3, name: "max_replicas", kind: "scalar", T: 5 /* ScalarType.INT32 */ },
    { no: 4, name: "nproc_per_node", kind: "scalar", T: 5 /* ScalarType.INT32 */ },
    { no: 5, name: "max_restarts", kind: "scalar", T: 5 /* ScalarType.INT32 */ },
  ]);

  static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): ElasticConfig {
    return new ElasticConfig().fromBinary(bytes, options);
  }

  static fromJson(jsonValue: JsonValue, options?: Partial<JsonReadOptions>): ElasticConfig {
    return new ElasticConfig().fromJson(jsonValue, options);
  }

  static fromJsonString(jsonString: string, options?: Partial<JsonReadOptions>): ElasticConfig {
    return new ElasticConfig().fromJsonString(jsonString, options);
  }

  static equals(a: ElasticConfig | PlainMessage<ElasticConfig> | undefined, b: ElasticConfig | PlainMessage<ElasticConfig> | undefined): boolean {
    return proto3.util.equals(ElasticConfig, a, b);
  }
}

/**
 * Proto for plugin that enables distributed training using https://github.com/kubeflow/pytorch-operator
 *
 * @generated from message flyteidl.plugins.kubeflow.DistributedPyTorchTrainingTask
 */
export class DistributedPyTorchTrainingTask extends Message<DistributedPyTorchTrainingTask> {
  /**
   * Worker replicas spec
   *
   * @generated from field: flyteidl.plugins.kubeflow.DistributedPyTorchTrainingReplicaSpec worker_replicas = 1;
   */
  workerReplicas?: DistributedPyTorchTrainingReplicaSpec;

  /**
   * Master replicas spec, master replicas can only have 1 replica
   *
   * @generated from field: flyteidl.plugins.kubeflow.DistributedPyTorchTrainingReplicaSpec master_replicas = 2;
   */
  masterReplicas?: DistributedPyTorchTrainingReplicaSpec;

  /**
   * 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.
   *
   * @generated from field: flyteidl.plugins.kubeflow.RunPolicy run_policy = 3;
   */
  runPolicy?: RunPolicy;

  /**
   * config for an elastic pytorch job
   *
   * @generated from field: flyteidl.plugins.kubeflow.ElasticConfig elastic_config = 4;
   */
  elasticConfig?: ElasticConfig;

  constructor(data?: PartialMessage<DistributedPyTorchTrainingTask>) {
    super();
    proto3.util.initPartial(data, this);
  }

  static readonly runtime: typeof proto3 = proto3;
  static readonly typeName = "flyteidl.plugins.kubeflow.DistributedPyTorchTrainingTask";
  static readonly fields: FieldList = proto3.util.newFieldList(() => [
    { no: 1, name: "worker_replicas", kind: "message", T: DistributedPyTorchTrainingReplicaSpec },
    { no: 2, name: "master_replicas", kind: "message", T: DistributedPyTorchTrainingReplicaSpec },
    { no: 3, name: "run_policy", kind: "message", T: RunPolicy },
    { no: 4, name: "elastic_config", kind: "message", T: ElasticConfig },
  ]);

  static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): DistributedPyTorchTrainingTask {
    return new DistributedPyTorchTrainingTask().fromBinary(bytes, options);
  }

  static fromJson(jsonValue: JsonValue, options?: Partial<JsonReadOptions>): DistributedPyTorchTrainingTask {
    return new DistributedPyTorchTrainingTask().fromJson(jsonValue, options);
  }

  static fromJsonString(jsonString: string, options?: Partial<JsonReadOptions>): DistributedPyTorchTrainingTask {
    return new DistributedPyTorchTrainingTask().fromJsonString(jsonString, options);
  }

  static equals(a: DistributedPyTorchTrainingTask | PlainMessage<DistributedPyTorchTrainingTask> | undefined, b: DistributedPyTorchTrainingTask | PlainMessage<DistributedPyTorchTrainingTask> | undefined): boolean {
    return proto3.util.equals(DistributedPyTorchTrainingTask, a, b);
  }
}

/**
 * @generated from message flyteidl.plugins.kubeflow.DistributedPyTorchTrainingReplicaSpec
 */
export class DistributedPyTorchTrainingReplicaSpec extends Message<DistributedPyTorchTrainingReplicaSpec> {
  /**
   * Number of replicas
   *
   * @generated from field: int32 replicas = 1;
   */
  replicas = 0;

  /**
   * Image used for the replica group
   *
   * @generated from field: string image = 2;
   */
  image = "";

  /**
   * Resources required for the replica group
   *
   * @generated from field: flyteidl.core.Resources resources = 3;
   */
  resources?: Resources;

  /**
   * RestartPolicy determines whether pods will be restarted when they exit
   *
   * @generated from field: flyteidl.plugins.kubeflow.RestartPolicy restart_policy = 4;
   */
  restartPolicy = RestartPolicy.NEVER;

  constructor(data?: PartialMessage<DistributedPyTorchTrainingReplicaSpec>) {
    super();
    proto3.util.initPartial(data, this);
  }

  static readonly runtime: typeof proto3 = proto3;
  static readonly typeName = "flyteidl.plugins.kubeflow.DistributedPyTorchTrainingReplicaSpec";
  static readonly fields: FieldList = proto3.util.newFieldList(() => [
    { no: 1, name: "replicas", kind: "scalar", T: 5 /* ScalarType.INT32 */ },
    { no: 2, name: "image", kind: "scalar", T: 9 /* ScalarType.STRING */ },
    { no: 3, name: "resources", kind: "message", T: Resources },
    { no: 4, name: "restart_policy", kind: "enum", T: proto3.getEnumType(RestartPolicy) },
  ]);

  static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): DistributedPyTorchTrainingReplicaSpec {
    return new DistributedPyTorchTrainingReplicaSpec().fromBinary(bytes, options);
  }

  static fromJson(jsonValue: JsonValue, options?: Partial<JsonReadOptions>): DistributedPyTorchTrainingReplicaSpec {
    return new DistributedPyTorchTrainingReplicaSpec().fromJson(jsonValue, options);
  }

  static fromJsonString(jsonString: string, options?: Partial<JsonReadOptions>): DistributedPyTorchTrainingReplicaSpec {
    return new DistributedPyTorchTrainingReplicaSpec().fromJsonString(jsonString, options);
  }

  static equals(a: DistributedPyTorchTrainingReplicaSpec | PlainMessage<DistributedPyTorchTrainingReplicaSpec> | undefined, b: DistributedPyTorchTrainingReplicaSpec | PlainMessage<DistributedPyTorchTrainingReplicaSpec> | undefined): boolean {
    return proto3.util.equals(DistributedPyTorchTrainingReplicaSpec, a, b);
  }
}

