feat: introduce worker gateway

Signed-off-by: thxCode <thxcode0824@gmail.com>

Author: thxCode

cmd/gpustack/main.go

@@ -9,6 +9,7 @@ import (
 	"gpustack.ai/gpustack/pkg/devicemanager"
 	"gpustack.ai/gpustack/pkg/utils/signalx"
 	"gpustack.ai/gpustack/pkg/worker"
+	"gpustack.ai/gpustack/pkg/workergateway"
 )
 
 func main() {
@@ -18,6 +19,7 @@ func main() {
 		Short: "manages GPUStack Kubernetes resources.",
 	}
 	c.AddCommand(worker.NewCommand())
+	c.AddCommand(workergateway.NewCommand())
 	c.AddCommand(devicemanager.NewCommand())
 	c = cmd.Harness(c)
 

pkg/utils/httpx/client.go

@@ -179,12 +179,6 @@ func NewTraceRequest(uri string) (*http.Request, error) {
 	return http.NewRequest(http.MethodTrace, uri, nil)
 }
 
-// Error is similar to http.Error,
-// but it can get the error message by the given code.
-func Error(rw http.ResponseWriter, code int) {
-	http.Error(rw, http.StatusText(code), code)
-}
-
 // Close closes the http response body without error.
 func Close(resp *http.Response) {
 	if resp != nil && resp.Body != nil {

pkg/utils/httpx/writer.go

@@ -36,3 +36,9 @@ func JSON(w http.ResponseWriter, code int, v any) {
 	w.WriteHeader(code)
 	_, _ = w.Write(buf.Bytes())
 }
+
+// Error is similar to http.Error,
+// but it can get the error message by the given code.
+func Error(w http.ResponseWriter, code int) {
+	http.Error(w, http.StatusText(code), code)
+}

pkg/webserver/server.go

@@ -33,15 +33,20 @@ func New(c *Config) (Server, error) {
 	if c.Listener == nil {
 		return nil, errors.New("listener is required")
 	}
-	tcpAddr, ok := c.Listener.Addr().(*net.TCPAddr)
-	if !ok {
-		return nil, errors.New("listener must be a TCP listener")
+
+	var (
+		host string
+		port int
+	)
+	if tcpAddr, ok := c.Listener.Addr().(*net.TCPAddr); ok {
+		host = tcpAddr.IP.String()
+		port = tcpAddr.Port
 	}
 
 	return &server{
 		listener: c.Listener,
-		host:     tcpAddr.IP.String(),
-		port:     tcpAddr.Port,
+		host:     host,
+		port:     port,
 		runners:  c.Runners,
 		mux:      http.NewServeMux(),
 	}, nil
@@ -103,6 +108,9 @@ func (s *server) WebhookMux() *http.ServeMux {
 }
 
 func (s *server) HostPort() (string, int, error) {
+	if s.host == "" || s.port == 0 {
+		return "", 0, errors.New("server is not listening on a TCP address")
+	}
 	return s.host, s.port, nil
 }
 

pkg/workergateway/apis/types.go

@@ -0,0 +1,71 @@
+package apis
+
+import (
+	"k8s.io/apimachinery/pkg/api/resource"
+
+	worker "gpustack.ai/gpustack/api/worker/v1"
+)
+
+type (
+	AggregatedInstanceTypeSpec     = worker.InstanceTypeSpec
+	AggregatedInstanceTypeResource = worker.InstanceTypeResource
+)
+
+type AggregatedInstanceType struct {
+	Name string `json:"name"`
+
+	Spec   AggregatedInstanceTypeSpec   `json:"spec"`
+	Status AggregatedInstanceTypeStatus `json:"status"`
+}
+
+type AggregatedInstanceTypeStatus struct {
+	// Accelerator is the accelerator resource of the AggregatedInstanceType, e.g. "1", "4".
+	Accelerator AggregatedInstanceTypeResource `json:"accelerator"`
+
+	// CPU is the CPU resource of the AggregatedInstanceType, e.g. "4", "8".
+	CPU AggregatedInstanceTypeResource `json:"cpu"`
+
+	// RAM is the RAM resource of the AggregatedInstanceType, e.g. "40G", "16G".
+	RAM AggregatedInstanceTypeResource `json:"ram"`
+
+	// LocalStorage is the local storage resource of the AggregatedInstanceType, e.g. "100G", "500G".
+	LocalStorage AggregatedInstanceTypeResource `json:"localStorage"`
+
+	// Tiers is the list of once max request tiers of the AggregatedInstanceType.
+	//
+	// Each tier represents a combination of the AggregatedInstanceTypeOnceMaxRequestCandidate those satisfy the once max request of the tier.
+	Tiers []AggregatedInstanceTypeOnceMaxRequestTier `json:"tiers"`
+}
+
+type AggregatedInstanceTypeOnceMaxRequestTier struct {
+	// OnceMaxRequest is the accelerator once max request resource of the tier, e.g. "4", "8".
+	OnceMaxRequest resource.Quantity `json:"onceMaxRequest"`
+
+	// Candidates is the list of candidates of the tier.
+	//
+	// All candidates in the same tier must satisfy the max request of the tier,
+	// but may have different resource values.
+	Candidates []AggregatedInstanceTypeOnceMaxRequestCandidate `json:"candidates"`
+}
+
+// AggregatedInstanceTypeOnceMaxRequestCandidate represents a candidate of the max request tier of the AggregatedInstanceType.
+type AggregatedInstanceTypeOnceMaxRequestCandidate struct {
+	// Cluster is the candidate belongs to, e.g. "cluster-a", "cluster-b".
+	Cluster string `json:"cluster"`
+
+	// InstanceType is the instance type name of the candidate, e.g. "nvidia-a100-40g", "nvidia-v100-32g".
+	InstanceType string `json:"instanceType"`
+
+	// CPU is the CPU once max request resource of the candidate, e.g. "4", "8".
+	CPU resource.Quantity `json:"cpu"`
+
+	// RAM is the RAM once max request resource of the candidate, e.g. "40G", "16G".
+	RAM resource.Quantity `json:"ram"`
+
+	// LocalStorage is the local storage once max request resource of the candidate, e.g. "100G", "500G".
+	LocalStorage resource.Quantity `json:"localStorage"`
+}
+
+type AggregatedInstanceTypeList struct {
+	Items []AggregatedInstanceType `json:"items,omitempty"`
+}

pkg/workergateway/cmd.go

@@ -0,0 +1,49 @@
+package workergateway
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/spf13/cobra"
+)
+
+// NewCommand returns a new cobra command for the worker gateway.
+func NewCommand() *cobra.Command {
+	o := NewOptions()
+
+	c := &cobra.Command{
+		Use: "worker-gateway",
+		Aliases: []string{
+			"wg",
+		},
+		Short: "aggregates resources from upstream Kubernetes clusters.",
+		PreRunE: func(c *cobra.Command, args []string) error {
+			return o.Validate(c.Context())
+		},
+		RunE: func(c *cobra.Command, args []string) error {
+			ctx := c.Context()
+
+			cfg, err := o.Complete(ctx)
+			if err != nil {
+				return fmt.Errorf("complete config: %w", err)
+			}
+			wg, err := cfg.Apply(ctx)
+			if err != nil {
+				return fmt.Errorf("apply config: %w", err)
+			}
+			if err = wg.Prepare(ctx); err != nil {
+				return fmt.Errorf("prepare: %w", err)
+			}
+			err = wg.Start(ctx)
+			if err != nil && !errors.Is(err, context.Canceled) {
+				return err
+			}
+			return nil
+		},
+	}
+
+	o.AddFlags(c.Flags())
+
+	return c
+}

pkg/workergateway/config.go

@@ -0,0 +1,33 @@
+package workergateway
+
+import (
+	"context"
+
+	"gpustack.ai/gpustack/pkg/webserver"
+	"gpustack.ai/gpustack/pkg/workergateway/manager"
+)
+
+type Config struct {
+	// Manager.
+	ManagerConfig *manager.Config
+
+	// Server.
+	ServerConfig *webserver.Config
+}
+
+func (c *Config) Apply(ctx context.Context) (*WorkerGateway, error) {
+	mgr, err := c.ManagerConfig.Apply(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	srv, err := c.ServerConfig.Apply(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return &WorkerGateway{
+		Manager: mgr,
+		Server:  srv,
+	}, nil
+}

pkg/workergateway/gateway.go

@@ -0,0 +1,124 @@
+package workergateway
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/pprof"
+	"runtime"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/collectors"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"k8s.io/apiserver/pkg/server/healthz"
+	"k8s.io/apiserver/pkg/server/routes"
+	"k8s.io/component-base/logs"
+	"k8s.io/component-base/metrics/legacyregistry"
+	klog "k8s.io/klog/v2"
+	ctrlhealthz "sigs.k8s.io/controller-runtime/pkg/healthz"
+	ctrlmetrics "sigs.k8s.io/controller-runtime/pkg/metrics"
+
+	"gpustack.ai/gpustack/pkg/utils/gox"
+	"gpustack.ai/gpustack/pkg/utils/httpx"
+	"gpustack.ai/gpustack/pkg/webserver"
+	"gpustack.ai/gpustack/pkg/workergateway/manager"
+)
+
+func init() {
+	ctrlmetrics.Registry = struct {
+		prometheus.Registerer
+		prometheus.Gatherer
+	}{
+		Registerer: legacyregistry.Registerer(),
+		Gatherer:   legacyregistry.DefaultGatherer,
+	}
+}
+
+type WorkerGateway struct {
+	// Manager.
+	Manager manager.Manager
+
+	// Server.
+	Server webserver.Server
+}
+
+// Prepare prepares the runtime for the worker gateway,
+// including installing system resources, etc.
+func (wg *WorkerGateway) Prepare(ctx context.Context) error {
+	// Register metric collectors.
+	{
+		reg := ctrlmetrics.Registry
+		cs := []prometheus.Collector{
+			collectors.NewBuildInfoCollector(),
+			gox.NewStatsCollector(),
+		}
+		for i := range cs {
+			err := reg.Register(cs[i])
+			if err != nil {
+				return fmt.Errorf("register metric collector: %w", err)
+			}
+		}
+	}
+
+	return nil
+}
+
+func (wg *WorkerGateway) Start(ctx context.Context) error {
+	mu := wg.Server
+
+	// Register /metrics.
+	{
+		h := promhttp.HandlerOpts{
+			ErrorLog:      klog.NewStandardLogger("WARNING"),
+			ErrorHandling: promhttp.HTTPErrorOnError,
+		}
+		mu.Register("/metrics", promhttp.HandlerFor(ctrlmetrics.Registry, h))
+	}
+
+	// Register /healthz.
+	{
+		p := "/readyz"
+		h := &ctrlhealthz.Handler{
+			Checks: map[string]ctrlhealthz.Checker{
+				"ping": ctrlhealthz.Ping,
+				"log":  healthz.LogHealthz.Check,
+			},
+		}
+		mu.Register(p, http.StripPrefix(p, h))
+	}
+
+	// Register /livez.
+	{
+		p := "/livez"
+		h := &ctrlhealthz.Handler{
+			Checks: map[string]ctrlhealthz.Checker{
+				"ping": ctrlhealthz.Ping,
+				"log":  healthz.LogHealthz.Check,
+				"gopool": func(r *http.Request) error {
+					return gox.IsHealthy()
+				},
+			},
+		}
+		mu.Register(p, http.StripPrefix(p, h))
+	}
+
+	// Register /debug.
+	{
+		runtime.SetBlockProfileRate(1)
+		mu.Register("/debug/pprof/", httpx.LoopbackAccessHandlerFunc(pprof.Index))
+		mu.Register("/debug/pprof/cmdline", httpx.LoopbackAccessHandlerFunc(pprof.Cmdline))
+		mu.Register("/debug/pprof/profile", httpx.LoopbackAccessHandlerFunc(pprof.Profile))
+		mu.Register("/debug/pprof/symbol", httpx.LoopbackAccessHandlerFunc(pprof.Symbol))
+		mu.Register("/debug/pprof/trace", httpx.LoopbackAccessHandlerFunc(pprof.Trace))
+		mu.Register("/debug/flags/v", httpx.LoopbackAccessHandlerFunc(routes.StringFlagPutHandler(logs.GlogSetter)))
+	}
+
+	// Register API routes.
+	{
+		p := "/apis"
+		mu.Register(p+"/", http.StripPrefix(p, wg.getHandleApis()))
+	}
+
+	klog.Info("starting worker gateway")
+	return mu.Start(ctx)
+}