diff --git a/cmflib/cmf.py b/cmflib/cmf.py index b45fcc0f..9c2ce140 100644 --- a/cmflib/cmf.py +++ b/cmflib/cmf.py @@ -94,6 +94,9 @@ _dvc_ingest, ) +# Import async proxy for factory pattern (shared subprocess architecture) +from cmflib.cmf_async_api import CmfAsyncProxy + class Cmf: """This class provides methods to log metadata for distributed AI pipelines. The class instance creates an ML metadata store to store the metadata. @@ -139,6 +142,44 @@ class Cmf: __neo4j_password = attr_dict.get("neo4j-password", "") __neo4j_user = attr_dict.get("neo4j-user", "") + def __new__( + cls, + filepath: str = "mlmd", + pipeline_name: str = "", + custom_properties: t.Optional[t.Dict] = None, + graph: bool = False, + is_server: bool = False, + async_logging: bool = True, + finalize_timeout: int = 300, + ): + """ + Factory method that returns appropriate CMF implementation. + + By default (async_logging=True), returns CmfAsyncProxy which uses a single + shared subprocess for ALL Cmf instances. + + This provides day-one design benefits: async mode doesn't waste + resources on database connections and initialization in main process. + + The shared subprocess architecture scales to 1000+ concurrent Cmf instances + without resource exhaustion. + """ + if async_logging and not is_server: + # Return async proxy with shared subprocess + logger.info("[Cmf Factory] Creating CmfAsyncProxy (shared subprocess) for async logging") + return CmfAsyncProxy( + filepath=filepath, + pipeline_name=pipeline_name, + custom_properties=custom_properties, + graph=graph, + finalize_timeout=finalize_timeout + ) + else: + # Return standard synchronous implementation + # Create instance normally (calls __init__) + instance = super(Cmf, cls).__new__(cls) + return instance + def __init__( self, filepath: str = "mlmd", @@ -146,6 +187,8 @@ def __init__( custom_properties: t.Optional[t.Dict] = None, graph: bool = False, is_server: bool = False, + async_logging: bool = True, + finalize_timeout: int = 300, ): #path to directory self.cmf_init_path = filepath.rsplit("/",1)[0] \ @@ -273,8 +316,12 @@ def __check_git_init(): f"Current Directory: {os.getcwd()}" ) sys.exit(1) - + def finalize(self): + """ + Finalize the CMF logging session. + Performs git commits. + """ commit_value = git_commit(self.execution_name) if self.execution: self.execution.properties["Git_End_Commit"].string_value = commit_value @@ -983,7 +1030,6 @@ def log_model( Returns: Artifact object from ML Metadata library associated with the new model artifact. """ - logging_dir = change_dir(self.cmf_init_path) # Assigning current file name as stage and execution name current_script = sys.argv[0] @@ -1243,7 +1289,6 @@ def commit_metrics(self, metrics_name: str): Returns: Artifact object from the ML Protocol Buffers library associated with the new metrics artifact. """ - logging_dir = change_dir(self.cmf_init_path) # code for nano cmf # Assigning current file name as stage and execution name diff --git a/cmflib/cmf_async_api.py b/cmflib/cmf_async_api.py new file mode 100644 index 00000000..f90851bd --- /dev/null +++ b/cmflib/cmf_async_api.py @@ -0,0 +1,254 @@ +""" +CMF Async API - Lightweight proxy for async CMF logging. + +This module provides CmfAsyncProxy, a lightweight wrapper that submits tasks +to the shared worker subprocess. Multiple CmfAsyncProxy instances share one subprocess. +""" + +import logging +import typing as t +import uuid + +from cmflib.cmf_subprocess_manager import get_manager + +logger = logging.getLogger(__name__) + + +class CmfAsyncProxy: + """ + Lightweight async proxy for CMF. + + This class provides the same interface as Cmf, but all operations are + executed in a shared worker subprocess. Multiple CmfAsyncProxy instances + share the same subprocess, avoiding resource exhaustion. + + Example: + ```python + # Create multiple Cmf instances - they all share one subprocess + cmf1 = Cmf(pipeline_name="pipeline1", async_logging=True) + cmf2 = Cmf(pipeline_name="pipeline2", async_logging=True) + + # Operations are queued to the shared subprocess + cmf1.log_dataset("data.csv", event="input") + cmf2.log_model("model.pkl", event="output") + + # Finalize and wait for completion + cmf1.finalize() + cmf2.finalize() + ``` + + Args: + filepath: Path to metadata store + pipeline_name: Name of the pipeline + custom_properties: Pipeline custom properties + graph: Whether to use graph database + finalize_timeout: Timeout for finalize() in seconds + """ + + def __init__( + self, + filepath: str = "mlmd", + pipeline_name: str = "", + custom_properties: t.Optional[t.Dict] = None, + graph: bool = False, + finalize_timeout: int = 300, + ): + # Generate unique session ID + self._session_id = f"{pipeline_name}_{uuid.uuid4().hex[:8]}" + + # Store parameters + self._worker_params = { + 'filepath': filepath, + 'pipeline_name': pipeline_name, + 'custom_properties': custom_properties or {}, + 'graph': graph, + } + + self.filepath = filepath + self.pipeline_name = pipeline_name + self.finalize_timeout = finalize_timeout + + # Get the singleton manager + self._manager = get_manager() + + # Start subprocess if not already started + self._manager.start() + + # Initialize this session in the worker subprocess + self._submit_task("_init_session", **self._worker_params) + + logger.info(f"[CmfAsyncProxy] Created session {self._session_id}") + + def _submit_task(self, method: str, **kwargs) -> t.Any: + """ + Submit a task to the worker subprocess. + + This is a blocking call - it waits for the subprocess to complete + the task and return the result. + + Args: + method: Name of the CMF method to call + **kwargs: Keyword arguments for the method + + Returns: + Result from the method execution + """ + return self._manager.submit_task( + session_id=self._session_id, + method=method, + kwargs=kwargs, + timeout=self.finalize_timeout + ) + + # Public API - mirrors Cmf interface + + def create_context( + self, pipeline_stage: str, custom_properties: t.Optional[t.Dict] = None + ): + """Create context (executed in worker subprocess).""" + return self._submit_task( + "create_context", + pipeline_stage=pipeline_stage, + custom_properties=custom_properties + ) + + def update_context( + self, + type_name: str, + context_name: str, + context_id: int, + properties: t.Optional[t.Dict] = None, + custom_properties: t.Optional[t.Dict] = None + ): + """Update context (executed in worker subprocess).""" + return self._submit_task( + "update_context", + type_name=type_name, + context_name=context_name, + context_id=context_id, + properties=properties, + custom_properties=custom_properties + ) + + def create_execution( + self, + execution_type: str, + custom_properties: t.Optional[t.Dict] = None, + cmd: t.Optional[str] = None, + create_new_execution: bool = True, + ): + """Create execution (executed in worker subprocess).""" + return self._submit_task( + "create_execution", + execution_type=execution_type, + custom_properties=custom_properties, + cmd=cmd, + create_new_execution=create_new_execution + ) + + def update_execution( + self, execution_id: int, custom_properties: t.Optional[t.Dict] = None + ): + """Update execution (executed in worker subprocess).""" + return self._submit_task( + "update_execution", + execution_id=execution_id, + custom_properties=custom_properties + ) + + def log_dataset( + self, + url: str, + event: str, + custom_properties: t.Optional[t.Dict] = None, + label: t.Optional[str] = None, + label_properties: t.Optional[t.Dict] = None, + external: bool = False, + ): + """Log dataset (executed in worker subprocess).""" + return self._submit_task( + "log_dataset", + url=url, + event=event, + custom_properties=custom_properties, + label=label, + label_properties=label_properties, + external=external + ) + + def log_model( + self, + path: str, + event: str, + model_framework: str = "Default", + model_type: str = "Default", + model_name: str = "Default", + custom_properties: t.Optional[t.Dict] = None, + ): + """Log model (executed in worker subprocess).""" + return self._submit_task( + "log_model", + path=path, + event=event, + model_framework=model_framework, + model_type=model_type, + model_name=model_name, + custom_properties=custom_properties + ) + + def log_execution_metrics( + self, metrics_name: str, custom_properties: t.Optional[t.Dict] = None + ): + """Log execution metrics (executed in worker subprocess).""" + return self._submit_task( + "log_execution_metrics", + metrics_name=metrics_name, + custom_properties=custom_properties + ) + + def log_metric( + self, metrics_name: str, custom_properties: t.Optional[t.Dict] = None + ): + """Log metric (executed in worker subprocess).""" + return self._submit_task( + "log_metric", + metrics_name=metrics_name, + custom_properties=custom_properties + ) + + def commit_metrics(self, metrics_name: str): + """Commit metrics (executed in worker subprocess).""" + return self._submit_task( + "commit_metrics", + metrics_name=metrics_name + ) + + def finalize(self): + """ + Finalize the CMF logging session. + + Blocks until the worker subprocess completes the finalize operation. + """ + logger.info(f"[CmfAsyncProxy] Finalizing session {self._session_id}") + + # Execute finalize in subprocess (blocking) + self._submit_task("finalize") + + # Cleanup this session + self._submit_task("_cleanup_session") + + logger.info(f"[CmfAsyncProxy] Session {self._session_id} finalized") + + +def shutdown_worker(timeout: float = 10.0): + """ + Shutdown the shared worker subprocess. + + This is called automatically at program exit. You can also call it manually + if you want to explicitly shutdown the subprocess before program exit. + + Args: + timeout: Maximum time to wait for subprocess to exit + """ + manager = get_manager() + manager.shutdown(timeout=timeout) diff --git a/cmflib/cmf_merger.py b/cmflib/cmf_merger.py index 0a9d15ba..90d24544 100644 --- a/cmflib/cmf_merger.py +++ b/cmflib/cmf_merger.py @@ -225,13 +225,15 @@ def parse_json_to_mlmd(mlmd_json, path_to_store: str, cmd: str, exec_uuid: Union graph = bool(os.getenv('NEO4J_URI', "")) # Initialize the cmf class with pipeline_name and graph_status + # Note: async_logging=False is required for metadata pull/push operations + # to ensure we get a regular Cmf instance with store attribute, not CmfAsyncProxy if cmd == "pull": cmf_class = Cmf(filepath=path_to_store, pipeline_name=pipeline_name, #intializing cmf - graph=graph) + graph=graph, async_logging=False) else: # in else, we are assuming cmd="push" cmf_class = Cmf(filepath=path_to_store, pipeline_name=pipeline_name, #intializing cmf - graph=graph, is_server=True) + graph=graph, is_server=True, async_logging=False) # Process each stage sequentially for stage in data["stages"]: diff --git a/cmflib/cmf_subprocess_manager.py b/cmflib/cmf_subprocess_manager.py new file mode 100644 index 00000000..3e046b5c --- /dev/null +++ b/cmflib/cmf_subprocess_manager.py @@ -0,0 +1,242 @@ +""" +CMF Subprocess Manager - Singleton that manages the shared worker subprocess. + +This module provides a singleton manager that: +1. Starts a single worker subprocess +2. Routes tasks from all Cmf instances to that subprocess +3. Handles results and errors +4. Gracefully shuts down the subprocess +""" + +import multiprocessing +import threading +import atexit +import logging +import time +import uuid +import typing as t + +logger = logging.getLogger(__name__) + + +class CmfSubprocessManager: + """ + Singleton manager for the shared CMF worker subprocess. + + This class ensures that only ONE subprocess is created, no matter how many + Cmf instances are created. All Cmf instances share this single subprocess. + + Thread-safe: Uses locks to prevent race conditions during initialization. + """ + + _instance: t.Optional['CmfSubprocessManager'] = None + _lock = threading.Lock() + + def __new__(cls): + """Singleton pattern: only one instance ever exists.""" + if cls._instance is None: + with cls._lock: + if cls._instance is None: # Double-check + cls._instance = super(CmfSubprocessManager, cls).__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self): + """Initialize the manager (only runs once due to singleton).""" + if self._initialized: + return + + # Communication channels + self.task_queue: t.Optional[multiprocessing.Queue] = None + self.result_queue: t.Optional[multiprocessing.Queue] = None + self.shutdown_event: t.Optional[multiprocessing.Event] = None + + # Subprocess + self.worker_process: t.Optional[multiprocessing.Process] = None + + # State + self._started = False + self._shutdown_requested = False + self._initialized = True + + logger.info("[CMF Manager] CmfSubprocessManager initialized") + + def start(self): + """ + Start the worker subprocess. + + This is called automatically by the first Cmf instance that's created. + Subsequent calls do nothing (subprocess already running). + """ + with self._lock: + if self._started: + logger.debug("[CMF Manager] Subprocess already started") + return + + logger.info("[CMF Manager] Starting worker subprocess") + + # Create communication channels + self.task_queue = multiprocessing.Queue() + self.result_queue = multiprocessing.Queue() + self.shutdown_event = multiprocessing.Event() + + # Import worker loop + from cmflib.cmf_worker_loop import worker_loop + + # Create and start subprocess + self.worker_process = multiprocessing.Process( + target=worker_loop, + args=(self.task_queue, self.result_queue, self.shutdown_event), + name="cmf-worker-subprocess", + daemon=False # Not daemon - we want graceful shutdown + ) + self.worker_process.start() + + self._started = True + + # Register cleanup on program exit + atexit.register(self.shutdown) + + logger.info(f"[CMF Manager] Worker subprocess started (PID: {self.worker_process.pid})") + + def submit_task( + self, + session_id: str, + method: str, + kwargs: t.Dict, + timeout: t.Optional[float] = None + ) -> t.Any: + """ + Submit a task to the worker subprocess and wait for the result. + + This is a blocking call - it waits for the subprocess to complete the task. + + Args: + session_id: Unique identifier for the CMF session + method: Name of the CMF method to call + kwargs: Keyword arguments for the method + timeout: Maximum time to wait for result (None = wait forever) + + Returns: + Result from the method execution + + Raises: + RuntimeError: If subprocess not started or task execution failed + TimeoutError: If result not received within timeout + """ + if not self._started: + raise RuntimeError("Subprocess not started. Call start() first.") + + # Generate unique task ID + task_id = str(uuid.uuid4()) + + # Create task + task = { + "task_id": task_id, + "session_id": session_id, + "method": method, + "kwargs": kwargs, + "timestamp": time.time() + } + + logger.debug(f"[CMF Manager] Submitting task {task_id}: {session_id}.{method}") + + # Send task to subprocess + self.task_queue.put(task) + + # Wait for result + start_time = time.time() + while True: + try: + result = self.result_queue.get(timeout=1.0) + + # Check if this is our result + if result.get("task_id") == task_id: + if result.get("status") == "success": + logger.debug(f"[CMF Manager] Task {task_id} completed successfully") + return result.get("result") + + elif result.get("status") == "error": + error_msg = result.get("error", "Unknown error") + error_trace = result.get("traceback", "") + logger.error(f"[CMF Manager] Task {task_id} failed: {error_msg}") + raise RuntimeError(f"Task execution failed: {error_msg}\n{error_trace}") + + else: + # Not our result - put it back (shouldn't happen with FIFO execution) + logger.warning(f"[CMF Manager] Received result for wrong task: {result.get('task_id')}") + self.result_queue.put(result) + + except multiprocessing.queues.Empty: + # Check timeout + if timeout is not None and (time.time() - start_time) > timeout: + raise TimeoutError(f"Task {task_id} timed out after {timeout}s") + + # Check if subprocess is still alive + if not self.worker_process.is_alive(): + raise RuntimeError(f"Worker subprocess died while executing task {task_id}") + + continue + + def shutdown(self, timeout: float = 10.0): + """ + Gracefully shutdown the worker subprocess. + + Args: + timeout: Maximum time to wait for subprocess to exit + """ + if not self._started or self._shutdown_requested: + return + + with self._lock: + if self._shutdown_requested: + return + + self._shutdown_requested = True + + logger.info("[CMF Manager] Shutting down worker subprocess") + + # Signal shutdown + self.shutdown_event.set() + + # Send poison pill + self.task_queue.put(None) + + # Wait for subprocess to exit + if self.worker_process.is_alive(): + self.worker_process.join(timeout=timeout) + + if self.worker_process.is_alive(): + logger.warning("[CMF Manager] Subprocess did not exit gracefully, terminating") + self.worker_process.terminate() + self.worker_process.join(timeout=2.0) + + if self.worker_process.is_alive(): + logger.error("[CMF Manager] Subprocess did not terminate, killing") + self.worker_process.kill() + self.worker_process.join() + + # Cleanup + if self.worker_process.exitcode == 0: + logger.info("[CMF Manager] Worker subprocess exited cleanly") + else: + logger.warning(f"[CMF Manager] Worker subprocess exited with code {self.worker_process.exitcode}") + + self._started = False + + +# Global singleton instance (lazy initialization) +_manager: t.Optional[CmfSubprocessManager] = None + + +def get_manager() -> CmfSubprocessManager: + """ + Get the global singleton manager instance. + + Returns: + The CmfSubprocessManager singleton + """ + global _manager + if _manager is None: + _manager = CmfSubprocessManager() + return _manager diff --git a/cmflib/cmf_worker_loop.py b/cmflib/cmf_worker_loop.py new file mode 100644 index 00000000..4ab2f3b6 --- /dev/null +++ b/cmflib/cmf_worker_loop.py @@ -0,0 +1,220 @@ +""" +CMF Worker Loop - Runs in the subprocess. + +This module contains the worker loop that executes CMF operations in a separate process. +All CMF instances in the main process share this single worker subprocess. +""" + +import multiprocessing +import traceback +import logging +import time +import typing as t + +logger = logging.getLogger(__name__) + + +def worker_loop( + task_queue: multiprocessing.Queue, + result_queue: multiprocessing.Queue, + shutdown_event: multiprocessing.Event +): + """ + Main worker loop that runs in the subprocess. + + This loop: + 1. Pulls tasks from task_queue + 2. Executes the task (CMF operations) + 3. Puts the result back in result_queue + 4. Continues until shutdown_event is set + + Args: + task_queue: Queue to receive tasks from main process + result_queue: Queue to send results back to main process + shutdown_event: Event to signal shutdown + """ + logger.info(f"[CMF Worker] Worker subprocess started (PID: {multiprocessing.current_process().pid})") + + # Import here to avoid circular dependency and to keep imports isolated in subprocess + from cmflib.cmf import Cmf + + # Map session_id -> Cmf instance + # Each Cmf object in main process gets a corresponding real Cmf instance here + cmf_sessions: t.Dict[str, Cmf] = {} + + try: + while not shutdown_event.is_set(): + try: + # Wait for next task (with timeout to check shutdown_event) + task = task_queue.get(timeout=1.0) + + if task is None: # Poison pill + logger.info("[CMF Worker] Received poison pill, shutting down") + break + + # Extract task details + task_id = task.get("task_id") + session_id = task.get("session_id") + method = task.get("method") + kwargs = task.get("kwargs", {}) + + logger.debug(f"[CMF Worker] Executing task {task_id}: {session_id}.{method}") + + # Execute the task + result = execute_task( + cmf_sessions=cmf_sessions, + session_id=session_id, + method=method, + kwargs=kwargs + ) + + # Send result back + result_queue.put({ + "task_id": task_id, + "status": "success", + "result": result + }) + + except multiprocessing.queues.Empty: + # Timeout - just check shutdown_event again + continue + + except Exception as e: + # Task execution error - send back error result + logger.error(f"[CMF Worker] Error executing task {task_id}: {e}") + result_queue.put({ + "task_id": task_id, + "status": "error", + "error": str(e), + "traceback": traceback.format_exc() + }) + + except KeyboardInterrupt: + logger.info("[CMF Worker] Interrupted by keyboard") + + except Exception as e: + logger.error(f"[CMF Worker] Fatal error in worker loop: {e}") + traceback.print_exc() + + finally: + # Cleanup: close all Cmf sessions + logger.info("[CMF Worker] Cleaning up CMF sessions") + for session_id, cmf in list(cmf_sessions.items()): + try: + if hasattr(cmf, 'driver') and cmf.graph and cmf.driver: + cmf.driver.close() + except Exception as e: + logger.error(f"[CMF Worker] Error closing session {session_id}: {e}") + + logger.info("[CMF Worker] Worker subprocess shutting down") + + +def execute_task( + cmf_sessions: t.Dict[str, t.Any], + session_id: str, + method: str, + kwargs: t.Dict +) -> t.Any: + """ + Execute a CMF task. + + Routes the method call to the appropriate handler based on the method name. + + Args: + cmf_sessions: Dictionary mapping session_id to Cmf instances + session_id: Unique identifier for the CMF session + method: Name of the CMF method to call + kwargs: Keyword arguments for the method + + Returns: + Result of the method execution (usually None for logging operations) + """ + # Special commands for session management + if method == "_init_session": + return _init_session_handler(cmf_sessions, session_id, kwargs) + + elif method == "_cleanup_session": + return _cleanup_session_handler(cmf_sessions, session_id) + + # Regular CMF methods + elif session_id not in cmf_sessions: + raise ValueError(f"Session {session_id} not found. Call _init_session first.") + + # Get the Cmf instance for this session + cmf = cmf_sessions[session_id] + + # Call the method on the Cmf instance + method_func = getattr(cmf, method) + result = method_func(**kwargs) + + return result + + +def _init_session_handler( + cmf_sessions: t.Dict[str, t.Any], + session_id: str, + kwargs: t.Dict +) -> str: + """ + Initialize a new CMF session. + + Creates a real Cmf instance (synchronous mode) for this session. + + Args: + cmf_sessions: Dictionary to store the new session + session_id: Unique identifier for this session + kwargs: Parameters to pass to Cmf constructor + + Returns: + Success message + """ + from cmflib.cmf import Cmf + + logger.info(f"[CMF Worker] Initializing session: {session_id}") + + # Create a real Cmf instance in synchronous mode + cmf_sessions[session_id] = Cmf( + **kwargs, + async_logging=False # IMPORTANT: Worker always uses synchronous mode + ) + + logger.info(f"[CMF Worker] Session {session_id} initialized") + return f"Session {session_id} initialized" + + +def _cleanup_session_handler( + cmf_sessions: t.Dict[str, t.Any], + session_id: str +) -> str: + """ + Clean up a CMF session. + + Closes resources and removes the session from the dictionary. + + Args: + cmf_sessions: Dictionary containing sessions + session_id: Session to clean up + + Returns: + Success message + """ + if session_id not in cmf_sessions: + logger.warning(f"[CMF Worker] Cleanup: Session {session_id} not found") + return f"Session {session_id} not found" + + logger.info(f"[CMF Worker] Cleaning up session: {session_id}") + + cmf = cmf_sessions[session_id] + + # Close graph driver if exists + try: + if hasattr(cmf, 'driver') and cmf.graph and cmf.driver: + cmf.driver.close() + except Exception as e: + logger.error(f"[CMF Worker] Error closing driver for {session_id}: {e}") + + # Remove from dictionary + del cmf_sessions[session_id] + + logger.info(f"[CMF Worker] Session {session_id} cleaned up") + return f"Session {session_id} cleaned up"