genesis-flow 1.0.6__py3-none-any.whl → 1.0.8__py3-none-any.whl

{genesis_flow-1.0.6.dist-info → genesis_flow-1.0.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: genesis-flow
-Version: 1.0.6
+Version: 1.0.8
 Summary: Genesis-Flow: MLflow v3.1.4 compatible fork for Genesis platform
 Maintainer-email: Databricks <mlflow-oss-maintainers@googlegroups.com>
 License: Copyright 2018 Databricks, Inc.  All rights reserved.
@@ -378,6 +378,10 @@ poetry install --with dev
 
 ```python
 import mlflow
+import os
+
+# Optional: Disable secure model loading if you encounter loading issues
+# os.environ["MLFLOW_ENABLE_SECURE_MODEL_LOADING"] = "false"
 
 # Set tracking URI (supports file, PostgreSQL, etc.)
 mlflow.set_tracking_uri("file:///path/to/mlruns")
@@ -504,7 +508,7 @@ export MLFLOW_POSTGRES_USERNAME="user@tenant"
 export GOOGLE_APPLICATION_CREDENTIALS="/path/to/service-account.json"
 
 # Security configuration
-export MLFLOW_ENABLE_SECURE_MODEL_LOADING=true
+export MLFLOW_ENABLE_SECURE_MODEL_LOADING=true  # Can be set to false if model loading fails
 export MLFLOW_STRICT_INPUT_VALIDATION=true
 ```
 
@@ -761,11 +765,32 @@ entry_points = {
 
 - ✅ **Input validation** against injection attacks
 - ✅ **Path traversal protection** for file operations  
-- ✅ **Secure pickle loading** with restricted unpickling
+- ✅ **Secure pickle loading** with restricted unpickling (configurable via `MLFLOW_ENABLE_SECURE_MODEL_LOADING`)
 - ✅ **Authentication hooks** for enterprise SSO integration
 - ✅ **Audit logging** for compliance requirements
 - ✅ **Encrypted communication** support
 
+#### Secure Model Loading Configuration
+
+Genesis-Flow includes enhanced security for model loading that prevents arbitrary code execution from untrusted pickle files. This feature can be configured using the `MLFLOW_ENABLE_SECURE_MODEL_LOADING` environment variable:
+
+```bash
+# Enable secure model loading (default - recommended for production)
+export MLFLOW_ENABLE_SECURE_MODEL_LOADING=true
+
+# Disable secure model loading (use with caution - only for trusted models)
+export MLFLOW_ENABLE_SECURE_MODEL_LOADING=false
+```
+
+**When to disable secure loading:**
+- When loading legacy models that contain custom classes not in the allowlist
+- During development with trusted model sources
+- When migrating from standard MLflow with complex custom models
+
+**Security implications:**
+- When enabled: Only whitelisted classes can be deserialized, preventing arbitrary code execution
+- When disabled: Standard cloudpickle behavior, any Python object can be loaded (potential security risk)
+
 ### Security Best Practices
 
 1. **Use MongoDB authentication** in production

{genesis_flow-1.0.6.dist-info → genesis_flow-1.0.8.dist-info}/RECORD

@@ -1,4 +1,4 @@
-genesis_flow-1.0.6.dist-info/licenses/LICENSE.txt,sha256=Y5U1Xebzka__NZlqMPtBsYm0mRpMtUmTrONatpoL-ig,11382
+genesis_flow-1.0.8.dist-info/licenses/LICENSE.txt,sha256=Y5U1Xebzka__NZlqMPtBsYm0mRpMtUmTrONatpoL-ig,11382
 mlflow/__init__.py,sha256=-_r__N5Afed81pLVtr2wKbHQIA0aj9u9n_7kWGxLWi4,11194
 mlflow/__main__.py,sha256=_PcdoxKehR_a2MI6GqBfzYzRCXZhVyDCSdbxDWVlWd4,39
 mlflow/cli.py,sha256=f1ObrWZ03HgRiRoVEE1Gffe-dGcSY7CxJyEFgb5VUMM,26137
@@ -311,7 +311,7 @@ mlflow/pyfunc/backend.py,sha256=5-tjUnY1nPecoNCxFjhLjPbhHsDGR9_0P9ym02_J3HA,2077
 mlflow/pyfunc/context.py,sha256=nXQcP61XR7cKx6CeFsdAM4TD-NOrIEt_GQSfJ990xc4,3007
 mlflow/pyfunc/dbconnect_artifact_cache.py,sha256=-Ji_GU-NLvGyxYrOFhR81SYn4E6ykiFcZl35hRQ9c6A,5769
 mlflow/pyfunc/mlserver.py,sha256=ER7tXcOZ-tAUblMbgqPVT2E8BtBGH1HXLSSOmsQUB_c,1271
-mlflow/pyfunc/model.py,sha256=aDwfNKgoVw4sl9W0WZfOa9st5WTsSF3NWbF5TgM3bDs,64507
+mlflow/pyfunc/model.py,sha256=by6xBLOmu1D1B0sH9B2LpVe6TjS14D0sStE60iJ9ofk,65686
 mlflow/pyfunc/spark_model_cache.py,sha256=sFFAi-LxXEJdqk8szm32XxK5IFsISpxxVjAWjWbHZ1U,2091
 mlflow/pyfunc/stdin_server.py,sha256=gNkWxlCz3KLu6jvZNeNytmKAy4j5gIXOpIehrizm43s,1362
 mlflow/pyfunc/loaders/__init__.py,sha256=W3Bny093PoREOCqavIzKJ3lYeZUBXhxk3EDkCqc_lBc,276
@@ -613,7 +613,7 @@ mlflow/utils/requirements_utils.py,sha256=79BL9iFqg9H4VRYiqk-gTLlrfLyUY7KVG4BDV8
 mlflow/utils/rest_utils.py,sha256=R5MZk0Ck358VeW-DKqpfSyDHbWjg_slQTpRfr4ZLDLM,26664
 mlflow/utils/search_logged_model_utils.py,sha256=JWczQxogo4u9Gr3gfoSVf7tVny1e-rxfAchh3Vqy6RU,4289
 mlflow/utils/search_utils.py,sha256=PgXtwlViHvmahY63otykh5iJVW0PR9DIwetrvCCBw2U,86142
-mlflow/utils/secure_loading.py,sha256=yJpKwrhzINmo0yVdp6fXvbNkZu_HEdwyB8deljGDXxk,10095
+mlflow/utils/secure_loading.py,sha256=X4AdKHDkIzMvriFwCvTPLhN-yPbqs6dgUOB244Q_gRM,11892
 mlflow/utils/security_validation.py,sha256=cLBrLNOoVJuAffO677v0m_v4i-01v7aLk-sJk87mpkQ,12830
 mlflow/utils/server_cli_utils.py,sha256=gbT5CVkOLorSw-y8bnNSBEP9mClcVTagtTN5SK5Azhw,2381
 mlflow/utils/spark_utils.py,sha256=zUbQIRAtwyU3rDJK8m7v42KO2TSGn28Coe_UYBEhIWA,395
@@ -641,8 +641,8 @@ mlflow/utils/autologging_utils/metrics_queue.py,sha256=bwpMX7Go6xFxrpYROi6rDBdVt
 mlflow/utils/autologging_utils/safety.py,sha256=IwbTbusyE87Hc4qkhhvMikoaZqX9kpr972FWS2B8goc,51465
 mlflow/utils/autologging_utils/versioning.py,sha256=2hSN4KXFWEJCcopDdLG6BiPeqSoqjETeNMuUBsCCwlI,3762
 mlflow/utils/import_hooks/__init__.py,sha256=werje98Woelkbwrhtlb8wmRdt3RtiL--LqGru7Xh3YU,13589
-genesis_flow-1.0.6.dist-info/METADATA,sha256=_ofi_wrmOHLBu-0gjHwRW97dMrrVpnksdsLhUr1WrlU,33401
-genesis_flow-1.0.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-genesis_flow-1.0.6.dist-info/entry_points.txt,sha256=YZZiTHSYQpv8vou6gdqWI17piK9Pm0P3BmZ-xepUtnw,449
-genesis_flow-1.0.6.dist-info/top_level.txt,sha256=wm8UqYyUHI21EvrTDHb3eYICy0dOVDLBhAL-jp5zbuI,7
-genesis_flow-1.0.6.dist-info/RECORD,,
+genesis_flow-1.0.8.dist-info/METADATA,sha256=tam5UxyijS6-7TN_92L9-e7BHQKsKtpyNoMBhCCNGLw,34630
+genesis_flow-1.0.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+genesis_flow-1.0.8.dist-info/entry_points.txt,sha256=YZZiTHSYQpv8vou6gdqWI17piK9Pm0P3BmZ-xepUtnw,449
+genesis_flow-1.0.8.dist-info/top_level.txt,sha256=wm8UqYyUHI21EvrTDHb3eYICy0dOVDLBhAL-jp5zbuI,7
+genesis_flow-1.0.8.dist-info/RECORD,,

mlflow/pyfunc/model.py

@@ -819,27 +819,52 @@ def _maybe_compress_cloudpickle_dump(python_model, path, compression):
 
 def _maybe_decompress_cloudpickle_load(path, compression):
     """
-    Genesis-Flow: Secure model loading with safety checks.
+    Genesis-Flow: Model loading with optional security checks.
+    
+    Security can be disabled by setting MLFLOW_ENABLE_SECURE_MODEL_LOADING=false
+    Default behavior is secure loading (enabled).
     """
-    from mlflow.utils.secure_loading import SecureModelLoader, SecurityError
+    import os
+    import cloudpickle
     
     _check_compression_supported(compression)
     
-    # For compressed files, we need to decompress first then load securely
-    if compression and compression != "none":
-        import tempfile
-        file_open = _COMPRESSION_INFO.get(compression, {}).get("open", open)
-        with file_open(path, "rb") as compressed_f:
-            with tempfile.NamedTemporaryFile(delete=False) as temp_f:
-                temp_f.write(compressed_f.read())
-                temp_path = temp_f.name
-        try:
-            return SecureModelLoader.safe_cloudpickle_load(temp_path)
-        finally:
-            os.unlink(temp_path)
+    # Check if secure loading should be used (default: true for security)
+    use_secure_loading = os.getenv('MLFLOW_ENABLE_SECURE_MODEL_LOADING', 'true').lower() != 'false'
+    
+    if use_secure_loading:
+        # Use secure loading for enhanced security
+        from mlflow.utils.secure_loading import SecureModelLoader, SecurityError
+        
+        # For compressed files, we need to decompress first then load securely
+        if compression and compression != "none":
+            import tempfile
+            file_open = _COMPRESSION_INFO.get(compression, {}).get("open", open)
+            with file_open(path, "rb") as compressed_f:
+                with tempfile.NamedTemporaryFile(delete=False) as temp_f:
+                    temp_f.write(compressed_f.read())
+                    temp_path = temp_f.name
+            try:
+                return SecureModelLoader.safe_cloudpickle_load(temp_path)
+            finally:
+                os.unlink(temp_path)
+        else:
+            # Direct secure loading for uncompressed files
+            return SecureModelLoader.safe_cloudpickle_load(path)
     else:
-        # Direct secure loading for uncompressed files
-        return SecureModelLoader.safe_cloudpickle_load(path)
+        # Use standard cloudpickle loading (original MLflow behavior)
+        # Warning: This bypasses security checks and may allow arbitrary code execution
+        import logging
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            "Secure model loading is disabled. Models will be loaded without security checks. "
+            "This may allow arbitrary code execution. Set MLFLOW_ENABLE_SECURE_MODEL_LOADING=true "
+            "to enable secure loading."
+        )
+        
+        file_open = _COMPRESSION_INFO.get(compression, {}).get("open", open)
+        with file_open(path, "rb") as f:
+            return cloudpickle.load(f)
 
 
 if IS_PYDANTIC_V2_OR_NEWER:

mlflow/utils/secure_loading.py

@@ -17,6 +17,10 @@ logger = logging.getLogger(__name__)
 
 # Allowlist of safe classes that can be unpickled
 SAFE_PICKLE_CLASSES = {
+
+    # STD types
+    'time.time',
+
     # NumPy types
     'numpy.ndarray',
     'numpy.dtype',
@@ -42,7 +46,7 @@ SAFE_PICKLE_CLASSES = {
     # Built-in types
     'builtins.dict', 'builtins.list', 'builtins.tuple', 'builtins.set',
     'builtins.str', 'builtins.int', 'builtins.float', 'builtins.bool',
-    'builtins.type',
+    'builtins.type', 'builtins.object',
 
     # Collections
     'collections.OrderedDict',
@@ -62,6 +66,15 @@ SAFE_PICKLE_CLASSES = {
     'cloudpickle.cloudpickle._function_setstate',
     'cloudpickle.cloudpickle._make_empty_cell',
     'cloudpickle.cloudpickle._make_cell',
+    'cloudpickle.cloudpickle.subimport',
+    'cloudpickle.cloudpickle_fast._class_setstate',
+    'cloudpickle.cloudpickle_fast._function_setstate',
+    'cloudpickle.cloudpickle_fast._builtin_type',
+    'cloudpickle.cloudpickle_fast._make_skeleton_class',
+    'cloudpickle.cloudpickle_fast._make_function',
+    'cloudpickle.cloudpickle_fast._make_empty_cell',
+    'cloudpickle.cloudpickle_fast._make_cell',
+    'cloudpickle.cloudpickle_fast.subimport',
 
     # Sentence Transformers
     'sentence_transformers.SentenceTransformer.SentenceTransformer',
@@ -103,6 +116,49 @@ SAFE_PICKLE_CLASSES = {
     'tokenizers.Tokenizer',
     'tokenizers.models.Model',
     'tokenizers.AddedToken',
+
+    # PyTorch
+    'model.load_checkpoint',
+    'torch.device',
+    'model.SepClassifier',
+    
+    # Generic model module classes (commonly used in custom ML models)
+    # These are general enough to be safe but specific to model loading
+    'model.Model',
+    'model.BaseModel',
+    'model.PythonModel',
+    'model.MLModel',
+    'model.NeuralNetwork',
+    'model.Classifier',
+    'model.Regressor',
+    'model.Predictor',
+    'model.Estimator',
+    'model.Pipeline',
+    'model.Transformer',
+    'model.Encoder',
+    'model.Decoder',
+    'model.Generator',
+    'model.Discriminator',
+    'model.load_model',
+    'model.save_model',
+    'model.predict',
+    'model.transform',
+    'model.fit',
+    'model.evaluate',
+    'model.train',
+    'model.inference',
+    'model.forward',
+    'model.backward',
+    
+    # Additional model-related common class names
+    'model.ModelConfig',
+    'model.ModelWrapper',
+    'model.ModelFactory',
+    'model.ModelRegistry',
+    'model.ModelLoader',
+    'model.ModelState',
+    'model.ModelCheckpoint',
+    'model.ModelArtifact',
 }
 
 
@@ -144,11 +200,14 @@ class RestrictedUnpickler(pickle.Unpickler):
             'numpy': ['ndarray', 'dtype', 'int32', 'int64', 'float32', 'float64', 'bool_'],
             'pandas': ['DataFrame', 'Series', 'Index'],
             'builtins': ['dict', 'list', 'tuple', 'set', 'str', 'int', 'float', 'bool'],
+            'model': ['*'],  # Allow all classes from 'model' module for custom models
         }
         
-        if module in safe_modules and name in safe_modules[module]:
-            logger.debug(f"Loading safe module class: {full_name}")
-            return super().find_class(module, name)
+        if module in safe_modules:
+            # Check if all classes are allowed ('*') or if specific class is in list
+            if safe_modules[module] == ['*'] or name in safe_modules[module]:
+                logger.debug(f"Loading safe module class: {full_name}")
+                return super().find_class(module, name)
         
         # Log and block unsafe class
         logger.warning(f"Blocked potentially unsafe class: {full_name}")