datamint 2.3.5__py3-none-any.whl → 2.4.0__py3-none-any.whl

datamint/api/base_api.py

@@ -61,22 +61,56 @@ class BaseApi:
             client: Optional HTTP client instance. If None, a new one will be created.
         """
         self.config = config
-        self.client = client or self._create_client()
+        self._owns_client = client is None  # Track if we created the client
+        self.client = client or BaseApi._create_client(config)
         self.semaphore = asyncio.Semaphore(20)
         self._api_instance: 'Api | None' = None  # Injected by Api class
 
-    def _create_client(self) -> httpx.Client:
-        """Create and configure HTTP client with authentication and timeouts."""
-        headers = None
-        if self.config.api_key:
-            headers = {"apikey": self.config.api_key}
+    @staticmethod
+    def _create_client(config: ApiConfig) -> httpx.Client:
+        """Create and configure HTTP client with authentication and timeouts.
+        
+        The client is designed to be long-lived and reused across multiple requests.
+        It maintains connection pooling for improved performance.
+        Default limits: max_keepalive_connections=20, max_connections=100
+        """
+        headers = {"apikey": config.api_key} if config.api_key else None
 
         return httpx.Client(
-            base_url=self.config.server_url,
+            base_url=config.server_url,
             headers=headers,
-            timeout=self.config.timeout
+            timeout=config.timeout,
+            limits=httpx.Limits(
+                max_keepalive_connections=5,  # Increased from default 20
+                max_connections=20,  # Increased from default 100
+                keepalive_expiry=8
+            )
         )
 
+    def close(self) -> None:
+        """Close the HTTP client and release resources.
+        
+        Should be called when the API instance is no longer needed.
+        Only closes the client if it was created by this instance.
+        """
+        if self._owns_client and self.client is not None:
+            self.client.close()
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - ensures client is closed."""
+        self.close()
+
+    def __del__(self):
+        """Destructor - ensures client is closed when instance is garbage collected."""
+        try:
+            self.close()
+        except Exception:
+            pass  # Ignore errors during cleanup
+
     def _stream_request(self, method: str, endpoint: str, **kwargs):
         """Make streaming HTTP request with error handling.
 

datamint/api/client.py

@@ -68,6 +68,8 @@ class Api:
                                     f" Please check your api_key and/or other configurations. {e}")
 
     def _get_endpoint(self, name: str):
+        if self._client is None:
+            self._client = BaseApi._create_client(self.config)
         if name not in self._endpoints:
             api_class = self._API_MAP[name]
             endpoint = api_class(self.config, self._client)

datamint/apihandler/base_api_handler.py

@@ -30,7 +30,6 @@ ResourceFields: TypeAlias = Literal['modality', 'created_by', 'published_by', 'p
 
 _PAGE_LIMIT = 5000
 
-
 @deprecated(reason="Please use `from datamint import Api` instead.", version="2.0.0")
 class BaseAPIHandler:
     """

datamint/apihandler/dto/annotation_dto.py

@@ -178,6 +178,8 @@ class CreateAnnotationDto:
         if model_id is not None:
             if is_model == False:
                 raise ValueError("model_id==False while self.model_id is provided.")
+            if not isinstance(model_id, str):
+                raise ValueError("model_id must be a string if provided.")
             is_model = True
         self.is_model = is_model
         self.geometry = geometry

datamint/dataset/base_dataset.py

@@ -307,6 +307,10 @@ class DatamintBaseDataset:
         self.image_lsets, self.image_lcodes = self._get_labels_set(framed=False)
         worklist_id = self.get_info()['worklist_id']
         groups: dict[str, dict] = self.api.annotationsets.get_segmentation_group(worklist_id)['groups']
+        if not groups:
+            self.seglabel_list = []
+            self.seglabel2code = {}
+            return
         # order by 'index' key
         max_index = max([g['index'] for g in groups.values()])
         self.seglabel_list : list[str] = ['UNKNOWN'] * max_index  # 1-based

datamint/lightning/__init__.py

@@ -0,0 +1 @@
+from .datamintdatamodule import DatamintDataModule

datamint/lightning/datamintdatamodule.py

