Skip to content

Commit f2e03aa

Browse files
AbirdcflyAbirdcfly
authored
[Bug] fix: correct non-inherited context (#763)
fix: correct non-inherited context Signed-off-by: Abirdcfly <[email protected]>
1 parent c225afc commit f2e03aa

File tree

8 files changed

+29
-29
lines changed

8 files changed

+29
-29
lines changed

‎pkg/controller/modeladapter/modeladapter_controller.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -302,7 +302,7 @@ func (r *ModelAdapterReconciler) Reconcile(ctx context.Context, req ctrl.Request
302302
// the finalizer is present, so let's unload lora from those inference engines
303303
// note: the base model pod could be deleted as well, so here we do best effort offloading
304304
// we do not need to reconcile the object if it encounters the unloading error.
305-
if err := r.unloadModelAdapter(modelAdapter); err != nil {
305+
if err := r.unloadModelAdapter(ctx, modelAdapter); err != nil {
306306
return ctrl.Result{}, err
307307
}
308308
if ok := controllerutil.RemoveFinalizer(modelAdapter, ModelAdapterFinalizer); !ok {
@@ -682,7 +682,7 @@ func (r *ModelAdapterReconciler) loadModelAdapter(url string, instance *modelv1a
682682

683683
// unloadModelAdapter unloads the loras from inference engines
684684
// base model pod could be deleted, in this case, we just do optimistic unloading. It only returns some necessary errors and http errors should not be returned.
685-
func (r *ModelAdapterReconciler) unloadModelAdapter(instance *modelv1alpha1.ModelAdapter) error {
685+
func (r *ModelAdapterReconciler) unloadModelAdapter(ctx context.Context, instance *modelv1alpha1.ModelAdapter) error {
686686
if len(instance.Status.Instances) == 0 {
687687
klog.Warningf("model adapter %s/%s has not been deployed to any pods yet, skip unloading", instance.GetNamespace(), instance.GetName())
688688
return nil
@@ -692,7 +692,7 @@ func (r *ModelAdapterReconciler) unloadModelAdapter(instance *modelv1alpha1.Mode
692692

693693
podName := instance.Status.Instances[0]
694694
targetPod := &corev1.Pod{}
695-
if err := r.Get(context.TODO(), types.NamespacedName{
695+
if err := r.Get(ctx, types.NamespacedName{
696696
Namespace: instance.Namespace,
697697
Name: podName,
698698
}, targetPod); err != nil {
@@ -803,7 +803,7 @@ func (r *ModelAdapterReconciler) reconcileEndpointSlice(ctx context.Context, ins
803803
// TODO: do necessary refactor to support multiple lora instance
804804
podName := instance.Status.Instances[0]
805805
pod := &corev1.Pod{}
806-
if err := r.Get(context.TODO(), types.NamespacedName{Namespace: instance.Namespace, Name: podName}, pod); err != nil {
806+
if err := r.Get(ctx, types.NamespacedName{Namespace: instance.Namespace, Name: podName}, pod); err != nil {
807807
if !apierrors.IsNotFound(err) {
808808
klog.Warning("Error getting Pod from lora instance list", err)
809809
return ctrl.Result{}, err

‎pkg/controller/podautoscaler/podautoscaler_controller.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -516,7 +516,7 @@ func (r *PodAutoscalerReconciler) scaleForResourceMappings(ctx context.Context,
516516
scale.SetNamespace(namespace)
517517
scale.SetName(name)
518518

519-
err := r.Get(context.TODO(), client.ObjectKey{Namespace: namespace, Name: name}, scale)
519+
err := r.Get(ctx, client.ObjectKey{Namespace: namespace, Name: name}, scale)
520520
if err == nil {
521521
return scale, targetGR, nil
522522
}

‎pkg/controller/rayclusterfleet/progress.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ func (r *RayClusterFleetReconciler) syncRolloutStatus(ctx context.Context, allRS
112112

113113
newDeployment := d
114114
newDeployment.Status = newStatus
115-
err := r.Status().Update(context.Background(), newDeployment)
115+
err := r.Status().Update(ctx, newDeployment)
116116
return err
117117
}
118118

‎pkg/controller/rayclusterfleet/rayclusterfleet_controller.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ func (r *RayClusterFleetReconciler) Reconcile(ctx context.Context, req ctrl.Requ
140140
return ctrl.Result{}, err
141141
}
142142

143-
clusterMap, err := r.getRayClusterMapForFleet(f, rsList)
143+
clusterMap, err := r.getRayClusterMapForFleet(ctx, f, rsList)
144144
if err != nil {
145145
return ctrl.Result{}, err
146146
}
@@ -209,7 +209,7 @@ func (r *RayClusterFleetReconciler) getReplicaSetsForFleet(ctx context.Context,
209209
return ownedReplicaSets, nil
210210
}
211211

212-
func (r *RayClusterFleetReconciler) getRayClusterMapForFleet(f *orchestrationv1alpha1.RayClusterFleet, rsList []*orchestrationv1alpha1.RayClusterReplicaSet) (map[types.UID][]*rayclusterv1.RayCluster, error) {
212+
func (r *RayClusterFleetReconciler) getRayClusterMapForFleet(ctx context.Context, f *orchestrationv1alpha1.RayClusterFleet, rsList []*orchestrationv1alpha1.RayClusterReplicaSet) (map[types.UID][]*rayclusterv1.RayCluster, error) {
213213
clusterList := &rayclusterv1.RayClusterList{}
214214

215215
// Get all RayClusters matches the fleet
@@ -218,7 +218,7 @@ func (r *RayClusterFleetReconciler) getRayClusterMapForFleet(f *orchestrationv1a
218218
return nil, err
219219
}
220220

221-
err = r.List(context.TODO(), clusterList, client.InNamespace(f.Namespace), client.MatchingLabelsSelector{Selector: selector})
221+
err = r.List(ctx, clusterList, client.InNamespace(f.Namespace), client.MatchingLabelsSelector{Selector: selector})
222222
if err != nil {
223223
return nil, err
224224
}

‎pkg/controller/rayclusterreplicaset/rayclusterreplicaset_controller.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ func (r *RayClusterReplicaSetReconciler) Reconcile(ctx context.Context, req ctrl
153153

154154
// status update if necessary
155155
newStatus := calculateStatus(replicaset, filteredClusters, scaleError)
156-
if err := r.updateReplicaSetStatus(replicaset, newStatus, rsKey); err != nil {
156+
if err := r.updateReplicaSetStatus(ctx, replicaset, newStatus, rsKey); err != nil {
157157
return reconcile.Result{}, err
158158
}
159159

@@ -202,7 +202,7 @@ func (r *RayClusterReplicaSetReconciler) scaleDown(ctx context.Context, replicas
202202
}
203203

204204
// updateReplicaSetStatus attempts to update the Status.Replicas of the given ReplicaSet, with a single GET/PUT retry.
205-
func (r *RayClusterReplicaSetReconciler) updateReplicaSetStatus(rs *orchestrationv1alpha1.RayClusterReplicaSet, newStatus orchestrationv1alpha1.RayClusterReplicaSetStatus, rsKey string) error {
205+
func (r *RayClusterReplicaSetReconciler) updateReplicaSetStatus(ctx context.Context, rs *orchestrationv1alpha1.RayClusterReplicaSet, newStatus orchestrationv1alpha1.RayClusterReplicaSetStatus, rsKey string) error {
206206
// Check if the expectations have been fulfilled for this ReplicaSet
207207
if !r.Expectations.SatisfiedExpectations(rsKey) {
208208
klog.V(4).Info("Expectations not yet fulfilled for ReplicaSet, delaying status update", "replicaSet", rsKey)
@@ -227,7 +227,7 @@ func (r *RayClusterReplicaSetReconciler) updateReplicaSetStatus(rs *orchestratio
227227
// Update ReplicaSet status if necessary
228228
newInstance := rs.DeepCopy()
229229
newInstance.Status = newStatus
230-
if err := r.Status().Update(context.Background(), newInstance); err != nil {
230+
if err := r.Status().Update(ctx, newInstance); err != nil {
231231
klog.ErrorS(err, "unable to update ReplicaSet status")
232232
return err
233233
}

‎pkg/metadata/users.go‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -62,12 +62,12 @@ func (s *httpServer) createUser(w http.ResponseWriter, r *http.Request) {
6262
return
6363
}
6464

65-
if utils.CheckUser(u, s.redisClient) {
65+
if utils.CheckUser(r.Context(), u, s.redisClient) {
6666
fmt.Fprintf(w, "User: %+v exists", u.Name)
6767
return
6868
}
6969

70-
if err := utils.SetUser(u, s.redisClient); err != nil {
70+
if err := utils.SetUser(r.Context(), u, s.redisClient); err != nil {
7171
http.Error(w, fmt.Sprintf("error occurred on creating user: %+v", err), http.StatusInternalServerError)
7272
return
7373
}
@@ -90,7 +90,7 @@ func (s *httpServer) readUser(w http.ResponseWriter, r *http.Request) {
9090
return
9191
}
9292

93-
user, err := utils.GetUser(u, s.redisClient)
93+
user, err := utils.GetUser(r.Context(), u, s.redisClient)
9494
if err != nil {
9595
fmt.Fprint(w, "user does not exists")
9696
return
@@ -114,12 +114,12 @@ func (s *httpServer) updateUser(w http.ResponseWriter, r *http.Request) {
114114
return
115115
}
116116

117-
if !utils.CheckUser(u, s.redisClient) {
117+
if !utils.CheckUser(r.Context(), u, s.redisClient) {
118118
fmt.Fprintf(w, "User: %+v does not exists", u.Name)
119119
return
120120
}
121121

122-
if err := utils.SetUser(u, s.redisClient); err != nil {
122+
if err := utils.SetUser(r.Context(), u, s.redisClient); err != nil {
123123
http.Error(w, fmt.Sprintf("error occurred on updating user: %+v", err), http.StatusInternalServerError)
124124
return
125125
}
@@ -142,12 +142,12 @@ func (s *httpServer) deleteUser(w http.ResponseWriter, r *http.Request) {
142142
return
143143
}
144144

145-
if !utils.CheckUser(u, s.redisClient) {
145+
if !utils.CheckUser(r.Context(), u, s.redisClient) {
146146
fmt.Fprintf(w, "User: %+v does not exists", u.Name)
147147
return
148148
}
149149

150-
if err := utils.DelUser(u, s.redisClient); err != nil {
150+
if err := utils.DelUser(r.Context(), u, s.redisClient); err != nil {
151151
http.Error(w, fmt.Sprintf("error occurred on deleting user: %+v", err), http.StatusInternalServerError)
152152
return
153153
}

‎pkg/plugins/gateway/gateway.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ import (
4545
envoyTypePb "github.com/envoyproxy/go-control-plane/envoy/type/v3"
4646
"github.com/vllm-project/aibrix/pkg/cache"
4747
routing "github.com/vllm-project/aibrix/pkg/plugins/gateway/algorithms"
48-
ratelimiter "github.com/vllm-project/aibrix/pkg/plugins/gateway/ratelimiter"
48+
"github.com/vllm-project/aibrix/pkg/plugins/gateway/ratelimiter"
4949
"github.com/vllm-project/aibrix/pkg/utils"
5050
healthPb "google.golang.org/grpc/health/grpc_health_v1"
5151
)
@@ -254,7 +254,7 @@ func (s *Server) HandleRequestHeaders(ctx context.Context, requestID string, req
254254
}
255255

256256
if username != "" {
257-
user, err = utils.GetUser(utils.User{Name: username}, s.redisClient)
257+
user, err = utils.GetUser(ctx, utils.User{Name: username}, s.redisClient)
258258
if err != nil {
259259
klog.ErrorS(err, "unable to process user info", "requestID", requestID, "username", username)
260260
return generateErrorResponse(

‎pkg/utils/users.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -30,17 +30,17 @@ type User struct {
3030
Tpm int64 `json:"tpm"`
3131
}
3232

33-
func CheckUser(u User, redisClient *redis.Client) bool {
34-
val, err := redisClient.Exists(context.Background(), genKey(u.Name)).Result()
33+
func CheckUser(ctx context.Context, u User, redisClient *redis.Client) bool {
34+
val, err := redisClient.Exists(ctx, genKey(u.Name)).Result()
3535
if err != nil {
3636
return false
3737
}
3838

3939
return val != 0
4040
}
4141

42-
func GetUser(u User, redisClient *redis.Client) (User, error) {
43-
val, err := redisClient.Get(context.Background(), genKey(u.Name)).Result()
42+
func GetUser(ctx context.Context, u User, redisClient *redis.Client) (User, error) {
43+
val, err := redisClient.Get(ctx, genKey(u.Name)).Result()
4444
if err != nil {
4545
return User{}, err
4646
}
@@ -53,7 +53,7 @@ func GetUser(u User, redisClient *redis.Client) (User, error) {
5353
return *user, nil
5454
}
5555

56-
func SetUser(u User, redisClient *redis.Client) error {
56+
func SetUser(ctx context.Context, u User, redisClient *redis.Client) error {
5757
if u.Rpm < 0 || u.Tpm < 0 {
5858
return fmt.Errorf("rpm or tpm can not negative")
5959
}
@@ -63,11 +63,11 @@ func SetUser(u User, redisClient *redis.Client) error {
6363
return err
6464
}
6565

66-
return redisClient.Set(context.Background(), genKey(u.Name), string(b), 0).Err()
66+
return redisClient.Set(ctx, genKey(u.Name), string(b), 0).Err()
6767
}
6868

69-
func DelUser(u User, redisClient *redis.Client) error {
70-
return redisClient.Del(context.Background(), genKey(u.Name)).Err()
69+
func DelUser(ctx context.Context, u User, redisClient *redis.Client) error {
70+
return redisClient.Del(ctx, genKey(u.Name)).Err()
7171
}
7272

7373
func genKey(s string) string {

0 commit comments

Comments
 (0)