Skip to content

Commit 2cd640f

Browse files
committed
enable remote ssh connection
1 parent d8a093f commit 2cd640f

File tree

9 files changed

+1230
-20
lines changed

9 files changed

+1230
-20
lines changed

src/snowflake/cli/_plugins/remote/commands.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,16 @@ def start(
6969
"--image-tag",
7070
help="Custom image tag to use for the remote development environment",
7171
),
72+
ssh: bool = typer.Option(
73+
False,
74+
"--ssh",
75+
help="Set up SSH configuration for connecting to the remote environment. This is a blocking command that keeps SSH connections alive.",
76+
),
77+
no_ssh_key: bool = typer.Option(
78+
False,
79+
"--no-ssh-key",
80+
help="When used with --ssh, skip SSH key generation and use token-only authentication (less secure)",
81+
),
7282
**options,
7383
) -> None:
7484
"""
@@ -83,9 +93,16 @@ def start(
8393
- Resume existing service: snow remote start myproject
8494
- Create new service: snow remote start --compute-pool my_pool
8595
- Create named service: snow remote start myproject --compute-pool my_pool
96+
- Start with SSH setup: snow remote start myproject --ssh
97+
- Start with SSH (no key): snow remote start myproject --ssh --no-ssh-key
8698
8799
The --compute-pool parameter is only required when creating a new service. For resuming
88100
existing services, the compute pool is not needed.
101+
102+
SSH Options:
103+
- Use --ssh to set up SSH configuration for secure terminal access
104+
- Use --no-ssh-key with --ssh for token-only authentication (less secure)
105+
- SSH setup is a blocking command that continuously refreshes authentication tokens
89106
"""
90107
try:
91108
manager = RemoteManager()
@@ -96,6 +113,8 @@ def start(
96113
external_access=eai_name,
97114
stage=stage,
98115
image_tag=image_tag,
116+
generate_ssh_key=ssh
117+
and not no_ssh_key, # Only generate SSH key if --ssh and not --no-ssh-key
99118
)
100119

101120
# Display appropriate success message based on what happened
@@ -132,6 +151,10 @@ def start(
132151
if image_tag:
133152
log.debug("Using custom image tag: %s", image_tag)
134153

154+
# Handle SSH setup if requested - this is a blocking operation
155+
if ssh:
156+
manager.setup_ssh_connection(service_name)
157+
135158
except ValueError as e:
136159
cc.warning(f"Error: {e}")
137160
raise typer.Exit(code=1)

src/snowflake/cli/_plugins/remote/constants.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ class ComputeResources:
4747
ENABLE_REMOTE_DEV_ENV_VAR = "IS_REMOTE_DEV"
4848
MEMORY_VOLUME_NAME = "dshm"
4949
USER_WORKSPACE_VOLUME_NAME = "user-workspace"
50-
USER_WORKSPACE_VOLUME_MOUNT_PATH = "/root/workspace"
50+
USER_WORKSPACE_VOLUME_MOUNT_PATH = "/root/user-default"
5151
USER_VSCODE_DATA_VOLUME_NAME = "user-vscode-data"
5252
USER_VSCODE_DATA_VOLUME_MOUNT_PATH = "/root/.vscode-server"
5353

@@ -78,7 +78,7 @@ class ComputeResources:
7878
DEFAULT_IMAGE_REPO = "/snowflake/images/snowflake_images"
7979
DEFAULT_IMAGE_CPU = "st_plat/runtime/x86/runtime_image/snowbooks"
8080
DEFAULT_IMAGE_GPU = "st_plat/runtime/x86/generic_gpu/runtime_image/snowbooks"
81-
DEFAULT_IMAGE_TAG = "1.7.1"
81+
DEFAULT_IMAGE_TAG = "1.7.2"
8282

8383
# Percent of container memory to allocate for /dev/shm volume
8484
MEMORY_VOLUME_SIZE = 0.3
@@ -92,6 +92,13 @@ class ComputeResources:
9292
WEBSOCKET_SSH_ENDPOINT_NAME = "websocket-ssh"
9393
RAY_DASHBOARD_ENDPOINT_NAME = "ray-dashboard"
9494

95+
# SSH Configuration
96+
DEFAULT_SSH_REFRESH_INTERVAL = 300 # 5 minutes
97+
SSH_RETRY_INTERVAL = 30 # 30 seconds
98+
SSH_COUNTDOWN_INTERVAL = 30 # Show countdown every 30 seconds
99+
SSH_KEY_DIR = "~/.ssh/snowflake-remote"
100+
SSH_CONFIG_PATH = "~/.ssh/config"
101+
95102
# ML runtime health check settings
96103
ML_RUNTIME_HEALTH_CHECK_PORT = "5001"
97104
ENABLE_HEALTH_CHECKS = "false"

0 commit comments

Comments
 (0)