experimaestro 1.5.1__py3-none-any.whl → 2.0.0a8__py3-none-any.whl

experimaestro/notifications.py

@@ -12,6 +12,7 @@ from tqdm.auto import tqdm as std_tqdm
 
 from .utils import logger
 from experimaestro.taskglobals import Env as TaskEnv
+from .progress import FileBasedProgressReporter
 
 # --- Progress and other notifications
 
@@ -41,7 +42,13 @@ class LevelInformation:
         return result
 
     def __repr__(self) -> str:
-        return f"[{self.level}] {self.desc} {int(self.progress*1000)/10}%"
+        return f"[{self.level}] {self.desc} {int(self.progress * 1000) / 10}%"
+
+
+class ListenerInformation:
+    def __init__(self, url: str):
+        self.url = url
+        self.error_count = 0
 
 
 class Reporter(threading.Thread):
@@ -59,7 +66,7 @@ class Reporter(threading.Thread):
         super().__init__(daemon=True)
         self.path = path / Reporter.NOTIFICATION_FOLDER
         self.path.mkdir(exist_ok=True)
-        self.urls: Dict[str, str] = {}
+        self.urls: Dict[str, ListenerInformation] = {}
 
         # Last check of notification URLs
         self.lastcheck = 0
@@ -72,15 +79,18 @@ class Reporter(threading.Thread):
 
         self.progress_threshold = 0.01
         self.cv = threading.Condition()
-        self.start()
+
+        # File-based progress reporter
+        self.file_reporter = FileBasedProgressReporter(task_path=path)
 
     def stop(self):
         self.stopping = True
         with self.cv:
-            self.cv.notifyAll()
+            # self.cv.notifyAll()
+            self.cv.notify_all()
 
     @staticmethod
-    def isfatal_httperror(e: Exception) -> bool:
+    def isfatal_httperror(e: Exception, info: ListenerInformation) -> bool:
         """Returns True if this HTTP error indicates that the server won't recover"""
         if isinstance(e, HTTPError):
             if e.code >= 400 and e.code < 500:
@@ -90,6 +100,13 @@ class Reporter(threading.Thread):
                 return True
             if isinstance(e.reason, socket.gaierror) and e.reason.errno == -2:
                 return True
+            if isinstance(e.reason, TimeoutError):
+                info.error_count += 1
+
+            # Too many errors
+            if info.error_count > 3:
+                logger.info("Too many errors with %s", info.error_count)
+                return True
 
         return False
 
@@ -97,11 +114,12 @@ class Reporter(threading.Thread):
         return any(level.modified(self) for level in self.levels)
 
     def check_urls(self):
+        """Check whether we have new schedulers to notify"""
         mtime = os.path.getmtime(self.path)
         if mtime > self.lastcheck:
             for f in self.path.iterdir():
-                self.urls[f.name] = f.read_text().strip()
-                logger.info("Added new notification URL: %s", self.urls[f.name])
+                self.urls[f.name] = ListenerInformation(f.read_text().strip())
+                logger.info("Added new notification URL: %s", self.urls[f.name].url)
                 f.unlink()
 
             self.lastcheck = os.path.getmtime(self.path)
@@ -128,7 +146,9 @@ class Reporter(threading.Thread):
                         params = level.report()
 
                         # Go over all URLs
-                        for key, baseurl in self.urls.items():
+                        for key, info in self.urls.items():
+                            baseurl = info.url
+
                             url = "{}/progress?{}".format(
                                 baseurl, urllib.parse.urlencode(params)
                             )
@@ -147,7 +167,7 @@ class Reporter(threading.Thread):
                                     url,
                                     e,
                                 )
-                                if Reporter.isfatal_httperror(e):
+                                if Reporter.isfatal_httperror(e, info):
                                     toremove.append(key)
 
                 # Removes unvalid URLs
@@ -165,12 +185,13 @@ class Reporter(threading.Thread):
             self.check_urls()
             if self.urls:
                 # Go over all URLs
