experimaestro 2.0.0b4__py3-none-any.whl → 2.0.0b17__py3-none-any.whl

experimaestro/core/partial_lock.py

@@ -0,0 +1,394 @@
+"""Partial directory locking for exclusive access.
+
+This module provides locking for partial directories to ensure only one job
+can write to a partial at a time.
+"""
+
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Dict, Type
+import fasteners
+import logging
+
+from experimaestro.locking import (
+    DynamicDependencyLock,
+    DynamicLockFile,
+    JobDependencyLock,
+    LockError,
+    TrackedDynamicResource,
+)
+from experimaestro.dynamic import DynamicDependency
+from experimaestro.scheduler.dependencies import Resource
+
+
+logger = logging.getLogger("xpm.partial")
+
+
+class PartialLockFile(DynamicLockFile):
+    """Lock file for partial directories.
+
+    Stores:
+    - job_uri: Reference to the job holding the lock
+    - information: {"partial_name": name_of_partial}
+
+    The lock file is created in the partial directory when a job acquires
+    exclusive access, enabling recovery if the scheduler or job crashes.
+    """
+
+    partial_name: str
+
+    def from_information(self, info) -> None:
+        """Set partial name from information dict."""
+        if info is None:
+            # Creating a new lock file
+            self.partial_name = ""
+        elif isinstance(info, dict):
+            self.partial_name = info.get("partial_name", "")
+        else:
+            raise ValueError(f"Invalid information format: {info}")
+
+    def to_information(self) -> dict:
+        """Return partial name for JSON serialization."""
+        return {"partial_name": self.partial_name}
+
+
+class PartialJobResource(Resource, TrackedDynamicResource):
+    """Resource for partial directory locking.
+
+    Tracks the number of jobs holding write locks on the partial directory.
+    Uses file-based tracking for recovery, similar to CounterToken.
+
+    File structure in partial directory:
+    - {partial_path}/.experimaestro/locks/ipc.lock: IPC lock
+    - {partial_path}/.experimaestro/locks/informations.json: {"max_write_locks": 1}
+    - {partial_path}/.experimaestro/locks/jobs/{task_id}/{identifier}.json: lock file
+    """
+
+    #: Lock file class for partial lock files
+    lock_file_class: Type[DynamicLockFile] = PartialLockFile
+
+    #: Maps partial paths to PartialJobResource instances (singleton pattern)
+    RESOURCES: Dict[str, "PartialJobResource"] = {}
+
+    @staticmethod
+    def forkhandler():
+        """Clear resources after fork to avoid sharing state."""
+        PartialJobResource.RESOURCES = {}
+
+    @staticmethod
+    def create(partial_path: Path, max_write_locks: int = 1) -> "PartialJobResource":
+        """Get or create a PartialJobResource for the given path.
+
+        This implements a singleton pattern per path to ensure only one
+        resource instance exists for each partial directory.
+
+        Args:
+            partial_path: Path to the partial directory
+            max_write_locks: Maximum number of concurrent write locks (default: 1)
+
+        Returns:
+            PartialJobResource for the path
+        """
+        key = str(partial_path)
+        resource = PartialJobResource.RESOURCES.get(key)
+        if resource is None:
+            resource = PartialJobResource(partial_path, max_write_locks)
+            PartialJobResource.RESOURCES[key] = resource
+        return resource
+
+    @property
+    def lock_folder(self) -> Path:
+        """Path to the lock folder within the partial directory."""
+        return self._partial_path / ".experimaestro" / "locks"
+
+    @property
+    def partial_path(self) -> Path:
+        """Path to the partial directory."""
+        return self._partial_path
+
+    def _write_informations(self, max_write_locks: int) -> None:
+        """Write partial informations to disk."""
+        self.lock_folder.mkdir(parents=True, exist_ok=True)
+        with self.informations_path.open("w") as f:
+            json.dump({"max_write_locks": max_write_locks}, f)
+
+    def _read_informations(self) -> int:
+        """Read max_write_locks from informations file."""
+        try:
+            with self.informations_path.open("r") as f:
+                data = json.load(f)
+                return data.get("max_write_locks", 1)
+        except FileNotFoundError:
+            return 1
+
+    def __init__(self, partial_path: Path, max_write_locks: int = 1):
+        """Initialize a partial job resource.
+
+        Args:
+            partial_path: Path to the partial directory
+            max_write_locks: Maximum number of concurrent write locks
+        """
+        self._partial_path = partial_path
+        self.max_write_locks = max_write_locks
+        self.write_locks = 0  # Current number of write locks
+
+        # Create informations file if needed
+        self.lock_folder.mkdir(parents=True, exist_ok=True)
+        if not self.informations_path.is_file():
+            self._write_informations(max_write_locks)
+
+        # Initialize base classes
+        Resource.__init__(self)
+        TrackedDynamicResource.__init__(self, str(partial_path))
+
+    def __str__(self):
+        return f"partial[{self._partial_path}]"
+
+    # --- TrackedDynamicResource abstract method implementations ---
+
+    def _reset_state(self) -> None:
+        """Reset state before re-reading lock files."""
+        self.max_write_locks = self._read_informations()
+        self.write_locks = 0
+
+    def _account_lock_file(self, lf: DynamicLockFile) -> None:
+        """Account for a lock file - increment write lock count."""
+        self.write_locks += 1
+
+    def _unaccount_lock_file(self, lf: DynamicLockFile) -> None:
+        """Unaccount a lock file - decrement write lock count."""
+        self.write_locks -= 1
+
+    def is_available(self, dependency: "PartialDependency") -> bool:
+        """Check if the partial is available (write locks not exhausted)."""
+        return self.write_locks < self.max_write_locks
+
+    def _do_acquire(self, dependency: "PartialDependency") -> None:
+        """Increment write lock count."""
+        self.write_locks += 1
+        logger.debug(
+            "Partial state [acquired]: write_locks %d, max %d, path %s",
+            self.write_locks,
+            self.max_write_locks,
+            self._partial_path,
+        )
+
+    def _do_release(self, dependency: "PartialDependency") -> None:
+        """Decrement write lock count."""
+        self.write_locks -= 1
+        logger.debug(
+            "Partial state [released]: write_locks %d, max %d, path %s",
+            self.write_locks,
+            self.max_write_locks,
+            self._partial_path,
+        )
+
+    def _get_lock_file_information(self, dependency: "PartialDependency"):
+        """Return partial name for lock file."""
+        return {"partial_name": dependency.partial_name}
+
+    def _handle_information_change(self) -> None:
+        """Handle max_write_locks changes from informations.json."""
+        max_write_locks = self._read_informations()
+        delta = max_write_locks - self.max_write_locks
+        self.max_write_locks = max_write_locks
+        logger.debug(
+            "Partial information modified: write_locks %d, max %d, path %s",
+            self.write_locks,
+            self.max_write_locks,
+            self._partial_path,
+        )
+
+        # Notify waiting tasks if more locks became available
+        if delta > 0:
+            self.available_condition.notify_all()
+
+    # --- Partial API ---
+
+    def dependency(self, partial_name: str) -> "PartialDependency":
+        """Create a dependency on this partial resource.
+
+        Args:
+            partial_name: Name of the partial (for symlink creation)
+
+        Returns:
+            PartialDependency for this resource
+        """
+        return PartialDependency(self, partial_name)
+
+
+# Register fork handler to clear resources after fork
+if sys.platform != "win32":
+    os.register_at_fork(after_in_child=PartialJobResource.forkhandler)
+
+
+class PartialJobLock(JobDependencyLock):
+    """Job-side lock for a partial directory.
+
+    Inherits from JobDependencyLock to participate in the dynamic lock lifecycle.
+    Uses IPC locking (fasteners.InterProcessLock) for exclusive access.
+
+    File structure:
+    - {partial_path}/.experimaestro/locks/ipc.lock: IPC lock
+    - {partial_path}/.experimaestro/locks/jobs/{task_id}/{identifier}.json: lock file
+
+    Lifecycle:
+    1. Scheduler acquires IPC lock and releases it in aio_job_started
+    2. Job acquires IPC lock (blocks until scheduler releases)
+    3. Job creates/updates lock file to track who holds the lock
+    4. On release: IPC lock released, lock file deleted (via base class)
+    """
+
+    def __init__(self, data: dict):
+        self.partial_path = Path(data["partial_path"])
+        self.partial_name = data["partial_name"]
+        self.job_uri = data["job_uri"]
+        self.lock_folder = self.partial_path / ".experimaestro" / "locks"
+        self.ipc_lock_path = self.lock_folder / "ipc.lock"
+        self.lock_file_path = Path(data["lock_file_path"])
+        self._lock = None
+
+    def acquire(self) -> None:
+        """Acquire exclusive lock on the partial directory.
+
+        Verifies the lock file exists (created by scheduler) and acquires
+        the IPC lock for exclusive access.
+        """
+        # Verify lock file exists (base class)
+        super().acquire()
+
+        logger.info("Acquiring partial lock: %s", self.partial_path)
+
+        # Acquire the IPC lock (blocking - waits for other jobs to finish)
+        self._lock = fasteners.InterProcessLock(str(self.ipc_lock_path))
+        if not self._lock.acquire(blocking=True):
+            raise LockError(f"Could not lock partial: {self.partial_path}")
+
+        logger.info("Acquired partial lock: %s", self.partial_path)
+
+    def release(self) -> None:
+        """Release the partial lock and delete the lock file."""
+        if self._lock is not None and self._lock.acquired:
+            logger.info("Releasing partial lock: %s", self.partial_path)
+            self._lock.release()
+            self._lock = None
+
+            # Delete the lock file (handled by base class)
+            super().release()
+
+
+class PartialLock(DynamicDependencyLock):
+    """Scheduler-side lock for a partial directory.
+
+    Inherits from DynamicDependencyLock to participate in the dynamic lock lifecycle.
+    Ensures exclusive access to a partial directory while a job is running.
+
+    Manages lock acquisition/release through the PartialJobResource, similar to
+    how CounterTokenLock manages through CounterToken.
+
+    File structure:
+    - {partial_path}/.experimaestro/locks/ipc.lock: IPC lock
+    - {partial_path}/.experimaestro/locks/jobs/{task_id}/{identifier}.json: lock file
+    """
+
+    dependency: "PartialDependency"
+
+    def __init__(self, dependency: "PartialDependency"):
+        super().__init__(dependency)
+
+    @property
+    def lock_folder(self) -> Path:
+        """Path to the lock folder within the partial directory."""
+        return self.dependency.resource.lock_folder
+
+    def _acquire(self):
+        """Acquire exclusive lock via the resource."""
+        self.dependency.resource.acquire(self.dependency)
+
+    def _release(self):
+        """Release the lock via the resource."""
+        self.dependency.resource.release(self.dependency)
+
+    def __str__(self):
+        return f"PartialLock({self.dependency.partial_name})"
+
+    async def aio_job_before_start(self, job) -> None:
+        """Create symlink before job starts."""
+        from experimaestro.scheduler.jobs import Job
+
+        assert isinstance(job, Job)
+
+        # Create symlink in job's .experimaestro/partials directory
+        # The symlink points to the partial data directory (not the lock folder)
+        partials_dir = job.experimaestro_path / "partials"
+        partials_dir.mkdir(parents=True, exist_ok=True)
+
+        symlink_path = partials_dir / self.dependency.partial_name
+        if not symlink_path.exists():
+            symlink_path.symlink_to(self.dependency.partial_path)
+            logger.debug(
+                "Created partial symlink: %s -> %s",
+                symlink_path,
+                self.dependency.partial_path,
+            )
+
+    def to_json(self) -> dict:
+        """Serialize lock for job process."""
+        data = super().to_json()
+        data.update(
+            {
+                "partial_path": str(self.dependency.partial_path),
+                "partial_name": self.dependency.partial_name,
+                "lock_file_path": str(self.lock_file_path),
+                "job_uri": str(self.dependency.target.basepath),
+            }
+        )
+        return data
+
+    @classmethod
+    def from_json(cls, data: dict) -> PartialJobLock:
+        """Create job-side lock from serialized data."""
+        return PartialJobLock(data)
+
+
+class PartialDependency(DynamicDependency):
+    """A dependency on a partial directory (dynamic - availability can change).
+
+    This ensures that only one job can write to a partial directory at a time.
+    Uses PartialJobResource for state tracking and recovery.
+    """
+
+    def __init__(self, resource: PartialJobResource, partial_name: str):
+        """Create a partial dependency.
+
+        Args:
+            resource: The PartialJobResource managing this partial
+            partial_name: Name of the partial (used for symlink creation)
+        """
+        super().__init__(resource)
+        self._resource = resource
+        self._partial_name = partial_name
+
+    def _create_lock(self) -> PartialLock:
+        """Create a partial lock for this dependency."""
+        return PartialLock(self)
+
+    @property
+    def resource(self) -> PartialJobResource:
+        """The resource managing this partial."""
+        return self._resource
+
+    @property
+    def partial_path(self) -> Path:
+        """Path to the partial directory."""
+        return self._resource.partial_path
+
+    @property
+    def partial_name(self) -> str:
+        """Name of the partial."""
+        return self._partial_name
+
+    def __repr__(self) -> str:
+        return f"PartialDep[{self.partial_name}]"

