json2vec 0.1.0__py3-none-any.whl

json2vec/logging/throughput.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import datetime
+from collections import defaultdict
+from functools import partialmethod
+from typing import TYPE_CHECKING
+
+import torch
+from lightning import Callback, Trainer
+
+from json2vec.structs.enums import Metric, Strata
+
+if TYPE_CHECKING:
+    from json2vec.architecture.root import JSON2Vec
+
+
+class ThroughputLogger(Callback):
+    def __init__(self):
+        super().__init__()
+
+        self.timestamp: dict[Strata, datetime.datetime] = defaultdict(lambda: datetime.datetime.now())
+
+    def start(self, trainer: Trainer, pl_module: JSON2Vec, batch, batch_idx, strata: Strata):
+        self.timestamp[strata] = datetime.datetime.now()
+
+    def end(self, trainer: Trainer, pl_module: JSON2Vec, outputs, batch, batch_idx, strata: Strata):
+        now = datetime.datetime.now()
+        then = self.timestamp[strata]
+        throughput = pl_module.session.structure.batch_size / (now - then).total_seconds()
+
+        pl_module.track((Metric.throughput, strata), torch.tensor(throughput))
+
+    on_train_batch_start = partialmethod(start, strata=Strata.train)
+    on_validation_batch_start = partialmethod(start, strata=Strata.validate)
+    on_test_batch_start = partialmethod(start, strata=Strata.test)
+
+    on_train_batch_end = partialmethod(end, strata=Strata.train)
+    on_validation_batch_end = partialmethod(end, strata=Strata.validate)
+    on_test_batch_end = partialmethod(end, strata=Strata.test)

json2vec/logging/tracking.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import enum
+from typing import Callable
+
+from lightning.pytorch.loggers import Logger
+from loguru import logger
+
+from json2vec.structs.environment import TrackingEnvironment
+
+
+class LoggingFramework(enum.StrEnum):
+    wandb = "wandb"
+    neptune = "neptune"
+    comet = "comet"
+    mlflow = "mlflow"
+    tensorboard = "tensorboard"
+    csv = "csv"
+
+
+class LoggerFactory:
+    AUTO_DETECTION_ORDER: tuple[LoggingFramework, ...] = (
+        LoggingFramework.wandb,
+        LoggingFramework.neptune,
+        LoggingFramework.comet,
+        LoggingFramework.mlflow,
+        LoggingFramework.tensorboard,
+        LoggingFramework.csv,
+    )
+    AUTO_LOGGER_FIELDS: dict[LoggingFramework, tuple[str, ...]] = {
+        LoggingFramework.wandb: ("wandb_api_key",),
+        LoggingFramework.neptune: ("neptune_api_token",),
+        LoggingFramework.comet: ("comet_api_key",),
+        LoggingFramework.mlflow: ("mlflow_tracking_uri",),
+        LoggingFramework.tensorboard: ("tensorboard_log_dir",),
+        LoggingFramework.csv: ("csv_log_dir",),
+    }
+
+    @staticmethod
+    def wandb(project: str, run: str, notes: str) -> Logger:
+        from lightning.pytorch.loggers import WandbLogger
+
+        tracker = WandbLogger(project=project, name=run)
+        if notes:
+            try:
+                tracker.experiment.notes = notes
+            except Exception:
+                logger.bind(component="tracking", backend=LoggingFramework.wandb.value).warning(
+                    "failed to attach run notes"
+                )
+        return tracker
+
+    @staticmethod
+    def neptune(project: str, run: str, notes: str) -> Logger:
+        from lightning.pytorch.loggers import NeptuneLogger
+
+        tracker = NeptuneLogger(project=project, name=run)
+        if notes:
+            try:
+                tracker.experiment["sys/notes"] = notes
+            except Exception:
+                logger.bind(component="tracking", backend=LoggingFramework.neptune.value).warning(
+                    "failed to attach run notes"
+                )
+        return tracker
+
+    @staticmethod
+    def comet(project: str, run: str, notes: str) -> Logger:
+        from lightning.pytorch.loggers import CometLogger
+
+        tracker = CometLogger(project_name=project, experiment_name=run)
+        if notes:
+            try:
+                tracker.experiment.log_other("notes", notes)
+            except Exception:
+                logger.bind(component="tracking", backend=LoggingFramework.comet.value).warning(
+                    "failed to attach run notes"
+                )
+        return tracker
+
+    @staticmethod
+    def mlflow(project: str, run: str, notes: str) -> Logger:
+        from lightning.pytorch.loggers import MLFlowLogger
+
+        tags = {"notes": notes} if notes else None
+        return MLFlowLogger(experiment_name=project, run_name=run, tags=tags)
+
+    @staticmethod
+    def tensorboard(project: str, run: str, _: str) -> Logger:
+        from lightning.pytorch.loggers import TensorBoardLogger
+
+        save_dir = TrackingEnvironment.from_env().resolved_tensorboard_log_dir
+        return TensorBoardLogger(save_dir=save_dir, name=project, version=run)
+
+    @staticmethod
+    def csv(project: str, run: str, _: str) -> Logger:
+        from lightning.pytorch.loggers import CSVLogger
+
+        save_dir = TrackingEnvironment.from_env().resolved_csv_log_dir
+        return CSVLogger(save_dir=save_dir, name=project, version=run)
+
+    @staticmethod
+    def _builders() -> dict[LoggingFramework, Callable[[str, str, str], Logger]]:
+        return {
+            LoggingFramework.wandb: LoggerFactory.wandb,
+            LoggingFramework.neptune: LoggerFactory.neptune,
+            LoggingFramework.comet: LoggerFactory.comet,
+            LoggingFramework.mlflow: LoggerFactory.mlflow,
+            LoggingFramework.tensorboard: LoggerFactory.tensorboard,
+            LoggingFramework.csv: LoggerFactory.csv,
+        }
+
+    @staticmethod
+    def _resolve_framework() -> LoggingFramework | None:
+        settings = TrackingEnvironment.from_env()
+        forced = settings.logger
+        if forced is not None:
+            forced = forced.lower()
+            if forced in {"none", "false", "off", "disabled"}:
+                return None
+
+            try:
+                return LoggingFramework(forced)
+            except ValueError:
+                logger.bind(component="tracking", backend=forced).warning("unsupported logger backend override")
+                return None
+
+        for backend in LoggerFactory.AUTO_DETECTION_ORDER:
+            if any(getattr(settings, field) is not None for field in LoggerFactory.AUTO_LOGGER_FIELDS[backend]):
+                return backend
+
+        return None
+
+    @staticmethod
+    def create(project: str, run: str, notes: str) -> Logger | bool:
+        backend = LoggerFactory._resolve_framework()
+        if backend is None:
+            return False
+
+        builder = LoggerFactory._builders().get(backend)
+        if builder is None:
+            logger.bind(component="tracking", backend=backend.value).warning("unsupported logger backend")
+            return False
+
+        try:
+            tracker = builder(project, run, notes)
+        except Exception:
+            logger.bind(component="tracking", backend=backend.value).exception("failed to initialize trainer logger")
+            return False
+
+        logger.bind(component="tracking", backend=backend.value, project=project, run=run).info("enabled trainer logger")
+        return tracker

