experimaestro 1.6.1__py3-none-any.whl → 1.7.0__py3-none-any.whl

experimaestro/scheduler/dynamic_outputs.py

@@ -0,0 +1,184 @@
+"""Handles dynamic task outputs"""
+
+import asyncio
+import json
+import logging
+import queue
+import threading
+from collections import defaultdict
+from functools import cached_property
+from pathlib import Path
+from typing import Callable, TYPE_CHECKING
+
+from watchdog.events import FileSystemEventHandler
+
+from experimaestro.ipc import ipcom
+from experimaestro.utils import logger
+
+from .base import Job, experiment
+
+if TYPE_CHECKING:
+    from experimaestro.core.objects import WatchedOutput
+
+
+class TaskOutputCallbackHandler:
+    def __init__(self, converter: Callable):
+        pass
+
+
+class TaskOutputs(FileSystemEventHandler):
+    """Represent and monitors dynamic outputs generated by one task"""
+
+    #: Global dictionary for handles
+    HANDLERS: dict[Path, "TaskOutputs"] = {}
+
+    #: Global lock to access current HANDLERS
+    LOCK = threading.Lock()
+
+    def create(job: Job):
+        with TaskOutputs.LOCK:
+            if instance := TaskOutputs.get(job.task_outputs_path, None):
+                return instance
+
+            instance = TaskOutputs(job.task_outputs_path)
+            TaskOutputs[job.task_outputs_path] = instance
+            return instance
+
+    def __init__(self, path: Path):
+        """Monitors an event path"""
+        logger.debug("Watching dynamic task outputs in %s", path)
+        self.path = path
+        self.handle = None
+        self.count = 0
+        self.lock = threading.Lock()
+        self.listeners: dict[str, dict[Callable, set[Callable]]] = defaultdict(
+            lambda: defaultdict(set)
+        )
+
+        #: The events registered so far
+        self.events = []
+
+    def __enter__(self):
+        """Starts monitoring task outputs"""
+        self.job.task_outputs_path.parent.mkdir(parents=True, exist_ok=True)
+        with self.lock:
+            if self.handle is None:
+                assert self.count == 0
+                self.handle = ipcom().fswatch(self, self.path.parent, False)
+            self.count += 1
+        return self
+
+    def __exit__(self, *args):
+        """Stops monitoring task outputs"""
+        with self.lock:
+            self.count -= 1
+            if self.count == 0:
+                ipcom().fsunwatch(self.handle)
+                self.fh.close()
+
+                self.handle = None
+                self._fh = None
+
+    def watch_output(self, watched: "WatchedOutput"):
+        """Add a new listener"""
+        key = f"{watched.config.__identifier__}/{watched.method_name}"
+        with self.lock:
+            # Process events so far
+            listener = self.listeners[key].get(watched.method, None)
+            if listener is None:
+                listener = TaskOutputCallbackHandler(watched.method)
+
+            # Register
+            self.listeners[key][watched.method].add(watched.callback)
+
+    #
+    # --- Events
+    #
+
+    @cached_property
+    def fh(self):
+        if self._fh is None:
+            self._fh = self.path.open("rt")
+        return self._fh
+
+    def on_modified(self, event):
+        self.handle(Path(event.src_path))
+
+    def on_created(self, event):
+        self.handle(Path(event.src_path))
+
+    def handle(self, path: Path):
+        if path != self.path:
+            return
+
+        with self.lock:
+            logger.debug("[TASK OUTPUT] Handling task output for %s", self.path)
+
+            while json_line := self.fh.readline():
+                # Read the event
+                event = json.loads(json_line)
+                logger.debug("Event: %s", event)
+
+                # FIXME: move elsewhere
+                # # Process the event
+                # event = self.config_method(
+                #     self.job.config.__xpm__.mark_output,
+                #     *event["args"],
+                #     **event["kwargs"],
+                # )
+
+                self.events.append(event)
+                # self.job.scheduler.xp.taskOutputsWorker.add(self, event)
+
+
+class TaskOutputsWorker(threading.Thread):
+    """This worker process dynamic output queue for one experiment"""
+
+    def __init__(self, xp: experiment):
+        super().__init__(name="task outputs worker", daemon=True)
+        self.queue = queue.Queue()
+        self.xp = xp
+
+    def watch_output(self, watched: "WatchedOutput"):
+        """Watch an output
+
+        :param watched: The watched output specification
+        """
+        logger.debug("Registering task output listener %s", watched)
+
+        # path = watched.job.tasks_output_path
+        TaskOutputs.create(watched.job).watch_output(watched)
+
+    def add(self, watcher, event):
+        asyncio.run_coroutine_threadsafe(
+            self.xp.update_task_output_count(1),
+            self.xp.scheduler.loop,
+        ).result()
+        self.queue.put((watcher, event))
+
+    def run(self):
+        logging.debug("Starting output listener queue")
+        while True:
+            # Get the next element in the queue
+            element = self.queue.get()
+            if element is None:
+                # end of processing
+                break
+
+            # Call all the listeners
+            logging.debug("Got one event: %s", element)
+            watcher, event = element
+            for listener in watcher.listeners:
+                try:
+                    logger.debug("Calling listener [%s] with %s", listener, event)
+                    listener(event)
+                    logger.debug(
+                        "[done] Calling listener [%s] with %s", listener, event
+                    )
+                except Exception:
+                    logging.exception("Exception while calling the listener")
+                self.queue.task_done()
+
+                asyncio.run_coroutine_threadsafe(
+                    self.xp.update_task_output_count(-1), self.xp.scheduler.loop
+                ).result()

