comet-ml--opik/apps/opik-python-backend/src/opik_backend/executor_isolated.py

"""Isolated subprocess executor with environment variable scoping."""
import json
import logging
import os
import resource
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Callable, List
from threading import Lock

from opik_backend.subprocess_log_config import SubprocessLogConfig

logger = logging.getLogger(__name__)

# Metrics setup
from opentelemetry import metrics

meter = metrics.get_meter("isolated_executor")
isolated_creation_histogram = meter.create_histogram(
    name="isolated_subprocess_creation_latency",
    description="Latency of isolated subprocess creation in milliseconds",
    unit="ms",
)
isolated_execution_histogram = meter.create_histogram(
    name="isolated_subprocess_execution_latency",
    description="Latency of isolated code execution in milliseconds",
    unit="ms",
)

active_process_counter = meter.create_up_down_counter(
    name="isolated_subprocess_active_count",
    description="Current number of active isolated subprocesses",
    unit="1",
)

# Memory limit for subprocesses in bytes (20MB)
SUBPROCESS_MEMORY_LIMIT_BYTES = 20 * 1024 * 1024  # 20MB

def _calculate_latency_ms(start_time):
    """Calculate elapsed time in milliseconds."""
    return (time.time() - start_time) * 1000

def _set_memory_limit():
    """Set memory limit for subprocess to 20MB.
    
    Uses RLIMIT_STACK to limit only stack size.
    This prevents deeply nested calls and excessive local variables
    while allowing the Python interpreter and runtime heap to function normally.
    """
    try:
        # RLIMIT_STACK limits stack size only (local variables, call stack depth)
        # Prevents stack overflow from deeply nested recursion
        # Does NOT limit heap or runtime data structures
        resource.setrlimit(resource.RLIMIT_STACK, (SUBPROCESS_MEMORY_LIMIT_BYTES, SUBPROCESS_MEMORY_LIMIT_BYTES))
    except Exception as e:
        logger.warning(f"Failed to set stack memory limit: {e}")

