cmsdials 0.0.0__py3-none-any.whl

cmsdials/clients/h2d/models.py

@@ -0,0 +1,50 @@
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel, Field
+
+from ...utils.base_model import OBaseModel, PaginatedBaseModel
+
+
+class LumisectionHistogram2D(BaseModel):
+    dataset: str = Field(..., max_length=255)
+    me: str = Field(..., max_length=255)
+    dataset_id: int
+    file_id: int
+    run_number: int
+    ls_number: int
+    me_id: int
+    x_min: float
+    x_max: float
+    x_bin: float
+    y_min: float
+    y_max: float
+    y_bin: float
+    entries: int
+    data: list[list[float]]
+
+
+class PaginatedLumisectionHistogram2DList(PaginatedBaseModel):
+    next: Optional[AnyUrl]
+    previous: Optional[AnyUrl]
+    results: list[LumisectionHistogram2D]
+
+
+class LumisectionHistogram2DFilters(OBaseModel):
+    next_token: Optional[str] = None
+    page_size: Optional[int] = None
+    dataset_id: Optional[int] = None
+    file_id: Optional[int] = None
+    run_number: Optional[int] = None
+    run_number__lte: Optional[int] = None
+    run_number__gte: Optional[int] = None
+    ls_number: Optional[int] = None
+    ls_numbet__lte: Optional[int] = None
+    ls_number__gte: Optional[int] = None
+    me_id: Optional[int] = None
+    entries__gte: Optional[int] = None
+    dataset: Optional[str] = None
+    dataset__regex: Optional[str] = None
+    logical_file_name: Optional[str] = None
+    logical_file_name__regex: Optional[str] = None
+    me: Optional[str] = None
+    me__regex: Optional[str] = None

cmsdials/clients/lumisection/client.py

@@ -0,0 +1,13 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import Lumisection, LumisectionFilters, PaginatedLumisectionList
+
+
+class LumisectionClient(BaseAuthorizedAPIClient):
+    data_model = Lumisection
+    pagination_model = PaginatedLumisectionList
+    filter_class = LumisectionFilters
+    lookup_url = "lumisection/"
+
+    def get(self, dataset_id: int, run_number: int, ls_number: int, **kwargs):
+        edp = f"{dataset_id}/{run_number}/{ls_number}/"
+        return super().get(edp, **kwargs)

cmsdials/clients/lumisection/models.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel, Field
+
+from ...utils.base_model import OBaseModel, PaginatedBaseModel
+
+
+class Lumisection(BaseModel):
+    dataset_id: int
+    dataset: str = Field(..., max_length=255)
+    run_number: int
+    ls_number: int
+    th1_count: int
+    th2_count: int
+
+
+class PaginatedLumisectionList(PaginatedBaseModel):
+    next: Optional[AnyUrl]
+    previous: Optional[AnyUrl]
+    results: list[Lumisection]
+
+
+class LumisectionFilters(OBaseModel):
+    next_token: Optional[str] = None
+    page_size: Optional[int] = None
+    dataset_id: Optional[int] = None
+    run_number: Optional[int] = None
+    run_number__lte: Optional[int] = None
+    run_number__gte: Optional[int] = None
+    ls_number: Optional[int] = None
+    ls_number__lte: Optional[int] = None
+    ls_number__gte: Optional[int] = None
+    dataset: Optional[str] = None
+    dataset__regex: Optional[str] = None

cmsdials/clients/mes/client.py

@@ -0,0 +1,12 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import MEFilters, MonitoringElement
+
+
+class MonitoringElementClient(BaseAuthorizedAPIClient):
+    data_model = MonitoringElement
+    filter_class = MEFilters
+    lookup_url = "mes/"
+
+    def get(self, id: int, **kwargs):  # noqa: A002
+        edp = f"{id}/"
+        return super().get(edp, **kwargs)

cmsdials/clients/mes/models.py

@@ -0,0 +1,18 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from ...utils.base_model import OBaseModel
+
+
+class MonitoringElement(BaseModel):
+    me_id: int
+    me: str = Field(..., max_length=255)
+    count: int
+    dim: int
+
+
+class MEFilters(OBaseModel):
+    me: Optional[str] = None
+    me__regex: Optional[str] = None
+    dim: Optional[int] = None

cmsdials/clients/ml_bad_lumisection/client.py

