Repository: skypilot-org/skypilot
Pull request: Fix db init race condition
<!-- Describe the changes in this PR -->

Close https://github.com/skypilot-org/skypilot/issues/6220

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: install pre-commit (auto-check on commit) or `bash format.sh`
- [x] Unit test which repros on master and fixed on PR branch
- [ ] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [ ] Relevant individual tests: `/smoke-test -k test_name` (CI) or `pytest tests/test_smoke.py::test_name` (local)
- [ ] Backward compatibility: `/quicktest-core` (CI) or `pytest tests/smoke_tests/test_backward_compat.py` (local)

<!-- CI commands (/-prefixed) can only be triggered by repo members -->


Problem statement:
[Core] DB init race condition
```
def initialize_and_get_db() -> sqlalchemy.engine.Engine:
    global _SQLALCHEMY_ENGINE
    if _SQLALCHEMY_ENGINE is not None:
        return _SQLALCHEMY_ENGINE
    with _DB_INIT_LOCK:
        if _SQLALCHEMY_ENGINE is None:
```

Problematic timeline:

```
thread 1 -> AcquireLock -> set _SQLALCHEMY_ENGINE ----------------------------------> create table
thread 2 -----------------------------------------------> check _SQLALCHEMY_ENGINE and return, table not initialized
```

For concurrent request, a thread may get `_SQLALCHEMY_ENGINE` with partial state, which cause table not found error.

Hints from issue discussion:

Commit message: Fix

Signed-off-by: Aylei <rayingecho@gmail.com>

Patch to review:
diff --git a/sky/global_user_state.py b/sky/global_user_state.py
index 6559b6a0ccc..14543ef75fa 100644
--- a/sky/global_user_state.py
+++ b/sky/global_user_state.py
@@ -220,17 +220,16 @@ def replace_char_class(match):
     return like_pattern
 
 
-def create_table():
+def create_table(engine: sqlalchemy.engine.Engine):
     # Enable WAL mode to avoid locking issues.
     # See: issue #1441 and PR #1509
     # https://github.com/microsoft/WSL/issues/2395
     # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
     #  This may cause the database locked problem from WSL issue #1441.
-    if (_SQLALCHEMY_ENGINE.dialect.name
-            == db_utils.SQLAlchemyDialect.SQLITE.value and
+    if (engine.dialect.name == db_utils.SQLAlchemyDialect.SQLITE.value and
             not common_utils.is_wsl()):
         try:
-            with orm.Session(_SQLALCHEMY_ENGINE) as session:
+            with orm.Session(engine) as session:
                 session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
                 session.commit()
         except sqlalchemy_exc.OperationalError as e:
@@ -240,12 +239,12 @@ def create_table():
             # is not critical and is likely to be enabled by other processes.
 
     # Create tables if they don't exist
-    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, _SQLALCHEMY_ENGINE)
+    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, engine)
 
     # For backward compatibility.
     # TODO(zhwu): Remove this function after all users have migrated to
     # the latest version of SkyPilot.
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(engine) as session:
         # Add autostop column to clusters table
         db_utils.add_column_to_table_sqlalchemy(session,
                                                 'clusters',
@@ -391,15 +390,15 @@ def initialize_and_get_db() -> sqlalchemy.engine.Engine:
                 conn_string = skypilot_config.get_nested(('db',), None)
             if conn_string:
                 logger.debug(f'using db URI from {conn_string}')
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine(
-                    conn_string, poolclass=sqlalchemy.NullPool)
+                engine = sqlalchemy.create_engine(conn_string,
+                                                  poolclass=sqlalchemy.NullPool)
             else:
                 db_path = os.path.expanduser('~/.sky/state.db')
                 pathlib.Path(db_path).parents[0].mkdir(parents=True,
                                                        exist_ok=True)
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine('sqlite:///' +
-                                                              db_path)
-            create_table()
+                engine = sqlalchemy.create_engine('sqlite:///' + db_path)
+            create_table(engine)
+            _SQLALCHEMY_ENGINE = engine
     return _SQLALCHEMY_ENGINE
 
 
diff --git a/sky/jobs/state.py b/sky/jobs/state.py
index aac6ca82f20..5e437686fb9 100644
--- a/sky/jobs/state.py
+++ b/sky/jobs/state.py
@@ -112,17 +112,16 @@
 )
 
 