json2vec/processors/__init__.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import importlib
+import pkgutil
+
+subpkg = importlib.import_module(".extensions", __name__)
+for _, fullname, _ in pkgutil.iter_modules(subpkg.__path__, subpkg.__name__ + "."):
+    importlib.import_module(fullname)

json2vec/processors/base.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import ast
+import enum
+import inspect
+import textwrap
+from functools import cache
+from typing import Any, Callable
+
+import pluggy
+import pydantic
+
+from json2vec.processors.spec import PluginSpec
+
+pm: pluggy.PluginManager = pluggy.PluginManager(project_name="processors")
+
+pm.add_hookspecs(module_or_class=PluginSpec)
+
+
+class ProcessorMode(enum.StrEnum):
+    yielding = "yield"
+    returning = "return"
+
+
+def has_yield_expression(node: ast.AST, root: bool = False) -> bool:
+    for child in ast.iter_child_nodes(node):
+        if isinstance(child, (ast.Yield, ast.YieldFrom)):
+            return True
+
+        if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
+            if root and has_yield_expression(child):
+                return True
+            continue
+
+        if has_yield_expression(child):
+            return True
+
+    return False
+
+
+def is_yielding_processor(func: Callable[..., Any]) -> bool:
+    try:
+        source: str = textwrap.dedent(inspect.getsource(func))
+    except (OSError, TypeError):
+        return inspect.isgeneratorfunction(func)
+
+    module: ast.Module = ast.parse(source)
+    candidates: list[ast.FunctionDef | ast.AsyncFunctionDef] = [
+        node
+        for node in module.body
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+    ]
+
+    target = next((node for node in candidates if node.name == func.__name__), None)
+    if target is None:
+        return inspect.isgeneratorfunction(func)
+
+    return has_yield_expression(target, root=True)
+
+
+class Processor(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True, frozen=True)
+    name: str
+    func: Callable[..., Any]
+    mode: ProcessorMode
+
+    def __call__(self, observation: dict, **kwargs) -> Any:
+        return self.func(observation, **_filter_supported_kwargs(self.func, kwargs))
+
+
+@cache
+def _accepted_kwargs(func: Callable[..., Any]) -> tuple[bool, frozenset[str]]:
+    signature = inspect.signature(func)
+    accepts_variadic_kwargs = any(
+        parameter.kind == inspect.Parameter.VAR_KEYWORD
+        for parameter in signature.parameters.values()
+    )
+    accepted = frozenset(signature.parameters.keys())
+    return accepts_variadic_kwargs, accepted
+
+
+def _filter_supported_kwargs(func: Callable[..., Any], kwargs: dict[str, Any]) -> dict[str, Any]:
+    accepts_variadic_kwargs, accepted = _accepted_kwargs(func)
+    if accepts_variadic_kwargs:
+        return kwargs
+
+    return {key: value for key, value in kwargs.items() if key in accepted}
+
+
+PROCESSORS: dict[str, Processor] = {}
+
+
+def register(func: Callable[..., Any]) -> Callable[..., Any]:
+    name = func.__name__
+
+    if name in PROCESSORS:
+        raise ValueError(f"Processor '{name}' is already registered.")
+
+    mode: ProcessorMode = ProcessorMode.yielding if is_yielding_processor(func) else ProcessorMode.returning
+    PROCESSORS[name] = Processor(name=name, func=func, mode=mode)
+
+    return func

