From 63d2927bb8fbf87269d6fb5501cf0b12f7c2bd33 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Wed, 28 Oct 2026 14:34:13 +0500 Subject: [PATCH 1/3] Mirror source layout in API and plugin tests --- .../_internal/core/compatibility/test_gpus.py | 23 ++++++++++++++++ .../_internal/core/compatibility/test_runs.py | 23 ++++++++++++++++ src/tests/api/_public/__init__.py | 0 src/tests/api/{ => _public}/test_runs.py | 0 src/tests/api/server/__init__.py | 0 src/tests/api/{ => server}/common.py | 0 .../test_api_client.py} | 0 src/tests/api/{ => server}/test_backends.py | 2 +- src/tests/api/{ => server}/test_projects.py | 2 +- src/tests/api/{ => server}/test_users.py | 2 +- src/tests/api/test_offers.py | 27 ------------------- src/tests/plugins/builtin/__init__.py | 0 .../plugins/builtin/rest_plugin/__init__.py | 0 .../rest_plugin/test_plugin.py} | 0 14 files changed, 49 insertions(+), 30 deletions(-) create mode 100644 src/tests/_internal/core/compatibility/test_gpus.py create mode 100644 src/tests/_internal/core/compatibility/test_runs.py create mode 100644 src/tests/api/_public/__init__.py rename src/tests/api/{ => _public}/test_runs.py (100%) create mode 100644 src/tests/api/server/__init__.py rename src/tests/api/{ => server}/common.py (100%) rename src/tests/api/{test_client.py => server/test_api_client.py} (100%) rename src/tests/api/{ => server}/test_backends.py (93%) rename src/tests/api/{ => server}/test_projects.py (97%) rename src/tests/api/{ => server}/test_users.py (97%) delete mode 100644 src/tests/api/test_offers.py create mode 100644 src/tests/plugins/builtin/__init__.py create mode 100644 src/tests/plugins/builtin/rest_plugin/__init__.py rename src/tests/plugins/{test_rest_plugin.py => builtin/rest_plugin/test_plugin.py} (100%) diff --git a/src/tests/_internal/core/compatibility/test_gpus.py b/src/tests/_internal/core/compatibility/test_gpus.py new file mode 100644 index 0000000000..e48895111c --- /dev/null +++ b/src/tests/_internal/core/compatibility/test_gpus.py @@ -0,0 +1,23 @@ +from typing import Optional + +import pytest + +from dstack._internal.core.compatibility.gpus import get_list_gpus_excludes +from dstack._internal.server.schemas.gpus import ListGpusRequest +from dstack._internal.server.testing.common import get_run_spec + + +class TestGetListGpusExcludes: + @pytest.mark.parametrize("reverse", [None, False, True]) + def test_omits_disabled_reverse(self, reverse: Optional[bool]): + kwargs = {} if reverse is None else {"reverse": reverse} + request = ListGpusRequest.model_validate( + {"run_spec": get_run_spec(repo_id="test-repo"), **kwargs} + ) + + body = request.model_dump(exclude=get_list_gpus_excludes(request)) + + if reverse: + assert body["reverse"] is True + else: + assert "reverse" not in body diff --git a/src/tests/_internal/core/compatibility/test_runs.py b/src/tests/_internal/core/compatibility/test_runs.py new file mode 100644 index 0000000000..2e29479fca --- /dev/null +++ b/src/tests/_internal/core/compatibility/test_runs.py @@ -0,0 +1,23 @@ +from typing import Optional + +import pytest + +from dstack._internal.core.compatibility.runs import get_get_plan_excludes +from dstack._internal.server.schemas.runs import GetRunPlanRequest +from dstack._internal.server.testing.common import get_run_spec + + +class TestGetGetPlanExcludes: + @pytest.mark.parametrize("reverse_offers", [None, False, True]) + def test_omits_disabled_reverse_offers(self, reverse_offers: Optional[bool]): + kwargs = {} if reverse_offers is None else {"reverse_offers": reverse_offers} + request = GetRunPlanRequest.model_validate( + {"run_spec": get_run_spec(repo_id="test-repo"), **kwargs} + ) + + body = request.model_dump(exclude=get_get_plan_excludes(request)) + + if reverse_offers: + assert body["reverse_offers"] is True + else: + assert "reverse_offers" not in body diff --git a/src/tests/api/_public/__init__.py b/src/tests/api/_public/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/tests/api/test_runs.py b/src/tests/api/_public/test_runs.py similarity index 100% rename from src/tests/api/test_runs.py rename to src/tests/api/_public/test_runs.py diff --git a/src/tests/api/server/__init__.py b/src/tests/api/server/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/tests/api/common.py b/src/tests/api/server/common.py similarity index 100% rename from src/tests/api/common.py rename to src/tests/api/server/common.py diff --git a/src/tests/api/test_client.py b/src/tests/api/server/test_api_client.py similarity index 100% rename from src/tests/api/test_client.py rename to src/tests/api/server/test_api_client.py diff --git a/src/tests/api/test_backends.py b/src/tests/api/server/test_backends.py similarity index 93% rename from src/tests/api/test_backends.py rename to src/tests/api/server/test_backends.py index c5af11e135..796f2d74fe 100644 --- a/src/tests/api/test_backends.py +++ b/src/tests/api/server/test_backends.py @@ -2,7 +2,7 @@ from dstack._internal.core.models.backends.base import BackendType from dstack.api.server._backends import BackendsAPIClient -from tests.api.common import RequestRecorder +from tests.api.server.common import RequestRecorder AWS_CONFIG_PAYLOAD = { "type": "aws", diff --git a/src/tests/api/test_projects.py b/src/tests/api/server/test_projects.py similarity index 97% rename from src/tests/api/test_projects.py rename to src/tests/api/server/test_projects.py index eca666f150..7a6ddc30ed 100644 --- a/src/tests/api/test_projects.py +++ b/src/tests/api/server/test_projects.py @@ -5,7 +5,7 @@ from dstack._internal.utils.common import get_or_error, render_datetime_as_api from dstack.api.server._projects import ProjectsAPIClient -from tests.api.common import RequestRecorder +from tests.api.server.common import RequestRecorder PROJECT_PAYLOAD = { "project_id": "1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e", diff --git a/src/tests/api/test_users.py b/src/tests/api/server/test_users.py similarity index 97% rename from src/tests/api/test_users.py rename to src/tests/api/server/test_users.py index 3e10536310..a5cfa50e5d 100644 --- a/src/tests/api/test_users.py +++ b/src/tests/api/server/test_users.py @@ -6,7 +6,7 @@ from dstack._internal.core.models.users import UsersInfoList from dstack._internal.utils.common import get_or_error, render_datetime_as_api from dstack.api.server._users import UsersAPIClient -from tests.api.common import RequestRecorder +from tests.api.server.common import RequestRecorder USER_PAYLOAD = { "id": "11111111-1111-4111-8111-111111111111", diff --git a/src/tests/api/test_offers.py b/src/tests/api/test_offers.py deleted file mode 100644 index b7139c67d3..0000000000 --- a/src/tests/api/test_offers.py +++ /dev/null @@ -1,27 +0,0 @@ -import pytest - -from dstack._internal.core.compatibility.gpus import get_list_gpus_excludes -from dstack._internal.core.compatibility.runs import get_get_plan_excludes -from dstack._internal.server.schemas.gpus import ListGpusRequest -from dstack._internal.server.schemas.runs import GetRunPlanRequest -from dstack._internal.server.testing.common import get_run_spec - - -class TestOfferRequestCompatibility: - @pytest.mark.parametrize( - ("request_type", "get_excludes", "field"), - [ - (GetRunPlanRequest, get_get_plan_excludes, "reverse_offers"), - (ListGpusRequest, get_list_gpus_excludes, "reverse"), - ], - ) - @pytest.mark.parametrize("reverse", [None, False, True]) - def test_omits_disabled_sort_flag(self, request_type, get_excludes, field, reverse): - kwargs = {} if reverse is None else {field: reverse} - request = request_type(run_spec=get_run_spec(repo_id="test-repo"), **kwargs) - body = request.model_dump(exclude=get_excludes(request)) - - if reverse: - assert body[field] is True - else: - assert field not in body diff --git a/src/tests/plugins/builtin/__init__.py b/src/tests/plugins/builtin/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/tests/plugins/builtin/rest_plugin/__init__.py b/src/tests/plugins/builtin/rest_plugin/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/tests/plugins/test_rest_plugin.py b/src/tests/plugins/builtin/rest_plugin/test_plugin.py similarity index 100% rename from src/tests/plugins/test_rest_plugin.py rename to src/tests/plugins/builtin/rest_plugin/test_plugin.py From 457f99fa2f427ab65d12eaffd1bfbb48caf36862 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Wed, 28 Oct 2026 14:54:17 +0500 Subject: [PATCH 2/3] Mirror source layout in server and CLI tests --- .../_internal/cli/commands/test_dstack.py | 19 - .../{configurators => }/test_profile.py | 0 src/tests/_internal/cli/test_main.py | 14 + .../core/backends/lium/test_compute.py | 97 +- .../core/backends/lium/test_volumes.py | 114 --- .../{test_instances => instances}/__init__.py | 0 .../{test_instances => instances}/conftest.py | 0 .../{test_instances => instances}/helpers.py | 0 .../test_check.py | 2 +- .../test_cloud_provisioning.py | 2 +- .../test_instances.py} | 2 +- .../test_ssh_deploy.py | 2 +- .../test_termination.py | 2 +- .../{test_runs => runs}/__init__.py | 0 .../{test_runs => runs}/conftest.py | 0 .../{test_runs => runs}/helpers.py | 0 .../{test_runs => runs}/test_active.py | 2 +- .../{test_runs => runs}/test_pending.py | 2 +- .../test_pipeline.py => runs/test_runs.py} | 0 .../test_terminating.py} | 2 +- ...t_running_jobs.py => test_jobs_running.py} | 0 ...bmitted_jobs.py => test_jobs_submitted.py} | 0 ...ating_jobs.py => test_jobs_terminating.py} | 0 .../services/gateways/test_gateway_update.py | 141 --- .../server/services/gateways/test_gateways.py | 133 +++ .../server/services/logs/__init__.py | 0 .../{test_logs.py => logs/test_aws.py} | 896 +---------------- .../server/services/logs/test_filelog.py | 904 ++++++++++++++++++ .../test_fluentbit.py} | 0 .../server/services/test_backend_configs.py | 240 ----- .../_internal/server/services/test_config.py | 234 ++++- 31 files changed, 1389 insertions(+), 1419 deletions(-) delete mode 100644 src/tests/_internal/cli/commands/test_dstack.py rename src/tests/_internal/cli/services/{configurators => }/test_profile.py (100%) delete mode 100644 src/tests/_internal/core/backends/lium/test_volumes.py rename src/tests/_internal/server/background/pipeline_tasks/{test_instances => instances}/__init__.py (100%) rename src/tests/_internal/server/background/pipeline_tasks/{test_instances => instances}/conftest.py (100%) rename src/tests/_internal/server/background/pipeline_tasks/{test_instances => instances}/helpers.py (100%) rename src/tests/_internal/server/background/pipeline_tasks/{test_instances => instances}/test_check.py (99%) rename src/tests/_internal/server/background/pipeline_tasks/{test_instances => instances}/test_cloud_provisioning.py (99%) rename src/tests/_internal/server/background/pipeline_tasks/{test_instances/test_pipeline.py => instances/test_instances.py} (99%) rename src/tests/_internal/server/background/pipeline_tasks/{test_instances => instances}/test_ssh_deploy.py (99%) rename src/tests/_internal/server/background/pipeline_tasks/{test_instances => instances}/test_termination.py (99%) rename src/tests/_internal/server/background/pipeline_tasks/{test_runs => runs}/__init__.py (100%) rename src/tests/_internal/server/background/pipeline_tasks/{test_runs => runs}/conftest.py (100%) rename src/tests/_internal/server/background/pipeline_tasks/{test_runs => runs}/helpers.py (100%) rename src/tests/_internal/server/background/pipeline_tasks/{test_runs => runs}/test_active.py (99%) rename src/tests/_internal/server/background/pipeline_tasks/{test_runs => runs}/test_pending.py (99%) rename src/tests/_internal/server/background/pipeline_tasks/{test_runs/test_pipeline.py => runs/test_runs.py} (100%) rename src/tests/_internal/server/background/pipeline_tasks/{test_runs/test_termination.py => runs/test_terminating.py} (99%) rename src/tests/_internal/server/background/pipeline_tasks/{test_running_jobs.py => test_jobs_running.py} (100%) rename src/tests/_internal/server/background/pipeline_tasks/{test_submitted_jobs.py => test_jobs_submitted.py} (100%) rename src/tests/_internal/server/background/pipeline_tasks/{test_terminating_jobs.py => test_jobs_terminating.py} (100%) delete mode 100644 src/tests/_internal/server/services/gateways/test_gateway_update.py create mode 100644 src/tests/_internal/server/services/logs/__init__.py rename src/tests/_internal/server/services/{test_logs.py => logs/test_aws.py} (56%) create mode 100644 src/tests/_internal/server/services/logs/test_filelog.py rename src/tests/_internal/server/services/{test_fluentbit_logs.py => logs/test_fluentbit.py} (100%) delete mode 100644 src/tests/_internal/server/services/test_backend_configs.py diff --git a/src/tests/_internal/cli/commands/test_dstack.py b/src/tests/_internal/cli/commands/test_dstack.py deleted file mode 100644 index 3179cf0c9a..0000000000 --- a/src/tests/_internal/cli/commands/test_dstack.py +++ /dev/null @@ -1,19 +0,0 @@ -from pathlib import Path - -from pytest import CaptureFixture - -from tests._internal.cli.common import run_dstack_cli - - -class TestDstack: - def test_prints_help_and_exists_with_0_exit_code(self, capsys: CaptureFixture): - exit_code = run_dstack_cli([]) - assert exit_code == 0 - assert capsys.readouterr().out.startswith("Usage: dstack [-h] [-v] COMMAND ...") - - def test_prints_errors_to_stderr(self, tmp_path: Path, capsys: CaptureFixture): - exit_code = run_dstack_cli(["ps", "--nonsense"], home_dir=tmp_path) - assert exit_code == 1 - captured = capsys.readouterr() - assert captured.out == "" - assert "Unrecognized arguments: --nonsense" in captured.err diff --git a/src/tests/_internal/cli/services/configurators/test_profile.py b/src/tests/_internal/cli/services/test_profile.py similarity index 100% rename from src/tests/_internal/cli/services/configurators/test_profile.py rename to src/tests/_internal/cli/services/test_profile.py diff --git a/src/tests/_internal/cli/test_main.py b/src/tests/_internal/cli/test_main.py index 7a5b42dbd7..0da5bf1246 100644 --- a/src/tests/_internal/cli/test_main.py +++ b/src/tests/_internal/cli/test_main.py @@ -1,10 +1,12 @@ import signal import subprocess import sys +from pathlib import Path from textwrap import dedent from unittest.mock import Mock import pytest +from pytest import CaptureFixture import dstack._internal.cli.main as main_module from dstack._internal.compat import IS_WINDOWS @@ -12,6 +14,18 @@ class TestMain: + def test_prints_help_and_exists_with_0_exit_code(self, capsys: CaptureFixture): + exit_code = run_dstack_cli([]) + assert exit_code == 0 + assert capsys.readouterr().out.startswith("Usage: dstack [-h] [-v] COMMAND ...") + + def test_prints_errors_to_stderr(self, tmp_path: Path, capsys: CaptureFixture): + exit_code = run_dstack_cli(["ps", "--nonsense"], home_dir=tmp_path) + assert exit_code == 1 + captured = capsys.readouterr() + assert captured.out == "" + assert "Unrecognized arguments: --nonsense" in captured.err + def test_reraises_keyboard_interrupt(self, monkeypatch: pytest.MonkeyPatch, capsys): monkeypatch.setattr(sys, "excepthook", sys.excepthook) monkeypatch.setattr(main_module, "check_for_updates", Mock(side_effect=KeyboardInterrupt)) diff --git a/src/tests/_internal/core/backends/lium/test_compute.py b/src/tests/_internal/core/backends/lium/test_compute.py index 48013ddb3c..2425f7f27c 100644 --- a/src/tests/_internal/core/backends/lium/test_compute.py +++ b/src/tests/_internal/core/backends/lium/test_compute.py @@ -4,6 +4,7 @@ import subprocess from datetime import datetime, timedelta, timezone from types import SimpleNamespace +from typing import Union, cast from unittest.mock import Mock, patch import gpuhunt @@ -19,6 +20,7 @@ _get_bootstrap_command, _get_image_env_and_cwd_command, _get_provider_ssh, + _get_volume_commands, ) from dstack._internal.core.backends.lium.models import LiumConfig, LiumCreds from dstack._internal.core.errors import ( @@ -31,8 +33,10 @@ from dstack._internal.core.models.common import RegistryAuth from dstack._internal.core.models.instances import InstanceAvailability, InstanceRuntime from dstack._internal.core.models.resources import ResourcesSpec -from dstack._internal.core.models.runs import JobProvisioningData, Requirements +from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements from dstack._internal.core.models.unix import UnixUser +from dstack._internal.core.models.volumes import LiumVolumeConfiguration, VolumeMountPoint +from dstack._internal.server.testing.common import get_volume from dstack._internal.utils.crypto import generate_rsa_key_pair_bytes COMPUTE_MODULE = "dstack._internal.core.backends.lium.compute" @@ -676,3 +680,94 @@ def test_oversized_environment_fails_without_starting_partial_runner(self, tmp_p assert result.returncode != 0 assert result.stdout == b"" + + +class TestVolumeLifecycle: + def test_create_unbounded_volume(self, compute): + compute.api_client.create_volume.return_value = {"id": "volume-id"} + with patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-name"): + data = compute.create_volume(Mock()) + + compute.api_client.create_volume.assert_called_once_with("unique-name") + assert data.volume_id == "volume-id" + assert data.size_gb is None + assert data.price is None + assert not data.attachable + assert not data.detachable + + def test_recovers_lost_create_response_by_unique_name(self, compute): + compute.api_client.create_volume.side_effect = LiumAPIError("timed out") + compute.api_client.get_volumes.return_value = [ + {"id": "other-id", "name": "other"}, + {"id": "volume-id", "name": "unique-name"}, + ] + with patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-name"): + data = compute.create_volume(Mock()) + + assert data.volume_id == "volume-id" + assert compute.api_client.create_volume.call_count == 1 + + def test_register_existing_volume(self, compute): + volume = SimpleNamespace(configuration=LiumVolumeConfiguration(volume_id="existing")) + compute.api_client.get_volume.return_value = {"id": "existing"} + + data = compute.register_volume(volume) + + assert data.volume_id == "existing" + compute.api_client.create_volume.assert_not_called() + + def test_register_missing_volume(self, compute): + compute.api_client.get_volume.return_value = None + with pytest.raises(ComputeError, match="not found"): + compute.register_volume( + SimpleNamespace(configuration=LiumVolumeConfiguration(volume_id="missing")) + ) + + def test_delete_volume(self, compute): + compute.delete_volume(SimpleNamespace(volume_id="volume-id")) + compute.api_client.delete_volume.assert_called_once_with("volume-id") + + +class TestVolumeMount: + @pytest.mark.parametrize( + "user", + [None, UnixUser(uid=10001), UnixUser(username="app")], + ) + def test_custom_path(self, user): + commands = _volume_commands(user=user, path="/checkpoints") + + assert commands[0] == "grep -q ' /mnt .* - fuse.s3fs ' /proc/self/mountinfo" + assert commands[1:] == ["mkdir -p -- /checkpoints", "mount --bind /mnt /checkpoints"] + + def test_native_mount_path_needs_no_bind(self): + assert len(_volume_commands(path="/mnt")) == 1 + + def test_quotes_mount_path(self): + commands = _volume_commands(path="/checkpoint dir; echo unsafe") + assert commands[-1] == "mount --bind /mnt '/checkpoint dir; echo unsafe'" + + @pytest.mark.parametrize("path", ["/", "/root", "/root/.dstack-lium", "/proc/a", "/mnt/a"]) + def test_rejects_paths_that_hide_bootstrap_or_nest_source(self, path): + with pytest.raises(ComputeError, match="cannot be mounted"): + _volume_commands(path=path) + + def test_two_volumes_rejected(self): + with pytest.raises(ComputeError, match="only one"): + _get_volume_commands(Mock(), Mock(), [Mock(), Mock()]) + + def test_selects_volume_from_alternatives(self): + assert _volume_commands(name=["elsewhere", "cache"], path="/data")[-1] == ( + "mount --bind /mnt /data" + ) + + +def _volume_commands(user=None, path="/checkpoints", name: Union[str, list[str]] = "cache"): + # A real `Job` needs a full `JobSpec`, while `_get_volume_commands()` only reads the user and + # the mount points + job = cast( + Job, + SimpleNamespace( + job_spec=SimpleNamespace(user=user, volumes=[VolumeMountPoint(name=name, path=path)]) + ), + ) + return _get_volume_commands(Mock(), job, [get_volume(name="cache")]) diff --git a/src/tests/_internal/core/backends/lium/test_volumes.py b/src/tests/_internal/core/backends/lium/test_volumes.py deleted file mode 100644 index 922c896bfc..0000000000 --- a/src/tests/_internal/core/backends/lium/test_volumes.py +++ /dev/null @@ -1,114 +0,0 @@ -from types import SimpleNamespace -from typing import Union, cast -from unittest.mock import Mock, patch - -import pytest - -from dstack._internal.core.backends.lium.api_client import LiumAPIError -from dstack._internal.core.backends.lium.compute import LiumCompute, _get_volume_commands -from dstack._internal.core.backends.lium.models import LiumConfig, LiumCreds -from dstack._internal.core.errors import ComputeError -from dstack._internal.core.models.runs import Job -from dstack._internal.core.models.unix import UnixUser -from dstack._internal.core.models.volumes import LiumVolumeConfiguration, VolumeMountPoint -from dstack._internal.server.testing.common import get_volume - -COMPUTE_MODULE = "dstack._internal.core.backends.lium.compute" - - -@pytest.fixture -def compute(): - compute = LiumCompute(LiumConfig(creds=LiumCreds(api_key="test-key"))) - compute.api_client = Mock() - return compute - - -class TestVolumeLifecycle: - def test_create_unbounded_volume(self, compute): - compute.api_client.create_volume.return_value = {"id": "volume-id"} - with patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-name"): - data = compute.create_volume(Mock()) - - compute.api_client.create_volume.assert_called_once_with("unique-name") - assert data.volume_id == "volume-id" - assert data.size_gb is None - assert data.price is None - assert not data.attachable - assert not data.detachable - - def test_recovers_lost_create_response_by_unique_name(self, compute): - compute.api_client.create_volume.side_effect = LiumAPIError("timed out") - compute.api_client.get_volumes.return_value = [ - {"id": "other-id", "name": "other"}, - {"id": "volume-id", "name": "unique-name"}, - ] - with patch(f"{COMPUTE_MODULE}.generate_unique_volume_name", return_value="unique-name"): - data = compute.create_volume(Mock()) - - assert data.volume_id == "volume-id" - assert compute.api_client.create_volume.call_count == 1 - - def test_register_existing_volume(self, compute): - volume = SimpleNamespace(configuration=LiumVolumeConfiguration(volume_id="existing")) - compute.api_client.get_volume.return_value = {"id": "existing"} - - data = compute.register_volume(volume) - - assert data.volume_id == "existing" - compute.api_client.create_volume.assert_not_called() - - def test_register_missing_volume(self, compute): - compute.api_client.get_volume.return_value = None - with pytest.raises(ComputeError, match="not found"): - compute.register_volume( - SimpleNamespace(configuration=LiumVolumeConfiguration(volume_id="missing")) - ) - - def test_delete_volume(self, compute): - compute.delete_volume(SimpleNamespace(volume_id="volume-id")) - compute.api_client.delete_volume.assert_called_once_with("volume-id") - - -class TestVolumeMount: - @pytest.mark.parametrize( - "user", - [None, UnixUser(uid=10001), UnixUser(username="app")], - ) - def test_custom_path(self, user): - commands = _volume_commands(user=user, path="/checkpoints") - - assert commands[0] == "grep -q ' /mnt .* - fuse.s3fs ' /proc/self/mountinfo" - assert commands[1:] == ["mkdir -p -- /checkpoints", "mount --bind /mnt /checkpoints"] - - def test_native_mount_path_needs_no_bind(self): - assert len(_volume_commands(path="/mnt")) == 1 - - def test_quotes_mount_path(self): - commands = _volume_commands(path="/checkpoint dir; echo unsafe") - assert commands[-1] == "mount --bind /mnt '/checkpoint dir; echo unsafe'" - - @pytest.mark.parametrize("path", ["/", "/root", "/root/.dstack-lium", "/proc/a", "/mnt/a"]) - def test_rejects_paths_that_hide_bootstrap_or_nest_source(self, path): - with pytest.raises(ComputeError, match="cannot be mounted"): - _volume_commands(path=path) - - def test_two_volumes_rejected(self): - with pytest.raises(ComputeError, match="only one"): - _get_volume_commands(Mock(), Mock(), [Mock(), Mock()]) - - def test_selects_volume_from_alternatives(self): - assert _volume_commands(name=["elsewhere", "cache"], path="/data")[-1] == ( - "mount --bind /mnt /data" - ) - - -def _volume_commands(user=None, path="/checkpoints", name: Union[str, list[str]] = "cache"): - # A real `Job` needs a full `JobSpec`, while `_get_volume_commands()` only reads the user and - # the mount points - job = cast( - Job, - SimpleNamespace( - job_spec=SimpleNamespace(user=user, volumes=[VolumeMountPoint(name=name, path=path)]) - ), - ) - return _get_volume_commands(Mock(), job, [get_volume(name="cache")]) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/__init__.py b/src/tests/_internal/server/background/pipeline_tasks/instances/__init__.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/test_instances/__init__.py rename to src/tests/_internal/server/background/pipeline_tasks/instances/__init__.py diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/conftest.py b/src/tests/_internal/server/background/pipeline_tasks/instances/conftest.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/test_instances/conftest.py rename to src/tests/_internal/server/background/pipeline_tasks/instances/conftest.py diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/helpers.py b/src/tests/_internal/server/background/pipeline_tasks/instances/helpers.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/test_instances/helpers.py rename to src/tests/_internal/server/background/pipeline_tasks/instances/helpers.py diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py b/src/tests/_internal/server/background/pipeline_tasks/instances/test_check.py similarity index 99% rename from src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py rename to src/tests/_internal/server/background/pipeline_tasks/instances/test_check.py index e4d08473dd..a1b61979bf 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py +++ b/src/tests/_internal/server/background/pipeline_tasks/instances/test_check.py @@ -56,7 +56,7 @@ list_events, ) from dstack._internal.utils.common import get_current_datetime -from tests._internal.server.background.pipeline_tasks.test_instances.helpers import ( +from tests._internal.server.background.pipeline_tasks.instances.helpers import ( process_instance, ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_cloud_provisioning.py b/src/tests/_internal/server/background/pipeline_tasks/instances/test_cloud_provisioning.py similarity index 99% rename from src/tests/_internal/server/background/pipeline_tasks/test_instances/test_cloud_provisioning.py rename to src/tests/_internal/server/background/pipeline_tasks/instances/test_cloud_provisioning.py index 45f29c8976..804f81e0fe 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_cloud_provisioning.py +++ b/src/tests/_internal/server/background/pipeline_tasks/instances/test_cloud_provisioning.py @@ -36,7 +36,7 @@ get_job_provisioning_data, get_placement_group_provisioning_data, ) -from tests._internal.server.background.pipeline_tasks.test_instances.helpers import ( +from tests._internal.server.background.pipeline_tasks.instances.helpers import ( instance_to_pipeline_item, lock_instance, process_instance, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_pipeline.py b/src/tests/_internal/server/background/pipeline_tasks/instances/test_instances.py similarity index 99% rename from src/tests/_internal/server/background/pipeline_tasks/test_instances/test_pipeline.py rename to src/tests/_internal/server/background/pipeline_tasks/instances/test_instances.py index 1c0ef2bca8..cf1639617d 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_pipeline.py +++ b/src/tests/_internal/server/background/pipeline_tasks/instances/test_instances.py @@ -21,7 +21,7 @@ create_project, ) from dstack._internal.utils.common import get_current_datetime -from tests._internal.server.background.pipeline_tasks.test_instances.helpers import ( +from tests._internal.server.background.pipeline_tasks.instances.helpers import ( instance_to_pipeline_item, lock_instance, process_instance, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_ssh_deploy.py b/src/tests/_internal/server/background/pipeline_tasks/instances/test_ssh_deploy.py similarity index 99% rename from src/tests/_internal/server/background/pipeline_tasks/test_instances/test_ssh_deploy.py rename to src/tests/_internal/server/background/pipeline_tasks/instances/test_ssh_deploy.py index c103458ed4..38f5bf3f59 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_ssh_deploy.py +++ b/src/tests/_internal/server/background/pipeline_tasks/instances/test_ssh_deploy.py @@ -19,7 +19,7 @@ get_remote_connection_info, ) from dstack._internal.utils.common import get_current_datetime -from tests._internal.server.background.pipeline_tasks.test_instances.helpers import ( +from tests._internal.server.background.pipeline_tasks.instances.helpers import ( process_instance, ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py b/src/tests/_internal/server/background/pipeline_tasks/instances/test_termination.py similarity index 99% rename from src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py rename to src/tests/_internal/server/background/pipeline_tasks/instances/test_termination.py index d0d273cae4..414518d0d1 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py +++ b/src/tests/_internal/server/background/pipeline_tasks/instances/test_termination.py @@ -22,7 +22,7 @@ create_project, get_job_provisioning_data, ) -from tests._internal.server.background.pipeline_tasks.test_instances.helpers import ( +from tests._internal.server.background.pipeline_tasks.instances.helpers import ( instance_to_pipeline_item, lock_instance, process_instance, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/__init__.py b/src/tests/_internal/server/background/pipeline_tasks/runs/__init__.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/test_runs/__init__.py rename to src/tests/_internal/server/background/pipeline_tasks/runs/__init__.py diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/conftest.py b/src/tests/_internal/server/background/pipeline_tasks/runs/conftest.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/test_runs/conftest.py rename to src/tests/_internal/server/background/pipeline_tasks/runs/conftest.py diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/helpers.py b/src/tests/_internal/server/background/pipeline_tasks/runs/helpers.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/test_runs/helpers.py rename to src/tests/_internal/server/background/pipeline_tasks/runs/helpers.py diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py b/src/tests/_internal/server/background/pipeline_tasks/runs/test_active.py similarity index 99% rename from src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py rename to src/tests/_internal/server/background/pipeline_tasks/runs/test_active.py index b2c2971725..0fbea23378 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py +++ b/src/tests/_internal/server/background/pipeline_tasks/runs/test_active.py @@ -52,7 +52,7 @@ get_run_spec, ) from dstack._internal.utils.common import get_current_datetime -from tests._internal.server.background.pipeline_tasks.test_runs.helpers import ( +from tests._internal.server.background.pipeline_tasks.runs.helpers import ( lock_run, run_to_pipeline_item, ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py b/src/tests/_internal/server/background/pipeline_tasks/runs/test_pending.py similarity index 99% rename from src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py rename to src/tests/_internal/server/background/pipeline_tasks/runs/test_pending.py index 6cb006821d..63f84d4bfc 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py +++ b/src/tests/_internal/server/background/pipeline_tasks/runs/test_pending.py @@ -31,7 +31,7 @@ get_run_spec, ) from dstack._internal.utils.common import get_current_datetime -from tests._internal.server.background.pipeline_tasks.test_runs.helpers import ( +from tests._internal.server.background.pipeline_tasks.runs.helpers import ( lock_run, run_to_pipeline_item, ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pipeline.py b/src/tests/_internal/server/background/pipeline_tasks/runs/test_runs.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pipeline.py rename to src/tests/_internal/server/background/pipeline_tasks/runs/test_runs.py diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_termination.py b/src/tests/_internal/server/background/pipeline_tasks/runs/test_terminating.py similarity index 99% rename from src/tests/_internal/server/background/pipeline_tasks/test_runs/test_termination.py rename to src/tests/_internal/server/background/pipeline_tasks/runs/test_terminating.py index 3a4eec04b1..ec8051ecba 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_termination.py +++ b/src/tests/_internal/server/background/pipeline_tasks/runs/test_terminating.py @@ -36,7 +36,7 @@ get_run_spec, ) from dstack._internal.utils.common import get_current_datetime -from tests._internal.server.background.pipeline_tasks.test_runs.helpers import ( +from tests._internal.server.background.pipeline_tasks.runs.helpers import ( lock_run, run_to_pipeline_item, ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_jobs_running.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py rename to src/tests/_internal/server/background/pipeline_tasks/test_jobs_running.py diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_jobs_submitted.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py rename to src/tests/_internal/server/background/pipeline_tasks/test_jobs_submitted.py diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_jobs_terminating.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py rename to src/tests/_internal/server/background/pipeline_tasks/test_jobs_terminating.py diff --git a/src/tests/_internal/server/services/gateways/test_gateway_update.py b/src/tests/_internal/server/services/gateways/test_gateway_update.py deleted file mode 100644 index 0d11d3de96..0000000000 --- a/src/tests/_internal/server/services/gateways/test_gateway_update.py +++ /dev/null @@ -1,141 +0,0 @@ -import shlex -import subprocess -from datetime import datetime, timedelta, timezone -from pathlib import Path -from unittest.mock import AsyncMock, Mock - -import pytest -from sqlalchemy.ext.asyncio import AsyncSession - -from dstack._internal import settings as core_settings -from dstack._internal.server.services import gateways -from dstack._internal.server.testing.common import ( - create_backend, - create_gateway_replica, - create_project, -) - - -@pytest.mark.asyncio -class TestInitGateways: - @pytest.mark.parametrize("target_version", ["0.22.2", "0.22.3"]) - async def test_release_version_controls_existing_gateway_update( - self, - session: AsyncSession, - monkeypatch: pytest.MonkeyPatch, - gateway_update_sandbox: tuple[Path, Path, Mock], - target_version: str, - ): - root, events_file, connection = gateway_update_sandbox - now = datetime(2026, 1, 1, tzinfo=timezone.utc) - previously_updated_at = now - timedelta(minutes=2) - monkeypatch.setattr(core_settings, "DSTACK_GATEWAY_PACKAGE_URL", None) - monkeypatch.setattr(core_settings, "DSTACK_VERSION", target_version) - monkeypatch.setattr(gateways.settings, "SKIP_GATEWAY_UPDATE", False) - monkeypatch.setattr(gateways, "get_current_datetime", lambda: now) - monkeypatch.setattr( - gateways.gateway_connections_pool, - "get_or_add", - AsyncMock(return_value=connection), - ) - monkeypatch.setattr( - gateways.gateway_connections_pool, - "all", - AsyncMock(return_value=[connection]), - ) - configure = AsyncMock() - monkeypatch.setattr(gateways, "configure_gateway_replica", configure) - project = await create_project(session=session) - backend = await create_backend(session=session, project_id=project.id) - replica = await create_gateway_replica(session=session, backend=backend) - replica.app_updated_at = previously_updated_at - await session.commit() - - await gateways.init_gateways(session) - - configure.assert_awaited_once() - events = events_file.read_text().splitlines() - if target_version == "0.22.2": - assert events == ["blue pip show dstack"] - assert (root / "version").read_text().strip() == "blue" - assert replica.app_updated_at == previously_updated_at - return - - assert (root / "version").read_text().strip() == "green" - assert replica.app_updated_at == now - install_package = f"green pip install dstack[gateway]=={target_version}" - install_service = "green python -m dstack._internal.proxy.gateway.systemd install" - reload_service = "systemctl daemon-reload active=green" - restart_service = "systemctl restart dstack.gateway active=green" - assert ( - events.index(install_package) - < events.index(install_service) - < events.index(reload_service) - < events.index(restart_service) - ) - - -@pytest.fixture -def gateway_update_sandbox(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): - # Execute the real shell: mocking aexec would miss quoting, version gating, and the order - # of service installation / venv switching / restart. Only the remote commands are stubbed; - # no SSH, package installation, sudo, or real systemd operations are performed. - root = tmp_path / "gateway root" - root.mkdir() - (root / "version").write_text("blue\n") - events = tmp_path / "events" - events.touch() - for color in ("blue", "green"): - venv_bin = root / color / "bin" - venv_bin.mkdir(parents=True) - for command in ("pip", "python"): - _write_executable( - venv_bin / command, - f'printf "%s\\n" "{color} {command} $*" >> "$TEST_EVENTS"\n' - + ( - 'if [ "$1" = show ]; then echo "Version: 0.22.2"; fi\n' - if command == "pip" - else "" - ), - ) - command_bin = tmp_path / "commands" - command_bin.mkdir() - _write_executable( - command_bin / "sudo", - 'if [ "$1" = systemctl ]; then\n' - ' printf "%s\\n" "$* active=$(cat "$TEST_ROOT/version")" >> "$TEST_EVENTS"\n' - "else\n" - ' exec "$@"\n' - "fi\n", - ) - # The only script substitution redirects its hard-coded remote root to this temporary dir. - remote_root = "root=/home/ubuntu/dstack" - assert gateways._GATEWAY_UPDATE_SCRIPT.count(remote_root) == 1 - monkeypatch.setattr( - gateways, - "_GATEWAY_UPDATE_SCRIPT", - gateways._GATEWAY_UPDATE_SCRIPT.replace(remote_root, f"root={shlex.quote(str(root))}"), - ) - - def execute(command: str, timeout: float) -> str: - return subprocess.run( - shlex.split(command), - check=True, - capture_output=True, - text=True, - timeout=min(timeout, 5), - env={ - "PATH": f"{command_bin}:/usr/bin:/bin", - "TEST_EVENTS": str(events), - "TEST_ROOT": str(root), - }, - ).stdout - - connection = Mock(ip_address="1.1.1.1") - connection.tunnel.aexec = AsyncMock(side_effect=execute) - return root, events, connection - - -def _write_executable(path: Path, script: str): - path.write_text("#!/bin/sh\nset -e\n" + script) - path.chmod(0o755) diff --git a/src/tests/_internal/server/services/gateways/test_gateways.py b/src/tests/_internal/server/services/gateways/test_gateways.py index 816619a4ed..a121f97edd 100644 --- a/src/tests/_internal/server/services/gateways/test_gateways.py +++ b/src/tests/_internal/server/services/gateways/test_gateways.py @@ -1,8 +1,16 @@ +import shlex +import subprocess +from datetime import datetime, timedelta, timezone +from pathlib import Path +from unittest.mock import AsyncMock, Mock + import pytest from sqlalchemy.ext.asyncio import AsyncSession +from dstack._internal import settings as core_settings from dstack._internal.proxy.gateway.const import SERVICE_SCALING_WINDOWS from dstack._internal.proxy.gateway.schemas.stats import Stat +from dstack._internal.server.services import gateways from dstack._internal.server.services.gateways import ( _merge_per_window_stats, get_gateway_replica_models, @@ -86,3 +94,128 @@ async def test_no_replicas_returns_empty(self, test_db, session: AsyncSession): await session.refresh(gateway, ["gateway_replicas", "gateway_replica"]) result = get_gateway_replica_models(gateway) assert result == [] + + +@pytest.mark.asyncio +class TestInitGateways: + @pytest.mark.parametrize("target_version", ["0.22.2", "0.22.3"]) + async def test_release_version_controls_existing_gateway_update( + self, + session: AsyncSession, + monkeypatch: pytest.MonkeyPatch, + gateway_update_sandbox: tuple[Path, Path, Mock], + target_version: str, + ): + root, events_file, connection = gateway_update_sandbox + now = datetime(2026, 1, 1, tzinfo=timezone.utc) + previously_updated_at = now - timedelta(minutes=2) + monkeypatch.setattr(core_settings, "DSTACK_GATEWAY_PACKAGE_URL", None) + monkeypatch.setattr(core_settings, "DSTACK_VERSION", target_version) + monkeypatch.setattr(gateways.settings, "SKIP_GATEWAY_UPDATE", False) + monkeypatch.setattr(gateways, "get_current_datetime", lambda: now) + monkeypatch.setattr( + gateways.gateway_connections_pool, + "get_or_add", + AsyncMock(return_value=connection), + ) + monkeypatch.setattr( + gateways.gateway_connections_pool, + "all", + AsyncMock(return_value=[connection]), + ) + configure = AsyncMock() + monkeypatch.setattr(gateways, "configure_gateway_replica", configure) + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + replica = await create_gateway_replica(session=session, backend=backend) + replica.app_updated_at = previously_updated_at + await session.commit() + + await gateways.init_gateways(session) + + configure.assert_awaited_once() + events = events_file.read_text().splitlines() + if target_version == "0.22.2": + assert events == ["blue pip show dstack"] + assert (root / "version").read_text().strip() == "blue" + assert replica.app_updated_at == previously_updated_at + return + + assert (root / "version").read_text().strip() == "green" + assert replica.app_updated_at == now + install_package = f"green pip install dstack[gateway]=={target_version}" + install_service = "green python -m dstack._internal.proxy.gateway.systemd install" + reload_service = "systemctl daemon-reload active=green" + restart_service = "systemctl restart dstack.gateway active=green" + assert ( + events.index(install_package) + < events.index(install_service) + < events.index(reload_service) + < events.index(restart_service) + ) + + +@pytest.fixture +def gateway_update_sandbox(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + # Execute the real shell: mocking aexec would miss quoting, version gating, and the order + # of service installation / venv switching / restart. Only the remote commands are stubbed; + # no SSH, package installation, sudo, or real systemd operations are performed. + root = tmp_path / "gateway root" + root.mkdir() + (root / "version").write_text("blue\n") + events = tmp_path / "events" + events.touch() + for color in ("blue", "green"): + venv_bin = root / color / "bin" + venv_bin.mkdir(parents=True) + for command in ("pip", "python"): + _write_executable( + venv_bin / command, + f'printf "%s\\n" "{color} {command} $*" >> "$TEST_EVENTS"\n' + + ( + 'if [ "$1" = show ]; then echo "Version: 0.22.2"; fi\n' + if command == "pip" + else "" + ), + ) + command_bin = tmp_path / "commands" + command_bin.mkdir() + _write_executable( + command_bin / "sudo", + 'if [ "$1" = systemctl ]; then\n' + ' printf "%s\\n" "$* active=$(cat "$TEST_ROOT/version")" >> "$TEST_EVENTS"\n' + "else\n" + ' exec "$@"\n' + "fi\n", + ) + # The only script substitution redirects its hard-coded remote root to this temporary dir. + remote_root = "root=/home/ubuntu/dstack" + assert gateways._GATEWAY_UPDATE_SCRIPT.count(remote_root) == 1 + monkeypatch.setattr( + gateways, + "_GATEWAY_UPDATE_SCRIPT", + gateways._GATEWAY_UPDATE_SCRIPT.replace(remote_root, f"root={shlex.quote(str(root))}"), + ) + + def execute(command: str, timeout: float) -> str: + return subprocess.run( + shlex.split(command), + check=True, + capture_output=True, + text=True, + timeout=min(timeout, 5), + env={ + "PATH": f"{command_bin}:/usr/bin:/bin", + "TEST_EVENTS": str(events), + "TEST_ROOT": str(root), + }, + ).stdout + + connection = Mock(ip_address="1.1.1.1") + connection.tunnel.aexec = AsyncMock(side_effect=execute) + return root, events, connection + + +def _write_executable(path: Path, script: str): + path.write_text("#!/bin/sh\nset -e\n" + script) + path.chmod(0o755) diff --git a/src/tests/_internal/server/services/logs/__init__.py b/src/tests/_internal/server/services/logs/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/tests/_internal/server/services/test_logs.py b/src/tests/_internal/server/services/logs/test_aws.py similarity index 56% rename from src/tests/_internal/server/services/test_logs.py rename to src/tests/_internal/server/services/logs/test_aws.py index 5154f903f1..3a80ee783a 100644 --- a/src/tests/_internal/server/services/test_logs.py +++ b/src/tests/_internal/server/services/logs/test_aws.py @@ -1,6 +1,5 @@ import logging from datetime import datetime, timedelta, timezone -from pathlib import Path from typing import List from unittest.mock import Mock, call from uuid import UUID @@ -11,8 +10,7 @@ import time_machine from sqlalchemy.ext.asyncio import AsyncSession -from dstack._internal.core.errors import ServerClientError -from dstack._internal.core.models.logs import LogEvent, LogEventSource, LogProducer +from dstack._internal.core.models.logs import LogEvent, LogEventSource from dstack._internal.server.models import ProjectModel from dstack._internal.server.schemas.logs import PollLogsRequest from dstack._internal.server.schemas.runner import LogEvent as RunnerLogEvent @@ -20,781 +18,9 @@ CloudWatchLogStorage, ) from dstack._internal.server.services.logs.base import LogStorageError -from dstack._internal.server.services.logs.filelog import FileLogStorage from dstack._internal.server.testing.common import create_project -class TestFileLogStorage: - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_writes_logs(self, test_db, session: AsyncSession, tmp_path: Path): - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513234, message=b"Hello"), - RunnerLogEvent(timestamp=1696586513235, message=b"World"), - ], - job_logs=[], - ) - runner_log_path = ( - tmp_path - / "projects" - / project.name - / "logs" - / "test_run" - / "1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e" - / "runner.log" - ) - assert runner_log_path.read_text() == ( - '{"timestamp":"2023-10-06T10:01:53.234000Z","log_source":"stdout","message":"Hello"}\n' - '{"timestamp":"2023-10-06T10:01:53.235000Z","log_source":"stdout","message":"World"}\n' - ) - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_basic(self, test_db, session: AsyncSession, tmp_path: Path): - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write test logs - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), - RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), - RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), - ], - job_logs=[], - ) - - # Test basic polling without pagination - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=10, - diagnose=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - assert len(job_submission_logs.logs) == 3 - assert job_submission_logs.next_token is None # No more logs, so no next_token - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_with_next_token_pagination( - self, test_db, session: AsyncSession, tmp_path: Path - ): - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write test logs - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), - RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), - RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), - RunnerLogEvent(timestamp=1696586513237, message=b"Log4"), - RunnerLogEvent(timestamp=1696586513238, message=b"Log5"), - ], - job_logs=[], - ) - - # First page: get 2 logs - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=2, - diagnose=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - assert len(job_submission_logs.logs) == 2 - assert job_submission_logs.logs[0].message == "Log1" - assert job_submission_logs.logs[1].message == "Log2" - assert job_submission_logs.next_token == "2" # Next line to read - - # Second page: use next_token - poll_request.next_token = job_submission_logs.next_token - job_submission_logs = log_storage.poll_logs(project, poll_request) - - assert len(job_submission_logs.logs) == 2 - assert job_submission_logs.logs[0].message == "Log3" - assert job_submission_logs.logs[1].message == "Log4" - assert job_submission_logs.next_token == "4" # Next line to read - - # Third page: get remaining log - poll_request.next_token = job_submission_logs.next_token - job_submission_logs = log_storage.poll_logs(project, poll_request) - - assert len(job_submission_logs.logs) == 1 - assert job_submission_logs.logs[0].message == "Log5" - assert job_submission_logs.next_token is None # No more logs - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_with_start_from_specific_line( - self, test_db, session: AsyncSession, tmp_path: Path - ): - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write test logs - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), - RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), - RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), - ], - job_logs=[], - ) - - # Start from line 1 (second log) - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - next_token="1", - limit=10, - diagnose=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - assert len(job_submission_logs.logs) == 2 - assert job_submission_logs.logs[0].message == "Log2" - assert job_submission_logs.logs[1].message == "Log3" - assert job_submission_logs.next_token is None - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_invalid_next_token_raises_error( - self, test_db, session: AsyncSession, tmp_path: Path - ): - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Test with non-integer next_token - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - next_token="invalid", - limit=10, - diagnose=True, - ) - with pytest.raises(ServerClientError): - log_storage.poll_logs(project, poll_request) - - # Test with negative next_token - poll_request.next_token = "-1" - with pytest.raises(ServerClientError): - log_storage.poll_logs(project, poll_request) - - # Test with float next_token - poll_request.next_token = "1.5" - with pytest.raises(ServerClientError): - log_storage.poll_logs(project, poll_request) - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_file_not_found_raises_no_error( - self, test_db, session: AsyncSession, tmp_path: Path - ): - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Test with non-existent log file - poll_request = PollLogsRequest( - run_name="nonexistent_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=10, - diagnose=True, - ) - log_storage.poll_logs(project, poll_request) - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_with_time_filtering_and_pagination( - self, test_db, session: AsyncSession, tmp_path: Path - ): - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write test logs with different timestamps - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent( - timestamp=1696586513234, message=b"Log1" - ), # 2023-10-06T10:01:53.234 - RunnerLogEvent( - timestamp=1696586513235, message=b"Log2" - ), # 2023-10-06T10:01:53.235 - RunnerLogEvent( - timestamp=1696586513236, message=b"Log3" - ), # 2023-10-06T10:01:53.236 - RunnerLogEvent( - timestamp=1696586513237, message=b"Log4" - ), # 2023-10-06T10:01:53.237 - ], - job_logs=[], - ) - - # Filter logs after 2023-10-06T10:01:53.235 with pagination - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - start_time=datetime(2023, 10, 6, 10, 1, 53, 235000, timezone.utc), - limit=1, - diagnose=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - # Should get Log3 first (timestamp > 235) - assert len(job_submission_logs.logs) == 1 - assert job_submission_logs.logs[0].message == "Log3" - assert job_submission_logs.next_token == "3" - - # Get next page - poll_request.next_token = job_submission_logs.next_token - job_submission_logs = log_storage.poll_logs(project, poll_request) - - # Should get Log4 - assert len(job_submission_logs.logs) == 1 - assert job_submission_logs.logs[0].message == "Log4" - # Should not have next_token since we reached end of file - assert job_submission_logs.next_token is None - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_empty_file_returns_empty_list( - self, test_db, session: AsyncSession, tmp_path: Path - ): - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Create empty log file - log_file_path = ( - tmp_path - / "projects" - / project.name - / "logs" - / "test_run" - / "1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e" - / "runner.log" - ) - log_file_path.parent.mkdir(parents=True, exist_ok=True) - log_file_path.write_text("") - - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=10, - diagnose=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - assert len(job_submission_logs.logs) == 0 - assert job_submission_logs.next_token is None - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_next_token_pagination_complete_workflow( - self, test_db, session: AsyncSession, tmp_path: Path - ): - """Test complete pagination workflow using next_token""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write 10 logs - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513000 + i, message=f"Log{i + 1}".encode()) - for i in range(10) - ], - job_logs=[], - ) - - # First page: get 3 logs - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=3, - diagnose=True, - ) - page1 = log_storage.poll_logs(project, poll_request) - - assert len(page1.logs) == 3 - assert page1.logs[0].message == "Log1" - assert page1.logs[1].message == "Log2" - assert page1.logs[2].message == "Log3" - assert page1.next_token == "3" # Next line to read - - # Second page: use next_token - poll_request.next_token = page1.next_token - page2 = log_storage.poll_logs(project, poll_request) - - assert len(page2.logs) == 3 - assert page2.logs[0].message == "Log4" - assert page2.logs[1].message == "Log5" - assert page2.logs[2].message == "Log6" - assert page2.next_token == "6" - - # Third page: get more logs - poll_request.next_token = page2.next_token - page3 = log_storage.poll_logs(project, poll_request) - - assert len(page3.logs) == 3 - assert page3.logs[0].message == "Log7" - assert page3.logs[1].message == "Log8" - assert page3.logs[2].message == "Log9" - assert page3.next_token == "9" - - # Fourth page: get last log - poll_request.next_token = page3.next_token - page4 = log_storage.poll_logs(project, poll_request) - - assert len(page4.logs) == 1 - assert page4.logs[0].message == "Log10" - assert page4.next_token is None - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_next_token_with_time_filtering( - self, test_db, session: AsyncSession, tmp_path: Path - ): - """Test next_token behavior with time filtering""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write logs with different timestamps - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513000, message=b"Log1"), # Before filter - RunnerLogEvent(timestamp=1696586513100, message=b"Log2"), # Before filter - RunnerLogEvent(timestamp=1696586513200, message=b"Log3"), # After filter - RunnerLogEvent(timestamp=1696586513300, message=b"Log4"), # After filter - RunnerLogEvent(timestamp=1696586513400, message=b"Log5"), # After filter - ], - job_logs=[], - ) - - # Filter logs after timestamp 150 with pagination - start_time = datetime.fromtimestamp(1696586513.150, tz=timezone.utc) - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - start_time=start_time, - limit=2, - diagnose=True, - ) - - page1 = log_storage.poll_logs(project, poll_request) - assert len(page1.logs) == 2 - assert page1.logs[0].message == "Log3" - assert page1.logs[1].message == "Log4" - assert page1.next_token == "4" - - # Get next page - poll_request.next_token = page1.next_token - page2 = log_storage.poll_logs(project, poll_request) - assert len(page2.logs) == 1 - assert page2.logs[0].message == "Log5" - assert page2.next_token is None - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_next_token_edge_cases(self, test_db, session: AsyncSession, tmp_path: Path): - """Test edge cases for next_token behavior""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write exactly one log - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513000, message=b"OnlyLog"), - ], - job_logs=[], - ) - - # Request with limit higher than available logs - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=10, - diagnose=True, - ) - result = log_storage.poll_logs(project, poll_request) - - assert len(result.logs) == 1 - assert result.logs[0].message == "OnlyLog" - assert result.next_token is None - - # Request with limit equal to available logs - poll_request.limit = 1 - result = log_storage.poll_logs(project, poll_request) - - assert len(result.logs) == 1 - assert result.logs[0].message == "OnlyLog" - assert result.next_token is None - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_next_token_beyond_file_end( - self, test_db, session: AsyncSession, tmp_path: Path - ): - """Test next_token that points beyond the end of file""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write 3 logs - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513000, message=b"Log1"), - RunnerLogEvent(timestamp=1696586513100, message=b"Log2"), - RunnerLogEvent(timestamp=1696586513200, message=b"Log3"), - ], - job_logs=[], - ) - - # Use next_token that points beyond the file - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - next_token="10", # Points beyond the 3 logs in file - limit=5, - diagnose=True, - ) - result = log_storage.poll_logs(project, poll_request) - - assert len(result.logs) == 0 - assert result.next_token is None - - -class TestPollLogsRequestValidation: - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_descending_basic( - self, test_db, session: AsyncSession, tmp_path: Path - ): - """Test basic descending log polling functionality.""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write test logs - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), - RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), - RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), - RunnerLogEvent(timestamp=1696586513237, message=b"Log4"), - RunnerLogEvent(timestamp=1696586513238, message=b"Log5"), - ], - job_logs=[], - ) - - # Test descending polling - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=10, - diagnose=True, - descending=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - # Should return logs in descending order (newest first) - assert len(job_submission_logs.logs) == 5 - assert job_submission_logs.logs[0].message == "Log5" - assert job_submission_logs.logs[1].message == "Log4" - assert job_submission_logs.logs[2].message == "Log3" - assert job_submission_logs.logs[3].message == "Log2" - assert job_submission_logs.logs[4].message == "Log1" - assert job_submission_logs.next_token is None # All logs returned - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_descending_with_limit( - self, test_db, session: AsyncSession, tmp_path: Path - ): - """Test descending log polling with limit smaller than total logs.""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write test logs - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), - RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), - RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), - RunnerLogEvent(timestamp=1696586513237, message=b"Log4"), - RunnerLogEvent(timestamp=1696586513238, message=b"Log5"), - ], - job_logs=[], - ) - - # Test with limit smaller than total logs - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=3, - diagnose=True, - descending=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - # Should return only the last 3 logs in descending order - assert len(job_submission_logs.logs) == 3 - assert job_submission_logs.logs[0].message == "Log5" - assert job_submission_logs.logs[1].message == "Log4" - assert job_submission_logs.logs[2].message == "Log3" - # Should have next_token for pagination - assert job_submission_logs.next_token is not None - - # Test next page - poll_request.next_token = job_submission_logs.next_token - job_submission_logs = log_storage.poll_logs(project, poll_request) - - # Should return remaining logs in descending order - assert len(job_submission_logs.logs) == 2 - assert job_submission_logs.logs[0].message == "Log2" - assert job_submission_logs.logs[1].message == "Log1" - assert job_submission_logs.next_token is None # No more logs - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_descending_with_time_filtering( - self, test_db, session: AsyncSession, tmp_path: Path - ): - """Test descending log polling with time filtering.""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write test logs with different timestamps - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent( - timestamp=1696586513234, message=b"Log1" - ), # 2023-10-06T10:01:53.234 - RunnerLogEvent( - timestamp=1696586513235, message=b"Log2" - ), # 2023-10-06T10:01:53.235 - RunnerLogEvent( - timestamp=1696586513236, message=b"Log3" - ), # 2023-10-06T10:01:53.236 - RunnerLogEvent( - timestamp=1696586513237, message=b"Log4" - ), # 2023-10-06T10:01:53.237 - RunnerLogEvent( - timestamp=1696586513238, message=b"Log5" - ), # 2023-10-06T10:01:53.238 - ], - job_logs=[], - ) - - # Filter logs between 2023-10-06T10:01:53.235 and 2023-10-06T10:01:53.237 - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - start_time=datetime(2023, 10, 6, 10, 1, 53, 235000, timezone.utc), - end_time=datetime(2023, 10, 6, 10, 1, 53, 237000, timezone.utc), - limit=10, - diagnose=True, - descending=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - # Should return logs in descending order within the time range - assert len(job_submission_logs.logs) == 2 - assert job_submission_logs.logs[0].message == "Log4" # timestamp 237 - assert job_submission_logs.logs[1].message == "Log3" # timestamp 236 - assert job_submission_logs.next_token is None # No more logs in range - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_descending_invalid_next_token( - self, test_db, session: AsyncSession, tmp_path: Path - ): - """Test descending log polling with invalid next_token.""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Test with non-integer next_token - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - next_token="invalid", - limit=10, - diagnose=True, - descending=True, - ) - with pytest.raises(ServerClientError): - log_storage.poll_logs(project, poll_request) - - # Test with negative next_token - poll_request.next_token = "-1" - with pytest.raises(ServerClientError): - log_storage.poll_logs(project, poll_request) - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_descending_empty_file( - self, test_db, session: AsyncSession, tmp_path: Path - ): - """Test descending log polling with empty log file.""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Test with non-existent log file - poll_request = PollLogsRequest( - run_name="nonexistent_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=10, - diagnose=True, - descending=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - # Should return empty logs without error - assert len(job_submission_logs.logs) == 0 - assert job_submission_logs.next_token is None - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_descending_pagination_workflow( - self, test_db, session: AsyncSession, tmp_path: Path - ): - """Test complete descending pagination workflow.""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Write test logs - log_storage.write_logs( - project=project, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - runner_logs=[ - RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), - RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), - RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), - RunnerLogEvent(timestamp=1696586513237, message=b"Log4"), - RunnerLogEvent(timestamp=1696586513238, message=b"Log5"), - ], - job_logs=[], - ) - - # First page: get last 2 logs - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=2, - diagnose=True, - descending=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - assert len(job_submission_logs.logs) == 2 - assert job_submission_logs.logs[0].message == "Log5" - assert job_submission_logs.logs[1].message == "Log4" - assert job_submission_logs.next_token is not None - - # Second page: get next 2 logs - poll_request.next_token = job_submission_logs.next_token - job_submission_logs = log_storage.poll_logs(project, poll_request) - - assert len(job_submission_logs.logs) == 2 - assert job_submission_logs.logs[0].message == "Log3" - assert job_submission_logs.logs[1].message == "Log2" - assert job_submission_logs.next_token is not None - - # Third page: get remaining log - poll_request.next_token = job_submission_logs.next_token - job_submission_logs = log_storage.poll_logs(project, poll_request) - - assert len(job_submission_logs.logs) == 1 - assert job_submission_logs.logs[0].message == "Log1" - assert job_submission_logs.next_token is None # No more logs - - @pytest.mark.asyncio - @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - async def test_poll_logs_descending_malformed_lines( - self, test_db, session: AsyncSession, tmp_path: Path - ): - """Test descending log polling with malformed log lines.""" - project = await create_project(session=session) - log_storage = FileLogStorage(tmp_path) - - # Create log file with malformed lines - log_file_path = log_storage._get_log_file_path( - project_name=project.name, - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - producer=LogProducer.RUNNER, - ) - log_file_path.parent.mkdir(exist_ok=True, parents=True) - - with open(log_file_path, "w") as f: - f.write( - '{"timestamp": "2023-10-06T10:01:53.234Z", "log_source": "stdout", "message": "Log1"}\n' - ) - f.write("invalid json line\n") - f.write( - '{"timestamp": "2023-10-06T10:01:53.235Z", "log_source": "stdout", "message": "Log2"}\n' - ) - f.write("another invalid line\n") - f.write( - '{"timestamp": "2023-10-06T10:01:53.236Z", "log_source": "stdout", "message": "Log3"}\n' - ) - - # Test descending polling - poll_request = PollLogsRequest( - run_name="test_run", - job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), - limit=10, - diagnose=True, - descending=True, - ) - job_submission_logs = log_storage.poll_logs(project, poll_request) - - # Should return only valid logs in descending order, skipping malformed lines - assert len(job_submission_logs.logs) == 3 - assert job_submission_logs.logs[0].message == "Log3" - assert job_submission_logs.logs[1].message == "Log2" - assert job_submission_logs.logs[2].message == "Log1" - - class TestCloudWatchLogStorage: FAKE_NOW = datetime(2023, 10, 6, 10, 1, 54, tzinfo=timezone.utc) @@ -1992,123 +1218,3 @@ async def test_poll_logs_descending_empty_response_with_same_token( assert job_submission_logs.logs == [] assert job_submission_logs.next_token is None assert mock_client.get_log_events.call_count == 1 - - -class TestFileLogStorageReadLinesReversed: - # No changes to the first 6 tests, they will now pass. - def test_basic_file(self, tmp_path: Path): - file = tmp_path / "test.txt" - content = b"line1\nline2\nline3\n" - file.write_bytes(content) - lines = list(FileLogStorage._read_lines_reversed(file)) - assert lines == [ - (b"", 18), - (b"line3", 12), - (b"line2", 6), - (b"line1", 0), - ] - - def test_file_without_trailing_newline(self, tmp_path: Path): - file = tmp_path / "test.txt" - content = b"line1\nline2" - file.write_bytes(content) - lines = list(FileLogStorage._read_lines_reversed(file)) - assert lines == [ - (b"line2", 6), - (b"line1", 0), - ] - - def test_empty_file(self, tmp_path: Path): - file = tmp_path / "test.txt" - file.touch() - lines = list(FileLogStorage._read_lines_reversed(file)) - assert lines == [] - - def test_single_line_file(self, tmp_path: Path): - file = tmp_path / "test.txt" - content = b"the only line" - file.write_bytes(content) - lines = list(FileLogStorage._read_lines_reversed(file)) - assert lines == [(b"the only line", 0)] - - def test_file_with_empty_lines(self, tmp_path: Path): - file = tmp_path / "test.txt" - content = b"lineA\n\nlineC\n" - file.write_bytes(content) - lines = list(FileLogStorage._read_lines_reversed(file)) - assert lines == [ - (b"", 13), - (b"lineC", 7), - (b"", 6), - (b"lineA", 0), - ] - - def test_file_with_only_newlines(self, tmp_path: Path): - file = tmp_path / "test.txt" - content = b"\n\n" - file.write_bytes(content) - lines = list(FileLogStorage._read_lines_reversed(file)) - assert lines == [ - (b"", 2), - (b"", 1), - ] - - def test_large_file_spanning_multiple_chunks(self, tmp_path: Path): - file = tmp_path / "large_file.txt" - line_content = b"abcdefghi" # 9 bytes + 1 newline = 10 bytes per line - num_lines = 5 - content = b"\n".join([line_content] * num_lines) - file.write_bytes(content) - # Pass the small chunk_size directly to the method - lines = list(FileLogStorage._read_lines_reversed(file, chunk_size=10)) - assert len(lines) == num_lines - assert lines[0] == (line_content, 40) - assert lines[1] == (line_content, 30) - assert lines[2] == (line_content, 20) - assert lines[3] == (line_content, 10) - assert lines[4] == (line_content, 0) - - # The rest of the tests will now pass without modification - def test_start_offset_in_middle_of_line(self, tmp_path: Path): - file = tmp_path / "test.txt" - content = b"line1\nline2\nline3\n" - file.write_bytes(content) - lines = list(FileLogStorage._read_lines_reversed(file, start_offset=10)) - assert lines == [ - (b"line2", 6), - (b"line1", 0), - ] - - def test_start_offset_at_line_boundary(self, tmp_path: Path): - file = tmp_path / "test.txt" - content = b"line1\nline2\nline3\n" - file.write_bytes(content) - lines = list(FileLogStorage._read_lines_reversed(file, start_offset=12)) - assert lines == [ - (b"line2", 6), - (b"line1", 0), - ] - - def test_start_offset_zero(self, tmp_path: Path): - file = tmp_path / "test.txt" - content = b"line1\nline2\nline3\n" - file.write_bytes(content) - lines = list(FileLogStorage._read_lines_reversed(file, start_offset=0)) - assert lines == [] - - def test_start_offset_larger_than_file(self, tmp_path: Path): - file = tmp_path / "test.txt" - content = b"line1\nline2\n" - file.write_bytes(content) - lines_with_offset = list(FileLogStorage._read_lines_reversed(file, start_offset=1000)) - lines_without_offset = list(FileLogStorage._read_lines_reversed(file)) - assert lines_with_offset == lines_without_offset - assert lines_with_offset == [(b"", 12), (b"line2", 6), (b"line1", 0)] - - def test_long_line_larger_than_chunk(self, tmp_path: Path): - file = tmp_path / "long_line.txt" - content = b"a" * 25 - file.write_bytes(content) - # Pass the small chunk_size directly - lines = list(FileLogStorage._read_lines_reversed(file, chunk_size=10)) - assert lines == [(content, 0)] diff --git a/src/tests/_internal/server/services/logs/test_filelog.py b/src/tests/_internal/server/services/logs/test_filelog.py new file mode 100644 index 0000000000..8fb41825e2 --- /dev/null +++ b/src/tests/_internal/server/services/logs/test_filelog.py @@ -0,0 +1,904 @@ +from datetime import datetime, timezone +from pathlib import Path +from uuid import UUID + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from dstack._internal.core.errors import ServerClientError +from dstack._internal.core.models.logs import LogProducer +from dstack._internal.server.schemas.logs import PollLogsRequest +from dstack._internal.server.schemas.runner import LogEvent as RunnerLogEvent +from dstack._internal.server.services.logs.filelog import FileLogStorage +from dstack._internal.server.testing.common import create_project + + +class TestFileLogStorage: + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_writes_logs(self, test_db, session: AsyncSession, tmp_path: Path): + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513234, message=b"Hello"), + RunnerLogEvent(timestamp=1696586513235, message=b"World"), + ], + job_logs=[], + ) + runner_log_path = ( + tmp_path + / "projects" + / project.name + / "logs" + / "test_run" + / "1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e" + / "runner.log" + ) + assert runner_log_path.read_text() == ( + '{"timestamp":"2023-10-06T10:01:53.234000Z","log_source":"stdout","message":"Hello"}\n' + '{"timestamp":"2023-10-06T10:01:53.235000Z","log_source":"stdout","message":"World"}\n' + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_basic(self, test_db, session: AsyncSession, tmp_path: Path): + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write test logs + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), + RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), + RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), + ], + job_logs=[], + ) + + # Test basic polling without pagination + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=10, + diagnose=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + assert len(job_submission_logs.logs) == 3 + assert job_submission_logs.next_token is None # No more logs, so no next_token + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_with_next_token_pagination( + self, test_db, session: AsyncSession, tmp_path: Path + ): + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write test logs + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), + RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), + RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), + RunnerLogEvent(timestamp=1696586513237, message=b"Log4"), + RunnerLogEvent(timestamp=1696586513238, message=b"Log5"), + ], + job_logs=[], + ) + + # First page: get 2 logs + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=2, + diagnose=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + assert len(job_submission_logs.logs) == 2 + assert job_submission_logs.logs[0].message == "Log1" + assert job_submission_logs.logs[1].message == "Log2" + assert job_submission_logs.next_token == "2" # Next line to read + + # Second page: use next_token + poll_request.next_token = job_submission_logs.next_token + job_submission_logs = log_storage.poll_logs(project, poll_request) + + assert len(job_submission_logs.logs) == 2 + assert job_submission_logs.logs[0].message == "Log3" + assert job_submission_logs.logs[1].message == "Log4" + assert job_submission_logs.next_token == "4" # Next line to read + + # Third page: get remaining log + poll_request.next_token = job_submission_logs.next_token + job_submission_logs = log_storage.poll_logs(project, poll_request) + + assert len(job_submission_logs.logs) == 1 + assert job_submission_logs.logs[0].message == "Log5" + assert job_submission_logs.next_token is None # No more logs + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_with_start_from_specific_line( + self, test_db, session: AsyncSession, tmp_path: Path + ): + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write test logs + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), + RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), + RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), + ], + job_logs=[], + ) + + # Start from line 1 (second log) + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + next_token="1", + limit=10, + diagnose=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + assert len(job_submission_logs.logs) == 2 + assert job_submission_logs.logs[0].message == "Log2" + assert job_submission_logs.logs[1].message == "Log3" + assert job_submission_logs.next_token is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_invalid_next_token_raises_error( + self, test_db, session: AsyncSession, tmp_path: Path + ): + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Test with non-integer next_token + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + next_token="invalid", + limit=10, + diagnose=True, + ) + with pytest.raises(ServerClientError): + log_storage.poll_logs(project, poll_request) + + # Test with negative next_token + poll_request.next_token = "-1" + with pytest.raises(ServerClientError): + log_storage.poll_logs(project, poll_request) + + # Test with float next_token + poll_request.next_token = "1.5" + with pytest.raises(ServerClientError): + log_storage.poll_logs(project, poll_request) + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_file_not_found_raises_no_error( + self, test_db, session: AsyncSession, tmp_path: Path + ): + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Test with non-existent log file + poll_request = PollLogsRequest( + run_name="nonexistent_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=10, + diagnose=True, + ) + log_storage.poll_logs(project, poll_request) + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_with_time_filtering_and_pagination( + self, test_db, session: AsyncSession, tmp_path: Path + ): + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write test logs with different timestamps + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent( + timestamp=1696586513234, message=b"Log1" + ), # 2023-10-06T10:01:53.234 + RunnerLogEvent( + timestamp=1696586513235, message=b"Log2" + ), # 2023-10-06T10:01:53.235 + RunnerLogEvent( + timestamp=1696586513236, message=b"Log3" + ), # 2023-10-06T10:01:53.236 + RunnerLogEvent( + timestamp=1696586513237, message=b"Log4" + ), # 2023-10-06T10:01:53.237 + ], + job_logs=[], + ) + + # Filter logs after 2023-10-06T10:01:53.235 with pagination + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + start_time=datetime(2023, 10, 6, 10, 1, 53, 235000, timezone.utc), + limit=1, + diagnose=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + # Should get Log3 first (timestamp > 235) + assert len(job_submission_logs.logs) == 1 + assert job_submission_logs.logs[0].message == "Log3" + assert job_submission_logs.next_token == "3" + + # Get next page + poll_request.next_token = job_submission_logs.next_token + job_submission_logs = log_storage.poll_logs(project, poll_request) + + # Should get Log4 + assert len(job_submission_logs.logs) == 1 + assert job_submission_logs.logs[0].message == "Log4" + # Should not have next_token since we reached end of file + assert job_submission_logs.next_token is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_empty_file_returns_empty_list( + self, test_db, session: AsyncSession, tmp_path: Path + ): + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Create empty log file + log_file_path = ( + tmp_path + / "projects" + / project.name + / "logs" + / "test_run" + / "1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e" + / "runner.log" + ) + log_file_path.parent.mkdir(parents=True, exist_ok=True) + log_file_path.write_text("") + + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=10, + diagnose=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + assert len(job_submission_logs.logs) == 0 + assert job_submission_logs.next_token is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_next_token_pagination_complete_workflow( + self, test_db, session: AsyncSession, tmp_path: Path + ): + """Test complete pagination workflow using next_token""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write 10 logs + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513000 + i, message=f"Log{i + 1}".encode()) + for i in range(10) + ], + job_logs=[], + ) + + # First page: get 3 logs + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=3, + diagnose=True, + ) + page1 = log_storage.poll_logs(project, poll_request) + + assert len(page1.logs) == 3 + assert page1.logs[0].message == "Log1" + assert page1.logs[1].message == "Log2" + assert page1.logs[2].message == "Log3" + assert page1.next_token == "3" # Next line to read + + # Second page: use next_token + poll_request.next_token = page1.next_token + page2 = log_storage.poll_logs(project, poll_request) + + assert len(page2.logs) == 3 + assert page2.logs[0].message == "Log4" + assert page2.logs[1].message == "Log5" + assert page2.logs[2].message == "Log6" + assert page2.next_token == "6" + + # Third page: get more logs + poll_request.next_token = page2.next_token + page3 = log_storage.poll_logs(project, poll_request) + + assert len(page3.logs) == 3 + assert page3.logs[0].message == "Log7" + assert page3.logs[1].message == "Log8" + assert page3.logs[2].message == "Log9" + assert page3.next_token == "9" + + # Fourth page: get last log + poll_request.next_token = page3.next_token + page4 = log_storage.poll_logs(project, poll_request) + + assert len(page4.logs) == 1 + assert page4.logs[0].message == "Log10" + assert page4.next_token is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_next_token_with_time_filtering( + self, test_db, session: AsyncSession, tmp_path: Path + ): + """Test next_token behavior with time filtering""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write logs with different timestamps + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513000, message=b"Log1"), # Before filter + RunnerLogEvent(timestamp=1696586513100, message=b"Log2"), # Before filter + RunnerLogEvent(timestamp=1696586513200, message=b"Log3"), # After filter + RunnerLogEvent(timestamp=1696586513300, message=b"Log4"), # After filter + RunnerLogEvent(timestamp=1696586513400, message=b"Log5"), # After filter + ], + job_logs=[], + ) + + # Filter logs after timestamp 150 with pagination + start_time = datetime.fromtimestamp(1696586513.150, tz=timezone.utc) + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + start_time=start_time, + limit=2, + diagnose=True, + ) + + page1 = log_storage.poll_logs(project, poll_request) + assert len(page1.logs) == 2 + assert page1.logs[0].message == "Log3" + assert page1.logs[1].message == "Log4" + assert page1.next_token == "4" + + # Get next page + poll_request.next_token = page1.next_token + page2 = log_storage.poll_logs(project, poll_request) + assert len(page2.logs) == 1 + assert page2.logs[0].message == "Log5" + assert page2.next_token is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_next_token_edge_cases(self, test_db, session: AsyncSession, tmp_path: Path): + """Test edge cases for next_token behavior""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write exactly one log + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513000, message=b"OnlyLog"), + ], + job_logs=[], + ) + + # Request with limit higher than available logs + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=10, + diagnose=True, + ) + result = log_storage.poll_logs(project, poll_request) + + assert len(result.logs) == 1 + assert result.logs[0].message == "OnlyLog" + assert result.next_token is None + + # Request with limit equal to available logs + poll_request.limit = 1 + result = log_storage.poll_logs(project, poll_request) + + assert len(result.logs) == 1 + assert result.logs[0].message == "OnlyLog" + assert result.next_token is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_next_token_beyond_file_end( + self, test_db, session: AsyncSession, tmp_path: Path + ): + """Test next_token that points beyond the end of file""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write 3 logs + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513000, message=b"Log1"), + RunnerLogEvent(timestamp=1696586513100, message=b"Log2"), + RunnerLogEvent(timestamp=1696586513200, message=b"Log3"), + ], + job_logs=[], + ) + + # Use next_token that points beyond the file + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + next_token="10", # Points beyond the 3 logs in file + limit=5, + diagnose=True, + ) + result = log_storage.poll_logs(project, poll_request) + + assert len(result.logs) == 0 + assert result.next_token is None + + +class TestPollLogsRequestValidation: + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_descending_basic( + self, test_db, session: AsyncSession, tmp_path: Path + ): + """Test basic descending log polling functionality.""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write test logs + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), + RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), + RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), + RunnerLogEvent(timestamp=1696586513237, message=b"Log4"), + RunnerLogEvent(timestamp=1696586513238, message=b"Log5"), + ], + job_logs=[], + ) + + # Test descending polling + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=10, + diagnose=True, + descending=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + # Should return logs in descending order (newest first) + assert len(job_submission_logs.logs) == 5 + assert job_submission_logs.logs[0].message == "Log5" + assert job_submission_logs.logs[1].message == "Log4" + assert job_submission_logs.logs[2].message == "Log3" + assert job_submission_logs.logs[3].message == "Log2" + assert job_submission_logs.logs[4].message == "Log1" + assert job_submission_logs.next_token is None # All logs returned + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_descending_with_limit( + self, test_db, session: AsyncSession, tmp_path: Path + ): + """Test descending log polling with limit smaller than total logs.""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write test logs + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), + RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), + RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), + RunnerLogEvent(timestamp=1696586513237, message=b"Log4"), + RunnerLogEvent(timestamp=1696586513238, message=b"Log5"), + ], + job_logs=[], + ) + + # Test with limit smaller than total logs + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=3, + diagnose=True, + descending=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + # Should return only the last 3 logs in descending order + assert len(job_submission_logs.logs) == 3 + assert job_submission_logs.logs[0].message == "Log5" + assert job_submission_logs.logs[1].message == "Log4" + assert job_submission_logs.logs[2].message == "Log3" + # Should have next_token for pagination + assert job_submission_logs.next_token is not None + + # Test next page + poll_request.next_token = job_submission_logs.next_token + job_submission_logs = log_storage.poll_logs(project, poll_request) + + # Should return remaining logs in descending order + assert len(job_submission_logs.logs) == 2 + assert job_submission_logs.logs[0].message == "Log2" + assert job_submission_logs.logs[1].message == "Log1" + assert job_submission_logs.next_token is None # No more logs + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_descending_with_time_filtering( + self, test_db, session: AsyncSession, tmp_path: Path + ): + """Test descending log polling with time filtering.""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write test logs with different timestamps + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent( + timestamp=1696586513234, message=b"Log1" + ), # 2023-10-06T10:01:53.234 + RunnerLogEvent( + timestamp=1696586513235, message=b"Log2" + ), # 2023-10-06T10:01:53.235 + RunnerLogEvent( + timestamp=1696586513236, message=b"Log3" + ), # 2023-10-06T10:01:53.236 + RunnerLogEvent( + timestamp=1696586513237, message=b"Log4" + ), # 2023-10-06T10:01:53.237 + RunnerLogEvent( + timestamp=1696586513238, message=b"Log5" + ), # 2023-10-06T10:01:53.238 + ], + job_logs=[], + ) + + # Filter logs between 2023-10-06T10:01:53.235 and 2023-10-06T10:01:53.237 + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + start_time=datetime(2023, 10, 6, 10, 1, 53, 235000, timezone.utc), + end_time=datetime(2023, 10, 6, 10, 1, 53, 237000, timezone.utc), + limit=10, + diagnose=True, + descending=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + # Should return logs in descending order within the time range + assert len(job_submission_logs.logs) == 2 + assert job_submission_logs.logs[0].message == "Log4" # timestamp 237 + assert job_submission_logs.logs[1].message == "Log3" # timestamp 236 + assert job_submission_logs.next_token is None # No more logs in range + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_descending_invalid_next_token( + self, test_db, session: AsyncSession, tmp_path: Path + ): + """Test descending log polling with invalid next_token.""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Test with non-integer next_token + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + next_token="invalid", + limit=10, + diagnose=True, + descending=True, + ) + with pytest.raises(ServerClientError): + log_storage.poll_logs(project, poll_request) + + # Test with negative next_token + poll_request.next_token = "-1" + with pytest.raises(ServerClientError): + log_storage.poll_logs(project, poll_request) + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_descending_empty_file( + self, test_db, session: AsyncSession, tmp_path: Path + ): + """Test descending log polling with empty log file.""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Test with non-existent log file + poll_request = PollLogsRequest( + run_name="nonexistent_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=10, + diagnose=True, + descending=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + # Should return empty logs without error + assert len(job_submission_logs.logs) == 0 + assert job_submission_logs.next_token is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_descending_pagination_workflow( + self, test_db, session: AsyncSession, tmp_path: Path + ): + """Test complete descending pagination workflow.""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Write test logs + log_storage.write_logs( + project=project, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + runner_logs=[ + RunnerLogEvent(timestamp=1696586513234, message=b"Log1"), + RunnerLogEvent(timestamp=1696586513235, message=b"Log2"), + RunnerLogEvent(timestamp=1696586513236, message=b"Log3"), + RunnerLogEvent(timestamp=1696586513237, message=b"Log4"), + RunnerLogEvent(timestamp=1696586513238, message=b"Log5"), + ], + job_logs=[], + ) + + # First page: get last 2 logs + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=2, + diagnose=True, + descending=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + assert len(job_submission_logs.logs) == 2 + assert job_submission_logs.logs[0].message == "Log5" + assert job_submission_logs.logs[1].message == "Log4" + assert job_submission_logs.next_token is not None + + # Second page: get next 2 logs + poll_request.next_token = job_submission_logs.next_token + job_submission_logs = log_storage.poll_logs(project, poll_request) + + assert len(job_submission_logs.logs) == 2 + assert job_submission_logs.logs[0].message == "Log3" + assert job_submission_logs.logs[1].message == "Log2" + assert job_submission_logs.next_token is not None + + # Third page: get remaining log + poll_request.next_token = job_submission_logs.next_token + job_submission_logs = log_storage.poll_logs(project, poll_request) + + assert len(job_submission_logs.logs) == 1 + assert job_submission_logs.logs[0].message == "Log1" + assert job_submission_logs.next_token is None # No more logs + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_poll_logs_descending_malformed_lines( + self, test_db, session: AsyncSession, tmp_path: Path + ): + """Test descending log polling with malformed log lines.""" + project = await create_project(session=session) + log_storage = FileLogStorage(tmp_path) + + # Create log file with malformed lines + log_file_path = log_storage._get_log_file_path( + project_name=project.name, + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + producer=LogProducer.RUNNER, + ) + log_file_path.parent.mkdir(exist_ok=True, parents=True) + + with open(log_file_path, "w") as f: + f.write( + '{"timestamp": "2023-10-06T10:01:53.234Z", "log_source": "stdout", "message": "Log1"}\n' + ) + f.write("invalid json line\n") + f.write( + '{"timestamp": "2023-10-06T10:01:53.235Z", "log_source": "stdout", "message": "Log2"}\n' + ) + f.write("another invalid line\n") + f.write( + '{"timestamp": "2023-10-06T10:01:53.236Z", "log_source": "stdout", "message": "Log3"}\n' + ) + + # Test descending polling + poll_request = PollLogsRequest( + run_name="test_run", + job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"), + limit=10, + diagnose=True, + descending=True, + ) + job_submission_logs = log_storage.poll_logs(project, poll_request) + + # Should return only valid logs in descending order, skipping malformed lines + assert len(job_submission_logs.logs) == 3 + assert job_submission_logs.logs[0].message == "Log3" + assert job_submission_logs.logs[1].message == "Log2" + assert job_submission_logs.logs[2].message == "Log1" + + +class TestFileLogStorageReadLinesReversed: + # No changes to the first 6 tests, they will now pass. + def test_basic_file(self, tmp_path: Path): + file = tmp_path / "test.txt" + content = b"line1\nline2\nline3\n" + file.write_bytes(content) + lines = list(FileLogStorage._read_lines_reversed(file)) + assert lines == [ + (b"", 18), + (b"line3", 12), + (b"line2", 6), + (b"line1", 0), + ] + + def test_file_without_trailing_newline(self, tmp_path: Path): + file = tmp_path / "test.txt" + content = b"line1\nline2" + file.write_bytes(content) + lines = list(FileLogStorage._read_lines_reversed(file)) + assert lines == [ + (b"line2", 6), + (b"line1", 0), + ] + + def test_empty_file(self, tmp_path: Path): + file = tmp_path / "test.txt" + file.touch() + lines = list(FileLogStorage._read_lines_reversed(file)) + assert lines == [] + + def test_single_line_file(self, tmp_path: Path): + file = tmp_path / "test.txt" + content = b"the only line" + file.write_bytes(content) + lines = list(FileLogStorage._read_lines_reversed(file)) + assert lines == [(b"the only line", 0)] + + def test_file_with_empty_lines(self, tmp_path: Path): + file = tmp_path / "test.txt" + content = b"lineA\n\nlineC\n" + file.write_bytes(content) + lines = list(FileLogStorage._read_lines_reversed(file)) + assert lines == [ + (b"", 13), + (b"lineC", 7), + (b"", 6), + (b"lineA", 0), + ] + + def test_file_with_only_newlines(self, tmp_path: Path): + file = tmp_path / "test.txt" + content = b"\n\n" + file.write_bytes(content) + lines = list(FileLogStorage._read_lines_reversed(file)) + assert lines == [ + (b"", 2), + (b"", 1), + ] + + def test_large_file_spanning_multiple_chunks(self, tmp_path: Path): + file = tmp_path / "large_file.txt" + line_content = b"abcdefghi" # 9 bytes + 1 newline = 10 bytes per line + num_lines = 5 + content = b"\n".join([line_content] * num_lines) + file.write_bytes(content) + # Pass the small chunk_size directly to the method + lines = list(FileLogStorage._read_lines_reversed(file, chunk_size=10)) + assert len(lines) == num_lines + assert lines[0] == (line_content, 40) + assert lines[1] == (line_content, 30) + assert lines[2] == (line_content, 20) + assert lines[3] == (line_content, 10) + assert lines[4] == (line_content, 0) + + # The rest of the tests will now pass without modification + def test_start_offset_in_middle_of_line(self, tmp_path: Path): + file = tmp_path / "test.txt" + content = b"line1\nline2\nline3\n" + file.write_bytes(content) + lines = list(FileLogStorage._read_lines_reversed(file, start_offset=10)) + assert lines == [ + (b"line2", 6), + (b"line1", 0), + ] + + def test_start_offset_at_line_boundary(self, tmp_path: Path): + file = tmp_path / "test.txt" + content = b"line1\nline2\nline3\n" + file.write_bytes(content) + lines = list(FileLogStorage._read_lines_reversed(file, start_offset=12)) + assert lines == [ + (b"line2", 6), + (b"line1", 0), + ] + + def test_start_offset_zero(self, tmp_path: Path): + file = tmp_path / "test.txt" + content = b"line1\nline2\nline3\n" + file.write_bytes(content) + lines = list(FileLogStorage._read_lines_reversed(file, start_offset=0)) + assert lines == [] + + def test_start_offset_larger_than_file(self, tmp_path: Path): + file = tmp_path / "test.txt" + content = b"line1\nline2\n" + file.write_bytes(content) + lines_with_offset = list(FileLogStorage._read_lines_reversed(file, start_offset=1000)) + lines_without_offset = list(FileLogStorage._read_lines_reversed(file)) + assert lines_with_offset == lines_without_offset + assert lines_with_offset == [(b"", 12), (b"line2", 6), (b"line1", 0)] + + def test_long_line_larger_than_chunk(self, tmp_path: Path): + file = tmp_path / "long_line.txt" + content = b"a" * 25 + file.write_bytes(content) + # Pass the small chunk_size directly + lines = list(FileLogStorage._read_lines_reversed(file, chunk_size=10)) + assert lines == [(content, 0)] diff --git a/src/tests/_internal/server/services/test_fluentbit_logs.py b/src/tests/_internal/server/services/logs/test_fluentbit.py similarity index 100% rename from src/tests/_internal/server/services/test_fluentbit_logs.py rename to src/tests/_internal/server/services/logs/test_fluentbit.py diff --git a/src/tests/_internal/server/services/test_backend_configs.py b/src/tests/_internal/server/services/test_backend_configs.py deleted file mode 100644 index 6558cc1358..0000000000 --- a/src/tests/_internal/server/services/test_backend_configs.py +++ /dev/null @@ -1,240 +0,0 @@ -import json -from pathlib import Path -from textwrap import dedent -from typing import cast -from unittest.mock import patch - -import yaml - -from dstack._internal.core.backends.kubernetes.backend import KubernetesBackend -from dstack._internal.core.backends.kubernetes.models import ( - KubernetesBackendConfigWithCreds, - KubernetesConfig, -) -from dstack._internal.server import settings -from dstack._internal.server.services.config import ( - ServerConfigManager, - config_yaml_to_backend_config, - file_config_to_config, -) - - -class TestCrusoeBackendConfig: - def test_config_parsing(self, tmp_path: Path): - config_yaml_path = tmp_path / "config.yml" - config_dict = { - "projects": [ - { - "name": "main", - "backends": [ - { - "type": "crusoe", - "project_id": "test-project-id", - "regions": ["us-east1-a"], - "creds": { - "type": "access_key", - "access_key": "test-access-key", - "secret_key": "test-secret-key", - }, - } - ], - } - ] - } - config_yaml_path.write_text(yaml.dump(config_dict)) - - with patch.object(settings, "SERVER_DIR_PATH", tmp_path): - m = ServerConfigManager() - assert m.load_config() - assert m.config is not None - assert m.config.projects is not None - assert len(m.config.projects) > 0 - assert m.config.projects[0].backends is not None - backend_file_cfg = m.config.projects[0].backends[0] - backend_cfg = file_config_to_config(backend_file_cfg) - - assert backend_cfg.type == "crusoe" - assert backend_cfg.project_id == "test-project-id" - assert backend_cfg.regions == ["us-east1-a"] - assert backend_cfg.creds.access_key == "test-access-key" - assert backend_cfg.creds.secret_key == "test-secret-key" - - -class TestNebiusBackendConfig: - def test_with_filename(self, tmp_path: Path): - creds_json = { - "subject-credentials": { - "type": "JWT", - "alg": "RS256", - "private-key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n", - "kid": "publickey-e00test", - "iss": "serviceaccount-e00test", - "sub": "serviceaccount-e00test", - } - } - creds_file = tmp_path / "nebius_creds.json" - creds_file.write_text(json.dumps(creds_json)) - - config_yaml_path = tmp_path / "config.yml" - config_dict = { - "projects": [ - { - "name": "main", - "backends": [ - { - "type": "nebius", - "creds": {"type": "service_account", "filename": str(creds_file)}, - } - ], - } - ] - } - config_yaml_path.write_text(yaml.dump(config_dict)) - - with patch.object(settings, "SERVER_DIR_PATH", tmp_path): - m = ServerConfigManager() - assert m.load_config() - assert m.config is not None - assert m.config.projects is not None - assert len(m.config.projects) > 0 - assert m.config.projects[0].backends is not None - backend_file_cfg = m.config.projects[0].backends[0] - backend_cfg = file_config_to_config(backend_file_cfg) - - assert backend_cfg.type == "nebius" - assert backend_cfg.creds.service_account_id == "serviceaccount-e00test" - assert backend_cfg.creds.public_key_id == "publickey-e00test" - assert ( - backend_cfg.creds.private_key_content - == "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n" - ) - - def test_with_private_key_file(self, tmp_path: Path): - pk_file = tmp_path / "private.key" - pk_file.write_text("TEST_PRIVATE_KEY") - - config_yaml_path = tmp_path / "config.yml" - config_dict = { - "projects": [ - { - "name": "main", - "backends": [ - { - "type": "nebius", - "projects": ["project-e00test"], - "creds": { - "type": "service_account", - "service_account_id": "serviceaccount-e00test", - "public_key_id": "publickey-e00test", - "private_key_file": str(pk_file), - }, - } - ], - } - ] - } - config_yaml_path.write_text(yaml.dump(config_dict)) - - with patch.object(settings, "SERVER_DIR_PATH", tmp_path): - m = ServerConfigManager() - assert m.load_config() - assert m.config is not None - assert m.config.projects is not None - assert len(m.config.projects) > 0 - assert m.config.projects[0].backends is not None - backend_file_cfg = m.config.projects[0].backends[0] - backend_cfg = file_config_to_config(backend_file_cfg) - - assert backend_cfg.type == "nebius" - assert backend_cfg.creds.service_account_id == "serviceaccount-e00test" - assert backend_cfg.creds.public_key_id == "publickey-e00test" - assert backend_cfg.creds.private_key_content == "TEST_PRIVATE_KEY" - - -class TestKubernetesBackendConfig: - def test_ui_config_embedded_kubeconfig_initializes_backend(self): - config_yaml = dedent( - """ - type: kubernetes - kubeconfig: - data: | - apiVersion: v1 - kind: Config - current-context: gpu-training - - clusters: - - name: gpu-training - cluster: - server: https://gpu-cluster.internal.example.com:6443 - insecure-skip-tls-verify: true - - users: - - name: ml-engineer - user: - token: test-token - - contexts: - - name: gpu-training - context: - cluster: gpu-training - user: ml-engineer - - proxy_jump: - hostname: 204.12.171.137 - port: 32000 - """ - ) - - backend_config = config_yaml_to_backend_config(config_yaml) - assert isinstance(backend_config, KubernetesBackendConfigWithCreds) - # `KubernetesConfig` adds no fields to the parsed config - backend = KubernetesBackend(cast(KubernetesConfig, backend_config)) - - cluster = backend.compute().region_cluster_map[""] - assert cluster.api_client.configuration.host == ( # pyright: ignore[reportAttributeAccessIssue] - "https://gpu-cluster.internal.example.com:6443" - ) - assert cluster.proxy_jump.hostname == "204.12.171.137" - assert cluster.proxy_jump.port == 32000 - - def test_kubeconfig_context_namespace_does_not_set_backend_namespace(self): - config_yaml = dedent( - """ - type: kubernetes - kubeconfig: - data: | - apiVersion: v1 - kind: Config - current-context: gpu-training - - clusters: - - name: gpu-training - cluster: - server: https://gpu-cluster.internal.example.com:6443 - insecure-skip-tls-verify: true - - users: - - name: ml-engineer - user: - token: test-token - - contexts: - - name: gpu-training - context: - cluster: gpu-training - user: ml-engineer - namespace: training-jobs - - proxy_jump: - hostname: 204.12.171.137 - port: 32000 - """ - ) - - backend_config = config_yaml_to_backend_config(config_yaml) - assert isinstance(backend_config, KubernetesBackendConfigWithCreds) - # `KubernetesConfig` adds no fields to the parsed config - backend = KubernetesBackend(cast(KubernetesConfig, backend_config)) - - cluster = backend.compute().region_cluster_map[""] - assert cluster.namespace == "default" diff --git a/src/tests/_internal/server/services/test_config.py b/src/tests/_internal/server/services/test_config.py index 6d09ce631a..bc07da3b57 100644 --- a/src/tests/_internal/server/services/test_config.py +++ b/src/tests/_internal/server/services/test_config.py @@ -1,5 +1,7 @@ import json from pathlib import Path +from textwrap import dedent +from typing import cast from unittest.mock import AsyncMock, Mock, patch import pytest @@ -8,9 +10,18 @@ from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.core.backends.aws.configurator import DEFAULT_REGIONS +from dstack._internal.core.backends.kubernetes.backend import KubernetesBackend +from dstack._internal.core.backends.kubernetes.models import ( + KubernetesBackendConfigWithCreds, + KubernetesConfig, +) from dstack._internal.server import settings from dstack._internal.server.models import BackendModel, ImportModel, ProjectModel -from dstack._internal.server.services.config import ServerConfigManager +from dstack._internal.server.services.config import ( + ServerConfigManager, + config_yaml_to_backend_config, + file_config_to_config, +) from dstack._internal.server.testing.common import ( create_backend, create_export, @@ -247,3 +258,224 @@ async def test_new_project_imports_global_exports( ) imports = imports_res.scalars().all() assert len(imports) == 1 + + +class TestCrusoeBackendConfig: + def test_config_parsing(self, tmp_path: Path): + config_yaml_path = tmp_path / "config.yml" + config_dict = { + "projects": [ + { + "name": "main", + "backends": [ + { + "type": "crusoe", + "project_id": "test-project-id", + "regions": ["us-east1-a"], + "creds": { + "type": "access_key", + "access_key": "test-access-key", + "secret_key": "test-secret-key", + }, + } + ], + } + ] + } + config_yaml_path.write_text(yaml.dump(config_dict)) + + with patch.object(settings, "SERVER_DIR_PATH", tmp_path): + m = ServerConfigManager() + assert m.load_config() + assert m.config is not None + assert m.config.projects is not None + assert len(m.config.projects) > 0 + assert m.config.projects[0].backends is not None + backend_file_cfg = m.config.projects[0].backends[0] + backend_cfg = file_config_to_config(backend_file_cfg) + + assert backend_cfg.type == "crusoe" + assert backend_cfg.project_id == "test-project-id" + assert backend_cfg.regions == ["us-east1-a"] + assert backend_cfg.creds.access_key == "test-access-key" + assert backend_cfg.creds.secret_key == "test-secret-key" + + +class TestNebiusBackendConfig: + def test_with_filename(self, tmp_path: Path): + creds_json = { + "subject-credentials": { + "type": "JWT", + "alg": "RS256", + "private-key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n", + "kid": "publickey-e00test", + "iss": "serviceaccount-e00test", + "sub": "serviceaccount-e00test", + } + } + creds_file = tmp_path / "nebius_creds.json" + creds_file.write_text(json.dumps(creds_json)) + + config_yaml_path = tmp_path / "config.yml" + config_dict = { + "projects": [ + { + "name": "main", + "backends": [ + { + "type": "nebius", + "creds": {"type": "service_account", "filename": str(creds_file)}, + } + ], + } + ] + } + config_yaml_path.write_text(yaml.dump(config_dict)) + + with patch.object(settings, "SERVER_DIR_PATH", tmp_path): + m = ServerConfigManager() + assert m.load_config() + assert m.config is not None + assert m.config.projects is not None + assert len(m.config.projects) > 0 + assert m.config.projects[0].backends is not None + backend_file_cfg = m.config.projects[0].backends[0] + backend_cfg = file_config_to_config(backend_file_cfg) + + assert backend_cfg.type == "nebius" + assert backend_cfg.creds.service_account_id == "serviceaccount-e00test" + assert backend_cfg.creds.public_key_id == "publickey-e00test" + assert ( + backend_cfg.creds.private_key_content + == "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n" + ) + + def test_with_private_key_file(self, tmp_path: Path): + pk_file = tmp_path / "private.key" + pk_file.write_text("TEST_PRIVATE_KEY") + + config_yaml_path = tmp_path / "config.yml" + config_dict = { + "projects": [ + { + "name": "main", + "backends": [ + { + "type": "nebius", + "projects": ["project-e00test"], + "creds": { + "type": "service_account", + "service_account_id": "serviceaccount-e00test", + "public_key_id": "publickey-e00test", + "private_key_file": str(pk_file), + }, + } + ], + } + ] + } + config_yaml_path.write_text(yaml.dump(config_dict)) + + with patch.object(settings, "SERVER_DIR_PATH", tmp_path): + m = ServerConfigManager() + assert m.load_config() + assert m.config is not None + assert m.config.projects is not None + assert len(m.config.projects) > 0 + assert m.config.projects[0].backends is not None + backend_file_cfg = m.config.projects[0].backends[0] + backend_cfg = file_config_to_config(backend_file_cfg) + + assert backend_cfg.type == "nebius" + assert backend_cfg.creds.service_account_id == "serviceaccount-e00test" + assert backend_cfg.creds.public_key_id == "publickey-e00test" + assert backend_cfg.creds.private_key_content == "TEST_PRIVATE_KEY" + + +class TestKubernetesBackendConfig: + def test_ui_config_embedded_kubeconfig_initializes_backend(self): + config_yaml = dedent( + """ + type: kubernetes + kubeconfig: + data: | + apiVersion: v1 + kind: Config + current-context: gpu-training + + clusters: + - name: gpu-training + cluster: + server: https://gpu-cluster.internal.example.com:6443 + insecure-skip-tls-verify: true + + users: + - name: ml-engineer + user: + token: test-token + + contexts: + - name: gpu-training + context: + cluster: gpu-training + user: ml-engineer + + proxy_jump: + hostname: 204.12.171.137 + port: 32000 + """ + ) + + backend_config = config_yaml_to_backend_config(config_yaml) + assert isinstance(backend_config, KubernetesBackendConfigWithCreds) + # `KubernetesConfig` adds no fields to the parsed config + backend = KubernetesBackend(cast(KubernetesConfig, backend_config)) + + cluster = backend.compute().region_cluster_map[""] + assert cluster.api_client.configuration.host == ( # pyright: ignore[reportAttributeAccessIssue] + "https://gpu-cluster.internal.example.com:6443" + ) + assert cluster.proxy_jump.hostname == "204.12.171.137" + assert cluster.proxy_jump.port == 32000 + + def test_kubeconfig_context_namespace_does_not_set_backend_namespace(self): + config_yaml = dedent( + """ + type: kubernetes + kubeconfig: + data: | + apiVersion: v1 + kind: Config + current-context: gpu-training + + clusters: + - name: gpu-training + cluster: + server: https://gpu-cluster.internal.example.com:6443 + insecure-skip-tls-verify: true + + users: + - name: ml-engineer + user: + token: test-token + + contexts: + - name: gpu-training + context: + cluster: gpu-training + user: ml-engineer + namespace: training-jobs + + proxy_jump: + hostname: 204.12.171.137 + port: 32000 + """ + ) + + backend_config = config_yaml_to_backend_config(config_yaml) + assert isinstance(backend_config, KubernetesBackendConfigWithCreds) + # `KubernetesConfig` adds no fields to the parsed config + backend = KubernetesBackend(cast(KubernetesConfig, backend_config)) + + cluster = backend.compute().region_cluster_map[""] + assert cluster.namespace == "default" From e3e9801553e5db8184a09980d9baf085e9490b77 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Wed, 28 Oct 2026 14:57:29 +0500 Subject: [PATCH 3/3] Name pipeline package tests test_pipeline.py --- .../instances/{test_instances.py => test_pipeline.py} | 0 .../pipeline_tasks/runs/{test_runs.py => test_pipeline.py} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename src/tests/_internal/server/background/pipeline_tasks/instances/{test_instances.py => test_pipeline.py} (100%) rename src/tests/_internal/server/background/pipeline_tasks/runs/{test_runs.py => test_pipeline.py} (100%) diff --git a/src/tests/_internal/server/background/pipeline_tasks/instances/test_instances.py b/src/tests/_internal/server/background/pipeline_tasks/instances/test_pipeline.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/instances/test_instances.py rename to src/tests/_internal/server/background/pipeline_tasks/instances/test_pipeline.py diff --git a/src/tests/_internal/server/background/pipeline_tasks/runs/test_runs.py b/src/tests/_internal/server/background/pipeline_tasks/runs/test_pipeline.py similarity index 100% rename from src/tests/_internal/server/background/pipeline_tasks/runs/test_runs.py rename to src/tests/_internal/server/background/pipeline_tasks/runs/test_pipeline.py