Spaces:
Running
Running
| """ | |
| Spatial Atlas — Safe Code Executor | |
| Runs generated ML pipeline scripts in a subprocess with timeout. | |
| Captures stdout/stderr for debugging and self-healing. | |
| """ | |
| import asyncio | |
| import logging | |
| import sys | |
| from pathlib import Path | |
| logger = logging.getLogger("spatial-atlas.mlebench.executor") | |
| class CodeExecutor: | |
| """Execute ML pipeline code safely in a subprocess.""" | |
| def __init__(self, timeout: int = 600): | |
| self.timeout = timeout | |
| self.last_stdout: str = "" | |
| self.last_stderr: str = "" | |
| self.last_error: str | None = None | |
| async def execute( | |
| self, | |
| code: str, | |
| working_dir: Path, | |
| submission_path: Path | None = None, | |
| ) -> bytes | None: | |
| """ | |
| Execute ML code in subprocess, return submission.csv bytes. | |
| Args: | |
| code: Complete Python script to execute | |
| working_dir: Directory to run in (contains data/) | |
| submission_path: Where to find submission.csv (default: working_dir/submission.csv) | |
| Returns: | |
| submission.csv bytes if produced, None if execution failed | |
| """ | |
| script_path = working_dir / "pipeline.py" | |
| script_path.write_text(code) | |
| if submission_path is None: | |
| submission_path = working_dir / "submission.csv" | |
| self.last_stdout = "" | |
| self.last_stderr = "" | |
| self.last_error = None | |
| logger.info(f"Executing pipeline.py in {working_dir} (timeout={self.timeout}s)") | |
| try: | |
| proc = await asyncio.create_subprocess_exec( | |
| sys.executable, str(script_path), | |
| cwd=str(working_dir), | |
| stdout=asyncio.subprocess.PIPE, | |
| stderr=asyncio.subprocess.PIPE, | |
| env=self._safe_env(), | |
| ) | |
| stdout, stderr = await asyncio.wait_for( | |
| proc.communicate(), timeout=self.timeout | |
| ) | |
| self.last_stdout = stdout.decode("utf-8", errors="replace") | |
| self.last_stderr = stderr.decode("utf-8", errors="replace") | |
| # Log output for debugging | |
| if self.last_stdout: | |
| logger.info(f"Pipeline stdout (last 500 chars): {self.last_stdout[-500:]}") | |
| if self.last_stderr: | |
| logger.warning(f"Pipeline stderr (last 500 chars): {self.last_stderr[-500:]}") | |
| if proc.returncode != 0: | |
| self.last_error = ( | |
| f"Script exited with code {proc.returncode}.\n" | |
| f"Stderr:\n{self.last_stderr[-2000:]}" | |
| ) | |
| logger.error(f"Pipeline failed: exit code {proc.returncode}") | |
| return None | |
| # Check for submission file | |
| if submission_path.exists(): | |
| csv_bytes = submission_path.read_bytes() | |
| logger.info(f"Submission produced: {len(csv_bytes)} bytes") | |
| return csv_bytes | |
| else: | |
| self.last_error = ( | |
| f"Script ran successfully but did not produce {submission_path.name}.\n" | |
| f"Stdout:\n{self.last_stdout[-1000:]}" | |
| ) | |
| logger.error("No submission.csv found after execution") | |
| return None | |
| except asyncio.TimeoutError: | |
| self.last_error = f"Code execution timed out after {self.timeout}s" | |
| logger.error(self.last_error) | |
| try: | |
| proc.kill() | |
| await proc.wait() | |
| except Exception: | |
| pass | |
| return None | |
| except Exception as e: | |
| self.last_error = f"Execution error: {e}" | |
| logger.error(f"Pipeline execution exception: {e}") | |
| return None | |
| def _safe_env(self) -> dict[str, str]: | |
| """Build a safe environment for subprocess execution.""" | |
| import os | |
| env = os.environ.copy() | |
| # Ensure reproducibility | |
| env["PYTHONHASHSEED"] = "42" | |
| # Suppress warnings that clutter stderr | |
| env["PYTHONWARNINGS"] = "ignore" | |
| return env | |