Apache YuniKorn Core

The schedulerclient provides a way to interact with the YuniKorn scheduler using the yunikorn-scheduler-interface. It establishes a gRPC connection to the scheduler (defaulting to localhost:3333) and uses the si.SchedulerClient to perform operations like registering a resource manager and streaming allocation updates.

To use the client, you must:

1. Establish a gRPC connection using grpc.NewClient.
2. Initialize the client using si.NewSchedulerClient(conn).
3. Call RegisterResourceManager to register the client with the scheduler.
4. Use UpdateAllocation to open a bidirectional stream for sending AllocationRequest objects and receiving responses.

import (
	"context"
	"time"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)

// ... setup connection
conn, err := grpc.NewClient("localhost:3333", grpc.WithTransportCredentials(insecure.NewCredentials()))
c := si.NewSchedulerClient(conn)

ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
defer cancel()

// Register the resource manager
_, err = c.RegisterResourceManager(ctx, &si.RegisterResourceManagerRequest{})

// Start streaming allocation updates
stream, err := c.UpdateAllocation(ctx)
// Use stream.Send(&si.AllocationRequest{}) to send requests
// Use stream.Recv() to receive responses

The NonBlockingGRPCServer interface provides a way to start and manage gRPC servers without blocking the main execution thread. This is useful for initializing scheduler services that need to run in the background while the rest of the application continues its setup.

To use it, call NewNonBlockingGRPCServer() to get an instance, then use Start(endpoint, ss) to launch the server. The endpoint must follow a specific URI format (see ParseEndpoint). The ss parameter is the si.SchedulerServer implementation you wish to register.

Lifecycle management methods:

• Start(endpoint string, ss si.SchedulerServer): Launches the server in a background goroutine.
• Wait(): Blocks the caller until the server has fully stopped.
• Stop(): Performs a graceful shutdown of the gRPC server.
• ForceStop(): Performs an immediate, forceful shutdown.

import (
	"github.com/apache/yunikorn-core/pkg/common"
	"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)

// Assuming 'myScheduler' implements si.SchedulerServer
server := common.NewNonBlockingGRPCServer()

// Start the server on a TCP endpoint
server.Start("tcp://0.0.0.0:8080", myScheduler)

// ... do other work ...

// Gracefully shut down
server.Stop()
server.Wait()

The SimpleScheduler is a reference implementation of the si.SchedulerServer interface. It is used to demonstrate how to handle gRPC streams for managing resource managers, allocations, applications, and nodes within the YuniKorn ecosystem.

To implement a custom scheduler, you must satisfy the methods defined in the si.SchedulerServer interface, which include:

• RegisterResourceManager: Handles the registration of a resource manager.
• UpdateAllocation: Manages a long-lived stream for allocation updates.
• UpdateApplication: Manages a long-lived stream for application updates.
• UpdateNode: Manages a long-lived stream for node updates.

Each stream-based method (UpdateAllocation, UpdateApplication, UpdateNode) follows a pattern of receiving data from the stream via conn.Recv(), processing it, and sending a response back via conn.Send() until the context is cancelled or an io.EOF is received.

type SimpleScheduler struct {
	si.UnimplementedSchedulerServer
}

// Example method implementation for handling allocation updates
func (scheduler *SimpleScheduler) UpdateAllocation(conn si.Scheduler_UpdateAllocationServer) error {
    ctx := conn.Context()
    for {
        select {
        case <-ctx.Done():
            return ctx.Err()
        default:
        }

        _, err := conn.Recv()
        if err == io.EOF {
            return nil
        }
        if err != nil {
            return err
        }

        resp := si.AllocationResponse{}
        if err := conn.Send(&resp); err != nil {
            return err
        }
    }
}