@@ -0,0 +1,29 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import MLBadLumisection, MLBadLumisectionFilters, PaginatedMLBadLumisectionList
+
+
+class MLBadLumisectionClient(BaseAuthorizedAPIClient):
+    data_model = MLBadLumisection
+    pagination_model = PaginatedMLBadLumisectionList
+    filter_class = MLBadLumisectionFilters
+    lookup_url = "ml-bad-lumisection/"
+
+    def get(self, model_id: int, dataset_id: int, run_number: int, ls_number: int, me_id: int, **kwargs):
+        edp = f"{model_id}/{dataset_id}/{run_number}/{ls_number}/{me_id}/"
+        return super().get(edp, **kwargs)
+
+    def cert_json(self, model_id__in: list[int], dataset_id__in: list[int], run_number__in: list[int], **kwargs):
+        edp = "cert-json/"
+        midin = ",".join(str(v) for v in model_id__in)
+        didin = ",".join(str(v) for v in dataset_id__in)
+        ridin = ",".join(str(v) for v in run_number__in)
+        params = {"model_id__in": midin, "dataset_id__in": didin, "run_number__in": ridin}
+        return super().get(edp, params=params, return_raw_json=True, **kwargs)
+
+    def golden_json(self, model_id__in: list[int], dataset_id__in: list[int], run_number__in: list[int], **kwargs):
+        edp = "golden-json/"
+        midin = ",".join(str(v) for v in model_id__in)
+        didin = ",".join(str(v) for v in dataset_id__in)
+        ridin = ",".join(str(v) for v in run_number__in)
+        params = {"model_id__in": midin, "dataset_id__in": didin, "run_number__in": ridin}
+        return super().get(edp, params=params, return_raw_json=True, **kwargs)

cmsdials/clients/ml_bad_lumisection/models.py

@@ -0,0 +1,38 @@
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel
+
+from ...utils.base_model import OBaseModel, PaginatedBaseModel
+
+
+class MLBadLumisection(BaseModel):
+    model_id: int
+    dataset_id: int
+    file_id: int
+    run_number: int
+    ls_number: int
+    me_id: int
+    mse: float
+
+
+class PaginatedMLBadLumisectionList(PaginatedBaseModel):
+    next: Optional[AnyUrl]
+    previous: Optional[AnyUrl]
+    results: list[MLBadLumisection]
+
+
+class MLBadLumisectionFilters(OBaseModel):
+    next_token: Optional[str] = None
+    page_size: Optional[int] = None
+    model_id: Optional[int] = None
+    model_id__in: Optional[list[int]] = None
+    dataset_id: Optional[int] = None
+    dataset_id__in: Optional[list[int]] = None
+    dataset: Optional[str] = None
+    dataset__regex: Optional[str] = None
+    me_id: Optional[int] = None
+    me: Optional[str] = None
+    me__regex: Optional[str] = None
+    run_number: Optional[int] = None
+    run_number__in: Optional[list[int]] = None
+    ls_number: Optional[int] = None

cmsdials/clients/ml_models_index/client.py

@@ -0,0 +1,13 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import MLModelsIndex, MLModelsIndexFilters, PaginatedMLModelsIndexList
+
+
+class MLModelsIndexClient(BaseAuthorizedAPIClient):
+    data_model = MLModelsIndex
+    pagination_model = PaginatedMLModelsIndexList
+    filter_class = MLModelsIndexFilters
+    lookup_url = "ml-models-index/"
+
+    def get(self, model_id: int, **kwargs):
+        edp = f"{model_id}/"
+        return super().get(edp, **kwargs)

cmsdials/clients/ml_models_index/models.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel
+
+from ...utils.base_model import OBaseModel, PaginatedBaseModel
+
+
+class MLModelsIndex(BaseModel):
+    model_id: int
+    filename: str
+    target_me: str
+    active: bool
+
+
+class PaginatedMLModelsIndexList(PaginatedBaseModel):
+    next: Optional[AnyUrl]
+    previous: Optional[AnyUrl]
+    results: list[MLModelsIndex]
+
+
+class MLModelsIndexFilters(OBaseModel):
+    next_token: Optional[str] = None
+    page_size: Optional[int] = None
+    model_id: Optional[int] = None
+    model_id__in: Optional[list[int]] = None
+    filename: Optional[str] = None
+    filename__regex: Optional[str] = None
+    target_me: Optional[str] = None
+    target_me__regex: Optional[str] = None
+    active: Optional[bool] = None

cmsdials/clients/oms_proxy/client.py