experimaestro/scheduler/workspace.py

@@ -2,7 +2,7 @@ from collections import ChainMap
 from enum import Enum
 from functools import cached_property
 from pathlib import Path
-from typing import Optional
+from typing import Iterator, Optional
 from experimaestro.settings import WorkspaceSettings, Settings
 
 
@@ -46,6 +46,7 @@ class Workspace:
             path = path.absolute()
         self.path = path
         self.run_mode = run_mode
+        self.python_path = []
         from ..launchers import Launcher
 
         self.launcher = launcher or Launcher.get(path)

experimaestro/scriptbuilder.py

@@ -51,6 +51,8 @@ class PythonScriptBuilder:
         self.lockfiles: List[Path] = []
         self.notificationURL: Optional[str] = None
         self.command: Optional[AbstractCommand] = None
+
+        # This is used to serialize the full process identifier on disk
         self.processtype = "local"
 
     def write(self, job: CommandLineJob):
@@ -63,7 +65,7 @@ class PythonScriptBuilder:
             job {CommandLineJob} -- [description]
 
         Returns:
-            [type] -- [description]
+            str -- The script path on disk
         """
         assert isinstance(
             job, CommandLineJob
@@ -94,6 +96,7 @@ class PythonScriptBuilder:
             out.write("# Experimaestro generated task\n\n")
             out.write(
                 """import logging\n"""
+                """import sys\n"""
                 """logging.basicConfig(level=logging.INFO, """
                 """format='%(levelname)s:%(process)d:%(asctime)s [%(name)s] %(message)s', datefmt='%y-%m-%d %H:%M:%S')\n\n"""
             )
@@ -112,9 +115,17 @@ class PythonScriptBuilder:
             out.write("    ]\n")
 
             for name, value in job.environ.items():
-                out.write(f"""    os.environ["{name}"] = "{shquote(value)}"\n""")
+                if name == "PYTHONPATH":
+                    # Handles properly python path
+                    for path in value.split(":"):
+                        out.write(f"""    sys.path.insert(0, "{shquote(path)}")\n""")
+                else:
+                    out.write(f"""    os.environ["{name}"] = "{shquote(value)}"\n""")
             out.write("\n")
 
+            for path in job.python_path:
+                out.write(f"""    sys.path.insert(0, "{shquote(str(path))}")\n""")
+
             out.write(
                 f"""    TaskRunner("{shquote(connector.resolve(scriptpath))}","""
                 """ lockfiles).run()\n"""