toolchemy 0.2.185__py3-none-any.whl

toolchemy/ai/trackers/mlflow_tracker.py

@@ -0,0 +1,221 @@
+import time
+from typing import Dict, Optional, Any
+
+import mlflow
+mlflow.autolog(disable=True)
+from mlflow.client import MlflowClient
+from mlflow.entities import RunStatus, Metric, Param
+from mlflow.tracking.context.registry import resolve_tags
+from mlflow.utils.mlflow_tags import MLFLOW_PARENT_RUN_ID, MLFLOW_RUN_NAME
+from mlflow import MlflowException
+
+from toolchemy.ai.trackers.common import TrackerBase
+from toolchemy.utils.logger import get_logger
+
+
+class MLFlowTracker(TrackerBase):
+    def __init__(self, tracking_uri: str, experiment_name: str, with_artifact_logging=True,
+                 tracking_client: MlflowClient | None = None):
+        super().__init__(experiment_name, with_artifact_logging)
+        self._tracking_client = None
+        self._active_run = None
+        self._active_run_id = None
+        self._experiment_id = None
+        self._reset_run()
+
+        if tracking_client:
+            self._tracking_client = tracking_client
+        else:
+            self._tracking_client = self._build_tracking_client(tracking_uri)
+
+        mlflow.set_tracking_uri(tracking_uri)
+
+        logger = get_logger()
+        logger.info(f"Mlflow tracker created")
+        logger.info(f"> tracking uri: {tracking_uri}")
+        logger.info(f"> artifact logging: {self._artifact_logging}")
+
+    @property
+    def run_name(self) -> str:
+        if not self._active_run:
+            raise RuntimeError("There is no active run!")
+        return self._active_run.info.run_name
+
+    def start_run(
+            self, run_id: str = None,
+            run_name: str = None,
+            parent_run_id: str = None,
+            user_specified_tags: Dict[str, str] = None
+    ):
+        """
+        Starts the run
+
+        :param run_id: If specified, get the run with the specified ID and log parameters and metrics under that run.
+        :param run_name: Name of new run. Used only when run_id is unspecified.
+        :param parent_run_id: If specified: current run will be nested into parent_run_id
+        :param user_specified_tags: dict of with custom tags
+        """
+        if self._disabled:
+            return
+
+        if user_specified_tags is None:
+            user_specified_tags = {}
+        if parent_run_id is not None:
+            user_specified_tags[MLFLOW_PARENT_RUN_ID] = parent_run_id
+        if run_name:
+            user_specified_tags[MLFLOW_RUN_NAME] = run_name
+
+        tags = resolve_tags(user_specified_tags)
+
+        experiment = self._tracking_client.get_experiment_by_name(self._experiment_name)
+        if experiment:
+            experiment_comment_msg = "(already exists)"
+            self._logger.debug(f"Experiment '{self._experiment_name}' already exists")
+            self._experiment_id = experiment.experiment_id
+            if experiment.lifecycle_stage == "deleted":
+                self._logger.info(f"Restoring deleted experiment")
+                self._tracking_client.restore_experiment(self._experiment_id)
+        else:
+            experiment_comment_msg = "(does not exist, creating a new one)"
+            self._logger.debug(f"Experiment '{self._experiment_name}' does not exist, creating a new one")
+            self._experiment_id = self._tracking_client.create_experiment(self._experiment_name)
+
+        self._logger.info(f"Starting the experiment tracking")
+        self._logger.info(f"> experiment name: {self._experiment_name} {experiment_comment_msg}")
+        self._logger.info(f"> experiment id: {self._experiment_id}")
+        self._logger.info(f"> run name: {run_name}")
+
+        self._active_run = self._tracking_client.create_run(
+            experiment_id=self._experiment_id,
+            start_time=None,
+            run_name=run_name,
+            tags=tags)
+
+        self._active_run_id = self._active_run.info.run_id
+
+    def end_run(self):
+        if self._disabled:
+            return
+
+        status = RunStatus.to_string(RunStatus.FINISHED)
+        self._tracking_client.set_terminated(self._active_run_id, status)
+        self._reset_run()
+
+    def log(self, name: str, value: Any):
+        if self._disabled:
+            return
+
+        if isinstance(value, dict):
+            self._tracking_client.log_dict(self._active_run_id, value, f"{name}.json")
+
+        raise ValueError(f"Unsupported logged object type: {type(value)}")
+
+    def log_param(self, name: str, value: Any):
+        if self._disabled:
+            return
+
+        self._store_param(name, value)
+        self._tracking_client.log_param(self._active_run_id, name, value)
+
+    def log_params(self, params: Dict[str, Any]):
+        if self._disabled:
+            return
+
+        params_to_store = []
+        for key, value in params.items():
+            if isinstance(value, list):
+                for v in value:
+                    params_to_store.append(Param(key, str(v)))
+            else:
+                params_to_store.append(Param(key, str(value)))
+            self._store_param(key, value)
+
+        self._tracking_client.log_batch(self._active_run_id, [], params_to_store)
+
+    def log_text(self, name: str, value: str):
+        if self._disabled:
+            return
+        try:
+            self._tracking_client.log_text(run_id=self._active_run_id, text=value, artifact_file=name)
+        except MlflowException as e:
+            self._logger.error(f"An error occurred during text logging: {e}")
+            self._logger.error(f"> tracking uri: {self._tracking_client.tracking_uri}")
+            self._logger.error(f"> artifact uri: {self._active_run.info.artifact_uri}")
+            raise e
+
+    def log_metric(self, name: str, value: float, step: int | None = None, metric_metadata: dict | None = None):
+        if self._disabled:
+            return
+
+        metric_value = self._store_metric(name, value, metric_metadata)
+        self._tracking_client.log_metric(self._active_run_id, name, metric_value, step)
+
+    def log_metrics(self, metrics: Dict[str, float | list], step: Optional[int] = None):
+        if self._disabled:
+            return
+
+        metrics_to_store = []
+        timestamp = int(time.time() * 1000)
+        for k, value in metrics.items():
+            if isinstance(value, list):
+                for v in value:
+                    metric_value = self._store_metric(k, v)
+                    metrics_to_store.append(Metric(k, metric_value, timestamp, step or 0))
+            else:
+                metric_value = self._store_metric(k, value)
+                metrics_to_store.append(Metric(k, metric_value, timestamp, step or 0))
+
+        self._tracking_client.log_batch(self._active_run_id, metrics_to_store)
+
+    def log_artifact(self, artifact_path: str, save_dir: str = None):
+        if self._disabled:
+            return
+
+        if self._artifact_logging:
+            self._tracking_client.log_artifact(self._active_run_id, artifact_path, save_dir)
+
+    def log_figure(self, figure, save_path: str):
+        if self._disabled:
+            return
+        self._tracking_client.log_figure(self._active_run_id, figure, save_path)
+
+    def set_run_tag(self, name: str, value: str | int | float):
+        self._store_tag(name, value, run_name=self.run_name)
+        self._tracking_client.set_tag(self._active_run_id, name, value)
+
+    def set_experiment_tag(self, name: str, value: str | int | float):
+        self._store_tag(name, value)
+        self._tracking_client.set_experiment_tag(self._experiment_id, name, value)
+
+    @staticmethod
+    def _build_tracking_client(tracking_uri: str) -> MlflowClient:
+        tracking_client = mlflow.tracking.MlflowClient(tracking_uri=tracking_uri, registry_uri=tracking_uri)
+        return tracking_client
+
+    def _reset_run(self):
+        self._active_run = None
+        self._active_run_id = None
+
+
+def play():
+    from toolchemy.utils.datestimes import current_datetime_str
+    from toolchemy.ai.prompter import PrompterMLflow
+    from toolchemy.utils import Locations
+
+    locations = Locations()
+    tracker = MLFlowTracker("http://hal:5000", f"test-{current_datetime_str()}")
+    tracker.start_run()
+    prompter = PrompterMLflow(locations.in_resources("tests/prompts_mlflow"))
+    print(prompter.render("test_prompt", foo="foo1", bar="bar1"))
+    tracker.log_param("param1", "param1value")
+    tracker.log_param("param2", 2)
+    tracker.log_metric("metric1", 1.0)
+    tracker.log_metric("metric2", 2.0, metric_metadata={"info": "metric 2 metadata"})
+    tracker.log_text("text_test", "some longer piece of text")
+    print(prompter.render("test_prompt", foo="foo1", bar="bar1"))
+    tracker.end_run()
+    print(prompter.render("test_prompt", foo="foo1", bar="bar1"))
+
+
+if __name__ == "__main__":
+    play()

