genesis-flow 1.0.0__py3-none-any.whl

mlflow/plugins/registry.py

@@ -0,0 +1,316 @@
+"""
+Plugin Registry for Genesis-Flow
+
+Central registry for tracking discovered and registered plugins.
+Provides thread-safe registration, lookup, and management of plugin classes.
+"""
+
+import logging
+import threading
+from typing import Dict, List, Optional, Type, Set
+
+from mlflow.plugins.base import BasePlugin, PluginType
+
+logger = logging.getLogger(__name__)
+
+class PluginRegistry:
+    """
+    Thread-safe registry for plugin classes.
+    
+    Maintains a central catalog of all discovered plugin classes,
+    their types, and metadata.
+    """
+    
+    def __init__(self):
+        self._plugins: Dict[str, Type[BasePlugin]] = {}
+        self._plugins_by_type: Dict[PluginType, Set[str]] = {}
+        self._lock = threading.RLock()
+        
+        # Initialize type sets
+        for plugin_type in PluginType:
+            self._plugins_by_type[plugin_type] = set()
+    
+    def register_plugin_class(self, name: str, plugin_class: Type[BasePlugin]) -> bool:
+        """
+        Register a plugin class.
+        
+        Args:
+            name: Plugin name
+            plugin_class: Plugin class
+            
+        Returns:
+            True if registered successfully, False otherwise
+        """
+        with self._lock:
+            try:
+                # Validate plugin class
+                if not self._validate_plugin_class(plugin_class):
+                    logger.error(f"Invalid plugin class for {name}")
+                    return False
+                
+                # Check for conflicts
+                if name in self._plugins:
+                    existing_class = self._plugins[name]
+                    if existing_class != plugin_class:
+                        logger.warning(f"Plugin {name} already registered with different class")
+                        # Allow override but log it
+                    else:
+                        logger.debug(f"Plugin {name} already registered with same class")
+                        return True
+                
+                # Register the plugin
+                self._plugins[name] = plugin_class
+                
+                # Add to type index
+                plugin_type = self._get_plugin_type(plugin_class)
+                if plugin_type:
+                    self._plugins_by_type[plugin_type].add(name)
+                
+                logger.debug(f"Registered plugin: {name}")
+                return True
+                
+            except Exception as e:
+                logger.error(f"Error registering plugin {name}: {e}")
+                return False
+    
+    def unregister_plugin(self, name: str) -> bool:
+        """
+        Unregister a plugin class.
+        
+        Args:
+            name: Plugin name
+            
+        Returns:
+            True if unregistered successfully, False otherwise
+        """
+        with self._lock:
+            try:
+                if name not in self._plugins:
+                    logger.warning(f"Plugin {name} not registered")
+                    return True  # Already unregistered
+                
+                plugin_class = self._plugins[name]
+                
+                # Remove from main registry
+                del self._plugins[name]
+                
+                # Remove from type index
+                plugin_type = self._get_plugin_type(plugin_class)
+                if plugin_type and name in self._plugins_by_type[plugin_type]:
+                    self._plugins_by_type[plugin_type].remove(name)
+                
+                logger.debug(f"Unregistered plugin: {name}")
+                return True
+                
+            except Exception as e:
+                logger.error(f"Error unregistering plugin {name}: {e}")
+                return False
+    
+    def get_plugin_class(self, name: str) -> Optional[Type[BasePlugin]]:
+        """
+        Get a plugin class by name.
+        
+        Args:
+            name: Plugin name
+            
+        Returns:
+            Plugin class or None if not found
+        """
+        with self._lock:
+            return self._plugins.get(name)
+    
+    def is_plugin_registered(self, name: str) -> bool:
+        """
+        Check if a plugin is registered.
+        
+        Args:
+            name: Plugin name
+            
+        Returns:
+            True if registered, False otherwise
+        """
+        with self._lock:
+            return name in self._plugins
+    
+    def list_plugins(self, plugin_type: Optional[PluginType] = None) -> List[str]:
+        """
+        List registered plugin names.
+        
+        Args:
+            plugin_type: Optional filter by plugin type
+            
+        Returns:
+            List of plugin names
+        """
+        with self._lock:
+            if plugin_type:
+                return list(self._plugins_by_type.get(plugin_type, set()))
+            else:
+                return list(self._plugins.keys())
+    
+    def get_plugins_by_type(self, plugin_type: PluginType) -> Dict[str, Type[BasePlugin]]:
+        """
+        Get all plugins of a specific type.
+        
+        Args:
+            plugin_type: Plugin type
+            
+        Returns:
+            Dict mapping plugin names to classes
+        """
+        with self._lock:
+            result = {}
+            plugin_names = self._plugins_by_type.get(plugin_type, set())
+            for name in plugin_names:
+                if name in self._plugins:
+                    result[name] = self._plugins[name]
+            return result
+    
+    def get_plugin_types(self) -> Dict[str, PluginType]:
+        """
+        Get plugin types for all registered plugins.
+        
+        Returns:
+            Dict mapping plugin names to types
+        """
+        with self._lock:
+            result = {}
+            for name, plugin_class in self._plugins.items():
+                plugin_type = self._get_plugin_type(plugin_class)
+                if plugin_type:
+                    result[name] = plugin_type
+            return result
+    
+    def clear(self):
+        """Clear all registered plugins."""
+        with self._lock:
+            self._plugins.clear()
+            for plugin_type in PluginType:
+                self._plugins_by_type[plugin_type].clear()
+            logger.info("Plugin registry cleared")
+    
+    def get_stats(self) -> Dict[str, int]:
+        """
+        Get registry statistics.
+        
+        Returns:
+            Statistics dictionary
+        """
+        with self._lock:
+            stats = {
+                "total": len(self._plugins),
+            }
+            
+            # Count by type
+            for plugin_type in PluginType:
+                count = len(self._plugins_by_type[plugin_type])
+                stats[plugin_type.value] = count
+            
+            return stats
+    
+    def _validate_plugin_class(self, plugin_class: Type) -> bool:
+        """
+        Validate a plugin class.
+        
+        Args:
+            plugin_class: Class to validate
+            
+        Returns:
+            True if valid, False otherwise
+        """
+        try:
+            # Must be a class
+            if not isinstance(plugin_class, type):
+                return False
+            
+            # Must be a subclass of BasePlugin
+            if not issubclass(plugin_class, BasePlugin):
+                return False
+            
+            # Must not be BasePlugin itself
+            if plugin_class == BasePlugin:
+                return False
+            
+            # Must have required methods
+            required_methods = ['load', 'enable', 'disable', 'unload']
+            for method in required_methods:
+                if not hasattr(plugin_class, method):
+                    return False
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"Error validating plugin class: {e}")
+            return False
+    
+    def _get_plugin_type(self, plugin_class: Type[BasePlugin]) -> Optional[PluginType]:
+        """
+        Determine the type of a plugin class.
+        
+        Args:
+            plugin_class: Plugin class
+            
+        Returns:
+            Plugin type or None if cannot determine
+        """
+        try:
+            # Try to get from class metadata
+            if hasattr(plugin_class, 'get_default_metadata'):
+                metadata = plugin_class.get_default_metadata()
+                return metadata.plugin_type
+            
+            # Try to infer from class hierarchy
+            from mlflow.plugins.base import (
+                FrameworkPlugin, LoggingPlugin, ModelRegistryPlugin,
+                DeploymentPlugin, ArtifactPlugin
+            )
+            
+            if issubclass(plugin_class, FrameworkPlugin):
+                return PluginType.FRAMEWORK
+            elif issubclass(plugin_class, LoggingPlugin):
+                return PluginType.LOGGING
+            elif issubclass(plugin_class, ModelRegistryPlugin):
+                return PluginType.MODEL_REGISTRY
+            elif issubclass(plugin_class, DeploymentPlugin):
+                return PluginType.DEPLOYMENT
+            elif issubclass(plugin_class, ArtifactPlugin):
+                return PluginType.ARTIFACT
+            else:
+                return PluginType.CUSTOM
+                
+        except Exception as e:
+            logger.error(f"Error determining plugin type for {plugin_class}: {e}")
+            return PluginType.CUSTOM
+    
+    def export_registry(self) -> Dict[str, Dict]:
+        """
+        Export registry contents for serialization.
+        
+        Returns:
+            Serializable registry data
+        """
+        with self._lock:
+            result = {}
+            for name, plugin_class in self._plugins.items():
+                result[name] = {
+                    "class_name": plugin_class.__name__,
+                    "module": plugin_class.__module__,
+                    "type": self._get_plugin_type(plugin_class).value if self._get_plugin_type(plugin_class) else None,
+                }
+            return result
+    
+    def find_plugins_by_pattern(self, pattern: str) -> List[str]:
+        """
+        Find plugins matching a name pattern.
+        
+        Args:
+            pattern: Pattern to match (supports * wildcard)
+            
+        Returns:
+            List of matching plugin names
+        """
+        import fnmatch
+        
+        with self._lock:
+            return [name for name in self._plugins.keys() 
+                   if fnmatch.fnmatch(name, pattern)]

