added environment variable to disable extensions autoload (#1095)

* added environment variable to disable extensions autoload
---------

Co-authored-by: zilto <tjean@DESKTOP-V6JDCS2>

docs/how-tos/extensions-autoloading.rst

@@ -0,0 +1,97 @@
+=====================
+Extension autoloading
+=====================
+
+Under ``hamilton.plugins``, there are many modules named ``*_extensions`` (e.g., ``hamilton.plugins.pandas_extensions``, ``hamilton.plugins.mlflow_extensions``). They implement Hamilton features for 3rd party libraries, including ``@extract_columns``, materializers (``to.parquet``, ``from_.mlflow``), and more.
+
+
+Autoloading behavior
+--------------------
+
+By default, Hamilton attempts to load all extensions one-by-one. This means that as you have more Python packages in your environment (e.g., ``pandas``, ``pyspark``, ``mlflow``, ``xgboost``), importing Hamilton appears to become slower because it actually imports many packages.
+
+This behavior can be less desirable when your Hamilton dataflow doesn't use any of these packages, but you need them in your Python environment nonetheless. For example, if only ``pandas`` is needed for your dataflow, but you have ``mlflow`` and ``xgboost`` in your environment their respective extensions will be loaded each time.
+
+
+Disable autoloading
+--------------------
+
+Disabling extension autoloading allows to import Hamilton without any extensions, which can reduce import time from 2-3 sec to less than 0.5 sec. This speedup is welcomed when you need to restart a notebook's kernel often or you're operating in a low RAM environment (some Python packages are larger than 50Mbs).
+
+There are three ways to opt-out: programmatically, environment variables, configuration file. You must opt-out before having any other ``hamilton`` import.
+
+1. Programmatically
+~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: python
+
+    from hamilton import registry
+    registry.disable_autoload()
+
+2. Environment variables
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+From the console
+
+.. code-block:: console
+
+    export HAMILTON_AUTOLOAD_EXTENSIONS=0
+
+Programmatically via Python ``os.environ``.
+
+.. code-block:: python
+
+    import os
+    os.environ["HAMILTON_AUTOLOAD_EXTENSIONS"] = "0"
+
+Programmatically in Jupyter notebooks
+
+.. code-block:: python
+
+    %env HAMILTON_AUTOLOAD_EXTENSIONS=0
+
+3. Configuration file
+~~~~~~~~~~~~~~~~~~~~~
+
+Using the following command disables autoloading via the configuration file ``./hamilton.conf``. Hamilton won't autoload extensions anymore (i.e., you won't need to use approach 1 or 2 each time).
+
+.. code-block:: console
+
+    hamilton-disable-autoload-extensions
+
+To revert this configuration use the following command
+
+.. code-block:: console
+
+    hamilton-enable-autoload-extensions
+
+To reenable autoloading in specific files, you can delete the environment variable or use ``registry.enable_autoload()`` before calling ``registry.initialize()``
+
+.. code-block:: python
+
+    from hamilton import registry
+    registry.enable_autoload()
+    registry.initialize()
+
+
+Manually loading extensions
+----------------------------
+
+If you disabled autoloading, extensions need to be loaded manually. You should load them before having any other ``hamilton`` import to avoid hard-to-track bugs. There are two ways.
+
+1. Importing the extension
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: python
+
+    from hamilton.plugins import pandas_extensions, mlflow_extensions
+
+2. Registering the extension
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This approach has good IDE support via ``typing.Literal``
+
+.. code-block:: python
+
+    from hamilton import registry
+    registry.load_extensions("mlflow")

docs/how-tos/index.rst

@@ -18,6 +18,7 @@ directory. If there's an example you want but don't see, reach out or open an is
    cache-nodes
    scale-up
    microservice
+   extensions-autoloading
    wrapping-driver
    cli-reference
    pre-commit-hooks

hamilton/function_modifiers/base.py

@@ -20,38 +20,7 @@
     # Trigger load of extensions here because decorators are the only thing that use the registry
     # right now. Side note: ray serializes things weirdly, so we need to do this here rather than in
     # in the other choice of hamilton/base.py.