-                for key, baseurl in self.urls.items():
+                for key, info in self.urls.items():
+                    baseurl = info.url
                     url = "{}?status=eoj".format(baseurl)
                     try:
                         with urlopen(url) as _:
                             logger.debug(
-                                "EOJ botification sent for %s",
+                                "EOJ notification sent for %s",
                                 baseurl,
                             )
                     except Exception:
@@ -178,6 +199,8 @@ class Reporter(threading.Thread):
                             "Could not report EOJ",
                         )
 
+            self.file_reporter.eoj()
+
     def set_progress(
         self, progress: float, level: int, desc: Optional[str], console=False
     ):
@@ -196,6 +219,8 @@ class Reporter(threading.Thread):
                     self.levels[level].desc = desc
                 self.levels[level].progress = progress
 
+                self.file_reporter.set_progress(progress, level, desc)
+
                 self.cv.notify_all()
 
     INSTANCE: ClassVar[Optional["Reporter"]] = None
@@ -206,6 +231,7 @@ class Reporter(threading.Thread):
             taskpath = TaskEnv.instance().taskpath
             assert taskpath is not None, "Task path is not defined"
             Reporter.INSTANCE = Reporter(taskpath)
+            Reporter.INSTANCE.start()
         return Reporter.INSTANCE
 
 
@@ -243,7 +269,7 @@ class xpm_tqdm(std_tqdm):
 
     def update(self, n=1):
         result = super().update(n)
-        if self.total is not None:
+        if self.total is not None and self.total > 0:
             progress(self.n / self.total, level=self.pos, console=False)
         return result
 

experimaestro/progress.py

@@ -0,0 +1,406 @@
+"""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:
+    # TODO: Implement buffering and flushing
+
+    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/run.py

@@ -8,6 +8,7 @@ import json
 from typing import List
 import fasteners
 from experimaestro.notifications import progress, report_eoj
+from experimaestro.utils.multiprocessing import delayed_shutdown
 from .core.types import ObjectType
 from experimaestro.utils import logger
 from experimaestro.core.objects import ConfigInformation
@@ -92,15 +93,27 @@ class TaskRunner:
             logger.info("Finished cleanup")
 
     def handle_error(self, code, frame_type):
-        logger.info("Handling error")
+        logger.info("Error handler: finished with code %d", code)
         self.failedpath.write_text(str(code))
         self.cleanup()
+        logger.info("Exiting")
+        delayed_shutdown(60, exit_code=code)
         sys.exit(1)
 
     def run(self):
         atexit.register(self.cleanup)
-        signal.signal(signal.SIGTERM, self.handle_error)
-        signal.signal(signal.SIGINT, self.handle_error)
+        sigterm_handler = signal.signal(signal.SIGTERM, self.handle_error)
+        sigint_handler = signal.signal(signal.SIGINT, self.handle_error)
+
+        def remove_signal_handlers(remove_cleanup=True):
+            """Removes cleanup in forked processes"""
+            signal.signal(signal.SIGTERM, sigterm_handler)
+            signal.signal(signal.SIGINT, sigint_handler)
+            atexit.unregister(self.cleanup)
+
+        if sys.platform != "win32":
+            os.register_at_fork(after_in_child=remove_signal_handlers)
+
         try:
             workdir = self.scriptpath.parent
             os.chdir(workdir)
@@ -128,7 +141,11 @@ class TaskRunner:
                 self.started = True
                 run(workdir / "params.json")
 
+                # ... remove the handlers
+                remove_signal_handlers(remove_cleanup=False)
+
                 # Everything went OK
+                logger.info("Task ended successfully")
                 sys.exit(0)
         except Exception:
             logger.exception("Got exception while running")
@@ -136,6 +153,10 @@ class TaskRunner:
 
         except SystemExit as e:
             if e.code == 0:
+                # Normal exit, just create the ".done" file
                 self.donepath.touch()
+
+                # ... and finish the exit process
+                raise
             else:
                 self.handle_error(e.code, 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",
+]