json2vec 0.1.0__py3-none-any.whl

json2vec/data/processing.py

@@ -0,0 +1,152 @@
+import inspect
+from collections.abc import Callable
+from functools import partial
+from typing import Any
+
+import numpy as np
+from beartype import beartype
+
+from json2vec.structs.enums import Tokens
+
+
+def apply(
+    values: Any,
+    function: Callable[..., Any],
+    /,
+    *args: Any,
+    leaf_depth: int | None = None,
+    **kwargs: Any,
+) -> Any:
+    """Apply a function recursively to nested list leaves.
+
+    When ``leaf_depth`` is set, the function is applied exactly at that depth;
+    higher-level non-list values are preserved so downstream padding can mark
+    them as incomplete.
+    """
+    if leaf_depth is not None and leaf_depth < 0:
+        raise ValueError("leaf_depth must be >= 0")
+
+    def walk(node: Any, depth: int) -> Any:
+        if leaf_depth is None:
+            if isinstance(node, list):
+                return [walk(item, depth + 1) for item in node]
+
+            if node is None:
+                return None
+
+            return function(node, *args, **kwargs)
+
+        if depth == leaf_depth:
+            if node is None:
+                return None
+
+            return function(node, *args, **kwargs)
+
+        if isinstance(node, list):
+            return [walk(item, depth + 1) for item in node]
+
+        return node
+
+    return walk(values, depth=0)
+
+
+def _iter_leaf_nodes(
+    nested: Any,
+    shape: tuple[int, ...],
+    strides: tuple[int, ...],
+):
+    ndim = len(shape)
+    stack: list[tuple[Any, int, int]] = [(nested, 0, 0)]
+
+    while stack:
+        node, depth, base = stack.pop()
+
+        if depth == ndim:
+            yield base, node
+            continue
+
+        if not isinstance(node, list):
+            continue
+
+        limit = min(len(node), shape[depth])
+        step = strides[depth]
+        for index in range(limit - 1, -1, -1):
+            stack.append((node[index], depth + 1, base + (index * step)))
+
+
+def _fill_python(
+    nested: Any,
+    flat_values: np.ndarray,
+    flat_flags: np.ndarray,
+    shape: tuple[int, ...],
+    strides: tuple[int, ...],
+) -> None:
+    for flat_index, node in _iter_leaf_nodes(nested=nested, shape=shape, strides=strides):
+        if node is None:
+            flat_flags[flat_index] = Tokens.null.value
+        else:
+            flat_values[flat_index] = node
+            flat_flags[flat_index] = Tokens.valued.value
+
+
+@beartype
+def pad(
+    nested: Any, shape: tuple[int, ...], dtype: type | str = object, pad_value: Any = None
+) -> tuple[np.ndarray, np.ndarray]:
+    resolved_dtype = np.dtype(dtype)
+    values = np.full(shape, pad_value, dtype=resolved_dtype)
+    flags = np.full(shape, Tokens.padded.value, dtype=np.int8)
+
+    ndim = len(shape)
+    if ndim == 0:
+        if nested is None:
+            flags[...] = Tokens.null.value
+        else:
+            values[...] = nested
+            flags[...] = Tokens.valued.value
+        return values, flags
+
+    strides = [1] * ndim
+    for depth in range(ndim - 2, -1, -1):
+        strides[depth] = strides[depth + 1] * shape[depth + 1]
+    stride_tuple = tuple(strides)
+
+    flat_values = values.reshape(-1)
+    flat_flags = flags.reshape(-1)
+
+    _fill_python(
+        nested=nested,
+        flat_values=flat_values,
+        flat_flags=flat_flags,
+        shape=shape,
+        strides=stride_tuple,
+    )
+
+    return values, flags
+
+
+@beartype
+class Pipeline:
+    def __init__(self, **arguments):
+        self.arguments: dict[str, Any] = arguments
+        self.steps: list[Callable] = []
+
+    def __or__(self, function: Callable) -> "Pipeline":
+        required = [name for name in inspect.signature(function).parameters.keys()]
+
+        available = set(required) & set(self.arguments.keys())
+
+        self.steps.append(partial(function, **{arg: self.arguments[arg] for arg in available}))
+
+        return self
+
+    def __repr__(self):
+        return f"Pipeline({repr(self.source)}, {repr(self.arguments)})"
+
+    def __iter__(self):
+        stream = self.steps[0]()
+
+        for step in self.steps[1:]:
+            stream = step(stream)
+
+        return iter(stream)

json2vec/entrypoints/__init__.py