-    plugins_modules = [
-        "yaml",
-        "matplotlib",
-        "numpy",
-        "pandas",
-        "plotly",
-        "polars",
-        "polars_lazyframe",
-        "pyspark_pandas",
-        "spark",
-        "dask",
-        "geopandas",
-        "xgboost",
-        "lightgbm",
-        "sklearn_plot",
-        "vaex",
-        "ibis",
-        "dlt",
-        "kedro",
-        "huggingface",
-        "mlflow",
-    ]
-    for plugin_module in plugins_modules:
-        try:
-            registry.load_extension(plugin_module)
-        except NotImplementedError as e:
-            logger.debug(f"Did not load {plugin_module} extension because {str(e)}.")
-        except ModuleNotFoundError as e:
-            logger.debug(f"Did not load {plugin_module} extension because {e.msg}.")
-        except ImportError as e:
-            logger.debug(f"Did not load {plugin_module} extension because {str(e)}.")
-    registry.INITIALIZED = True
+    registry.initialize()
 
 
 def sanitize_function_name(name: str) -> str:

hamilton/registry.py

@@ -1,13 +1,43 @@
 import collections
+import configparser
 import functools
 import importlib
 import logging
-from typing import Any, Dict, Optional, Type
+import os
+import pathlib
+from typing import Any, Dict, Literal, Optional, Tuple, Type, get_args
 
 logger = logging.getLogger(__name__)
 
 # Use this to ensure the registry is loaded only once.
 INITIALIZED = False
+ExtensionName = Literal[
+    "yaml",
+    "matplotlib",
+    "numpy",
+    "pandas",
+    "plotly",
+    "polars",
+    "polars_lazyframe",
+    "pyspark_pandas",
+    "spark",
+    "dask",
+    "geopandas",
+    "xgboost",
+    "lightgbm",
+    "sklearn_plot",
+    "vaex",
+    "ibis",
+    "dlt",
+    "kedro",
+    "huggingface",
+    "mlflow",
+]
+HAMILTON_EXTENSIONS: Tuple[ExtensionName, ...] = get_args(ExtensionName)
+HAMILTON_AUTOLOAD_ENV = "HAMILTON_AUTOLOAD_EXTENSIONS"
+# NOTE the variable DEFAULT_CONFIG_LOCAITON is redundant with `hamilton.telemetry`
+# but this `registry` module must avoid circular imports
+DEFAULT_CONFIG_LOCATION = pathlib.Path("~/.hamilton.conf").expanduser()
 
 # This is a dictionary of extension name -> dict with dataframe and column types.
 DF_TYPE_AND_COLUMN_TYPES: Dict[str, Dict[str, Type]] = {}
@@ -16,6 +46,108 @@
 DATAFRAME_TYPE = "dataframe_type"
 
 
