diff --git a/tests/test_smoke.py b/tests/test_smoke.py index f37467417fa..79920147adb 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -37,7 +37,7 @@ import tempfile import textwrap import time -from typing import Dict, List, NamedTuple, Optional, Tuple +from typing import Dict, List, NamedTuple, Optional, TextIO, Tuple import urllib.parse import uuid @@ -1333,34 +1333,68 @@ def test_using_file_mounts_with_env_vars(generic_cloud: str): # ---------- storage ---------- + + +def _storage_mounts_commands_generator(f: TextIO, cluster_name: str, + storage_name: str, ls_hello_command: str, + cloud: str, only_mount: bool): + template_str = pathlib.Path( + 'tests/test_yamls/test_storage_mounting.yaml.j2').read_text() + template = jinja2.Template(template_str) + content = template.render( + storage_name=storage_name, + cloud=cloud, + only_mount=only_mount, + ) + f.write(content) + f.flush() + file_path = f.name + test_commands = [ + *STORAGE_SETUP_COMMANDS, + f'sky launch -y -c {cluster_name} --cloud {cloud} {file_path}', + f'sky logs {cluster_name} 1 --status', # Ensure job succeeded. + ls_hello_command, + f'sky stop -y {cluster_name}', + f'sky start -y {cluster_name}', + # Check if hello.txt from mounting bucket exists after restart in + # the mounted directory + f'sky exec {cluster_name} -- "set -ex; ls /mount_private_mount/hello.txt"', + ] + clean_command = f'sky down -y {cluster_name}; sky storage delete -y {storage_name}' + return test_commands, clean_command + + @pytest.mark.aws def test_aws_storage_mounts_with_stop(): name = _get_cluster_name() cloud = 'aws' storage_name = f'sky-test-{int(time.time())}' - template_str = pathlib.Path( - 'tests/test_yamls/test_storage_mounting.yaml.j2').read_text() - template = jinja2.Template(template_str) - content = template.render(storage_name=storage_name, cloud=cloud) + ls_hello_command = f'aws s3 ls {storage_name}/hello.txt' with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f: - f.write(content) - f.flush() - file_path = f.name - test_commands = [ - *STORAGE_SETUP_COMMANDS, - f'sky launch -y -c {name} --cloud {cloud} {file_path}', - f'sky logs {name} 1 --status', # Ensure job succeeded. - f'aws s3 ls {storage_name}/hello.txt', - f'sky stop -y {name}', - f'sky start -y {name}', - # Check if hello.txt from mounting bucket exists after restart in - # the mounted directory - f'sky exec {name} -- "set -ex; ls /mount_private_mount/hello.txt"' - ] + test_commands, clean_command = _storage_mounts_commands_generator( + f, name, storage_name, ls_hello_command, cloud, False) test = Test( 'aws_storage_mounts', test_commands, - f'sky down -y {name}; sky storage delete -y {storage_name}', + clean_command, + timeout=20 * 60, # 20 mins + ) + run_one_test(test) + + +@pytest.mark.aws +def test_aws_storage_mounts_with_stop_only_mount(): + name = _get_cluster_name() + cloud = 'aws' + storage_name = f'sky-test-{int(time.time())}' + ls_hello_command = f'aws s3 ls {storage_name}/hello.txt' + with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f: + test_commands, clean_command = _storage_mounts_commands_generator( + f, name, storage_name, ls_hello_command, cloud, True) + test = Test( + 'aws_storage_mounts_only_mount', + test_commands, + clean_command, timeout=20 * 60, # 20 mins ) run_one_test(test) @@ -1371,29 +1405,14 @@ def test_gcp_storage_mounts_with_stop(): name = _get_cluster_name() cloud = 'gcp' storage_name = f'sky-test-{int(time.time())}' - template_str = pathlib.Path( - 'tests/test_yamls/test_storage_mounting.yaml.j2').read_text() - template = jinja2.Template(template_str) - content = template.render(storage_name=storage_name, cloud=cloud) + ls_hello_command = f'gsutil ls gs://{storage_name}/hello.txt' with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f: - f.write(content) - f.flush() - file_path = f.name - test_commands = [ - *STORAGE_SETUP_COMMANDS, - f'sky launch -y -c {name} --cloud {cloud} {file_path}', - f'sky logs {name} 1 --status', # Ensure job succeeded. - f'gsutil ls gs://{storage_name}/hello.txt', - f'sky stop -y {name}', - f'sky start -y {name}', - # Check if hello.txt from mounting bucket exists after restart in - # the mounted directory - f'sky exec {name} -- "set -ex; ls /mount_private_mount/hello.txt"' - ] + test_commands, clean_command = _storage_mounts_commands_generator( + f, name, storage_name, ls_hello_command, cloud, False) test = Test( 'gcp_storage_mounts', test_commands, - f'sky down -y {name}; sky storage delete -y {storage_name}', + clean_command, timeout=20 * 60, # 20 mins ) run_one_test(test) @@ -1409,31 +1428,19 @@ def test_azure_storage_mounts_with_stop(): get_default_storage_account_name(default_region)) storage_account_key = data_utils.get_az_storage_account_key( storage_account_name) - template_str = pathlib.Path( - 'tests/test_yamls/test_storage_mounting.yaml.j2').read_text() - template = jinja2.Template(template_str) - content = template.render(storage_name=storage_name, cloud=cloud) + # if the file does not exist, az storage blob list returns '[]' + ls_hello_command = (f'output=$(az storage blob list -c {storage_name} ' + f'--account-name {storage_account_name} ' + f'--account-key {storage_account_key} ' + f'--prefix hello.txt) ' + f'[ "$output" = "[]" ] && exit 1 || exit 0') with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f: - f.write(content) - f.flush() - file_path = f.name - test_commands = [ - *STORAGE_SETUP_COMMANDS, - f'sky launch -y -c {name} --cloud {cloud} {file_path}', - f'sky logs {name} 1 --status', # Ensure job succeeded. - f'output=$(az storage blob list -c {storage_name} --account-name {storage_account_name} --account-key {storage_account_key} --prefix hello.txt)' - # if the file does not exist, az storage blob list returns '[]' - f'[ "$output" = "[]" ] && exit 1;' - f'sky stop -y {name}', - f'sky start -y {name}', - # Check if hello.txt from mounting bucket exists after restart in - # the mounted directory - f'sky exec {name} -- "set -ex; ls /mount_private_mount/hello.txt"' - ] + test_commands, clean_command = _storage_mounts_commands_generator( + f, name, storage_name, ls_hello_command, cloud, False) test = Test( 'azure_storage_mounts', test_commands, - f'sky down -y {name}; sky storage delete -y {storage_name}', + clean_command, timeout=20 * 60, # 20 mins ) run_one_test(test) @@ -1446,25 +1453,15 @@ def test_kubernetes_storage_mounts(): # built for x86_64 only. name = _get_cluster_name() storage_name = f'sky-test-{int(time.time())}' - template_str = pathlib.Path( - 'tests/test_yamls/test_storage_mounting.yaml.j2').read_text() - template = jinja2.Template(template_str) - content = template.render(storage_name=storage_name) + ls_hello_command = (f'aws s3 ls {storage_name}/hello.txt || ' + f'gsutil ls gs://{storage_name}/hello.txt') with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f: - f.write(content) - f.flush() - file_path = f.name - test_commands = [ - *STORAGE_SETUP_COMMANDS, - f'sky launch -y -c {name} --cloud kubernetes {file_path}', - f'sky logs {name} 1 --status', # Ensure job succeeded. - f'aws s3 ls {storage_name}/hello.txt || ' - f'gsutil ls gs://{storage_name}/hello.txt', - ] + test_commands, clean_command = _storage_mounts_commands_generator( + f, name, storage_name, ls_hello_command, 'kubernetes', False) test = Test( 'kubernetes_storage_mounts', test_commands, - f'sky down -y {name}; sky storage delete -y {storage_name}', + clean_command, timeout=20 * 60, # 20 mins ) run_one_test(test) diff --git a/tests/test_yamls/test_storage_mounting.yaml.j2 b/tests/test_yamls/test_storage_mounting.yaml.j2 index 4241c63409e..651fb190012 100644 --- a/tests/test_yamls/test_storage_mounting.yaml.j2 +++ b/tests/test_yamls/test_storage_mounting.yaml.j2 @@ -20,13 +20,13 @@ file_mounts: /mount_private_copy: name: {{storage_name}} source: ~/tmp-workdir - mode: COPY + mode: {% if only_mount | default(false) %}MOUNT{% else %}COPY{% endif %} # Mounting private buckets in COPY mode with a list of files as source /mount_private_copy_lof: name: {{storage_name}} source: ['~/tmp-workdir/tmp file', '~/tmp-workdir/tmp file2'] - mode: COPY + mode: {% if only_mount | default(false) %}MOUNT{% else %}COPY{% endif %} {% if include_private_mount | default(True) %} # Mounting private buckets in MOUNT mode @@ -38,7 +38,7 @@ file_mounts: run: | set -ex - + # Check public bucket contents ls -ltr /mount_public_s3/corpora ls -ltr /mount_public_gcp/tiles @@ -55,12 +55,12 @@ run: | ls -ltr /mount_private_mount/foo ls -ltr /mount_private_mount/tmp\ file {% endif %} - + # Symlinks are not copied to buckets ! ls /mount_private_copy/circle-link {% if include_private_mount | default(True) %} ! ls /mount_private_mount/circle-link - + # Write to private bucket in MOUNT mode should pass echo "hello" > /mount_private_mount/hello.txt {% endif %}