From 7de0f6fa58584e37fcfa65325cd166bd5d451fbf Mon Sep 17 00:00:00 2001 From: First Second Date: Sun, 26 Apr 2026 06:24:50 +0000 Subject: [PATCH 1/3] Initial commit for introducing child process --- cmflib/cmf.py | 49 ++++- cmflib/cmf_async_proxy.py | 377 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 423 insertions(+), 3 deletions(-) create mode 100644 cmflib/cmf_async_proxy.py diff --git a/cmflib/cmf.py b/cmflib/cmf.py index 0fa11209..cd30e951 100644 --- a/cmflib/cmf.py +++ b/cmflib/cmf.py @@ -94,6 +94,9 @@ _dvc_ingest, ) +# Import async proxy for factory pattern +from cmflib.cmf_async_proxy 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,41 @@ 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 (lightweight proxy). + If async_logging=False, returns standard Cmf instance (full implementation). + + This provides day-one design benefits: async mode doesn't waste + resources on database connections and initialization in main process. + """ + if async_logging and not is_server: + # Return lightweight async proxy + logger.info("[Cmf Factory] Creating CmfAsyncProxy for async logging") + return CmfAsyncProxy( + filepath=filepath, + pipeline_name=pipeline_name, + custom_properties=custom_properties, + graph=graph, + finalize_timeout=finalize_timeout + ) + else: + # Return standard implementation + # Create instance normally (calls __init__) + instance = super(Cmf, cls).__new__(cls) + return instance + def __init__( self, filepath: str = "mlmd", @@ -146,6 +184,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 +313,13 @@ def __check_git_init(): f"Current Directory: {os.getcwd()}" ) sys.exit(1) - + def finalize(self): + """ + Finalize the CMF logging session. + Performs git commits. + """ + # Perform git commit git_commit(self.execution_name) if self.graph: self.driver.close() @@ -980,7 +1025,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] @@ -1240,7 +1284,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_proxy.py b/cmflib/cmf_async_proxy.py new file mode 100644 index 00000000..13aa6fdf --- /dev/null +++ b/cmflib/cmf_async_proxy.py @@ -0,0 +1,377 @@ +""" +Asynchronous proxy for CMF metadata logging. +This module provides a lightweight proxy that queues commands to a worker process. +""" + +### +# Copyright (2024) Hewlett Packard Enterprise Development LP +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +### + +import os +import sys +import time +import traceback +import multiprocessing +import atexit +import typing as t +import logging + +logger = logging.getLogger(__name__) + + +class CmfAsyncProxy: + """ + Lightweight async proxy for CMF. + + This class runs in the main process and only queues commands to a worker process. + It does NOT open database connections, create contexts, or do any heavy lifting. + All actual work happens in the worker process. + + 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, + ): + # Store only parameters needed for worker + 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 + + # Initialize IPC and spawn worker + self._command_queue: t.Optional[multiprocessing.Queue] = None + self._child_process: t.Optional[multiprocessing.Process] = None + self._call_id_counter = 0 + + self._setup_worker() + + # Register cleanup on exit + atexit.register(self._emergency_cleanup) + + def _setup_worker(self): + """Initialize worker process and communication queue.""" + logger.info("[CmfAsyncProxy] Initializing async logging with worker process") + + # Create communication queue + self._command_queue = multiprocessing.Queue() + + # Spawn worker process + self._child_process = multiprocessing.Process( + target=_worker_process_main, + args=(self._command_queue, self._worker_params), + name=f"cmf-worker-{self.pipeline_name}" + ) + + # Start the worker + self._child_process.start() + logger.info(f"[CmfAsyncProxy] Worker process started (PID: {self._child_process.pid})") + + def _queue_command(self, method_name: str, **kwargs): + """ + Queue a command to be executed by the worker process. + + Args: + method_name: Name of the CMF method to execute + **kwargs: Keyword arguments for the method + """ + if self._command_queue is None: + raise RuntimeError("Worker not initialized") + + self._call_id_counter += 1 + message = { + "method": method_name, + "kwargs": kwargs, + "timestamp": time.time(), + "call_id": self._call_id_counter + } + + try: + self._command_queue.put(message, block=False) + except Exception as e: + logger.error(f"[CmfAsyncProxy] Failed to queue command {method_name}: {e}") + raise + + def _check_error_log(self): + """Check error log and raise exception if errors were encountered.""" + error_log_path = f"{self.filepath}_errors.log" + if os.path.exists(error_log_path): + with open(error_log_path, "r") as f: + errors = f.read() + if errors.strip(): + error_count = errors.count("ERROR in") + raise RuntimeError( + f"CMF metadata logging encountered {error_count} error(s). " + f"Check {error_log_path} for details:\n\n{errors[:500]}..." + ) + + def _emergency_cleanup(self): + """Emergency cleanup if program exits without calling finalize().""" + if self._child_process and self._child_process.is_alive(): + logger.warning("[CmfAsyncProxy] Emergency cleanup: terminating worker process") + try: + if self._command_queue: + self._command_queue.put({"method": "shutdown"}, block=False) + self._child_process.join(timeout=5) + if self._child_process.is_alive(): + logger.warning("[CmfAsyncProxy] Emergency cleanup: force terminating") + self._child_process.terminate() + # Add join() after terminate() to ensure cleanup and prevent zombie process + self._child_process.join(timeout=2) + except Exception as e: + logger.error(f"[CmfAsyncProxy] Emergency cleanup error: {e}") + + # Public API methods - all just queue commands + + def create_context(self, pipeline_stage: str, custom_properties: t.Optional[t.Dict] = None): + """Create context (queued to worker).""" + self._queue_command("create_context", + pipeline_stage=pipeline_stage, + custom_properties=custom_properties) + return None + + 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 (queued to worker).""" + self._queue_command("update_context", + type_name=type_name, + context_name=context_name, + context_id=context_id, + properties=properties, + custom_properties=custom_properties) + return None + + 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 (queued to worker).""" + self._queue_command("create_execution", + execution_type=execution_type, + custom_properties=custom_properties, + cmd=cmd, + create_new_execution=create_new_execution) + return None + + def update_execution(self, execution_id: int, custom_properties: t.Optional[t.Dict] = None): + """Update execution (queued to worker).""" + self._queue_command("update_execution", + execution_id=execution_id, + custom_properties=custom_properties) + return None + + 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 (queued to worker).""" + self._queue_command("log_dataset", + url=url, + event=event, + custom_properties=custom_properties, + label=label, + label_properties=label_properties, + external=external) + return None + + 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 (queued to worker).""" + self._queue_command("log_model", + path=path, + event=event, + model_framework=model_framework, + model_type=model_type, + model_name=model_name, + custom_properties=custom_properties) + return None + + def log_execution_metrics( + self, metrics_name: str, custom_properties: t.Optional[t.Dict] = None + ): + """Log execution metrics (queued to worker).""" + self._queue_command("log_execution_metrics", + metrics_name=metrics_name, + custom_properties=custom_properties) + return None + + def log_metric( + self, metrics_name: str, custom_properties: t.Optional[t.Dict] = None + ): + """Log metric (queued to worker).""" + self._queue_command("log_metric", + metrics_name=metrics_name, + custom_properties=custom_properties) + return None + + def commit_metrics(self, metrics_name: str): + """Commit metrics (queued to worker).""" + self._queue_command("commit_metrics", + metrics_name=metrics_name) + return None + + def finalize(self): + """ + Finalize the CMF logging session. + + Blocks until all queued operations complete in the worker process. + """ + if not self._child_process: + return + + logger.info("[CmfAsyncProxy] Finalizing async logging session") + + # Queue finalize command + self._queue_command("finalize") + + # Queue shutdown command + self._command_queue.put({"method": "shutdown"}) + + # Wait for worker to complete + logger.info(f"[CmfAsyncProxy] Waiting for worker (timeout: {self.finalize_timeout}s)") + self._child_process.join(timeout=self.finalize_timeout) + + # Check if process is still alive + if self._child_process.is_alive(): + logger.error("[CmfAsyncProxy] Worker timeout - terminating") + self._child_process.terminate() + self._child_process.join(timeout=5) + raise TimeoutError( + f"Worker process did not complete within {self.finalize_timeout} seconds" + ) + + # Check exit code + if self._child_process.exitcode != 0: + logger.warning(f"[CmfAsyncProxy] Worker exited with code {self._child_process.exitcode}") + + # Check for errors + self._check_error_log() + + logger.info("[CmfAsyncProxy] Async logging session finalized successfully") + + +def _worker_process_main(command_queue: multiprocessing.Queue, worker_params: dict): + """ + Main entry point for worker process. + + This function runs in a separate process and handles CMF commands from the queue. + It creates a full Cmf instance (synchronous) and processes commands until shutdown. + + Args: + command_queue: Queue to receive commands from main process + worker_params: Parameters for creating Cmf instance + """ + logger.info(f"[Worker] Process started (PID: {os.getpid()})") + + # Import here to avoid circular dependency + from cmflib.cmf import Cmf + + try: + # Create full CMF instance in worker process (synchronous mode) + logger.info("[Worker] Initializing Cmf...") + cmf_worker = Cmf( + **worker_params, + async_logging=False # Worker always uses synchronous mode + ) + logger.info("[Worker] Cmf initialized, entering command loop") + + error_log_path = f"{worker_params['filepath']}_errors.log" + + # Main loop - process commands until shutdown + while True: + try: + # Block waiting for next command + message = command_queue.get() + + method_name = message.get("method") + + if method_name == "shutdown": + logger.info("[Worker] Received shutdown signal") + break + + # Execute the command + logger.debug(f"[Worker] Executing {method_name}") + method = getattr(cmf_worker, method_name) + method(**message.get("kwargs", {})) + + except KeyboardInterrupt: + logger.info("[Worker] Received keyboard interrupt") + break + + except Exception as e: + # Log error but continue processing + logger.error(f"[Worker] Error executing {method_name}: {e}") + try: + with open(error_log_path, "a") as f: + timestamp = time.strftime("%Y-%m-%d %H:%M:%S") + f.write(f"\n[{timestamp}] ERROR in {method_name}\n") + f.write(f" Call ID: {message.get('call_id', 'N/A')}\n") + f.write(f" Error: {type(e).__name__}: {str(e)}\n") + f.write(f" Traceback:\n") + f.write(traceback.format_exc()) + f.write("\n" + "="*80 + "\n") + except: + pass + + except Exception as e: + logger.error(f"[Worker] Fatal error in worker process: {e}") + traceback.print_exc() + sys.exit(1) + + finally: + # Cleanup + try: + if hasattr(cmf_worker, 'driver') and cmf_worker.graph and cmf_worker.driver: + cmf_worker.driver.close() + except: + pass + logger.info("[Worker] Process shutting down") From 7d9d42a290d0e25ad94028624c0b8c6bff769145 Mon Sep 17 00:00:00 2001 From: First Second Date: Tue, 23 Jun 2026 07:07:48 +0000 Subject: [PATCH 2/3] Addressed a bug found while testing --- cmflib/cmf_merger.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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"]: From 74d7b15485894fbeeffe01f221733ccd32da8e9d Mon Sep 17 00:00:00 2001 From: First Second Date: Thu, 16 Jul 2026 23:08:49 +0000 Subject: [PATCH 3/3] updated the async logging architecture --- cmflib/cmf.py | 17 +- cmflib/cmf_async_api.py | 254 +++++++++++++++++++++ cmflib/cmf_async_proxy.py | 377 ------------------------------- cmflib/cmf_subprocess_manager.py | 242 ++++++++++++++++++++ cmflib/cmf_worker_loop.py | 220 ++++++++++++++++++ 5 files changed, 726 insertions(+), 384 deletions(-) create mode 100644 cmflib/cmf_async_api.py delete mode 100644 cmflib/cmf_async_proxy.py create mode 100644 cmflib/cmf_subprocess_manager.py create mode 100644 cmflib/cmf_worker_loop.py diff --git a/cmflib/cmf.py b/cmflib/cmf.py index f0b65ac5..9c2ce140 100644 --- a/cmflib/cmf.py +++ b/cmflib/cmf.py @@ -94,8 +94,8 @@ _dvc_ingest, ) -# Import async proxy for factory pattern -from cmflib.cmf_async_proxy import CmfAsyncProxy +# 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. @@ -155,15 +155,18 @@ def __new__( """ Factory method that returns appropriate CMF implementation. - By default (async_logging=True), returns CmfAsyncProxy (lightweight proxy). - If async_logging=False, returns standard Cmf instance (full 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 lightweight async proxy - logger.info("[Cmf Factory] Creating CmfAsyncProxy for async logging") + # 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, @@ -172,7 +175,7 @@ def __new__( finalize_timeout=finalize_timeout ) else: - # Return standard implementation + # Return standard synchronous implementation # Create instance normally (calls __init__) instance = super(Cmf, cls).__new__(cls) return instance 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_async_proxy.py b/cmflib/cmf_async_proxy.py deleted file mode 100644 index 13aa6fdf..00000000 --- a/cmflib/cmf_async_proxy.py +++ /dev/null @@ -1,377 +0,0 @@ -""" -Asynchronous proxy for CMF metadata logging. -This module provides a lightweight proxy that queues commands to a worker process. -""" - -### -# Copyright (2024) Hewlett Packard Enterprise Development LP -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# You may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -### - -import os -import sys -import time -import traceback -import multiprocessing -import atexit -import typing as t -import logging - -logger = logging.getLogger(__name__) - - -class CmfAsyncProxy: - """ - Lightweight async proxy for CMF. - - This class runs in the main process and only queues commands to a worker process. - It does NOT open database connections, create contexts, or do any heavy lifting. - All actual work happens in the worker process. - - 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, - ): - # Store only parameters needed for worker - 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 - - # Initialize IPC and spawn worker - self._command_queue: t.Optional[multiprocessing.Queue] = None - self._child_process: t.Optional[multiprocessing.Process] = None - self._call_id_counter = 0 - - self._setup_worker() - - # Register cleanup on exit - atexit.register(self._emergency_cleanup) - - def _setup_worker(self): - """Initialize worker process and communication queue.""" - logger.info("[CmfAsyncProxy] Initializing async logging with worker process") - - # Create communication queue - self._command_queue = multiprocessing.Queue() - - # Spawn worker process - self._child_process = multiprocessing.Process( - target=_worker_process_main, - args=(self._command_queue, self._worker_params), - name=f"cmf-worker-{self.pipeline_name}" - ) - - # Start the worker - self._child_process.start() - logger.info(f"[CmfAsyncProxy] Worker process started (PID: {self._child_process.pid})") - - def _queue_command(self, method_name: str, **kwargs): - """ - Queue a command to be executed by the worker process. - - Args: - method_name: Name of the CMF method to execute - **kwargs: Keyword arguments for the method - """ - if self._command_queue is None: - raise RuntimeError("Worker not initialized") - - self._call_id_counter += 1 - message = { - "method": method_name, - "kwargs": kwargs, - "timestamp": time.time(), - "call_id": self._call_id_counter - } - - try: - self._command_queue.put(message, block=False) - except Exception as e: - logger.error(f"[CmfAsyncProxy] Failed to queue command {method_name}: {e}") - raise - - def _check_error_log(self): - """Check error log and raise exception if errors were encountered.""" - error_log_path = f"{self.filepath}_errors.log" - if os.path.exists(error_log_path): - with open(error_log_path, "r") as f: - errors = f.read() - if errors.strip(): - error_count = errors.count("ERROR in") - raise RuntimeError( - f"CMF metadata logging encountered {error_count} error(s). " - f"Check {error_log_path} for details:\n\n{errors[:500]}..." - ) - - def _emergency_cleanup(self): - """Emergency cleanup if program exits without calling finalize().""" - if self._child_process and self._child_process.is_alive(): - logger.warning("[CmfAsyncProxy] Emergency cleanup: terminating worker process") - try: - if self._command_queue: - self._command_queue.put({"method": "shutdown"}, block=False) - self._child_process.join(timeout=5) - if self._child_process.is_alive(): - logger.warning("[CmfAsyncProxy] Emergency cleanup: force terminating") - self._child_process.terminate() - # Add join() after terminate() to ensure cleanup and prevent zombie process - self._child_process.join(timeout=2) - except Exception as e: - logger.error(f"[CmfAsyncProxy] Emergency cleanup error: {e}") - - # Public API methods - all just queue commands - - def create_context(self, pipeline_stage: str, custom_properties: t.Optional[t.Dict] = None): - """Create context (queued to worker).""" - self._queue_command("create_context", - pipeline_stage=pipeline_stage, - custom_properties=custom_properties) - return None - - 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 (queued to worker).""" - self._queue_command("update_context", - type_name=type_name, - context_name=context_name, - context_id=context_id, - properties=properties, - custom_properties=custom_properties) - return None - - 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 (queued to worker).""" - self._queue_command("create_execution", - execution_type=execution_type, - custom_properties=custom_properties, - cmd=cmd, - create_new_execution=create_new_execution) - return None - - def update_execution(self, execution_id: int, custom_properties: t.Optional[t.Dict] = None): - """Update execution (queued to worker).""" - self._queue_command("update_execution", - execution_id=execution_id, - custom_properties=custom_properties) - return None - - 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 (queued to worker).""" - self._queue_command("log_dataset", - url=url, - event=event, - custom_properties=custom_properties, - label=label, - label_properties=label_properties, - external=external) - return None - - 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 (queued to worker).""" - self._queue_command("log_model", - path=path, - event=event, - model_framework=model_framework, - model_type=model_type, - model_name=model_name, - custom_properties=custom_properties) - return None - - def log_execution_metrics( - self, metrics_name: str, custom_properties: t.Optional[t.Dict] = None - ): - """Log execution metrics (queued to worker).""" - self._queue_command("log_execution_metrics", - metrics_name=metrics_name, - custom_properties=custom_properties) - return None - - def log_metric( - self, metrics_name: str, custom_properties: t.Optional[t.Dict] = None - ): - """Log metric (queued to worker).""" - self._queue_command("log_metric", - metrics_name=metrics_name, - custom_properties=custom_properties) - return None - - def commit_metrics(self, metrics_name: str): - """Commit metrics (queued to worker).""" - self._queue_command("commit_metrics", - metrics_name=metrics_name) - return None - - def finalize(self): - """ - Finalize the CMF logging session. - - Blocks until all queued operations complete in the worker process. - """ - if not self._child_process: - return - - logger.info("[CmfAsyncProxy] Finalizing async logging session") - - # Queue finalize command - self._queue_command("finalize") - - # Queue shutdown command - self._command_queue.put({"method": "shutdown"}) - - # Wait for worker to complete - logger.info(f"[CmfAsyncProxy] Waiting for worker (timeout: {self.finalize_timeout}s)") - self._child_process.join(timeout=self.finalize_timeout) - - # Check if process is still alive - if self._child_process.is_alive(): - logger.error("[CmfAsyncProxy] Worker timeout - terminating") - self._child_process.terminate() - self._child_process.join(timeout=5) - raise TimeoutError( - f"Worker process did not complete within {self.finalize_timeout} seconds" - ) - - # Check exit code - if self._child_process.exitcode != 0: - logger.warning(f"[CmfAsyncProxy] Worker exited with code {self._child_process.exitcode}") - - # Check for errors - self._check_error_log() - - logger.info("[CmfAsyncProxy] Async logging session finalized successfully") - - -def _worker_process_main(command_queue: multiprocessing.Queue, worker_params: dict): - """ - Main entry point for worker process. - - This function runs in a separate process and handles CMF commands from the queue. - It creates a full Cmf instance (synchronous) and processes commands until shutdown. - - Args: - command_queue: Queue to receive commands from main process - worker_params: Parameters for creating Cmf instance - """ - logger.info(f"[Worker] Process started (PID: {os.getpid()})") - - # Import here to avoid circular dependency - from cmflib.cmf import Cmf - - try: - # Create full CMF instance in worker process (synchronous mode) - logger.info("[Worker] Initializing Cmf...") - cmf_worker = Cmf( - **worker_params, - async_logging=False # Worker always uses synchronous mode - ) - logger.info("[Worker] Cmf initialized, entering command loop") - - error_log_path = f"{worker_params['filepath']}_errors.log" - - # Main loop - process commands until shutdown - while True: - try: - # Block waiting for next command - message = command_queue.get() - - method_name = message.get("method") - - if method_name == "shutdown": - logger.info("[Worker] Received shutdown signal") - break - - # Execute the command - logger.debug(f"[Worker] Executing {method_name}") - method = getattr(cmf_worker, method_name) - method(**message.get("kwargs", {})) - - except KeyboardInterrupt: - logger.info("[Worker] Received keyboard interrupt") - break - - except Exception as e: - # Log error but continue processing - logger.error(f"[Worker] Error executing {method_name}: {e}") - try: - with open(error_log_path, "a") as f: - timestamp = time.strftime("%Y-%m-%d %H:%M:%S") - f.write(f"\n[{timestamp}] ERROR in {method_name}\n") - f.write(f" Call ID: {message.get('call_id', 'N/A')}\n") - f.write(f" Error: {type(e).__name__}: {str(e)}\n") - f.write(f" Traceback:\n") - f.write(traceback.format_exc()) - f.write("\n" + "="*80 + "\n") - except: - pass - - except Exception as e: - logger.error(f"[Worker] Fatal error in worker process: {e}") - traceback.print_exc() - sys.exit(1) - - finally: - # Cleanup - try: - if hasattr(cmf_worker, 'driver') and cmf_worker.graph and cmf_worker.driver: - cmf_worker.driver.close() - except: - pass - logger.info("[Worker] Process shutting down") 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"