json2vec/processors/extensions/example.py

@@ -0,0 +1,6 @@
+from json2vec.processors.base import register
+
+
+@register
+def default(item):
+    return item

json2vec/processors/spec.py

@@ -0,0 +1,8 @@
+import pluggy
+
+hookspec = pluggy.HookspecMarker("processors")
+
+
+class PluginSpec:
+    @hookspec
+    def plugin_class(self) -> None: ...

json2vec/structs/enums.py

@@ -0,0 +1,84 @@
+import enum
+
+
+class Tokens(enum.IntEnum):
+    valued = 0
+    null = 1
+    padded = 2
+    masked = 3
+    pruned = 4
+    other = 5
+
+
+class Stage(enum.StrEnum):
+    fit = "fit"
+    validate = "validate"
+    test = "test"
+    predict = "predict"
+
+
+class Strata(enum.StrEnum):
+    train = "train"
+    validate = "validate"
+    test = "test"
+    predict = "predict"
+
+    @classmethod
+    def from_stage(cls, stage: Stage | str) -> list["Strata"]:
+        match stage:
+            case Stage.fit:
+                return [cls.train, cls.validate]
+            case Stage.validate:
+                return [cls.validate]
+            case Stage.test:
+                return [cls.test]
+            case Stage.predict:
+                return [cls.predict]
+            case _:
+                raise ValueError(f"Unknown stage: {stage}")
+
+
+class Suffix(enum.StrEnum):
+    feather = "feather"
+    parquet = "parquet"
+    ndjson = "ndjson"
+    avro = "avro"
+    csv = "csv"
+    orc = "orc"
+    json = "json"
+
+
+class TensorKey(enum.StrEnum):
+    value = "value"
+    content = "content"
+    state = "state"
+    intervals = "intervals"
+    probability = "probability"
+    topk = "topk"
+    embedding = "embedding"
+
+
+class Metric(enum.StrEnum):
+    accuracy = "accuracy"
+    precision = "precision"
+    recall = "recall"
+    loss = "loss"
+    sigma = "sigma"
+    throughput = "throughput"
+    mae = "mae"
+    rmse = "rmse"
+
+
+class ShardingStrategy(enum.StrEnum):
+    file = "file"
+    chunk = "chunk"
+    record = "record"
+
+
+class Component(enum.StrEnum):
+    Request = "Request"
+    Embedder = "Embedder"
+    Decoder = "Decoder"
+    TensorField = "TensorField"
+    loss = "loss"
+    write = "write"

