utils

package
v0.6.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Mar 3, 2026 License: Apache-2.0 Imports: 28 Imported by: 0

Documentation

Overview

Copyright 2024 The Aibrix Team.

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at 

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

Index

Constants

View Source 
const (
	DefaultLLMEngine = "vllm"
)

https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken

View Source 
const (
	NAMESPACE = "aibrix-system"
)

Variables

View Source 
var (
	ReplicaSetDeploymentFinder = regexp.MustCompile(`^(.*)-\w+$`)     // Deployment-[random name]
	RayClusterFleetFinder      = regexp.MustCompile(`^(.*)-\w+-\w+$`) // RayClusterFleet-[random name]-[random name]
)
View Source 
var (
	RayClusterIdentifier      string = "ray.io/is-ray-node"
	RayClusterIdentifierYes   string = "yes"
	RayClusterNodeType        string = "ray.io/node-type"
	ReyClusterFleetIdentifier string = getReyClusterFleetIdentifier()
)
View Source 
var AllowedSchemas = []string{"s3://", "gcs://", "tos://", "huggingface://", "hf://", "/"}
View Source 
var DeploymentIdentifier string = getDeploymentIdentifier()

Functions

func AddLabel 

func AddLabel(labels map[string]string, labelKey, labelValue string) map[string]string

AddLabel returns a map with the given key and value added to the given map.

func CheckRedisHealth added in v0.5.0 

func CheckRedisHealth(ctx context.Context, redisClient *redis.Client) error

func CheckUser 

func CheckUser(ctx context.Context, u User, redisClient *redis.Client) bool

func CloneAndAddLabel 

func CloneAndAddLabel(labels map[string]string, labelKey, labelValue string) map[string]string

Clones the given map and returns a new map with the given key and value added. Returns the given map, if labelKey is empty.

func CloneAndRemoveLabel 

func CloneAndRemoveLabel(labels map[string]string, labelKey string) map[string]string

CloneAndRemoveLabel clones the given map and returns a new map with the given key removed. Returns the given map, if labelKey is empty.

func CloneSelectorAndAddLabel 

func CloneSelectorAndAddLabel(selector *metav1.LabelSelector, labelKey, labelValue string) *metav1.LabelSelector

Use DeepCopy to clone the selector, or create new one if nil Returns the given selector, if labelKey is empty.

func CountReadyPods 

func CountReadyPods(podList *v1.PodList) (int64, error)

func CountRoutablePods 

func CountRoutablePods(pods []*v1.Pod) (cnt int)

CountRoutablePods filters and returns the number of pods that are routable. A pod is routable if it have a valid PodIP and not in terminating state.

func CryptoShuffle added in v0.6.0 

func CryptoShuffle[T any](slice []T)

func DelUser 

func DelUser(ctx context.Context, u User, redisClient *redis.Client) error

func DeploymentNameFromPod added in v0.4.0 

func DeploymentNameFromPod(pod *v1.Pod) string

DeploymentNameFromPod extracts the deployment name from the pod using two methods: 1. If the pod has a label with the key "app.kubernetes.io/name", its value is considered the deployment name. 2. If the pod has an owner reference of kind "ReplicaSet", the deployment name is extracted from the owner reference's name. Alternatively, if the pod is a ray cluster node, we check: 1. If the pod has a label with the key "orchestration.aibrix.ai/raycluster-fleet-name", its value is considered the deployment name. 2. "app.kubernetes.io/name" is discared for ray cluster node identifid by the label "ray.io/is-ray-node" 3. If the pod has an owner reference of kind "RayCluster", the deployment name is extracted from the owner reference's name.

func DetokenizeText 

func DetokenizeText(tokenIds []int) (string, error)

func FilterActivePods 

func FilterActivePods(pods []v1.Pod) []v1.Pod

FilterActivePods returns active pods.

func FilterPodByName 

func FilterPodByName(podname string, pods []*v1.Pod) (*v1.Pod, bool)

FilterPodByName returns the pod with the given name.

func FilterPods 

func FilterPods(pods []v1.Pod, filterFn filterPod) []v1.Pod

