Skip to content

Commit 8fb3579

Browse files
Code2LifeCode2Life
committed
fix: add platform support for container images and enhance architecture detection in Docker and Colima backends
1 parent e26db57 commit 8fb3579

5 files changed

Lines changed: 170 additions & 22 deletions

File tree

‎cmd/ggo/studio/studio.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ var (
3333
outputFormat string
3434
command []string
3535
endpoint string
36+
platform string // container platform (e.g., linux/amd64, linux/arm64)
3637
)
3738

3839
// NewStudioCmd creates the studio command
@@ -193,6 +194,7 @@ Examples:
193194
cmd.Flags().StringVar(&dockerHost, "docker-host", "", "Custom Docker socket path (e.g., unix:///path/to/docker.sock)")
194195
cmd.Flags().StringArrayVarP(&command, "command", "c", nil, "Container startup command or ENTRYPOINT args (can be specified multiple times)")
195196
cmd.Flags().StringVar(&endpoint, "endpoint", "", "Override GPU worker endpoint URL")
197+
cmd.Flags().StringVar(&platform, "platform", "", "Container image platform (e.g., linux/amd64, linux/arm64). Default: auto-detect from VM/host")
196198

197199
return cmd
198200
}
@@ -359,6 +361,7 @@ func buildCreateOptions(name string, shareInfo *api.SharePublicInfo) (*studio.Cr
359361
},
360362
Command: command,
361363
Endpoint: endpointOverride,
364+
Platform: platform,
362365
}, nil
363366
}
364367

‎internal/studio/backend_colima.go‎

Lines changed: 66 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,20 +47,63 @@ func (b *ColimaBackend) SetDockerHost(dockerHost string) {
4747
b.dockerBackend = NewDockerBackend()
4848
}
4949