-def create_table():
+def create_table(engine: sqlalchemy.engine.Engine):
     # Enable WAL mode to avoid locking issues.
     # See: issue #3863, #1441 and PR #1509
     # https://github.com/microsoft/WSL/issues/2395
     # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
     #  This may cause the database locked problem from WSL issue #1441.
-    if (_SQLALCHEMY_ENGINE.dialect.name
-            == db_utils.SQLAlchemyDialect.SQLITE.value and
+    if (engine.dialect.name == db_utils.SQLAlchemyDialect.SQLITE.value and
             not common_utils.is_wsl()):
         try:
-            with orm.Session(_SQLALCHEMY_ENGINE) as session:
+            with orm.Session(engine) as session:
                 session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
                 session.commit()
         except sqlalchemy_exc.OperationalError as e:
@@ -132,10 +131,10 @@ def create_table():
             # is not critical and is likely to be enabled by other processes.
 
     # Create tables if they don't exist
-    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, _SQLALCHEMY_ENGINE)
+    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, engine)
 
     # Backward compatibility: add columns that not exist in older databases
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(engine) as session:
         db_utils.add_column_to_table_sqlalchemy(session, 'spot',
                                                 'failure_reason',
                                                 sqlalchemy.Text())
@@ -228,15 +227,15 @@ def initialize_and_get_db() -> sqlalchemy.engine.Engine:
                 conn_string = skypilot_config.get_nested(('db',), None)
             if conn_string:
                 logger.debug(f'using db URI from {conn_string}')
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine(
-                    conn_string, poolclass=sqlalchemy.NullPool)
+                engine = sqlalchemy.create_engine(conn_string,
+                                                  poolclass=sqlalchemy.NullPool)
             else:
                 db_path = os.path.expanduser('~/.sky/spot_jobs.db')
                 pathlib.Path(db_path).parents[0].mkdir(parents=True,
                                                        exist_ok=True)
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine('sqlite:///' +
-                                                              db_path)
-            create_table()
+                engine = sqlalchemy.create_engine('sqlite:///' + db_path)
+            create_table(engine)
+            _SQLALCHEMY_ENGINE = engine
     return _SQLALCHEMY_ENGINE
 
 
diff --git a/tests/test_failover.py b/tests/test_failover.py
index 40408887315..1f8a8e09c17 100644
--- a/tests/test_failover.py
+++ b/tests/test_failover.py
@@ -28,7 +28,7 @@ def _mock_db_conn(tmp_path, monkeypatch):
     monkeypatch.setattr(global_user_state, '_SQLALCHEMY_ENGINE',
                         sqlalchemy_engine)
 
-    global_user_state.create_table()
+    global_user_state.create_table(sqlalchemy_engine)
 
 
 @pytest.mark.parametrize('enable_all_clouds', [[sky.AWS()]], indirect=True)
diff --git a/tests/test_global_user_state.py b/tests/test_global_user_state.py
index 7189e02a336..fb8f29b68ef 100644
--- a/tests/test_global_user_state.py
+++ b/tests/test_global_user_state.py
@@ -1,8 +1,11 @@
 import sys
+import threading
+from unittest import mock
 
 import pytest
 
 import sky
+from sky import global_user_state
 
 
 @pytest.mark.skipif(sys.platform != 'linux', reason='Only test in CI.')
@@ -10,3 +13,54 @@ def test_enabled_clouds_empty():
     # In test environment, no cloud should be enabled.
     assert sky.global_user_state.get_cached_enabled_clouds(
         sky.clouds.cloud.CloudCapability.COMPUTE, workspace='default') == []