@@ -0,0 +1,3 @@
+from json2vec.entrypoints.pipeline import build, execute, fit, predict, test, validate
+
+__all__ = ["build", "fit", "validate", "test", "predict", "execute"]

json2vec/entrypoints/pipeline.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import torch
+from lightning.pytorch.callbacks import Callback, EarlyStopping, ModelCheckpoint
+from lightning.pytorch.trainer.trainer import Trainer
+from loguru import logger
+
+from json2vec.architecture.root import JSON2Vec
+from json2vec.inference.callback import Writer
+from json2vec.logging.epoch import EpochLifecycleLogger
+from json2vec.logging.throughput import ThroughputLogger
+from json2vec.logging.tracking import LoggerFactory
+from json2vec.structs.enums import Metric, Stage, Strata
+from json2vec.structs.experiment import Experiment, PatchOp, Session
+
+
+def build(model: JSON2Vec, callbacks: list[Callback], names: list[str] | None = None) -> Trainer:
+    active_callbacks: list[Callback] = list(callbacks)
+    if not any(isinstance(callback, EpochLifecycleLogger) for callback in active_callbacks):
+        active_callbacks.append(EpochLifecycleLogger())
+
+    logger.bind(
+        component="trainer",
+        session=model.session.name,
+        stage=model.session.task,
+        callbacks=[type(callback).__name__ for callback in active_callbacks],
+    ).info("building lightning trainer")
+
+    return Trainer(
+        accelerator="auto" if torch.cuda.is_available() else "cpu",
+        precision="bf16-mixed" if torch.cuda.is_available() else None,
+        logger=LoggerFactory.create(*names) if names is not None else False,
+        enable_model_summary=False,
+        enable_progress_bar=False,
+        callbacks=active_callbacks,
+        **model.session.trainer,
+    )
+
+
+def fit(
+    names: list[str],
+    session: Session | None = None,
+    checkpoint: str | os.PathLike[str] | None = None,
+    patches: list[PatchOp] | None = None,
+) -> Path:
+    logger.bind(component="task", task="fit", session=session.name if session else None).info("starting fit task")
+
+    checkpoint_path = str(checkpoint) if checkpoint is not None else None
+    model: JSON2Vec = JSON2Vec.get_or_create(session=session, checkpoint=checkpoint_path)
+    model.session = model.session.patch(patches)
+
+    monitor = f"{Metric.loss}/{Strata.validate}"
+    filename: str = f"{model.session.structure.name}-{model.session.name}-" + "{epoch}-{step}-{val_loss:.2f}"
+
+    checkpoint_dir = Path("models")
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    checkpointer: ModelCheckpoint = ModelCheckpoint(dirpath=checkpoint_dir, filename=filename, monitor=monitor)
+    callbacks: list[Callback] = [ThroughputLogger(), checkpointer]
+
+    if (patience := model.session.patience) is not None:
+        callbacks.append(EarlyStopping(patience=patience, monitor=monitor))
+
+    trainer: Trainer = build(model=model, callbacks=callbacks, names=names)
+    trainer.fit(model=model)
+
+    best_path = Path(str(checkpointer.best_model_path))
+    logger.bind(component="task", task="fit", session=model.session.name, checkpoint=str(best_path)).info(
+        "finished fit task"
+    )
+    return best_path
+
+
+def validate(
+    names: list[str],
+    checkpoint: str | os.PathLike[str],
+    session: Session | None = None,
+    patches: list[PatchOp] | None = None,
+) -> None:
+    logger.bind(component="task", task="validate", session=session.name if session else None).info(
+        "starting validate task"
+    )
+
+    model: JSON2Vec = JSON2Vec.get_or_create(session=session, checkpoint=str(checkpoint))
+    model.session = model.session.patch(patches)
+
+    callbacks: list[Callback] = [ThroughputLogger()]
+    trainer: Trainer = build(model=model, callbacks=callbacks, names=names)
+    trainer.validate(model=model)
+    logger.bind(component="task", task="validate", session=model.session.name).info("finished validate task")
+
+
+def test(
+    names: list[str],
+    checkpoint: str | os.PathLike[str],
+    session: Session | None = None,
+    patches: list[PatchOp] | None = None,
+) -> None:
+    logger.bind(component="task", task="test", session=session.name if session else None).info("starting test task")
+
+    model: JSON2Vec = JSON2Vec.get_or_create(session=session, checkpoint=str(checkpoint))
+    model.session = model.session.patch(patches)
+
+    callbacks: list[Callback] = [ThroughputLogger()]
+    trainer: Trainer = build(model=model, callbacks=callbacks, names=names)
+    trainer.test(model=model)
+    logger.bind(component="task", task="test", session=model.session.name).info("finished test task")
+
+
+def predict(
+    session: Session | None,
+    names: list[str] | None,
+    checkpoint: str | os.PathLike[str],
+    patches: list[PatchOp] | None = None,
+) -> Path:
+    logger.bind(component="task", task="predict", session=session.name if session else None).info("starting predict task")
+
+    model: JSON2Vec = JSON2Vec.get_or_create(session=session, checkpoint=str(checkpoint))
+    model.session = model.session.patch(patches)
+
+    os.makedirs(name=(outpath := "tmp/predictions"), exist_ok=True)
+    callbacks: list[Callback] = [Writer(outpath)]
+    trainer: Trainer = build(model=model, callbacks=callbacks, names=names)
+    trainer.predict(model=model, return_predictions=False)
+
+    output_path = Path(outpath)
+    logger.bind(component="task", task="predict", session=model.session.name, output=str(output_path)).info(
+        "finished predict task"
+    )
+    return output_path
+
+
+def execute(experiment: Experiment) -> dict[str, Any]:
+    logger.bind(
+        component="pipeline",
+        project=experiment.project,
+        run=experiment.name,
+        sessions=len(experiment.sessions),
+    ).info("starting experiment execution")
+
+    checkpoint: str | os.PathLike[str] | None = experiment.checkpoint
+    names: list[str] = [experiment.project, experiment.name, experiment.notes]
+
+    tasks: dict[Stage, Any] = {
+        Stage.fit: fit,
+        Stage.validate: validate,
+        Stage.test: test,
+        Stage.predict: predict,
+    }
+
+    outputs: dict[str, Any] = {}
+
+    for session in experiment.sessions:
+        logger.bind(component="pipeline", session=session.name, stage=session.task).info("dispatching session")
+        task = tasks[session.task]
+
+        output = task(
+            session=session,
+            checkpoint=checkpoint,
+            names=names,
+        )
+        outputs[session.name] = output
+
+        if isinstance(output, (str, os.PathLike)) and session.task == Stage.fit:
+            checkpoint = output
+
+    logger.bind(component="pipeline", project=experiment.project, run=experiment.name).info(
+        "finished experiment execution"
+    )
+
+    return outputs