mlflow/plugins/templates/framework_plugin_template.py

@@ -0,0 +1,329 @@
+"""
+Framework Plugin Template for Genesis-Flow
+
+This template provides a starting point for creating new ML framework plugins.
+Copy this file and customize it for your specific framework.
+
+Example: Creating a TensorFlow plugin
+1. Copy this file to tensorflow_plugin.py
+2. Replace "MyFramework" with "TensorFlow"
+3. Update metadata with TensorFlow-specific information
+4. Implement TensorFlow-specific methods
+5. Register in entry points or built-in plugins
+"""
+
+import logging
+from typing import Dict, Callable
+
+from mlflow.plugins.base import FrameworkPlugin, PluginMetadata, PluginType
+
+logger = logging.getLogger(__name__)
+
+class MyFrameworkPlugin(FrameworkPlugin):
+    """
+    Framework plugin template for Genesis-Flow.
+    
+    Replace "MyFramework" with your framework name throughout this file.
+    """
+    
+    @classmethod
+    def get_default_metadata(cls) -> PluginMetadata:
+        """
+        Get default metadata for the framework plugin.
+        
+        Customize all fields below for your framework.
+        """
+        return PluginMetadata(
+            name="myframework",  # Change to your framework name (lowercase)
+            version="1.0.0",     # Your plugin version
+            description="My Framework integration for Genesis-Flow",  # Framework description
+            author="Your Name or Organization",  # Plugin author
+            plugin_type=PluginType.FRAMEWORK,   # Keep as FRAMEWORK
+            dependencies=["myframework"],        # Required Python packages
+            optional_dependencies=[              # Optional packages that enhance functionality
+                "numpy", 
+                "scipy",
+                # Add framework-specific optional dependencies
+            ],
+            min_genesis_flow_version="1.0.0",   # Minimum Genesis-Flow version
+            max_genesis_flow_version=None,      # Maximum version (None for no limit)
+            homepage="https://myframework.org", # Framework homepage
+            documentation="https://docs.myframework.org/mlflow",  # MLflow integration docs
+            license="Apache 2.0",               # Framework license
+            tags=[                               # Descriptive tags
+                "machine-learning", 
+                "myframework",
+                # Add framework-specific tags like "deep-learning", "nlp", etc.
+            ],
+        )
+    
+    def get_module_path(self) -> str:
+        """
+        Get the module path for framework integration.
+        
+        This should match the MLflow integration module for your framework.
+        Example: "mlflow.tensorflow" for TensorFlow
+        """
+        return "mlflow.myframework"  # Change to your framework module
+    
+    def get_autolog_functions(self) -> Dict[str, Callable]:
+        """
+        Get autologging functions provided by this framework.
+        
+        Returns dictionary mapping function names to callables.
+        """
+        try:
+            # Import your framework's MLflow integration
+            from mlflow.myframework import autolog
+            
+            return {
+                "autolog": autolog,
+                # Add other autolog-related functions if available
+            }
+            
+        except ImportError:
+            logger.warning("MyFramework autolog not available")
+            return {}
+    
+    def get_save_functions(self) -> Dict[str, Callable]:
+        """
+        Get model saving functions provided by this framework.
+        
+        Returns dictionary mapping function names to callables.
+        """
+        try:
+            # Import your framework's MLflow integration
+            from mlflow.myframework import log_model, save_model
+            
+            return {
+                "log_model": log_model,
+                "save_model": save_model,
+                # Add other save-related functions if available
+            }
+            
+        except ImportError:
+            logger.warning("MyFramework save functions not available")
+            return {}
+    
+    def get_load_functions(self) -> Dict[str, Callable]:
+        """
+        Get model loading functions provided by this framework.
+        
+        Returns dictionary mapping function names to callables.
+        """
+        try:
+            # Import your framework's MLflow integration
+            from mlflow.myframework import load_model
+            
+            return {
+                "load_model": load_model,
+                # Add other load-related functions if available
+            }
+            
+        except ImportError:
+            logger.warning("MyFramework load functions not available")
+            return {}
+    
+    def enable(self) -> bool:
+        """
+        Enable the framework plugin with enhanced functionality.
+        
+        Add framework-specific setup and configuration here.
+        """
+        try:
+            # Call parent enable method
+            if not super().enable():
+                return False
+            
+            # Add framework-specific setup
+            self._setup_framework_environment()
+            
+            self._logger.info("MyFramework plugin enabled with full functionality")
+            return True
+            
+        except Exception as e:
+            self._logger.error(f"Failed to enable MyFramework plugin: {e}")
+            return False
+    
+    def _setup_framework_environment(self):
+        """
+        Setup framework-specific environment and configurations.
+        
+        Customize this method for your framework's initialization needs.
+        """
+        try:
+            # Import your framework
+            import myframework  # Replace with your framework import
+            
+            # Log framework version and capabilities
+            self._logger.info(f"MyFramework version: {myframework.__version__}")
+            
+            # Check for GPU availability if relevant
+            if hasattr(myframework, 'is_gpu_available'):
+                gpu_available = myframework.is_gpu_available()
+                self._logger.info(f"GPU available: {gpu_available}")
+            
+            # Setup hooks for automatic logging
+            self.register_hook("model_save", self._on_model_save)
+            self.register_hook("model_load", self._on_model_load)
+            self.register_hook("training_start", self._on_training_start)
+            self.register_hook("training_end", self._on_training_end)
+            
+            # Framework-specific configuration
+            self._configure_framework_settings()
+            
+        except Exception as e:
+            self._logger.warning(f"MyFramework environment setup warning: {e}")
+    
+    def _configure_framework_settings(self):
+        """Configure framework-specific settings for optimal MLflow integration."""
+        try:
+            # Add framework-specific configuration here
+            # Example: Setting logging levels, performance options, etc.
+            pass
+            
+        except Exception as e:
+            self._logger.warning(f"MyFramework configuration warning: {e}")
+    
+    def _on_model_save(self, model, *args, **kwargs):
+        """
+        Hook called when a model is saved.
+        
+        Customize this to add framework-specific model save processing.
+        """
+        try:
+            # Check if it's your framework's model type
+            import myframework
+            
+            if isinstance(model, myframework.BaseModel):  # Replace with your model base class
+                self._logger.debug("MyFramework model save detected")
+                
+                # Add custom processing here
+                # Example: Log model architecture, parameters, etc.
+                self._log_model_metadata(model)
+                
+        except Exception as e:
+            self._logger.warning(f"MyFramework model save hook error: {e}")
+    
+    def _on_model_load(self, model_path, *args, **kwargs):
+        """
+        Hook called when a model is loaded.
+        
+        Customize this to add framework-specific model load processing.
+        """
+        try:
+            self._logger.debug(f"MyFramework model load detected: {model_path}")
+            
+            # Add custom processing here
+            # Example: Validate model, setup runtime environment, etc.
+            
+        except Exception as e:
+            self._logger.warning(f"MyFramework model load hook error: {e}")
+    
+    def _on_training_start(self, *args, **kwargs):
+        """Hook called when training starts."""
+        try:
+            self._logger.debug("MyFramework training start detected")
+            
+            # Add custom processing here
+            # Example: Setup training monitoring, resource tracking, etc.
+            
+        except Exception as e:
+            self._logger.warning(f"MyFramework training start hook error: {e}")
+    
+    def _on_training_end(self, *args, **kwargs):
+        """Hook called when training ends."""
+        try:
+            self._logger.debug("MyFramework training end detected")
+            
+            # Add custom processing here
+            # Example: Log final metrics, cleanup resources, etc.
+            
+        except Exception as e:
+            self._logger.warning(f"MyFramework training end hook error: {e}")
+    
+    def _log_model_metadata(self, model):
+        """Log additional metadata about the model."""
+        try:
+            # Example: Log model-specific information
+            # This is where you'd add framework-specific metadata logging
+            pass
+            
+        except Exception as e:
+            self._logger.warning(f"MyFramework model metadata logging error: {e}")
+    
+    def check_dependencies(self) -> bool:
+        """
+        Enhanced dependency checking for the framework.
+        
+        Customize this to add framework-specific dependency validation.
+        """
+        # Check basic dependencies first
+        if not super().check_dependencies():
+            return False
+        
+        try:
+            # Import and validate your framework
+            import myframework  # Replace with your framework import
+            
+            # Check minimum framework version
+            from packaging.version import Version
+            min_version = "1.0.0"  # Set your minimum version requirement
+            
+            if Version(myframework.__version__) < Version(min_version):
+                self._logger.error(f"MyFramework {min_version}+ required, found {myframework.__version__}")
+                return False
+            
+            # Add framework-specific validation
+            if not self._validate_framework_environment():
+                return False
+            
+            self._logger.debug(f"MyFramework {myframework.__version__} dependency satisfied")
+            return True
+            
+        except ImportError:
+            self._logger.error("MyFramework not installed")
+            return False
+        except Exception as e:
+            self._logger.error(f"MyFramework dependency check failed: {e}")
+            return False
+    
+    def _validate_framework_environment(self) -> bool:
+        """
+        Validate framework-specific environment requirements.
+        
+        Add custom validation logic here.
+        """
+        try:
+            # Add framework-specific validation
+            # Example: Check for required system libraries, GPU drivers, etc.
+            
+            # Example validation:
+            # import myframework
+            # if not myframework.is_properly_configured():
+            #     self._logger.error("MyFramework is not properly configured")
+            #     return False
+            
+            return True
+            
+        except Exception as e:
+            self._logger.error(f"MyFramework environment validation failed: {e}")
+            return False
+
+# Example usage and registration:
+
+# 1. To use this plugin directly:
+# metadata = MyFrameworkPlugin.get_default_metadata()
+# plugin = MyFrameworkPlugin(metadata)
+# plugin.enable()
+
+# 2. To register in built-in plugins, add to mlflow/plugins/builtin/__init__.py:
+# from mlflow.plugins.builtin.myframework_plugin import MyFrameworkPlugin
+# BUILTIN_PLUGINS["myframework"] = MyFrameworkPlugin
+
+# 3. To register via entry points, add to setup.py:
+# entry_points={
+#     "genesis_flow.frameworks": [
+#         "myframework = mypackage.myframework_plugin:MyFrameworkPlugin",
+#     ],
+# }

mlflow/prompt/constants.py

@@ -0,0 +1,20 @@
+# A special tag in RegisteredModel to indicate that it is a prompt
+import re
+
+IS_PROMPT_TAG_KEY = "mlflow.prompt.is_prompt"
+# A special tag in ModelVersion to store the prompt text
+PROMPT_TEXT_TAG_KEY = "mlflow.prompt.text"
+
+LINKED_PROMPTS_TAG_KEY = "mlflow.linkedPrompts"
+
+# A special tag to store associated run IDs for prompts
+PROMPT_ASSOCIATED_RUN_IDS_TAG_KEY = "mlflow.prompt.associatedRunIds"
+
+PROMPT_TEMPLATE_VARIABLE_PATTERN = re.compile(
+    r"\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*(?:\.[a-zA-Z_][a-zA-Z0-9_]*)*)\s*\}\}"
+)
+
+PROMPT_TEXT_DISPLAY_LIMIT = 30
+
+# Alphanumeric, underscore, hyphen, and dot are allowed in prompt name
+PROMPT_NAME_RULE = re.compile(r"^[a-zA-Z0-9_.-]+$")