learning-loop-node 0.24.0__tar.gz → 0.25.0__tar.gz

{learning_loop_node-0.24.0 → learning_loop_node-0.25.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: learning_loop_node
-Version: 0.24.0
+Version: 0.25.0
 Summary: Python Library for Nodes which connect to the Zauberzeug Learning Loop
 Author-email: Zauberzeug GmbH <info@zauberzeug.com>
 License: MIT
@@ -67,6 +67,7 @@ You can configure connection to our Learning Loop by specifying the following en
 | RESTART_AFTER_TRAINING   | -            | Restart the trainer after training (set to 1)                | Trainer (opt.)            | 0            |
 | KEEP_OLD_TRAININGS       | -            | Do not delete old trainings (set to 1)                       | Trainer (opt.)            | 0            |
 | TRAINER_IDLE_TIMEOUT_SEC | -            | Automatically shutdown trainer after timeout (in seconds)    | Trainer (opt.)            | 0 (disabled) |
+| TRAINING_SLOT_LOCK       | -            | Lock file shared by the trainers on one GPU (see Trainer Node) | Trainer (opt.)          | - (disabled) |
 | USE_BACKDOOR_CONTROLS    | -            | Always enable backdoor controls (set to 1)                   | Trainer / Detector (opt.) | 0            |
 
 Note that organization and project IDs are always lower case and may differ from the names in the Learning Loop which can have uppercase letters.
@@ -250,6 +251,20 @@ Trainers fetch the images and anntoations from the Learning Loop to train new mo
 
 - if the command line tool "jpeginfo" is installed, the downloader will drop corrupted images automatically
 
+### Several trainers on one GPU
+
+Trainers that share a GPU (say, one per model architecture on the same machine) coordinate through
+a lock file: mount one host directory into each of them and point `TRAINING_SLOT_LOCK` at a file
+inside it, e.g. `/slot/training.lock`. A trainer takes the lock when it accepts a training and holds
+it until the training is cleaned up, including detection and upload. While a sibling holds it, the
+others report the state `blocked` instead of `idle`, so the Learning Loop neither offers them for a
+training nor dispatches queued jobs to them; a trainer that is sent `begin_training` regardless
+refuses it with `{"success": false, "error_msg": ...}` in the acknowledgement and sets nothing up. A
+training resumed after a restart cannot be refused — it waits for the lock instead, reporting
+`waiting_for_slot` meanwhile. The lock is an `flock`, so the kernel releases it whenever the holding
+process ends, however it ends — there is nothing to clean up after a crash. Mount the directory,
+not the file, and never replace the file: the lock belongs to its inode.
+
 ## Converter Node
 
 A Conveter Node converts models from one format into another.

{learning_loop_node-0.24.0 → learning_loop_node-0.25.0}/README.md

@@ -34,6 +34,7 @@ You can configure connection to our Learning Loop by specifying the following en
 | RESTART_AFTER_TRAINING   | -            | Restart the trainer after training (set to 1)                | Trainer (opt.)            | 0            |
 | KEEP_OLD_TRAININGS       | -            | Do not delete old trainings (set to 1)                       | Trainer (opt.)            | 0            |
 | TRAINER_IDLE_TIMEOUT_SEC | -            | Automatically shutdown trainer after timeout (in seconds)    | Trainer (opt.)            | 0 (disabled) |
+| TRAINING_SLOT_LOCK       | -            | Lock file shared by the trainers on one GPU (see Trainer Node) | Trainer (opt.)          | - (disabled) |
 | USE_BACKDOOR_CONTROLS    | -            | Always enable backdoor controls (set to 1)                   | Trainer / Detector (opt.) | 0            |
 
 Note that organization and project IDs are always lower case and may differ from the names in the Learning Loop which can have uppercase letters.
@@ -217,6 +218,20 @@ Trainers fetch the images and anntoations from the Learning Loop to train new mo
 
 - if the command line tool "jpeginfo" is installed, the downloader will drop corrupted images automatically
 
+### Several trainers on one GPU
+
+Trainers that share a GPU (say, one per model architecture on the same machine) coordinate through
+a lock file: mount one host directory into each of them and point `TRAINING_SLOT_LOCK` at a file
+inside it, e.g. `/slot/training.lock`. A trainer takes the lock when it accepts a training and holds
+it until the training is cleaned up, including detection and upload. While a sibling holds it, the
+others report the state `blocked` instead of `idle`, so the Learning Loop neither offers them for a
+training nor dispatches queued jobs to them; a trainer that is sent `begin_training` regardless
+refuses it with `{"success": false, "error_msg": ...}` in the acknowledgement and sets nothing up. A
+training resumed after a restart cannot be refused — it waits for the lock instead, reporting
+`waiting_for_slot` meanwhile. The lock is an `flock`, so the kernel releases it whenever the holding
+process ends, however it ends — there is nothing to clean up after a crash. Mount the directory,
+not the file, and never replace the file: the lock belongs to its inode.
+
 ## Converter Node
 
 A Conveter Node converts models from one format into another.

{learning_loop_node-0.24.0 → learning_loop_node-0.25.0}/learning_loop_node/enums/trainer.py

@@ -4,7 +4,9 @@ from enum import Enum
 
 class TrainerState(str, Enum):
     Idle = 'idle'
+    Blocked = 'blocked'  # idle, but unable to take a training
     Initialized = 'initialized'
+    WaitingForSlot = 'waiting_for_slot'
     Preparing = 'preparing'
     DataDownloading = 'data_downloading'
     DataDownloaded = 'data_downloaded'

{learning_loop_node-0.24.0 → learning_loop_node-0.25.0}/learning_loop_node/helpers/misc.py

@@ -150,7 +150,7 @@ def ensure_socket_response(func):
             if isinstance(value, bool):
                 return asdict(SocketResponse.from_bool(value))
             if isinstance(value, SocketResponse):
-                return value
+                return asdict(value)
             if (args[0] in ['connect', 'disconnect', 'connect_error']):
                 return value
             if value is None:

{learning_loop_node-0.24.0 → learning_loop_node-0.25.0}/learning_loop_node/trainer/exceptions.py

@@ -5,6 +5,18 @@ class CriticalError(Exception):
     """
 
 
+class TrainingRefusedError(Exception):
+    """Raised when a trainer cannot take on a training it was asked to begin; nothing was set up."""
+
+
+class TrainingSlotTakenError(TrainingRefusedError):
+    """Raised when a sibling node holds the training slot."""
+
+    def __init__(self, holder: str) -> None:
+        super().__init__(f'the training slot is held by {holder}')
+        self.holder = holder
+
+
 class InsufficientMemoryError(RuntimeError):
     """Raised when not even the smallest unit of work fits in memory."""
 

{learning_loop_node-0.24.0 → learning_loop_node-0.25.0}/learning_loop_node/trainer/trainer_logic_generic.py

@@ -26,8 +26,9 @@ from ..helpers.misc import (
     is_valid_uuid4,
 )
 from .downloader import TrainingsDownloader
-from .exceptions import CriticalError, NodeNeedsRestartError
+from .exceptions import CriticalError, NodeNeedsRestartError, TrainingRefusedError
 from .io_helpers import ActiveTrainingIO, EnvironmentVars, LastTrainingIO
+from .training_slot import TrainingSlot, training_slot_from_env
 
 if TYPE_CHECKING:
     from .trainer_node import TrainerNode
@@ -37,7 +38,7 @@ logger = logging.getLogger('learning_loop_node.trainer_logic_generic')
 
 class TrainerLogicGeneric(ABC):
 
-    def __init__(self, model_format: str) -> None:
+    def __init__(self, model_format: str, training_slot: TrainingSlot | None = None) -> None:
 
         # NOTE: model_format is used in the file path for the model on the server:
         # It acts as a key for list of files (cf. _get_latest_model_files)
@@ -54,6 +55,7 @@ class TrainerLogicGeneric(ABC):
         self._training: Optional[Training] = None
         self._active_training_io: Optional[ActiveTrainingIO] = None
         self._environment_vars = EnvironmentVars()
+        self.training_slot = training_slot or training_slot_from_env()
 
     # ---------------------------------------- PROPERTIES TO AVOID CHECKING FOR NONE ----------------------------------------
 
@@ -98,9 +100,17 @@ class TrainerLogicGeneric(ABC):
 
     @property
     def state(self) -> str:
-        """Returns the current state of the training. Used solely by the node in send_status().
+        """Returns the current state of the training.
+
+        Without a training the state is `idle`, or `blocked` while a sibling node holds the training slot.
+        A training resumed after a restart reports `waiting_for_slot` until it gets the slot: it has a
+        training, so it is never `blocked`.
         """
-        if (not self.training_active) or (self.training.training_state is None):
+        if not self.training_active:
+            return TrainerState.Blocked.value if self.training_slot.taken_by_sibling else TrainerState.Idle.value
+        if not self.training_slot.held:
+            return TrainerState.WaitingForSlot.value
+        if self.training.training_state is None:
             return TrainerState.Idle.value
         return self.training.training_state
 
@@ -213,8 +223,21 @@ class TrainerLogicGeneric(ABC):
 
     async def begin_training(self, organization: str, project: str, training_config: Dict) -> None:
         """Called on `begin_training` event from the Learning Loop.
+
+        The training slot is taken before anything is set up, so a refused training leaves the
+        previous training folders in place.
+
+        :raise TrainingRefusedError: if a training is active or a sibling node holds the training slot.
         """
-        self._init_new_training(Context(organization=organization, project=project), training_config)
+        if self.training_active:
+            raise TrainingRefusedError(f'training {self.training.id} is still active')
+        context = Context(organization=organization, project=project)
+        self.training_slot.claim(self._slot_claim(context, training_config.get('training_number')))
+        try:
+            self._init_new_training(context, training_config)
+        except Exception:
+            self.training_slot.release()
+            raise
         self._begin_training_task()
 
     def _begin_training_task(self) -> None:
@@ -242,6 +265,7 @@ class TrainerLogicGeneric(ABC):
         """
         self.errors.reset_all()
         try:
+            await self.training_slot.wait_for(self._slot_claim(self.training.context, self.training.training_number))
             await self._training_loop()
         except asyncio.CancelledError:
             if not self.shutdown_event.is_set():
@@ -254,6 +278,12 @@ class TrainerLogicGeneric(ABC):
                 logger.info('CancelledError in _run - shutting down')
         except Exception:
             logger.exception('(Ignored) exception in trainer_logic._run:')
+        finally:
+            self._training = None
+            self.training_slot.release()
+
+    def _slot_claim(self, context: Context, training_number: Any) -> str:
+        return f'node {self.node.uuid}, training #{training_number} of {context.organization}/{context.project}'
 
     # ---------------------------------------- TRAINING STATES ----------------------------------------
 

{learning_loop_node-0.24.0 → learning_loop_node-0.25.0}/learning_loop_node/trainer/trainer_node.py

@@ -7,7 +7,10 @@ from typing import Dict, Optional
 from fastapi.encoders import jsonable_encoder
 from socketio import AsyncClient, exceptions
 
+from ..data_classes import SocketResponse
+from ..enums import TrainerState
 from ..node import Node
+from .exceptions import TrainingRefusedError
 from .io_helpers import LastTrainingIO
 from .rest import backdoor_controls
 from .trainer_logic_generic import TrainerLogicGeneric
@@ -64,7 +67,11 @@ class TrainerNode(Node):
         @sio_client.event
         async def begin_training(organization: str, project: str, details: Dict):
             self.log.info('received begin_training from server')
-            await self.trainer_logic.begin_training(organization, project, details)
+            try:
+                await self.trainer_logic.begin_training(organization, project, details)
+            except TrainingRefusedError as e:
+                self.log.warning('refused begin_training: %s', e)
+                return SocketResponse(success=False, error_msg=str(e), payload='refused')
             return True
 
         @sio_client.event
@@ -102,7 +109,7 @@ class TrainerNode(Node):
         if not self._idle_timeout:
             return
 
-        if self.trainer_logic.state == 'idle':
+        if self.trainer_logic.state in (TrainerState.Idle, TrainerState.Blocked):
             if self._first_idle_time is None:
                 self._first_idle_time = time.time()
             idle_time = time.time() - self._first_idle_time

learning_loop_node-0.25.0/learning_loop_node/trainer/training_slot.py

@@ -0,0 +1,207 @@
+"""The right to use the GPU, shared between the trainer nodes on one machine.
+
+A node holds the slot from the moment it accepts a training until the training is cleared.
+"""
+import asyncio
+import fcntl
+import logging
+import os
+import socket
+import weakref
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from .exceptions import TrainingSlotTakenError
+
+ENV_VAR = 'TRAINING_SLOT_LOCK'
+_CLAIM_ATTEMPTS = 3
+
+logger = logging.getLogger('learning_loop_node.training_slot')
+
+
+class TrainingSlot(ABC):
+    """The slot as a trainer uses it; an adapter implements `_acquire`, `_sibling_holder`, `release` and `held`."""
+
+    def claim(self, training: str) -> None:
+        """Take the slot for `training` without waiting; a no-op if this slot already holds it.
+
+        :raise TrainingSlotTakenError: if a sibling node holds the slot.
+        """
+        holder = self._describe(training)
+        for _ in range(_CLAIM_ATTEMPTS - 1):
+            try:
+                self._acquire(holder)
+                return
+            except TrainingSlotTakenError:
+                if self._sibling_holder() is not None:
+                    raise
+        self._acquire(holder)
+
+    async def wait_for(self, training: str) -> None:
+        """Take the slot for `training`, waiting as long as a sibling node holds it."""
+        try:
+            self.claim(training)
+            return
+        except TrainingSlotTakenError as e:
+            logger.info('%s - waiting', e)
+        while True:
+            await asyncio.sleep(1)
+            try:
+                self.claim(training)
+            except TrainingSlotTakenError:
+                continue
+            logger.info('acquired training slot')
+            return
+
+    @abstractmethod
+    def release(self) -> None:
+        """Give the slot back; a no-op if this slot does not hold it."""
+
+    @property
+    @abstractmethod
+    def held(self) -> bool:
+        """Whether this slot object holds the slot."""
+
+    @property
+    def taken_by_sibling(self) -> bool:
+        return not self.held and self._sibling_holder() is not None
+
+    @abstractmethod
+    def _acquire(self, holder: str) -> None:
+        """Take the slot without waiting; a no-op if this slot already holds it.
+
+        :raise TrainingSlotTakenError: if the slot is held elsewhere, including by a momentary probe.
+        """
+
+    @abstractmethod
+    def _sibling_holder(self) -> str | None:
+        """Who holds the slot, None if it is free or cannot be inspected; only asked while this slot does not hold it."""
+
+    @staticmethod
+    def _describe(training: str) -> str:
+        return f'pid {os.getpid()} on {socket.gethostname()}, {training}'
+
+
+class AlwaysFreeSlot(TrainingSlot):
+    """The slot of a node that has the GPU to itself."""
+
+    def __init__(self) -> None:
+        self._held = False
+
+    def release(self) -> None:
+        self._held = False
+
+    @property
+    def held(self) -> bool:
+        return self._held
+
+    def _acquire(self, holder: str) -> None:
+        self._held = True
+
+    def _sibling_holder(self) -> str | None:
+        return None
+
+
+class FileTrainingSlot(TrainingSlot):
+    """A slot backed by an exclusive `flock` on a file every sibling node mounts.
+
+    The kernel drops the lock with the process, so a node that dies in any way frees the slot.
+    The lock belongs to the inode: the file is only ever truncated and rewritten in place, never
+    replaced, and the mounted path is its directory.
+
+    The file's content names the holder and is diagnostic only — `_sibling_holder` tells a free slot
+    from a taken one by trying the lock, and that probe holds the lock for an instant, so a
+    claim that loses against a probe tries again. The file is made writable for everyone,
+    since sibling containers may run as different users.
+
+    A forked child shares the parent's lock and would keep it alive after the parent releases it,
+    so the child closes its copy right after the fork; the parent's lock is unaffected.
+    """
+
+    _instances: 'weakref.WeakSet[FileTrainingSlot]' = weakref.WeakSet()
+
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+        self._fd: int | None = None
+        FileTrainingSlot._instances.add(self)
+
+    def _acquire(self, holder: str) -> None:
+        if self._fd is not None:
+            return
+        fd = self._open()
+        try:
+            self._lock(fd)
+            os.ftruncate(fd, 0)
+            os.write(fd, holder.encode())
+        except BaseException:
+            os.close(fd)
+            raise
+        self._fd = fd
+
+    def release(self) -> None:
+        if self._fd is None:
+            return
+        try:
+            os.ftruncate(self._fd, 0)
+        except OSError:
+            logger.warning('cannot clear the holder from training slot %s', self.path, exc_info=True)
+        finally:
+            os.close(self._fd)
+            self._fd = None
+
+    @property
+    def held(self) -> bool:
+        return self._fd is not None
+
+    def _sibling_holder(self) -> str | None:
+        try:
+            fd = os.open(self.path, os.O_RDONLY)
+        except FileNotFoundError:
+            return None
+        except OSError:
+            logger.warning('cannot inspect training slot %s - treating it as free', self.path, exc_info=True)
+            return None
+        try:
+            self._lock(fd)
+        except TrainingSlotTakenError as e:
+            return e.holder
+        finally:
+            os.close(fd)
+        return None
+
+    @classmethod
+    def _drop_inherited_locks(cls) -> None:
+        for slot in list(cls._instances):
+            if slot._fd is not None:
+                os.close(slot._fd)
+                slot._fd = None
+
+    def _open(self) -> int:
+        fd = os.open(self.path, os.O_RDWR | os.O_CREAT, 0o666)
+        try:
+            os.fchmod(fd, 0o666)
+        except PermissionError:
+            pass
+        return fd
+
+    @staticmethod
+    def _lock(fd: int) -> None:
+        try:
+            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError as e:
+            raise TrainingSlotTakenError(os.read(fd, 4096).decode(errors='replace') or 'unknown') from e
+
+os.register_at_fork(after_in_child=FileTrainingSlot._drop_inherited_locks)  # pylint: disable=protected-access
+
+
+def training_slot_from_env() -> TrainingSlot:
+    """The slot `TRAINING_SLOT_LOCK` names, or a free one if the variable is unset.
+
+    :raise FileNotFoundError: if the lock file's directory does not exist.
+    """
+    path = os.environ.get(ENV_VAR)
+    if not path:
+        return AlwaysFreeSlot()
+    if not Path(path).parent.is_dir():
+        raise FileNotFoundError(f'{ENV_VAR}={path}: the directory does not exist; mount it from the host')
+    return FileTrainingSlot(path)

{learning_loop_node-0.24.0 → learning_loop_node-0.25.0}/learning_loop_node.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: learning_loop_node
-Version: 0.24.0
+Version: 0.25.0
 Summary: Python Library for Nodes which connect to the Zauberzeug Learning Loop
 Author-email: Zauberzeug GmbH <info@zauberzeug.com>
 License: MIT
@@ -67,6 +67,7 @@ You can configure connection to our Learning Loop by specifying the following en
 | RESTART_AFTER_TRAINING   | -            | Restart the trainer after training (set to 1)                | Trainer (opt.)            | 0            |
 | KEEP_OLD_TRAININGS       | -            | Do not delete old trainings (set to 1)                       | Trainer (opt.)            | 0            |
 | TRAINER_IDLE_TIMEOUT_SEC | -            | Automatically shutdown trainer after timeout (in seconds)    | Trainer (opt.)            | 0 (disabled) |
+| TRAINING_SLOT_LOCK       | -            | Lock file shared by the trainers on one GPU (see Trainer Node) | Trainer (opt.)          | - (disabled) |
 | USE_BACKDOOR_CONTROLS    | -            | Always enable backdoor controls (set to 1)                   | Trainer / Detector (opt.) | 0            |
 
 Note that organization and project IDs are always lower case and may differ from the names in the Learning Loop which can have uppercase letters.
@@ -250,6 +251,20 @@ Trainers fetch the images and anntoations from the Learning Loop to train new mo
 
 - if the command line tool "jpeginfo" is installed, the downloader will drop corrupted images automatically
 
+### Several trainers on one GPU
+
+Trainers that share a GPU (say, one per model architecture on the same machine) coordinate through
+a lock file: mount one host directory into each of them and point `TRAINING_SLOT_LOCK` at a file
+inside it, e.g. `/slot/training.lock`. A trainer takes the lock when it accepts a training and holds
+it until the training is cleaned up, including detection and upload. While a sibling holds it, the
+others report the state `blocked` instead of `idle`, so the Learning Loop neither offers them for a
+training nor dispatches queued jobs to them; a trainer that is sent `begin_training` regardless
+refuses it with `{"success": false, "error_msg": ...}` in the acknowledgement and sets nothing up. A
+training resumed after a restart cannot be refused — it waits for the lock instead, reporting
+`waiting_for_slot` meanwhile. The lock is an `flock`, so the kernel releases it whenever the holding
+process ends, however it ends — there is nothing to clean up after a crash. Mount the directory,
+not the file, and never replace the file: the lock belongs to its inode.
+
 ## Converter Node
 
 A Conveter Node converts models from one format into another.

{learning_loop_node-0.24.0 → learning_loop_node-0.25.0}/learning_loop_node.egg-info/SOURCES.txt

@@ -68,5 +68,6 @@ learning_loop_node/trainer/test_executor.py
 learning_loop_node/trainer/trainer_logic.py
 learning_loop_node/trainer/trainer_logic_generic.py
 learning_loop_node/trainer/trainer_node.py
+learning_loop_node/trainer/training_slot.py
 learning_loop_node/trainer/rest/__init__.py
 learning_loop_node/trainer/rest/backdoor_controls.py

{learning_loop_node-0.24.0 → learning_loop_node-0.25.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "learning_loop_node"
-version = "0.24.0"
+version = "0.25.0"
 description = "Python Library for Nodes which connect to the Zauberzeug Learning Loop"
 authors = [{ name = "Zauberzeug GmbH", email = "info@zauberzeug.com" }]
 license = { text = "MIT" }