nsc: add --builder_shape flag to buildx setup (#1805)

Allows requesting a specific builder shape (e.g. 4x16) when setting up
the server-side buildx proxy. Wired through BuilderConfiguration.Shape
to GetBuilderConfigurationRequest.BuilderShape.

Requires builder customization to be enabled for the workspace.

Amp-Thread-ID:
https://ampcode.com/threads/T-019e2090-db2c-726b-b8a7-c25ee50d3a65

Co-authored-by: Amp <amp@ampcode.com>

Author: n-g

internal/cli/cmd/cluster/buildx.go

@@ -14,10 +14,12 @@ import (
 	"os/signal"
 	"path"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"syscall"
 	"time"
 
+	computev1beta "buf.build/gen/go/namespace/cloud/protocolbuffers/go/proto/namespace/cloud/compute/v1beta"
 	"github.com/cenkalti/backoff"
 	"github.com/docker/buildx/store"
 	"github.com/docker/buildx/store/storeutil"
@@ -83,6 +85,8 @@ func newSetupBuildxCmd() *cobra.Command {
 	_ = cmd.Flags().MarkHidden("use_server_side_proxy")
 	experimental := cmd.Flags().String("experimental", "", "A set of experimental features to be passed at construction time.")
 	_ = cmd.Flags().MarkHidden("experimental")
+	builderShape := cmd.Flags().String("builder_shape", "", "If set, requests a specific builder shape in the form <vcpu>x<memGB> (e.g. 4x16). Requires builder customization to be enabled for your workspace — reach out to the Namespace team.")
+	_ = cmd.Flags().MarkHidden("builder_shape")
 
 	cmd.RunE = fncobra.RunE(func(ctx context.Context, args []string) error {
 		if *debugDir != "" && !*background {
@@ -125,10 +129,16 @@ func newSetupBuildxCmd() *cobra.Command {
 		}
 
 		if *useServerSideProxy {
+			shape, err := parseBuilderShape(*builderShape)
+			if err != nil {
+				return err
+			}
+
 			if err := setupServerSideBuildxProxy(ctx, state, *name, *use, *defaultLoad, dockerCli, available, api.BuilderConfiguration{
 				SkipPrespawn: !*createAtStartup,
 				Name:         *tag,
 				Experimental: *experimental,
+				Shape:        shape,
 			}); err != nil {
 				return err
 			}
@@ -792,6 +802,32 @@ func newStatusBuildxCommand() *cobra.Command {
 	return cmd
 }
 
+func parseBuilderShape(spec string) (*computev1beta.InstanceShape, error) {
+	if spec == "" {
+		return nil, nil
+	}
+
+	parts := strings.Split(spec, "x")
+	if len(parts) != 2 {
+		return nil, fnerrors.Newf("invalid --builder_shape %q: expected format <vcpu>x<memGB> (e.g. 4x16)", spec)
+	}
+
+	vcpu, err := strconv.Atoi(parts[0])
+	if err != nil || vcpu <= 0 {
+		return nil, fnerrors.Newf("invalid --builder_shape %q: vcpu must be a positive integer", spec)
+	}
+
+	memGB, err := strconv.Atoi(parts[1])
+	if err != nil || memGB <= 0 {
+		return nil, fnerrors.Newf("invalid --builder_shape %q: memory (GB) must be a positive integer", spec)
+	}
+
+	return &computev1beta.InstanceShape{
+		VirtualCpu:      int32(vcpu),
+		MemoryMegabytes: int32(memGB) * 1024,
+	}, nil
+}
+
 func readRemoteBuilderConfigs(dockerCli *command.DockerCli, name string) ([]BuilderConfig, error) {
 	var res []BuilderConfig
 

internal/providers/nscloud/api/rpc.go

@@ -14,6 +14,7 @@ import (
 	"time"
 
 	builderv1beta "buf.build/gen/go/namespace/cloud/protocolbuffers/go/proto/namespace/cloud/builder/v1beta"
+	computev1beta "buf.build/gen/go/namespace/cloud/protocolbuffers/go/proto/namespace/cloud/compute/v1beta"
 	"github.com/bcicen/jstream"
 	"github.com/cenkalti/backoff"
 	"go.uber.org/atomic"
@@ -365,6 +366,7 @@ type BuilderConfiguration struct {
 	SkipPrespawn bool
 	Name         string
 	Experimental string
+	Shape        *computev1beta.InstanceShape
 }
 
 func GetBuilderConfiguration(ctx context.Context, platform BuildPlatform, conf BuilderConfiguration) (*builderv1beta.GetBuilderConfigurationResponse, error) {
@@ -391,6 +393,7 @@ func GetBuilderConfiguration(ctx context.Context, platform BuildPlatform, conf B
 				SkipBuilderPreSpawn: conf.SkipPrespawn,
 				BuilderName:         conf.Name,
 				Experimental:        conf.Experimental,
+				BuilderShape:        conf.Shape,
 			})
 			if err != nil {
 				return nil, fnerrors.Newf("failed while creating %v build cluster: %w", platform, err)