Repository: skypilot-org/skypilot
Pull request #3765 (2024-07-19 16:48:00)

Supporting encrypted disks for AWS
<!-- Describe the changes in this PR -->

Fixes https://github.com/skypilot-org/skypilot/issues/2483

Users can now specify a disk_encrypted option under the resources stanza to enable disk encryption at rest. This is currently only setup for AWS but can easily be extended for other providers given their capability.

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [x] Relevant individual smoke tests: `pytest tests/test_config.py` and `tests/unit_tests/test_resources.py`
- [ ] Backward compatibility tests: `conda deactivate; bash -i tests/backward_compatibility_tests.sh`

Manually updated my `~/.sky/config` with 

```
aws:
  disk_encrypted: true
```

Then launched 

```
sky launch -c minimal examples/minimal.yaml
```

and confirmed that the attached EBS disks were encrypted.


Patch:
diff --git a/.gitignore b/.gitignore
index efa74dd74..31b37d1ea 100644
--- a/.gitignore
+++ b/.gitignore
@@ -12,3 +12,6 @@ sky/clouds/service_catalog/data_fetchers/*.csv
 .vscode/
 .idea/
 .env
+
+# For editor files
+*.swp
diff --git a/README.md b/README.md
index 2d2e471dc..b803a9989 100644
--- a/README.md
+++ b/README.md
@@ -58,7 +58,7 @@ SkyPilot is a framework for running LLMs, AI, and batch jobs on any cloud, offer
 SkyPilot **abstracts away cloud infra burdens**:
 - Launch jobs & clusters on any cloud
 - Easy scale-out: queue and run many jobs, automatically managed
-- Easy access to object stores (S3, GCS, R2)
+- Easy access to object stores (S3, GCS, Azure, R2, IBM)
 
 SkyPilot **maximizes GPU availability for your jobs**:
 * Provision in all zones/regions/clouds you have access to ([the _Sky_](https://arxiv.org/abs/2205.07147)), with automatic failover
diff --git a/docs/source/docs/index.rst b/docs/source/docs/index.rst
index 5a648dbcd..b4bd66fba 100644
--- a/docs/source/docs/index.rst
+++ b/docs/source/docs/index.rst
@@ -33,7 +33,7 @@ SkyPilot **abstracts away cloud infra burdens**:
 
 - Launch jobs & clusters on any cloud
 - Easy scale-out: queue and run many jobs, automatically managed
-- Easy access to object stores (S3, GCS, R2)
+- Easy access to object stores (S3, GCS, Azure, R2, IBM)
 
 SkyPilot **maximizes GPU availability for your jobs**:
 
diff --git a/docs/source/reference/config.rst b/docs/source/reference/config.rst
index 7f24c5906..cb06a28cd 100644
--- a/docs/source/reference/config.rst
+++ b/docs/source/reference/config.rst
@@ -368,6 +368,18 @@ Available fields and semantics:
     # Default: 'LOCAL_CREDENTIALS'.
     remote_identity: LOCAL_CREDENTIALS
 
+  # Advanced Azure configurations (optional).
+  # Apply to all new instances but not existing ones.
+  azure:
+    # Specify an existing Azure storage account for SkyPilot-managed containers.
+    # If not set, SkyPilot will use its default naming convention to create and
+    # use the storage account unless container endpoint URI is used as source.
+    # Note: SkyPilot cannot create new storage accounts with custom names; it
+    # can only use existing ones or create accounts with its default naming
+    # scheme.
+    # Reference: https://learn.microsoft.com/en-us/azure/storage/common/storage-account-overview
+    storage_account: user-storage-account-name
+
   # Advanced Kubernetes configurations (optional).
   kubernetes:
     # The networking mode for accessing SSH jump pod (optional).
diff --git a/docs/source/reference/storage.rst b/docs/source/reference/storage.rst
index 20d7ca468..3c54680e7 100644
--- a/docs/source/reference/storage.rst
+++ b/docs/source/reference/storage.rst
@@ -28,7 +28,7 @@ Object storages are specified using the :code:`file_mounts` field in a SkyPilot
           # Mount an existing S3 bucket
           file_mounts:
             /my_data:
-              source: s3://my-bucket/  # or gs://, r2://, cos://<region>/<bucket>
+              source: s3://my-bucket/  # or gs://, https://<azure_storage_account>.blob.core.windows.net/<container>, r2://, cos://<region>/<bucket>
               mode: MOUNT  # Optional: either MOUNT or COPY. Defaults to MOUNT.
 
         This will `mount <storage-mounting-modes_>`__ the contents of the bucket at ``s3://my-bucket/`` to the remote VM at ``/my_data``.
@@ -45,7 +45,7 @@ Object storages are specified using the :code:`file_mounts` field in a SkyPilot
           file_mounts:
             /my_data:
               name: my-sky-bucket
-              store: gcs  # Optional: either of s3, gcs, r2, ibm
+              store: gcs  # Optional: either of s3, gcs, azure, r2, ibm
 
         SkyPilot will create an empty GCS bucket called ``my-sky-bucket`` and mount it at ``/my_data``.
         This bucket can be used to write checkpoints, logs or other outputs directly to the cloud.
@@ -68,7 +68,7 @@ Object storages are specified using the :code:`file_mounts` field in a SkyPilot
             /my_data:
               name: my-sky-bucket
               source: ~/dataset  # Optional: path to local data to upload to the bucket
-              store: s3  # Optional: either of s3, gcs, r2, ibm
+              store: s3  # Optional: either of s3, gcs, azure, r2, ibm
               mode: MOUNT  # Optional: either MOUNT or COPY. Defaults to MOUNT.
 
         SkyPilot will create a S3 bucket called ``my-sky-bucket`` and upload the
@@ -281,14 +281,21 @@ Storage YAML reference
 
         source: str
           The source attribute specifies the path that must be made available
-          in the storage object. It can either be a local path or a list of local
-          paths or it can be a remote path (s3://, gs://, r2://, cos://<region_name>).
+          in the storage object. It can either be:
+          - A local path
+          - A list of local paths
+          - A remote path using one of the following formats:
+            - s3://<bucket_name>
+            - gs://<bucket_name>
+            - https://<azure_storage_account>.blob.core.windows.net/<container_name>
+            - r2://<bucket_name>
+            - cos://<region_name>/<bucket_name>
 
           If the source is local, data is uploaded to the cloud to an appropriate
-          bucket (s3, gcs, r2, or ibm). If source is bucket URI,
+          bucket (s3, gcs, azure, r2, or ibm). If source is bucket URI,
           the data is copied or mounted directly (see mode flag below).
 
-        store: str; either of 's3', 'gcs', 'r2', 'ibm'
+        store: str; either of 's3', 'gcs', 'azure', 'r2', 'ibm'
           If you wish to force sky.Storage to be backed by a specific cloud object
           storage, you can specify it here. If not specified, SkyPilot chooses the
           appropriate object storage based on the source path and task's cloud provider.
diff --git a/docs/source/reference/yaml-spec.rst b/docs/source/reference/yaml-spec.rst
index 35e56726a..0354d3d03 100644
--- a/docs/source/reference/yaml-spec.rst
+++ b/docs/source/reference/yaml-spec.rst
@@ -300,8 +300,8 @@ Available fields:
       # Mounts the bucket at /datasets-storage on every node of the cluster.
       /datasets-storage:
         name: sky-dataset  # Name of storage, optional when source is bucket URI
-        source: /local/path/datasets  # Source path, can be local or s3/gcs URL. Optional, do not specify to create an empty bucket.
-        store: s3  # Could be either 's3', 'gcs' or 'r2'; default: None. Optional.
+        source: /local/path/datasets  # Source path, can be local or bucket URI. Optional, do not specify to create an empty bucket.
+        store: s3  # Could be either 's3', 'gcs', 'azure', 'r2', or 'ibm'; default: None. Optional.
         persistent: True  # Defaults to True; can be set to false to delete bucket after cluster is downed. Optional.
         mode: MOUNT  # Either MOUNT or COPY. Defaults to MOUNT. Optional.
 
diff --git a/llm/vicuna-llama-2/scripts/hardcoded_questions.py b/llm/vicuna-llama-2/scripts/hardcoded_questions.py
index 9ed7490ca..bfb8494b0 100644
--- a/llm/vicuna-llama-2/scripts/hardcoded_questions.py
+++ b/llm/vicuna-llama-2/scripts/hardcoded_questions.py
@@ -190,7 +190,7 @@ def identity_questions():
         SkyPilot abstracts away cloud infra burdens:
         * Launch jobs & clusters on any cloud
         * Easy scale-out: queue and run many jobs, automatically managed
-        * Easy access to object stores (S3, GCS, R2)
+        * Easy access to object stores (S3, GCS, Azure, R2, IBM)
 
         SkyPilot maximizes GPU availability for your jobs:
         * Provision in all zones/regions/clouds you have access to (the Sky), with automatic failover
diff --git a/sky/clouds/aws.py b/sky/clouds/aws.py
index cb09a3c6b..fc001ea75 100644
--- a/sky/clouds/aws.py
+++ b/sky/clouds/aws.py
@@ -399,6 +399,8 @@ class AWS(clouds.Cloud):
         image_id = self._get_image_id(image_id_to_use, region_name,
                                       r.instance_type)
 
+        disk_encrypted = skypilot_config.get_nested(('aws', 'disk_encrypted'),
+                                                    False)
         user_security_group_config = skypilot_config.get_nested(
             ('aws', 'security_group_name'), None)
         user_security_group = None
@@ -429,6 +431,7 @@ class AWS(clouds.Cloud):
         return {
             'instance_type': r.instance_type,
             'custom_resources': custom_resources,
+            'disk_encrypted': disk_encrypted,
             'use_spot': r.use_spot,
             'region': region_name,
             'zones': ','.join(zone_names),
diff --git a/sky/provision/kubernetes/network_utils.py b/sky/provision/kubernetes/network_utils.py
index 844f84a04..5486214b7 100644
--- a/sky/provision/kubernetes/network_utils.py
+++ b/sky/provision/kubernetes/network_utils.py
@@ -75,6 +75,10 @@ def fill_loadbalancer_template(namespace: str, service_name: str,
 
     with open(template_path, 'r', encoding='utf-8') as fin:
         template = fin.read()
+    annotations = skypilot_config.get_nested(
+        ('kubernetes', 'custom_metadata', 'annotations'), {})
+    labels = skypilot_config.get_nested(
+        ('kubernetes', 'custom_metadata', 'labels'), {})
     j2_template = jinja2.Template(template)
     cont = j2_template.render(
         namespace=namespace,
@@ -82,6 +86,8 @@ def fill_loadbalancer_template(namespace: str, service_name: str,
         ports=ports,
         selector_key=selector_key,
         selector_value=selector_value,
+        annotations=annotations,
+        labels=labels,
     )
     content = yaml.safe_load(cont)
     return content
@@ -98,6 +104,10 @@ def fill_ingress_template(namespace: str, service_details: List[Tuple[str, int,
             f'Template "{_INGRESS_TEMPLATE_NAME}" does not exist.')
     with open(template_path, 'r', encoding='utf-8') as fin:
         template = fin.read()
+    annotations = skypilot_config.get_nested(
+        ('kubernetes', 'custom_metadata', 'annotations'), {})
+    labels = skypilot_config.get_nested(
+        ('kubernetes', 'custom_metadata', 'labels'), {})
     j2_template = jinja2.Template(template)
     cont = j2_template.render(
         namespace=namespace,
@@ -109,6 +119,8 @@ def fill_ingress_template(namespace: str, service_details: List[Tuple[str, int,
         ingress_name=ingress_name,
         selector_key=selector_key,
         selector_value=selector_value,
+        annotations=annotations,
+        labels=labels,
     )
     content = yaml.safe_load(cont)
 
diff --git a/sky/templates/aws-ray.yml.j2 b/sky/templates/aws-ray.yml.j2
index ac84f8a4f..0a6b0bcc0 100644
--- a/sky/templates/aws-ray.yml.j2
+++ b/sky/templates/aws-ray.yml.j2
@@ -75,6 +75,7 @@ available_node_types:
           Ebs:
             VolumeSize: {{disk_size}}
             VolumeType: {{disk_tier}}
+            Encrypted: {{disk_encrypted}}
             {% if custom_disk_perf %}
             Iops: {{disk_iops}}
             Throughput: {{disk_throughput}}
diff --git a/sky/templates/kubernetes-ingress.yml.j2 b/sky/templates/kubernetes-ingress.yml.j2
index 84b3e9d69..d419b6b28 100644
--- a/sky/templates/kubernetes-ingress.yml.j2
+++ b/sky/templates/kubernetes-ingress.yml.j2
@@ -2,9 +2,16 @@ ingress_spec:
   apiVersion: networking.k8s.io/v1
   kind: Ingress
   metadata:
+    labels:
+      {%- for label_key, label_value in labels.items() %}
+      {{ label_key }}: {{ label_value|tojson }}
+      {%- endfor %}
     annotations:
       nginx.ingress.kubernetes.io/use-regex: "true"
       nginx.ingress.kubernetes.io/rewrite-target: /$2
+      {%- for key, value in annotations.items() %}
+      {{ key }}: {{ value|tojson }}
+      {%- endfor %}
     name: {{ ingress_name }}
     namespace: {{ namespace }}
   spec:
diff --git a/sky/templates/kubernetes-loadbalancer.yml.j2 b/sky/templates/kubernetes-loadbalancer.yml.j2
index 08d8b0cc6..7afc35ab3 100644
--- a/sky/templates/kubernetes-loadbalancer.yml.j2
+++ b/sky/templates/kubernetes-loadbalancer.yml.j2
@@ -5,6 +5,13 @@ service_spec:
     name: {{ service_name }}
     labels:
       parent: skypilot
+      {%- for label_key, label_value in labels.items() %}
+      {{ label_key }}: {{ label_value|tojson }}
+      {%- endfor %}
+    annotations:
+      {%- for key, value in annotations.items() %}
+      {{ key }}: {{ value|tojson }}
+      {%- endfor %}
   spec:
     type: LoadBalancer
     selector:
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index a7bfe8f9f..0fa1e8d34 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -706,6 +706,9 @@ def get_config_schema():
             'required': [],
             'additionalProperties': False,
             'properties': {
+                'disk_encrypted': {
+                    'type': 'boolean',
+                },
                 'security_group_name':
                     (_PRORPERTY_NAME_OR_CLUSTER_NAME_TO_PROPERTY),
                 **_LABELS_SCHEMA,


Test patch:
diff --git a/tests/test_config.py b/tests/test_config.py
index c01f06d6f..0cae5f9be 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -10,6 +10,7 @@ from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import kubernetes_enums
 
+DISK_ENCRYPTED = True
 VPC_NAME = 'vpc-12345678'
 PROXY_COMMAND = 'ssh -W %h:%p -i ~/.ssh/id_rsa -o StrictHostKeyChecking=no'
 NODEPORT_MODE_NAME = kubernetes_enums.KubernetesNetworkingMode.NODEPORT.value
@@ -42,6 +43,7 @@ def _create_config_file(config_file_path: pathlib.Path) -> None:
                 vpc_name: {VPC_NAME}
                 use_internal_ips: true
                 ssh_proxy_command: {PROXY_COMMAND}
+                disk_encrypted: {DISK_ENCRYPTED}
 
             gcp:
                 vpc_name: {VPC_NAME}
@@ -215,6 +217,7 @@ def test_config_get_set_nested(monkeypatch, tmp_path) -> None:
     # Check that the config is loaded with the expected values
     assert skypilot_config.loaded()
     assert skypilot_config.get_nested(('aws', 'vpc_name'), None) == VPC_NAME
+    assert skypilot_config.get_nested(('aws', 'disk_encrypted'), None)
     assert skypilot_config.get_nested(('aws', 'use_internal_ips'), None)
     assert skypilot_config.get_nested(('aws', 'ssh_proxy_command'),
                                       None) == PROXY_COMMAND
diff --git a/tests/unit_tests/test_resources.py b/tests/unit_tests/test_resources.py
index 450ca692f..6fb9f1bcd 100644
--- a/tests/unit_tests/test_resources.py
+++ b/tests/unit_tests/test_resources.py
@@ -121,6 +121,7 @@ def test_aws_make_deploy_variables(*mocks) -> None:
         'use_spot': False,
         'region': 'fake-region',
         'image_id': 'fake-image',
+        'disk_encrypted': False,
         'disk_tier': 'gp3',
         'disk_throughput': 218,
         'disk_iops': 3500,