Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/snowflake/cli/_plugins/remote/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,10 @@ class ComputeResources:
USER_VSCODE_DATA_VOLUME_NAME = "user-vscode-data"
USER_VSCODE_DATA_VOLUME_MOUNT_PATH = "/root/.vscode-server/data"

# Block storage constants
BLOCK_STORAGE_VOLUME_NAME = "block-storage"
DEFAULT_BLOCK_STORAGE_SIZE = "10Gi" # Default size for block storage volumes

# Service naming constants
SERVICE_NAME_PREFIX = "SNOW_REMOTE"

Expand Down
29 changes: 26 additions & 3 deletions src/snowflake/cli/_plugins/remote/container_spec.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@
import yaml
from snowflake import snowpark
from snowflake.cli._plugins.remote.constants import (
BLOCK_STORAGE_VOLUME_NAME,
DEFAULT_BLOCK_STORAGE_SIZE,
DEFAULT_CONTAINER_NAME,
DEFAULT_IMAGE_CPU,
DEFAULT_IMAGE_GPU,
Expand Down Expand Up @@ -97,7 +99,8 @@ def generate_service_spec(
Args:
session: Snowflake session
compute_pool: Compute pool for service execution
stage: Optional internal Snowflake stage to mount (e.g., @my_stage)
stage: Optional internal Snowflake stage to mount (e.g., @my_stage).
If not provided, a block storage volume will be created for persistent storage.
image: Optional custom image (can be full path like 'repo/image:tag' or just tag like '1.7.1')
ssh_public_key: Optional SSH public key to inject for secure authentication

Expand Down Expand Up @@ -168,7 +171,7 @@ def generate_service_spec(
}
)

# Mount user stage as volume if provided
# Mount user stage as volume if provided, otherwise use block storage
if stage:
# Mount user workspace volume
user_workspace_mount = PurePath(USER_WORKSPACE_VOLUME_MOUNT_PATH)
Expand Down Expand Up @@ -207,6 +210,25 @@ def generate_service_spec(
"source": vscode_data_source,
}
)
else:
# Use block storage for persistent storage when no stage is provided
# Mount user workspace volume using block storage
user_workspace_mount = PurePath(USER_WORKSPACE_VOLUME_MOUNT_PATH)
volume_mounts.append(
{
"name": BLOCK_STORAGE_VOLUME_NAME,
"mountPath": user_workspace_mount.as_posix(),
}
)

# Define block storage volume
volumes.append(
{
"name": BLOCK_STORAGE_VOLUME_NAME,
"source": "block",
"size": DEFAULT_BLOCK_STORAGE_SIZE,
}
)

