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

import type { BinaryReadOptions, FieldList, JsonReadOptions, JsonValue, PartialMessage, PlainMessage } from "@bufbuild/protobuf";
import { Message, proto3 } from "@bufbuild/protobuf";

/**
 * 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.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.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);
  }
}

/**
 * Custom proto for plugin that enables distributed training using https://github.com/kubeflow/pytorch-operator
 *
 * @generated from message flyteidl.plugins.DistributedPyTorchTrainingTask
 */
export class DistributedPyTorchTrainingTask extends Message<DistributedPyTorchTrainingTask> {
  /**
   * number of worker replicas spawned in the cluster for this job
   *
   * @generated from field: int32 workers = 1;
   */
  workers = 0;

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

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

  static readonly runtime: typeof proto3 = proto3;
  static readonly typeName = "flyteidl.plugins.DistributedPyTorchTrainingTask";
  static readonly fields: FieldList = proto3.util.newFieldList(() => [
    { no: 1, name: "workers", kind: "scalar", T: 5 /* ScalarType.INT32 */ },
    { no: 2, 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);
  }
}