json2vec/inference/callback.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Any
+
+import lightning.pytorch as lit
+import polars as pl
+import pyarrow as pa
+import pyarrow.parquet as pq
+from lightning.pytorch import callbacks
+from tensordict import TensorDict
+
+from json2vec.structs.packages import Prediction
+from json2vec.structs.tree import Address
+from json2vec.tensorfields.base import TensorFieldBase
+
+if TYPE_CHECKING:
+    from json2vec.architecture.root import JSON2Vec
+
+
+class Writer(callbacks.BasePredictionWriter):
+
+    def __init__(self, path: os.PathLike | str, flush_every_n_batches: int | None = None):
+
+        super().__init__(write_interval="batch")
+
+        self.path: os.PathLike = path
+        self.flush_every_n_batches: int | None = flush_every_n_batches
+        self.schema: pa.schema | None = None
+        self.writer: pq.ParquetWriter | None = None
+
+    @staticmethod
+    def _as_struct_frame(
+        values_by_address: dict[Address, dict[str, Any]], alias: str, num_rows: int
+    ) -> pl.DataFrame:
+        if len(values_by_address) == 0:
+            return pl.DataFrame({alias: [None] * num_rows})
+
+        columns: list[pl.DataFrame] = []
+        for address, values in values_by_address.items():
+            field_frame = pl.DataFrame(data=values)
+            columns.append(field_frame.select(pl.struct(pl.all()).alias(name=address)))
+
+        nested: pl.DataFrame = pl.concat(items=columns, how="horizontal")
+        return nested.select(pl.struct(pl.all()).alias(name=alias))
+
+    def write_on_batch_end(
+        self,
+        trainer: lit.Trainer,
+        pl_module: JSON2Vec,
+        output: dict[str, list[Prediction]],
+        batch_indices: list[int]|None,
+        batch: TensorDict[Address, TensorFieldBase],
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        num_rows = len(batch["metadata"])
+
+        supervised: dict[Address, dict[TensorKey, Any]]
+        embeddings: dict[Address, dict[TensorKey, Any]]
+
+        supervised, embeddings = pl_module.write(predictions=output["predictions"])
+
+        items = [
+            pl.from_records(data=batch["metadata"], schema=["inputs"], orient="row"),
+            self._as_struct_frame(values_by_address=supervised, alias="predictions", num_rows=num_rows),
+        ]
+
+        if len(embeddings) > 0:
+            items.append(self._as_struct_frame(values_by_address=embeddings, alias="embeddings", num_rows=num_rows))
+
+        table: pa.Table = pl.concat(
+            items=items,
+            how="horizontal"
+        ).to_arrow()
+
+        if self.writer is None:
+
+            os.makedirs(self.path, exist_ok=True)
+            self.schema: pa.schema = table.schema
+
+            self.writer: pq.ParquetWriter = pq.ParquetWriter(
+                where=os.path.join(self.path, f"rank-{trainer.local_rank}.parquet"),
+                schema=self.schema
+            )
+
+        if table.schema != self.schema:
+            table = table.cast(self.schema)
+
+        self.writer.write_table(table)
+
+        if self.flush_every_n_batches and (batch_idx + 1) % self.flush_every_n_batches == 0 and hasattr(self.writer, "flush"):
+            self.writer.flush()
+
+    def on_predict_end(self, trainer: lit.Trainer, pl_module: lit.LightningModule) -> None:
+        if self.writer:
+            self.writer.close()
+            self.writer: None = None

json2vec/inference/deployment.py

@@ -0,0 +1,175 @@
+from dataclasses import dataclass
+from typing import Any, Type, TypeAlias
+
+import litserve as ls
+import pydantic
+import torch
+from beartype import beartype
+from tensordict import TensorDict
+
+from json2vec.architecture.root import JSON2Vec
+from json2vec.data.datasets import encode, process
+from json2vec.structs.enums import Strata
+from json2vec.structs.environment import DeploymentEnvironment
+from json2vec.structs.packages import Prediction
+from json2vec.structs.tree import Address
+from json2vec.tensorfields.base import TensorFieldBase
+
+Input: TypeAlias = TensorDict[Address, TensorFieldBase]
+
+
+@dataclass
+class ErrorItem:
+    status_code: int
+    message: str
+
+
+@dataclass
+class BatchItem:
+    data: Input | None
+    valid_indices: list[int]
+    items: list[Input | ErrorItem]
+
+
+class Deployment(ls.LitAPI):
+    def __init__(self, checkpoint: str, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.checkpoint = checkpoint
+
+    def setup(self, device: str) -> None:
+        self.model: JSON2Vec = JSON2Vec.get_or_create(checkpoint=self.checkpoint).to(device)
+        self.model.eval()
+        self.state = self.model.state
+
+    @beartype
+    def decode_request(self, request: dict[str, Any] | pydantic.BaseModel) -> Input | ErrorItem:
+
+        if isinstance(request, pydantic.BaseModel):
+            request = request.model_dump()
+
+        try:
+            observations: list[Any] = list(
+                process(
+                    pipe=[request],
+                    session=self.model.session,
+                    strata=Strata.predict,
+                    state=self.state,
+                )
+            )
+
+        except Exception as exception:
+            return ErrorItem(status_code=422, message=str(exception))
+
+
+        if len(observations) == 0 or any(x is None for x in observations):
+            return ErrorItem(status_code=422, message="processor returned no observations for request")
+    
+
+        encoded = encode(
+            batch=observations,
+            session=self.model.session,
+            strata=Strata.predict,
+            state=self.state,
+        )
+
+        if encoded is None:
+            return ErrorItem(status_code=422, message="processor eliminated observation (filter)")
+
+        return encoded
+
+    @beartype
+    def batch(self, inputs: list[Input | ErrorItem]) -> BatchItem:
+        valid_indices: list[int] = []
+        valid_inputs: list[Input] = []
+
+        for index, item in enumerate(inputs):
+            if isinstance(item, ErrorItem):
+                continue
+
+            valid_indices.append(index)
+            valid_inputs.append(item)
+
+        data = torch.stack(valid_inputs, dim=0) if len(valid_inputs) > 0 else None
+        return BatchItem(data=data, valid_indices=valid_indices, items=inputs)
+
+    @beartype
+    def unbatch(self, outputs: list[Any]) -> list[Any]:
+        return list(outputs)
+
+    @beartype
+    def predict(self, data: BatchItem | Input | ErrorItem) -> list[list[Prediction] | ErrorItem] | list[Prediction] | ErrorItem:
+        if isinstance(data, ErrorItem):
+            return data
+
+        if isinstance(data, TensorDict):
+            with torch.inference_mode():
+                return self.model(data.to(self.device))
+
+        outputs: list[Any] = list(data.items)
+
+        if data.data is None:
+            return outputs
+
+        with torch.inference_mode():
+            predictions = self.model(data.data.to(self.device))
+
+        unbatched = Prediction.unbatch(predictions=predictions)
+
+        for index, item_predictions in zip(data.valid_indices, unbatched):
+            outputs[index] = item_predictions
+
+        return outputs
+
+    @beartype
+    def encode_response(self, response: list[Prediction] | ErrorItem) -> dict[str, Any] | pydantic.BaseModel:
+        if isinstance(response, ErrorItem):
+            return {
+                "predictions": {},
+                "error": {
+                    "status_code": response.status_code,
+                    "message": response.message,
+                },
+            }
+
+        predictions, embeddings = self.model.write(predictions=response)
+
+        payload = dict(predictions = predictions)
+
+        if len(embeddings) > 0:
+            payload["embeddings"] = embeddings
+
+        return Prediction.denest(payload)
+
+    @classmethod
+    @beartype
+    def forge(
+        cls,
+        request: Type[pydantic.BaseModel]|None=None,
+        response: Type[pydantic.BaseModel]|None=None,
+    ) -> Type["Deployment"]:
+
+        if request is not None:
+            cls.decode_request.__annotations__["request"] = request
+
+        if response is not None:
+            cls.encode_response.__annotations__["return"] = response
+
+        return cls
+
+    @classmethod
+    def serve(cls):
+
+        environment = DeploymentEnvironment()
+
+        server: ls.LitServer = ls.LitServer(
+            lit_api=Deployment(
+                checkpoint=environment.checkpoint,
+                max_batch_size=environment.max_batch_size,
+                batch_timeout=environment.batch_timeout,
+            ),
+            accelerator=environment.accelerator,
+            track_requests=environment.track_requests,
+            workers_per_device=environment.workers_per_device,
+        )
+
+        server.run(generate_client_file=False)

json2vec/logging/config.py

@@ -0,0 +1,27 @@
+import json
+import os
+import sys
+
+from loguru import logger
+from rich.console import Console
+from rich.json import JSON
+
+console = Console(file=sys.stdout)
+LOG_LEVEL: str = os.getenv("JSON2VEC_LOG_LEVEL", "DEBUG").upper()
+
+
+def sink(message):
+    record = message.record
+    extras = {k: str(v) for k, v in record["extra"].items()}
+    payload = {
+        "timestamp": record["time"].strftime("%Y-%m-%d %H:%M:%S"),
+        "level": record["level"].name,
+        **extras,
+        "message": record["message"],
+    }
+    console.print(JSON(json.dumps(payload), indent=None))
+
+
+logger.remove()
+logger.add(sink=sink, level=LOG_LEVEL, enqueue=True, backtrace=True, diagnose=False)
+logger.bind(component="logging", level=LOG_LEVEL).info("configured loguru sink")

json2vec/logging/epoch.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from functools import partialmethod
+from typing import TYPE_CHECKING, Literal
+
+from lightning import Callback, Trainer
+from loguru import logger
+
+from json2vec.structs.enums import Strata
+
+if TYPE_CHECKING:
+    from json2vec.architecture.root import JSON2Vec
+
+
+class EpochLifecycleLogger(Callback):
+    def info(
+        self,
+        trainer: Trainer,
+        pl_module: JSON2Vec,
+        strata: Strata,
+        hook: Literal["start", "end"],
+    ):
+        logger.bind(
+            source="lightning",
+            rank=pl_module.global_rank,
+            epoch=pl_module.current_epoch,
+            step=pl_module.global_step,
+            hook=hook,
+            strata=str(strata),
+        ).info(f"{hook}ing {strata} epoch {pl_module.current_epoch}")
+
+    on_train_epoch_start = partialmethod(info, strata=Strata.train, hook="start")
+    on_train_epoch_end = partialmethod(info, strata=Strata.train, hook="end")
+
+    on_validation_epoch_start = partialmethod(info, strata=Strata.validate, hook="start")
+    on_validation_epoch_end = partialmethod(info, strata=Strata.validate, hook="end")
+
+    on_test_epoch_start = partialmethod(info, strata=Strata.test, hook="start")
+    on_test_epoch_end = partialmethod(info, strata=Strata.test, hook="end")
+
+    on_predict_epoch_start = partialmethod(info, strata=Strata.predict, hook="start")
+    on_predict_epoch_end = partialmethod(info, strata=Strata.predict, hook="end")