dstack/src/tests/_internal/server/background/tasks/test_process_fleets.py at 8051e85e7ff6ad2ad6eca973691acd720b93dcfa · dstackai/dstack · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from dstack._internal.core.models.fleets import FleetNodesSpec, FleetStatus
from dstack._internal.core.models.instances import InstanceStatus
from dstack._internal.core.models.runs import RunStatus
from dstack._internal.core.models.users import GlobalRole, ProjectRole
from dstack._internal.server.background.tasks.process_fleets import process_fleets
from dstack._internal.server.models import InstanceModel
from dstack._internal.server.services.projects import add_project_member
from dstack._internal.server.testing.common import (
    create_fleet,
    create_instance,
    create_project,
    create_repo,
    create_run,
    create_user,
    get_fleet_spec,
)


class TestProcessEmptyFleets:
    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_deletes_empty_autocreated_fleet(self, test_db, session: AsyncSession):
        project = await create_project(session)
        spec = get_fleet_spec()
        spec.autocreated = True
        fleet = await create_fleet(
            session=session,
            project=project,
            spec=spec,
        )
        await process_fleets()
        await session.refresh(fleet)
        assert fleet.deleted

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_deletes_terminating_user_fleet(self, test_db, session: AsyncSession):
        project = await create_project(session)
        spec = get_fleet_spec()
        spec.autocreated = False
        fleet = await create_fleet(
            session=session,
            project=project,
            status=FleetStatus.TERMINATING,
        )
        await process_fleets()
        await session.refresh(fleet)
        assert fleet.deleted

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_does_not_delete_fleet_with_active_run(self, test_db, session: AsyncSession):
        project = await create_project(session)
        fleet = await create_fleet(
            session=session,
            project=project,
        )
        user = await create_user(session=session, global_role=GlobalRole.USER)
        await add_project_member(
            session=session, project=project, user=user, project_role=ProjectRole.USER
        )
        repo = await create_repo(
            session=session,
            project_id=project.id,
        )
        run = await create_run(
            session=session,
            project=project,
            repo=repo,
            user=user,
            status=RunStatus.RUNNING,
        )
        fleet.runs.append(run)
        await session.commit()
        await process_fleets()
        await session.refresh(fleet)
        assert not fleet.deleted

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_does_not_delete_fleet_with_instance(self, test_db, session: AsyncSession):
        project = await create_project(session)
        fleet = await create_fleet(
            session=session,
            project=project,
        )
        user = await create_user(session=session, global_role=GlobalRole.USER)
        await add_project_member(
            session=session, project=project, user=user, project_role=ProjectRole.USER
        )
        instance = await create_instance(
            session=session,
            project=project,
            fleet=fleet,
            status=InstanceStatus.IDLE,
        )
        fleet.instances.append(instance)
        await session.commit()
        await process_fleets()
        await session.refresh(fleet)
        assert not fleet.deleted

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_consolidation_creates_missing_instances(self, test_db, session: AsyncSession):
        project = await create_project(session)
        spec = get_fleet_spec()
        spec.configuration.nodes = FleetNodesSpec(min=2, target=2, max=2)
        fleet = await create_fleet(
            session=session,
            project=project,
            spec=spec,
        )
        await create_instance(
            session=session,
            project=project,
            fleet=fleet,
            status=InstanceStatus.IDLE,
            instance_num=1,
        )
        await process_fleets()
        instances = (await session.execute(select(InstanceModel))).scalars().all()
        assert len(instances) == 2
        assert {i.instance_num for i in instances} == {0, 1}  # uses 0 for next instance num

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_consolidation_terminates_redundant_instances(
        self, test_db, session: AsyncSession
    ):
        project = await create_project(session)
        spec = get_fleet_spec()
        spec.configuration.nodes = FleetNodesSpec(min=1, target=1, max=1)
        fleet = await create_fleet(
            session=session,
            project=project,
            spec=spec,
        )
        instance1 = await create_instance(
            session=session,
            project=project,
            fleet=fleet,
            status=InstanceStatus.BUSY,
            instance_num=0,
        )
        instance2 = await create_instance(
            session=session,
            project=project,
            fleet=fleet,
            status=InstanceStatus.IDLE,
            instance_num=1,
        )
        await process_fleets()
        await session.refresh(instance1)
        await session.refresh(instance2)
        assert instance1.status == InstanceStatus.BUSY
        assert instance2.status == InstanceStatus.TERMINATING