volcano.sh/volcano@v1.9.0/example/extender/extender.go (about)

     1  /*
     2  Copyright 2022 The Volcano Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package main
    18  
    19  import (
    20  	"encoding/json"
    21  	"io"
    22  	"net/http"
    23  
    24  	"volcano.sh/volcano/pkg/scheduler/api"
    25  	"volcano.sh/volcano/pkg/scheduler/plugins/extender"
    26  )
    27  
    28  var snapshot *api.ClusterInfo
    29  
    30  func onSessionOpen(w http.ResponseWriter, r *http.Request) {
    31  	content, err := io.ReadAll(r.Body)
    32  	if err != nil {
    33  		w.WriteHeader(http.StatusInternalServerError)
    34  		return
    35  	}
    36  
    37  	r.Body.Close()
    38  
    39  	req := &extender.OnSessionOpenRequest{}
    40  	if err := json.Unmarshal(content, req); err != nil {
    41  		w.WriteHeader(http.StatusBadRequest)
    42  		return
    43  	}
    44  
    45  	snapshot = &api.ClusterInfo{
    46  		Jobs:           req.Jobs,
    47  		Nodes:          req.Nodes,
    48  		Queues:         req.Queues,
    49  		NamespaceInfo:  req.NamespaceInfo,
    50  		RevocableNodes: req.RevocableNodes,
    51  	}
    52  }
    53  
    54  func onSessionClose(w http.ResponseWriter, r *http.Request) {
    55  	snapshot = nil
    56  }
    57  
    58  func predicate(w http.ResponseWriter, r *http.Request) {
    59  	content, err := io.ReadAll(r.Body)
    60  	if err != nil {
    61  		w.WriteHeader(http.StatusInternalServerError)
    62  		return
    63  	}
    64  
    65  	r.Body.Close()
    66  
    67  	req := &extender.PredicateRequest{}
    68  	if err := json.Unmarshal(content, req); err != nil || req.Task == nil || req.Node == nil {
    69  		w.WriteHeader(http.StatusBadRequest)
    70  		return
    71  	}
    72  
    73  	resp := &extender.PredicateResponse{}
    74  	if req.Task.BestEffort && len(req.Node.Tasks) > 10 {
    75  		sts := api.Status{}
    76  		sts.Code = api.Unschedulable
    77  		sts.Reason = "Too many tasks on the node"
    78  		resp.Status = append(resp.Status, &sts)
    79  	}
    80  	response, err := json.Marshal(resp)
    81  	if err != nil {
    82  		w.WriteHeader(http.StatusInternalServerError)
    83  		return
    84  	}
    85  
    86  	if _, err = w.Write(response); err != nil {
    87  		w.WriteHeader(http.StatusInternalServerError)
    88  	}
    89  }
    90  
    91  func prioritize(w http.ResponseWriter, r *http.Request) {
    92  	content, err := io.ReadAll(r.Body)
    93  	if err != nil {
    94  		w.WriteHeader(http.StatusInternalServerError)
    95  		return
    96  	}
    97  
    98  	r.Body.Close()
    99  
   100  	req := &extender.PrioritizeRequest{}
   101  	if err := json.Unmarshal(content, req); err != nil || req.Task == nil || len(req.Nodes) == 0 {
   102  		w.WriteHeader(http.StatusBadRequest)
   103  		return
   104  	}
   105  
   106  	resp := &extender.PrioritizeResponse{NodeScore: map[string]float64{}}
   107  	for i := range req.Nodes {
   108  		if req.Task.BestEffort && len(req.Nodes[i].Tasks) > 5 {
   109  			resp.NodeScore[req.Nodes[i].Name] = 0
   110  		} else {
   111  			resp.NodeScore[req.Nodes[i].Name] = 1
   112  		}
   113  	}
   114  
   115  	response, err := json.Marshal(resp)
   116  	if err != nil {
   117  		w.WriteHeader(http.StatusInternalServerError)
   118  		return
   119  	}
   120  
   121  	if _, err = w.Write(response); err != nil {
   122  		w.WriteHeader(http.StatusInternalServerError)
   123  	}
   124  }
   125  
   126  func main() {
   127  	http.HandleFunc("/onSessionOpen", onSessionOpen)
   128  	http.HandleFunc("/onSessionClose", onSessionClose)
   129  	http.HandleFunc("/predicate", predicate)
   130  	http.HandleFunc("/prioritize", prioritize)
   131  
   132  	if err := http.ListenAndServe(":8713", nil); err != nil {
   133  		panic(err)
   134  	}
   135  }