experimaestro 1.11.1__py3-none-any.whl → 2.0.0b4__py3-none-any.whl

experimaestro/progress.py

@@ -0,0 +1,404 @@
+"""File-based progress tracking system for experimaestro tasks."""
+
+import json
+import threading
+import time
+from dataclasses import dataclass, asdict
+from pathlib import Path
+from typing import Optional, List, Iterator, Dict, Any
+from datetime import datetime, timedelta
+import fcntl
+import os
+
+from .utils import logger
+
+DEFAULT_MAX_ENTRIES_PER_FILE = 10_000
+
+
+@dataclass
+class ProgressEntry:
+    """A single progress entry in the JSONL file"""
+
+    timestamp: float
+    level: int
+    progress: float
+    desc: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization"""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ProgressEntry":
+        """Create from dictionary"""
+        return cls(**data)
+
+
+class StateFile:
+    """Represents the state file for progress tracking.
+    Checks if the state must be written based on time and progress changes.
+    By default, it writes every second or when progress changes significantly (>1%)"""
+
+    def __init__(self, filename: Path):
+        self.filename = filename
+        self.state: Dict[int, ProgressEntry] = {}
+
+        # Write threshold to avoid too frequent writes
+        self._time_threshold = timedelta(seconds=1.0)
+        self._last_write_time: datetime = datetime.now()
+        # Minimum progress change to trigger write
+        self._progress_threshold = 0.01
+        self._last_write_progress: Optional[Dict[int, float]] = None
+
+        self.filename.parent.mkdir(parents=True, exist_ok=True)
+        self.load()
+
+    def _allow_write(self) -> bool:
+        """Check if the state should be written based on time and progress changes.
+        Allows writing if:
+        - BOTH: More than 1 second has passed since last write
+        - AND: Progress has changed significantly (>1%)
+        - OR: All entries are done (progress >= 1.0)"""
+        time_check = datetime.now() - self._last_write_time > self._time_threshold
+        progress_check = self._last_write_progress is None or any(
+            abs(entry.progress - self._last_write_progress.get(entry.level, 0.0))
+            > self._progress_threshold
+            for entry in self.state.values()
+        )
+        all_entries_done = all(entry.progress >= 1.0 for entry in self.state.values())
+        return all_entries_done or (time_check and progress_check)
+
+    def write(self, force: bool = False):
+        """Write the current state to the file."""
+        if self._allow_write() or force:
+            with open(self.filename, "w") as f:
+                json.dump({k: v.to_dict() for k, v in self.state.items()}, f)
+            self._last_write_time = datetime.now()
+            self._last_write_progress = {k: v.progress for k, v in self.state.items()}
+
+    def update(self, entry: ProgressEntry):
+        self.state[entry.level] = entry
+
+    def load(self):
+        """Load the state from the file"""
+        if self.filename.exists():
+            with self.filename.open("r") as f:
+                try:
+                    data = json.load(f)
+                    self.state = {
+                        int(k): ProgressEntry.from_dict(v) for k, v in data.items()
+                    }
+                except (json.JSONDecodeError, IOError):
+                    logger.warning(f"Failed to load state from {self.filename}")
+
+    def read(self) -> Dict[int, ProgressEntry]:
+        """Read the state from the file"""
+        self.load()
+        return self.state
+
+    # flush on exit
+    def __del__(self):
+        """Ensure state is written on exit"""
+        try:
+            self.write(force=True)
+        except Exception as e:
+            logger.error(f"Failed to write state on exit: {e}")
+
+
+class ProgressFileWriter:
+    def __init__(
+        self, task_path: Path, max_entries_per_file: int = DEFAULT_MAX_ENTRIES_PER_FILE
+    ):
+        self.task_path = task_path
+        self.progress_dir = task_path / ".experimaestro"
+        self.max_entries_per_file = max_entries_per_file
+        self.current_file_index = 0
+        self.current_file_entries = 0
+        self.lock = threading.Lock()
+
+        # Ensure directory exists
+        self.progress_dir.mkdir(exist_ok=True)
+
+        # State is the latest entry per level
+        self.state = StateFile(self.progress_dir / "progress_state.json")
+
+        # Find the latest file index
+        self._find_latest_file()
+
+    def _find_latest_file(self):
+        """Find the latest progress file and entry count"""
+        progress_files = list(self.progress_dir.glob("progress-*.jsonl"))
+        if not progress_files:
+            self.current_file_index = 0
+            self.current_file_entries = 0
+            return
+
+        # Sort by file index
+        max_index = None
+        for f in progress_files:
+            try:
+                index = int(f.stem.split("-")[1])
+                if max_index is None or index > max_index:
+                    max_index = index
+            except (ValueError, IndexError):
+                continue
+
+        if max_index is not None:
+            self.current_file_index = max_index
+            # Count entries in current file
+            current_file = self._get_current_file_path()
+            if current_file.exists():
+                with current_file.open("r") as f:
+                    self.current_file_entries = sum(1 for _ in f.readlines())
+            else:
+                self.current_file_entries = 0
+        else:
+            self.current_file_index = 0
+            self.current_file_entries = 0
+
+    def _get_current_file_path(self) -> Path:
+        """Get path to current progress file"""
+        return self.progress_dir / f"progress-{self.current_file_index:04d}.jsonl"
+
+    def _get_latest_symlink_path(self) -> Path:
+        """Get path to latest progress symlink"""
+        return self.progress_dir / "progress-latest.jsonl"
+
+    def _rotate_file_if_needed(self):
+        """Create new file if current one is full"""
+        if self.current_file_entries >= self.max_entries_per_file:
+            self.current_file_index += 1
+            self.current_file_entries = 0
+            logger.debug(f"Rotating to new progress file: {self.current_file_index}")
+
+    def _update_latest_symlink(self):
+        """Update symlink to point to latest file"""
+        current_file = self._get_current_file_path()
+        latest_symlink = self._get_latest_symlink_path()
+
+        # Remove existing symlink
+        if latest_symlink.exists() or latest_symlink.is_symlink():
+            latest_symlink.unlink()
+
+        # Create new symlink
+        latest_symlink.symlink_to(current_file.name)
+
+    def write_progress(self, level: int, progress: float, desc: Optional[str] = None):
+        """Write a progress entry to the file
+
+        Args:
+            level: Progress level (0 is top level)
+            progress: Progress value between 0.0 and 1.0
+            desc: Optional description
+        """
+        with self.lock:
+            # Eventually rotate internal state if needed
+            self._rotate_file_if_needed()
+
+            entry = ProgressEntry(
+                timestamp=time.time(), level=level, progress=progress, desc=desc
+            )
+            self.state.update(entry)
+            self.state.write(force=level == -1)  # Force write on EOJ
+
+            current_file = self._get_current_file_path()
+
+            # Write with file locking for concurrent access
+            with current_file.open("a") as f:
+                try:
+                    fcntl.flock(f.fileno(), fcntl.LOCK_EX)
+                    f.write(json.dumps(entry.to_dict()) + "\n")
+                    f.flush()  # Flush the file buffer
+                    os.fsync(f.fileno())  # Ensure data is written to disk
+                finally:
+                    fcntl.flock(f.fileno(), fcntl.LOCK_UN)
+
+            self.current_file_entries += 1
+            self._update_latest_symlink()
+
+            logger.debug(
+                f"Progress written: level={level}, progress={progress}, desc={desc}"
+            )
+
+    def __del__(self):
+        """Ensure state is written on exit"""
+        try:
+            self.state.write(force=True)
+        except Exception as e:
+            logger.error(f"Failed to write state on exit: {e}")
+
+
+class ProgressFileReader:
+    """Reads progress entries from JSONL files"""
+
+    def __init__(self, task_path: Path):
+        """Initialize progress file reader
+
+        Args:
+            task_path: Path to the task directory
+        """
+        self.task_path = task_path
+        self.progress_dir = task_path / ".experimaestro"
+        self.max_entries_per_file: Optional[int] = None
+        self.state = StateFile(self.progress_dir / "progress_state.json")
+
+    def get_progress_files(self) -> List[Path]:
+        """Get all progress files sorted by index"""
+        if not self.progress_dir.exists():
+            return []
+
+        progress_files = list(self.progress_dir.glob("progress-*.jsonl"))
+
+        # Filter out symlinks to avoid duplicates
+        progress_files = [f for f in progress_files if not f.is_symlink()]
+
+        # Sort by file index
+        # Alternatively, we could simply sort by filename
+        def get_index(path: Path) -> int:
+            try:
+                return int(path.stem.split("-")[1])
+            except (ValueError, IndexError):
+                return 0
+
+        return sorted(progress_files, key=get_index)
+
+    def get_latest_file(self) -> Optional[Path]:
+        """Get the latest progress file via symlink"""
+        latest_symlink = self.progress_dir / "progress-latest.jsonl"
+        if latest_symlink.exists() and latest_symlink.is_symlink():
+            return latest_symlink.resolve()
+
+        # Fallback to finding latest manually
+        files = self.get_progress_files()
+        return files[-1] if files else None
+
+    def read_entries(self, file_path: Path) -> Iterator[ProgressEntry]:
+        """Read progress entries from a file
+
+        Args:
+            file_path: Path to progress file
+
+        Yields:
+            ProgressEntry objects
+        """
+        if not file_path.exists():
+            return
+
+        try:
+            with file_path.open("r") as f:
+                fcntl.flock(f.fileno(), fcntl.LOCK_SH)
+                try:
+                    for line in f:
+                        line = line.strip()
+                        if line:
+                            try:
+                                data = json.loads(line)
+                                yield ProgressEntry.from_dict(data)
+                            except json.JSONDecodeError as e:
+                                logger.warning(
+                                    f"Invalid JSON in progress file {file_path}: {e}"
+                                )
+                finally:
+                    fcntl.flock(f.fileno(), fcntl.LOCK_UN)
+        except IOError as e:
+            logger.warning(f"Could not read progress file {file_path}: {e}")
+
+    def read_all_entries(self) -> Iterator[ProgressEntry]:
+        """Read all progress entries from all files in order
+
+        Yields:
+            ProgressEntry objects in chronological order
+        """
+        logger.warning("Reading all progress entries, this may be slow for large jobs.")
+        for file_path in self.get_progress_files():
+            yield from self.read_entries(file_path)
+
+    def read_latest_entries(self, count: Optional[int] = None) -> List[ProgressEntry]:
+        """Read the latest N progress entries"""
+        entries = []
+
+        # Read files in reverse order to get latest entries first
+        files = self.get_progress_files()
+        # Fetch the max length of files, in lines
+        if files and count is None:
+            # Fetch the number of entries in the first file
+            # This is the most likely to be the longest file
+            count = sum(1 for _ in self.read_entries(files[0]))
+        if count is None:
+            count = DEFAULT_MAX_ENTRIES_PER_FILE
+
+        for file_path in reversed(files):
+            file_entries = list(self.read_entries(file_path))
+            entries.extend(reversed(file_entries))
+
+            if len(entries) >= count:
+                break
+
+        # Return latest entries in chronological order
+        return list(reversed(entries[:count]))
+
+    def get_current_progress(
+        self, count: Optional[int] = None
+    ) -> Dict[int, ProgressEntry]:
+        """Get the current progress for each level"""
+        logger.warning(
+            "Reading current progress from progress logs, this may be slow for large jobs."
+        )
+        return {entry.level: entry for entry in self.read_latest_entries(count)}
+
+    def get_current_state(self) -> Optional[Dict[int, ProgressEntry]]:
+        """Fetch the latest progress entry from the state file"""
+        current_state = self.state.read()
+        return current_state or self.get_current_progress()
+
+    def is_done(self) -> bool:
+        """Check if the task is done by looking for a special 'done' file.
+        Fallback to checking for end-of-job (EOJ) entries."""
+
+        task_name = self.task_path.parent.stem.split(".")[-1]
+        job_done_file = self.task_path / f"{task_name}.done"
+        if job_done_file.exists() and job_done_file.is_file():
+            return True
+
+        # Check if any progress file has a level -1 entry indicating EOJ
+        return any(entry.level == -1 for entry in self.read_all_entries())
+
+
+class FileBasedProgressReporter:
+    """File-based progress reporter that replaces the socket-based Reporter"""
+
+    def __init__(self, task_path: Path):
+        """Initialize file-based progress reporter
+
+        Args:
+            task_path: Path to the task directory
+        """
+        self.task_path = task_path
+        self.writer = ProgressFileWriter(task_path)
+        self.current_progress = {}  # level -> (progress, desc)
+        self.lock = threading.Lock()
+
+    def set_progress(self, progress: float, level: int = 0, desc: Optional[str] = None):
+        """Set progress for a specific level
+
+        Args:
+            progress: Progress value between 0.0 and 1.0
+            level: Progress level (0 is top level)
+            desc: Optional description
+        """
+        with self.lock:
+            # Check if progress has changed significantly
+            current = self.current_progress.get(level, (None, None))
+            if (
+                current[0] is None
+                or abs(progress - current[0]) > 0.01
+                or desc != current[1]
+            ):
+                self.current_progress[level] = (progress, desc)
+                self.writer.write_progress(level, progress, desc)
+
+    def eoj(self):
+        """End of job notification"""
+        with self.lock:
+            # Write a special end-of-job marker
+            self.writer.write_progress(-1, 1.0, "EOJ")

experimaestro/rpyc.py

@@ -1,5 +1,4 @@
 import atexit
-import shutil
 import tempfile
 from pathlib import Path
 from subprocess import Popen, PIPE, run

experimaestro/run.py

@@ -9,6 +9,7 @@ from typing import List
 import fasteners
 from experimaestro.notifications import progress, report_eoj
 from experimaestro.utils.multiprocessing import delayed_shutdown
+from experimaestro.exceptions import GracefulTimeout
 from .core.types import ObjectType
 from experimaestro.utils import logger
 from experimaestro.core.objects import ConfigInformation
@@ -41,9 +42,6 @@ def run(parameters: Path):
         task = ConfigInformation.fromParameters(params["objects"])
         task.__taskdir__ = Path.cwd()
 
-        # Set the tags
-        task.__tags__ = params["tags"]
-
         # Notify that the task has started
         progress(0)
 
@@ -92,9 +90,20 @@ class TaskRunner:
                 report_eoj()
             logger.info("Finished cleanup")
 
-    def handle_error(self, code, frame_type):
-        logger.info("Error handler: finished with code %d", code)
-        self.failedpath.write_text(str(code))
+    def handle_error(self, code, frame_type, reason: str = "failed", message: str = ""):
+        """Handle task error and write failure information.
+
+        Args:
+            code: Exit code
+            frame_type: Signal frame type (unused)
+            reason: Failure reason (e.g., "failed", "timeout")
+            message: Optional message with details
+        """
+        logger.info("Error handler: finished with code %d, reason=%s", code, reason)
+        failure_info = {"code": code, "reason": reason}
+        if message:
+            failure_info["message"] = message
+        self.failedpath.write_text(json.dumps(failure_info))
         self.cleanup()
         logger.info("Exiting")
         delayed_shutdown(60, exit_code=code)
@@ -147,6 +156,10 @@ class TaskRunner:
                 # Everything went OK
                 logger.info("Task ended successfully")
                 sys.exit(0)
+        except GracefulTimeout as e:
+            logger.info("Task requested graceful timeout: %s", e.message)
+            self.handle_error(1, None, reason="timeout", message=e.message)
+
         except Exception:
             logger.exception("Got exception while running")
             self.handle_error(1, None)

experimaestro/scheduler/__init__.py

@@ -1 +1,18 @@
-from .base import *
+from .base import Scheduler, Listener
+from .workspace import Workspace, RunMode
+from .experiment import experiment, FailedExperiment
+from .jobs import Job, JobState, JobFailureStatus, JobDependency, JobContext
+
+__all__ = [
+    "Scheduler",
+    "Listener",
+    "Workspace",
+    "RunMode",
+    "experiment",
+    "FailedExperiment",
+    "Job",
+    "JobState",
+    "JobFailureStatus",
+    "JobDependency",
+    "JobContext",
+]