dstack 0.19.1__py3-none-any.whl → 0.19.3__py3-none-any.whl

tests/_internal/server/background/tasks/test_process_instances.py

@@ -8,7 +8,7 @@ import pytest
 from freezegun import freeze_time
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from dstack._internal.core.errors import BackendError, ProvisioningError
+from dstack._internal.core.errors import BackendError, NotYetTerminated, ProvisioningError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import (
     Gpu,
@@ -384,7 +384,9 @@ class TestTerminate:
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    @pytest.mark.parametrize("error", [BackendError("err"), RuntimeError("err")])
+    @pytest.mark.parametrize(
+        "error", [BackendError("err"), RuntimeError("err"), NotYetTerminated("")]
+    )
     async def test_terminate_retry(self, test_db, session: AsyncSession, error: Exception):
         project = await create_project(session=session)
         instance = await create_instance(

tests/_internal/server/background/tasks/test_process_submitted_jobs.py

@@ -536,6 +536,35 @@ class TestProcessSubmittedJobs:
         assert instance.total_blocks == 4
         assert instance.busy_blocks == 2
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_assigns_job_to_specific_fleet(self, test_db, session: AsyncSession):
+        project = await create_project(session)
+        user = await create_user(session)
+        repo = await create_repo(session=session, project_id=project.id)
+        fleet_a = await create_fleet(session=session, project=project, name="a")
+        await create_instance(session=session, project=project, fleet=fleet_a, price=1.0)
+        fleet_b = await create_fleet(session=session, project=project, name="b")
+        await create_instance(session=session, project=project, fleet=fleet_b, price=2.0)
+        fleet_c = await create_fleet(session=session, project=project, name="c")
+        await create_instance(session=session, project=project, fleet=fleet_c, price=3.0)
+        run_spec = get_run_spec(run_name="test-run", repo_id=repo.name)
+        # When more than one fleet is requested, the cheapest one is selected
+        run_spec.configuration.fleets = ["c", "b"]
+        run = await create_run(
+            session=session, project=project, repo=repo, user=user, run_spec=run_spec
+        )
+        job = await create_job(session=session, run=run)
+
+        await process_submitted_jobs()
+
+        await session.refresh(job)
+        res = await session.execute(select(JobModel).options(joinedload(JobModel.instance)))
+        job = res.unique().scalar_one()
+        assert job.status == JobStatus.SUBMITTED
+        assert job.instance is not None
+        assert job.instance.fleet == fleet_b
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_creates_new_instance_in_existing_fleet(self, test_db, session: AsyncSession):

tests/_internal/server/routers/test_backends.py

@@ -1,4 +1,5 @@
 import json
+import sys
 from datetime import datetime, timezone
 from unittest.mock import Mock, patch
 
@@ -25,7 +26,14 @@ from dstack._internal.server.testing.common import (
     get_auth_headers,
     get_volume_provisioning_data,
 )
+from dstack._internal.utils.crypto import generate_rsa_key_pair_bytes
 
+FAKE_NEBIUS_SERVICE_ACCOUNT_CREDS = {
+    "type": "service_account",
+    "service_account_id": "serviceaccount-e00test",
+    "public_key_id": "publickey-e00test",
+    "private_key_content": generate_rsa_key_pair_bytes()[0].decode(),
+}
 FAKE_OCI_CLIENT_CREDS = {
     "type": "client",
     "user": "ocid1.user.oc1..aaaaaaaa",
@@ -62,6 +70,7 @@ class TestListBackendTypes:
             "gcp",
             "kubernetes",
             "lambda",
+            *(["nebius"] if sys.version_info >= (3, 10) else []),
             "oci",
             "runpod",
             "tensordock",
@@ -182,6 +191,113 @@ class TestCreateBackend:
         res = await session.execute(select(BackendModel))
         assert len(res.scalars().all()) == 1
 
+    @pytest.mark.asyncio
+    @pytest.mark.skipif(sys.version_info < (3, 10), reason="Nebius requires Python 3.10")
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    class TestNebius:
+        async def test_creates(self, test_db, session: AsyncSession, client: AsyncClient):
+            user = await create_user(session=session, global_role=GlobalRole.USER)
+            project = await create_project(session=session, owner=user)
+            await add_project_member(
+                session=session, project=project, user=user, project_role=ProjectRole.ADMIN
+            )
+            body = {
+                "type": "nebius",
+                "creds": FAKE_NEBIUS_SERVICE_ACCOUNT_CREDS,
+            }
+            with patch(
+                "dstack._internal.core.backends.nebius.resources.get_region_to_project_id_map"
+            ) as get_region_to_project_id_map:
+                get_region_to_project_id_map.return_value = {"eu-north1": "project-e00test"}
+                response = await client.post(
+                    f"/api/project/{project.name}/backends/create",
+                    headers=get_auth_headers(user.token),
+                    json=body,
+                )
+            assert response.status_code == 200, response.json()
+            res = await session.execute(select(BackendModel))
+            assert len(res.scalars().all()) == 1
+
+        async def test_not_creates_with_invalid_creds(
+            self, test_db, session: AsyncSession, client: AsyncClient
+        ):
+            user = await create_user(session=session, global_role=GlobalRole.USER)
+            project = await create_project(session=session, owner=user)
+            await add_project_member(
+                session=session, project=project, user=user, project_role=ProjectRole.ADMIN
+            )
+            body = {
+                "type": "nebius",
+                "creds": FAKE_NEBIUS_SERVICE_ACCOUNT_CREDS,
+            }
+            with patch(
+                "dstack._internal.core.backends.nebius.resources.get_region_to_project_id_map"
+            ) as get_region_to_project_id_map:
+                get_region_to_project_id_map.side_effect = ValueError()
+                response = await client.post(
+                    f"/api/project/{project.name}/backends/create",
+                    headers=get_auth_headers(user.token),
+                    json=body,
+                )
+            assert response.status_code == 400, response.json()
+            res = await session.execute(select(BackendModel))
+            assert len(res.scalars().all()) == 0
+
+        async def test_creates_with_regions(
+            self, test_db, session: AsyncSession, client: AsyncClient
+        ):
+            user = await create_user(session=session, global_role=GlobalRole.USER)
+            project = await create_project(session=session, owner=user)
+            await add_project_member(
+                session=session, project=project, user=user, project_role=ProjectRole.ADMIN
+            )
+            body = {
+                "type": "nebius",
+                "creds": FAKE_NEBIUS_SERVICE_ACCOUNT_CREDS,
+                "regions": ["eu-north1"],
+            }
+            with patch(
+                "dstack._internal.core.backends.nebius.resources.get_region_to_project_id_map"
+            ) as get_region_to_project_id_map:
+                get_region_to_project_id_map.return_value = {
+                    "eu-north1": "project-e00test",
+                    "eu-west1": "project-e01test",
+                }
+                response = await client.post(
+                    f"/api/project/{project.name}/backends/create",
+                    headers=get_auth_headers(user.token),
+                    json=body,
+                )
+            assert response.status_code == 200, response.json()
+            res = await session.execute(select(BackendModel))
+            assert len(res.scalars().all()) == 1
+
+        async def test_not_creates_with_invalid_regions(
+            self, test_db, session: AsyncSession, client: AsyncClient
+        ):
+            user = await create_user(session=session, global_role=GlobalRole.USER)
+            project = await create_project(session=session, owner=user)
+            await add_project_member(
+                session=session, project=project, user=user, project_role=ProjectRole.ADMIN
+            )
+            body = {
+                "type": "nebius",
+                "creds": FAKE_NEBIUS_SERVICE_ACCOUNT_CREDS,
+                "regions": ["xx-xxxx1"],
+            }
+            with patch(
+                "dstack._internal.core.backends.nebius.resources.get_region_to_project_id_map"
+            ) as get_region_to_project_id_map:
+                get_region_to_project_id_map.return_value = {"eu-north1": "project-e00test"}
+                response = await client.post(
+                    f"/api/project/{project.name}/backends/create",
+                    headers=get_auth_headers(user.token),
+                    json=body,
+                )
+            assert response.status_code == 400, response.json()
+            res = await session.execute(select(BackendModel))
+            assert len(res.scalars().all()) == 0
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_creates_oci_backend(self, test_db, session: AsyncSession, client: AsyncClient):

tests/_internal/server/routers/test_fleets.py

@@ -366,6 +366,7 @@ class TestCreateFleet:
                     "name": "",
                     "default": False,
                     "reservation": None,
+                    "fleets": None,
                 },
                 "autocreated": False,
             },
@@ -484,6 +485,7 @@ class TestCreateFleet:
                     "name": "",
                     "default": False,
                     "reservation": None,
+                    "fleets": None,
                 },
                 "autocreated": False,
             },

