Add default block volume for persistent storage

Author: sfc-gh-zzhu

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"
 

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
 

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