code-loader 1.0.127__tar.gz → 1.0.128__tar.gz

{code_loader-1.0.127 → code_loader-1.0.128}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: code-loader
-Version: 1.0.127
+Version: 1.0.128
 Summary: 
 Home-page: https://github.com/tensorleap/code-loader
 License: MIT

{code_loader-1.0.127 → code_loader-1.0.128}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -1,5 +1,6 @@
 # mypy: ignore-errors
 import os
+from functools import lru_cache
 from typing import Optional, Union, Callable, List, Dict
 
 import numpy as np
@@ -21,6 +22,7 @@ import inspect
 import functools
 
 _called_from_inside_tl_decorator = 0
+_called_from_inside_tl_integration_test_decorator = False
 
 
 def _add_mapping_connection(user_unique_name, connection_destinations, arg_names, name, node_mapping_type):
@@ -46,17 +48,22 @@ def integration_test():
         leap_binder.integration_test_func = integration_test_function
 
         def inner(*args, **kwargs):
-            ret = integration_test_function(*args, **kwargs)
-
+            global _called_from_inside_tl_integration_test_decorator
             try:
-                os.environ[mapping_runtime_mode_env_var_mame] = 'True'
-                integration_test_function(None, None)
-            except Exception as e:
-                print(f'Error during integration test: Make sure to disable any non tensorleap decorators '
-                      f'functions before pushing a new TL version')
+                _called_from_inside_tl_integration_test_decorator = True
+                ret = integration_test_function(*args, **kwargs)
+
+                try:
+                    os.environ[mapping_runtime_mode_env_var_mame] = 'True'
+                    integration_test_function(None, None)
+                except Exception as e:
+                    print(f'Error during integration test: Make sure to disable any non tensorleap decorators '
+                          f'functions before pushing a new TL version')
+                finally:
+                    if mapping_runtime_mode_env_var_mame in os.environ:
+                        del os.environ[mapping_runtime_mode_env_var_mame]
             finally:
-                if mapping_runtime_mode_env_var_mame in os.environ:
-                    del os.environ[mapping_runtime_mode_env_var_mame]
+                _called_from_inside_tl_integration_test_decorator = False
 
         return inner
 
@@ -71,6 +78,7 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
         class TempMapping:
             pass
 
+        @lru_cache()
         def inner():
             class ModelPlaceholder:
                 def __init__(self):
@@ -762,7 +770,7 @@ def tensorleap_input_encoder(name: str, channel_dim=-1, model_input_index=None):
 
             _validate_result(result)
 
-            if _called_from_inside_tl_decorator == 0:
+            if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
                 result = np.expand_dims(result, axis=0)
 
             return result
@@ -846,7 +854,7 @@ def tensorleap_gt_encoder(name: str):
 
             _validate_result(result)
 
-            if _called_from_inside_tl_decorator == 0:
+            if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
                 result = np.expand_dims(result, axis=0)
 
             return result

{code_loader-1.0.127 → code_loader-1.0.128}/code_loader/plot_functions/plot_functions.py

@@ -29,7 +29,6 @@ def run_only_on_non_mapping_mode():
     def decorator(func):
         def wrapper(*args, **kwargs):
             if os.environ.get(mapping_runtime_mode_env_var_mame):
-                print(f"Skipping {func.__name__} in mapping mode.")
                 return
             return func(*args, **kwargs)
         return wrapper

{code_loader-1.0.127 → code_loader-1.0.128}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "code-loader"
-version = "1.0.127"
+version = "1.0.128"
 description = ""
 authors = ["dorhar <doron.harnoy@tensorleap.ai>"]
 license = "MIT"