# Setup environment variables
env_vars = {
Expand Down Expand Up @@ -294,7 +316,8 @@ def generate_service_spec_yaml(
Args:
session: Snowflake session
compute_pool: Compute pool for service execution
stage: Optional internal Snowflake stage to mount (e.g., @my_stage)
stage: Optional internal Snowflake stage to mount (e.g., @my_stage).
If not provided, a block storage volume will be created for persistent storage.
image: Optional custom image (can be full path like 'repo/image:tag' or just tag like '1.7.1')
ssh_public_key: Optional SSH public key to inject for secure authentication

Expand Down
2 changes: 1 addition & 1 deletion src/snowflake/cli/_plugins/remote/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -1071,4 +1071,4 @@ def delete(self, name_input: str) -> SnowflakeCursor:
This permanently deletes the service and all associated data.
"""
service_name = self._resolve_service_name(name_input)
return self.execute_query(f"DROP SERVICE {service_name}")
return self.execute_query(f"DROP SERVICE {service_name} FORCE")
119 changes: 119 additions & 0 deletions tests/remote/test_container_spec.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,12 @@
import pytest
from snowflake import snowpark
from snowflake.cli._plugins.remote.constants import (
BLOCK_STORAGE_VOLUME_NAME,
DEFAULT_BLOCK_STORAGE_SIZE,
DEFAULT_IMAGE_TAG,
RAY_DASHBOARD_ENDPOINT_NAME,
SERVER_UI_ENDPOINT_NAME,
USER_WORKSPACE_VOLUME_MOUNT_PATH,
WEBSOCKET_SSH_ENDPOINT_NAME,
)
from snowflake.cli._plugins.remote.container_spec import (
Expand Down Expand Up @@ -418,3 +421,119 @@ def test_generate_spec_with_gpu_resources(self, mock_session):
assert resources["limits"]["cpu"] == "8000m"
assert resources["requests"]["memory"] == "32Gi"
assert resources["limits"]["memory"] == "32Gi"

def test_generate_spec_without_stage_uses_block_storage(self, mock_session):
"""Test that block storage is created when no stage is provided."""
spec = generate_service_spec(session=mock_session, compute_pool="test_pool")

# Check volumes
volumes = spec["spec"]["volumes"]
block_volume = next(
(v for v in volumes if v["name"] == BLOCK_STORAGE_VOLUME_NAME), None
)

# Block storage volume should exist
assert block_volume is not None
assert block_volume["source"] == "block"
assert block_volume["size"] == DEFAULT_BLOCK_STORAGE_SIZE

# Check volume mounts
container = spec["spec"]["containers"][0]
volume_mounts = container["volumeMounts"]
block_mount = next(
(vm for vm in volume_mounts if vm["name"] == BLOCK_STORAGE_VOLUME_NAME),
None,
)

# Block storage should be mounted at user workspace path
assert block_mount is not None
assert block_mount["mountPath"] == USER_WORKSPACE_VOLUME_MOUNT_PATH

def test_generate_spec_with_stage_no_block_storage(self, mock_session):
"""Test that block storage is NOT created when stage is provided."""
spec = generate_service_spec(
session=mock_session, compute_pool="test_pool", stage="@my_stage"
)

# Check volumes
volumes = spec["spec"]["volumes"]
block_volume = next(
(v for v in volumes if v["name"] == BLOCK_STORAGE_VOLUME_NAME), None
)

# Block storage volume should NOT exist when stage is provided
assert block_volume is None

# Check that stage volumes are present instead
workspace_volume = next(v for v in volumes if v["name"] == "user-workspace")
assert workspace_volume["source"] == "@my_stage/user-default"

def test_generate_spec_block_storage_yaml_format(self, mock_session):
"""Test that block storage is correctly formatted in YAML output."""
from snowflake.cli._plugins.remote.container_spec import (
generate_service_spec_yaml,
)

yaml_output = generate_service_spec_yaml(
session=mock_session, compute_pool="test_pool"
)

# Verify it's valid YAML
import yaml

parsed_spec = yaml.safe_load(yaml_output)

# Check block storage volume exists
volumes = parsed_spec["spec"]["volumes"]
block_volume = next(
(v for v in volumes if v["name"] == BLOCK_STORAGE_VOLUME_NAME), None
)

assert block_volume is not None
assert block_volume["source"] == "block"
assert block_volume["size"] == DEFAULT_BLOCK_STORAGE_SIZE

# Check that YAML string contains block storage references
assert BLOCK_STORAGE_VOLUME_NAME in yaml_output
assert "source: block" in yaml_output
assert DEFAULT_BLOCK_STORAGE_SIZE in yaml_output

def test_generate_spec_block_storage_size(self, mock_session):
"""Test that block storage has correct default size."""
spec = generate_service_spec(session=mock_session, compute_pool="test_pool")

volumes = spec["spec"]["volumes"]
block_volume = next(
v for v in volumes if v["name"] == BLOCK_STORAGE_VOLUME_NAME
)

# Verify the default size is set correctly (10Gi as per constant)
assert block_volume["size"] == "10Gi"

def test_generate_spec_block_storage_only_one_workspace_mount(self, mock_session):
"""Test that only block storage OR stage volumes are present, not both."""
# Test without stage - should have block storage
spec_no_stage = generate_service_spec(
session=mock_session, compute_pool="test_pool"
)

volumes_no_stage = spec_no_stage["spec"]["volumes"]
volume_names_no_stage = [v["name"] for v in volumes_no_stage]

# Should have block storage, but not stage-based volumes
assert BLOCK_STORAGE_VOLUME_NAME in volume_names_no_stage
assert "user-workspace" not in volume_names_no_stage
assert "user-vscode-data" not in volume_names_no_stage

# Test with stage - should have stage volumes
spec_with_stage = generate_service_spec(
session=mock_session, compute_pool="test_pool", stage="@my_stage"
)

volumes_with_stage = spec_with_stage["spec"]["volumes"]
volume_names_with_stage = [v["name"] for v in volumes_with_stage]

# Should have stage volumes, but not block storage
assert "user-workspace" in volume_names_with_stage
assert "user-vscode-data" in volume_names_with_stage
assert BLOCK_STORAGE_VOLUME_NAME not in volume_names_with_stage
Loading