class IsolatedSubprocessExecutor:
    """
    Executes Python code in isolated subprocesses with environment variable scoping.

    Each execution creates a fresh subprocess, ensuring that:
    - Environment variables are scoped to each execution (no leakage between concurrent runs)
    - Subprocesses are completely independent
    - No shared state exists between executions
    - Resources are properly cleaned up after each execution

    This differs from ProcessExecutor which maintains a pool of reusable workers.
    Use this when you need true isolation with custom environment variables per execution.
    """

    def __init__(self, timeout_secs: int = 30):
        """
        Initialize the isolated subprocess executor.

        Args:
            timeout_secs: Timeout for each execution in seconds (default: 30)
        """
        self.timeout_secs = timeout_secs
        self.logger = logging.getLogger(__name__)
        self._active_processes: List[subprocess.Popen] = []  # Track active processes for cleanup
        self._process_lock = Lock()
        self._teardown_callbacks: List[Callable[[], None]] = []  # Callbacks to run on teardown
        self._log_collectors = {}  # Map process PID to log collector instance

    def execute(
        self,
        file_path: str,
        data: dict = {},
        env_vars: Optional[dict] = None,
        timeout_secs: Optional[int] = None,
        payload_type: Optional[str] = None,
        optimization_id: Optional[str] = None,
        job_id: Optional[str] = None,
    ) -> dict:
        """
        Execute Python file in an isolated subprocess with scoped environment variables.

        Each call creates a fresh subprocess with its own isolated environment.
        Environment variables passed in env_vars are scoped to the subprocess
        and don't affect other concurrent executions.

        Args:
            file_path: Path to Python file to execute (e.g., '/path/to/metric.py')
            data: Data dictionary to pass to the file via stdin
            env_vars: Environment variables to scope to this subprocess (optional).
                     These override/augment the parent environment for this execution only.
            timeout_secs: Execution timeout in seconds (uses default if not provided)
            payload_type: Type of payload being executed (e.g., 'trace_thread')
            optimization_id: Optimization identifier for log correlation
            job_id: Job identifier for log correlation

        Returns:
            dict: Result dictionary with format:
                - {"scores": [...]} on success
                - {"code": error_code, "error": message} on failure
        """
        timeout_secs = timeout_secs or self.timeout_secs
        if data is None:
            data = {}
        creation_start = time.time()
        process = None  # Initialize to None for exception handling
        result = None

        try:
            if env_vars is None:
                env_vars = {}
            # Prepare environment for subprocess
            subprocess_env = self._prepare_environment(env_vars)

            # Create subprocess with python to execute the file directly
            process = subprocess.Popen(
                [sys.executable, "-u", file_path],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                env=subprocess_env,
                text=True,
                bufsize=1,
                preexec_fn=_set_memory_limit,  # Apply memory limit to subprocess
            )

            # Track active process
            self._active_processes.append(process)

            creation_latency = _calculate_latency_ms(creation_start)
            isolated_creation_histogram.record(creation_latency)
            self.logger.debug(
                f"Created isolated subprocess. pid={process.pid}, creation_latency_ms={creation_latency:.3f}"
            )

            # Execute code in subprocess
            result = self._execute_in_subprocess(
                process,
                data,
                payload_type,
                timeout_secs,
                optimization_id=optimization_id,
                job_id=job_id,
                env_vars=env_vars or {},
            )

            execution_latency = _calculate_latency_ms(creation_start)
            isolated_execution_histogram.record(execution_latency)
            self.logger.debug(
                f"Isolated subprocess execution completed. total_latency_ms={execution_latency:.3f}"
            )

            return result

        except subprocess.TimeoutExpired:
            self.logger.error(
                f"Subprocess execution timed out. timeout_secs={timeout_secs}"
            )
            return {
                "code": 500,
                "error": f"Execution timed out after {timeout_secs} seconds",
            }
        except Exception as e:
            self.logger.error(
                f"Error during subprocess execution. error={str(e)}", exc_info=True
            )
            return {"code": 500, "error": f"Failed to execute file: {str(e)}"}
        finally:
            # Always remove process from active list and measure total latency
            self._remove_active_process(process)
            total_latency = _calculate_latency_ms(creation_start)
            self.logger.debug(
                f"Subprocess execution finished. total_latency_ms={total_latency:.3f}"
            )

    def _remove_active_process(self, process: Optional[subprocess.Popen]) -> None:
        """Remove process from active processes list if it exists."""
        if process in self._active_processes:
            self._active_processes.remove(process)

    def _prepare_environment(self, env_vars: Optional[dict] = None) -> dict:
        """
        Prepare environment variables for the subprocess.

        Starts with a copy of the parent environment and applies overrides.
        This ensures the subprocess has all necessary environment variables
        while allowing specific variables to be scoped to this execution.

        Args:
            env_vars: Environment variables to override/add

        Returns:
            dict: Complete environment dictionary for the subprocess
        """
        env = os.environ.copy()
        if env_vars:
            # Filter out None values and log them - subprocess.Popen requires all values to be strings
            for key, value in env_vars.items():
                if value is None:
                    self.logger.warning(f"Skipping environment variable '{key}' with None value")
                else:
                    env[key] = value

        env["PYTHONUNBUFFERED"] = '1'
        env["LOG_FORMAT"] = 'json'
        return env

    @staticmethod
    def _parse_last_json_line(stdout: str) -> tuple[Optional[dict], Optional[str]]:
        """Parse the last non-empty stdout line as JSON.

        Returns (result, None) on success, or (None, error_message) on failure.
        """
        lines = [line for line in stdout.split('\n') if line.strip()]
        if not lines:
            return None, "No output produced by subprocess"
        try:
            return json.loads(lines[-1]), None
        except (json.JSONDecodeError, ValueError) as e:
            return None, f"Invalid JSON response from subprocess: {str(e)}"

    def _execute_in_subprocess(
        self,
        process: subprocess.Popen,
        data: dict,
        payload_type: Optional[str],
        timeout_secs: int,
        optimization_id: str = None,
        job_id: str = None,
        env_vars: dict = {},
    ) -> dict:
        """
        Execute code in the subprocess and collect result.

        Uses python -c to execute code inline with stdin for data passing.
        Streams stderr to the logging backend in real-time if configured.

        Args:
            process: Subprocess Popen instance
            data: Data dictionary
            payload_type: Type of payload
            timeout_secs: Execution timeout
            optimization_id: Optional ID for log correlation
            job_id: Optional ID for log correlation

        Returns:
            dict: Execution result
        """
        # Prepare input as JSON to pass via stdin
        input_json = json.dumps(
            {
                "data": data,
                "payload_type": payload_type,
            }
        )

        try:
            # Initialize logger BEFORE process starts
            # Check if a log collector was pre-registered (e.g., for optimization jobs)
            pre_registered_collector = self._log_collectors.pop(0, None)  # 0 is placeholder key
            
            if pre_registered_collector is not None:
                # Use pre-registered collector (e.g., RedisBatchLogCollector for optimizations)
                self._log_collectors[process.pid] = pre_registered_collector
            elif SubprocessLogConfig.is_fully_configured():
                # Fallback to HTTP-based logging if configured
                try:
                    from opik_backend.subprocess_logger import BatchLogCollector
                    self._log_collectors[process.pid] = BatchLogCollector(
                        backend_url=SubprocessLogConfig.get_backend_url(),
                        optimization_id=optimization_id or "",
                        job_id=job_id or "",
                        api_key=env_vars.get("OPIK_API_KEY", ""),
                        workspace=env_vars.get("OPIK_WORKSPACE", ""),
                    )
                except ImportError as e:
                    self.logger.error(f"Subprocess logging is configured but BatchLogCollector import failed: {e}")
                    raise
                except Exception as e:
                    self.logger.error(f"Failed to initialize subprocess logging: {e}")
            
            # Decide execution strategy based on logging configuration
            if self._log_collectors.get(process.pid):
                # Real-time streaming: start log collector threads, then wait for process
                self._log_collectors[process.pid].start_stream_from_process(process)
                
                # Send input to the process
                process.stdin.write(input_json)
                process.stdin.close()  # Signal EOF so process knows input is done
                
                # Wait for process to complete
                process.wait(timeout=timeout_secs)
                
                # Wait for reader threads to finish reading all output
                # Note: log collector may have been removed if process was killed externally
                log_collector = self._log_collectors.get(process.pid)
                if log_collector:
                    log_collector.wait_for_reader_threads(timeout=SubprocessLogConfig.get_log_reader_timeout_secs())
                    
                    # Get stdout/stderr from last lines (no memory accumulation)
                    # Safe to do now that threads have finished
                    last_lines = log_collector.get_last_lines()
                    stdout = last_lines.get('stdout', '')
                    stderr = last_lines.get('stderr', '')
                else:
                    # Process was killed externally (e.g., cancellation)
                    stdout = ''
                    stderr = ''
            else:
                # Simple mode: use communicate() without logging overhead
                stdout, stderr = process.communicate(input=input_json, timeout=timeout_secs)

            # Parse result from stdout
            if process.returncode == 0:
                result, parse_error = self._parse_last_json_line(stdout)
                if result is None:
                    self.logger.error(f"Failed to parse subprocess output: {stdout}")
                    return {
                        "code": 500,
                        "error": parse_error,
                    }
                return result
            else:
                self.logger.error(
                    f"Subprocess exited with code {process.returncode}. "
                    f"stdout: {stdout}. stderr: {stderr}"
                )
                # optimizer_runner emits error JSON on stdout; prefer it over generic message
                result, _ = self._parse_last_json_line(stdout)
                if isinstance(result, dict) and "error" in result:
                    return result
                return {
                    "code": 500,
                    "error": f"Subprocess execution failed: {stderr}",
                }
        
        except subprocess.TimeoutExpired:
            process.kill()
            try:
                process.wait(timeout=2)
            except subprocess.TimeoutExpired:
                process.kill()
            raise
        finally:
            # Close log collector immediately after execution completes to flush all logs
            # This ensures cleanup happens even if teardown() is not called
            self._close_log_collector(process.pid)

    def _close_log_collector(self, pid: int):
        """Close the log collector for a specific process if it exists.
        
        This properly signals shutdown, waits for any pending flushes, 
        and cleans up all threads.
        """
        try:
            with self._process_lock:
                if pid in self._log_collectors:
                    # close() handles: signal stop -> shutdown executor -> final flush -> cleanup threads
                    self._log_collectors[pid].close()
                    del self._log_collectors[pid]
        except Exception as e:
            self.logger.warning(f"Error closing log collector for PID {pid}: {e}")

    def register_teardown_callback(self, callback):
        """
        Register a callback to run during teardown.

        Args:
            callback: A callable that takes no arguments and runs cleanup logic
        """
        self._teardown_callbacks.append(callback)

    def kill_process(self, pid: int, timeout: int = 2):
        """
        Terminate a specific process by PID.

        Args:
            pid: Process ID to terminate
            timeout: Seconds to wait before force killing

        Returns:
            bool: True if process was terminated, False if not found
        """
        try:
            # Atomically find and remove process from active list
            with self._process_lock:
                process = next((p for p in self._active_processes if p.pid == pid), None)
                if process is None:
                    self.logger.warning(f"Process with PID {pid} not found in active processes")
                    return False
                self._active_processes.remove(process)
            
            # Now kill the process (outside lock to avoid long hold times)
            process.terminate()
            try:
                process.wait(timeout=timeout)
                self.logger.info(f"Process {pid} terminated gracefully")
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait()
                self.logger.warning(f"Process {pid} force killed after timeout")
            
            # Close log collector after process is terminated to capture any final logs
            self._close_log_collector(pid)
            
            return True
        except Exception as e:
            self.logger.error(f"Error killing process {pid}: {e}")
            return False

    def kill_all_processes(self, timeout: int = 2):
        """
        Terminate all active processes in parallel.

        Args:
            timeout: Total seconds to wait for all processes (distributed across them)
        """
        # Collect PIDs first to avoid modification during iteration
        with self._process_lock:
            pids = [p.pid for p in list(self._active_processes)]
        
        if not pids:
            self.logger.info("No active processes to terminate")
            return
        
        # Kill processes in parallel using ThreadPoolExecutor
        with ThreadPoolExecutor(max_workers=len(pids)) as executor:
            futures = [
                executor.submit(self.kill_process, pid, timeout)
                for pid in pids
            ]
            # Wait for all to complete (with timeout for safety)
            for future in futures:
                try:
                    future.result(timeout=timeout)
                except Exception as e:
                    self.logger.warning(f"Error waiting for kill_process result: {e}")
        
        self.logger.info("All active processes terminated")

    def teardown(self):
        """
        Run cleanup logic including killing all processes and executing teardown callbacks.
        """
        self.logger.info("Running teardown...")
        
        # Kill all active processes (which will also close their log collectors)
        self.kill_all_processes()
        
        # Close any remaining log collectors that weren't cleaned up
        for pid in list(self._log_collectors.keys()):
            try:
                self._close_log_collector(pid)
            except Exception as e:
                self.logger.error(f"Error closing remaining log collector for PID {pid}: {e}")
        
        # Execute all teardown callbacks
        for callback in self._teardown_callbacks:
            try:
                callback()
            except Exception as e:
                self.logger.error(f"Error in teardown callback: {e}")
        
        self.logger.info("Teardown complete")

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - runs teardown."""
        self.teardown()
        return False