diff --git a/docs/kubernetes/ingress.mdx b/docs/kubernetes/ingress.mdx index faf08482a..f2dcabde2 100644 --- a/docs/kubernetes/ingress.mdx +++ b/docs/kubernetes/ingress.mdx @@ -257,7 +257,7 @@ helm upgrade oci://ghcr.io/nvidia/openshell/helm-chart \ --set grpcRoute.replicaRouting.enabled=true ``` -The chart renders one Service per replica and a GRPCRoute rule that matches the `x-openshell-replica` request header. The CLI and the Rust SDK set that header from the owner the gateway reports and retry without it if the named replica is unavailable. Replica routing requires `workload.kind=statefulset`, `workload.allowMultiReplicaStatefulSet=true`, and at most 15 replicas. +The chart renders one Service per replica and a GRPCRoute rule that matches the `x-openshell-replica` request header. The CLI and SDKs send `x-openshell-want-owner: 1` on `GetSandbox`, read the owner from the `x-openshell-owner` response header, and set it as `x-openshell-replica`. If the proxy answers `UNAVAILABLE` with an empty message, `no healthy upstream`, or `connection failure`, they retry once without it. Replica routing requires `workload.kind=statefulset`, `workload.allowMultiReplicaStatefulSet=true`, and at most 15 replicas. If the release uses `workload.kind=deployment`, this upgrade replaces the Deployment with a StatefulSet, so every gateway pod restarts at once. The StatefulSet also requests a 1Gi PersistentVolumeClaim per replica, which stays `Pending` on clusters without a default StorageClass. diff --git a/sdk/go/openshell/v1/exec_client.go b/sdk/go/openshell/v1/exec_client.go index 1adc4c29d..56b13d928 100644 --- a/sdk/go/openshell/v1/exec_client.go +++ b/sdk/go/openshell/v1/exec_client.go @@ -27,7 +27,8 @@ func (e *execClient) Run(ctx context.Context, workspace, sandboxName string, com if sandboxName == "" { return nil, &StatusError{Code: ErrorInvalidArgument, Message: "sandbox name must not be empty"} } - if _, err := e.sandboxes.Get(ctx, workspace, sandboxName); err != nil { + owner, err := sandboxOwner(ctx, e.sandboxes, workspace, sandboxName) + if err != nil { return nil, err } var opt *ExecOptions @@ -37,7 +38,9 @@ func (e *execClient) Run(ctx context.Context, workspace, sandboxName string, com req := converter.ExecRequestToProto(sandboxName, command, opt) req.WorkspaceScope = namedWorkspaceScope(workspace) - stream, err := e.client.ExecSandbox(ctx, req) + stream, err := openRouted(ctx, owner, func(ctx context.Context) (grpc.ServerStreamingClient[pb.ExecSandboxEvent], error) { + return e.client.ExecSandbox(ctx, req) + }) if err != nil { return nil, converter.FromGRPCError(err) } @@ -61,7 +64,8 @@ func (e *execClient) Stream(ctx context.Context, workspace, sandboxName string, if sandboxName == "" { return nil, &StatusError{Code: ErrorInvalidArgument, Message: "sandbox name must not be empty"} } - if _, err := e.sandboxes.Get(ctx, workspace, sandboxName); err != nil { + owner, err := sandboxOwner(ctx, e.sandboxes, workspace, sandboxName) + if err != nil { return nil, err } var opt *ExecOptions @@ -72,7 +76,9 @@ func (e *execClient) Stream(ctx context.Context, workspace, sandboxName string, req.WorkspaceScope = namedWorkspaceScope(workspace) streamCtx, cancel := context.WithCancel(ctx) - stream, err := e.client.ExecSandbox(streamCtx, req) + stream, err := openRouted(streamCtx, owner, func(ctx context.Context) (grpc.ServerStreamingClient[pb.ExecSandboxEvent], error) { + return e.client.ExecSandbox(ctx, req) + }) if err != nil { cancel() return nil, converter.FromGRPCError(err) @@ -85,7 +91,8 @@ func (e *execClient) Interactive(ctx context.Context, workspace, sandboxName str if sandboxName == "" { return nil, &StatusError{Code: ErrorInvalidArgument, Message: "sandbox name must not be empty"} } - if _, err := e.sandboxes.Get(ctx, workspace, sandboxName); err != nil { + owner, err := sandboxOwner(ctx, e.sandboxes, workspace, sandboxName) + if err != nil { return nil, err } var opt *ExecOptions @@ -93,22 +100,17 @@ func (e *execClient) Interactive(ctx context.Context, workspace, sandboxName str opt = &opts[0] } + startReq := converter.ExecInteractiveRequestToProto(sandboxName, command, cols, rows, opt) + startReq.WorkspaceScope = namedWorkspaceScope(workspace) + streamCtx, cancel := context.WithCancel(ctx) - stream, err := e.client.ExecSandboxInteractive(streamCtx) + start := &pb.ExecSandboxInput{Payload: &pb.ExecSandboxInput_Start{Start: startReq}} + stream, err := openBidiRouted(streamCtx, owner, start, e.client.ExecSandboxInteractive) if err != nil { cancel() return nil, converter.FromGRPCError(err) } - startReq := converter.ExecInteractiveRequestToProto(sandboxName, command, cols, rows, opt) - startReq.WorkspaceScope = namedWorkspaceScope(workspace) - if sendErr := stream.Send(&pb.ExecSandboxInput{ - Payload: &pb.ExecSandboxInput_Start{Start: startReq}, - }); sendErr != nil { - cancel() - return nil, converter.FromGRPCError(sendErr) - } - return newInteractiveSession(streamCtx, cancel, stream), nil } diff --git a/sdk/go/openshell/v1/routing.go b/sdk/go/openshell/v1/routing.go new file mode 100644 index 000000000..8c9076202 --- /dev/null +++ b/sdk/go/openshell/v1/routing.go @@ -0,0 +1,95 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package v1 + +import ( + "context" + "io" + "strings" + + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" +) + +// Metadata that lets an edge proxy send a sandbox's long-lived calls straight +// to the gateway replica holding its supervisor session. The headers and the +// retry rule are documented in docs/kubernetes/ingress.mdx. +const ( + ownerReplicaHeader = "x-openshell-owner" + ownerRequestHeader = "x-openshell-want-owner" + routeReplicaHeader = "x-openshell-replica" +) + +// ownerGetter is implemented by the gateway-backed sandbox client. +type ownerGetter interface { + getWithOwner(ctx context.Context, workspace, name string) (*Sandbox, string, error) +} + +// sandboxOwner checks that the sandbox exists and returns the replica that +// owns its supervisor session, or "" when the gateway reports none. +func sandboxOwner(ctx context.Context, sandboxes SandboxInterface, workspace, name string) (string, error) { + if g, ok := sandboxes.(ownerGetter); ok { + _, owner, err := g.getWithOwner(ctx, workspace, name) + return owner, err + } + _, err := sandboxes.Get(ctx, workspace, name) + return "", err +} + +// openRouted opens a stream tagged with the owner replica. It opens once more +// untagged when the proxy reports that the owner's route has no ready endpoints. +// A stale owner is safe: any replica relays the call to the real owner. +// With an owner, it returns only after the gateway answers, as the CLI does. +func openRouted[S grpc.ClientStream](ctx context.Context, owner string, open func(context.Context) (S, error)) (S, error) { + if owner == "" { + return open(ctx) + } + stream, err := answered(open(metadata.AppendToOutgoingContext(ctx, routeReplicaHeader, owner))) + if retryUnrouted(err) { + return answered(open(ctx)) + } + return stream, err +} + +// answered waits for the response headers. Nil metadata means the stream ended +// without headers, and RecvMsg then returns its status. +func answered[S grpc.ClientStream](stream S, err error) (S, error) { + if err == nil { + var md metadata.MD + if md, err = stream.Header(); err == nil && md == nil { + err = stream.RecvMsg(nil) + } + } + return stream, err +} + +// openBidiRouted is openRouted for a bidi stream whose first frame is sent on open. +func openBidiRouted[I, O any](ctx context.Context, owner string, first *I, open func(context.Context, ...grpc.CallOption) (grpc.BidiStreamingClient[I, O], error)) (grpc.BidiStreamingClient[I, O], error) { + return openRouted(ctx, owner, func(ctx context.Context) (grpc.BidiStreamingClient[I, O], error) { + s, err := open(ctx) + if err != nil { + return nil, err + } + // io.EOF means the stream already ended; answered reports its status. + if err := s.Send(first); err == io.EOF { + return answered(s, nil) + } else if err != nil { + return s, err + } + return s, nil + }) +} + +// retryUnrouted matches what Envoy Gateway answers for a replica route with no +// ready endpoints: UNAVAILABLE with no message. Gateway errors always carry one. +func retryUnrouted(err error) bool { + s, ok := status.FromError(err) + if !ok || s.Code() != codes.Unavailable { + return false + } + msg := s.Message() + return msg == "" || strings.Contains(msg, "no healthy upstream") || strings.Contains(msg, "connection failure") +} diff --git a/sdk/go/openshell/v1/routing_test.go b/sdk/go/openshell/v1/routing_test.go new file mode 100644 index 000000000..115ef924e --- /dev/null +++ b/sdk/go/openshell/v1/routing_test.go @@ -0,0 +1,192 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package v1 + +import ( + "context" + "io" + "net" + "sync" + "testing" + + pb "github.com/NVIDIA/OpenShell/sdk/go/proto/openshellv1" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" +) + +// routingServer reports owner from GetSandbox and records the replica header +// of every streaming call ("" when unrouted). Routed calls fail with routedErr. +type routingServer struct { + pb.UnimplementedOpenShellServer + owner string + routedErr error + + mu sync.Mutex + wantedBy []string + routes []string +} + +func (s *routingServer) GetSandbox(ctx context.Context, _ *pb.GetSandboxRequest) (*pb.SandboxResponse, error) { + md, _ := metadata.FromIncomingContext(ctx) + s.mu.Lock() + s.wantedBy = append(s.wantedBy, firstValue(md, ownerRequestHeader)) + s.mu.Unlock() + if s.owner != "" { + _ = grpc.SetHeader(ctx, metadata.Pairs(ownerReplicaHeader, s.owner)) + } + return &pb.SandboxResponse{Sandbox: &pb.Sandbox{}}, nil +} + +// route records the call and returns the error a routed call should fail with. +func (s *routingServer) route(ctx context.Context) error { + md, _ := metadata.FromIncomingContext(ctx) + replica := firstValue(md, routeReplicaHeader) + s.mu.Lock() + defer s.mu.Unlock() + s.routes = append(s.routes, replica) + if replica != "" { + return s.routedErr + } + return nil +} + +func (s *routingServer) ExecSandbox(_ *pb.ExecSandboxRequest, stream grpc.ServerStreamingServer[pb.ExecSandboxEvent]) error { + if err := s.route(stream.Context()); err != nil { + return err + } + return stream.Send(exitEvent()) +} + +func (s *routingServer) ExecSandboxInteractive(stream grpc.BidiStreamingServer[pb.ExecSandboxInput, pb.ExecSandboxEvent]) error { + if _, err := stream.Recv(); err != nil { + return err + } + if err := s.route(stream.Context()); err != nil { + return err + } + return stream.Send(exitEvent()) +} + +func (s *routingServer) ForwardTcp(stream grpc.BidiStreamingServer[pb.TcpForwardFrame, pb.TcpForwardFrame]) error { //nolint:revive // proto-generated method name + if _, err := stream.Recv(); err != nil { + return err + } + if err := s.route(stream.Context()); err != nil { + return err + } + return stream.Send(&pb.TcpForwardFrame{Payload: &pb.TcpForwardFrame_Data{Data: []byte("pong")}}) +} + +func exitEvent() *pb.ExecSandboxEvent { + return &pb.ExecSandboxEvent{Payload: &pb.ExecSandboxEvent_Exit{Exit: &pb.ExecSandboxExit{}}} +} + +func firstValue(md metadata.MD, key string) string { + if v := md.Get(key); len(v) > 0 { + return v[0] + } + return "" +} + +func routingClients(t *testing.T, srv *routingServer) (*execClient, *tcpClient) { + t.Helper() + lis := bufconn.Listen(bufSize) + gs := grpc.NewServer() + pb.RegisterOpenShellServer(gs, srv) + go func() { _ = gs.Serve(lis) }() + conn, err := grpc.NewClient("passthrough:///bufconn", + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return lis.Dial() }), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + require.NoError(t, err) + t.Cleanup(func() { + _ = conn.Close() + gs.Stop() + }) + sandboxes := newSandboxClient(conn) + return newExecClient(conn, sandboxes), newTCPClient(conn, sandboxes, nil) +} + +func TestExecRoutesToOwner(t *testing.T) { + srv := &routingServer{owner: "openshell-1"} + exec, _ := routingClients(t, srv) + + _, err := exec.Run(context.Background(), "default", "box", []string{"true"}) + require.NoError(t, err) + + assert.Equal(t, []string{"1"}, srv.wantedBy) + assert.Equal(t, []string{"openshell-1"}, srv.routes) +} + +func TestExecWithoutOwnerIsUnrouted(t *testing.T) { + srv := &routingServer{} + exec, _ := routingClients(t, srv) + + _, err := exec.Run(context.Background(), "default", "box", []string{"true"}) + require.NoError(t, err) + + assert.Equal(t, []string{""}, srv.routes) +} + +func TestRoutedCallsRetryOnceWithoutHintWhenOwnerHasNoEndpoints(t *testing.T) { + srv := &routingServer{owner: "openshell-1", routedErr: status.Error(codes.Unavailable, "")} + exec, tcp := routingClients(t, srv) + ctx := context.Background() + + _, err := exec.Run(ctx, "default", "box", []string{"true"}) + require.NoError(t, err) + + stream, err := exec.Stream(ctx, "default", "box", []string{"true"}) + require.NoError(t, err) + code, err := stream.ExitCode() + require.NoError(t, err) + assert.Zero(t, code) + + session, err := exec.Interactive(ctx, "default", "box", []string{"sh"}, 80, 24) + require.NoError(t, err) + _, err = session.ExitCode() + require.NoError(t, err) + + conn, err := tcp.Forward(ctx, "default", "box", 8080) + require.NoError(t, err) + buf := make([]byte, 4) + _, err = io.ReadFull(conn, buf) + require.NoError(t, err) + assert.Equal(t, "pong", string(buf)) + _ = conn.Close() + + // Each call is attempted routed, then once unrouted. + assert.Equal(t, []string{ + "openshell-1", "", "openshell-1", "", "openshell-1", "", "openshell-1", "", + }, srv.routes) +} + +func TestRoutedCallDoesNotRetryOnGatewayError(t *testing.T) { + srv := &routingServer{owner: "openshell-1", routedErr: status.Error(codes.Unavailable, "supervisor not connected")} + exec, tcp := routingClients(t, srv) + ctx := context.Background() + + _, err := exec.Run(ctx, "default", "box", []string{"true"}) + require.Error(t, err) + _, err = tcp.Forward(ctx, "default", "box", 8080) + require.Error(t, err) + + assert.Equal(t, []string{"openshell-1", "openshell-1"}, srv.routes) +} + +func TestRetryUnrouted(t *testing.T) { + assert.True(t, retryUnrouted(status.Error(codes.Unavailable, ""))) + assert.True(t, retryUnrouted(status.Error(codes.Unavailable, "no healthy upstream"))) + assert.True(t, retryUnrouted(status.Error(codes.Unavailable, "upstream connect error: connection failure"))) + assert.False(t, retryUnrouted(status.Error(codes.Unavailable, "supervisor not connected"))) + assert.False(t, retryUnrouted(status.Error(codes.NotFound, ""))) + assert.False(t, retryUnrouted(io.EOF)) + assert.False(t, retryUnrouted(nil)) +} diff --git a/sdk/go/openshell/v1/sandbox_client.go b/sdk/go/openshell/v1/sandbox_client.go index d3b36cbdd..5d042b210 100644 --- a/sdk/go/openshell/v1/sandbox_client.go +++ b/sdk/go/openshell/v1/sandbox_client.go @@ -13,6 +13,7 @@ import ( "github.com/NVIDIA/OpenShell/sdk/go/openshell/v1/types" pb "github.com/NVIDIA/OpenShell/sdk/go/proto/openshellv1" "google.golang.org/grpc" + "google.golang.org/grpc/metadata" ) const defaultPollInterval = 500 * time.Millisecond @@ -116,10 +117,28 @@ func validateTemplateCreateSpec(spec *SandboxSpec) error { } func (s *sandboxClient) Get(ctx context.Context, workspace, name string) (*Sandbox, error) { + return s.get(ctx, workspace, name) +} + +// getWithOwner is Get that also asks the gateway for the owner replica. +func (s *sandboxClient) getWithOwner(ctx context.Context, workspace, name string) (*Sandbox, string, error) { + var md metadata.MD + sb, err := s.get(metadata.AppendToOutgoingContext(ctx, ownerRequestHeader, "1"), workspace, name, grpc.Header(&md)) + if err != nil { + return nil, "", err + } + var owner string + if v := md.Get(ownerReplicaHeader); len(v) > 0 { + owner = v[0] + } + return sb, owner, nil +} + +func (s *sandboxClient) get(ctx context.Context, workspace, name string, opts ...grpc.CallOption) (*Sandbox, error) { resp, err := s.client.GetSandbox(ctx, &pb.GetSandboxRequest{ Name: name, WorkspaceScope: namedWorkspaceScope(workspace), - }) + }, opts...) if err != nil { return nil, converter.FromGRPCError(err) } diff --git a/sdk/go/openshell/v1/ssh_client.go b/sdk/go/openshell/v1/ssh_client.go index 4a0e7cc15..aa19db4f2 100644 --- a/sdk/go/openshell/v1/ssh_client.go +++ b/sdk/go/openshell/v1/ssh_client.go @@ -65,7 +65,8 @@ func (s *sshClient) Tunnel(ctx context.Context, workspace, sandboxName string, p Message: fmt.Sprintf("port must be in range 1-65535, got %d", port), } } - if _, err := s.sandboxes.Get(ctx, workspace, sandboxName); err != nil { + owner, err := sandboxOwner(ctx, s.sandboxes, workspace, sandboxName) + if err != nil { return nil, err } @@ -84,13 +85,6 @@ func (s *sshClient) Tunnel(ctx context.Context, workspace, sandboxName string, p } }() - streamCtx, cancel := context.WithCancel(ctx) - stream, err := s.client.ForwardTcp(streamCtx) - if err != nil { - cancel() - return nil, converter.FromGRPCError(err) - } - initFrame := &pb.TcpForwardFrame{ Payload: &pb.TcpForwardFrame_Init{ Init: &pb.TcpForwardInit{ @@ -105,7 +99,9 @@ func (s *sshClient) Tunnel(ctx context.Context, workspace, sandboxName string, p }, } - if err := stream.Send(initFrame); err != nil { + streamCtx, cancel := context.WithCancel(ctx) + stream, err := openBidiRouted(streamCtx, owner, initFrame, s.client.ForwardTcp) //nolint:revive // proto-generated method name + if err != nil { cancel() return nil, converter.FromGRPCError(err) } diff --git a/sdk/go/openshell/v1/tcp_client.go b/sdk/go/openshell/v1/tcp_client.go index 24f806206..010658e33 100644 --- a/sdk/go/openshell/v1/tcp_client.go +++ b/sdk/go/openshell/v1/tcp_client.go @@ -37,20 +37,14 @@ func (t *tcpClient) Forward(ctx context.Context, workspace, sandboxName string, Message: fmt.Sprintf("port must be in range 1-65535, got %d", port), } } - if _, err := t.sandboxes.Get(ctx, workspace, sandboxName); err != nil { + owner, err := sandboxOwner(ctx, t.sandboxes, workspace, sandboxName) + if err != nil { return nil, err } var cfg forwardConfig options.Apply(&cfg, opts) - streamCtx, cancel := context.WithCancel(ctx) - stream, err := t.client.ForwardTcp(streamCtx) - if err != nil { - cancel() - return nil, converter.FromGRPCError(err) - } - initFrame := &pb.TcpForwardFrame{ Payload: &pb.TcpForwardFrame_Init{ Init: &pb.TcpForwardInit{ @@ -67,7 +61,9 @@ func (t *tcpClient) Forward(ctx context.Context, workspace, sandboxName string, }, } - if err := stream.Send(initFrame); err != nil { + streamCtx, cancel := context.WithCancel(ctx) + stream, err := openBidiRouted(streamCtx, owner, initFrame, t.client.ForwardTcp) //nolint:revive // proto-generated method name + if err != nil { cancel() return nil, converter.FromGRPCError(err) }