Disable gRPC fork support in the server

src/dstack/_internal/cli/commands/server.py

@@ -80,9 +80,6 @@ def _command(self, args: argparse.Namespace):
             os.environ["DSTACK_DO_NOT_UPDATE_DEFAULT_PROJECT"] = "1"
         if args.token:
             os.environ["DSTACK_SERVER_ADMIN_TOKEN"] = args.token
-        # Hide noisy "Other threads are currently calling into gRPC, skipping fork() handlers"
-        # messages in server logs. Users can still change this with GRPC_VERBOSITY.
-        os.environ.setdefault("GRPC_VERBOSITY", "ERROR")
         uvicorn_log_level = os.getenv("DSTACK_SERVER_UVICORN_LOG_LEVEL", "ERROR").lower()
         reload_disabled = os.getenv("DSTACK_SERVER_RELOAD_DISABLED") is not None
 

src/dstack/_internal/server/app.py

@@ -18,7 +18,7 @@
 from dstack._internal.core.services.configs import update_default_project
 from dstack._internal.proxy.lib.deps import get_injector_from_app
 from dstack._internal.proxy.lib.routers import model_proxy
-from dstack._internal.server import settings
+from dstack._internal.server import bootstrap, settings
 from dstack._internal.server.background.pipeline_tasks import start_pipeline_tasks
 from dstack._internal.server.background.scheduled_tasks import start_scheduled_tasks
 from dstack._internal.server.background.scheduled_tasks.probes import PROBES_SCHEDULER
@@ -105,6 +105,11 @@ def create_app() -> FastAPI:
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     configure_logging()
+    if not bootstrap.is_process_set_up():
+        logger.warning(
+            "dstack._internal.server.bootstrap.set_up_process() was not called."
+            " Server entry points must call it before importing other dstack modules."
+        )
     if settings.SENTRY_DSN is not None:
         sentry_sdk.init(
             dsn=settings.SENTRY_DSN,

src/dstack/_internal/server/bootstrap.py

@@ -0,0 +1,45 @@
+"""
+Process-wide setup that must run before the server modules are imported.
+
+Keep this module free of imports that could pull in the libraries it configures.
+"""
+
+import os
+import sys
+
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+_is_process_set_up = False
+
+
+def set_up_process() -> None:
+    """
+    Configures process-wide settings that must be applied before dstack server modules,
+    and the libraries they use, are imported. Server entry points, including those of
+    applications extending the server, must call it before importing anything else from dstack.
+    """
+    global _is_process_set_up
+    _set_grpc_defaults()
+    _is_process_set_up = True
+
+
+def is_process_set_up() -> bool:
+    return _is_process_set_up
+
+
+def _set_grpc_defaults() -> None:
+    # gRPC's fork handlers restart its threads in every forked child, which can crash the
+    # child before it execs and stall the parent while it waits for the exec. The server
+    # forks only to exec subprocesses, so it doesn't need them. gRPC reads this on import.
+    if "grpc" in sys.modules and "GRPC_ENABLE_FORK_SUPPORT" not in os.environ:
+        logger.warning(
+            "grpc was imported before bootstrap.set_up_process(), so gRPC fork support stays"
+            " enabled. This may crash subprocesses the server starts and stall the server while"
+            " it waits for them. Call set_up_process() before other imports or set"
+            " GRPC_ENABLE_FORK_SUPPORT=false in the environment."
+        )
+    os.environ.setdefault("GRPC_ENABLE_FORK_SUPPORT", "false")
+    # Keeps gRPC's own log messages out of server logs. Users can still change this.
+    os.environ.setdefault("GRPC_VERBOSITY", "ERROR")

src/dstack/_internal/server/main.py

@@ -1,4 +1,8 @@
-from dstack._internal.server.app import create_app, register_routes
+from dstack._internal.server import bootstrap
+
+bootstrap.set_up_process()
+
+from dstack._internal.server.app import create_app, register_routes  # noqa: E402
 
 app = create_app()
 register_routes(app)

src/tests/_internal/server/test_bootstrap.py

@@ -0,0 +1,47 @@
+import logging
+import os
+import sys
+
+import pytest
+
+from dstack._internal.server import bootstrap
+
+
+class TestSetUpProcess:
+    @pytest.fixture(autouse=True)
+    def isolate_process_state(self, monkeypatch: pytest.MonkeyPatch):
+        monkeypatch.setattr(bootstrap, "_is_process_set_up", False)
+        monkeypatch.delenv("GRPC_ENABLE_FORK_SUPPORT", raising=False)
+        monkeypatch.delenv("GRPC_VERBOSITY", raising=False)
+        monkeypatch.delitem(sys.modules, "grpc", raising=False)
+
+    def test_sets_grpc_defaults(self):
+        bootstrap.set_up_process()
+        assert os.environ["GRPC_ENABLE_FORK_SUPPORT"] == "false"
+        assert os.environ["GRPC_VERBOSITY"] == "ERROR"
+        assert bootstrap.is_process_set_up()
+
+    def test_keeps_grpc_settings_from_environment(self, monkeypatch: pytest.MonkeyPatch):
+        monkeypatch.setenv("GRPC_ENABLE_FORK_SUPPORT", "true")
+        monkeypatch.setenv("GRPC_VERBOSITY", "DEBUG")
+        bootstrap.set_up_process()
+        assert os.environ["GRPC_ENABLE_FORK_SUPPORT"] == "true"
+        assert os.environ["GRPC_VERBOSITY"] == "DEBUG"
+
+    def test_warns_if_grpc_already_imported(
+        self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+    ):
+        monkeypatch.setitem(sys.modules, "grpc", object())
+        with caplog.at_level(logging.WARNING, logger=bootstrap.__name__):
+            bootstrap.set_up_process()
+        assert "grpc was imported before bootstrap.set_up_process()" in caplog.text
+        assert bootstrap.is_process_set_up()
+
+    def test_no_warning_if_grpc_already_imported_with_setting_from_environment(
+        self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+    ):
+        monkeypatch.setitem(sys.modules, "grpc", object())
+        monkeypatch.setenv("GRPC_ENABLE_FORK_SUPPORT", "false")
+        with caplog.at_level(logging.WARNING, logger=bootstrap.__name__):
+            bootstrap.set_up_process()
+        assert caplog.text == ""