Skip to content

Commit 9b1e66e

Browse files
Code2LifeCode2Life
committed
fix: add unit tests for SSH configuration management in the studio manager
1 parent e7b6957 commit 9b1e66e

14 files changed

Lines changed: 288 additions & 86 deletions

File tree

‎.vscode/launch.json‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,18 @@
4141
"agent", "status"
4242
],
4343
"program": "${workspaceFolder}/cmd/ggo/main.go",
44+
},
45+
{
46+
"name": "Debug GPU Go Studio Create",
47+
"type": "go",
48+
"request": "launch",
49+
"mode": "auto",
50+
"console": "integratedTerminal",
51+
"args": [
52+
"studio", "create", "my-studio-abc", "-s", "https://gpu.tf/s/bhchnf",
53+
// "-m", "colima", "--platform", "linux/arm64", "-i", "nvcr.io/nvidia/pytorch:26.01-py3-igpu", "-c", "sleep 1d"
54+
],
55+
"program": "${workspaceFolder}/cmd/ggo/main.go",
4456
}
4557
]
46-
}
58+
}

‎cmd/ggo/launch/launch_linux.go‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -148,8 +148,8 @@ func runLaunch(args []string, shareLink, serverURL string, verbose bool) error {
148148

149149
// Setup log path (consistent with ggo use)
150150
studioName := "current-os"
151-
logPath := paths.StudioLogsDir(studioName)
152-
if err := os.MkdirAll(logPath, 0755); err != nil {
151+
logPath := paths.StudioLogFilePath(studioName)
152+
if err := os.MkdirAll(filepath.Dir(logPath), 0755); err != nil {
153153
klog.Warningf("Failed to create log directory: %v", err)
154154
}
155155

@@ -192,7 +192,16 @@ func runLaunch(args []string, shareLink, serverURL string, verbose bool) error {
192192
env = setEnvVar(env, "TENSOR_FUSION_OPERATOR_CONNECTION_INFO", connectionInfo)
193193
env = setEnvVar(env, "TF_LOG_PATH", logPath)
194194
env = setEnvVar(env, "TF_LOG_LEVEL", getEnvDefault("TF_LOG_LEVEL", "info"))
195-
env = setEnvVar(env, "TF_ENABLE_LOG", getEnvDefault("TF_ENABLE_LOG", "0"))
195+
env = setEnvVar(env, "TF_ENABLE_LOG", getEnvDefault("TF_ENABLE_LOG", "1"))
196+
197+
// Add GPU bin directory to PATH (for nvidia-smi, amdsmi, etc.)
198+
binDir := filepath.Join(cacheDir, "bin")
199+
existingPath := os.Getenv("PATH")
200+
if existingPath != "" {
201+
env = setEnvVar(env, "PATH", binDir+":"+existingPath)
202+
} else {
203+
env = setEnvVar(env, "PATH", binDir)
204+
}
196205

197206
// Build LD_LIBRARY_PATH
198207
existingLDPath := os.Getenv("LD_LIBRARY_PATH")

‎cmd/ggo/launch/launch_windows.go‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -223,8 +223,8 @@ func runLaunch(args []string, shareLink, serverURL string, verbose bool) error {
223223

224224
// Setup log path (consistent with ggo use)
225225
studioName := "current-os"
226-
logPath := paths.StudioLogsDir(studioName)
227-
if err := os.MkdirAll(logPath, 0755); err != nil {
226+
logPath := paths.StudioLogFilePath(studioName)
227+
if err := os.MkdirAll(filepath.Dir(logPath), 0755); err != nil {
228228
klog.Warningf("Failed to create log directory: %v", err)
229229
}
230230

@@ -265,7 +265,17 @@ func runLaunch(args []string, shareLink, serverURL string, verbose bool) error {
265265
env = setEnvVar(env, "TENSOR_FUSION_OPERATOR_CONNECTION_INFO", connectionInfo)
266266
env = setEnvVar(env, "TF_LOG_PATH", logPath)
267267
env = setEnvVar(env, "TF_LOG_LEVEL", getEnvDefault("TF_LOG_LEVEL", "info"))
268-
env = setEnvVar(env, "TF_ENABLE_LOG", getEnvDefault("TF_ENABLE_LOG", "0"))
268+
env = setEnvVar(env, "TF_ENABLE_LOG", getEnvDefault("TF_ENABLE_LOG", "1"))
269+
270+
// Add GPU bin directory to PATH (for nvidia-smi, etc.)
271+
binDir := filepath.Join(paths.CacheDir(), "bin")
272+
existingPath := os.Getenv("PATH")
273+
if existingPath != "" {
274+
env = setEnvVar(env, "PATH", binDir+";"+existingPath)
275+
} else {
276+
env = setEnvVar(env, "PATH", binDir)
277+
}
278+
269279
execCmd.Env = env
270280

271281
// Run the command

‎cmd/ggo/use/use.go‎

Lines changed: 46 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"fmt"
88
"os"
99
"os/exec"
10+
"os/signal"
1011
"path/filepath"
1112
"runtime"
1213
"strings"
@@ -332,7 +333,7 @@ func renderUnixEnv(shareInfo *api.SharePublicInfo, config *studio.GPUEnvConfig,
332333
out.Println()
333334
out.Printf(" Connection URL: %s\n", shareInfo.ConnectionURL)
334335
out.Printf(" Hardware: %s\n", shareInfo.HardwareVendor)
335-
out.Printf(" Log Path: %s\n", config.LogPath)
336+
out.Printf(" Log Path: %s\n", envResult.EnvVars["TF_LOG_PATH"])
336337
out.Println()
337338
}
338339

@@ -448,10 +449,26 @@ func launchGPUShell(config *studio.GPUEnvConfig, envResult *studio.GPUEnvResult,
448449

449450
// Print GPU environment banner
450451
fmt.Printf("\n%s GPU environment activated %s\n", styles().Success.Render("✓"), styles().Muted.Render("(type 'exit' to deactivate)"))
451-
fmt.Println()
452+
fmt.Printf("%s\n\n", styles().Muted.Render("Ctrl+C interrupts the current command, not the GPU environment."))
453+
454+
// Ignore SIGINT in parent process - let the child shell handle Ctrl+C naturally.
455+
// Without this, Ctrl+C kills both the parent ggo process and the child shell,
456+
// causing the entire rGPU environment to exit unexpectedly.
457+
sigChan := make(chan os.Signal, 1)
458+
signal.Notify(sigChan, os.Interrupt)
459+
go func() {
460+
for range sigChan {
461+
// Intentionally ignored - the child shell handles Ctrl+C
462+
}
463+
}()
452464

453465
// Run the shell and wait for it to exit
454-
return cmd.Run()
466+
err := cmd.Run()
467+
468+
// Restore default signal handling after shell exits
469+
signal.Stop(sigChan)
470+
471+
return err
455472
}
456473

457474
// styles returns the default TUI styles
@@ -837,7 +854,7 @@ func renderWindowsEnv(shareInfo *api.SharePublicInfo, config *studio.GPUEnvConfi
837854
out.Println()
838855
out.Printf(" Connection URL: %s\n", shareInfo.ConnectionURL)
839856
out.Printf(" Hardware: %s\n", shareInfo.HardwareVendor)
840-
out.Printf(" Log Path: %s\n", config.LogPath)
857+
out.Printf(" Log Path: %s\n", envResult.EnvVars["TF_LOG_PATH"])
841858
out.Println()
842859
}
843860

@@ -990,12 +1007,13 @@ func launchGPUShellWindows(config *studio.GPUEnvConfig, envResult *studio.GPUEnv
9901007
env = append(env, fmt.Sprintf("%s=%s", k, v))
9911008
}
9921009

993-
// Add libs path to PATH - libs directory contains only .dll files
1010+
// Add libs path and GPU bin directory to PATH
1011+
binDir := getGPUBinDir()
9941012
existingPath := os.Getenv("PATH")
9951013
if existingPath != "" {
996-
env = append(env, fmt.Sprintf("PATH=%s;%s", libsPath, existingPath))
1014+
env = append(env, fmt.Sprintf("PATH=%s;%s;%s", binDir, libsPath, existingPath))
9971015
} else {
998-
env = append(env, fmt.Sprintf("PATH=%s", libsPath))
1016+
env = append(env, fmt.Sprintf("PATH=%s;%s", binDir, libsPath))
9991017
}
10001018

10011019
// Set CUDA_PATH - point to libs directory
@@ -1008,16 +1026,33 @@ func launchGPUShellWindows(config *studio.GPUEnvConfig, envResult *studio.GPUEnv
10081026
// Mark as GPU Go activated
10091027
env = append(env, "_GGO_ACTIVE=1")
10101028
env = append(env, fmt.Sprintf("_GGO_LIBS_PATH=%s", libsPath))
1029+
env = append(env, fmt.Sprintf("_GGO_BIN_PATH=%s", binDir))
10111030

10121031
cmd.Env = env
10131032

10141033
// Print GPU environment banner
10151034
styles := tui.DefaultStyles()
10161035
fmt.Printf("\n%s GPU environment activated %s\n", styles.Success.Render("✓"), styles.Muted.Render("(type 'exit' to deactivate)"))
1017-
fmt.Println()
1036+
fmt.Printf("%s\n\n", styles.Muted.Render("Ctrl+C interrupts the current command, not the GPU environment."))
1037+
1038+
// Ignore SIGINT (Ctrl+C) in parent process - let the child shell handle it.
1039+
// On Windows, CTRL_C_EVENT is sent to all processes in the console.
1040+
// Without this, the Go runtime terminates ggo on Ctrl+C, killing the shell.
1041+
sigChan := make(chan os.Signal, 1)
1042+
signal.Notify(sigChan, os.Interrupt)
1043+
go func() {
1044+
for range sigChan {
1045+
// Intentionally ignored - the child shell handles Ctrl+C
1046+
}
1047+
}()
10181048

10191049
// Run the shell and wait for it to exit
1020-
return cmd.Run()
1050+
err := cmd.Run()
1051+
1052+
// Restore default signal handling after shell exits
1053+
signal.Stop(sigChan)
1054+
1055+
return err
10211056
}
10221057

10231058
// tempEnvResultWindows implements Renderable for Windows temporary env setup
@@ -1132,7 +1167,7 @@ func setupLongTermUnix(shareInfo *api.SharePublicInfo, config *studio.GPUEnvConf
11321167
out.Printf(" Config directory: %s\n", outputDir)
11331168
out.Printf(" Connection URL: %s\n", shareInfo.ConnectionURL)
11341169
out.Printf(" Hardware: %s\n", shareInfo.HardwareVendor)
1135-
out.Printf(" Log Path: %s\n", config.LogPath)
1170+
out.Printf(" Log Path: %s\n", envResult.EnvVars["TF_LOG_PATH"])
11361171
out.Println()
11371172
}
11381173

@@ -1265,7 +1300,7 @@ func setupLongTermWindows(shareInfo *api.SharePublicInfo, config *studio.GPUEnvC
12651300
out.Printf(" Config directory: %s\n", outputDir)
12661301
out.Printf(" Connection URL: %s\n", shareInfo.ConnectionURL)
12671302
out.Printf(" Hardware: %s\n", shareInfo.HardwareVendor)
1268-
out.Printf(" Log Path: %s\n", config.LogPath)
1303+
out.Printf(" Log Path: %s\n", envResult.EnvVars["TF_LOG_PATH"])
12691304
out.Println()
12701305
}
12711306

‎docs/studio-guide.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ ggo use abc123 --long-term
3636
| `TENSOR_FUSION_OPERATOR_CONNECTION_INFO` | 远程 GPU 连接信息（必需） |
3737
| `TF_LOG_LEVEL` | 日志级别，默认 `info` |
3838
| `TF_ENABLE_LOG` | 启用日志，默认 `1` |
39-
| `TF_LOG_PATH` | 日志输出路径 |
39+
| `TF_LOG_PATH` | 日志输出路径（格式：`logs-YYYY-mm-dd.txt`） |
4040

4141
### Linux 库配置
4242

‎internal/agent/sse.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import (
1111
)
1212

1313
const (
14-
sseEndpoint = "https://sse.tensor-fusion.ai/stream"
14+
sseEndpoint = "https://sse.tensor-fusion.ai/v1/stream"
1515
sseReconnectMin = 1 * time.Second
1616
sseReconnectMax = 30 * time.Second
1717
sseDebounceDelay = 500 * time.Millisecond

‎internal/platform/paths.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"regexp"
77
"runtime"
88
"strings"
9+
"time"
910
)
1011

1112
// Platform constants for runtime.GOOS comparisons
@@ -214,6 +215,14 @@ func (p *Paths) StudioLogsDir(name string) string {
214215
return filepath.Join(p.StudioDir(), NormalizeName(name), "logs")
215216
}
216217

218+
// StudioLogFilePath returns the path to the daily log file for a studio
219+
// Format: ~/.gpugo/studio/{name}/logs/logs-YYYY-mm-dd.txt
220+
// Caller must ensure the parent directory exists (e.g. via EnsureStudioDirs or MkdirAll on StudioLogsDir)
221+
func (p *Paths) StudioLogFilePath(name string) string {
222+
fileName := "logs-" + time.Now().Format("2006-01-02") + ".txt"
223+
return filepath.Join(p.StudioLogsDir(name), fileName)
224+
}
225+
217226
// StudioConfigDir returns the config directory for a specific studio
218227
// Used for storing ld.so.preload and ld.so.conf content for the studio
219228
// All platforms: ~/.gpugo/studio/{name}/config

‎internal/studio/backend_apple.go‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ func (b *AppleContainerBackend) Create(ctx context.Context, opts *CreateOptions)
132132
StudioName: opts.Name,
133133
GPUWorkerURL: gpuWorkerURL,
134134
HardwareVendor: opts.HardwareVendor,
135-
MountUserHome: !opts.NoUserVolume,
135+
MountUserHome: false, // /Users is mounted directly into the container
136136
SkipFileMounts: true,
137137
}
138138

@@ -157,6 +157,11 @@ func (b *AppleContainerBackend) Create(ctx context.Context, opts *CreateOptions)
157157
args = append(args, "-v", mountOpt)
158158
}
159159

160+
// Mount /Users into container so host files are accessible at the same path
161+
if !opts.NoUserVolume {
162+
args = append(args, "-v", "/Users:/Users")
163+
}
164+
160165
// Add resource limits
161166
if opts.Resources.CPUs > 0 {
162167
cpus := int64(math.Ceil(opts.Resources.CPUs))
@@ -166,15 +171,13 @@ func (b *AppleContainerBackend) Create(ctx context.Context, opts *CreateOptions)
166171
args = append(args, "--cpus", strconv.FormatInt(cpus, 10))
167172
}
168173

169-
// Set default memory to 1/4 of system memory if not specified
174+
// Set default memory to 1/4 of system memory minus 1GB (reserve for VM overhead)
170175
memoryLimit := opts.Resources.Memory
171176
if memoryLimit == "" {
172177
memGB, err := getSystemMemoryGB()
173178
if err == nil && memGB > 0 {
174-
memAllocated := (memGB + 3) / 4
175-
if memAllocated < 1 {
176-
memAllocated = 1
177-
}
179+
// Subtract 1GB for VM overhead
180+
memAllocated := max((memGB+3)/4-1, 1)
178181
memoryLimit = fmt.Sprintf("%dG", memAllocated)
179182
}
180183
}

‎internal/studio/backend_colima.go‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -425,7 +425,7 @@ func (b *ColimaBackend) Create(ctx context.Context, opts *CreateOptions) (*Envir
425425
StudioName: opts.Name,
426426
GPUWorkerURL: gpuWorkerURL,
427427
HardwareVendor: opts.HardwareVendor,
428-
MountUserHome: !opts.NoUserVolume,
428+
MountUserHome: false, // /Users is mounted directly into the container
429429
}
430430

431431
setupResult, err := SetupContainerGPUEnv(ctx, setupConfig)
@@ -449,18 +449,23 @@ func (b *ColimaBackend) Create(ctx context.Context, opts *CreateOptions) (*Envir
449449
args = append(args, "-v", mountOpt)
450450
}
451451

452+
// Mount /Users into container so host files are accessible at the same path
453+
if !opts.NoUserVolume && runtime.GOOS == OSDarwin {
454+
args = append(args, "-v", "/Users:/Users")
455+
}
456+
452457
// Add resource limits
453458
if opts.Resources.CPUs > 0 {
454459
args = append(args, "--cpus", fmt.Sprintf("%.2f", opts.Resources.CPUs))
455460
}
456461

457-
// Set default memory to 1/4 of system memory if not specified
462+
// Set default memory to 1/4 of system memory minus 1GB (reserve for VM overhead)
458463
memoryLimit := opts.Resources.Memory
459464
if memoryLimit == "" {
460465
memGB, err := getSystemMemoryGB()
461466
if err == nil && memGB > 0 {
462-
// Calculate 1/4 of system memory (round up)
463-
memAllocated := (memGB + 3) / 4 // Round up division
467+
// Calculate 1/4 of system memory (round up), then subtract 1GB for VM overhead
468+
memAllocated := (memGB+3)/4 - 1
464469
if memAllocated < 1 {
465470
memAllocated = 1 // Minimum 1GB
466471
}

0 commit comments

Comments
 (0)