FilterPods returns replica sets that are filtered by filterFn (all returned ones should match filterFn).

func FilterPodsByLabel added in v0.5.0 

func FilterPodsByLabel(pods []*v1.Pod, labelKey, labelValue string) []*v1.Pod

FilterPodsByLabel filters pods that have a specific label key-value pair

func FilterPodsByLabelSelector added in v0.6.0 

func FilterPodsByLabelSelector(pods []*v1.Pod, labelSelector string) ([]*v1.Pod, error)

FilterPodsByLabelSelector filter pod by k8s labelSelector

func FilterReadyPod 

func FilterReadyPod(pod *v1.Pod) bool

func FilterRoutablePods 

func FilterRoutablePods(pods []*v1.Pod) []*v1.Pod

FilterRoutablePods filters and returns a list of pods that are routable. A pod is routable if it have a valid PodIP and not in terminating state.

func FilterRoutablePodsInPlace 

func FilterRoutablePodsInPlace(pods []*v1.Pod) []*v1.Pod

FilterRoutablePodsInPlace filters a list of pods that are routable. A pod is routable if it have a valid PodIP and not in terminating state.

func GVKCheckExists added in v0.6.0 

func GVKCheckExists(cfg *rest.Config, gvk schema.GroupVersionKind) (bool, error)

GVKCheckExists check if gvk is support

func GeneratePodKey 

func GeneratePodKey(podNamespace, podName string) string

GeneratePodKey generates a key in the format "namespace/name" for a given pod.

func GetLLMEngine added in v0.6.0 

func GetLLMEngine(pod *v1.Pod, labelName string, defaultValue string) string

func GetModelPortForPod added in v0.4.0 

func GetModelPortForPod(requestID string, pod *v1.Pod) int64

func GetPodCondition 

func GetPodCondition(status *v1.PodStatus, conditionType v1.PodConditionType) (int, *v1.PodCondition)

GetPodCondition extracts the provided condition from the given status and returns that. Returns nil and -1 if the condition is not present, and the index of the located condition.

func GetPodConditionFromList 

func GetPodConditionFromList(conditions []v1.PodCondition, conditionType v1.PodConditionType) (int, *v1.PodCondition)

GetPodConditionFromList extracts the provided condition from the given list of condition and returns the index of the condition and the condition. Returns -1 and nil if the condition is not present.

func GetPodEnv added in v0.6.0 

func GetPodEnv(pod *v1.Pod, envName, defaultValue string) string

func GetPodListByLabelSelector 

func GetPodListByLabelSelector(ctx context.Context, podLister client.Client, namespace string, selector labels.Selector) (*v1.PodList, error)

func GetPodReadyCondition 

func GetPodReadyCondition(status v1.PodStatus) *v1.PodCondition

GetPodReadyCondition extracts the pod ready condition from the given status and returns that. Returns nil if the condition is not present.

func GetPortAnnotationOrDefault 

func GetPortAnnotationOrDefault(annotations map[string]string, key string, defaultValue int) int

func GetPortsForPod added in v0.6.0 

func GetPortsForPod(pod *v1.Pod) []int

GetPortsForPod returns all ports for a pod based on its configuration. For distributed data-parallel (DP) inference, a pod may expose multiple ports: - Base port from pod labels (e.g., "model-port") - Additional ports for each DP rank (base_port + rank_index) Returns nil if the pod doesn't have the required configuration.

func GetPositiveIntAnnotationOrDefault 

func GetPositiveIntAnnotationOrDefault(annotations map[string]string, key string, defaultValue int) int

func GetRedisClient 

func GetRedisClient() *redis.Client

func GetStringAnnotationOrDefault 

func GetStringAnnotationOrDefault(annotations map[string]string, key string, defaultValue string) string

func HasVolume added in v0.6.0 

func HasVolume(vols []v1.Volume, name string) bool

func HasVolumeMount added in v0.6.0 

func HasVolumeMount(mounts []v1.VolumeMount, name, path string) bool

func IsDataParallelPod added in v0.6.0 

func IsDataParallelPod(pod *v1.Pod) bool