toolchemy/ai/trackers/neptune_tracker.py

@@ -0,0 +1,135 @@
+from typing import Dict, Optional, Any
+from neptune_scale import Run
+
+from toolchemy.ai.trackers.common import TrackerBase
+from toolchemy.utils.datestimes import current_datetime_str
+
+
+class NeptuneAITracker(TrackerBase):
+    def __init__(self, project_name: str, experiment_name: str, api_token: str, with_artifact_logging: bool = True,
+                 disabled: bool = False):
+        super().__init__(experiment_name=experiment_name, with_artifact_logging=with_artifact_logging, disabled=disabled)
+        self._api_token = api_token
+        self._project_name = project_name
+        self._active_run: Run | None = None
+        self._active_run_id: str | None = None
+
+    @property
+    def run_name(self) -> str:
+        if not self._active_run or not self._active_run_id:
+            raise RuntimeError("There is no active run!")
+        return self._active_run_id
+
+
+    def start_run(
+            self, run_id: str = None,
+            run_name: str = None,
+            parent_run_id: str = None,
+            user_specified_tags: Dict[str, str] = None
+    ):
+        if self._disabled:
+            return
+        if self._active_run or self._active_run_id:
+            raise RuntimeError(f"Cannot start a new run, there is already an active run")
+        if run_name is not None:
+            self._logger.warning(f"Neptune tracker uses 'run_id' as the run name. Use 'run_id' for the custom run name.")
+            if run_id is None:
+                run_id = run_name
+
+        if run_id is None:
+            run_id = self._generate_run_name()
+
+        self._active_run_id = run_id
+
+        self._active_run = Run(
+            project=self._project_name,
+            api_token=self._api_token,
+            experiment_name=self._experiment_name,
+            run_id=self._active_run_id,
+            enable_console_log_capture=True,
+        )
+        self._logger.info(f"Neptune tracking run started. Experiment name: {self.experiment_name}")
+
+    def end_run(self):
+        if self._disabled:
+            return
+        if self._active_run is None:
+            raise ValueError(f"No active run to stop")
+        self._active_run.close()
+        self._active_run = None
+        self._active_run_id = None
+
+    def log(self, name: str, value: Any):
+        if self._disabled:
+            return
+        raise NotImplemented()
+
+    def log_param(self, name: str, val, step: Optional[int] = None):
+        if self._disabled:
+            return
+        self._store_param(name, val)
+        self._active_run.log_configs({name: val})
+
+    def log_params(self, params: Dict[str, Any]):
+        if self._disabled:
+            return
+        for param_name, param_value in params.items():
+            self._store_param(param_name, param_value)
+        self._active_run.log_configs(params)
+
+    def log_text(self, name: str, value: str):
+        if self._disabled:
+            return
+        self._active_run.log_configs({name: value})
+
+    def log_metric(self, name: str, value: float, step: int | None = None, metric_metadata: dict | None = None):
+        if self._disabled:
+            return
+        self.log_metrics({name: value}, step)
+
+    def log_metrics(self, metrics: Dict[str, float | list], step: Optional[int] = None):
+        if self._disabled:
+            return
+
+        for k, v in metrics.items():
+            metric_value = self._store_metric(k, v)
+            metrics[k] = metric_value
+
+        self._active_run.log_metrics(metrics, step=step)
+
+    def log_artifact(self, artifact_path: str, save_dir: str = None):
+        if self._disabled:
+            return
+        if not self._artifact_logging:
+            return
+        if save_dir is None:
+            save_dir = ""
+
+        self._active_run.assign_files({save_dir: artifact_path})
+
+    def log_figure(self, figure, save_path: str):
+        if self._disabled:
+            return
+        if save_path is None:
+            save_path = ""
+
+        self._active_run.assign_files({save_path: figure})
+
+    def set_run_tag(self, name: str, value: str | int | float):
+        self._store_tag(name, value, run_name=self.run_name)
+        self._active_run.add_tags([f"{name}__{value}"])
+
+    def set_experiment_tag(self, name: str, value: str | int | float):
+        self._store_tag(name, value)
+        if self._active_run:
+            self.set_run_tag(name, value)
+
+    def get_id(self):
+        if self._disabled:
+            return ""
+        if not self._active_run:
+            raise ValueError(f"No active run")
+        return self._active_run
+
+    def _generate_run_name(self) -> str:
+        return f"{self.experiment_name}__{current_datetime_str('%Y_%m_%d__%H_%M_%S')}"