@@ -0,0 +1,103 @@
+from torch.utils.data import DataLoader
+from datamint import Dataset
+import lightning as L
+from typing import Any
+from copy import copy
+import numpy as np
+
+
+class DatamintDataModule(L.LightningDataModule):
+    """
+    LightningDataModule for Datamint datasets with train/val split.
+    TODO: Add support for test and predict dataloaders.
+    """
+
+    def __init__(
+        self,
+        project_name: str = "./",
+        batch_size: int = 32,
+        image_transform=None,
+        mask_transform=None,
+        alb_transform=None,
+        alb_train_transform=None,
+        alb_val_transform=None,
+        train_split: float = 0.9,
+        val_split: float = 0.1,
+        seed: int = 42,
+        num_workers: int = 4,
+        **dataset_kwargs: Any,
+    ):
+        super().__init__()
+        self.project_name = project_name
+        self.batch_size = batch_size
+        self.image_transform = image_transform
+        self.mask_transform = mask_transform
+
+        if alb_transform is not None and (alb_train_transform is not None or alb_val_transform is not None):
+            raise ValueError("You cannot specify both `alb_transform` and `alb_train_transform`/`alb_val_transform`.")
+
+        # Handle backward compatibility for alb_transform
+        if alb_transform is not None:
+            self.alb_train_transform = alb_transform
+            self.alb_val_transform = alb_transform
+        else:
+            self.alb_train_transform = alb_train_transform
+            self.alb_val_transform = alb_val_transform
+
+        self.train_split = train_split
+        self.val_split = val_split
+        self.seed = seed
+        self.dataset_kwargs = dataset_kwargs
+        self.num_workers = num_workers
+
+        self.dataset = None
+
+    def prepare_data(self) -> None:
+        """Download or update data if needed."""
+        Dataset(
+            project_name=self.project_name,
+            auto_update=True,
+        )
+
+    def setup(self, stage: str = None) -> None:
+        """Set up datasets and perform train/val split."""
+        if self.dataset is None:
+            # Create base dataset for getting indices
+            self.dataset = Dataset(
+                return_as_semantic_segmentation=True,
+                semantic_seg_merge_strategy="union",
+                return_frame_by_frame=True,
+                include_unannotated=False,
+                project_name=self.project_name,
+                image_transform=self.image_transform,
+                mask_transform=self.mask_transform,
+                alb_transform=None,  # No transform for base dataset
+                auto_update=False,
+                **self.dataset_kwargs,
+            )
+
+            indices = list(copy(self.dataset.subset_indices))
+            rs = np.random.RandomState(self.seed)
+            rs.shuffle(indices)
+            train_end = int(self.train_split * len(indices))
+            train_idx = indices[:train_end]
+            val_idx = indices[train_end:]
+
+            self.train_dataset = copy(self.dataset).subset(train_idx)
+            self.train_dataset.alb_transform = self.alb_train_transform
+            self.val_dataset = copy(self.dataset).subset(val_idx)
+            self.val_dataset.alb_transform = self.alb_val_transform
+
+    def train_dataloader(self) -> DataLoader:
+        return self.train_dataset.get_dataloader(batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True)
+
+    def val_dataloader(self) -> DataLoader:
+        return self.val_dataset.get_dataloader(batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
+
+    def test_dataloader(self):
+        # Use the same dataloader as validation for testing, because we have so few samples
+        return self.val_dataset.get_dataloader(batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
+
+    def predict_dataloader(self):
+        # Use the same dataloader as validation for testing, because we have so few samples
+        return self.val_dataset.get_dataloader(batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)

datamint/mlflow/__init__.py

@@ -0,0 +1,46 @@
+# Monkey patch mlflow.tracking._tracking_service.utils.get_tracking_uri
+from .tracking.fluent import set_project
+import mlflow.tracking._tracking_service.utils as mlflow_utils
+from functools import wraps
+import logging
+from .env_utils import setup_mlflow_environment, ensure_mlflow_configured
+
+_LOGGER = logging.getLogger(__name__)
+
+# Store reference to original function
+_original_get_tracking_uri = mlflow_utils.get_tracking_uri
+_SETUP_CALLED_SUCCESSFULLY = False
+
+
+@wraps(_original_get_tracking_uri)
+def _patched_get_tracking_uri(*args, **kwargs):
+    """Patched version of get_tracking_uri that ensures MLflow environment is set up first.
+
+    This wrapper ensures that setup_mlflow_environment is called before any tracking
+    URI operations, guaranteeing proper MLflow configuration.
+
+    Args:
+        *args: Arguments passed to the original get_tracking_uri function.
+        **kwargs: Keyword arguments passed to the original get_tracking_uri function.
+
+    Returns:
+        The result of the original get_tracking_uri function.
+    """
+    global _SETUP_CALLED_SUCCESSFULLY
+    if _SETUP_CALLED_SUCCESSFULLY:
+        return _original_get_tracking_uri(*args, **kwargs)
+    try:
+        _SETUP_CALLED_SUCCESSFULLY = setup_mlflow_environment(set_mlflow=True)
+    except Exception as e:
+        _SETUP_CALLED_SUCCESSFULLY = False
+        _LOGGER.error("Failed to set up MLflow environment: %s", e)
+    ret = _original_get_tracking_uri(*args, **kwargs)
+    return ret
+
+
+setup_mlflow_environment(set_mlflow=False)
+# Replace the original function with our patched version
+mlflow_utils.get_tracking_uri = _patched_get_tracking_uri
+
+
+__all__ = ['set_project', 'setup_mlflow_environment', 'ensure_mlflow_configured']

datamint/mlflow/artifact/__init__.py

@@ -0,0 +1 @@
+from .datamint_artifacts_repo import DatamintArtifactsRepository

datamint/mlflow/artifact/datamint_artifacts_repo.py

@@ -0,0 +1,8 @@
+from mlflow.store.artifact.mlflow_artifacts_repo import MlflowArtifactsRepository
+
+
+class DatamintArtifactsRepository(MlflowArtifactsRepository):
+    @classmethod
+    def resolve_uri(cls, artifact_uri, tracking_uri):
+        tracking_uri = tracking_uri.split('datamint://', maxsplit=1)[-1]
+        return super().resolve_uri(artifact_uri, tracking_uri)

datamint/mlflow/env_utils.py

@@ -0,0 +1,109 @@
+"""
+Utility functions for automatically configuring MLflow environment variables
+based on Datamint configuration.
+"""
+
+import os
+import logging
+from typing import Optional
+from urllib.parse import urlparse
+from datamint import configs
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def get_datamint_api_url() -> Optional[str]:
+    """Get the Datamint API URL from configuration or environment variables."""
+    # First check environment variable
+    api_url = os.getenv('DATAMINT_API_URL')
+    if api_url:
+        return api_url
+
+    # Then check configuration
+    api_url = configs.get_value(configs.APIURL_KEY)
+    if api_url:
+        return api_url
+
+    return None
+
+
+def get_datamint_api_key() -> Optional[str]:
+    """Get the Datamint API key from configuration or environment variables."""
+    # First check environment variable
+    api_key = os.getenv('DATAMINT_API_KEY')
+    if api_key:
+        return api_key
+
+    # Then check configuration
+    api_key = configs.get_value(configs.APIKEY_KEY)
+    if api_key:
+        return api_key
+
+    return None
+
+
+def _get_mlflowdatamint_uri() -> Optional[str]:
+    api_url = get_datamint_api_url()
+    if not api_url:
+        return None
+    _LOGGER.debug(f"Retrieved Datamint API URL: {api_url}")
+
+    # Remove trailing slash if present
+    api_url = api_url.rstrip('/')
+    # api_url samples:
+    # https://api.datamint.io
+    # http://localhost:3001
+
+    parsed_url = urlparse(api_url)
+    base_url = f"{parsed_url.scheme}://{parsed_url.hostname}"
+    _LOGGER.debug(f"Derived base URL for MLflow Datamint: {base_url}")
+    # FIXME: It should work with https or datamint-api server should forward https requests.
+    base_url = base_url.replace('https://', 'http://')
+    if len(base_url.replace('http:', '')) == 0:
+        return None
+
+    mlflow_uri = f"{base_url}:5000"
+    return mlflow_uri
+
+
+def setup_mlflow_environment(overwrite: bool = False,
+                             set_mlflow: bool = True) -> bool:
+    """
+    Automatically set up MLflow environment variables based on Datamint configuration.
+
+    Returns:
+        bool: True if MLflow environment was successfully configured, False otherwise.
+    """
+    _LOGGER.debug("Setting up MLflow environment variables from Datamint configuration")
+    api_key = get_datamint_api_key()
+    mlflow_uri = _get_mlflowdatamint_uri()
+    if not mlflow_uri or not api_key:
+        _LOGGER.warning("Datamint configuration incomplete, cannot auto-configure MLflow")
+        return False
+
+    if overwrite or not os.getenv('MLFLOW_TRACKING_TOKEN'):
+        os.environ['MLFLOW_TRACKING_TOKEN'] = api_key
+    if overwrite or not os.getenv('MLFLOW_TRACKING_URI'):
+        os.environ['MLFLOW_TRACKING_URI'] = mlflow_uri
+
+    if set_mlflow:
+        import mlflow
+        mlflow.set_tracking_uri(mlflow_uri)
+
+    return True
+
+
+def ensure_mlflow_configured() -> None:
+    """
+    Ensure MLflow environment is properly configured.
+    Raises an exception if configuration is incomplete.
+    """
+    if not setup_mlflow_environment():
+        if not os.getenv('MLFLOW_TRACKING_URI') or not os.getenv('MLFLOW_TRACKING_TOKEN'):
+            raise ValueError(
+                "MLflow environment not configured. Please either:\n"
+                "1. Run 'datamint-config' to set up Datamint configuration, or\n"
+                "2. Set DATAMINT_API_URL and DATAMINT_API_KEY environment variables, or\n"
+                "3. Manually set MLFLOW_TRACKING_URI and MLFLOW_TRACKING_TOKEN environment variables"
+            )

datamint/mlflow/env_vars.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+class EnvVars(Enum):
+    DATAMINT_PROJECT_ID = "DATAMINT_PROJECT_ID"
+    DATAMINT_PROJECT_NAME = "DATAMINT_PROJECT_NAME"

datamint/mlflow/lightning/callbacks/__init__.py

@@ -0,0 +1 @@
+from .modelcheckpoint import MLFlowModelCheckpoint