IsDataParallelPod checks if the pod is configured for Distributed Data Parallel. It returns true if "data-parallel-size" environment variable is present and its value is greater than 1.

func IsPodActive added in v0.4.0 

func IsPodActive(p *v1.Pod) bool

func IsPodPhaseTerminal added in v0.4.0 

func IsPodPhaseTerminal(phase v1.PodPhase) bool

IsPodPhaseTerminal returns true if the pod's phase is terminal.

func IsPodReady 

func IsPodReady(pod *v1.Pod) bool

IsPodReady returns true if a pod is ready; false otherwise.

func IsPodReadyConditionTrue 

func IsPodReadyConditionTrue(status v1.PodStatus) bool

IsPodReadyConditionTrue returns true if a pod is ready; false otherwise.

func IsPodTerminal added in v0.4.0 

func IsPodTerminal(pod *v1.Pod) bool

IsPodTerminal returns true if a pod is terminal, all containers are stopped and cannot ever regress.

func IsPodTerminating 

func IsPodTerminating(p *v1.Pod) bool

IsPodTerminating check if pod is in terminating status via whether the deletion timestamp is set

func IsRayClusterAvailable 

func IsRayClusterAvailable(cluster *rayclusterv1.RayCluster, minReadySeconds int32, now metav1.Time) bool

func IsRayClusterReady 

func IsRayClusterReady(cluster *rayclusterv1.RayCluster) bool

IsRayClusterReady returns true if a RayCluster is ready; false otherwise.

func IsRayClusterStateReady 

func IsRayClusterStateReady(status rayclusterv1.RayClusterStatus) bool

func LoadEnv 

func LoadEnv(key, defaultValue string) string

LoadEnv loads an environment variable or returns a default value if not set.

func LoadEnvBool added in v0.4.0 

func LoadEnvBool(key string, defaultValue bool) bool

LoadEnvBool loads a boolean environment variable or returns a default value if not set.

func LoadEnvDuration added in v0.4.0 

func LoadEnvDuration(key string, defaultValue time.Duration) time.Duration

LoadEnvDuration loads a duration environment variable or returns a default value if not set.

func LoadEnvFloat 

func LoadEnvFloat(key string, defaultValue float64) float64

func LoadEnvInt 

func LoadEnvInt(key string, defaultValue int) int

func LookupEnv 

func LookupEnv(key string) (string, bool)

LookupEnv retrieves an environment variable and returns whether it exists. It returns the value and a boolean indicating its existence.

func ParsePodKey 

func ParsePodKey(key string) (string, string, bool)

ParsePodKey parses a key in the format "namespace/podName". Returns (namespace, podName, success).

func SelectRandomPod 

func SelectRandomPod(pods []*v1.Pod, randomFn func(int) int) (*v1.Pod, error)

SelectRandomPod selects a random pod from the provided list, ensuring it's routable. It returns an error if no ready pods are available.

func SetConditionInList 

func SetConditionInList(inputList []metav1.Condition, conditionType string, status metav1.ConditionStatus, reason, message string, args ...interface{}) []metav1.Condition

SetConditionInList sets the specific condition type on the given PodAutoscaler to the specified value with the given reason and message. The message and args are treated like a format string. The condition will be added if it is not present. The new list will be returned.

func SetUser 

func SetUser(ctx context.Context, u User, redisClient *redis.Client) error

func TokenizeInputText 

func TokenizeInputText(text string) ([]int, error)

func TrimMessage 

func TrimMessage(message string) string

func ValidateArtifactURL 

func ValidateArtifactURL(artifactURL string) error