tests/_internal/server/routers/test_prometheus.py

@@ -28,6 +28,7 @@ from dstack._internal.server.testing.common import (
     create_repo,
     create_run,
     create_user,
+    get_auth_headers,
     get_instance_offer_with_availability,
     get_job_provisioning_data,
     get_job_runtime_data,
@@ -38,6 +39,7 @@ from dstack._internal.server.testing.common import (
 @pytest.fixture
 def enable_metrics(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr("dstack._internal.server.settings.ENABLE_PROMETHEUS_METRICS", True)
+    monkeypatch.setattr("dstack._internal.server.routers.prometheus._auth._token", None)
 
 
 FAKE_NOW = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)
@@ -289,6 +291,25 @@ class TestGetPrometheusMetrics:
         response = await client.get("/metrics")
         assert response.status_code == 404
 
+    @pytest.mark.parametrize("token", [None, "foo"])
+    async def test_returns_403_if_not_authenticated(
+        self, monkeypatch: pytest.MonkeyPatch, client: AsyncClient, token: Optional[str]
+    ):
+        monkeypatch.setattr("dstack._internal.server.routers.prometheus._auth._token", "secret")
+        if token is not None:
+            headers = get_auth_headers(token)
+        else:
+            headers = None
+        response = await client.get("/metrics", headers=headers)
+        assert response.status_code == 403
+
+    async def test_returns_200_if_token_is_valid(
+        self, monkeypatch: pytest.MonkeyPatch, client: AsyncClient
+    ):
+        monkeypatch.setattr("dstack._internal.server.routers.prometheus._auth._token", "secret")
+        response = await client.get("/metrics", headers=get_auth_headers("secret"))
+        assert response.status_code == 200
+
 
 async def _create_project(session: AsyncSession, name: str, user: UserModel) -> ProjectModel:
     project = await create_project(session=session, owner=user, name=name)

tests/_internal/server/routers/test_runs.py

@@ -124,6 +124,7 @@ def get_dev_env_run_plan_dict(
                 "idle_duration": None,
                 "utilization_policy": None,
                 "reservation": None,
+                "fleets": None,
             },
             "configuration_path": "dstack.yaml",
             "profile": {
@@ -142,6 +143,7 @@ def get_dev_env_run_plan_dict(
                 "idle_duration": None,
                 "utilization_policy": None,
                 "reservation": None,
+                "fleets": None,
             },
             "repo_code_hash": None,
             "repo_data": {"repo_dir": "/repo", "repo_type": "local"},
@@ -274,6 +276,7 @@ def get_dev_env_run_dict(
                 "idle_duration": None,
                 "utilization_policy": None,
                 "reservation": None,
+                "fleets": None,
             },
             "configuration_path": "dstack.yaml",
             "profile": {
@@ -292,6 +295,7 @@ def get_dev_env_run_dict(
                 "idle_duration": None,
                 "utilization_policy": None,
                 "reservation": None,
+                "fleets": None,
             },
             "repo_code_hash": None,
             "repo_data": {"repo_dir": "/repo", "repo_type": "local"},

tests/_internal/utils/test_common.py

@@ -67,21 +67,36 @@ class TestPrettyDate:
         past_time = now - timedelta(days=5)
         assert pretty_date(past_time) == "5 days ago"
 
+    def test_week_ago(self):
+        now = datetime.now(tz=timezone.utc)
+        past_time = now - timedelta(days=7)
+        assert pretty_date(past_time) == "1 week ago"
+
     def test_weeks_ago(self):
         now = datetime.now(tz=timezone.utc)
         past_time = now - timedelta(days=21)
         assert pretty_date(past_time) == "3 weeks ago"
 
+    def test_month_ago(self):
+        now = datetime.now(tz=timezone.utc)
+        past_time = now - timedelta(days=31)
+        assert pretty_date(past_time) == "1 month ago"
+
     def test_months_ago(self):
         now = datetime.now(tz=timezone.utc)
         past_time = now - timedelta(days=90)
         assert pretty_date(past_time) == "3 months ago"
 
-    def test_years_ago(self):
+    def test_year_ago(self):
         now = datetime.now(tz=timezone.utc)
         past_time = now - timedelta(days=400)
         assert pretty_date(past_time) == "1 year ago"
 
+    def test_years_ago(self):
+        now = datetime.now(tz=timezone.utc)
+        past_time = now - timedelta(days=700)
+        assert pretty_date(past_time) == "2 years ago"
+
     def test_future_time(self):
         now = datetime.now(tz=timezone.utc)
         future_time = now + timedelta(hours=1)

tests/_internal/utils/test_event_loop.py

@@ -0,0 +1,18 @@
+import asyncio
+
+from dstack._internal.utils.event_loop import DaemonEventLoop
+
+
+def test_daemon_event_loop():
+    q = asyncio.Queue()
+
+    async def worker(i):
+        await q.put(i)
+
+    async def all_workers():
+        await asyncio.gather(*[worker(i) for i in range(3)])
+
+    loop = DaemonEventLoop()
+    loop.await_(all_workers())
+    assert q.qsize() == 3
+    assert {loop.await_(q.get()) for _ in range(3)} == {0, 1, 2}

dstack/_internal/core/backends/datacrunch/api_client.py

@@ -1,77 +0,0 @@
-from typing import Optional
-
-from datacrunch import DataCrunchClient
-from datacrunch.exceptions import APIException
-from datacrunch.instances.instances import Instance
-
-from dstack._internal.core.errors import NoCapacityError
-from dstack._internal.utils.ssh import get_public_key_fingerprint
-
-
-class DataCrunchAPIClient:
-    def __init__(self, client_id: str, client_secret: str):
-        self.client = DataCrunchClient(client_id, client_secret)
-
-    def delete_instance(self, instance_id: str) -> None:
-        try:
-            self.client.instances.action(id_list=[instance_id], action="delete")
-        except APIException:
-            pass
-
-    def get_or_create_ssh_key(self, name: str, public_key: str) -> str:
-        fingerprint = get_public_key_fingerprint(public_key)
-        keys = self.client.ssh_keys.get()
-        found_keys = [
-            key for key in keys if fingerprint == get_public_key_fingerprint(key.public_key)
-        ]
-        if found_keys:
-            key = found_keys[0]
-            return key.id
-
-        key = self.client.ssh_keys.create(name, public_key)
-        return key.id
-
-    def get_or_create_startup_scrpit(self, name: str, script: str) -> str:
-        scripts = self.client.startup_scripts.get()
-        found_scripts = [startup_script for startup_script in scripts if script == startup_script]
-        if found_scripts:
-            startup_script = found_scripts[0]
-            return startup_script.id
-
-        startup_script = self.client.startup_scripts.create(name, script)
-        return startup_script.id
-
-    def get_instance_by_id(self, instance_id: str) -> Optional[Instance]:
-        try:
-            return self.client.instances.get_by_id(instance_id)
-        except APIException:
-            return None
-
-    def deploy_instance(
-        self,
-        instance_type,
-        image,
-        ssh_key_ids,
-        hostname,
-        description,
-        startup_script_id,
-        disk_size,
-        is_spot=True,
-        location="FIN-01",
-    ) -> Instance:
-        try:
-            instance = self.client.instances.create(
-                instance_type=instance_type,
-                image=image,
-                ssh_key_ids=ssh_key_ids,
-                hostname=hostname,
-                description=description,
-                startup_script_id=startup_script_id,
-                is_spot=is_spot,
-                location=location,
-                os_volume={"name": "OS volume", "size": disk_size},
-            )
-        except APIException as e:
-            raise NoCapacityError(f"DataCrunch API error: {e.message}")
-
-        return instance