Merge pull request #213 from AllenNeuralDynamics:feat-resource-monitor-from-rig

Add constraint to check available storage from rig schema

Author: bruno-f-cruz

src/clabe/resource_monitor/__init__.py

@@ -1,9 +1,15 @@
 from ._base import ResourceMonitor
-from ._constraints import Constraint, available_storage_constraint_factory, remote_dir_exists_constraint_factory
+from ._constraints import (
+    Constraint,
+    available_storage_constraint_factory,
+    available_storage_constraint_factory_from_rig,
+    remote_dir_exists_constraint_factory,
+)
 
 __all__ = [
     "ResourceMonitor",
     "Constraint",
     "available_storage_constraint_factory",
     "remote_dir_exists_constraint_factory",
+    "available_storage_constraint_factory_from_rig",
 ]

src/clabe/resource_monitor/_constraints.py

@@ -1,10 +1,39 @@
 import os
 import shutil
 from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from aind_behavior_services import AindBehaviorRigModel
 
 from ._base import Constraint
 
 
+def available_storage_constraint_factory_from_rig(rig: "AindBehaviorRigModel", min_bytes: float = 2e11) -> Constraint:
+    """
+    Creates a constraint to check if the rig's data directory has sufficient available storage.
+
+    This factory function creates a constraint that validates whether the data directory
+    specified in the rig model has enough free space to meet the minimum requirements.
+
+    Args:
+        rig: The rig model containing the data directory path
+        min_bytes: Minimum required free space in bytes. Defaults to 200GB
+    Returns:
+        Constraint: A constraint object for available storage validation
+    Example:
+        ```python
+        from aind_behavior_services import AindBehaviorRigModel
+        from clabe.resource_monitor import available_storage_constraint_factory_from_rig
+        rig = AindBehaviorRigModel(rig_name="example_rig", version="1.0", data_directory="D:/data")
+        storage_constraint = available_storage_constraint_factory_from_rig(rig, min_bytes=5e11)  # 500GB
+        monitor = ResourceMonitor()
+        monitor.add_constraint(storage_constraint)
+        ```
+    """
+    return available_storage_constraint_factory(drive=rig.data_directory, min_bytes=min_bytes)
+
+
 def available_storage_constraint_factory(drive: os.PathLike = Path(r"C:\\"), min_bytes: float = 2e11) -> Constraint:
     """
     Creates a constraint to check if a drive has sufficient available storage.