@@ -0,0 +1,49 @@
+from typing import Optional
+
+import requests
+from requests.exceptions import HTTPError
+
+from ...auth._base import BaseCredentials
+from ...utils.api_client import BaseAPIClient
+from ...utils.logger import logger
+from .models import OMSFilter, OMSPage
+
+
+class OMSProxyClient(BaseAPIClient):
+    default_timeout = 30
+    lookup_url = "oms-proxy/"
+
+    def __init__(
+        self,
+        creds: BaseCredentials,
+        *args: str,
+        **kwargs: str,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.creds = creds
+
+    def _build_headers(self) -> dict:
+        base = {
+            "accept": "application/json",
+            "User-Agent": self._build_user_agent()
+        }
+        self.creds.before_request(base)
+        return base
+
+    def query(self, endpoint: str, filters: list[OMSFilter], pages: Optional[list[OMSPage]] = None):
+        headers = self._build_headers()
+        endpoint_url = self.api_url + self.lookup_url
+
+        # Format filters
+        filters = {f"filter[{_filter.attribute_name}][{_filter.operator}]": _filter.value for _filter in filters}
+        pages = {f"page[{page.attribute_name}]": page.value for page in pages} if pages and len(pages) > 0 else {}
+        params = {"endpoint": endpoint, **filters, **pages}
+
+        response = requests.get(endpoint_url, headers=headers, params=params, timeout=self.default_timeout)
+        try:
+            response.raise_for_status()
+        except HTTPError as err:
+            logger.info(f"Api raw response: {response.text}")
+            raise err
+
+        return response.json()

cmsdials/clients/oms_proxy/models.py

@@ -0,0 +1,14 @@
+from typing import Literal, Union
+
+from pydantic import BaseModel
+
+
+class OMSFilter(BaseModel):
+    attribute_name: str
+    value: Union[str, int, float]
+    operator: Literal["EQ", "NEQ", "LT", "GT", "LE", "GE", "LIKE"]
+
+
+class OMSPage(BaseModel):
+    attribute_name: str
+    value: int

cmsdials/clients/run/client.py

@@ -0,0 +1,13 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import PaginatedRunList, Run, RunFilters
+
+
+class RunClient(BaseAuthorizedAPIClient):
+    data_model = Run
+    pagination_model = PaginatedRunList
+    filter_class = RunFilters
+    lookup_url = "run/"
+
+    def get(self, dataset_id: int, run_number: int, **kwargs):
+        edp = f"{dataset_id}/{run_number}/"
+        return super().get(edp, **kwargs)

cmsdials/clients/run/models.py

@@ -0,0 +1,29 @@
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel, Field
+
+from ...utils.base_model import OBaseModel, PaginatedBaseModel
+
+
+class Run(BaseModel):
+    dataset_id: int
+    dataset: str = Field(..., max_length=255)
+    run_number: int
+    ls_count: int
+
+
+class PaginatedRunList(PaginatedBaseModel):
+    next: Optional[AnyUrl]
+    previous: Optional[AnyUrl]
+    results: list[Run]
+
+
+class RunFilters(OBaseModel):
+    next_token: Optional[str] = None
+    page_size: Optional[int] = None
+    dataset_id: Optional[int] = None
+    run_number: Optional[int] = None
+    run_number__lte: Optional[int] = None
+    run_number__gte: Optional[int] = None
+    dataset: Optional[str] = None
+    dataset__regex: Optional[str] = None

cmsdials/cmsdials.py

@@ -0,0 +1,27 @@
+from typing import Optional
+
+from .auth._base import BaseCredentials
+from .clients.dataset_index.client import DatasetIndexClient
+from .clients.file_index.client import FileIndexClient
+from .clients.h1d.client import LumisectionHistogram1DClient
+from .clients.h2d.client import LumisectionHistogram2DClient
+from .clients.lumisection.client import LumisectionClient
+from .clients.mes.client import MonitoringElementClient
+from .clients.ml_bad_lumisection.client import MLBadLumisectionClient
+from .clients.ml_models_index.client import MLModelsIndexClient
+from .clients.oms_proxy.client import OMSProxyClient
+from .clients.run.client import RunClient
+
+
+class Dials:
+    def __init__(self, creds: BaseCredentials, workspace: Optional[str] = None, *args, **kwargs) -> None:
+        self.dataset_index = DatasetIndexClient(creds, workspace, *args, **kwargs)
+        self.file_index = FileIndexClient(creds, workspace, *args, **kwargs)
+        self.h1d = LumisectionHistogram1DClient(creds, workspace, *args, **kwargs)
+        self.h2d = LumisectionHistogram2DClient(creds, workspace, *args, **kwargs)
+        self.lumi = LumisectionClient(creds, workspace, *args, **kwargs)
+        self.run = RunClient(creds, workspace, *args, **kwargs)
+        self.mes = MonitoringElementClient(creds, workspace, *args, **kwargs)
+        self.oms = OMSProxyClient(creds, *args, **kwargs)
+        self.ml_bad_lumis = MLBadLumisectionClient(creds, workspace, *args, **kwargs)
+        self.ml_models_index = MLModelsIndexClient(creds, workspace, *args, **kwargs)

cmsdials/filters.py

@@ -0,0 +1,25 @@
+from .clients.dataset_index.models import DatasetIndexFilters
+from .clients.file_index.models import FileIndexFilters
+from .clients.h1d.models import LumisectionHistogram1DFilters
+from .clients.h2d.models import LumisectionHistogram2DFilters
+from .clients.lumisection.models import LumisectionFilters
+from .clients.mes.models import MEFilters
+from .clients.ml_bad_lumisection.models import MLBadLumisectionFilters
+from .clients.ml_models_index.models import MLModelsIndexFilters
+from .clients.oms_proxy.models import OMSFilter, OMSPage
+from .clients.run.models import RunFilters
+
+
+__all__ = [
+    "DatasetIndexFilters",
+    "FileIndexFilters",
+    "LumisectionHistogram1DFilters",
+    "LumisectionHistogram2DFilters",
+    "LumisectionFilters",
+    "RunFilters",
+    "MEFilters",
+    "MLBadLumisectionFilters",
+    "MLModelsIndexFilters",
+    "OMSFilter",
+    "OMSPage",
+]

cmsdials/utils/_json.py

@@ -0,0 +1,24 @@
+import json
+from datetime import datetime
+
+
+class TokenEncoder(json.JSONEncoder):
+    def default(self, o):
+        if isinstance(o, datetime):
+            return o.isoformat()
+
+        return json.JSONEncoder.default(self, o)
+
+
+class TokenDecoder(json.JSONDecoder):
+    def __init__(self, *args, **kwargs):
+        json.JSONDecoder.__init__(self, *args, **kwargs, object_hook=self.object_hook)
+
+    def object_hook(self, obj):
+        ret = {}
+        for key, value in obj.items():
+            if key in ("expires_at", "refresh_expires_at"):
+                ret[key] = datetime.fromisoformat(value)
+            else:
+                ret[key] = value
+        return ret

cmsdials/utils/api_client.py

@@ -0,0 +1,226 @@
+import sys
+from importlib import util as importlib_util
+from traceback import format_exc
+from typing import Optional
+from warnings import warn
+
+import requests
+from pydantic import AnyUrl
+from requests import Response, Session
+from requests.adapters import DEFAULT_RETRIES, HTTPAdapter
+from requests.exceptions import HTTPError
+
+from .._version import __version__
+from ..auth._base import BaseCredentials
+from ..utils.logger import logger
+
+
+if importlib_util.find_spec("tqdm"):
+    from tqdm.auto import tqdm
+
+    TQDM_INSTALLED = True
+else:
+    TQDM_INSTALLED = False
+
+
+class BaseAPIClient:
+    PRODUCTION_BASE_URL = "https://cmsdials-api.web.cern.ch/"
+    PRODUCTION_API_ROUTE = "api/"
+    PRODUCTION_API_VERSION = "v1/"
+
+    def __init__(
+        self, base_url: Optional[str] = None, route: Optional[str] = None, version: Optional[str] = None
+    ) -> None:
+        self.base_url = self.__endswithslash(base_url or self.PRODUCTION_BASE_URL)
+        self.route = self.__endswithslash(route or self.PRODUCTION_API_ROUTE)
+        self.version = self.__endswithslash(version or self.PRODUCTION_API_VERSION)
+
+    @staticmethod
+    def __endswithslash(value: str) -> str:
+        if value.endswith("/") is False:
+            return value + "/"
+        return value
+
+    @classmethod
+    def _requests_get_retriable(cls, *args, retries=DEFAULT_RETRIES, **kwargs) -> Response:
+        """
+        requests.get() with an additional `retries` parameter.
+
+        Specify retries=<number of attempts - 1> for simple use cases.
+        For advanced usage, see https://docs.python-requests.org/en/latest/user/advanced/
+        """
+        with Session() as s:
+            s.mount(cls.PRODUCTION_BASE_URL, HTTPAdapter(max_retries=retries))
+            ret = s.get(*args, **kwargs)
+        return ret
+
+    @staticmethod
+    def _build_user_agent():
+        py_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
+        req_version = requests.__version__
+        return f"dials-py/{__version__} (python/{py_version}, python-requests/{req_version})"
+
+    @property
+    def api_url(self):
+        return self.base_url + self.route + self.version
+
+
+class BaseAuthorizedAPIClient(BaseAPIClient):
+    data_model = None
+    pagination_model = None
+    filter_class = None
+    lookup_url = None
+    default_timeout = 30  # seconds
+
+    def __init__(
+        self,
+        creds: BaseCredentials,
+        workspace: Optional[str] = None,
+        *args: str,
+        **kwargs: str,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.creds = creds
+        self.workspace = workspace
+
+    def _build_headers(self) -> dict:
+        base = {
+            "accept": "application/json",
+            "User-Agent": self._build_user_agent()
+        }
+        if self.workspace is not None:
+            base["Workspace"] = self.workspace
+        self.creds.before_request(base)
+        return base
+
+    def get(self, edp: str, retries=DEFAULT_RETRIES, params: Optional[dict] = None, return_raw_json: bool = False):
+        endpoint_url = self.api_url + self.lookup_url + edp
+        headers = self._build_headers()
+        response = self._requests_get_retriable(
+            endpoint_url,
+            headers=headers,
+            params=params,
+            timeout=self.default_timeout,
+            retries=retries,
+        )
+
+        try:
+            response.raise_for_status()
+        except HTTPError as err:
+            logger.info(f"Api raw response: {response.text}")
+            raise err
+
+        response = response.json()
+        if return_raw_json:
+            return response
+        return self.data_model(**response)
+
+    def list(self, filters=None, retries=DEFAULT_RETRIES):
+        filters = filters or self.filter_class()
+        endpoint_url = self.api_url + self.lookup_url
+        headers = self._build_headers()
+        response = self._requests_get_retriable(
+            endpoint_url,
+            headers=headers,
+            params=filters.cleandict(),
+            timeout=self.default_timeout,
+            retries=retries,
+        )
+
+        try:
+            response.raise_for_status()
+        except HTTPError as err:
+            logger.info(f"Api raw response: {response.text}")
+            raise err
+
+        response = response.json()
+        if self.pagination_model:
+            return self.pagination_model(**response)
+        elif isinstance(response, list):
+            return [self.data_model(**res) for res in response]
+
+        raise ValueError("pagination model is None and response is not a list.")
+
+    def __list_sync(
+        self,
+        filters,
+        max_pages: Optional[int] = None,
+        enable_progress: bool = False,
+        retries=DEFAULT_RETRIES,
+        keep_failed: bool = False,
+        resume_from=None,
+    ):
+        next_string: Optional[AnyUrl] = None
+        results = []
+        is_last_page = False
+
+        if resume_from is not None:
+            results = resume_from.results
+            next_string: AnyUrl = resume_from.next
+            if next_string is None and len(results):
+                warn(
+                    "resume_from.next is None while resume_from.result is not empty, doing nothing.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                is_last_page = True
+
+        total_pages = 0
+        use_tqdm = TQDM_INSTALLED and enable_progress
+
+        if use_tqdm:
+            progress = tqdm(desc="Progress", total=1)
+
+        while is_last_page is False:
+            next_token = dict(next_string.query_params()).get("next_token") if next_string else None
+            curr_filters = self.filter_class(**filters.model_dump())
+            curr_filters.next_token = next_token
+            try:
+                response = self.list(curr_filters, retries=retries)
+            except Exception as e:  # noqa: BLE001
+                if use_tqdm:
+                    progress.close()
+
+                if not keep_failed:
+                    raise e
+
+                exc_formatted = format_exc()
+                warn(
+                    "HTTP request failed, returning partial results. Exception: " + exc_formatted,
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                return self.pagination_model(
+                    next=next_string,
+                    previous=None,
+                    results=results,
+                    exc_type=e.__class__.__name__,
+                    exc_formatted=exc_formatted,
+                )
+            results.extend(response.results)
+            next_string: AnyUrl = response.next
+            is_last_page = next_string is None
+            total_pages += 1
+            max_pages_reached = max_pages and total_pages >= max_pages
+            if use_tqdm:
+                if is_last_page or max_pages_reached:
+                    progress.update()
+                else:
+                    progress.total = total_pages + 1
+                    progress.update(1)
+            if max_pages_reached:
+                break
+
+        if use_tqdm:
+            progress.close()
+
+        return self.pagination_model(
+            next=next_string,
+            previous=None,
+            results=results,  # No problem re-using last response count
+        )
+
+    def list_all(self, filters, max_pages: Optional[int] = None, **kwargs):
+        if self.pagination_model is None:
+            return self.list(filters, retries=kwargs.get("retries", DEFAULT_RETRIES))
+        return self.__list_sync(filters, max_pages, **kwargs)