feat: enhance platform handling in studio commands and library management, default to linux/amd64 for container images

Author: Code2Life

cmd/ggo/studio/studio.go

@@ -223,7 +223,7 @@ Examples:
 	cmd.Flags().StringVar(&dockerHost, "docker-host", "", "Custom Docker socket path (e.g., unix:///path/to/docker.sock)")
 	cmd.Flags().StringArrayVarP(&command, "command", "c", nil, "Container startup command or ENTRYPOINT args (can be specified multiple times)")
 	cmd.Flags().StringVar(&endpoint, "endpoint", "", "Override GPU worker endpoint URL")
-	cmd.Flags().StringVar(&platform, "platform", "", "Container image platform (e.g., linux/amd64, linux/arm64). Default: auto-detect from VM/host")
+	cmd.Flags().StringVar(&platform, "platform", "", "Container image platform (e.g., linux/amd64, linux/arm64). Default: linux/amd64")
 
 	return cmd
 }
@@ -253,9 +253,17 @@ func runCreate(cmd *cobra.Command, args []string) error {
 		klog.Infof("Resolved share link: worker_id=%s vendor=%s connection_url=%s",
 			shareInfo.WorkerID, shareInfo.HardwareVendor, shareInfo.ConnectionURL)
 
+		// Determine target arch from platform flag (default: amd64)
+		targetArch := "amd64"
+		if platform != "" {
+			if parts := strings.SplitN(platform, "/", 2); len(parts) == 2 {
+				targetArch = parts[1]
+			}
+		}
+
 		// Download required GPU client libraries before creating studio
 		// Filter by vendor from share info to avoid downloading unnecessary libraries
-		if err := ensureRemoteGPUClientLibs(ctx, out, shareInfo.HardwareVendor); err != nil {
+		if err := ensureRemoteGPUClientLibs(ctx, out, shareInfo.HardwareVendor, targetArch); err != nil {
 			cmd.SilenceUsage = true
 			klog.Errorf("Failed to ensure GPU client libraries: error=%v", err)
 			return fmt.Errorf("failed to download GPU client libraries: %w", err)
@@ -307,16 +315,16 @@ func runCreate(cmd *cobra.Command, args []string) error {
 
 // ensureRemoteGPUClientLibs downloads remote-gpu-client libraries if not already present
 // vendorSlug filters by vendor (e.g., "nvidia", "amd") to avoid downloading unnecessary libraries
+// targetArch specifies the CPU architecture (e.g., "amd64", "arm64") for the target container platform
 // Note: Studio environments run in Linux containers, so we always download Linux libraries
-// using the current CPU architecture (arm64 or amd64)
-func ensureRemoteGPUClientLibs(ctx context.Context, out *tui.Output, vendorSlug string) error {
+func ensureRemoteGPUClientLibs(ctx context.Context, out *tui.Output, vendorSlug, targetArch string) error {
 	depsMgr := deps.NewManager()
 
 	// Target library types that are needed for GPU client functionality
 	targetTypes := []string{deps.LibraryTypeRemoteGPUClient, deps.LibraryTypeVGPULibrary}
 
 	if !out.IsJSON() {
-		out.Printf("Downloading GPU client libraries for %s...\n", vendorSlug)
+		out.Printf("Downloading GPU client libraries for %s (linux/%s)...\n", vendorSlug, targetArch)
 	}
 
 	progressFn := func(lib deps.Library, downloaded, total int64) {
@@ -326,9 +334,9 @@ func ensureRemoteGPUClientLibs(ctx context.Context, out *tui.Output, vendorSlug
 		}
 	}
 
-	// Studio environments always run in Linux containers, so we need Linux libraries
-	// Use current CPU architecture (arm64/amd64) since container arch matches host
-	libs, err := depsMgr.EnsureLibrariesByTypesForPlatform(ctx, targetTypes, vendorSlug, "linux", "", progressFn)
+	// Studio environments always run in Linux containers
+	// Download libraries for the specified target architecture
+	libs, err := depsMgr.EnsureLibrariesByTypesForPlatform(ctx, targetTypes, vendorSlug, "linux", targetArch, progressFn)
 	if err != nil {
 		return fmt.Errorf("failed to ensure GPU client libraries: %w", err)
 	}
@@ -396,6 +404,12 @@ func buildCreateOptions(name string, shareInfo *api.SharePublicInfo) (*studio.Cr
 		}
 	}
 
+	// Default platform to linux/amd64 for studio containers
+	effectivePlatform := platform
+	if effectivePlatform == "" {
+		effectivePlatform = "linux/amd64"
+	}
+
 	return &studio.CreateOptions{
 		Name:           name,
 		Mode:           studioMode,
@@ -412,7 +426,7 @@ func buildCreateOptions(name string, shareInfo *api.SharePublicInfo) (*studio.Cr
 		},
 		Command:     command,
 		Endpoint:    endpointOverride,
-		Platform:    platform,
+		Platform:    effectivePlatform,
 		UseLocalGPU: gpuWorkerURL == "" && (studioMode == studio.ModeDocker || studioMode == studio.ModeWSL || studioMode == studio.ModeAuto),
 	}, nil
 }

internal/deps/deps.go

@@ -661,29 +661,38 @@ func (m *Manager) ComputeUpdateDiff() (*UpdateDiff, error) {
 	return diff, nil
 }
 
-// DownloadLibrary downloads a library to the cache directory
+// DownloadLibrary downloads a library to the default (flat) libs directory
 func (m *Manager) DownloadLibrary(ctx context.Context, lib Library, progressFn func(downloaded, total int64)) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	return m.downloadLibraryUnsafe(ctx, lib, progressFn)
+	return m.downloadLibraryToDir(ctx, lib, m.paths.LibsDir(), progressFn)
 }
 
-func (m *Manager) downloadLibraryUnsafe(ctx context.Context, lib Library, progressFn func(downloaded, total int64)) error {
+// DownloadLibraryToDir downloads a library to a specific libs directory.
+func (m *Manager) DownloadLibraryToDir(ctx context.Context, lib Library, libsDir string, progressFn func(downloaded, total int64)) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	return m.downloadLibraryToDir(ctx, lib, libsDir, progressFn)
+}
+
+// downloadLibraryToDir downloads a library to a specific libs directory.
+// Shared libraries (.so/.dll) go to libsDir; binaries go to cache root.
+func (m *Manager) downloadLibraryToDir(ctx context.Context, lib Library, libsDir string, progressFn func(downloaded, total int64)) error {
 	// Ensure cache directory exists
 	cacheDir := m.paths.CacheDir()
 	if err := os.MkdirAll(cacheDir, 0755); err != nil {
 		return fmt.Errorf("failed to create cache directory: %w", err)
 	}
 
 	// Ensure libs directory exists for shared libraries
-	libsDir := m.paths.LibsDir()
 	if err := os.MkdirAll(libsDir, 0755); err != nil {
 		return fmt.Errorf("failed to create libs directory: %w", err)
 	}
 
 	// Destination path - shared libraries go to libs dir, binaries to cache dir
-	destPath := m.GetLibraryPath(lib.Name)
+	destPath := m.GetLibraryPathInDir(lib.Name, libsDir)
 	tmpPath := destPath + ".tmp"
 
 	// Check if already downloaded with correct version
@@ -866,7 +875,7 @@ func (m *Manager) DownloadAllRequired(ctx context.Context, progressFn func(lib L
 			}
 		}
 
-		if err := m.downloadLibraryUnsafe(ctx, lib, libProgressFn); err != nil {
+		if err := m.downloadLibraryToDir(ctx, lib, m.paths.LibsDir(), libProgressFn); err != nil {
 			result.Status = DownloadStatusFailed
 			result.Error = err.Error()
 			klog.Errorf("Failed to download library: name=%s error=%v", lib.Name, err)
@@ -917,6 +926,15 @@ func (m *Manager) GetLibraryPath(name string) string {
 	return filepath.Join(m.paths.CacheDir(), name)
 }
 
+// GetLibraryPathInDir returns the path to a library in a specific libs directory.
+// For shared libraries, uses the given libsDir; for binaries, uses cache root.
+func (m *Manager) GetLibraryPathInDir(name string, libsDir string) string {
+	if isSharedLibrary(name) {
+		return filepath.Join(libsDir, name)
+	}
+	return filepath.Join(m.paths.CacheDir(), name)
+}
+
 // GetLibsDir returns the directory for .so/.dll library files
 // This directory is used for LD_LIBRARY_PATH, ld.so.conf, and ld.so.preload
 func (m *Manager) GetLibsDir() string {
@@ -1115,9 +1133,28 @@ func (m *Manager) EnsureLibrariesByTypes(ctx context.Context, libTypes []string,
 
 // EnsureLibrariesByTypesForPlatform ensures ALL libraries of the specified types exist and are downloaded for a specific platform
 // targetOS and targetArch specify the target platform (e.g., "linux", "arm64")
-// If both are empty, uses the current platform
-// This is useful when running on macOS but needing Linux libraries for containers
+// If both are empty, uses the current platform and the flat libs directory (agent/worker path).
+// If targetOS is specified, uses arch-specific subdirectory (e.g., libs/linux-amd64/) to avoid
+// collisions between different architectures (studio/use/launch path).
 func (m *Manager) EnsureLibrariesByTypesForPlatform(ctx context.Context, libTypes []string, vendorSlug, targetOS, targetArch string, progressFn func(lib Library, downloaded, total int64)) ([]Library, error) {
+	// Resolve effective platform values
+	effectiveOS := targetOS
+	effectiveArch := targetArch
+	if effectiveOS == "" {
+		effectiveOS = runtime.GOOS
+	}
+	if effectiveArch == "" {
+		effectiveArch = runtime.GOARCH
+	}
+
+	// Determine the libs directory:
+	// - If targetOS was explicitly specified, use arch-specific subdir (studio/use/launch)
+	// - Otherwise, use flat libs dir (agent/worker)
+	libsDir := m.paths.LibsDir()
+	if targetOS != "" {
+		libsDir = m.paths.LibsDirForPlatform(effectiveOS, effectiveArch)
+	}
+
 	// Ensure deps manifest exists and is up to date for the target platform
 	if err := m.ensureDepsManifestForPlatform(ctx, targetOS, targetArch); err != nil {
 		return nil, err
@@ -1167,7 +1204,7 @@ func (m *Manager) EnsureLibrariesByTypesForPlatform(ctx context.Context, libType
 	var toDownload []Library
 
 	for _, lib := range targetLibs {
-		filePath := m.GetLibraryPath(lib.Name)
+		filePath := m.GetLibraryPathInDir(lib.Name, libsDir)
 		downloadedLib, exists := downloaded.Libraries[lib.Key()]
 
 		needsDownload := false
@@ -1186,7 +1223,7 @@ func (m *Manager) EnsureLibrariesByTypesForPlatform(ctx context.Context, libType
 		}
 	}
 
-	klog.V(4).Infof("Libraries to download: %d out of %d total (libs will go to: %s)", len(toDownload), len(targetLibs), m.paths.LibsDir())
+	klog.V(4).Infof("Libraries to download: %d out of %d total (libs will go to: %s)", len(toDownload), len(targetLibs), libsDir)
 
 	// Download missing libraries
 	for _, lib := range toDownload {
@@ -1196,8 +1233,8 @@ func (m *Manager) EnsureLibrariesByTypesForPlatform(ctx context.Context, libType
 			}
 		}
 
-		klog.Infof("Downloading library: name=%s version=%s type=%s", lib.Name, lib.Version, lib.Type)
-		if err := m.DownloadLibrary(ctx, lib, libProgressFn); err != nil {
+		klog.Infof("Downloading library: name=%s version=%s type=%s to=%s", lib.Name, lib.Version, lib.Type, libsDir)
+		if err := m.DownloadLibraryToDir(ctx, lib, libsDir, libProgressFn); err != nil {
 			return nil, fmt.Errorf("failed to download library %s: %w", lib.Name, err)
 		}
 	}

internal/platform/paths.go

@@ -72,10 +72,19 @@ func (p *Paths) LibDir() string {
 // This is separate from CacheDir which contains all downloaded files (including binaries)
 // The libs directory is used for LD_LIBRARY_PATH, ld.so.conf, and ld.so.preload
 // All platforms: ~/.gpugo/cache/libs (or CacheDir/libs)
+// Used by agent/worker which always runs on the native host architecture.
 func (p *Paths) LibsDir() string {
 	return filepath.Join(p.cacheDir, "libs")
 }
 
+// LibsDirForPlatform returns the arch-specific libs directory for a given OS and architecture.
+// This is used by studio/use/launch which may download libraries for a different platform
+// (e.g., linux/amd64 libs on an arm64 Mac).
+// Layout: ~/.gpugo/cache/libs/{os}-{arch} (e.g., ~/.gpugo/cache/libs/linux-amd64)
+func (p *Paths) LibsDirForPlatform(targetOS, arch string) string {
+	return filepath.Join(p.cacheDir, "libs", targetOS+"-"+arch)
+}
+
 // BinDir returns the directory for binaries
 // All platforms: ~/.gpugo/bin (or UserDir/bin)
 func (p *Paths) BinDir() string {

internal/studio/backend_apple.go

@@ -132,6 +132,7 @@ func (b *AppleContainerBackend) Create(ctx context.Context, opts *CreateOptions)
 		StudioName:     opts.Name,
 		GPUWorkerURL:   gpuWorkerURL,
 		HardwareVendor: opts.HardwareVendor,
+		Platform:       opts.Platform,
 		MountUserHome:  false, // /Users is mounted directly into the container
 		SkipFileMounts: true,
 	}

internal/studio/backend_colima.go

@@ -425,6 +425,7 @@ func (b *ColimaBackend) Create(ctx context.Context, opts *CreateOptions) (*Envir
 		StudioName:     opts.Name,
 		GPUWorkerURL:   gpuWorkerURL,
 		HardwareVendor: opts.HardwareVendor,
+		Platform:       opts.Platform,
 		MountUserHome:  false, // /Users is mounted directly into the container
 	}
 

internal/studio/backend_docker.go

@@ -203,6 +203,7 @@ func (b *DockerBackend) Create(ctx context.Context, opts *CreateOptions) (*Envir
 		StudioName:     opts.Name,
 		GPUWorkerURL:   gpuWorkerURL,
 		HardwareVendor: opts.HardwareVendor,
+		Platform:       opts.Platform,
 		MountUserHome:  !opts.NoUserVolume,
 	}
 

internal/studio/backend_wsl.go

@@ -266,6 +266,7 @@ func (b *WSLBackend) Create(ctx context.Context, opts *CreateOptions) (*Environm
 		StudioName:     opts.Name,
 		GPUWorkerURL:   gpuWorkerURL,
 		HardwareVendor: opts.HardwareVendor,
+		Platform:       opts.Platform,
 		MountUserHome:  !opts.NoUserVolume,
 	}
 

internal/studio/container_setup.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
-	"runtime"
 	"strings"
 
 	"github.com/NexusGPU/gpu-go/internal/deps"
@@ -21,6 +20,10 @@ type ContainerSetupConfig struct {
 	GPUWorkerURL string
 	// HardwareVendor is the GPU vendor (nvidia, amd, hygon)
 	HardwareVendor string
+	// Platform is the container platform (e.g., "linux/amd64", "linux/arm64")
+	// Used to determine which arch-specific libs to download and mount.
+	// If empty, defaults to linux/amd64.
+	Platform string
 	// MountUserHome indicates whether to mount the user's home directory
 	MountUserHome bool
 	// SkipSSHMounts disables mounting SSH key files into the container
@@ -43,7 +46,7 @@ type ContainerSetupResult struct {
 
 // SetupContainerGPUEnv sets up the GPU environment for a container
 // This performs the same setup as `ggo use` does for Linux:
-// 1. Downloads GPU client libraries for Linux (using host CPU arch)
+// 1. Downloads GPU client libraries for Linux (using target platform arch)
 // 2. Sets up GPU environment (env vars, ld.so.preload, ld.so.conf.d)
 // 3. Optionally mounts user home directory
 func SetupContainerGPUEnv(ctx context.Context, config *ContainerSetupConfig) (*ContainerSetupResult, error) {
@@ -56,10 +59,20 @@ func SetupContainerGPUEnv(ctx context.Context, config *ContainerSetupConfig) (*C
 	normalizedName := platform.NormalizeName(config.StudioName)
 	vendor := ParseVendor(config.HardwareVendor)
 
+	// Parse target arch from platform string (e.g., "linux/amd64" → "amd64")
+	targetArch := "amd64" // default
+	if config.Platform != "" {
+		if parts := strings.SplitN(config.Platform, "/", 2); len(parts) == 2 {
+			targetArch = parts[1]
+		}
+	}
+	// Arch-specific libs directory (e.g., ~/.gpugo/cache/libs/linux-amd64/)
+	libsDir := paths.LibsDirForPlatform("linux", targetArch)
+
 	// Step 1: Download GPU client libraries for Linux (container target)
-	// Libraries are downloaded for Linux with host CPU architecture
+	// Libraries are downloaded for Linux with the target CPU architecture
 	if config.GPUWorkerURL != "" {
-		if err := ensureGPUClientLibraries(ctx, vendor); err != nil {
+		if err := ensureGPUClientLibraries(ctx, vendor, targetArch); err != nil {
 			klog.Warningf("Failed to download GPU client libraries: %v (continuing anyway)", err)
 		} else {
 			result.LibrariesDownloaded = true
@@ -72,6 +85,7 @@ func SetupContainerGPUEnv(ctx context.Context, config *ContainerSetupConfig) (*C
 			Vendor:        vendor,
 			ConnectionURL: config.GPUWorkerURL,
 			CachePath:     paths.CacheDir(),
+			LibsPath:      libsDir,
 			LogPath:       paths.StudioLogsDir(normalizedName),
 			StudioName:    normalizedName,
 			IsContainer:   true,
@@ -89,7 +103,7 @@ func SetupContainerGPUEnv(ctx context.Context, config *ContainerSetupConfig) (*C
 
 		if config.SkipFileMounts {
 			result.EnvVars["LD_LIBRARY_PATH"] = "/opt/gpugo/libs"
-			if preload := buildContainerLDPreload(paths.LibsDir(), vendor); preload != "" {
+			if preload := buildContainerLDPreload(libsDir, vendor); preload != "" {
 				result.EnvVars["LD_PRELOAD"] = preload
 			}
 		}
@@ -134,7 +148,7 @@ func SetupContainerGPUEnv(ctx context.Context, config *ContainerSetupConfig) (*C
 
 	// Step 5: Download and mount GPU binary (like nvidia-smi) to /usr/local/bin/
 	if !config.SkipFileMounts && config.GPUWorkerURL != "" && config.HardwareVendor != "" {
-		gpuBinMount, err := ensureAndMountGPUBinary(ctx, paths, config.HardwareVendor)
+		gpuBinMount, err := ensureAndMountGPUBinary(ctx, paths, config.HardwareVendor, targetArch)
 		if err != nil {
 			klog.Warningf("Failed to setup GPU binary mount: %v (continuing without it)", err)
 		} else if gpuBinMount != nil {
@@ -152,9 +166,9 @@ func SetupContainerGPUEnv(ctx context.Context, config *ContainerSetupConfig) (*C
 
 // ensureAndMountGPUBinary downloads GPU binary (like nvidia-smi) and returns a volume mount
 // The binary is mounted to /usr/local/bin/ in the container
-func ensureAndMountGPUBinary(ctx context.Context, paths *platform.Paths, vendorSlug string) (*VolumeMount, error) {
-	// Studios run in Linux containers, so download Linux binary with host CPU arch
-	binPath, err := deps.EnsureGPUBinaryForPlatform(ctx, paths, vendorSlug, "linux", runtime.GOARCH)
+func ensureAndMountGPUBinary(ctx context.Context, paths *platform.Paths, vendorSlug, targetArch string) (*VolumeMount, error) {
+	// Studios run in Linux containers, so download Linux binary with target CPU arch
+	binPath, err := deps.EnsureGPUBinaryForPlatform(ctx, paths, vendorSlug, "linux", targetArch)
 	if err != nil {
 		return nil, fmt.Errorf("failed to ensure GPU binary: %w", err)
 	}
@@ -180,8 +194,8 @@ func ensureAndMountGPUBinary(ctx context.Context, paths *platform.Paths, vendorS
 }
 
 // ensureGPUClientLibraries downloads GPU client libraries for Linux containers
-// Libraries are downloaded for Linux platform with the current host CPU architecture
-func ensureGPUClientLibraries(ctx context.Context, vendor GPUVendor) error {
+// Libraries are downloaded for Linux platform with the specified target CPU architecture
+func ensureGPUClientLibraries(ctx context.Context, vendor GPUVendor, targetArch string) error {
 	depsMgr := deps.NewManager()
 
 	// Target library types needed for GPU client functionality
@@ -198,11 +212,11 @@ func ensureGPUClientLibraries(ctx context.Context, vendor GPUVendor) error {
 		vendorSlug = "hygon"
 	}
 
-	klog.Infof("Downloading GPU client libraries for %s (linux/%s)...", vendorSlug, runtime.GOARCH)
+	klog.Infof("Downloading GPU client libraries for %s (linux/%s)...", vendorSlug, targetArch)
 
 	// Studios run in Linux containers, so we always download Linux libraries
-	// CPU architecture matches the host (arm64 on Apple Silicon, amd64 on Intel/AMD)
-	_, err := depsMgr.EnsureLibrariesByTypesForPlatform(ctx, targetTypes, vendorSlug, "linux", "", nil)
+	// Use the specified target architecture from the platform flag
+	_, err := depsMgr.EnsureLibrariesByTypesForPlatform(ctx, targetTypes, vendorSlug, "linux", targetArch, nil)
 	if err != nil {
 		return fmt.Errorf("failed to ensure GPU client libraries: %w", err)
 	}

internal/studio/libdownloader.go

@@ -64,10 +64,10 @@ func DetectArchitecture() Architecture {
 
 	// Normalize architecture names
 	switch arch.Arch {
-	case "amd64":
-		arch.Arch = "amd64"
-	case "arm64":
-		arch.Arch = "arm64"
+	case ArchAmd64:
+		arch.Arch = ArchAmd64
+	case ArchArm64:
+		arch.Arch = ArchArm64
 	case "386":
 		arch.Arch = "386"
 	default: