api

type ComputeUtilization struct {
	ProcessID          string
	DeviceUUID         string
	UtilizationPercent float64
}

type DataResponse[T any] struct {
	Data T `json:"data"`
}

type DeviceInfo struct {
	UUID                       string
	Vendor                     string
	Model                      string
	Index                      int32
	NUMANode                   int32
	TotalMemoryBytes           uint64
	MaxTflops                  float64
	VirtualizationCapabilities VirtualizationCapabilities
	Properties                 map[string]string
	Healthy                    bool

	ParentUUID string

	// Host - Guest device node mapping, eg /dev/nvidia0 -> /dev/nvidia0
	// When multiple device allocated, deduplicated by device node
	DeviceNode map[string]string

	// Env to inject to guest
	DeviceEnv map[string]string

	IsolationMode IsolationMode
}

type DeviceSpec struct {
	GuestPath string `json:"guestPath,omitempty"`

	HostPath string `json:"hostPath,omitempty"`

	Permissions string `json:"permissions,omitempty"`
}

type ErrorResponse struct {
	Error string `json:"error"`
}

type GPUUsageMetrics struct {
	DeviceUUID        string
	MemoryBytes       uint64
	MemoryPercentage  float64
	ComputePercentage float64
	ComputeTflops     float64
	Rx                float64 // PCIe RX in KB
	Tx                float64 // PCIe TX in KB
	Temperature       float64
	PowerUsage        int64 // in watts
	ExtraMetrics      map[string]float64
}

type IsolationMode = tfv1.IsolationModeType

type LimiterInfo struct {
	WorkerUID string         `json:"worker_uid"`
	Requests  *tfv1.Resource `json:"requests,omitempty"`
	Limits    *tfv1.Resource `json:"limits,omitempty"`
}

type ListLimitersResponse struct {
	Limiters []LimiterInfo `json:"limiters"`
}

type ListPodsResponse struct {
	Pods []PodInfo `json:"pods"`
}

type ListProcessesResponse struct {
	Processes []ProcessInfo `json:"processes"`
}

type MemoryUtilization struct {
	ProcessID     string
	DeviceUUID    string
	UsedBytes     uint64
	ReservedBytes uint64
}

type MessageAndDataResponse[T any] struct {
	Message string `json:"message"`
	Data    T      `json:"data"`
}

type Mount struct {
	GuestPath string `json:"guestPath,omitempty"`

	HostPath string `json:"hostPath,omitempty"`
}

type NodeInfo struct {
	// Extra metadata for centralized management
	RAMSizeBytes  int64
	DataDiskBytes int64

	// Aggregated info of whole Node
	TotalTFlops    float64
	TotalVRAMBytes int64
	DeviceIDs      []string
}

type PodInfo struct {
	PodName     string        `json:"pod_name"`
	Namespace   string        `json:"namespace"`
	GPUIDs      []string      `json:"gpu_uuids"`
	TflopsLimit *float64      `json:"tflops_limit,omitempty"`
	VramLimit   *uint64       `json:"vram_limit,omitempty"`
	QoSLevel    tfv1.QoSLevel `json:"qos_level,omitempty"`
}

type ProcessInfo struct {
	WorkerUID      string            `json:"worker_uid"`
	ProcessMapping map[string]string `json:"process_mapping"` // container PID -> host PID
}

type ProcessInformation struct {
	ProcessID                 string
	DeviceUUID                string
	ComputeUtilizationPercent float64
	ActiveSMs                 uint64
	TotalSMs                  uint64
	MemoryUsedBytes           uint64
	MemoryReservedBytes       uint64
	MemoryUtilizationPercent  float64
}

type StatusResponse struct {
	Status string `json:"status"`
}

type TrapResponse struct {
	Message       string `json:"message"`
	SnapshotCount int    `json:"snapshot_count"`
}

type VirtualizationCapabilities struct {
	SupportsPartitioning  bool
	SupportsSoftIsolation bool
	SupportsHardIsolation bool
	SupportsSnapshot      bool
	SupportsMetrics       bool
	MaxPartitions         uint32
	MaxWorkersPerDevice   uint32
}

type WorkerAllocation struct {
	WorkerInfo *WorkerInfo

	// the complete or partitioned device info
	DeviceInfos []*DeviceInfo

	Envs map[string]string

	Mounts []*Mount

	Devices []*DeviceSpec
}

type WorkerInfo struct {
	WorkerUID        string
	Namespace        string
	WorkerName       string
	AllocatedDevices []string
	Status           WorkerStatus

	QoS           tfv1.QoSLevel
	IsolationMode IsolationMode

	Requests tfv1.Resource
	Limits   tfv1.Resource

	WorkloadName      string
	WorkloadNamespace string

	// Only set for partitioned mode
	PartitionTemplateID string

	// Extra information from backend
	Labels      map[string]string
	Annotations map[string]string

	DeletedAt int64
}

type WorkerMetrics struct {
	DeviceUUID        string
	WorkerUID         string
	ProcessID         string
	MemoryBytes       uint64
	MemoryPercentage  float64
	ComputeTflops     float64
	ComputePercentage float64
	ExtraMetrics      map[string]float64
}

type WorkerStatus string