+
+
+def test_concurrent_database_initializationd(tmp_path):
+    """Test that concurrent database initialization."""
+    # Store original state to restore later
+    with mock.patch.object(global_user_state, '_SQLALCHEMY_ENGINE', None):
+
+        temp_db_path = tmp_path / 'state.db'
+        results = []
+        num_threads = 5
+
+        def worker_thread(thread_id):
+            """Worker thread that initializes DB and performs operations"""
+            try:
+                with mock.patch('os.path.expanduser',
+                                return_value=str(temp_db_path)):
+                    # Force database initialization
+                    global_user_state.initialize_and_get_db()
+
+                    # Immediately try to use the database
+                    # This should work if tables are properly created
+                    global_user_state.get_user(f"test_user_{thread_id}")
+
+                    results.append((thread_id, "SUCCESS", None))
+            except Exception as e:
+                results.append((thread_id, "FAILED", e))
+
+        threads = []
+        for i in range(num_threads):
+            t = threading.Thread(target=worker_thread, args=(i,))
+            threads.append(t)
+
+        for t in threads:
+            t.start()
+
+        for t in threads:
+            t.join(timeout=10)
+
+        successes = [r for r in results if r[1] == "SUCCESS"]
+        failures = [r for r in results if r[1] == "FAILED"]
+
+        print(f"Results: {len(successes)} successes, {len(failures)} failures")
+        if failures:
+            print("Failures:")
+            for thread_id, status, error in failures:
+                print(f"  Thread {thread_id}: {error}")
+
+        assert len(failures) == 0, (
+            f"Race condition detected: {len(failures)} threads failed. "
+            f"This indicates the database initialization has race conditions. "
+            f"Failures: {[(f[0], str(f[2])) for f in failures]}")
diff --git a/tests/test_jobs.py b/tests/test_jobs.py
index f2e33729741..1136d9878de 100644
--- a/tests/test_jobs.py
+++ b/tests/test_jobs.py
@@ -19,7 +19,7 @@ def _mock_db_conn(tmp_path, monkeypatch):
     monkeypatch.setattr(global_user_state, '_SQLALCHEMY_ENGINE',
                         sqlalchemy_engine)
 
-    global_user_state.create_table()
+    global_user_state.create_table(sqlalchemy_engine)
 
 
 @pytest.fixture
diff --git a/tests/test_jobs_and_serve.py b/tests/test_jobs_and_serve.py
index d13431ef22f..e9e8f2dfa53 100644
--- a/tests/test_jobs_and_serve.py
+++ b/tests/test_jobs_and_serve.py
@@ -47,7 +47,7 @@ def _mock_db_conn(tmp_path, monkeypatch):
     monkeypatch.setattr(global_user_state, '_SQLALCHEMY_ENGINE',
                         sqlalchemy_engine)
 
-    global_user_state.create_table()
+    global_user_state.create_table(sqlalchemy_engine)
 
 
 def _generate_tmp_yaml(tmp_path, filename: str) -> str:


Code review comments:
diff_hunk: @@ -1,12 +1,66 @@
 import sys
+import threading
+from unittest import mock
 
 import pytest
 
 import sky
+from sky import global_user_state
 
 
 @pytest.mark.skipif(sys.platform != 'linux', reason='Only test in CI.')
 def test_enabled_clouds_empty():
     # In test environment, no cloud should be enabled.
     assert sky.global_user_state.get_cached_enabled_clouds(
         sky.clouds.cloud.CloudCapability.COMPUTE, workspace='default') == []
+
+
+def test_concurrent_database_initializationd(tmp_path):
+    """Test that concurrent database initialization."""
+    # Store original state to restore later
+    with mock.patch.object(global_user_state, '_SQLALCHEMY_ENGINE', None):
+
+        temp_db_path = tmp_path / 'state.db'
+        results = []
+        num_threads = 5
+
+        def worker_thread(thread_id):
+            """Worker thread that initializes DB and performs operations"""
+            try:
+                with mock.patch('os.path.expanduser',
+                                return_value=str(temp_db_path)):
+                    # Force database initialization
+                    global_user_state.initialize_and_get_db()
+
+                    # Immediately try to use the database
+                    # This should work if tables are properly created
+                    global_user_state.get_user(f"test_user_{thread_id}")
path: tests/test_global_user_state.py
text: @user1:
I'm not certain this query will succeed given the `_SQLALCHEMY_ENGINE` is patched to None at this point