toolchemy/db/lightdb.py

@@ -0,0 +1,260 @@
+import logging
+import uuid
+from tqdm import tqdm
+from collections import defaultdict
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+from tinydb import TinyDB, Query
+
+from toolchemy.utils.logger import get_logger
+from toolchemy.utils.utils import hash_dict, pp
+from toolchemy.utils.cacher import Cacher, ICacher
+
+
+class NotFoundError(Exception):
+    pass
+
+
+class FilterOp(Enum):
+    GREATER = 1
+    GREATER_OR_EQUAL = 2
+    LESS = 3
+    LESS_OR_EQUAL = 4
+    EQUAL = 5
+
+
+@dataclass
+class Filter:
+    key: str
+    value: Any
+    op: FilterOp = FilterOp.EQUAL
+
+    def __str__(self) -> str:
+        return f"Filter(key={self.key}, op={self.op.name}, value={self.value})"
+
+
+class ILightDB(ABC):
+    ID_FIELD = "id_"
+    HASH_FIELD = "hash_"
+
+    @abstractmethod
+    def insert(self, doc: dict) -> str:
+        pass
+
+    @abstractmethod
+    def insert_batch(self, docs: list[dict]) -> list[str]:
+        pass
+
+    @abstractmethod
+    def update(self, doc: dict) -> str:
+        pass
+
+    @abstractmethod
+    def upsert(self, doc: dict) -> str:
+        pass
+
+    @abstractmethod
+    def retrieve(self, doc_id: str) -> dict | None:
+        pass
+
+    @abstractmethod
+    def search(self, query_filter: Filter) -> list[dict]:
+        pass
+
+    @abstractmethod
+    def all(self) -> list[dict]:
+        pass
+
+    @abstractmethod
+    def remove(self, ids: list[str]) -> int:
+        pass
+
+
+class LightTinyDB(ILightDB):
+    def __init__(self, db_file_path: str, indexes: list[str] | None = None, log_level: int = logging.INFO, cacher: ICacher | None = None):
+        self._logger = get_logger(level=log_level)
+        self._db = TinyDB(db_file_path)
+
+        self._cacher = cacher or Cacher()
+
+        self._indexes = {}
+
+        for index_name in (indexes or []):
+            self._create_and_rebuild_index(index_name)
+
+    def all(self) -> list[dict]:
+        return [dict(document) for document in self._db.all()]
+
+    def insert(self, doc: dict) -> str:
+        doc_prepared = self._prepare_doc_for_store(doc)
+        cache_key = self._cacher.create_cache_key("exists", [doc])
+        if not self._cacher.exists(cache_key):
+            self._logger.debug(f"Inserting doc: {doc_prepared}")
+            self._db.insert(doc_prepared)
+            self._handle_index(doc_prepared)
+            self._cacher.set(cache_key, doc_prepared[self.ID_FIELD])
+        return doc_prepared[self.ID_FIELD]
+
+    def insert_batch(self, docs: list[dict]) -> list[str]:
+        docs_prepared = [self._prepare_doc_for_store(doc) for doc in docs]
+        docs_prepared = [doc for doc in docs_prepared if not self._cacher.exists(self._cacher.create_cache_key("exists", [doc]))]
+        self._logger.debug(f"Inserting doc batch: {docs_prepared}")
+        self._db.insert_multiple(docs_prepared)
+        for doc_prepared in docs_prepared:
+            self._handle_index(doc_prepared)
+        for doc in docs_prepared:
+            self._cacher.set(self._cacher.create_cache_key("exists", [doc]), doc[self.ID_FIELD])
+        return [doc[self.ID_FIELD] for doc in docs_prepared]
+
+    def update(self, doc: dict) -> str:
+        doc = self._prepare_doc_for_store(doc)
+        q = Query()
+        self._db.update(doc, q.id_ == doc[self.ID_FIELD])
+        self._handle_index(doc)
+        return doc[self.ID_FIELD]
+
+    def upsert(self, doc: dict) -> str:
+        doc = self._prepare_doc_for_store(doc)
+        q = Query()
+        self._db.upsert(doc, q.id_ == doc[self.ID_FIELD])
+        self._handle_index(doc)
+        return doc[self.ID_FIELD]
+
+    def retrieve(self, doc_id: str) -> dict | None:
+        docs = self.search(Filter(self.ID_FIELD, doc_id))
+        if len(docs) < 1:
+            return None
+        return docs[0]
+
+    def search(self, query_filter: Filter) -> list[dict]:
+        self._logger.debug(f"Searching with filter: {str(query_filter)}")
+        if query_filter.op == FilterOp.EQUAL and query_filter.value is not None:
+            if self._has_index(query_filter.key):
+                return self._search_index(query_filter.key, query_filter.value)
+        tinydb_query = Query()
+        documents = self._db.search(tinydb_query[query_filter.key].test(self._filter_to_test_fn(query_filter)))
+        self._logger.debug(f"> documents found: {len(documents)}")
+        return [self._prepare_doc_for_return(dict(document)) for document in documents]
+
+    def _create_and_rebuild_index(self, field_name: str) -> None:
+        self._logger.debug(f"Creating index for: '{field_name}'...")
+        self._indexes[field_name] = defaultdict(list)
+        docs = self._db.all()
+        for doc in tqdm(docs, desc=f"recreating index '{field_name}':'"):
+            self._add_to_index(field_name, doc)
+        self._logger.debug(f"> indexed {len(docs)} documents")
+
+    def _add_to_index(self, field_name: str, doc: dict) -> None:
+        # self._logger.debug(f"Adding to index '{field_name}': {doc.get(field_name, 'MISSING')}")
+        if field_name not in doc:
+            return
+        self._indexes[field_name][doc[field_name]].append(doc)
+
+    def _remove_from_index(self, field_name: str, doc: dict) -> None:
+        if field_name not in doc:
+            return
+        for indexed_doc in (self._indexes[field_name][doc[field_name]] or []):
+            if indexed_doc[self.ID_FIELD] == doc[self.ID_FIELD]:
+                self._indexes[field_name][doc[field_name]].remove(indexed_doc)
+
+    def _handle_index(self, doc: dict, remove: bool = False) -> None:
+        if remove:
+            self._handle_index_remove(doc)
+        else:
+            self._handle_index_add(doc)
+
+    def _handle_index_add(self, doc: dict) -> None:
+        for field_name in doc.keys():
+            if field_name in self._indexes:
+                self._add_to_index(field_name, doc)
+
+    def _handle_index_remove(self, doc: dict):
+        self._logger.debug(f"_handler_index_remove| doc: {doc}")
+        if len(doc.keys()) == 1 and list(doc.keys())[0] == self.ID_FIELD:
+            self._logger.debug(f"> the doc has a single key, trying to get the full document")
+            doc = self.retrieve(doc[self.ID_FIELD])
+            if doc is None:
+                return
+            self._logger.debug(f"> the full document: {doc}")
+        for doc_field in doc.keys():
+            if doc_field in self._indexes:
+                self._remove_from_index(doc_field, doc)
+
+    def _has_index(self, key: str) -> bool:
+        if key not in self._indexes:
+            self._logger.debug(f"> there is no index: '{key}'")
+            return False
+        return True
+
+    def _search_index(self, key: str, value: Any) -> list[dict]:
+        if value not in self._indexes[key]:
+            return []
+        return self._indexes[key][value]
+
+    def _filter_to_test_fn(self, query_filter: Filter):
+        def test_func(val):
+            if not isinstance(val, str) and not isinstance(val, int) and not isinstance(val, float):
+                return False
+            if query_filter.op == FilterOp.GREATER:
+                return val > query_filter.value
+            if query_filter.op == FilterOp.GREATER_OR_EQUAL:
+                return val >= query_filter.value
+            if query_filter.op == FilterOp.LESS:
+                return val < query_filter.value
+            if query_filter.op == FilterOp.LESS_OR_EQUAL:
+                return val <= query_filter.value
+            if query_filter.op == FilterOp.EQUAL:
+                return val == query_filter.value
+            raise ValueError(f"Unknown operator {query_filter.op}")
+        return test_func
+
+    def remove(self, ids: list[str] | str) -> int:
+        if isinstance(ids, str):
+            ids = [ids]
+        removed_elements_ids = []
+        for doc_id in ids:
+            self._handle_index({self.ID_FIELD: doc_id}, remove=True)
+            q = Query()
+            removed_elements = self._db.remove(q.id_ == doc_id)
+            removed_elements_ids.extend(removed_elements)
+        return len(list(set(removed_elements_ids)))
+
+    def _prepare_doc_for_store(self, doc: dict) -> dict:
+        self._ensure_created_at(doc)
+        return self._ensure_hash(self._ensure_id(doc))
+
+    def _prepare_doc_for_return(self, doc: dict) -> dict:
+        doc_copy = doc.copy()
+        self._ensure_created_at(doc_copy)
+        return doc_copy
+
+    def _ensure_created_at(self, doc: dict):
+        if "created_at" not in doc:
+            err_msg = f"There is no created_at in doc!\n{pp(doc, print_msg=False)}"
+            self._logger.error(err_msg)
+            raise ValueError(err_msg)
+
+    def _ensure_id(self, doc: dict) -> dict:
+        if self.ID_FIELD not in doc:
+            doc[self.ID_FIELD] = self._generate_id()
+        return doc
+
+    def _generate_id(self) -> str:
+        return str(uuid.uuid4())
+
+    def _ensure_hash(self, doc: dict) -> dict:
+        if self.HASH_FIELD in doc and doc[self.HASH_FIELD]:
+            self._logger.debug(f"the doc already has the hash property, keeping it as is: {doc[self.HASH_FIELD]}")
+            return doc
+        doc[self.HASH_FIELD] = self._generate_hash(doc)
+        return doc
+
+    def _generate_hash(self, doc: dict) -> str:
+        doc_copy = doc.copy()
+        if self.ID_FIELD in doc_copy:
+            del doc_copy[self.ID_FIELD]
+        if self.HASH_FIELD in doc_copy:
+            del doc_copy[self.HASH_FIELD]
+        return hash_dict(doc_copy)

toolchemy/utils/__init__.py

@@ -0,0 +1,19 @@
+from .cacher import ICacher, CacherPickle, Cacher, CacheEntryDoesNotExistError, DummyCacher
+from .locations import Locations
+from .logger import get_logger
+from .utils import pp, pp_cast, ff, hash_dict, to_json, truncate
+from .timer import Timer
+from .at_exit_collector import ICollectable, AtExitCollector
+
+
+__all__ = [
+    "ICollectable", "AtExitCollector",
+    "ICacher",
+    "CacherPickle",
+    "Cacher",
+    "CacheEntryDoesNotExistError",
+    "DummyCacher",
+    "get_logger",
+    "Locations",
+    "pp", "pp_cast", "ff", "hash_dict", "to_json", "truncate",
+    "Timer"]