diff --git a/AGENTS.md b/AGENTS.md index 17adfe8ad..bc935dadb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -32,6 +32,7 @@ Before touching a subsystem, read the relevant notes in `contributing/`: `ARCHIT - Never make network calls inside a DB session or transaction. Fetch what you need before opening the session, or commit and close it before the call. - Don't use function-level (inner) imports to break circular imports. Inject the dependency or move the shared code to a lower-level module instead. - Never edit a migration that has already been applied or released; add a new migration instead. +- Derive paths under `SERVER_DIR_PATH` on access (a `get_*` function), not as module-level constants, so that patching `settings.SERVER_DIR_PATH` redirects all of them. Tests rely on this to keep server state out of the real `~/.dstack`. ## Testing Guidelines - Default to `uv run pytest`. Use markers from `src/tests/conftest.py` like `--runpostgres` if need to include specific tests. diff --git a/src/dstack/_internal/server/app.py b/src/dstack/_internal/server/app.py index 6380961e3..c3e0fab1a 100644 --- a/src/dstack/_internal/server/app.py +++ b/src/dstack/_internal/server/app.py @@ -64,9 +64,10 @@ from dstack._internal.server.settings import ( DEFAULT_PROJECT_NAME, DO_NOT_UPDATE_DEFAULT_PROJECT, - SERVER_CONFIG_FILE_PATH, SERVER_URL, UPDATE_DEFAULT_PROJECT, + get_server_config_file_path, + init_server_data_dir, ) from dstack._internal.server.utils import otel, sentry_utils from dstack._internal.server.utils.logging import configure_logging @@ -118,6 +119,7 @@ async def lifespan(app: FastAPI): ) server_executor = ThreadPoolExecutor(max_workers=settings.SERVER_EXECUTOR_MAX_WORKERS) asyncio.get_running_loop().set_default_executor(server_executor) + init_server_data_dir() await migrate() _print_dstack_logo() if not check_required_ssh_version(): @@ -141,17 +143,18 @@ async def lifespan(app: FastAPI): user=admin, ) if server_config_manager is not None: + server_config_file_path = get_server_config_file_path() server_config_dir = _get_server_config_dir() if not server_config_loaded: logger.info("Initializing the default configuration...", {"show_path": False}) await server_config_manager.init_config(session=session) logger.info( - f"Initialized the default configuration at [link=file://{SERVER_CONFIG_FILE_PATH}]{server_config_dir}[/link]", + f"Initialized the default configuration at [link=file://{server_config_file_path}]{server_config_dir}[/link]", {"show_path": False}, ) else: logger.info( - f"Applying [link=file://{SERVER_CONFIG_FILE_PATH}]{server_config_dir}[/link]...", + f"Applying [link=file://{server_config_file_path}]{server_config_dir}[/link]...", {"show_path": False}, ) await server_config_manager.apply_config(session=session, owner=admin) @@ -435,4 +438,4 @@ def _print_dstack_logo(): def _get_server_config_dir() -> str: - return str(SERVER_CONFIG_FILE_PATH).replace(os.path.expanduser("~"), "~", 1) + return str(get_server_config_file_path()).replace(os.path.expanduser("~"), "~", 1) diff --git a/src/dstack/_internal/server/db.py b/src/dstack/_internal/server/db.py index 2eb18a3f3..3136b4e6b 100644 --- a/src/dstack/_internal/server/db.py +++ b/src/dstack/_internal/server/db.py @@ -14,7 +14,6 @@ from dstack._internal.server import settings from dstack._internal.server.services.locking import advisory_lock_ctx -from dstack._internal.server.settings import DATABASE_URL class Database: @@ -63,13 +62,16 @@ def get_new_db() -> Database: Use this when you need to access the DB in a new thread instead of calling Database directly since it's easier to monkey-patch. """ - return Database(url=DATABASE_URL) + return Database(url=settings.get_database_url()) -_db = get_new_db() +_db: Optional[Database] = None def get_db() -> Database: + global _db + if _db is None: + _db = get_new_db() return _db @@ -79,17 +81,18 @@ def override_db(new_db: Database): async def migrate(): - async with _db.engine.connect() as connection: + db = get_db() + async with db.engine.connect() as connection: async with advisory_lock_ctx( bind=connection, - dialect_name=_db.dialect_name, + dialect_name=db.dialect_name, resource="migrations", ): await connection.run_sync(_run_alembic_upgrade) async def get_session(): - async with _db.get_session() as session: + async with get_db().get_session() as session: yield session await session.commit() diff --git a/src/dstack/_internal/server/migrations/env.py b/src/dstack/_internal/server/migrations/env.py index c7c27f1f8..0ab4b713a 100644 --- a/src/dstack/_internal/server/migrations/env.py +++ b/src/dstack/_internal/server/migrations/env.py @@ -7,6 +7,7 @@ from dstack._internal.server.db import get_db from dstack._internal.server.models import BaseModel, EnumAsString +from dstack._internal.server.settings import init_server_data_dir config = context.config @@ -94,6 +95,7 @@ def run_migrations(connection: Connection): async def run_async_migrations(): + init_server_data_dir() engine = get_db().engine async with engine.connect() as connection: await connection.run_sync(run_migrations) diff --git a/src/dstack/_internal/server/services/config.py b/src/dstack/_internal/server/services/config.py index 81f677638..bba6e5d1f 100644 --- a/src/dstack/_internal/server/services/config.py +++ b/src/dstack/_internal/server/services/config.py @@ -217,7 +217,7 @@ async def _init_config(self, session: AsyncSession) -> Optional[ServerConfig]: def _load_config(self) -> Optional[ServerConfig]: try: - with open(settings.SERVER_CONFIG_FILE_PATH) as f: + with open(settings.get_server_config_file_path()) as f: content = f.read() except OSError: return @@ -225,7 +225,7 @@ def _load_config(self) -> Optional[ServerConfig]: return ServerConfig.model_validate(config_dict) def _save_config(self, config: ServerConfig): - with open(settings.SERVER_CONFIG_FILE_PATH, "w+") as f: + with open(settings.get_server_config_file_path(), "w+") as f: f.write(config_to_yaml(config)) diff --git a/src/dstack/_internal/server/services/gateways/connection.py b/src/dstack/_internal/server/services/gateways/connection.py index dada5bea6..fe8187188 100644 --- a/src/dstack/_internal/server/services/gateways/connection.py +++ b/src/dstack/_internal/server/services/gateways/connection.py @@ -1,6 +1,7 @@ import contextlib import shutil import uuid +from pathlib import Path from typing import AsyncIterator, Optional import aiorwlock @@ -17,13 +18,16 @@ SERVER_CONNECTIONS_DIR_ON_GATEWAY, ) from dstack._internal.proxy.gateway.schemas.stats import PerWindowStats +from dstack._internal.server import settings from dstack._internal.server.services.gateways.client import GatewayClient -from dstack._internal.server.settings import SERVER_DIR_PATH from dstack._internal.utils.logging import get_logger from dstack._internal.utils.path import FileContent, make_tmp_symlink_to_dir logger = get_logger(__name__) -CONNECTIONS_DIR = SERVER_DIR_PATH / "gateway-connections" + + +def _get_connections_dir() -> Path: + return settings.SERVER_DIR_PATH / "gateway-connections" class GatewayConnection: @@ -42,7 +46,7 @@ def __init__(self, ip_address: str, id_rsa: str, server_port: int): self.server_port = server_port # a persistent connection_dir is needed to discover and close leftover connections # in case of server restarts w/o graceful shutdown - self.connection_dir = CONNECTIONS_DIR / ip_address + self.connection_dir = _get_connections_dir() / ip_address # connection_dir can have a long path that won't be accepted by the ssh command, # so we create a short temporary symlink self.temp_dir, self.connection_symlink_dir = make_tmp_symlink_to_dir( diff --git a/src/dstack/_internal/server/services/jobs/server_connection.py b/src/dstack/_internal/server/services/jobs/server_connection.py index ccfed0482..047a53ffa 100644 --- a/src/dstack/_internal/server/services/jobs/server_connection.py +++ b/src/dstack/_internal/server/services/jobs/server_connection.py @@ -27,12 +27,15 @@ logger = get_logger(__name__) -CONNECTIONS_DIR = settings.SERVER_DIR_PATH / "job-server-connections" _MIN_ALIVE_CHECK_INTERVAL = 30 _PROBE_TIMEOUT = 3 _REMOTE_SOCKET_PATH = Path(DSTACK_RUN_SERVER_SOCKET_PATH) +def _get_connections_dir() -> Path: + return settings.SERVER_DIR_PATH / "job-server-connections" + + def _get_server_socket() -> IPSocket: # The server may be bound to a specific address, making loopback unreachable host = settings.SERVER_HOST @@ -51,7 +54,7 @@ def __init__(self, job: JobModel, job_runtime_data: Optional[JobRuntimeData]) -> self._last_verified_at = 0.0 # Keep the control socket discoverable across server process restarts. The temporary # symlink keeps its effective path below OpenSSH's Unix-socket length limit. - self._connection_dir = CONNECTIONS_DIR / str(job.id) + self._connection_dir = _get_connections_dir() / str(job.id) self._connection_dir.mkdir(parents=True, exist_ok=True) self._temp_dir, effective_dir = make_tmp_symlink_to_dir( self._connection_dir, @@ -207,7 +210,7 @@ async def remove(self, job_id: uuid.UUID) -> None: if connection is not None: await self._close(connection) self._failure_started_at.pop(job_id, None) - shutil.rmtree(CONNECTIONS_DIR / str(job_id), ignore_errors=True) + shutil.rmtree(_get_connections_dir() / str(job_id), ignore_errors=True) async def remove_all(self) -> None: job_ids = set(self._connections).union(self._failure_started_at) diff --git a/src/dstack/_internal/server/services/runner/pool.py b/src/dstack/_internal/server/services/runner/pool.py index e3a012967..93ba015f7 100644 --- a/src/dstack/_internal/server/services/runner/pool.py +++ b/src/dstack/_internal/server/services/runner/pool.py @@ -19,10 +19,7 @@ SSHTunnel, UnixSocket, ) -from dstack._internal.server.settings import ( - SERVER_DIR_PATH, - SERVER_SSH_CONNECT_TIMEOUT, -) +from dstack._internal.server import settings from dstack._internal.utils.logging import get_logger from dstack._internal.utils.path import FileContent, make_tmp_symlink_to_dir @@ -31,12 +28,14 @@ PrivateKeyOrPair = Union[str, tuple[str, Optional[str]]] """A host private key or pair of (host private key, optional proxy jump private key)""" -CONNECTIONS_DIR = SERVER_DIR_PATH / "instance-connections" - MIN_ALIVE_CHECK_INTERVAL = 30 """How often (at most) `InstanceConnection.is_alive()` runs `ssh -O check`, in seconds.""" +def _get_connections_dir() -> Path: + return settings.SERVER_DIR_PATH / "instance-connections" + + @dataclass(frozen=True) class InstanceConnectionKey: hostname: str @@ -142,7 +141,7 @@ def startup_cleanup(self) -> None: Must be called on server startup before the pool is used. Leftover live masters are reaped by `ControlPersist`. """ - shutil.rmtree(CONNECTIONS_DIR, ignore_errors=True) + shutil.rmtree(_get_connections_dir(), ignore_errors=True) def close_all(self) -> None: """ @@ -218,7 +217,7 @@ def __init__( ssh_proxies=InstanceConnection._get_proxies(ssh_private_key, jpd), options={ **SSH_DEFAULT_OPTIONS, - "ConnectTimeout": str(SERVER_SSH_CONNECT_TIMEOUT), + "ConnectTimeout": str(settings.SERVER_SSH_CONNECT_TIMEOUT), # Auto-close half-opened connections (the instance not responding). "ServerAliveInterval": "10", "ServerAliveCountMax": "3", @@ -276,7 +275,7 @@ def close(self) -> None: self._tunnel.close() # Remove a stale control.sock left by a killed master, forwarded UDS files # (ssh does not unlink them on exit), and the dir itself, so that - # CONNECTIONS_DIR does not accumulate dirs of gone instances. + # the connections dir does not accumulate dirs of gone instances. # A master that survives close() because it is unreachable via a deleted # symlink is reaped by ControlPersist. shutil.rmtree(self._real_conn_dir, ignore_errors=True) @@ -311,7 +310,7 @@ def _resolve_conn_dir( return temp_dir, path, path conn_dir = ( - CONNECTIONS_DIR + _get_connections_dir() / f"{key.hostname}:{key.port},{','.join(map(str, key.ports_to_forward))}" ) conn_dir.mkdir(parents=True, exist_ok=True) diff --git a/src/dstack/_internal/server/services/templates.py b/src/dstack/_internal/server/services/templates.py index 752f45a49..1f633b486 100644 --- a/src/dstack/_internal/server/services/templates.py +++ b/src/dstack/_internal/server/services/templates.py @@ -48,7 +48,7 @@ def _list_templates_sync(repo_key: str, repo_url: str) -> List[UITemplate]: def _fetch_templates_repo(repo_key: str, repo_url: str) -> Path: - repo_dir = settings.SERVER_DATA_DIR_PATH / "templates-repos" / repo_key + repo_dir = settings.get_server_data_dir_path() / "templates-repos" / repo_key if repo_dir.exists(): try: repo = git.Repo(str(repo_dir)) diff --git a/src/dstack/_internal/server/settings.py b/src/dstack/_internal/server/settings.py index b7be0bf3b..360572535 100644 --- a/src/dstack/_internal/server/settings.py +++ b/src/dstack/_internal/server/settings.py @@ -16,14 +16,38 @@ SERVER_DIR_PATH = Path(os.getenv("DSTACK_SERVER_DIR", DSTACK_DIR_PATH / "server")).resolve() -SERVER_CONFIG_FILE_PATH = SERVER_DIR_PATH / "config.yml" -SERVER_DATA_DIR_PATH = SERVER_DIR_PATH / "data" -SERVER_DATA_DIR_PATH.mkdir(parents=True, exist_ok=True) +# Paths under `SERVER_DIR_PATH` are derived on access rather than at import time, so that +# patching `SERVER_DIR_PATH` redirects all of them. Tests rely on this to keep each worker's +# server state out of the real `~/.dstack`. +# +# TODO: Turn module level-constants into a ServerSettings class instance so that +# all settings can be properties and there is no constant/function distinction. + + +def get_server_config_file_path() -> Path: + return SERVER_DIR_PATH / "config.yml" + + +def get_server_data_dir_path() -> Path: + return SERVER_DIR_PATH / "data" + + +def init_server_data_dir() -> Path: + """ + Creates the server data dir. Call before connecting to the default SQLite database. + """ + data_dir = get_server_data_dir_path() + data_dir.mkdir(parents=True, exist_ok=True) + return data_dir + + +def get_database_url() -> str: + return os.getenv( + "DSTACK_DATABASE_URL", + f"sqlite+aiosqlite:///{get_server_data_dir_path()}/sqlite.db", + ) -DATABASE_URL = os.getenv( - "DSTACK_DATABASE_URL", f"sqlite+aiosqlite:///{str(SERVER_DATA_DIR_PATH.absolute())}/sqlite.db" -) SERVER_HOST = os.getenv("DSTACK_SERVER_HOST", "localhost") SERVER_PORT = int(os.getenv("DSTACK_SERVER_PORT", "8000")) diff --git a/src/dstack/_internal/server/testing/conf.py b/src/dstack/_internal/server/testing/conf.py index 19aa51d53..054c5c840 100644 --- a/src/dstack/_internal/server/testing/conf.py +++ b/src/dstack/_internal/server/testing/conf.py @@ -11,6 +11,25 @@ SQLITE_URL = "sqlite+aiosqlite://" +@pytest.fixture(scope="session", autouse=True) +def server_dir(tmp_path_factory: pytest.TempPathFactory): + """ + Points the server dir at a tmp dir private to this pytest process. + + The real `~/.dstack/server` is shared by every process on the machine, including a + developer's running server and, under `pytest -n auto`, the other workers. Sharing it + makes unrelated tests race over the same connection dirs and SSH control sockets. + + Everything the server derives from `SERVER_DIR_PATH` is resolved on access, so patching + it here redirects all of it. `DSTACK_SERVER_DIR` is set too, for subprocesses. + """ + path = tmp_path_factory.mktemp("server-dir") + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setattr(settings, "SERVER_DIR_PATH", path) + monkeypatch.setenv("DSTACK_SERVER_DIR", str(path)) + yield path + + @pytest.fixture(scope="session") def postgres_container(): with PostgresContainer( diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index ffb7d1802..034dbf7c2 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -1073,7 +1073,6 @@ async def test_provisioning_server_access_failure_terminates_job_after_retry_tim test_db, session: AsyncSession, worker: JobRunningWorker, - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ): project = await create_project(session=session) @@ -1107,7 +1106,6 @@ async def test_provisioning_server_access_failure_terminates_job_after_retry_tim instance_assigned=True, ) last_processed_at = job.last_processed_at - monkeypatch.setattr(server_connection, "CONNECTIONS_DIR", tmp_path) failing_connection = MagicMock() failing_connection.job_id = job.id failing_connection.open = AsyncMock(side_effect=SSHError("cannot open tunnel")) diff --git a/src/tests/_internal/server/conftest.py b/src/tests/_internal/server/conftest.py index 125cc5de1..5275b6be0 100644 --- a/src/tests/_internal/server/conftest.py +++ b/src/tests/_internal/server/conftest.py @@ -12,6 +12,7 @@ from dstack._internal.server.services.logs.filelog import FileLogStorage from dstack._internal.server.testing.conf import ( # noqa: F401 postgres_container, + server_dir, session, test_db, ) diff --git a/src/tests/_internal/server/services/jobs/test_server_connection.py b/src/tests/_internal/server/services/jobs/test_server_connection.py index fb855f40d..260feb893 100644 --- a/src/tests/_internal/server/services/jobs/test_server_connection.py +++ b/src/tests/_internal/server/services/jobs/test_server_connection.py @@ -17,8 +17,7 @@ @pytest.fixture -def tunnel_mock(tmp_path, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(server_connection, "CONNECTIONS_DIR", tmp_path) +def tunnel_mock(monkeypatch: pytest.MonkeyPatch): tunnel = MagicMock() tunnel.forwarded_sockets = [] tunnel.reverse_forwarded_sockets = [] diff --git a/src/tests/_internal/server/services/test_backend_configs.py b/src/tests/_internal/server/services/test_backend_configs.py index d99bdcb98..5b2c15e1b 100644 --- a/src/tests/_internal/server/services/test_backend_configs.py +++ b/src/tests/_internal/server/services/test_backend_configs.py @@ -38,7 +38,7 @@ def test_config_parsing(self, tmp_path: Path): } config_yaml_path.write_text(yaml.dump(config_dict)) - with patch.object(settings, "SERVER_CONFIG_FILE_PATH", config_yaml_path): + with patch.object(settings, "SERVER_DIR_PATH", tmp_path): m = ServerConfigManager() assert m.load_config() assert m.config is not None @@ -86,7 +86,7 @@ def test_with_filename(self, tmp_path: Path): } config_yaml_path.write_text(yaml.dump(config_dict)) - with patch.object(settings, "SERVER_CONFIG_FILE_PATH", config_yaml_path): + with patch.object(settings, "SERVER_DIR_PATH", tmp_path): m = ServerConfigManager() assert m.load_config() assert m.config is not None @@ -130,7 +130,7 @@ def test_with_private_key_file(self, tmp_path: Path): } config_yaml_path.write_text(yaml.dump(config_dict)) - with patch.object(settings, "SERVER_CONFIG_FILE_PATH", config_yaml_path): + with patch.object(settings, "SERVER_DIR_PATH", tmp_path): m = ServerConfigManager() assert m.load_config() assert m.config is not None diff --git a/src/tests/_internal/server/services/test_config.py b/src/tests/_internal/server/services/test_config.py index 81265445e..6d09ce631 100644 --- a/src/tests/_internal/server/services/test_config.py +++ b/src/tests/_internal/server/services/test_config.py @@ -63,7 +63,7 @@ async def test_creates_backend(self, test_db, session: AsyncSession, tmp_path: P yaml.dump(config, f) with ( patch("boto3.session.Session"), - patch.object(settings, "SERVER_CONFIG_FILE_PATH", config_filepath), + patch.object(settings, "SERVER_DIR_PATH", tmp_path), patch( "dstack._internal.core.backends.aws.compute.get_vpc_id_subnets_ids_or_error" ), @@ -105,7 +105,7 @@ async def test_skips_update_when_source_config_matches( with open(config_filepath, "w+") as f: yaml.dump(config, f) with ( - patch.object(settings, "SERVER_CONFIG_FILE_PATH", config_filepath), + patch.object(settings, "SERVER_DIR_PATH", tmp_path), patch( "dstack._internal.server.services.backends.update_backend", new_callable=AsyncMock, @@ -143,7 +143,7 @@ async def test_populates_source_config_for_legacy_backend( mock_session = Mock() mock_session.client.return_value = Mock() with ( - patch.object(settings, "SERVER_CONFIG_FILE_PATH", config_filepath), + patch.object(settings, "SERVER_DIR_PATH", tmp_path), patch( "dstack._internal.core.backends.aws.auth.authenticate", return_value=mock_session, @@ -161,7 +161,7 @@ async def test_populates_source_config_for_legacy_backend( assert json.loads(backend.source_config)["regions"] is None assert json.loads(backend.source_auth.get_plaintext_or_error()) == creds with ( - patch.object(settings, "SERVER_CONFIG_FILE_PATH", config_filepath), + patch.object(settings, "SERVER_DIR_PATH", tmp_path), patch( "dstack._internal.server.services.backends.update_backend", new_callable=AsyncMock, @@ -199,7 +199,7 @@ async def test_forces_update_when_current_backend_config_is_unavailable( with open(config_filepath, "w+") as f: yaml.dump(config, f) with ( - patch.object(settings, "SERVER_CONFIG_FILE_PATH", config_filepath), + patch.object(settings, "SERVER_DIR_PATH", tmp_path), patch( "dstack._internal.server.services.backends.get_backend_config", new_callable=AsyncMock, @@ -234,7 +234,7 @@ async def test_new_project_imports_global_exports( config = {"projects": [{"name": "new-project"}]} with open(config_filepath, "w+") as f: yaml.dump(config, f) - with patch.object(settings, "SERVER_CONFIG_FILE_PATH", config_filepath): + with patch.object(settings, "SERVER_DIR_PATH", tmp_path): manager = ServerConfigManager() manager.load_config() await manager.apply_config(session, owner)