add: druidfetcher support for different configId (#307)

Update config to support multiple DruidFetcher

---------

Signed-off-by: s0nicboOm <[email protected]>

numalogic/connectors/_config.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass, field
 from enum import IntEnum
-
-from omegaconf import MISSING
+from typing import Optional
 
 
 class ConnectorType(IntEnum):
@@ -60,8 +59,12 @@ class DruidConf(ConnectorConf):
  endpoint: Druid endpoint
  delay_hrs: Delay in hours for fetching data from Druid
  fetcher: DruidFetcherConf
+ id_fetcher: dict of DruidFetcherConf for fetching ids
+
+ Note: Either one of the fetcher or id_fetcher should be provided.
  """
 
  endpoint: str
  delay_hrs: float = 3.0
- fetcher: DruidFetcherConf = MISSING
+ fetcher: Optional[DruidFetcherConf] = None
+ id_fetcher: Optional[dict[str, DruidFetcherConf]] = None

numalogic/tools/types.py

@@ -17,11 +17,11 @@
 from numalogic.base import TorchModel, BaseThresholdModel, BaseTransformer
 
 try:
- from redis.client import AbstractRedis
+ from redis.client import Redis
 except ImportError:
  redis_client_t = TypeVar("redis_client_t")
 else:
- redis_client_t = TypeVar("redis_client_t", bound=AbstractRedis, covariant=True)
+ redis_client_t = TypeVar("redis_client_t", bound=Redis, covariant=True)
 
 artifact_t = TypeVar(
  "artifact_t",

numalogic/udfs/trainer.py

@@ -15,6 +15,7 @@
 from numalogic.config import PreprocessFactory, ModelFactory, ThresholdFactory, RegistryFactory
 from numalogic.config._config import NumalogicConf
 from numalogic.config.factory import ConnectorFactory
+from numalogic.connectors import DruidFetcherConf
 from numalogic.models.autoencoder import AutoencoderTrainer
 from numalogic.tools.data import StreamingDataset
 from numalogic.tools.exceptions import ConfigNotFoundError, RedisRegistryError
@@ -101,6 +102,36 @@ def get_conf(self, config_id: str) -> StreamConf:
  except KeyError as err:
  raise ConfigNotFoundError(f"Config with ID {config_id} not found!") from err
 
+ def register_druid_fetcher_conf(self, config_id: str, conf: DruidFetcherConf) -> None:
+ """
+ Register DruidFetcherConf with the UDF.
+
+ Args:
+ config_id: Config ID
+ conf: DruidFetcherConf object
+ """
+ self.pl_conf.druid_conf.id_fetcher[config_id] = conf
+
+ def get_druid_fetcher_conf(self, config_id: str) -> DruidFetcherConf:
+ """
+ Get DruidFetcherConf with the given ID.
+
+ Args:
+ config_id: Config ID
+
+ Returns
+ -------
+ DruidFetcherConf object
+
+ Raises
+ ------
+ ConfigNotFoundError: If config with the given ID is not found
+ """
+ try:
+ return self.pl_conf.druid_conf.id_fetcher[config_id]
+ except KeyError as err:
+ raise ConfigNotFoundError(f"Config with ID {config_id} not found!") from err
+
  @classmethod
  def compute(
  cls,
@@ -316,18 +347,25 @@ def fetch_data(self, payload: TrainerPayload) -> pd.DataFrame:
  """
  _start_time = time.perf_counter()
  _conf = self.get_conf(payload.config_id)
+ _fetcher_conf = self.druid_conf.fetcher or (
+ self.get_druid_fetcher_conf(payload.config_id) if self.druid_conf.id_fetcher else None
+ )
+ if not _fetcher_conf:
+ raise ConfigNotFoundError(
+ f"Druid fetcher config not found for config_id: {payload.config_id}!"
+ )
 
  try:
  _df = self.data_fetcher.fetch(
- datasource=self.druid_conf.fetcher.datasource,
+ datasource=_fetcher_conf.datasource,
  filter_keys=_conf.composite_keys,
  filter_values=payload.composite_keys,
- dimensions=list(self.druid_conf.fetcher.dimensions),
+ dimensions=list(_fetcher_conf.dimensions),
  delay=self.druid_conf.delay_hrs,
- granularity=self.druid_conf.fetcher.granularity,
- aggregations=dict(self.druid_conf.fetcher.aggregations),
- group_by=list(self.druid_conf.fetcher.group_by),
- pivot=self.druid_conf.fetcher.pivot,
+ granularity=_fetcher_conf.granularity,
+ aggregations=dict(_fetcher_conf.aggregations),
+ group_by=list(_fetcher_conf.group_by),
+ pivot=_fetcher_conf.pivot,
  hours=_conf.numalogic_conf.trainer.train_hours,
  )
  except Exception:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "numalogic"
-version = "0.6.0a10"
+version = "0.6.0a11"
 description = "Collection of operational Machine Learning models and tools."
 authors = ["Numalogic Developers"]
 packages = [{ include = "numalogic" }]

tests/udfs/resources/_config2.yaml

@@ -31,9 +31,10 @@ redis_conf:
 druid_conf:
  url: "druid-endpoint"
  endpoint: "endpoint"
- fetcher:
- dimensions: [ "col1" ]
- datasource: "table-name"
- group_by: [ "timestamp", "col1" ]
- pivot:
- columns: [ "col2" ]
+ id_fetcher:
+ druid-config:
+ dimensions: [ "col1" ]
+ datasource: "table-name"
+ group_by: [ "timestamp", "col1" ]
+ pivot:
+ columns: [ "col2" ]