⚡️ Speed up ln.connect() (#871)

* ⚡️ Profile connect

* ⚡️ More profiling

* ♻️ Simplify reloading of schema modules

* ♻️ More simplification

* ♻️ Prettier code, same performance

* 🔥 Remove debug logging statements

* ➕ Add line_profiler to dev dependencies

* ♻️ Prettify code, same performance

* ⬆️ Check out wetlab and bionty branches

* 💚 Try to fix tests

* 💚 Deal in lamindb

* 🏗️ Load bionty sources in bionty instead of in ln.connect

* 🔊 Profile more

* 🔥 Remove profiling statements

lamindb_setup/_check_setup.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import importlib as il
 import os
 from typing import TYPE_CHECKING, Optional
 
@@ -52,9 +53,18 @@ def _get_current_instance_settings() -> InstanceSettings | None:
 
 
 # we make this a private function because in all the places it's used,
-# users should see it
-def _check_instance_setup(from_lamindb: bool = False) -> bool:
+# users should not see it
+def _check_instance_setup(
+    from_lamindb: bool = False, from_module: str | None = None
+) -> bool:
+    reload_module = from_lamindb or from_module is not None
+    from ._init_instance import get_schema_module_name, reload_schema_modules
+
     if django.IS_SETUP:
+        # reload logic here because module might not yet have been imported
+        # upon first setup
+        if from_module is not None:
+            il.reload(il.import_module(from_module))
         return True
     silence_loggers()
     if os.environ.get("LAMINDB_MULTI_INSTANCE") == "true":
@@ -65,15 +75,18 @@ def _check_instance_setup(from_lamindb: bool = False) -> bool:
         return True
     isettings = _get_current_instance_settings()
     if isettings is not None:
-        if from_lamindb and settings.auto_connect:
+        if reload_module and settings.auto_connect:
             if not django.IS_SETUP:
-                from ._init_instance import reload_schema_modules
-
                 django.setup_django(isettings)
-                reload_schema_modules(isettings)
+                if from_module is not None:
+                    # this only reloads `from_module`
+                    il.reload(il.import_module(from_module))
+                else:
+                    # this bulk reloads all schema modules
+                    reload_schema_modules(isettings)
                 logger.important(f"connected lamindb: {isettings.slug}")
         return django.IS_SETUP
     else:
-        if from_lamindb and settings.auto_connect:
+        if reload_module and settings.auto_connect:
             logger.warning(InstanceNotSetupError.default_message)
         return False

lamindb_setup/_connect_instance.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import sys
 from typing import TYPE_CHECKING
 from uuid import UUID
 
@@ -10,8 +11,7 @@
 from ._close import close as close_instance
 from ._init_instance import MESSAGE_NO_MULTIPLE_INSTANCE, load_from_isettings
 from ._silence_loggers import silence_loggers
-from .core._hub_core import connect_instance as load_instance_from_hub
-from .core._hub_core import connect_instance_new as load_instance_from_hub_edge
+from .core._hub_core import connect_instance as connect_instance_from_hub
 from .core._hub_utils import (
     LaminDsn,
     LaminDsnModel,
@@ -128,14 +128,9 @@ def _connect_instance(
         # on the hub
         # do not call hub if the user is anonymous
         if owner != "anonymous":
-            if settings.user.handle in {"Koncopd", "sunnyosun", "falexwolf"}:
-                hub_result = load_instance_from_hub_edge(
-                    owner=owner, name=name, access_token=access_token
-                )
-            else:
-                hub_result = load_instance_from_hub(
-                    owner=owner, name=name, access_token=access_token
-                )
+            hub_result = connect_instance_from_hub(
+                owner=owner, name=name, access_token=access_token
+            )
         else:
             hub_result = "anonymous-user"
         # if hub_result is not a string, it means it made a request

lamindb_setup/_init_instance.py

@@ -97,8 +97,9 @@ def register_user_and_storage_in_instance(isettings: InstanceSettings, usettings
         logger.warning(f"instance seems not set up ({error})")
 
 
-def reload_schema_modules(isettings: InstanceSettings):
-    schema_names = ["core"] + list(isettings.schema)
+def reload_schema_modules(isettings: InstanceSettings, include_core: bool = True):
+    schema_names = ["core"] if include_core else []
+    # schema_names += list(isettings.schema)
     schema_module_names = [get_schema_module_name(n) for n in schema_names]
 
     for schema_module_name in schema_module_names:
@@ -114,16 +115,10 @@ def reload_lamindb_itself(isettings) -> bool:
 
         importlib.reload(lamindb)
         reloaded = True
-    if "bionty" in isettings.schema and "bionty" in sys.modules:
-        schema_module = importlib.import_module("bionty")
-        importlib.reload(schema_module)
-        reloaded = True
     return reloaded
 
 
 def reload_lamindb(isettings: InstanceSettings):
-    # only touch lamindb if we're operating from lamindb
-    reload_schema_modules(isettings)
     log_message = settings.auto_connect
     if not reload_lamindb_itself(isettings):
         log_message = True
@@ -352,7 +347,7 @@ def load_from_isettings(
     user: UserSettings | None = None,
     write_settings: bool = True,
 ) -> None:
-    from .core._setup_bionty_sources import load_bionty_sources, write_bionty_sources
+    from .core._setup_bionty_sources import write_bionty_sources
 
     user = settings.user if user is None else user
 
@@ -370,7 +365,6 @@ def load_from_isettings(
         # yet be registered
         if not isettings._get_settings_file().exists():
             register_user(user)
-        load_bionty_sources(isettings)
     isettings._persist(write_to_disk=write_settings)
     reload_lamindb(isettings)
 

lamindb_setup/core/django.py

@@ -8,6 +8,7 @@
 from lamin_utils import logger
 from ._settings_store import current_instance_settings_file
 from ._settings_instance import InstanceSettings
+import sys
 
 IS_RUN_FROM_IPYTHON = getattr(builtins, "__IPYTHON__", False)
 IS_SETUP = False

noxfile.py

@@ -27,8 +27,8 @@ def lint(session: nox.Session) -> None:
     ["hub-local", "hub-prod", "hub-cloud", "storage", "docs"],
 )
 def install(session: nox.Session, group: str) -> None:
-    no_deps_packages = "git+https://github.com/laminlabs/lnschema-core git+https://github.com/laminlabs/wetlab git+https://github.com/laminlabs/lamin-cli"
-    schema_deps = f"""uv pip install --system git+https://github.com/laminlabs/bionty git+https://github.com/laminlabs/lamindb@main
+    no_deps_packages = "git+https://github.com/laminlabs/lnschema-core git+https://github.com/laminlabs/wetlab@lazyreload git+https://github.com/laminlabs/lamin-cli"
+    schema_deps = f"""uv pip install --system git+https://github.com/laminlabs/bionty@lazyconnect git+https://github.com/laminlabs/lamindb@improvelaminconnect
 uv pip install --system --no-deps {no_deps_packages}
 """
     if group == "hub-cloud":

pyproject.toml

@@ -39,6 +39,7 @@ erdiagram = [
     "django-schema-graph",
 ]
 dev = [
+    "line_profiler",
     "pyjwt<3.0.0",
     "psycopg2-binary",
     "python-dotenv",