Fix lazy import (#2987) (#2990)

Signed-off-by: Kevin Su <[email protected]>
Co-authored-by: Kevin Su <[email protected]>

Author: wild-endeavor

flytekit/core/type_engine.py

@@ -1146,7 +1146,6 @@ class TypeEngine(typing.Generic[T]):
     _RESTRICTED_TYPES: typing.List[type] = []
     _DATACLASS_TRANSFORMER: TypeTransformer = DataclassTransformer()  # type: ignore
     _ENUM_TRANSFORMER: TypeTransformer = EnumTransformer()  # type: ignore
-    has_lazy_import = False
     lazy_import_lock = threading.Lock()
 
     @classmethod
@@ -1255,16 +1254,7 @@ def lazy_import_transformers(cls):
             # Avoid a race condition where concurrent threads may exit lazy_import_transformers before the transformers
             # have been imported. This could be implemented without a lock if you assume python assignments are atomic
             # and re-registering transformers is acceptable, but I decided to play it safe.
-            if cls.has_lazy_import:
-                return
-            cls.has_lazy_import = True
-            from flytekit.types.structured import (
-                register_arrow_handlers,
-                register_bigquery_handlers,
-                register_pandas_handlers,
-                register_snowflake_handlers,
-            )
-            from flytekit.types.structured.structured_dataset import DuplicateHandlerError
+            from flytekit.types.structured import lazy_import_structured_dataset_handler
 
             if is_imported("tensorflow"):
                 from flytekit.extras import tensorflow  # noqa: F401
@@ -1279,29 +1269,11 @@ def lazy_import_transformers(cls):
                     from flytekit.types.schema.types_pandas import PandasSchemaReader, PandasSchemaWriter  # noqa: F401
                 except ValueError:
                     logger.debug("Transformer for pandas is already registered.")
-                try:
-                    register_pandas_handlers()
-                except DuplicateHandlerError:
-                    logger.debug("Transformer for pandas is already registered.")
-            if is_imported("pyarrow"):
-                try:
-                    register_arrow_handlers()
-                except DuplicateHandlerError:
-                    logger.debug("Transformer for arrow is already registered.")
-            if is_imported("google.cloud.bigquery"):
-                try:
-                    register_bigquery_handlers()
-                except DuplicateHandlerError:
-                    logger.debug("Transformer for bigquery is already registered.")
             if is_imported("numpy"):
                 from flytekit.types import numpy  # noqa: F401
             if is_imported("PIL"):
                 from flytekit.types.file import image  # noqa: F401
-            if is_imported("snowflake.connector"):
-                try:
-                    register_snowflake_handlers()
-                except DuplicateHandlerError:
-                    logger.debug("Transformer for snowflake is already registered.")
+            lazy_import_structured_dataset_handler()
 
     @classmethod
     def to_literal_type(cls, python_type: Type[T]) -> LiteralType:

flytekit/lazy_import/lazy_module.py

@@ -2,10 +2,12 @@
 import sys
 import types
 
-LAZY_MODULES = []
 
+class _LazyModule(types.ModuleType):
+    """
+    `lazy_module` returns an instance of this class if the module is not found in the python environment.
+    """
 
-class LazyModule(types.ModuleType):
     def __init__(self, module_name: str):
         super().__init__(module_name)
         self._module_name = module_name
@@ -17,8 +19,12 @@ def __getattribute__(self, attr):
 def is_imported(module_name):
     """
     This function is used to check if a module has been imported by the regular import.
+    Return false if module is lazy imported and not used yet.
     """
-    return module_name in sys.modules and module_name not in LAZY_MODULES
+    return (
+        module_name in sys.modules
+        and object.__getattribute__(lazy_module(module_name), "__class__").__name__ != "_LazyModule"
+    )
 
 
 def lazy_module(fullname):
@@ -37,11 +43,12 @@ def lazy_module(fullname):
     if spec is None or spec.loader is None:
         # Return a lazy module if the module is not found in the python environment,
         # so that we can raise a proper error when the user tries to access an attribute in the module.
-        return LazyModule(fullname)
+        # The reason to do this is because importlib.util.LazyLoader still requires
+        # the module to be installed even if you don't use it.
+        return _LazyModule(fullname)
     loader = importlib.util.LazyLoader(spec.loader)
     spec.loader = loader
     module = importlib.util.module_from_spec(spec)
     sys.modules[fullname] = module
-    LAZY_MODULES.append(module)
     loader.exec_module(module)
     return module

flytekit/types/structured/__init__.py

@@ -13,15 +13,15 @@
 """
 
 from flytekit.deck.renderer import ArrowRenderer, TopFrameRenderer
+from flytekit.lazy_import.lazy_module import is_imported
 from flytekit.loggers import logger
 
 from .structured_dataset import (
+    DuplicateHandlerError,
     StructuredDataset,
     StructuredDatasetDecoder,
     StructuredDatasetEncoder,
-    StructuredDatasetMetadata,
     StructuredDatasetTransformerEngine,
-    StructuredDatasetType,
 )
 
 
@@ -84,3 +84,27 @@ def register_snowflake_handlers():
             "We won't register snowflake handler for structured dataset because "
             "we can't find package snowflake-connector-python"
         )
+
+
+def lazy_import_structured_dataset_handler():
+    if is_imported("pandas"):
+        try:
+            register_pandas_handlers()
+            register_csv_handlers()
+        except DuplicateHandlerError:
+            logger.debug("Transformer for pandas is already registered.")
+    if is_imported("pyarrow"):
+        try:
+            register_arrow_handlers()
+        except DuplicateHandlerError:
+            logger.debug("Transformer for arrow is already registered.")
+    if is_imported("google.cloud.bigquery"):
+        try:
+            register_bigquery_handlers()
+        except DuplicateHandlerError:
+            logger.debug("Transformer for bigquery is already registered.")
+    if is_imported("snowflake.connector"):
+        try:
+            register_snowflake_handlers()
+        except DuplicateHandlerError:
+            logger.debug("Transformer for snowflake is already registered.")

flytekit/types/structured/structured_dataset.py

@@ -170,6 +170,18 @@ def literal(self) -> Optional[literals.StructuredDataset]:
         return self._literal_sd
 
     def open(self, dataframe_type: Type[DF]):
+        from flytekit.types.structured import lazy_import_structured_dataset_handler
+
+        """
+        Load the handler if needed. For the use case like:
+        @task
+        def t1(sd: StructuredDataset):
+          import pandas as pd
+          sd.open(pd.DataFrame).all()
+
+        pandas is imported inside the task, so pandnas handler won't be loaded during deserialization in type engine.
+        """
+        lazy_import_structured_dataset_handler()
         self._dataframe_type = dataframe_type
         return self
 

tests/flytekit/unit/core/test_generice_idl_type_engine.py

@@ -3439,35 +3439,6 @@ def outer_workflow(input: OuterWorkflowInput) -> OuterWorkflowOutput:
     assert none_value_output is None, f"None value was {none_value_output}, not None as expected"
 
 
-@pytest.mark.serial
-def test_lazy_import_transformers_concurrently():
-    # Ensure that next call to TypeEngine.lazy_import_transformers doesn't skip the import. Mark as serial to ensure
-    # this achieves what we expect.
-    TypeEngine.has_lazy_import = False
-
-    # Configure the mocks similar to https://stackoverflow.com/questions/29749193/python-unit-testing-with-two-mock-objects-how-to-verify-call-order
-    after_import_mock, mock_register = mock.Mock(), mock.Mock()
-    mock_wrapper = mock.Mock()
-    mock_wrapper.mock_register = mock_register
-    mock_wrapper.after_import_mock = after_import_mock
-
-    with mock.patch.object(StructuredDatasetTransformerEngine, "register", new=mock_register):
-        def run():
-            TypeEngine.lazy_import_transformers()
-            after_import_mock()
-
-        N = 5
-        with ThreadPoolExecutor(max_workers=N) as executor:
-            futures = [executor.submit(run) for _ in range(N)]
-            [f.result() for f in futures]
-
-        # Assert that all the register calls come before anything else.
-        assert mock_wrapper.mock_calls[-N:] == [mock.call.after_import_mock()] * N
-        expected_number_of_register_calls = len(mock_wrapper.mock_calls) - N
-        assert all([mock_call[0] == "mock_register" for mock_call in
-                    mock_wrapper.mock_calls[:expected_number_of_register_calls]])
-
-
 @pytest.mark.skipif(sys.version_info < (3, 10), reason="PEP604 requires >=3.10, 585 requires >=3.9")
 def test_option_list_with_pipe():
     pt = list[int] | None

tests/flytekit/unit/core/test_type_engine.py

@@ -3451,10 +3451,6 @@ def outer_workflow(input: OuterWorkflowInput) -> OuterWorkflowOutput:
 
 @pytest.mark.serial
 def test_lazy_import_transformers_concurrently():
-    # Ensure that next call to TypeEngine.lazy_import_transformers doesn't skip the import. Mark as serial to ensure
-    # this achieves what we expect.
-    TypeEngine.has_lazy_import = False
-
     # Configure the mocks similar to https://stackoverflow.com/questions/29749193/python-unit-testing-with-two-mock-objects-how-to-verify-call-order
     after_import_mock, mock_register = mock.Mock(), mock.Mock()
     mock_wrapper = mock.Mock()
@@ -3471,11 +3467,11 @@ def run():
             futures = [executor.submit(run) for _ in range(N)]
             [f.result() for f in futures]
 
-        # Assert that all the register calls come before anything else.
-        assert mock_wrapper.mock_calls[-N:] == [mock.call.after_import_mock()] * N
+        assert mock_wrapper.mock_calls[-1] == mock.call.after_import_mock()
         expected_number_of_register_calls = len(mock_wrapper.mock_calls) - N
+        assert sum([mock_call[0] == "mock_register" for mock_call in mock_wrapper.mock_calls]) == expected_number_of_register_calls
         assert all([mock_call[0] == "mock_register" for mock_call in
-                    mock_wrapper.mock_calls[:expected_number_of_register_calls]])
+                    mock_wrapper.mock_calls[:int(len(mock_wrapper.mock_calls)/N)-1]])
 
 
 @pytest.mark.skipif(sys.version_info < (3, 10), reason="PEP604 requires >=3.10, 585 requires >=3.9")

tests/flytekit/unit/lazy_module/test_lazy_module.py

@@ -1,12 +1,24 @@
-import pytest
+import sys
+from unittest.mock import Mock
 
-from flytekit.lazy_import.lazy_module import LazyModule, lazy_module
+import pytest
+from flytekit.lazy_import.lazy_module import _LazyModule, lazy_module, is_imported
 
 
 def test_lazy_module():
     mod = lazy_module("click")
     assert mod.__name__ == "click"
     mod = lazy_module("fake_module")
-    assert isinstance(mod, LazyModule)
+
+    sys.modules["fake_module"] = mod
+    assert not is_imported("fake_module")
+    assert isinstance(mod, _LazyModule)
     with pytest.raises(ImportError, match="Module fake_module is not yet installed."):
         print(mod.attr)
+
+    non_lazy_module = Mock()
+    non_lazy_module.__name__ = 'NonLazyModule'
+    sys.modules["fake_module"] = non_lazy_module
+    assert is_imported("fake_module")
+
+    assert is_imported("dataclasses")

tests/flytekit/unit/types/structured_dataset/test_structured_dataset.py

@@ -22,7 +22,6 @@
 from flytekit.models.literals import StructuredDatasetMetadata
 from flytekit.models.types import LiteralType, SchemaType, SimpleType, StructuredDatasetType
 from flytekit.tools.translator import get_serializable
-from flytekit.types.file import FlyteFile
 from flytekit.types.structured.structured_dataset import (
     PARQUET,
     StructuredDataset,