50+
// GetVMArch returns the architecture of the Colima VM
51+
// Returns "amd64", "arm64", or empty string if detection fails
52+
func (b *ColimaBackend) GetVMArch(ctx context.Context) string {
53+
// Use colima status to get VM info
54+
cmd := exec.CommandContext(ctx, "colima", "status", "-p", b.profile, "--json")
55+
output, err := cmd.Output()
56+
if err != nil {
57+
// Fallback: try running uname -m inside docker
58+
unameCmd := exec.CommandContext(ctx, "docker", "run", "--rm", "alpine", "uname", "-m")
59+
unameCmd.Env = append(os.Environ(), fmt.Sprintf("DOCKER_HOST=%s", b.dockerHost))
60+
unameOutput, unameErr := unameCmd.Output()
61+
if unameErr == nil {
62+
arch := strings.TrimSpace(string(unameOutput))
63+
return NormalizeArch(arch)
64+
}
65+
return ""
66+
}
67+
68+
var status struct {
69+
Arch string `json:"arch"`
70+
}
71+
72+
if err := json.Unmarshal(output, &status); err != nil {
73+
return ""
74+
}
75+
76+
return NormalizeArch(status.Arch)
77+
}
78+
5079
// pullImageWithProgress pulls a Docker image with progress output to stderr
51-
func (b *ColimaBackend) pullImageWithProgress(ctx context.Context, image string) error {
52-
// Check if image already exists locally
80+
// If platform is specified, it will use --platform flag
81+
func (b *ColimaBackend) pullImageWithProgress(ctx context.Context, image, platform string) error {
82+
// Check if image already exists locally with correct platform
5383
checkCmd := exec.CommandContext(ctx, "docker", "image", "inspect", image)
5484
checkCmd.Env = append(os.Environ(), fmt.Sprintf("DOCKER_HOST=%s", b.dockerHost))
5585
if err := checkCmd.Run(); err == nil {
56-
return nil // Image exists
86+
// Image exists, but we should re-pull if platform is specified
87+
// to ensure we get the right architecture
88+
if platform == "" {
89+
return nil // Image exists and no specific platform requested
90+
}
5791
}
5892

59-
// Image doesn't exist, pull it with progress
93+
// Image doesn't exist or we need to ensure correct platform, pull it with progress
6094
fmt.Fprintf(os.Stderr, "\n Pulling image: %s\n", image)
95+
if platform != "" {
96+
fmt.Fprintf(os.Stderr, " Platform: %s\n", platform)
97+
}
6198
fmt.Fprintf(os.Stderr, " This may take a few minutes for large images...\n\n")
6299

63-
pullCmd := exec.CommandContext(ctx, "docker", "pull", image)
100+
pullArgs := []string{"pull"}
101+
if platform != "" {
102+
pullArgs = append(pullArgs, "--platform", platform)
103+
}
104+
pullArgs = append(pullArgs, image)
105+
106+
pullCmd := exec.CommandContext(ctx, "docker", pullArgs...)
64107
pullCmd.Env = append(os.Environ(), fmt.Sprintf("DOCKER_HOST=%s", b.dockerHost))
65108
// Stream output to stderr so user can see progress
66109
pullCmd.Stdout = os.Stderr
@@ -315,12 +358,29 @@ func (b *ColimaBackend) Create(ctx context.Context, opts *CreateOptions) (*Envir
315358
return nil, err
316359
}
317360

361+
// Determine the platform to use
362+
// If user specified platform, use it; otherwise detect from VM
363+
platform := opts.Platform
364+
if platform == "" {
365+
// Auto-detect VM architecture
366+
vmArch := b.GetVMArch(ctx)
367+
if vmArch != "" {
368+
platform = "linux/" + vmArch
369+
fmt.Fprintf(os.Stderr, " Auto-detected VM architecture: %s\n", platform)
370+
}
371+
}
372+
318373
// Generate container name with random suffix to avoid conflicts
319374
containerName := GenerateContainerName(opts.Name)
320375

321376
// Build docker run command
322377
args := []string{"run", "-d", "--name", containerName}
323378

379+
// Add platform flag if specified
380+
if platform != "" {
381+
args = append(args, "--platform", platform)
382+
}
383+
324384
// Add labels
325385
args = append(args, "--label", "ggo.managed=true")
326386
args = append(args, "--label", fmt.Sprintf("ggo.name=%s", opts.Name))
@@ -417,7 +477,7 @@ func (b *ColimaBackend) Create(ctx context.Context, opts *CreateOptions) (*Envir
417477
}
418478

419479
// Pull image first with progress visible to user
420-
if err := b.pullImageWithProgress(ctx, image); err != nil {
480+
if err := b.pullImageWithProgress(ctx, image, platform); err != nil {
421481
return nil, fmt.Errorf("failed to pull image: %w", err)
422482
}
423483

‎internal/studio/backend_docker.go‎

Lines changed: 51 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -71,21 +71,49 @@ func (b *DockerBackend) setDockerEnv(cmd *exec.Cmd) {
7171
}
7272
}
7373

74+
// GetHostArch returns the architecture of the Docker host
75+
// Returns "amd64", "arm64", or empty string if detection fails
76+
func (b *DockerBackend) GetHostArch(ctx context.Context) string {
77+
cmd := exec.CommandContext(ctx, b.dockerCmd, "info", "--format", "{{.Architecture}}")
78+
b.setDockerEnv(cmd)
79+
output, err := cmd.Output()
80+
if err != nil {
81+
return ""
82+
}
83+
84+
arch := strings.TrimSpace(string(output))
85+
return NormalizeArch(arch)
86+
}
87+
7488
// pullImageWithProgress pulls a Docker image with progress output to stderr
75-
func (b *DockerBackend) pullImageWithProgress(ctx context.Context, image string) error {
89+
// If platform is specified, it will use --platform flag
90+
func (b *DockerBackend) pullImageWithProgress(ctx context.Context, image, platform string) error {
7691
// Check if image already exists locally
7792
checkCmd := exec.CommandContext(ctx, b.dockerCmd, "image", "inspect", image)
7893
b.setDockerEnv(checkCmd)
7994
if err := checkCmd.Run(); err == nil {
80-
klog.V(2).Infof("Image %s already exists locally", image)
81-
return nil // Image exists
95+
// Image exists, but we should re-pull if platform is specified
96+
// to ensure we get the right architecture
97+
if platform == "" {
98+
klog.V(2).Infof("Image %s already exists locally", image)
99+
return nil // Image exists and no specific platform requested
100+
}
82101
}
83102

84-
// Image doesn't exist, pull it with progress
103+
// Image doesn't exist or we need to ensure correct platform, pull it with progress
85104
fmt.Fprintf(os.Stderr, "\n Pulling image: %s\n", image)
105+
if platform != "" {
106+
fmt.Fprintf(os.Stderr, " Platform: %s\n", platform)
107+
}
86108
fmt.Fprintf(os.Stderr, " This may take a few minutes for large images...\n\n")
87109

88-
pullCmd := exec.CommandContext(ctx, b.dockerCmd, "pull", image)
110+
pullArgs := []string{"pull"}
111+
if platform != "" {
112+
pullArgs = append(pullArgs, "--platform", platform)
113+
}
114+
pullArgs = append(pullArgs, image)
115+
116+
pullCmd := exec.CommandContext(ctx, b.dockerCmd, pullArgs...)
89117
b.setDockerEnv(pullCmd)
90118
// Stream output to stderr so user can see progress
91119
pullCmd.Stdout = os.Stderr
@@ -100,12 +128,29 @@ func (b *DockerBackend) pullImageWithProgress(ctx context.Context, image string)
100128
}
101129

102130
func (b *DockerBackend) Create(ctx context.Context, opts *CreateOptions) (*Environment, error) {
131+
// Determine the platform to use
132+
// If user specified platform, use it; otherwise detect from host
133+
platform := opts.Platform
134+
if platform == "" {
135+
// Auto-detect host architecture
136+
hostArch := b.GetHostArch(ctx)
137+
if hostArch != "" {
138+
platform = "linux/" + hostArch
139+
fmt.Fprintf(os.Stderr, " Auto-detected host architecture: %s\n", platform)
140+
}
141+
}
142+
103143
// Generate container name with random suffix to avoid conflicts
104144
containerName := GenerateContainerName(opts.Name)
105145

106146
// Build docker run command
107147
args := []string{"run", "-d", "--name", containerName}
108148

149+
// Add platform flag if specified
150+
if platform != "" {
151+
args = append(args, "--platform", platform)
152+
}
153+
109154
// Add labels
110155
args = append(args, "--label", "ggo.managed=true")
111156
args = append(args, "--label", fmt.Sprintf("ggo.name=%s", opts.Name))
@@ -209,7 +254,7 @@ func (b *DockerBackend) Create(ctx context.Context, opts *CreateOptions) (*Envir
209254
klog.V(2).Infof("Running docker command: %s %v", b.dockerCmd, args)
210255

211256
// Pull image first with progress visible to user
212-
if err := b.pullImageWithProgress(ctx, image); err != nil {
257+
if err := b.pullImageWithProgress(ctx, image, platform); err != nil {
213258
return nil, fmt.Errorf("failed to pull image: %w", err)
214259
}
215260

‎internal/studio/container_setup.go‎

Lines changed: 27 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -335,8 +335,11 @@ func setupSSHAuthorizedKeys(paths *platform.Paths, studioName string) (string, e
335335

336336
// getSSHVolumeMounts returns volume mounts for SSH access to the container
337337
// This includes:
338-
// 1. User's .ssh directory mounted to /root/.ssh (read-only for security)
338+
// 1. Individual SSH key files mounted to /root/.ssh/ (for git operations etc.)
339339
// 2. Generated authorized_keys file mounted to /root/.ssh/authorized_keys
340+
//
341+
// Note: We mount individual files instead of the entire .ssh directory to avoid
342+
// conflicts when also mounting authorized_keys (can't mount a file inside a read-only directory mount)
340343
func getSSHVolumeMounts(paths *platform.Paths, studioName string) []VolumeMount {
341344
var mounts []VolumeMount
342345

@@ -352,24 +355,38 @@ func getSSHVolumeMounts(paths *platform.Paths, studioName string) []VolumeMount
352355
return mounts
353356
}
354357

355-
// Mount user's .ssh directory to /root/.ssh (read-only for security)
356-
// This allows the container to use user's SSH keys for git operations etc.
357-
mounts = append(mounts, VolumeMount{
358-
HostPath: sshDir,
359-
ContainerPath: "/root/.ssh",
360-
ReadOnly: true,
361-
})
358+
// Mount individual SSH key files (for git operations etc.)
359+
// These are mounted as read-only for security
360+
sshKeyFiles := []string{
361+
"id_ed25519",
362+
"id_ecdsa",
363+
"id_rsa",
364+
"id_dsa",
365+
"config",
366+
"known_hosts",
367+
}
368+
369+
for _, keyFile := range sshKeyFiles {
370+
keyPath := filepath.Join(sshDir, keyFile)
371+
if _, err := os.Stat(keyPath); err == nil {
372+
mounts = append(mounts, VolumeMount{
373+
HostPath: keyPath,
374+
ContainerPath: "/root/.ssh/" + keyFile,
375+
ReadOnly: true,
376+
})
377+
klog.V(2).Infof("Mounting SSH file: %s", keyFile)
378+
}
379+
}
362380

363381
// Create and mount authorized_keys for SSH access to the container
364382
authorizedKeysPath, err := setupSSHAuthorizedKeys(paths, studioName)
365383
if err != nil {
366384
klog.V(2).Infof("Could not setup authorized_keys: %v (SSH access to container may require password)", err)
367385
} else {
368-
// Mount authorized_keys file - this overwrites the read-only .ssh mount for this specific file
369386
mounts = append(mounts, VolumeMount{
370387
HostPath: authorizedKeysPath,
371388
ContainerPath: "/root/.ssh/authorized_keys",
372-
ReadOnly: true,
389+
ReadOnly: false, // needs to be writable for SSH daemon to read
373390
})
374391
}
375392

‎internal/studio/types.go‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,26 @@ const (
104104
OSWindows = "windows"
105105
)
106106

107+
// Architecture constants for CPU architecture detection
108+
const (
109+
ArchX86_64 = "x86_64"
110+
ArchAarch64 = "aarch64"
111+
ArchAmd64 = "amd64"
112+
ArchArm64 = "arm64"
113+
)
114+
115+
// NormalizeArch normalizes architecture names to Docker/OCI format (amd64, arm64)
116+
func NormalizeArch(arch string) string {
117+
switch arch {
118+
case ArchX86_64:
119+
return ArchAmd64
120+
case ArchAarch64:
121+
return ArchArm64
122+
default:
123+
return arch
124+
}
125+
}
126+
107127
// Docker-related constants
108128
const (
109129
DefaultProtocolTCP = "tcp"
@@ -135,6 +155,9 @@ type CreateOptions struct {
135155
Command []string `json:"command,omitempty"`
136156
// Endpoint overrides the GPU worker endpoint URL
137157
Endpoint string `json:"endpoint,omitempty"`
158+
// Platform specifies the container image platform (e.g., linux/amd64, linux/arm64)
159+
// If empty, auto-detect from VM/host architecture
160+
Platform string `json:"platform,omitempty"`
138161
}
139162

140163
// PortMapping represents a port mapping

0 commit comments

Comments
 (0)