learning-loop-node 0.24.0__py3-none-any.whl → 0.25.0__py3-none-any.whl

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/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/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/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/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/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.dist-info → learning_loop_node-0.25.0.dist-info}/METADATA

@@ -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.dist-info → learning_loop_node-0.25.0.dist-info}/RECORD

@@ -39,31 +39,32 @@ learning_loop_node/enums/__init__.py,sha256=tjSrhztIQ8W656_QuXfTbbVNtH_wDXP5hpYZ
 learning_loop_node/enums/annotator.py,sha256=mtTAw-8LJIrHcYkBjYHCZuhYEEHS6QzSK8k6BhLusvQ,285
 learning_loop_node/enums/detector.py,sha256=Qvm5LWWR9BfsDxHEQ8YzaPaUuSmp4BescYuV4X4ikwE,512
 learning_loop_node/enums/general.py,sha256=OSpcIpdFjuRByON2xBoiIXl1SgDREIqkwnaxkTFjnHU,164
-learning_loop_node/enums/trainer.py,sha256=VaD63guLO4aKgVfXT0EryPlXKQGegSET3Cp4R0uC_Ew,746
+learning_loop_node/enums/trainer.py,sha256=VdamD-k6Q8kMm_cwPA8iaEUfZvmTxLfsBHQpEi43Ahk,849
 learning_loop_node/helpers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 learning_loop_node/helpers/background_tasks.py,sha256=sbbS8JO1ls0ZCbNW68KOT9QLOFyFoJRXwwAq78iOmoA,3256
 learning_loop_node/helpers/entrypoint.py,sha256=U4vHFqgke6mXjDL-E4MND8RO7P1k2VBGwdDOOoJvD9Y,3915
 learning_loop_node/helpers/environment_reader.py,sha256=zsoFysxNxBCOJzUwX4Pbm3ftYhxx5zsbOHQfD9VcfvA,1937
 learning_loop_node/helpers/gdrive_downloader.py,sha256=zeYJciTAJVRpu_eFjwgYLCpIa6hU1d71anqEBb564Rk,1145
 learning_loop_node/helpers/log_conf.py,sha256=hqVAa_9NnYEU6N0dcOKmph82p7MpgKqeF_eomTLYzWY,961
-learning_loop_node/helpers/misc.py,sha256=qdN18Ukbw6eR6sFL4BxNq0EZzg2tVm2Cm8cTbFyy8QU,8931
+learning_loop_node/helpers/misc.py,sha256=zfy1SumqNBwBw6-BYignQJpwMeMf1k_MA3xhYt_GiG4,8939
 learning_loop_node/helpers/run.py,sha256=_uox-j3_K_bL3yCAwy3JYSOiIxrnhzVxyxWpCe8_J9U,876
 learning_loop_node/trainer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 learning_loop_node/trainer/batch_size.py,sha256=LjA-zyY4hdnSMy6ZsdQOhWbmd5H_DI-9LG3yFInv04E,5476
 learning_loop_node/trainer/cuda.py,sha256=ehpfWqPBCVFu06SnGsjmWbSDRJ8sZc-L_nCtemRLUOU,12567
 learning_loop_node/trainer/downloader.py,sha256=AIyEKM4XYtd6VgmXrP3VayV9DpJzdURK1Brx81ePNSM,1470
-learning_loop_node/trainer/exceptions.py,sha256=IRT7qvn7-RjXaqL8mF21ySm7JWhQtqm1Cl0Gi3DChVY,540
+learning_loop_node/trainer/exceptions.py,sha256=hhd23WvuqV-6fiJo3vCNLN0hmF6TxAH9C5Bc7pCtmUM,940
 learning_loop_node/trainer/executor.py,sha256=XQ1TwgOF2pQv7bgzSaljDWG4I1ySjYXL8r9c_p9MX1Q,3947
 learning_loop_node/trainer/io_helpers.py,sha256=zWl0-HLP8zxOG3ybm76vXgEozULj_1QRl95jkrF126k,7255
 learning_loop_node/trainer/metrics.py,sha256=H1jQh5fHN3-cTdMTuH8mF_uODV0T9DBdqcPUD-ZqFAI,1439
 learning_loop_node/trainer/subprocess.py,sha256=WxuTly0zhpClEjTd9wqoI_-7YTNMaiZpRHxu00vttKw,2647
 learning_loop_node/trainer/test_executor.py,sha256=6BVGDN_6f5GEMMEvDLSG1yzMybSvgXaP5uYpSfsVPP0,2224
 learning_loop_node/trainer/trainer_logic.py,sha256=Yaw1GNTC05rVLQczY6vKUgkymBYfm3_0XaR5C0q-CP4,9159
-learning_loop_node/trainer/trainer_logic_generic.py,sha256=IRP68ZGpnOX4OS03OwB7HmTUvk6nbBRmCjMCVupGE_Q,27196
-learning_loop_node/trainer/trainer_node.py,sha256=ZnydfHJ2ehAhb821Xke_09qthCE3vDf7iqeYUrliSSo,4938
+learning_loop_node/trainer/trainer_logic_generic.py,sha256=OtjdZlixz-2UXZkG_o5PaCdXCFvwgGuD5J4I9n5bROY,28812
+learning_loop_node/trainer/trainer_node.py,sha256=GU4LyFA-JC9s4d-PbyKhKWyGMMIzTMqKJFNxM6tDJE4,5316
+learning_loop_node/trainer/training_slot.py,sha256=Njq6Z0T2Lqlr24ftoIp11Pw2OxCo9Z2RGn1tuFhHqLY,6737
 learning_loop_node/trainer/rest/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 learning_loop_node/trainer/rest/backdoor_controls.py,sha256=ZnK8ypY5r_q0-YZbtaOxhQThzuZvMsQHM5gJGESd_dE,5131
-learning_loop_node-0.24.0.dist-info/METADATA,sha256=rQNtas26mLdkkjI37mJ_jteRe7RtUO_3DpReanEUC2U,18609
-learning_loop_node-0.24.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
-learning_loop_node-0.24.0.dist-info/top_level.txt,sha256=nS5qiFc4stVF_VXVCVJxZnrqVOv3IGX5-_E3sLby6OE,19
-learning_loop_node-0.24.0.dist-info/RECORD,,
+learning_loop_node-0.25.0.dist-info/METADATA,sha256=UOGaOXyynobyOFSEROGxOPH49odFh2-V2IGZ_LhQNsI,19855
+learning_loop_node-0.25.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+learning_loop_node-0.25.0.dist-info/top_level.txt,sha256=nS5qiFc4stVF_VXVCVJxZnrqVOv3IGX5-_E3sLby6OE,19
+learning_loop_node-0.25.0.dist-info/RECORD,,