From e7e75e98138bee4a212252bb6afc600cb5c8042d Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Sat, 2 Aug 2025 00:12:40 -0700 Subject: [PATCH 01/10] [Core feature] Reuse same literals in the dynamic task Signed-off-by: Barry Wu --- flytekit/core/type_engine.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 58ba0b8556..f6406a1cd2 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -21,6 +21,7 @@ from typing import Any, Dict, List, NamedTuple, Optional, Type, cast import msgpack +from cachetools import LRUCache from dataclasses_json import DataClassJsonMixin, dataclass_json from flyteidl.core import literals_pb2 from fsspec.asyn import _run_coros_in_chunks # pylint: disable=W0212 @@ -1174,6 +1175,7 @@ class TypeEngine(typing.Generic[T]): _DATACLASS_TRANSFORMER: TypeTransformer = DataclassTransformer() # type: ignore _ENUM_TRANSFORMER: TypeTransformer = EnumTransformer() # type: ignore lazy_import_lock = threading.Lock() + _CACHE = LRUCache(maxsize=128) @classmethod def register( @@ -1377,6 +1379,10 @@ def calculate_hash(cls, python_val: typing.Any, python_type: Type[T]) -> Optiona break return hsh + @lru_cache(typed=True) + def make_key(python_val: typing.Any, python_type: Type[T]) -> tuple: + return (repr(python_val), python_type.__name__) + @classmethod def to_literal( cls, ctx: FlyteContext, python_val: typing.Any, python_type: Type[T], expected: LiteralType @@ -1386,6 +1392,10 @@ def to_literal( to_literal function, and allowing this to_literal function, to then invoke yet another async function, namely an async transformer. """ + key = cls.make_key(python_val, python_type) + if key in cls._CACHE: + return cls._CACHE[key] + from flytekit.core.promise import Promise cls.to_literal_checks(python_val, python_type, expected) @@ -1406,6 +1416,8 @@ def to_literal( modify_literal_uris(lv) lv.hash = cls.calculate_hash(python_val, python_type) + + cls._CACHE[key] = lv return lv @classmethod From 4ba1485f25035835ab5283699a2ed55d6e27dba5 Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Tue, 5 Aug 2025 11:32:14 -0700 Subject: [PATCH 02/10] Fix the cache key generation, test_list fail and add unit tests for cache Signed-off-by: Barry Wu --- flytekit/core/type_engine.py | 32 ++++-- tests/flytekit/unit/core/test_list.py | 4 +- tests/flytekit/unit/core/test_type_engine.py | 100 +++++++++++++++++++ 3 files changed, 129 insertions(+), 7 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index f6406a1cd2..cab7cb5058 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -1175,7 +1175,7 @@ class TypeEngine(typing.Generic[T]): _DATACLASS_TRANSFORMER: TypeTransformer = DataclassTransformer() # type: ignore _ENUM_TRANSFORMER: TypeTransformer = EnumTransformer() # type: ignore lazy_import_lock = threading.Lock() - _CACHE = LRUCache(maxsize=128) + _CACHE: LRUCache = LRUCache(maxsize=128) @classmethod def register( @@ -1379,9 +1379,27 @@ def calculate_hash(cls, python_val: typing.Any, python_type: Type[T]) -> Optiona break return hsh - @lru_cache(typed=True) - def make_key(python_val: typing.Any, python_type: Type[T]) -> tuple: - return (repr(python_val), python_type.__name__) + @classmethod + def make_key(cls, python_val: typing.Any, python_type: Type[T]) -> Optional[tuple]: + import cloudpickle + + val_hash: typing.Any + type_hash: typing.Any + try: + try: + val_hash = hash(python_val) + except Exception: + val_hash = hash(cloudpickle.dumps(python_val)) + + try: + type_hash = hash(python_type) + except Exception: + type_hash = hash(cloudpickle.dumps(python_type)) + + return (val_hash, type_hash) + + except Exception: + return None @classmethod def to_literal( @@ -1393,7 +1411,7 @@ def to_literal( namely an async transformer. """ key = cls.make_key(python_val, python_type) - if key in cls._CACHE: + if key is not None and key in cls._CACHE: return cls._CACHE[key] from flytekit.core.promise import Promise @@ -1417,7 +1435,9 @@ def to_literal( modify_literal_uris(lv) lv.hash = cls.calculate_hash(python_val, python_type) - cls._CACHE[key] = lv + if key is not None: + cls._CACHE[key] = lv + return lv @classmethod diff --git a/tests/flytekit/unit/core/test_list.py b/tests/flytekit/unit/core/test_list.py index 96ee2efe78..c7f2f2d415 100644 --- a/tests/flytekit/unit/core/test_list.py +++ b/tests/flytekit/unit/core/test_list.py @@ -73,6 +73,8 @@ async def test_coroutine_batching_of_list_transformer(): lt = LiteralType(simple=SimpleType.INTEGER) python_val = [MyInt(10), MyInt(11), MyInt(12), MyInt(13), MyInt(14)] + # Use the different python_val to avoid hitting the cache + python_val_2 = [MyInt(11), MyInt(10), MyInt(12), MyInt(13), MyInt(14)] ctx = FlyteContext.current_context() with mock.patch("flytekit.core.type_engine._TYPE_ENGINE_COROS_BATCH_SIZE", 2): @@ -80,6 +82,6 @@ async def test_coroutine_batching_of_list_transformer(): with mock.patch("flytekit.core.type_engine._TYPE_ENGINE_COROS_BATCH_SIZE", 5): with pytest.raises(ValueError): - TypeEngine.to_literal(ctx, python_val, typing.List[MyInt], lt) + TypeEngine.to_literal(ctx, python_val_2, typing.List[MyInt], lt) del TypeEngine._REGISTRY[MyInt] diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 93d5d6af67..9be872b7c2 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -3851,3 +3851,103 @@ async def test_dict_transformer_annotated_type(): literal3 = await TypeEngine.async_to_literal(ctx, nested_dict, nested_dict_type, expected_type) assert literal3.map.literals["outer"].map.literals["inner"].scalar.primitive.integer == 42 + +def test_type_engine_cache(): + # Clear cache before test + TypeEngine._CACHE.clear() + + # Test data + ctx = FlyteContext.current_context() + python_val = [1, 2, 3, 4, 5] + python_type = typing.List[int] + expected = TypeEngine.to_literal_type(python_type) + + # First call - should not use cache + literal1 = TypeEngine.to_literal(ctx, python_val, python_type, expected) + + # Verify cache is populated + key = TypeEngine.make_key(python_val, python_type) + assert key is not None + assert key in TypeEngine._CACHE + + # Second call with same parameters - should use cache + literal2 = TypeEngine.to_literal(ctx, python_val, python_type, expected) + + # Verify both literals are identical (same object from cache) + assert literal1 is literal2 + + # Test with different data - should not use cache + different_val = [2, 1, 3, 4, 5] + literal3 = TypeEngine.to_literal(ctx, different_val, python_type, expected) + key_different = TypeEngine.make_key(different_val, python_type) + + assert key_different is not key + assert key_different is not None + assert key_different in TypeEngine._CACHE + + # Verify different literals are different objects + assert literal1 is not literal3 + + # Test cache with unhashable objects (should not cache) + python_val = {"a": [1, 2, 3]} # dict with list is unhashable + python_type = typing.Dict[str, typing.List[int]] + expected = TypeEngine.to_literal_type(python_type) + + # First call + literal4 = TypeEngine.to_literal(ctx, python_val, python_type, expected) + key = TypeEngine.make_key(python_val, python_type) + assert key is not None + assert key in TypeEngine._CACHE + + # Second call with same unhashable data + literal5 = TypeEngine.to_literal(ctx, python_val, python_type, expected) + + # Should be different objects since unhashable objects can't be cached + assert literal4 is literal5 + + # Add many different values to test cache size limit + for i in range(200): # More than the default maxsize of 128 + test_val = [i, i+1, i+2] + test_type = typing.List[int] + test_expected = TypeEngine.to_literal_type(test_type) + TypeEngine.to_literal(ctx, test_val, test_type, test_expected) + + # Cache should not exceed maxsize + assert len(TypeEngine._CACHE) == 128 + + # Test cache with pandas DataFrame (if available) + try: + import pandas as pd + + # Create DataFrame + df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + df_type = pd.DataFrame + df_expected = TypeEngine.to_literal_type(df_type) + + # First call + literal6 = TypeEngine.to_literal(ctx, df, df_type, df_expected) + + # Second call with same DataFrame + literal7 = TypeEngine.to_literal(ctx, df, df_type, df_expected) + + # Should be same object (DataFrame should be hashable via cloudpickle) + assert literal6 is literal7 + + except ImportError: + # Skip pandas test if not available + pass + + # Clean up + TypeEngine._CACHE.clear() + +def test_make_key_with_annotated_types(): + # Test with Annotated type + annotated_val = [1, 2, 3] + annotated_type = typing.Annotated[typing.List[int], "test_annotation"] + + key = TypeEngine.make_key(annotated_val, annotated_type) + key_without_annotation = TypeEngine.make_key(annotated_val, typing.List[int]) + # Should handle Annotated types correctly + assert key is not None + assert key_without_annotation is not None + assert key != key_without_annotation From edd9c8d1b91c5b0690b72575afaf82f31bfbcfd2 Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Tue, 5 Aug 2025 15:45:25 -0700 Subject: [PATCH 03/10] Add types-cachetools in dev-requirements.in to pass lint Signed-off-by: Barry Wu --- dev-requirements.in | 1 + 1 file changed, 1 insertion(+) diff --git a/dev-requirements.in b/dev-requirements.in index 34f11ff34f..cef6ce1929 100644 --- a/dev-requirements.in +++ b/dev-requirements.in @@ -46,6 +46,7 @@ types-protobuf<5 types-croniter types-decorator types-mock +types-cachetools autoflake pillow From c13abeb5ef45f0f54438822a4bc910992d69c2eb Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Wed, 6 Aug 2025 23:46:54 -0700 Subject: [PATCH 04/10] Fix some typos Signed-off-by: Barry Wu --- tests/flytekit/unit/core/test_type_engine.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 9be872b7c2..c884df6ff2 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -3888,8 +3888,8 @@ def test_type_engine_cache(): # Verify different literals are different objects assert literal1 is not literal3 - # Test cache with unhashable objects (should not cache) - python_val = {"a": [1, 2, 3]} # dict with list is unhashable + # Test cache with unhashable objects + python_val = {"a": [1, 2, 3]} python_type = typing.Dict[str, typing.List[int]] expected = TypeEngine.to_literal_type(python_type) @@ -3902,7 +3902,7 @@ def test_type_engine_cache(): # Second call with same unhashable data literal5 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - # Should be different objects since unhashable objects can't be cached + # Should be the same object since unhashable objects will fallback to cloudpickle to hash assert literal4 is literal5 # Add many different values to test cache size limit From 07cec6f88288ea3d1fe83a5ff484aa00be914c11 Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Fri, 15 Aug 2025 02:51:02 +0800 Subject: [PATCH 05/10] Add Flytefile test and use mock path and fixture in the test case Signed-off-by: Barry Wu --- flytekit/core/type_engine.py | 16 +-- tests/flytekit/unit/core/test_list.py | 3 + tests/flytekit/unit/core/test_type_engine.py | 142 ++++++++++++------- 3 files changed, 102 insertions(+), 59 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index cab7cb5058..e06d268f8f 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -1383,22 +1383,16 @@ def calculate_hash(cls, python_val: typing.Any, python_type: Type[T]) -> Optiona def make_key(cls, python_val: typing.Any, python_type: Type[T]) -> Optional[tuple]: import cloudpickle - val_hash: typing.Any - type_hash: typing.Any + val_hash: int + type_hash: int try: - try: - val_hash = hash(python_val) - except Exception: - val_hash = hash(cloudpickle.dumps(python_val)) - - try: - type_hash = hash(python_type) - except Exception: - type_hash = hash(cloudpickle.dumps(python_type)) + val_hash = hash(cloudpickle.dumps(python_val)) + type_hash = hash(cloudpickle.dumps(python_type)) return (val_hash, type_hash) except Exception: + logger.warning(f"Could not hash python_val: {python_val} or python_type: {python_type}") return None @classmethod diff --git a/tests/flytekit/unit/core/test_list.py b/tests/flytekit/unit/core/test_list.py index c7f2f2d415..e254f7139b 100644 --- a/tests/flytekit/unit/core/test_list.py +++ b/tests/flytekit/unit/core/test_list.py @@ -84,4 +84,7 @@ async def test_coroutine_batching_of_list_transformer(): with pytest.raises(ValueError): TypeEngine.to_literal(ctx, python_val_2, typing.List[MyInt], lt) + with mock.patch("flytekit.core.type_engine._TYPE_ENGINE_COROS_BATCH_SIZE", 5): + TypeEngine.to_literal(ctx, python_val, typing.List[MyInt], lt) + del TypeEngine._REGISTRY[MyInt] diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index c884df6ff2..f65c32f7b2 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -3852,29 +3852,37 @@ async def test_dict_transformer_annotated_type(): literal3 = await TypeEngine.async_to_literal(ctx, nested_dict, nested_dict_type, expected_type) assert literal3.map.literals["outer"].map.literals["inner"].scalar.primitive.integer == 42 -def test_type_engine_cache(): - # Clear cache before test +@pytest.fixture(autouse=True) +def clear_type_engine_cache(): + """Clear TypeEngine cache before and after each test""" + TypeEngine._CACHE.clear() + yield TypeEngine._CACHE.clear() - # Test data +def test_type_engine_cache_with_list(): ctx = FlyteContext.current_context() python_val = [1, 2, 3, 4, 5] python_type = typing.List[int] expected = TypeEngine.to_literal_type(python_type) + list_transformer = TypeEngine.get_transformer(typing.List[int]) + with mock.patch.object(list_transformer, 'async_to_literal', + wraps=list_transformer.async_to_literal) as mock_async_to_literal: - # First call - should not use cache - literal1 = TypeEngine.to_literal(ctx, python_val, python_type, expected) + # First call + literal1 = TypeEngine.to_literal(ctx, python_val, python_type, expected) + assert mock_async_to_literal.call_count == 1 - # Verify cache is populated - key = TypeEngine.make_key(python_val, python_type) - assert key is not None - assert key in TypeEngine._CACHE + key = TypeEngine.make_key(python_val, python_type) + assert key is not None + assert key in TypeEngine._CACHE - # Second call with same parameters - should use cache - literal2 = TypeEngine.to_literal(ctx, python_val, python_type, expected) + # Second call with same DataFrame + literal2 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - # Verify both literals are identical (same object from cache) - assert literal1 is literal2 + # Verify async_to_literal was called + assert mock_async_to_literal.call_count == 1 + + assert literal1 is literal2 # Test with different data - should not use cache different_val = [2, 1, 3, 4, 5] @@ -3888,23 +3896,6 @@ def test_type_engine_cache(): # Verify different literals are different objects assert literal1 is not literal3 - # Test cache with unhashable objects - python_val = {"a": [1, 2, 3]} - python_type = typing.Dict[str, typing.List[int]] - expected = TypeEngine.to_literal_type(python_type) - - # First call - literal4 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - key = TypeEngine.make_key(python_val, python_type) - assert key is not None - assert key in TypeEngine._CACHE - - # Second call with same unhashable data - literal5 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - - # Should be the same object since unhashable objects will fallback to cloudpickle to hash - assert literal4 is literal5 - # Add many different values to test cache size limit for i in range(200): # More than the default maxsize of 128 test_val = [i, i+1, i+2] @@ -3915,30 +3906,30 @@ def test_type_engine_cache(): # Cache should not exceed maxsize assert len(TypeEngine._CACHE) == 128 - # Test cache with pandas DataFrame (if available) - try: - import pandas as pd - - # Create DataFrame - df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) - df_type = pd.DataFrame - df_expected = TypeEngine.to_literal_type(df_type) +def test_type_engine_cache_with_dict(): + ctx = FlyteContext.current_context() + python_val = {"a": [1, 2, 3]} + python_type = typing.Dict[str, typing.List[int]] + expected = TypeEngine.to_literal_type(python_type) + dict_transformer = TypeEngine.get_transformer(typing.Dict[str, typing.List[int]]) + with mock.patch.object(dict_transformer, 'async_to_literal', + wraps=dict_transformer.async_to_literal) as mock_async_to_literal: # First call - literal6 = TypeEngine.to_literal(ctx, df, df_type, df_expected) + literal1 = TypeEngine.to_literal(ctx, python_val, python_type, expected) + assert mock_async_to_literal.call_count == 1 - # Second call with same DataFrame - literal7 = TypeEngine.to_literal(ctx, df, df_type, df_expected) + key = TypeEngine.make_key(python_val, python_type) + assert key is not None + assert key in TypeEngine._CACHE - # Should be same object (DataFrame should be hashable via cloudpickle) - assert literal6 is literal7 + # Second call with same DataFrame + literal2 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - except ImportError: - # Skip pandas test if not available - pass + # Verify async_to_literal was called + assert mock_async_to_literal.call_count == 1 - # Clean up - TypeEngine._CACHE.clear() + assert literal1 is literal2 def test_make_key_with_annotated_types(): # Test with Annotated type @@ -3951,3 +3942,58 @@ def test_make_key_with_annotated_types(): assert key is not None assert key_without_annotation is not None assert key != key_without_annotation + +def test_type_engine_cache_with_pandas(): + import pandas as pd + ctx = FlyteContext.current_context() + # Create DataFrame + df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + df_type = pd.DataFrame + df_expected = TypeEngine.to_literal_type(df_type) + + # Get the transformer for DataFrame + df_transformer = TypeEngine._REGISTRY[pd.DataFrame] + + # Mock the async_to_literal method with wraps to track calls + with mock.patch.object(df_transformer, 'async_to_literal', + wraps=df_transformer.async_to_literal) as mock_async_to_literal: + + # First call + literal1 = TypeEngine.to_literal(ctx, df, df_type, df_expected) + assert mock_async_to_literal.call_count == 1 + + # Second call with same DataFrame + literal2 = TypeEngine.to_literal(ctx, df, df_type, df_expected) + + # Verify async_to_literal was called + assert mock_async_to_literal.call_count == 1 + + assert literal1 is literal2 + +def test_type_engine_cache_with_flytefile(): + + transformer = TypeEngine.get_transformer(FlyteFile) + ctx = FlyteContext.current_context() + + temp_dir = tempfile.mkdtemp(prefix="temp_example_") + file_path = os.path.join(temp_dir, "file.txt") + with open(file_path, "w") as file1: + file1.write("hello world") + + lt = TypeEngine.to_literal_type(FlyteFile) + + # Mock the file upload + with mock.patch.object(transformer, 'async_to_literal', + wraps=transformer.async_to_literal) as mock_async_to_literal: + + # Test 1: Upload local file to remote + lv1 = TypeEngine.to_literal(ctx, file_path, FlyteFile, lt) + assert mock_async_to_literal.call_count == 1 + + # Second call with same DataFrame + lv2 = TypeEngine.to_literal(ctx, file_path, FlyteFile, lt) + + # Verify async_to_literal was called + assert mock_async_to_literal.call_count == 1 + + assert lv1 is lv2 From 266caabcde6482b3c8cf4a6ccc1e63f069fbbdf1 Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Fri, 15 Aug 2025 14:41:34 +0800 Subject: [PATCH 06/10] Fix some nit Signed-off-by: Barry Wu --- tests/flytekit/unit/core/test_type_engine.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index f65c32f7b2..e0c21999ef 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -3870,7 +3870,6 @@ def test_type_engine_cache_with_list(): # First call literal1 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - assert mock_async_to_literal.call_count == 1 key = TypeEngine.make_key(python_val, python_type) assert key is not None @@ -3879,7 +3878,7 @@ def test_type_engine_cache_with_list(): # Second call with same DataFrame literal2 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - # Verify async_to_literal was called + # Verify async_to_literal was only called once assert mock_async_to_literal.call_count == 1 assert literal1 is literal2 @@ -3917,7 +3916,6 @@ def test_type_engine_cache_with_dict(): # First call literal1 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - assert mock_async_to_literal.call_count == 1 key = TypeEngine.make_key(python_val, python_type) assert key is not None @@ -3926,7 +3924,7 @@ def test_type_engine_cache_with_dict(): # Second call with same DataFrame literal2 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - # Verify async_to_literal was called + # Verify async_to_literal was only called once assert mock_async_to_literal.call_count == 1 assert literal1 is literal2 @@ -3960,7 +3958,6 @@ def test_type_engine_cache_with_pandas(): # First call literal1 = TypeEngine.to_literal(ctx, df, df_type, df_expected) - assert mock_async_to_literal.call_count == 1 # Second call with same DataFrame literal2 = TypeEngine.to_literal(ctx, df, df_type, df_expected) @@ -3988,7 +3985,6 @@ def test_type_engine_cache_with_flytefile(): # Test 1: Upload local file to remote lv1 = TypeEngine.to_literal(ctx, file_path, FlyteFile, lt) - assert mock_async_to_literal.call_count == 1 # Second call with same DataFrame lv2 = TypeEngine.to_literal(ctx, file_path, FlyteFile, lt) From 1576060a4ed848283e179258b7b36bbf3c665bb1 Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Tue, 19 Aug 2025 16:48:23 +0800 Subject: [PATCH 07/10] Fix unit_test_codecov-No module named 'pandas' Signed-off-by: Barry Wu --- tests/flytekit/unit/core/test_type_engine.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index e0c21999ef..57413622e6 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -3941,6 +3941,7 @@ def test_make_key_with_annotated_types(): assert key_without_annotation is not None assert key != key_without_annotation +@pytest.mark.skipif("pandas" not in sys.modules, reason="Pandas is not installed.") def test_type_engine_cache_with_pandas(): import pandas as pd ctx = FlyteContext.current_context() From f3877869efddfaf8a88eeb93db1b23af0cc59f03 Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Fri, 22 Aug 2025 15:14:47 +0800 Subject: [PATCH 08/10] Use importorskip to replace @pytest.mark.skipif Signed-off-by: Barry Wu --- tests/flytekit/unit/core/test_type_engine.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 57413622e6..a3800b634b 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -3941,9 +3941,8 @@ def test_make_key_with_annotated_types(): assert key_without_annotation is not None assert key != key_without_annotation -@pytest.mark.skipif("pandas" not in sys.modules, reason="Pandas is not installed.") def test_type_engine_cache_with_pandas(): - import pandas as pd + pd = pytest.importorskip("pandas") ctx = FlyteContext.current_context() # Create DataFrame df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) From 9597122d6186d8931f1954c7bd86b497bcf5887a Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 26 Aug 2025 21:42:05 -0700 Subject: [PATCH 09/10] Apply suggestion from @machichima Co-authored-by: Nary Yeh <60069744+machichima@users.noreply.github.com> --- tests/flytekit/unit/core/test_list.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/flytekit/unit/core/test_list.py b/tests/flytekit/unit/core/test_list.py index e254f7139b..75917c3a0c 100644 --- a/tests/flytekit/unit/core/test_list.py +++ b/tests/flytekit/unit/core/test_list.py @@ -84,6 +84,7 @@ async def test_coroutine_batching_of_list_transformer(): with pytest.raises(ValueError): TypeEngine.to_literal(ctx, python_val_2, typing.List[MyInt], lt) + # Cache hit for python_val prevents async_to_literal calls, avoiding the batch size limit of 2 error defined in MyIntAsyncTransformer with mock.patch("flytekit.core.type_engine._TYPE_ENGINE_COROS_BATCH_SIZE", 5): TypeEngine.to_literal(ctx, python_val, typing.List[MyInt], lt) From 8d65a9fe26d98d015354cff8021e6c3ac6129ea9 Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Thu, 28 Aug 2025 00:03:22 +0800 Subject: [PATCH 10/10] Fix some nit Signed-off-by: Barry Wu --- flytekit/core/type_engine.py | 12 +++++------ tests/flytekit/unit/core/test_type_engine.py | 22 ++++++++++---------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index e06d268f8f..2895147a06 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -1175,7 +1175,7 @@ class TypeEngine(typing.Generic[T]): _DATACLASS_TRANSFORMER: TypeTransformer = DataclassTransformer() # type: ignore _ENUM_TRANSFORMER: TypeTransformer = EnumTransformer() # type: ignore lazy_import_lock = threading.Lock() - _CACHE: LRUCache = LRUCache(maxsize=128) + _LITERAL_CACHE: LRUCache = LRUCache(maxsize=128) @classmethod def register( @@ -1380,7 +1380,7 @@ def calculate_hash(cls, python_val: typing.Any, python_type: Type[T]) -> Optiona return hsh @classmethod - def make_key(cls, python_val: typing.Any, python_type: Type[T]) -> Optional[tuple]: + def _get_literal_cache_key(cls, python_val: typing.Any, python_type: Type[T]) -> Optional[tuple]: import cloudpickle val_hash: int @@ -1404,9 +1404,9 @@ def to_literal( to_literal function, and allowing this to_literal function, to then invoke yet another async function, namely an async transformer. """ - key = cls.make_key(python_val, python_type) - if key is not None and key in cls._CACHE: - return cls._CACHE[key] + key = cls._get_literal_cache_key(python_val, python_type) + if key is not None and key in cls._LITERAL_CACHE: + return cls._LITERAL_CACHE[key] from flytekit.core.promise import Promise @@ -1430,7 +1430,7 @@ def to_literal( lv.hash = cls.calculate_hash(python_val, python_type) if key is not None: - cls._CACHE[key] = lv + cls._LITERAL_CACHE[key] = lv return lv diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index a3800b634b..7e04ab0214 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -3855,9 +3855,9 @@ async def test_dict_transformer_annotated_type(): @pytest.fixture(autouse=True) def clear_type_engine_cache(): """Clear TypeEngine cache before and after each test""" - TypeEngine._CACHE.clear() + TypeEngine._LITERAL_CACHE.clear() yield - TypeEngine._CACHE.clear() + TypeEngine._LITERAL_CACHE.clear() def test_type_engine_cache_with_list(): ctx = FlyteContext.current_context() @@ -3871,9 +3871,9 @@ def test_type_engine_cache_with_list(): # First call literal1 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - key = TypeEngine.make_key(python_val, python_type) + key = TypeEngine._get_literal_cache_key(python_val, python_type) assert key is not None - assert key in TypeEngine._CACHE + assert key in TypeEngine._LITERAL_CACHE # Second call with same DataFrame literal2 = TypeEngine.to_literal(ctx, python_val, python_type, expected) @@ -3886,11 +3886,11 @@ def test_type_engine_cache_with_list(): # Test with different data - should not use cache different_val = [2, 1, 3, 4, 5] literal3 = TypeEngine.to_literal(ctx, different_val, python_type, expected) - key_different = TypeEngine.make_key(different_val, python_type) + key_different = TypeEngine._get_literal_cache_key(different_val, python_type) assert key_different is not key assert key_different is not None - assert key_different in TypeEngine._CACHE + assert key_different in TypeEngine._LITERAL_CACHE # Verify different literals are different objects assert literal1 is not literal3 @@ -3903,7 +3903,7 @@ def test_type_engine_cache_with_list(): TypeEngine.to_literal(ctx, test_val, test_type, test_expected) # Cache should not exceed maxsize - assert len(TypeEngine._CACHE) == 128 + assert len(TypeEngine._LITERAL_CACHE) == 128 def test_type_engine_cache_with_dict(): ctx = FlyteContext.current_context() @@ -3917,9 +3917,9 @@ def test_type_engine_cache_with_dict(): # First call literal1 = TypeEngine.to_literal(ctx, python_val, python_type, expected) - key = TypeEngine.make_key(python_val, python_type) + key = TypeEngine._get_literal_cache_key(python_val, python_type) assert key is not None - assert key in TypeEngine._CACHE + assert key in TypeEngine._LITERAL_CACHE # Second call with same DataFrame literal2 = TypeEngine.to_literal(ctx, python_val, python_type, expected) @@ -3934,8 +3934,8 @@ def test_make_key_with_annotated_types(): annotated_val = [1, 2, 3] annotated_type = typing.Annotated[typing.List[int], "test_annotation"] - key = TypeEngine.make_key(annotated_val, annotated_type) - key_without_annotation = TypeEngine.make_key(annotated_val, typing.List[int]) + key = TypeEngine._get_literal_cache_key(annotated_val, annotated_type) + key_without_annotation = TypeEngine._get_literal_cache_key(annotated_val, typing.List[int]) # Should handle Annotated types correctly assert key is not None assert key_without_annotation is not None