v1alpha1

type AcceleratorConfig struct {
	Volumes []AcceleratorVolume
	EnvVars []EnvironmentVariableConfig
}

type AcceleratorVolume struct {
	Name      string
	HostPath  string
	MountPath string
}

type ControllerConfig struct {
	// Accelerators is a map from the name of the accelerator to the config for that accelerator.
	// This should match the value specified as a container limit.
	// e.g. alpha.kubernetes.io/nvidia-gpu
	Accelerators map[string]AcceleratorConfig

	// Path to the file containing the grpc server source
	GrpcServerFilePath string
}

type EnvironmentVariableConfig struct {
	Name  string
	Value string
}

type MasterSpec struct {
	ReplicaName string `json:"replicaName"`
	ReplicaRank int    `json:"replicaRank"`
}

type PyTorchJob struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`
	Spec              PyTorchJobSpec   `json:"spec"`
	Status            PyTorchJobStatus `json:"status"`
}

type PyTorchJobList struct {
	metav1.TypeMeta `json:",inline"`
	// Standard list metadata
	// More info: http://releases.k8s.io/HEAD/docs/devel/api-conventions.md#metadata
	metav1.ListMeta `json:"metadata,omitempty"`
	// Items is a list of PyTorchJobs
	Items []PyTorchJob `json:"items"`
}

type PyTorchJobPhase string

type PyTorchJobSpec struct {
	// TODO(jlewi): Can we we get rid of this and use some value from Kubernetes or a random ide.
	RuntimeId string

	// ReplicaSpecs specifies the PyTorch replicas to run.
	ReplicaSpecs []*PyTorchReplicaSpec `json:"replicaSpecs"`

	// PyTorchImage defines the tensorflow docker image that should be used for default parameter server
	PyTorchImage string `json:"pytorchImage,omitempty"`

	// TerminationPolicy specifies the condition that the pytorchjob should be considered finished.
	TerminationPolicy *TerminationPolicySpec `json:"terminationPolicy,omitempty"`

	// SchedulerName specifies the name of scheduler which should handle the PyTorchJob
	SchedulerName string `json:"schedulerName,omitempty"`
}

type PyTorchJobStatus struct {
	// Phase is the PyTorchJob running phase
	Phase  PyTorchJobPhase `json:"phase"`
	Reason string          `json:"reason"`

	// State indicates the state of the job.
	State State `json:"state"`

	// ReplicaStatuses specifies the status of each PyTorch replica.
	ReplicaStatuses []*PyTorchReplicaStatus `json:"replicaStatuses"`
}

type PyTorchReplicaSpec struct {
	// Replicas is the number of desired replicas.
	// This is a pointer to distinguish between explicit zero and unspecified.
	// Defaults to 1.
	// More info: http://kubernetes.io/docs/user-guide/replication-controller#what-is-a-replication-controller
	// +optional
	Replicas *int32              `json:"replicas,omitempty" protobuf:"varint,1,opt,name=replicas"`
	Template *v1.PodTemplateSpec `json:"template,omitempty" protobuf:"bytes,3,opt,name=template"`
	// MasterPort is the port to use for PyTorch services.
	MasterPort         *int32 `json:"masterPort,omitempty" protobuf:"varint,1,opt,name=masterPort"`
	PyTorchReplicaType `json:"replicaType"`
}

type PyTorchReplicaStatus struct {
	PyTorchReplicaType `json:"replica_type"`

	// State is the overall state of the replica
	State ReplicaState `json:"state"`

	// ReplicasStates provides the number of replicas in each status.
	ReplicasStates map[ReplicaState]int
}

type PyTorchReplicaType string

type ReplicaState string

type State string

type TerminationPolicySpec struct {
	// Master policy waits for a particular process (which is the master) to exit.
	Master *MasterSpec `json:"master,omitempty"`
}