+def load_autoload_config() -> configparser.ConfigParser:
+    """Load the Hamilton config file and set the autoloading environment variable"""
+    config = configparser.ConfigParser()
+    config.read(DEFAULT_CONFIG_LOCATION)
+
+    if config.has_option("DEFAULT", HAMILTON_AUTOLOAD_ENV):
+        os.environ[HAMILTON_AUTOLOAD_ENV] = config.get("DEFAULT", HAMILTON_AUTOLOAD_ENV)
+
+    return config
+
+
+load_autoload_config()
+
+
+def load_extension(plugin_module: ExtensionName):
+    """Given a module name, loads it for Hamilton to use.
+
+    :param plugin_module: the module name sans .py. e.g. pandas, polars, pyspark_pandas.
+    """
+    mod = importlib.import_module(f"hamilton.plugins.{plugin_module}_extensions")
+    # We have various plugin extensions. We default to assuming it's a dataframe extension with columns,
+    # unless it explicitly says it's not.
+    # We need to check the following if we are to enable `@extract_columns` for example.
+    extractable = getattr(mod, "COLUMN_FRIENDLY_DF_TYPE", True)
+    if extractable:
+        assert hasattr(mod, "register_types"), "Error extension missing function register_types()"
+        assert hasattr(
+            mod, f"get_column_{plugin_module}"
+        ), f"Error extension missing get_column_{plugin_module}"
+        assert hasattr(
+            mod, f"fill_with_scalar_{plugin_module}"
+        ), f"Error extension missing fill_with_scalar_{plugin_module}"
+        logger.info(f"Detected {plugin_module} and successfully loaded Hamilton extensions.")
+
+
+def initialize():
+    """Iterate over all extensions and try to load them"""
+    logger.debug(f"{HAMILTON_AUTOLOAD_ENV}={os.environ.get(HAMILTON_AUTOLOAD_ENV)}")
+    for extension_name in HAMILTON_EXTENSIONS:
+        # skip modules that aren't explicitly imported by the user
+        if str(os.environ.get(HAMILTON_AUTOLOAD_ENV)) == "0":
+            continue
+
+        try:
+            load_extension(extension_name)
+        except NotImplementedError as e:
+            logger.debug(f"Did not load {extension_name} extension because {str(e)}.")
+        except ModuleNotFoundError as e:
+            logger.debug(f"Did not load {extension_name} extension because {e.msg}.")
+        except ImportError as e:
+            logger.debug(f"Did not load {extension_name} extension because {str(e)}.")
+
+    global INITIALIZED
+    INITIALIZED = True
+
+
+def disable_autoload():
+    """Disable extension autoloading by setting an environment variable.
+    This needs to be done before hamilton.driver is imported.
+    """
+    os.environ[HAMILTON_AUTOLOAD_ENV] = "0"
+
+
+def enable_autoload():
+    """Enable extension autoloading by deleting an environment variable.
+    This needs to be done before hamilton.driver is imported.
+    """
+    del os.environ[HAMILTON_AUTOLOAD_ENV]
+
+
+def config_enable_autoload():
+    """Modify the Hamilton config file to enable extension autoloading.
+    Autoloading can be disabled manually via `hamilton.registry.disable_autoload()`
+    before importing `hamilton.driver`.
+
+    NOTE the function name is tied to an entrypoint in `pyproject.toml`
+    """
+    config = load_autoload_config()
+    if "DEFAULT" not in config:
+        config.add_section("DEFAULT")
+
+    config.remove_option("DEFAULT", HAMILTON_AUTOLOAD_ENV)
+    with DEFAULT_CONFIG_LOCATION.open("w") as f:
+        config.write(f)
+
+
+def config_disable_autoload():
+    """Modify the Hamilton config file to disable extension autoloading.
+    Autoloading can be enabled manually via `hamilton.registry.enable_autoload()`
+    before importing `hamilton.driver`.
+
+    NOTE the function name is tied to an entrypoint in `pyproject.toml`
+    """
+    config = load_autoload_config()
+    if "DEFAULT" not in config:
+        config.add_section("DEFAULT")
+
+    config.set("DEFAULT", HAMILTON_AUTOLOAD_ENV, "0")
+    with DEFAULT_CONFIG_LOCATION.open("w") as f:
+        config.write(f)
+
+
 def register_types(extension_name: str, dataframe_type: Type, column_type: Optional[Type]):
     """Registers the dataframe and column types for the extension. Note that column types are optional
     as some extensions may not have a column type (E.G. spark). In this case, this is not included
@@ -73,27 +205,6 @@ def get_column_type_from_df_type(dataframe_type: Type) -> Type:
     )
 
 
-def load_extension(plugin_module: str):
-    """Given a module name, loads it for Hamilton to use.
-
-    :param plugin_module: the module name sans .py. e.g. pandas, polars, pyspark_pandas.
-    """
-    mod = importlib.import_module(f"hamilton.plugins.{plugin_module}_extensions")
-    # We have various plugin extensions. We default to assuming it's a dataframe extension with columns,
-    # unless it explicitly says it's not.
-    # We need to check the following if we are to enable `@extract_columns` for example.
-    extractable = getattr(mod, "COLUMN_FRIENDLY_DF_TYPE", True)
-    if extractable:
-        assert hasattr(mod, "register_types"), "Error extension missing function register_types()"
-        assert hasattr(
-            mod, f"get_column_{plugin_module}"
-        ), f"Error extension missing get_column_{plugin_module}"
-        assert hasattr(
-            mod, f"fill_with_scalar_{plugin_module}"
-        ), f"Error extension missing fill_with_scalar_{plugin_module}"
-        logger.info(f"Detected {plugin_module} and successfully loaded Hamilton extensions.")
-
-
 LOADER_REGISTRY = collections.defaultdict(list)
 SAVER_REGISTRY = collections.defaultdict(list)
 

pyproject.toml

@@ -153,6 +153,8 @@ h_experiments = "hamilton.plugins.h_experiments.__main__:main"
 hamilton = "hamilton.cli.__main__:cli"
 hamilton-admin-build-ui = "hamilton.admin:build_ui"
 hamilton-admin-build-and-publish = "hamilton.admin:build_and_publish"
+hamilton-disable-autoload-extensions = "hamilton.registry:config_disable_autoload"
+hamilton-enable-autoload-extensions = "hamilton.registry:config_enable_autoload"
 
 [project.urls]
 homepage = "https://www.tryhamilton.dev/"

tests/test_registry.py

@@ -0,0 +1,13 @@
+import pytest
+
+from hamilton import registry
+
+
+@pytest.mark.parametrize("entrypoint", ["config_disable_autoload", "config_enable_autoload"])
+def test_command_entrypoints_arent_renamed(entrypoint: str):
+    """Ensures that functions associated with an entrypoint in
+    pyproject.toml aren't renamed.
+
+    This doesn't prevent the entrypoints from being renamed
+    """
+    assert hasattr(registry, entrypoint)