diff --git a/src/snowflake/cli/_plugins/remote/constants.py b/src/snowflake/cli/_plugins/remote/constants.py index b4e21aabf4..d60cadd5d6 100644 --- a/src/snowflake/cli/_plugins/remote/constants.py +++ b/src/snowflake/cli/_plugins/remote/constants.py @@ -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" diff --git a/src/snowflake/cli/_plugins/remote/container_spec.py b/src/snowflake/cli/_plugins/remote/container_spec.py index 377ac10c49..9c0c87674f 100644 --- a/src/snowflake/cli/_plugins/remote/container_spec.py +++ b/src/snowflake/cli/_plugins/remote/container_spec.py @@ -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, @@ -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 @@ -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) @@ -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 = { @@ -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 diff --git a/src/snowflake/cli/_plugins/remote/manager.py b/src/snowflake/cli/_plugins/remote/manager.py index 0e06f5f3a2..6010fbfbd4 100644 --- a/src/snowflake/cli/_plugins/remote/manager.py +++ b/src/snowflake/cli/_plugins/remote/manager.py @@ -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") diff --git a/tests/remote/test_container_spec.py b/tests/remote/test_container_spec.py index 69752f8c0b..6b39c82ffa 100644 --- a/tests/remote/test_container_spec.py +++ b/tests/remote/test_container_spec.py @@ -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 ( @@ -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