json2vec/structs/environment.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import os
+from typing import Literal, Self
+from urllib.parse import urlparse
+
+from pydantic import AliasChoices, Field, ValidationInfo, field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from json2vec.structs.enums import ShardingStrategy
+
+
+class DataLoaderEnvironment(BaseSettings):
+    model_config = SettingsConfigDict(extra="ignore", case_sensitive=False)
+
+    num_workers: int | None = Field(
+        default=None,
+        ge=0,
+        validation_alias=AliasChoices("JSON2VEC_NUM_WORKERS", "NUM_WORKERS"),
+    )
+    persistent_workers: bool = Field(
+        default=True,
+        validation_alias=AliasChoices("JSON2VEC_PERSISTENT_WORKERS", "PERSISTENT_WORKERS"),
+    )
+    pin_memory: bool = Field(
+        default=True,
+        validation_alias=AliasChoices("JSON2VEC_PIN_MEMORY", "PIN_MEMORY"),
+    )
+    sharding: ShardingStrategy = Field(
+        default=ShardingStrategy.file,
+        validation_alias=AliasChoices("JSON2VEC_SHARDING", "JSON2VEC_SHARDING_STRATEGY", "SHARDING_STRATEGY"),
+    )
+    chunk_batch_size: int = Field(
+        default=4096,
+        ge=1,
+        validation_alias=AliasChoices("JSON2VEC_CHUNK_BATCH_SIZE", "JSON2VEC_PYARROW_BATCH_SIZE", "CHUNK_BATCH_SIZE"),
+    )
+
+    @field_validator("sharding", mode="before")
+    @classmethod
+    def normalize_sharding(cls, value: ShardingStrategy | str) -> ShardingStrategy | str:
+        if isinstance(value, str):
+            normalized = value.strip().lower()
+            return normalized
+
+        return value
+
+    @classmethod
+    def from_env(cls) -> Self:
+        return cls()
+
+
+class TrackingEnvironment(BaseSettings):
+    model_config = SettingsConfigDict(extra="ignore", case_sensitive=False)
+
+    logger: str | None = Field(default=None, validation_alias=AliasChoices("JSON2VEC_LOGGER"))
+    wandb_api_key: str | None = Field(default=None, validation_alias=AliasChoices("WANDB_API_KEY"))
+    neptune_api_token: str | None = Field(default=None, validation_alias=AliasChoices("NEPTUNE_API_TOKEN"))
+    comet_api_key: str | None = Field(default=None, validation_alias=AliasChoices("COMET_API_KEY"))
+    mlflow_tracking_uri: str | None = Field(default=None, validation_alias=AliasChoices("MLFLOW_TRACKING_URI"))
+    tensorboard_log_dir: str | None = Field(
+        default=None,
+        validation_alias=AliasChoices("JSON2VEC_TENSORBOARD_LOG_DIR", "TENSORBOARD_LOG_DIR"),
+    )
+    csv_log_dir: str | None = Field(
+        default=None,
+        validation_alias=AliasChoices("JSON2VEC_CSV_LOG_DIR", "CSV_LOG_DIR"),
+    )
+
+    @field_validator("*", mode="before")
+    @classmethod
+    def strip_string_values(cls, value):
+        if isinstance(value, str):
+            stripped = value.strip()
+            if stripped == "":
+                return None
+            return stripped
+
+        return value
+
+    @property
+    def resolved_tensorboard_log_dir(self) -> str:
+        return self.tensorboard_log_dir or "logs/tensorboard"
+
+    @property
+    def resolved_csv_log_dir(self) -> str:
+        return self.csv_log_dir or "logs/csv"
+
+    @classmethod
+    def from_env(cls) -> Self:
+        return cls()
+
+
+class DeploymentEnvironment(BaseSettings):
+    model_config = SettingsConfigDict(extra="ignore", case_sensitive=False)
+
+    checkpoint: str = Field(
+        default="model.ckpt",
+        validation_alias=AliasChoices("JSON2VEC_CHECKPOINT", "CHECKPOINT"),
+    )
+    max_batch_size: int = Field(
+        default=128,
+        ge=1,
+        validation_alias=AliasChoices("JSON2VEC_MAX_BATCH_SIZE", "MAX_BATCH_SIZE"),
+    )
+    batch_timeout: float = Field(
+        default=0.0,
+        ge=0.0,
+        validation_alias=AliasChoices("JSON2VEC_BATCH_TIMEOUT", "BATCH_TIMEOUT"),
+    )
+    workers_per_device: int = Field(
+        default=1,
+        ge=1,
+        validation_alias=AliasChoices("JSON2VEC_WORKERS_PER_DEVICE", "JSON2VEC_N_WORKERS", "N_WORKERS"),
+    )
+    accelerator: Literal["auto", "cpu", "cuda", "mps"] = Field(
+        default="auto",
+        validation_alias=AliasChoices("JSON2VEC_ACCELERATOR", "ACCELERATOR"),
+    )
+    track_requests: bool = Field(
+        default=False,
+        validation_alias=AliasChoices("JSON2VEC_TRACK_REQUESTS", "TRACK_REQUESTS"),
+    )
+
+    @field_validator("checkpoint", "accelerator", mode="before")
+    @classmethod
+    def strip_required_strings(cls, value: str | None, info: ValidationInfo) -> str | None:
+        if isinstance(value, str):
+            stripped = value.strip()
+            if stripped == "":
+                raise ValueError(f"{info.field_name} must not be blank")
+            return stripped
+
+        return value
+
+    @classmethod
+    def from_env(cls) -> Self:
+        return cls()