ValidateArtifactURL checks if the ArtifactURL has a valid schema (s3://, gcs://, tos://, huggingface://, https://, /)

func ValidateKVCacheBackend added in v0.4.0 

func ValidateKVCacheBackend(kv *orchestrationv1alpha1.KVCache) error

ValidateKVCacheBackend validates that the backend specified in annotations is valid.

Types

type CustomizedRegistry 

type CustomizedRegistry[V any, A comparable] struct {
	Registry[V]
	// contains filtered or unexported fields
}

CustomizedRegistry extends Registry to provide a customized array provider.

func NewRegistryWithArrayProvider 

func NewRegistryWithArrayProvider[V any, A comparable](provider func([]V) A) *CustomizedRegistry[V, A]

func (*CustomizedRegistry[V, A]) Array 

func (reg *CustomizedRegistry[V, A]) Array() (arr A)

func (*CustomizedRegistry[V, A]) Delete 

func (reg *CustomizedRegistry[V, A]) Delete(key string)

func (*CustomizedRegistry[V, A]) Len 

func (reg *CustomizedRegistry[V, A]) Len() int

func (*CustomizedRegistry[V, A]) Store 

func (reg *CustomizedRegistry[V, A]) Store(key string, value V)

type Message 

type Message struct {
	Content string `json:"content"`
	Role    string `json:"role"`
}

type PodArray 

type PodArray struct {
	Pods []*v1.Pod
	// contains filtered or unexported fields
}

PodArray is a simple implementation of types.PodList indexed by deployment names

func (*PodArray) All 

func (arr *PodArray) All() []*v1.Pod

func (*PodArray) Indexes 

func (arr *PodArray) Indexes() []string

func (*PodArray) Len 

func (arr *PodArray) Len() int

func (*PodArray) ListByIndex 

func (arr *PodArray) ListByIndex(deploymentName string) []*v1.Pod

func (*PodArray) ListPortsForPod added in v0.6.0 

func (arr *PodArray) ListPortsForPod() map[string][]int

type Registry 

type Registry[V any] struct {
	// contains filtered or unexported fields
}

Registry is a generic hash set focus on storing values with string keys. Array output is optimized by offering a cached copy.

func NewRegistry 

func NewRegistry[V any]() *Registry[V]

func (*Registry[V]) Array 

func (reg *Registry[V]) Array() (arr []V)

func (*Registry[V]) Delete 

func (reg *Registry[V]) Delete(key string)

func (*Registry[V]) Len 

func (reg *Registry[V]) Len() int

func (*Registry[V]) Load 

func (reg *Registry[V]) Load(key string) (value V, ok bool)

func (*Registry[V]) Store 

func (reg *Registry[V]) Store(key string, value V)

type SyncMap 

type SyncMap[K any, V any] struct {
	// contains filtered or unexported fields
}

func (*SyncMap[K, V]) CompareAndDelete 

func (sm *SyncMap[K, V]) CompareAndDelete(key K, old V) (deleted bool)

func (*SyncMap[K, V]) CompareAndSwap 

func (sm *SyncMap[K, V]) CompareAndSwap(key K, old V, new V) bool

func (*SyncMap[K, V]) Delete 

func (sm *SyncMap[K, V]) Delete(key K)

func (*SyncMap[K, V]) Keys 

func (sm *SyncMap[K, V]) Keys() []K

func (*SyncMap[K, V]) Len 

func (sm *SyncMap[K, V]) Len() int

func (*SyncMap[K, V]) Load 

func (sm *SyncMap[K, V]) Load(key K) (typedVal V, ok bool)

func (*SyncMap[K, V]) LoadAndDelete 

func (sm *SyncMap[K, V]) LoadAndDelete(key K) (typedVal V, loaded bool)

func (*SyncMap[K, V]) LoadOrStore 

func (sm *SyncMap[K, V]) LoadOrStore(key K, value V) (V, bool)

func (*SyncMap[K, V]) Range 

func (sm *SyncMap[K, V]) Range(f func(key K, value V) bool)

func (*SyncMap[K, V]) Store 

func (sm *SyncMap[K, V]) Store(key K, value V)

func (*SyncMap[K, V]) Swap 

func (sm *SyncMap[K, V]) Swap(key K, value V) (V, bool)

func (*SyncMap[K, V]) Values 

func (sm *SyncMap[K, V]) Values() []V

type User 

type User struct {
	Name string `json:"name" validate:"required"`
	Rpm  int64  `json:"rpm"`
	Tpm  int64  `json:"tpm"`
}

func GetUser 

func GetUser(ctx context.Context, u User, redisClient *redis.Client) (User, error)

Source Files

View all Source files

Directories

Path Synopsis

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.