|
4 | 4 | import subprocess |
5 | 5 | from datetime import datetime, timedelta, timezone |
6 | 6 | from types import SimpleNamespace |
| 7 | +from typing import Union, cast |
7 | 8 | from unittest.mock import Mock, patch |
8 | 9 |
|
9 | 10 | import gpuhunt |
|
19 | 20 | _get_bootstrap_command, |
20 | 21 | _get_image_env_and_cwd_command, |
21 | 22 | _get_provider_ssh, |
| 23 | + _get_volume_commands, |
22 | 24 | ) |
23 | 25 | from dstack._internal.core.backends.lium.models import LiumConfig, LiumCreds |
24 | 26 | from dstack._internal.core.errors import ( |
|
31 | 33 | from dstack._internal.core.models.common import RegistryAuth |
32 | 34 | from dstack._internal.core.models.instances import InstanceAvailability, InstanceRuntime |
33 | 35 | from dstack._internal.core.models.resources import ResourcesSpec |
34 | | -from dstack._internal.core.models.runs import JobProvisioningData, Requirements |
| 36 | +from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements |
35 | 37 | from dstack._internal.core.models.unix import UnixUser |
| 38 | +from dstack._internal.core.models.volumes import LiumVolumeConfiguration, VolumeMountPoint |
| 39 | +from dstack._internal.server.testing.common import get_volume |
36 | 40 | from dstack._internal.utils.crypto import generate_rsa_key_pair_bytes |
37 | 41 |
|
38 | 42 | COMPUTE_MODULE = "dstack._internal.core.backends.lium.compute" |
@@ -676,3 +680,94 @@ def test_oversized_environment_fails_without_starting_partial_runner(self, tmp_p |
676 | 680 |
|
677 | 681 | assert result.returncode != 0 |
678 | 682 | assert result.stdout == b"" |
| 683 | + |
| 684 | + |
| 685 | +class TestVolumeLifecycle: |
| 686 | + def test_create_unbounded_volume(self, compute): |
| 687 | + compute.api_client.create_volume.return_value = {"id": "volume-id"} |
| 688 | + with patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-name"): |
| 689 | + data = compute.create_volume(Mock()) |
| 690 | + |
| 691 | + compute.api_client.create_volume.assert_called_once_with("unique-name") |
| 692 | + assert data.volume_id == "volume-id" |
| 693 | + assert data.size_gb is None |
| 694 | + assert data.price is None |
| 695 | + assert not data.attachable |
| 696 | + assert not data.detachable |
| 697 | + |
| 698 | + def test_recovers_lost_create_response_by_unique_name(self, compute): |
| 699 | + compute.api_client.create_volume.side_effect = LiumAPIError("timed out") |
| 700 | + compute.api_client.get_volumes.return_value = [ |
| 701 | + {"id": "other-id", "name": "other"}, |
| 702 | + {"id": "volume-id", "name": "unique-name"}, |
| 703 | + ] |
| 704 | + with patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-name"): |
| 705 | + data = compute.create_volume(Mock()) |
| 706 | + |
| 707 | + assert data.volume_id == "volume-id" |
| 708 | + assert compute.api_client.create_volume.call_count == 1 |
| 709 | + |
| 710 | + def test_register_existing_volume(self, compute): |
| 711 | + volume = SimpleNamespace(configuration=LiumVolumeConfiguration(volume_id="existing")) |
| 712 | + compute.api_client.get_volume.return_value = {"id": "existing"} |
| 713 | + |
| 714 | + data = compute.register_volume(volume) |
| 715 | + |
| 716 | + assert data.volume_id == "existing" |
| 717 | + compute.api_client.create_volume.assert_not_called() |
| 718 | + |
| 719 | + def test_register_missing_volume(self, compute): |
| 720 | + compute.api_client.get_volume.return_value = None |
| 721 | + with pytest.raises(ComputeError, match="not found"): |
| 722 | + compute.register_volume( |
| 723 | + SimpleNamespace(configuration=LiumVolumeConfiguration(volume_id="missing")) |
| 724 | + ) |
| 725 | + |
| 726 | + def test_delete_volume(self, compute): |
| 727 | + compute.delete_volume(SimpleNamespace(volume_id="volume-id")) |
| 728 | + compute.api_client.delete_volume.assert_called_once_with("volume-id") |
| 729 | + |
| 730 | + |
| 731 | +class TestVolumeMount: |
| 732 | + @pytest.mark.parametrize( |
| 733 | + "user", |
| 734 | + [None, UnixUser(uid=10001), UnixUser(username="app")], |
| 735 | + ) |
| 736 | + def test_custom_path(self, user): |
| 737 | + commands = _volume_commands(user=user, path="/checkpoints") |
| 738 | + |
| 739 | + assert commands[0] == "grep -q ' /mnt .* - fuse.s3fs ' /proc/self/mountinfo" |
| 740 | + assert commands[1:] == ["mkdir -p -- /checkpoints", "mount --bind /mnt /checkpoints"] |
| 741 | + |
| 742 | + def test_native_mount_path_needs_no_bind(self): |
| 743 | + assert len(_volume_commands(path="/mnt")) == 1 |
| 744 | + |
| 745 | + def test_quotes_mount_path(self): |
| 746 | + commands = _volume_commands(path="/checkpoint dir; echo unsafe") |
| 747 | + assert commands[-1] == "mount --bind /mnt '/checkpoint dir; echo unsafe'" |
| 748 | + |
| 749 | + @pytest.mark.parametrize("path", ["/", "/root", "/root/.dstack-lium", "/proc/a", "/mnt/a"]) |
| 750 | + def test_rejects_paths_that_hide_bootstrap_or_nest_source(self, path): |
| 751 | + with pytest.raises(ComputeError, match="cannot be mounted"): |
| 752 | + _volume_commands(path=path) |
| 753 | + |
| 754 | + def test_two_volumes_rejected(self): |
| 755 | + with pytest.raises(ComputeError, match="only one"): |
| 756 | + _get_volume_commands(Mock(), Mock(), [Mock(), Mock()]) |
| 757 | + |
| 758 | + def test_selects_volume_from_alternatives(self): |
| 759 | + assert _volume_commands(name=["elsewhere", "cache"], path="/data")[-1] == ( |
| 760 | + "mount --bind /mnt /data" |
| 761 | + ) |
| 762 | + |
| 763 | + |
| 764 | +def _volume_commands(user=None, path="/checkpoints", name: Union[str, list[str]] = "cache"): |
| 765 | + # A real `Job` needs a full `JobSpec`, while `_get_volume_commands()` only reads the user and |
| 766 | + # the mount points |
| 767 | + job = cast( |
| 768 | + Job, |
| 769 | + SimpleNamespace( |
| 770 | + job_spec=SimpleNamespace(user=user, volumes=[VolumeMountPoint(name=name, path=path)]) |
| 771 | + ), |
| 772 | + ) |
| 773 | + return _get_volume_commands(Mock(), job, [get_volume(name="cache")]) |
0 commit comments