experimaestro/core/types.py

@@ -30,7 +30,7 @@ if typing.TYPE_CHECKING:
     from experimaestro.scheduler.base import Job
     from experimaestro.launchers import Launcher
     from experimaestro.core.objects import Config
-    from experimaestro.core.subparameters import Subparameters
+    from experimaestro.core.partial import Partial
 
 
 @dataclass
@@ -310,8 +310,8 @@ class ObjectType(Type):
         # --- Value class (for external value types, e.g., nn.Module subclasses)
         self._original_type: type = tp  # Keep reference to original config class
 
-        # --- Subparameters for partial identifier computation
-        self._subparameters: Dict[str, "Subparameters"] = {}
+        # --- Partial for partial identifier computation
+        self._partials: Dict[str, "Partial"] = {}
 
     def set_value_type(self, value_class: type) -> None:
         """Register an explicit value class for this configuration.
@@ -357,15 +357,13 @@ class ObjectType(Type):
         else:
             self._file = Path(inspect.getfile(self.originaltype)).absolute()
 
-        assert (
-            self._module and self._package
-        ) or self._file, f"Could not detect module/file for {self.originaltype}"
+        assert (self._module and self._package) or self._file, (
+            f"Could not detect module/file for {self.originaltype}"
+        )
 
         # The class of the object
 
-        self._arguments = ChainMap(
-            {}, *(tp.arguments for tp in self.parents())
-        )  # type: ChainMap[Argument, Any]
+        self._arguments = ChainMap({}, *(tp.arguments for tp in self.parents()))  # type: ChainMap[Argument, Any]
 
         # Add arguments from annotations
         for annotation in self.annotations:
@@ -409,15 +407,15 @@ class ObjectType(Type):
                                 )
                                 raise
 
-        # Collect subparameters from class attributes
-        from .subparameters import Subparameters as SubparametersClass
+        # Collect partial from class attributes
+        from .partial import Partial as PartialClass
 
         for name, value in self._original_type.__dict__.items():
-            if isinstance(value, SubparametersClass):
+            if isinstance(value, PartialClass):
                 # Auto-set name from attribute name if not already set
                 if value.name is None:
                     value.name = name
-                self._subparameters[name] = value
+                self._partials[name] = value
 
     def name(self):
         return f"{self.value_type.__module__}.{self.value_type.__qualname__}"
@@ -488,8 +486,7 @@ class ObjectType(Type):
             # Legacy mechanism: parent class is the target
             if len(self.value_type.__bases__) != 1:
                 raise RuntimeError(
-                    "Deprecated configurations must have "
-                    "only one parent (the new configuration)"
+                    "Deprecated configurations must have only one parent (the new configuration)"
                 )
             parent = self.value_type.__bases__[0].__getxpmtype__()
             self.identifier = parent.identifier