unaiverse.modules.utils

What this module does 🔴

Shared infrastructure for the modules package: the ModuleWrapper base class that all networks extend, processor/module helpers (HumanModule, LoggerModule, MultiIdentity), agent processor validation, image/transform factories, training utilities, and an API gateway server/scheduler with Featherless integration.

utils

█████ █████ ██████ █████ █████ █████ █████ ██████████ ███████████ █████████ ██████████ ░░███ ░░███ ░░██████ ░░███ ░░███ ░░███ ░░███ ░░███░░░░░█░░███░░░░░███ ███░░░░░███░░███░░░░░█ ░███ ░███ ░███░███ ░███ ██████ ░███ ░███ ░███ ░███ █ ░ ░███ ░███ ░███ ░░░ ░███ █ ░ ░███ ░███ ░███░░███░███ ░░░░░███ ░███ ░███ ░███ ░██████ ░██████████ ░░█████████ ░██████
░███ ░███ ░███ ░░██████ ███████ ░███ ░░███ ███ ░███░░█ ░███░░░░░███ ░░░░░░░░███ ░███░░█
░███ ░███ ░███ ░░█████ ███░░███ ░███ ░░░█████░ ░███ ░ █ ░███ ░███ ███ ░███ ░███ ░ █ ░░████████ █████ ░░█████░░████████ █████ ░░███ ██████████ █████ █████░░█████████ ██████████ ░░░░░░░░ ░░░░░ ░░░░░ ░░░░░░░░ ░░░░░ ░░░ ░░░░░░░░░░ ░░░░░ ░░░░░ ░░░░░░░░░ ░░░░░░░░░░ A Collectionless AI Project (https://collectionless.ai) Registration/Login: https://unaiverse.io Code Repositories: https://github.com/collectionlessai/ Main Developers: Stefano Melacci (Project Leader), Christian Di Maio, Tommaso Guidi

OPENAI_NATIVE_SAMPLER_PARAMS module-attribute

OPENAI_NATIVE_SAMPLER_PARAMS: frozenset[str] = frozenset({'max_tokens', 'temperature', 'top_p', 'frequency_penalty', 'presence_penalty', 'stop', 'seed', 'n', 'logit_bias', 'logprobs', 'top_logprobs', 'response_format'})

MultiIdentity

MultiIdentity()

Bases: Module

Identity module that passes one or more inputs through unchanged.

Used by the framework as the default processor when no explicit module is supplied. It accepts any number of positional arguments and returns either the single input unchanged (when called with exactly one argument) or the full argument tuple (when called with more than one argument).

Examples:

>>> m = MultiIdentity()
>>> import torch
>>> t = torch.tensor([1.0, 2.0])
>>> m(t) is t  # single input: returned as-is
True
>>> out = m(t, t)  # multiple inputs: returned as a tuple
>>> len(out)
2

Initialize the MultiIdentity module with no learnable parameters.

Source code in unaiverse/modules/utils.py

def __init__(self) -> None:
    """Initialize the MultiIdentity module with no learnable parameters."""
    super().__init__()

forward

forward(*args: object, **kwargs: object)

Return the single input unchanged, or all positional inputs as a tuple.

Parameters:

Name	Type	Description	Default
*args	object	One or more input tensors or arbitrary objects to pass through.	()
**kwargs	object	Accepted but ignored; present for interface compatibility.	{}

Returns:

Type	Description
	The single element args[0] if exactly one positional argument is
	provided, otherwise the full args tuple.

Source code in unaiverse/modules/utils.py

def forward(self, *args: object, **kwargs: object):
    """Return the single input unchanged, or all positional inputs as a tuple.

    Args:
        *args: One or more input tensors or arbitrary objects to pass through.
        **kwargs: Accepted but ignored; present for interface compatibility.

    Returns:
        The single element ``args[0]`` if exactly one positional argument is
        provided, otherwise the full ``args`` tuple.
    """
    if len(args) == 1:
        return args[0]
    return args

HumanModule

HumanModule()

Bases: Module

Placeholder module representing a human-in-the-loop processor.

When an agent is configured with proc="human", the framework wraps this module in a ModuleWrapper and uses it as the processor. It has no learnable parameters and simply echoes its text and image inputs back unchanged, acting as a neutral stand-in until the real human interaction layer overrides the output at a higher level.

See has_human_processor to check whether an agent uses this module.

Initialize the HumanModule with no learnable parameters.

Source code in unaiverse/modules/utils.py

def __init__(self) -> None:
    """Initialize the HumanModule with no learnable parameters."""
    super().__init__()

forward

forward(text: str | None = None, img: Image | None = None, whatever: object | None = None) -> tuple[str | None, Image | None]

Return the text and image inputs unchanged.

Acts as a transparent pass-through. The whatever argument is accepted for interface compatibility but is not forwarded to the output.

Parameters:

Name	Type	Description	Default
text	str | None	Optional text string to pass through. Defaults to None.	None
img	Image | None	Optional PIL image to pass through. Defaults to None.	None
whatever	object | None	Optional additional input; accepted but ignored. Defaults to None.	None

Returns:

Type	Description
tuple[str | None, Image | None]	A two-element tuple (text, img) containing the unchanged inputs.

Source code in unaiverse/modules/utils.py

def forward(self, text: str | None = None, img: Image.Image | None = None,
            whatever: object | None = None) -> tuple[str | None, Image.Image | None]:
    """Return the text and image inputs unchanged.

    Acts as a transparent pass-through. The ``whatever`` argument is accepted for
    interface compatibility but is not forwarded to the output.

    Args:
        text: Optional text string to pass through. Defaults to ``None``.
        img: Optional PIL image to pass through. Defaults to ``None``.
        whatever: Optional additional input; accepted but ignored. Defaults to ``None``.

    Returns:
        A two-element tuple ``(text, img)`` containing the unchanged inputs.
    """
    return text, img

LoggerModule

LoggerModule(log_file: str | None = None)

Bases: Module

A processor module that logs inputs and outputs to a file and returns cycling dummy responses.

LoggerModule is intended for debugging and development: it records every forward call (the received text, the image size if present, and the generated response) to an optional log file, while producing a deterministic but varied stream of dummy text outputs by cycling through an internally shuffled vocabulary of words. No image is ever returned.

The file handler is created lazily on the first forward call, so the log file is only opened when the module is actually invoked.

Examples:

>>> import tempfile, os
>>> with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f:
...     log_path = f.name
>>> m = LoggerModule(log_file=log_path)
>>> response_text, response_img = m("hello")
>>> response_img is None
True
>>> os.unlink(log_path)

Initialize the LoggerModule.

Sets up the internal cycling vocabulary (randomly shuffled) and, if a log file path is given, prepares the logging infrastructure (the file handler is opened lazily on the first forward call).

Parameters:

Name	Type	Description	Default
log_file	str | None	Path to the log file that records inputs and outputs. The parent directory is created automatically if it does not exist. If None, no file logging is performed. Defaults to None.	None

Source code in unaiverse/modules/utils.py

def __init__(self, log_file: str | None = None) -> None:
    """Initialize the LoggerModule.

    Sets up the internal cycling vocabulary (randomly shuffled) and, if a log file
    path is given, prepares the logging infrastructure (the file handler is opened
    lazily on the first ``forward`` call).

    Args:
        log_file: Path to the log file that records inputs and outputs. The parent
            directory is created automatically if it does not exist. If ``None``,
            no file logging is performed. Defaults to ``None``.
    """
    super().__init__()
    self.log_file = log_file
    self._initialized = False
    self._logger = logging.getLogger("CallableLogger")
    self._logger.setLevel(logging.INFO)
    self.__handler = None
    self._idx = 0
    self._objects = ["telescope", "hammer", "compass", "anchor", "lantern", "keyboard", "cat", "dog", "tiger",
                     "zebra", "batman", "superman", "candy", "table", "chair", "balloon", "kitchen", "sofa", "lamp",
                     "arrow", "green", "red", "blue", "yellow", "magenta", "brown", "pink", "orange", "white",
                     "paris", "rome", "boston", "york", "berlin", "singapore", "taiwan", "japan", "china",
                     "turkey", "italy", "france", "germany", "spain", "madrid", "barcelona", "portugal",
                     "norway", "sweden", "belgium", "romania", "sunny", "snowy", "rainy"]
    random.shuffle(self._objects)

log_file instance-attribute

log_file = log_file

forward

forward(text: str, img: Image | None = None) -> tuple[str | None, Image | None]

Log the received inputs and return a cycling dummy text response.

On each call the received text and the size of img (or None if no image was provided) are written to the log file (if configured). The method then returns the next word from the internal cycling vocabulary as the response text, advancing the internal index modulo the vocabulary size. The image output is always None.

If the log file path was provided at construction time but the file handler has not yet been opened, it is initialized during this call (the parent directory is created if necessary).

Parameters:

Name	Type	Description	Default
text	str	The text input to log and acknowledge.	required
img	Image | None	An optional PIL image whose size (img.size) is recorded in the log. Not forwarded to the output. Defaults to None.	None

Returns:

Type	Description
str | None	A two-element tuple (response_text, None) where response_text is a
Image | None	string of the form "<index>_<word>" taken from the internal vocabulary
tuple[str | None, Image | None]	cycle.

Source code in unaiverse/modules/utils.py

def forward(self, text: str, img: Image.Image | None = None) -> tuple[str | None, Image.Image | None]:
    """Log the received inputs and return a cycling dummy text response.

    On each call the received ``text`` and the size of ``img`` (or ``None`` if no
    image was provided) are written to the log file (if configured). The method then
    returns the next word from the internal cycling vocabulary as the response text,
    advancing the internal index modulo the vocabulary size. The image output is
    always ``None``.

    If the log file path was provided at construction time but the file handler has
    not yet been opened, it is initialized during this call (the parent directory is
    created if necessary).

    Args:
        text: The text input to log and acknowledge.
        img: An optional PIL image whose size (``img.size``) is recorded in the log.
            Not forwarded to the output. Defaults to ``None``.

    Returns:
        A two-element tuple ``(response_text, None)`` where ``response_text`` is a
        string of the form ``"<index>_<word>"`` taken from the internal vocabulary
        cycle.
    """
    if self.log_file is not None:
        if not self._initialized:
            self.__setup_logger()
        self._logger.info("-------------------------------------------------------------------------------")
        self._logger.info(f"[INPUT] text={text if text is not None else None}, "
                          f"img={img.size if img is not None else None}")

    text = f"{self._idx}_{self._objects[self._idx]}"
    self._idx = (self._idx + 1) % len(self._objects)
    img = None

    if self.log_file is not None:
        self._logger.info(f"[OUTPUT] text={text}, img={None}")
        self._logger.info("-------------------------------------------------------------------------------")
        self.__handler.flush()
    return text, img

ModuleWrapper

ModuleWrapper(module: Module | Callable, proc_inputs: list[StreamType] | None = None, proc_outputs: list[StreamType] | None = None, proc_opts: dict | None = None, seed: int = -1, agent: object = None, device=None)

Bases: Module

Wrap a callable or torch.nn.Module with stream-based pre/post-processing and optional learning.

ModuleWrapper is the core adapter between the UNaIVERSE streaming pipeline and a bare PyTorch module or any callable. It holds the StreamType descriptors for each input and output slot, applies check_and_preprocess on incoming data and check_and_postprocess on outgoing data in forward, and can run a full supervised or unsupervised learning step (optimizer + backward pass) via learn_backward.

The framework-controlled keyword arguments first and last are intercepted inside forward and stripped before the wrapped module is called, unless the module's own signature declares them.

The device is resolved in the following priority order: the device argument (if given), then the PROC_DEVICE environment variable, and finally "cpu".

Attributes:

Name	Type	Description
device		The torch.device on which the wrapped module runs.
module	Callable | None	The wrapped callable or torch.nn.Module.
proc_inputs		List of StreamType objects describing the processor inputs.
proc_outputs		List of StreamType objects describing the processor outputs.
proc_opts		Dictionary containing at least "optimizer" and "losses" keys.
proc_optional_inputs		Per-input default info derived from the module's signature.
agent		The agent that owns this wrapper.

Examples:

>>> import torch
>>> linear = torch.nn.Linear(4, 2)
>>> from unaiverse.streams.dataprops import StreamType
>>> proc_in = [StreamType(data_type="tensor",
...                       tensor_shape=(None, 4),
...                       tensor_dtype=torch.float32,
...                       pubsub=False, private_only=False)]
>>> proc_out = [StreamType(data_type="tensor",
...                        tensor_shape=(None, 2),
...                        tensor_dtype=torch.float32,
...                        pubsub=False, private_only=False)]
>>> wrapper = ModuleWrapper(module=linear, proc_inputs=proc_in, proc_outputs=proc_out)
>>> x = torch.randn(1, 4)
>>> out = wrapper(x)
>>> out[0].shape
torch.Size([1, 2])

Initialize the ModuleWrapper.

Resolves the target device, wraps the given module (moving it to the device if it is a torch.nn.Module), inspects the module's forward signature to determine which framework kwargs (first, last) it accepts and to cache per-argument default information, then sets the stream descriptors and optimization options.

Parameters:

Name	Type	Description	Default
module	Module | Callable	The torch.nn.Module or callable to wrap. If a torch.nn.Module, it is moved to the resolved device.	required
proc_inputs	list[StreamType] | None	List of StreamType objects describing the input streams. Defaults to None.	None
proc_outputs	list[StreamType] | None	List of StreamType objects describing the output streams. Defaults to None.	None
proc_opts	dict | None	Dictionary with at least the keys "optimizer" (a torch.optim.Optimizer or None) and "losses" (a list of loss callables). Defaults to None.	None
seed	int	Random seed applied via set_seed before the module is wrapped. Negative values disable seeding. Defaults to -1.	-1
agent	object	The agent that owns this processor; stored as self.agent. Defaults to None.	None
device		The torch.device or device string on which to run the module. If None, the PROC_DEVICE environment variable is consulted; falling back to "cpu" if neither is set. Defaults to None.	None

Raises:

Type	Description
GenException	If the module does not expose a callable forward method and is not itself callable.

Source code in unaiverse/modules/utils.py

def __init__(self,
             module: torch.nn.Module | Callable,
             proc_inputs: list[StreamType] | None = None,
             proc_outputs: list[StreamType] | None = None,
             proc_opts: dict | None = None,
             seed: int = -1,
             agent: object = None,
             device=None) -> None:
    """Initialize the ModuleWrapper.

    Resolves the target device, wraps the given module (moving it to the device if
    it is a ``torch.nn.Module``), inspects the module's forward signature to
    determine which framework kwargs (``first``, ``last``) it accepts and to cache
    per-argument default information, then sets the stream descriptors and
    optimization options.

    Args:
        module: The ``torch.nn.Module`` or callable to wrap. If a
            ``torch.nn.Module``, it is moved to the resolved device.
        proc_inputs: List of ``StreamType`` objects describing the input streams.
            Defaults to ``None``.
        proc_outputs: List of ``StreamType`` objects describing the output streams.
            Defaults to ``None``.
        proc_opts: Dictionary with at least the keys ``"optimizer"`` (a
            ``torch.optim.Optimizer`` or ``None``) and ``"losses"`` (a list of loss
            callables). Defaults to ``None``.
        seed: Random seed applied via ``set_seed`` before the module is wrapped.
            Negative values disable seeding. Defaults to ``-1``.
        agent: The agent that owns this processor; stored as ``self.agent``.
            Defaults to ``None``.
        device: The ``torch.device`` or device string on which to run the module.
            If ``None``, the ``PROC_DEVICE`` environment variable is consulted;
            falling back to ``"cpu"`` if neither is set. Defaults to ``None``.

    Raises:
        GenException: If the module does not expose a callable ``forward`` method
            and is not itself callable.
    """
    super(ModuleWrapper, self).__init__()
    self.device = None
    self.module: Callable | None = None
    self.proc_inputs = None  # The list of Stream objects describing the input types of the module
    self.proc_outputs = None  # The list of Stream objects describing the output types of the module
    self.proc_opts = None
    self.proc_optional_inputs = []
    self.agent = None
    self._has_first = False
    self._has_last = False
    self._sig_defaults = []
    self._last_raw_outputs = None

    # Working
    set_seed(seed)

    self.guess_device(device)
    self.set_module(module)
    self.set_proc_inputs_and_outputs(proc_inputs, proc_outputs)
    self.set_proc_opts(proc_opts)
    self.set_agent(agent)

device instance-attribute

device = None

module instance-attribute

module: Callable | None = None

proc_inputs instance-attribute

proc_inputs = None

proc_outputs instance-attribute

proc_outputs = None

proc_opts instance-attribute

proc_opts = None

proc_optional_inputs instance-attribute

proc_optional_inputs = []

agent instance-attribute

agent = None

guess_device

guess_device(device)

Resolve and set the compute device for the wrapped module.

Determines self.device using the following priority order: the device argument (if not None), then the PROC_DEVICE environment variable, and finally "cpu" as the default fallback.

Parameters:

Name	Type	Description	Default
device		A torch.device, a device string such as "cuda:0", or None to defer to the PROC_DEVICE environment variable. If None and PROC_DEVICE is also unset, "cpu" is used.	required

Source code in unaiverse/modules/utils.py

def guess_device(self, device):
    """Resolve and set the compute device for the wrapped module.

    Determines ``self.device`` using the following priority order: the ``device``
    argument (if not ``None``), then the ``PROC_DEVICE`` environment variable, and
    finally ``"cpu"`` as the default fallback.

    Args:
        device: A ``torch.device``, a device string such as ``"cuda:0"``, or ``None``
            to defer to the ``PROC_DEVICE`` environment variable. If ``None`` and
            ``PROC_DEVICE`` is also unset, ``"cpu"`` is used.
    """
    device_env = os.getenv("PROC_DEVICE", None)
    self.device = torch.device("cpu") if device_env is None else torch.device(device_env)

    # Override if it was given
    if device is not None:
        self.device = torch.device(device) if isinstance(device, str) else device

set_agent

set_agent(agent)

Set the agent that owns this processor wrapper.

Stores a reference to the owning agent so that downstream components (such as learn_backward and AgentProcessorChecker) can reach back to the agent context when needed.

Parameters:

Name	Type	Description	Default
agent		The agent object that owns this ModuleWrapper. May be None when no owning agent has been assigned yet.	required

Source code in unaiverse/modules/utils.py

def set_agent(self, agent):
    """Set the agent that owns this processor wrapper.

    Stores a reference to the owning agent so that downstream components (such as
    ``learn_backward`` and ``AgentProcessorChecker``) can reach back to the agent
    context when needed.

    Args:
        agent: The agent object that owns this ``ModuleWrapper``. May be ``None``
            when no owning agent has been assigned yet.
    """
    self.agent = agent

set_proc_inputs_and_outputs

set_proc_inputs_and_outputs(proc_inputs, proc_outputs)

Update the input and/or output stream descriptors.

When proc_inputs is not None, replaces self.proc_inputs with the provided list and rebuilds proc_optional_inputs so it stays aligned with the new input count (see __refresh_proc_optional_inputs). When proc_outputs is not None, replaces self.proc_outputs. Either argument may be None to leave the corresponding attribute unchanged.

Parameters:

Name	Type	Description	Default
proc_inputs		A list of StreamType objects describing the new input streams, or None to leave self.proc_inputs unchanged.	required
proc_outputs		A list of StreamType objects describing the new output streams, or None to leave self.proc_outputs unchanged.	required

Source code in unaiverse/modules/utils.py

def set_proc_inputs_and_outputs(self, proc_inputs, proc_outputs):
    """Update the input and/or output stream descriptors.

    When ``proc_inputs`` is not ``None``, replaces ``self.proc_inputs`` with the
    provided list and rebuilds ``proc_optional_inputs`` so it stays aligned with the
    new input count (see ``__refresh_proc_optional_inputs``). When ``proc_outputs``
    is not ``None``, replaces ``self.proc_outputs``. Either argument may be ``None``
    to leave the corresponding attribute unchanged.

    Args:
        proc_inputs: A list of ``StreamType`` objects describing the new input
            streams, or ``None`` to leave ``self.proc_inputs`` unchanged.
        proc_outputs: A list of ``StreamType`` objects describing the new output
            streams, or ``None`` to leave ``self.proc_outputs`` unchanged.
    """
    if proc_inputs is not None:
        self.proc_inputs = proc_inputs
        self.__refresh_proc_optional_inputs()  # rebuilds, truncating or padding as needed
    if proc_outputs is not None:
        self.proc_outputs = proc_outputs

set_proc_opts

set_proc_opts(proc_opts)

Update the optimization options dictionary.

Replaces self.proc_opts with the provided dictionary. The call is a no-op when proc_opts is None, preserving any previously assigned options.

The dictionary is expected to contain at least the keys "optimizer" (a torch.optim.Optimizer or None) and "losses" (a list of loss callables, one per output slot). See learn_backward for how these entries are consumed.

Parameters:

Name	Type	Description	Default
proc_opts		A dictionary of optimization options, or None to leave self.proc_opts unchanged.	required

Source code in unaiverse/modules/utils.py

def set_proc_opts(self, proc_opts):
    """Update the optimization options dictionary.

    Replaces ``self.proc_opts`` with the provided dictionary. The call is a no-op
    when ``proc_opts`` is ``None``, preserving any previously assigned options.

    The dictionary is expected to contain at least the keys ``"optimizer"`` (a
    ``torch.optim.Optimizer`` or ``None``) and ``"losses"`` (a list of loss
    callables, one per output slot). See ``learn_backward`` for how these entries
    are consumed.

    Args:
        proc_opts: A dictionary of optimization options, or ``None`` to leave
            ``self.proc_opts`` unchanged.
    """
    if proc_opts is not None:
        self.proc_opts = proc_opts

set_module

set_module(mod: Callable | Module)

Set the wrapped module and cache its forward-signature metadata.

Inspects the callable's signature to determine whether the framework-controlled keyword arguments first and last are declared (stored in _has_first and _has_last respectively). The per-parameter default information is cached in _sig_defaults so that proc_optional_inputs can be rebuilt efficiently whenever proc_inputs changes.

If mod is a torch.nn.Module it is moved to self.device via .to(). If mod has a forward method it is inspected via that method's signature; otherwise mod itself must be directly callable.

Parameters:

Name	Type	Description	Default
mod	Callable | Module	A torch.nn.Module or any callable to wrap. The object must expose either a callable forward attribute or be callable itself.	required

Raises:

Type	Description
GenException	If mod neither has a callable forward attribute nor is itself callable, making it impossible to inspect the call signature.

Source code in unaiverse/modules/utils.py

def set_module(self, mod: Callable | torch.nn.Module):
    """Set the wrapped module and cache its forward-signature metadata.

    Inspects the callable's signature to determine whether the framework-controlled
    keyword arguments ``first`` and ``last`` are declared (stored in ``_has_first``
    and ``_has_last`` respectively). The per-parameter default information is cached
    in ``_sig_defaults`` so that ``proc_optional_inputs`` can be rebuilt efficiently
    whenever ``proc_inputs`` changes.

    If ``mod`` is a ``torch.nn.Module`` it is moved to ``self.device`` via ``.to()``.
    If ``mod`` has a ``forward`` method it is inspected via that method's signature;
    otherwise ``mod`` itself must be directly callable.

    Args:
        mod: A ``torch.nn.Module`` or any callable to wrap. The object must expose
            either a callable ``forward`` attribute or be callable itself.

    Raises:
        GenException: If ``mod`` neither has a callable ``forward`` attribute nor is
            itself callable, making it impossible to inspect the call signature.
    """
    # Case 1: torch.nn.Module / object with .forward
    if hasattr(mod, "forward") and callable(getattr(mod, "forward")):
        sig = inspect.signature(mod.forward)
    elif callable(mod):
        sig = inspect.signature(mod)
    else:
        raise GenException("Invalid processor module, cannot guess the signature of the callable method, if any")

    self.module = mod.to(self.device) if isinstance(mod, torch.nn.Module) else mod
    self._has_first = "first" in sig.parameters  # B8: rename below
    self._has_last = "last" in sig.parameters

    # Cache the signature-derived defaults ONCE, excluding framework-handled kwargs and *args/**kwargs
    # catch-alls. This is the canonical per-positional-data-arg default list of the wrapped module's forward.
    framework_kwargs = {"first", "last"}
    self._sig_defaults: list[dict] = []
    for name, param in sig.parameters.items():
        if name in framework_kwargs:
            continue
        if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
            continue
        if param.default is not inspect.Parameter.empty:
            self._sig_defaults.append({"has_default": True, "default_value": param.default})
        else:
            self._sig_defaults.append({"has_default": False, "default_value": None})

    # Align proc_optional_inputs to current proc_inputs (if already known)
    self.__refresh_proc_optional_inputs()

forward

forward(*args, **kwargs) -> tuple

Preprocess inputs, run the wrapped module, and post-process the outputs.

Each positional argument is preprocessed by the corresponding StreamType entry in proc_inputs (via check_and_preprocess), then the wrapped module is called with the preprocessed data. The raw outputs are stored in _last_raw_outputs (used later by learn_backward) and then passed through the corresponding proc_outputs post-processing step (check_and_postprocess) before being returned.

The framework-controlled keyword arguments "first" and "last" are intercepted and stripped unless the wrapped module's own signature declares them (see _has_first and _has_last).

Parameters:

Name	Type	Description	Default
*args		Positional inputs, one per entry in proc_inputs. Each element is preprocessed by the corresponding StreamType descriptor and moved to self.device.	()
**kwargs		Additional keyword arguments forwarded to the wrapped module. The keys "first" and "last" are silently removed when the wrapped module does not declare them.	{}

Returns:

Type	Description
tuple	A tuple of post-processed output values, one per entry in proc_outputs.

Raises:

Type	Description
AssertionError	If self.module, self.proc_inputs, or self.proc_outputs is None (i.e., the wrapper has not been fully configured before calling forward).

Examples:

>>> import torch
>>> from unaiverse.streams.dataprops import StreamType
>>> linear = torch.nn.Linear(4, 2)
>>> proc_in = [StreamType(data_type="tensor", tensor_shape=(None, 4),
...                       tensor_dtype=torch.float32,
...                       pubsub=False, private_only=False)]
>>> proc_out = [StreamType(data_type="tensor", tensor_shape=(None, 2),
...                        tensor_dtype=torch.float32,
...                        pubsub=False, private_only=False)]
>>> wrapper = ModuleWrapper(module=linear, proc_inputs=proc_in, proc_outputs=proc_out)
>>> x = torch.randn(1, 4)
>>> out = wrapper(x)
>>> out[0].shape
torch.Size([1, 2])

Source code in unaiverse/modules/utils.py

def forward(self, *args, **kwargs) -> tuple:
    """Preprocess inputs, run the wrapped module, and post-process the outputs.

    Each positional argument is preprocessed by the corresponding ``StreamType``
    entry in ``proc_inputs`` (via ``check_and_preprocess``), then the wrapped module
    is called with the preprocessed data. The raw outputs are stored in
    ``_last_raw_outputs`` (used later by ``learn_backward``) and then passed through
    the corresponding ``proc_outputs`` post-processing step
    (``check_and_postprocess``) before being returned.

    The framework-controlled keyword arguments ``"first"`` and ``"last"`` are
    intercepted and stripped unless the wrapped module's own signature declares them
    (see ``_has_first`` and ``_has_last``).

    Args:
        *args: Positional inputs, one per entry in ``proc_inputs``. Each element is
            preprocessed by the corresponding ``StreamType`` descriptor and moved to
            ``self.device``.
        **kwargs: Additional keyword arguments forwarded to the wrapped module. The
            keys ``"first"`` and ``"last"`` are silently removed when the wrapped
            module does not declare them.

    Returns:
        A tuple of post-processed output values, one per entry in ``proc_outputs``.

    Raises:
        AssertionError: If ``self.module``, ``self.proc_inputs``, or
            ``self.proc_outputs`` is ``None`` (i.e., the wrapper has not been fully
            configured before calling ``forward``).

    Examples:
        >>> import torch
        >>> from unaiverse.streams.dataprops import StreamType
        >>> linear = torch.nn.Linear(4, 2)
        >>> proc_in = [StreamType(data_type="tensor", tensor_shape=(None, 4),
        ...                       tensor_dtype=torch.float32,
        ...                       pubsub=False, private_only=False)]
        >>> proc_out = [StreamType(data_type="tensor", tensor_shape=(None, 2),
        ...                        tensor_dtype=torch.float32,
        ...                        pubsub=False, private_only=False)]
        >>> wrapper = ModuleWrapper(module=linear, proc_inputs=proc_in, proc_outputs=proc_out)
        >>> x = torch.randn(1, 4)
        >>> out = wrapper(x)
        >>> out[0].shape
        torch.Size([1, 2])
    """
    assert self.module is not None
    assert self.proc_inputs is not None
    assert self.proc_outputs is not None

    # The forward signature expected by who calls this method is:
    # forward(self, *args, first: bool, last: bool, **kwargs)
    # so we have to discard 'first' and 'last' that are not used by an external module not designed for this library
    if 'first' in kwargs and not self._has_first:
        del kwargs['first']
    if 'last' in kwargs and not self._has_last:
        del kwargs['last']
    self._last_raw_outputs = None

    # Preprocessing data
    args = [self.proc_inputs[i].check_and_preprocess(args[i], device=self.device)
            for i in range(0, len(self.proc_inputs))]  # Don't try to build a tuple here, keep a list!

    # Calling the module
    outputs = self.module(*args, **kwargs)  # Forward

    if not isinstance(outputs, tuple):
        outputs = (outputs,)
    self._last_raw_outputs = outputs

    # Postprocessing data (this does not affect the raw outputs, that might be used while learning)
    outputs = [self.proc_outputs[i].check_and_postprocess(outputs[i])
               for i in range(0, len(self.proc_outputs))]  # Don't try to build a tuple here, keep a list!

    return tuple(outputs)

learn_backward

learn_backward(targets: list | None = None) -> list

Run a supervised or unsupervised backward pass and apply an optimizer step.

Uses the raw outputs stored by the most recent forward call (_last_raw_outputs) together with the provided targets to compute the per-slot losses defined in proc_opts["losses"], sums them into a scalar, calls loss.backward(), and steps the optimizer from proc_opts["optimizer"].

When targets is None or all entries are None, the method runs each loss function with only the corresponding raw output as argument (unsupervised mode). When at least one target is provided, each slot whose target is not None is treated as supervised (the target is preprocessed via proc_outputs[i].check_and_preprocess); slots with a None target contribute a zero loss to the sum.

If the wrapped module has a y attribute and targets is not None, the first target tensor is written to module.y after the optimizer step (autoregressive-teacher-forcing support).

The method is a no-op and returns [] when proc_opts is absent or does not contain both "losses" and "optimizer" keys, or when proc_outputs is None.

Parameters:

Name	Type	Description	Default
targets	list | None	A list of target tensors aligned with proc_outputs, or None for unsupervised learning. Individual list entries may be None to mark that slot as unsupervised even in an otherwise supervised call. Defaults to None.	None

Returns:

Type	Description
list	A list of per-output torch.Tensor loss values (one per output slot)
list	if a backward pass was performed, or an empty list when the method is a
list	no-op (no optimizer configured).

Source code in unaiverse/modules/utils.py

def learn_backward(self, targets: list | None = None) -> list:
    """Run a supervised or unsupervised backward pass and apply an optimizer step.

    Uses the raw outputs stored by the most recent ``forward`` call
    (``_last_raw_outputs``) together with the provided ``targets`` to compute the
    per-slot losses defined in ``proc_opts["losses"]``, sums them into a scalar,
    calls ``loss.backward()``, and steps the optimizer from
    ``proc_opts["optimizer"]``.

    When ``targets`` is ``None`` or all entries are ``None``, the method runs each
    loss function with only the corresponding raw output as argument (unsupervised
    mode). When at least one target is provided, each slot whose target is not
    ``None`` is treated as supervised (the target is preprocessed via
    ``proc_outputs[i].check_and_preprocess``); slots with a ``None`` target
    contribute a zero loss to the sum.

    If the wrapped module has a ``y`` attribute and ``targets`` is not ``None``,
    the first target tensor is written to ``module.y`` after the optimizer step
    (autoregressive-teacher-forcing support).

    The method is a no-op and returns ``[]`` when ``proc_opts`` is absent or does
    not contain both ``"losses"`` and ``"optimizer"`` keys, or when ``proc_outputs``
    is ``None``.

    Args:
        targets: A list of target tensors aligned with ``proc_outputs``, or
            ``None`` for unsupervised learning. Individual list entries may be
            ``None`` to mark that slot as unsupervised even in an otherwise
            supervised call. Defaults to ``None``.

    Returns:
        A list of per-output ``torch.Tensor`` loss values (one per output slot)
        if a backward pass was performed, or an empty list when the method is a
        no-op (no optimizer configured).
    """
    if (self.proc_opts is None or len(self.proc_opts) == 0 or
            ('losses' not in self.proc_opts and 'optimizer' not in self.proc_opts) or self.proc_outputs is None):
        return []

    loss_functions: list = self.proc_opts['losses']
    optimizer: torch.optim.optimizer.Optimizer = self.proc_opts['optimizer']

    # Evaluating loss function(s), one for each processor output slot (they are set to 0. if no targets are there)
    if targets is not None and len(targets) > 0 and any(x is not None for x in targets):

        # Preprocessing targets
        targets = [self.proc_outputs[i].check_and_preprocess(targets[i], device=self.device,
                                                             allow_class_ids=True, targets=True)
                   for i in range(0, len(self.proc_outputs))]

        # Supervised or partly supervised learning
        loss_values = [loss_fcn(self._last_raw_outputs[i], targets[i]) if targets[i] is not None else
                       torch.tensor(0., device=self.device)
                       for i, loss_fcn in enumerate(loss_functions)]
        loss = torch.stack(loss_values).sum()  # Sum of losses
    else:

        # Unsupervised learning
        loss_values = [loss_fcn(self._last_raw_outputs[i]) for i, loss_fcn in enumerate(loss_functions)]
        loss = torch.stack(loss_values).sum()  # Sum of losses

    # Learning step
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Teaching (for autoregressive models, expected to have attribute "y")
    if targets is not None and hasattr(self.module, 'y'):
        self.module.y = targets[0]

    return loss_values

AgentProcessorChecker

AgentProcessorChecker(agent: object)

Validate and normalize the processor-related attributes of an agent-like container.

AgentProcessorChecker is a one-shot validation object: constructing it immediately validates the processor configuration of the supplied agent, performs all necessary type conversions and heuristic inference, and writes the normalized values back onto the agent. No state is kept after construction.

The checker enforces the following invariants after it completes:

• agent.proc is a fully configured ModuleWrapper (never None, never a raw torch.nn.Module).
• agent.proc_inputs and agent.proc_outputs are non-empty lists of StreamType objects with canonical names ("proc_input_N" / "proc_output_N").
• agent.proc_opts contains both "optimizer" and "losses" keys.
• agent.proc_optional_inputs is aligned with agent.proc_inputs.

The special string value "human" for proc is recognized and converted into a ModuleWrapper around a HumanModule with text and image stream types. When proc is None, a ModuleWrapper around MultiIdentity is created.

See ModuleWrapper, HumanModule, and MultiIdentity for details on the wrapped modules.

Validate and normalize the processor attributes of the given agent.

Checks that agent exposes the required processor attributes, performs type validation, converts string StreamType shorthands to proper StreamType objects, auto-wraps plain torch.nn.Module or callable processors in ModuleWrapper, and heuristically infers missing proc_inputs, proc_outputs, and proc_opts. The validated and completed values are written back onto agent so callers always receive a fully configured processor on agent.proc.

The heuristic inference pipeline runs in this order:

1. __guess_proc_inputs - infers input StreamType from the first layer of the wrapped torch.nn.Module (Conv2d, Linear, Conv1d, or Embedding).
2. __fix_proc_inputs - assigns canonical names to inputs.
3. __guess_proc_outputs - runs a dummy forward pass to infer output shapes.
4. __fix_proc_outputs - assigns canonical names to outputs.
5. __guess_proc_opts - fills in missing "optimizer" and "losses" entries.
6. __fix_proc_opts - normalizes the options dict and replaces cross_entropy with target_shape_fixed_cross_entropy.

Parameters:

Name	Type	Description	Default
agent	object	Any object that exposes the attributes proc, proc_inputs, proc_outputs, proc_opts, and proc_optional_inputs. In practice this is always an AgentBasics subclass instance.	required

Raises:

Type	Description
AssertionError	If agent is missing any of the required processor attributes.
GenException	If proc is not None, a torch.nn.Module, a ModuleWrapper, or a callable; if proc_inputs or proc_outputs are not None or valid lists of StreamType/str; if proc_opts is not None or a dict; if any string shorthand in proc_inputs or proc_outputs is not a recognized StreamType data type; or if proc_inputs / proc_outputs cannot be inferred automatically.

Source code in unaiverse/modules/utils.py

def __init__(self, agent: object) -> None:
    """Validate and normalize the processor attributes of the given agent.

    Checks that ``agent`` exposes the required processor attributes, performs type
    validation, converts string ``StreamType`` shorthands to proper ``StreamType``
    objects, auto-wraps plain ``torch.nn.Module`` or callable processors in
    ``ModuleWrapper``, and heuristically infers missing ``proc_inputs``,
    ``proc_outputs``, and ``proc_opts``. The validated and completed values are
    written back onto ``agent`` so callers always receive a fully configured
    processor on ``agent.proc``.

    The heuristic inference pipeline runs in this order:

    1. ``__guess_proc_inputs`` - infers input ``StreamType`` from the first layer of
       the wrapped ``torch.nn.Module`` (``Conv2d``, ``Linear``, ``Conv1d``, or
       ``Embedding``).
    2. ``__fix_proc_inputs`` - assigns canonical names to inputs.
    3. ``__guess_proc_outputs`` - runs a dummy forward pass to infer output shapes.
    4. ``__fix_proc_outputs`` - assigns canonical names to outputs.
    5. ``__guess_proc_opts`` - fills in missing ``"optimizer"`` and ``"losses"``
       entries.
    6. ``__fix_proc_opts`` - normalizes the options dict and replaces
       ``cross_entropy`` with ``target_shape_fixed_cross_entropy``.

    Args:
        agent: Any object that exposes the attributes ``proc``, ``proc_inputs``,
            ``proc_outputs``, ``proc_opts``, and ``proc_optional_inputs``. In
            practice this is always an ``AgentBasics`` subclass instance.

    Raises:
        AssertionError: If ``agent`` is missing any of the required processor
            attributes.
        GenException: If ``proc`` is not ``None``, a ``torch.nn.Module``, a
            ``ModuleWrapper``, or a callable; if ``proc_inputs`` or ``proc_outputs``
            are not ``None`` or valid lists of ``StreamType``/``str``; if
            ``proc_opts`` is not ``None`` or a ``dict``; if any string shorthand in
            ``proc_inputs`` or ``proc_outputs`` is not a recognized ``StreamType``
            data type; or if ``proc_inputs`` / ``proc_outputs`` cannot be inferred
            automatically.
    """

    # Right now, the processor has its own attributes, but they are also directly referenced from the agent object.
    # In the future, we might drop the ones on the agent side.
    assert hasattr(agent, 'proc'), "Invalid processor container (agent) object"
    assert hasattr(agent, 'proc_inputs'), "Invalid processor container (agent) object"
    assert hasattr(agent, 'proc_outputs'), "Invalid processor container (agent) object"
    assert hasattr(agent, 'proc_opts'), "Invalid processor container (agent) object"
    assert hasattr(agent, 'proc_optional_inputs'), "Invalid processor container (agent) object"

    # Getting processor-related info from the main object which collects processor and its properties
    proc: torch.nn.Module | ModuleWrapper | Callable | None = agent.proc
    proc_inputs: list[StreamType] | list[str] | None = agent.proc_inputs
    proc_outputs: list[StreamType] | list[str] | None = agent.proc_outputs
    proc_opts: dict | None = agent.proc_opts

    # The agent itself could be a processor
    if proc is None and callable(agent):
        proc = agent

    if not (proc is None or isinstance_fcn(proc, torch.nn.Module) or isinstance_fcn(proc, ModuleWrapper) or
            callable(proc)):
        raise GenException("Processor (proc) must be either None or a torch.nn.Module, "
                           "or a ModuleWrapper, or a callable object")
    if not ((proc_inputs is None or (
            isinstance_fcn(proc_inputs, list) and (len(proc_inputs) == 0 or
                                                   (len(proc_inputs) > 0 and
                                                    isinstance_fcn(proc_inputs[0],
                                                                   (StreamType, str))))))):
        raise GenException("Invalid proc_inputs: it must be None or a list of StreamType/str")
    if not ((proc_outputs is None or (
            isinstance_fcn(proc_outputs, list) and (len(proc_outputs) == 0 or
                                                    (len(proc_outputs) > 0 and
                                                     isinstance_fcn(proc_outputs[0],
                                                                    (StreamType, str))))))):
        raise GenException("Invalid proc_outputs: it must be None or a list of StreamType/str")
    if not (proc_opts is None or isinstance_fcn(proc_opts, dict)):
        raise GenException("Invalid proc_opts: it must be None or a dictionary")

    # From strings to StreamType (proc_inputs)
    if proc_inputs is not None:
        proc_inputs_copy = proc_inputs.copy()
        for i, p in enumerate(proc_inputs_copy):
            if isinstance(p, str):
                try:
                    proc_inputs_copy[i] = StreamType(data_type=p, pubsub=False, private_only=False)
                except Exception as e:
                    raise GenException(f"Invalid stream type {p}: {e}")
        proc_inputs = proc_inputs_copy

    # From strings to StreamType (proc_outputs)
    if proc_outputs is not None:
        proc_outputs_copy = proc_outputs.copy()
        for i, p in enumerate(proc_outputs_copy):
            if isinstance(p, str):
                try:
                    proc_outputs_copy[i] = StreamType(data_type=p, pubsub=False, private_only=False)
                except Exception as e:
                    raise GenException(f"Invalid stream type {p}: {e}")
        proc_outputs = proc_outputs_copy

    # From None to empty dictionary (proc_opts)
    if proc_opts is None:
        proc_opts = {}

    # Type checking
    if proc_inputs is not None:
        proc_inputs: list[StreamType]
    if proc_outputs is not None:
        proc_outputs: list[StreamType]

    # Dummy processor (if no processor was provided)
    if proc is None:
        if proc_inputs is None:
            proc_inputs: list[StreamType] = [StreamType(data_type="all", pubsub=False, private_only=False)]
        if proc_outputs is None:
            proc_outputs: list[StreamType] = [StreamType(data_type="all", pubsub=False, private_only=False)]
        assert proc_outputs is not None
        proc_opts = {'optimizer': None, 'losses': [None] * len(proc_outputs)}
        proc = ModuleWrapper(module=MultiIdentity(), proc_inputs=proc_inputs,
                             proc_outputs=proc_outputs, proc_opts=proc_opts,
                             agent=agent, device=torch.device("cpu"))
    else:

        # String telling it is a human converted to a ModuleWrapper with HumanModule processor
        if isinstance(proc, str) and proc.lower().strip() == "human":
            proc_inputs: list[StreamType] = [StreamType(data_type="text", pubsub=False, private_only=False),
                                             StreamType(data_type="img", pubsub=False, private_only=False)]
            proc_outputs: list[StreamType] = [StreamType(data_type="text", pubsub=False, private_only=False),
                                              StreamType(data_type="img", pubsub=False, private_only=False)]
            proc = ModuleWrapper(module=HumanModule(), proc_inputs=proc_inputs,
                                 proc_outputs=proc_outputs, agent=agent, device=torch.device("cpu"))

        # Creating a ModuleWrapper with the given processor
        elif not isinstance(proc, ModuleWrapper):
            proc: torch.Module | Callable
            proc = ModuleWrapper(module=proc, proc_opts=proc_opts, proc_inputs=proc_inputs,
                                 proc_outputs=proc_outputs, agent=agent, device=torch.device("cpu"))
        else:

            # If it is already a ModuleWrapper, we force the reference to the agent and the references to the
            # types of inputs and outputs: actually, the inputs/outpus/opts might have been already defined by the
            # wrapper creator, but we give priority to the user options
            proc.set_proc_inputs_and_outputs(proc_inputs, proc_outputs)
            proc.set_proc_opts(proc_opts)
            proc.set_agent(agent)

    # From here after, the proc is not None, and it is a ModuleWrapper.
    # However, proc_inputs, proc_outputs, and proc_opts might still not be specified
    assert proc is not None
    assert isinstance(proc, ModuleWrapper)

    # Guessing inputs, fixing attributes
    AgentProcessorChecker.__guess_proc_inputs(proc)
    if proc.proc_inputs is None:
        raise GenException("The inputs of the processor were not described and cannot be automatically guessed")

    # Fixing naming conventions
    AgentProcessorChecker.__fix_proc_inputs(proc)

    # Guessing outputs, fixing attributes
    AgentProcessorChecker.__guess_proc_outputs(proc)
    if proc.proc_outputs is None:
        raise GenException("The outputs of the processor were not described and cannot be automatically guessed")

    # Fixing naming conventions (should not be needed)
    AgentProcessorChecker.__fix_proc_outputs(proc)

    # Guessing optimization-related options and stuff, fixing attributes
    AgentProcessorChecker.__guess_proc_opts(proc)

    # Fixing attribute-presence conventions
    AgentProcessorChecker.__fix_proc_opts(proc)

    # Ensuring all is OK
    assert proc.proc_inputs is not None
    assert proc.proc_outputs is not None
    assert proc.proc_opts is not None
    assert proc.proc_optional_inputs is not None

    # Updating processor container object (agent)
    agent.proc = proc
    agent.proc_inputs = proc.proc_inputs
    agent.proc_outputs = proc.proc_outputs
    agent.proc_opts = proc.proc_opts
    agent.proc_optional_inputs = proc.proc_optional_inputs

APIGatewayServer

APIGatewayServer(call: Callable = featherless_call)

A shared LLM API gateway that owns the Featherless connection and runs the scheduler.

APIGatewayServer is a single-process asyncio TCP server that acts as a multiplexing proxy between multiple agent processes and the Featherless AI chat-completions API. It enforces a shared concurrency budget across all clients using the APIGatewayScheduler (strict round-robin, unit-based admission).

The server speaks a newline-delimited JSON protocol over TCP (HOST:PORT, localhost by default). The message contract is:

• Client registration (one persistent connection per client process, kept alive for the process lifetime)::

  client -> server: {"op": "hello"}

• Generation request (sent on a separate short-lived connection)::

  client -> server: {"op": "generate", "process_id": "", "sys_prompt": "", "prompt": "", "cost": , "model": "", "sampler": {...}} server -> client: {"ok": true, "result": ""} | {"ok": false, "error": ""}

The open registration socket is the client's liveness token. When a client process dies for any reason the OS closes the socket, the server detects the drop, and when the last registered client disconnects the server shuts itself down cleanly.

Startup arbitration is handled via the TCP bind: if the port is already taken, asyncio.start_server raises and the caller knows another server is already running. A configurable grace period (GW_GRACE env var, default 10 s) causes the server to exit if no client registers shortly after boot.

Class attributes (all configurable via environment variables): HOST: Listening address. Defaults to GW_HOST or "127.0.0.1". PORT: Listening port. Defaults to GW_PORT or 11321. TOTAL_UNITS: Total concurrency budget shared across all clients. Defaults to GW_UNITS or 8. VALID_COSTS: Tuple of accepted per-query unit costs. Defaults to GW_VALID_COSTS or (1, 2, 4). LOCKFILE: Path to an advisory lock file. Defaults to GW_LOCKFILE or a temp-directory path llm_api_gateway.lock.

Examples:

Starting the gateway server programmatically (normally done via serve_api_gateway)::

>>> import asyncio
>>> from unaiverse.modules.utils import APIGatewayServer, featherless_call
>>> async def main():
...     server = APIGatewayServer(call=featherless_call)
...     await server.run()
>>> # asyncio.run(main())  # Blocks until all clients disconnect.

Initialize the gateway server with a scheduler and an empty client registry.

Creates the APIGatewayScheduler, stores the model-call callable, initializes the set of live registration writers (live), and creates the asyncio shutdown event that run waits on.

Parameters:

Name	Type	Description	Default
call	Callable	An asynchronous callable with the signature async (sys_prompt, prompt, sampler, model) -> str that performs the actual LLM API call. Defaults to featherless_call.	featherless_call

Source code in unaiverse/modules/utils.py

def __init__(self, call: Callable = featherless_call):
    """Initialize the gateway server with a scheduler and an empty client registry.

    Creates the ``APIGatewayScheduler``, stores the model-call callable, initializes
    the set of live registration writers (``live``), and creates the asyncio shutdown
    event that ``run`` waits on.

    Args:
        call: An asynchronous callable with the signature
            ``async (sys_prompt, prompt, sampler, model) -> str`` that performs the
            actual LLM API call. Defaults to ``featherless_call``.
    """
    self.sched: APIGatewayScheduler = APIGatewayScheduler()
    self.call: Callable = call
    self.live: set = set()  # Set of registration writers (one per live client)
    self.shutdown: asyncio.Event = asyncio.Event()

HOST class-attribute instance-attribute

HOST: str = get('GW_HOST', '127.0.0.1')

PORT class-attribute instance-attribute

PORT: int = int(get('GW_PORT', '11321'))

TOTAL_UNITS class-attribute instance-attribute

TOTAL_UNITS = int(get('GW_UNITS', '8'))

VALID_COSTS class-attribute instance-attribute

VALID_COSTS = tuple((int(x)) for x in (split(',')))

LOCKFILE class-attribute instance-attribute

LOCKFILE: str = get('GW_LOCKFILE', join(gettempdir(), 'llm_api_gateway.lock'))

sched instance-attribute

sched: APIGatewayScheduler = APIGatewayScheduler()

call instance-attribute

call: Callable = call

live instance-attribute

live: set = set()

shutdown instance-attribute

shutdown: Event = Event()

handle async

handle(reader: StreamReader, writer: StreamWriter) -> None

Serve a single client TCP connection until the peer closes it (async).

Reads newline-delimited JSON messages in a loop. Two "op" values are recognized:

• "hello": marks this connection as the client's liveness registration. The writer is added to self.live. When the connection drops, the writer is removed from self.live and, if the set is now empty, self.shutdown is set to trigger server teardown.
• "generate": extracts process_id, sys_prompt, prompt, cost, model, and sampler from the message, validates the cost against VALID_COSTS, submits the query to the scheduler via self.sched.submit, and writes the result back as a JSON line {"ok": true, "result": "..."} or {"ok": false, "error": "..."}.

Malformed JSON lines are silently skipped. ConnectionResetError and asyncio.IncompleteReadError are swallowed so that abrupt client disconnections do not propagate as unhandled exceptions.

Parameters:

Name	Type	Description	Default
reader	StreamReader	The asyncio stream reader for the client connection.	required
writer	StreamWriter	The asyncio stream writer for the client connection.	required

Source code in unaiverse/modules/utils.py

async def handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    """Serve a single client TCP connection until the peer closes it (async).

    Reads newline-delimited JSON messages in a loop. Two ``"op"`` values are
    recognized:

    - ``"hello"``: marks this connection as the client's liveness registration.
      The writer is added to ``self.live``. When the connection drops,
      the writer is removed from ``self.live`` and, if the set is now empty,
      ``self.shutdown`` is set to trigger server teardown.
    - ``"generate"``: extracts ``process_id``, ``sys_prompt``, ``prompt``,
      ``cost``, ``model``, and ``sampler`` from the message, validates the cost
      against ``VALID_COSTS``, submits the query to the scheduler via
      ``self.sched.submit``, and writes the result back as a JSON line
      ``{"ok": true, "result": "..."}`` or ``{"ok": false, "error": "..."}``.

    Malformed JSON lines are silently skipped. ``ConnectionResetError`` and
    ``asyncio.IncompleteReadError`` are swallowed so that abrupt client
    disconnections do not propagate as unhandled exceptions.

    Args:
        reader: The asyncio stream reader for the client connection.
        writer: The asyncio stream writer for the client connection.
    """
    registered = False
    try:
        while True:
            line = await reader.readline()
            if not line:
                break  # Connection closed by the client (or it died)
            try:
                msg = json.loads(line)
            except json.JSONDecodeError:
                continue
            op = msg.get("op")
            if op == "hello":

                # This connection is the client's liveness registration
                registered = True
                self.live.add(writer)
            elif op == "generate":
                pid = msg["process_id"]
                sys_prompt = msg.get("sys_prompt", "")
                prompt = msg["prompt"]
                cost = int(msg.get("cost", 1))
                model = msg.get("model")
                sampler = msg.get("sampler") or {}
                if cost not in APIGatewayServer.VALID_COSTS:
                    out = {"ok": False, "error": f"bad cost {cost}"}
                else:
                    try:
                        result = await self.sched.submit(pid, sys_prompt, prompt, cost, model, sampler)
                        out = {"ok": True, "result": result}
                    except Exception as e:
                        out = {"ok": False, "error": str(e)}
                writer.write((json.dumps(out) + "\n").encode())
                await writer.drain()
    except (ConnectionResetError, asyncio.IncompleteReadError):
        pass
    finally:
        if registered:
            self.live.discard(writer)

            # When the last interested client disconnects, shut down
            if not self.live:
                self.shutdown.set()
        try:
            writer.close()
        except Exception:
            pass

run async

run() -> None

Bind the listening socket, start the dispatcher, and run until shutdown is requested.

Binding also arbitrates the startup race: if the port is already taken, start_server raises and this process is not the server. A grace period guards against an orphan server whose starter died before registering.

Source code in unaiverse/modules/utils.py

async def run(self) -> None:
    """Bind the listening socket, start the dispatcher, and run until shutdown is requested.

    Binding also arbitrates the startup race: if the port is already taken, ``start_server`` raises and this process
    is not the server. A grace period guards against an orphan server whose starter died before registering.
    """

    # Bind first (this is also the startup-race arbiter): if the port is taken, bind() raises and we are NOT
    # the server
    server = await asyncio.start_server(self.handle, APIGatewayServer.HOST, APIGatewayServer.PORT)
    asyncio.create_task(self.sched.dispatcher(self.call))

    # Grace period: if no client registers shortly after boot, exit, so a server spawned by a starter that
    # immediately died does not linger
    grace_secs = float(os.environ.get("GW_GRACE", "10"))

    async def grace():
        await asyncio.sleep(grace_secs)
        if not self.live:
            self.shutdown.set()
    asyncio.create_task(grace())

    async with server:
        await self.shutdown.wait()

APIGatewayScheduler

APIGatewayScheduler()

Strict round-robin scheduler with a shared unit-budget admission, used by the gateway server.

The scheduler serves one query per process per turn (round-robin), with equal importance regardless of how many units a query costs: the cost only affects the budget admission check, never the turn order. Within a process, queries are served in submission order (FIFO).

Two design decisions are worth highlighting: - Strict hold, never skip. When it is a process's turn but its next query does not fit the remaining budget, the dispatcher STALLS and holds the free units until that query can run; it does not skip ahead to another process. This is intentional (chosen over skip-to-fill), accepting the throughput cost. - At most one held "head" per process. The dispatcher pops a single query and keeps it in the held dict instead of peeking the private asyncio.Queue deque (which caused an IndexError race across await boundaries). A single asyncio.Condition drives every wakeup (new work submitted, or units freed by a finished job).

Create a Scheduler.

Source code in unaiverse/modules/utils.py

def __init__(self):
    """Create a Scheduler."""
    self.available: int = APIGatewayServer.TOTAL_UNITS
    self.cond: asyncio.Condition = asyncio.Condition()
    self.queues: dict[str, asyncio.Queue] = {}  # Process ID -> queue of (prompt, cost, model, fut) - ORDERED
    self.rr: int = 0  # Round-robin cursor into the (insertion-ordered) list of process IDs

available instance-attribute

available: int = TOTAL_UNITS

cond instance-attribute

cond: Condition = Condition()

queues instance-attribute

queues: dict[str, Queue] = {}

rr instance-attribute

rr: int = 0

submit async

submit(pid: str, sys_prompt: str, prompt: str, cost: int, model: str | None, sampler: dict | None = None) -> str

Enqueue a query for a process and wait until the dispatcher runs it.

Parameters:

Name	Type	Description	Default
pid	str	The process ID the query belongs to (round-robin fairness is per process).	required
sys_prompt	str	The system prompt string (can be "").	required
prompt	str	The prompt string to send to the model.	required
cost	int	The query's unit cost (one of VALID_COSTS).	required
model	str | None	The model identifier to use (None lets the call fall back to its own default).	required
sampler	dict | None	Sampling parameters to forward to the model call (None means use API defaults).	None

Returns:

Type	Description
str	The model's textual result.

Source code in unaiverse/modules/utils.py

async def submit(self, pid: str, sys_prompt: str, prompt: str, cost: int, model: str | None,
                 sampler: dict | None = None) -> str:
    """Enqueue a query for a process and wait until the dispatcher runs it.

    Args:
        pid: The process ID the query belongs to (round-robin fairness is per process).
        sys_prompt: The system prompt string (can be "").
        prompt: The prompt string to send to the model.
        cost: The query's unit cost (one of VALID_COSTS).
        model: The model identifier to use (None lets the call fall back to its own default).
        sampler: Sampling parameters to forward to the model call (None means use API defaults).

    Returns:
        The model's textual result.
    """
    fut: asyncio.Future = asyncio.get_running_loop().create_future()
    await self._queue(pid).put((sys_prompt, prompt, cost, model, sampler, fut))
    async with self.cond:
        self.cond.notify_all()  # Wake the dispatcher: there is new work
    return await fut

dispatcher async

dispatcher(call: Callable) -> None

Forever serve queued queries in strict round-robin order under the unit budget.

Parameters:

Name	Type	Description	Default
call	Callable	The asynchronous function performing the actual model call (passed on to each job).	required

Source code in unaiverse/modules/utils.py

async def dispatcher(self, call: Callable) -> None:
    """Forever serve queued queries in strict round-robin order under the unit budget.

    Args:
        call: The asynchronous function performing the actual model call (passed on to each job).
    """

    # held[pid] holds a (prompt, cost, model, fut) we popped for a process but could not launch yet (its turn came
    # but the budget did not fit). Strict hold keeps it here until it can run; never put it back nor skip past it.
    held: dict = {}
    while True:
        async with self.cond:
            pids = list(self.queues.keys())
            if not pids:
                await self._wait_locked()
                continue

            pid = pids[self.rr % len(pids)]
            q = self.queues[pid]

            # Determine this process's current head: either the one we are holding, or the next one off its queue
            if pid in held:
                sys_prompt, prompt, cost, model, sampler, fut = held[pid]
            elif not q.empty():
                sys_prompt, prompt, cost, model, sampler, fut = q.get_nowait()
            else:

                # The current process has nothing: advance the turn. If nobody has work at all, wait; otherwise
                # keep rotating until we reach a process that does
                self.rr = (self.rr + 1) % len(pids)
                if all(self.queues[p].empty() for p in pids) and not held:
                    await self._wait_locked()
                continue

            if self.available >= cost:
                self.available -= cost
                held.pop(pid, None)
                asyncio.create_task(
                    self._run_job(sys_prompt, prompt, cost, model, sampler, fut, call))
                self.rr = (self.rr + 1) % len(pids)
                continue
            else:

                # Strict hold: it is this process's turn and its head does not fit. Hold the head and wait for units
                # to free up. Do NOT advance the cursor and do NOT serve anyone else.
                held[pid] = (sys_prompt, prompt, cost, model, sampler, fut)
                await self._wait_locked()
                continue

has_human_processor

has_human_processor(agent: object) -> bool

Return True if the given agent's processor wraps a HumanModule.

Inspects the proc attribute of the agent, checks that it is not None, and tests whether its inner module is an instance of HumanModule. This is the canonical check used by the framework to decide whether a human is expected to supply the agent's output interactively.

Parameters:

Name	Type	Description	Default
agent	object	The agent object to inspect. Must expose a proc attribute (raises AssertionError if the attribute is absent).	required

Returns:

Type	Description
bool	True if agent.proc is not None and agent.proc.module is a
bool	HumanModule instance; False otherwise.

Raises:

Type	Description
AssertionError	If agent does not have a proc attribute.

Source code in unaiverse/modules/utils.py

def has_human_processor(agent: object) -> bool:
    """Return ``True`` if the given agent's processor wraps a ``HumanModule``.

    Inspects the ``proc`` attribute of the agent, checks that it is not ``None``,
    and tests whether its inner ``module`` is an instance of ``HumanModule``.
    This is the canonical check used by the framework to decide whether a human
    is expected to supply the agent's output interactively.

    Args:
        agent: The agent object to inspect. Must expose a ``proc`` attribute
            (raises ``AssertionError`` if the attribute is absent).

    Returns:
        ``True`` if ``agent.proc`` is not ``None`` and ``agent.proc.module`` is a
        ``HumanModule`` instance; ``False`` otherwise.

    Raises:
        AssertionError: If ``agent`` does not have a ``proc`` attribute.
    """
    assert hasattr(agent, "proc")
    return getattr(agent, "proc") is not None and isinstance(getattr(agent, "proc").module, HumanModule)

transforms_factory

transforms_factory(trans_type: str, add_batch_dim: bool = True, return_inverse: bool = False) -> Compose | None

Build and return a torchvision transform pipeline for the requested image type.

Each trans_type identifier maps to a fixed, opinionated pipeline:

• "rgb<N>" / "gray<N>" - convert color mode, resize and center-crop to <N> pixels, convert to tensor, then apply ImageNet (or grayscale) normalization. <N> must be a positive integer suffix (e.g. "rgb224").
• "rgb-no_norm<N>" / "gray-no_norm<N>" - same as above but without normalization; the output tensor is uint8 in [0, 255].
• "rgb" / "gray" - no resize or crop; normalize with ImageNet statistics.
• "rgb-no_norm" / "gray-no_norm" - no resize, crop, or normalization.
• "gray_mnist" - MNIST-specific pipeline: convert to grayscale, resize and center-crop to 28x28, convert to tensor, normalize with MNIST mean (0.1307) and standard deviation (0.3081).

When add_batch_dim=True, a batch dimension is added by unsqueeze(0) at the end of the forward pipeline and removed by squeeze(0) at the start of the inverse pipeline, so the output is a 4-D tensor ready for direct model inference.

When return_inverse=True, the inverse (de-normalization and to-PIL) pipeline is returned instead of the forward one, which is useful for visualization.

Parameters:

Name	Type	Description	Default
trans_type	str	A string identifying the desired transform pipeline. Supported base names are "rgb", "gray", "rgb-no_norm", "gray-no_norm", and "gray_mnist"; append an integer to any of the first four for a resize-and-crop variant (e.g. "rgb224").	required
add_batch_dim	bool	If True, appends an unsqueeze(0) step to the forward pipeline and inserts a squeeze(0) step at the start of the inverse pipeline. Defaults to True.	True
return_inverse	bool	If True, returns the inverse (de-normalization) pipeline instead of the forward one. Defaults to False.	False

Returns:

Type	Description
Compose | None	The requested transforms.Compose pipeline, or None if the combination of
Compose | None	arguments produces no pipeline (should not occur for valid inputs).

Raises:

Type	Description
ValueError	If trans_type does not match any supported pattern.

Examples:

>>> fwd = transforms_factory("rgb224")
>>> from PIL import Image
>>> img = Image.new("RGB", (300, 300))
>>> tensor = fwd(img)  # shape: (1, 3, 224, 224)
>>> inv = transforms_factory("rgb224", return_inverse=True)
>>> recovered = inv(tensor.squeeze(0))  # PIL Image

Source code in unaiverse/modules/utils.py

def transforms_factory(trans_type: str, add_batch_dim: bool = True,
                       return_inverse: bool = False) -> transforms.Compose | None:
    """Build and return a torchvision transform pipeline for the requested image type.

    Each ``trans_type`` identifier maps to a fixed, opinionated pipeline:

    - ``"rgb<N>"`` / ``"gray<N>"`` - convert color mode, resize and center-crop to ``<N>``
      pixels, convert to tensor, then apply ImageNet (or grayscale) normalization. ``<N>``
      must be a positive integer suffix (e.g. ``"rgb224"``).
    - ``"rgb-no_norm<N>"`` / ``"gray-no_norm<N>"`` - same as above but without normalization;
      the output tensor is ``uint8`` in ``[0, 255]``.
    - ``"rgb"`` / ``"gray"`` - no resize or crop; normalize with ImageNet statistics.
    - ``"rgb-no_norm"`` / ``"gray-no_norm"`` - no resize, crop, or normalization.
    - ``"gray_mnist"`` - MNIST-specific pipeline: convert to grayscale, resize and
      center-crop to 28x28, convert to tensor, normalize with MNIST mean (0.1307) and
      standard deviation (0.3081).

    When ``add_batch_dim=True``, a batch dimension is added by ``unsqueeze(0)`` at the
    end of the forward pipeline and removed by ``squeeze(0)`` at the start of the inverse
    pipeline, so the output is a 4-D tensor ready for direct model inference.

    When ``return_inverse=True``, the inverse (de-normalization and to-PIL) pipeline is
    returned instead of the forward one, which is useful for visualization.

    Args:
        trans_type: A string identifying the desired transform pipeline. Supported base
            names are ``"rgb"``, ``"gray"``, ``"rgb-no_norm"``, ``"gray-no_norm"``, and
            ``"gray_mnist"``; append an integer to any of the first four for a
            resize-and-crop variant (e.g. ``"rgb224"``).
        add_batch_dim: If ``True``, appends an ``unsqueeze(0)`` step to the forward
            pipeline and inserts a ``squeeze(0)`` step at the start of the inverse
            pipeline. Defaults to ``True``.
        return_inverse: If ``True``, returns the inverse (de-normalization) pipeline
            instead of the forward one. Defaults to ``False``.

    Returns:
        The requested ``transforms.Compose`` pipeline, or ``None`` if the combination of
        arguments produces no pipeline (should not occur for valid inputs).

    Raises:
        ValueError: If ``trans_type`` does not match any supported pattern.

    Examples:
        >>> fwd = transforms_factory("rgb224")
        >>> from PIL import Image
        >>> img = Image.new("RGB", (300, 300))
        >>> tensor = fwd(img)  # shape: (1, 3, 224, 224)
        >>> inv = transforms_factory("rgb224", return_inverse=True)
        >>> recovered = inv(tensor.squeeze(0))  # PIL Image
    """
    supported_types = {"rgb*",          "gray*",
                       "rgb-no_norm*",  "gray-no_norm*",
                       "rgb",           "gray",
                       "rgb-no_norm",   "gray-no_norm",
                                        "gray_mnist"}

    found = False
    num = -1
    for _type in supported_types:
        if _type.endswith("*"):
            has_star = True
            __type = _type[0:-1]
        else:
            has_star = False
            __type = _type

        if has_star and trans_type.startswith(__type) and len(trans_type) > len(__type):
            try:
                num = int(trans_type[len(__type):])
                trans_type = _type
                found = True
                break
            except ValueError:
                pass
        elif trans_type == _type:
            found = True
            break

    if not found:
        raise ValueError(f"Invalid transformation type '{trans_type}': must be one of {supported_types}, "
                         f"where * is an integer number")

    trans = None
    inverse_trans = None

    if trans_type == "rgb*":
        trans = transforms.Compose([
            transforms.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),  # Ensure 3 channels
            transforms.Resize(num),
            transforms.CenterCrop(num),
            transforms.ToTensor(),  # Convert PIL to tensor (3, H, W), float [0,1]
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),
        ])
        inverse_trans = transforms.Compose([
            transforms.Normalize(mean=[0., 0., 0.],
                                 std=[1. / 0.229, 1. / 0.224, 1. / 0.225]),
            transforms.Lambda(lambda x: x + torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)),
            transforms.ToPILImage()
        ])
    elif trans_type == "gray*":
        trans = transforms.Compose([
            transforms.Lambda(lambda img: img.convert("L") if img.mode != "L" else img),  # Ensure 1 channel
            transforms.Resize(num),
            transforms.CenterCrop(num),
            transforms.ToTensor(),  # Convert PIL to tensor (1, H, W), float [0,1]
            transforms.Normalize(mean=[0.45],
                                 std=[0.225])
        ])
        inverse_trans = transforms.Compose([
            transforms.Normalize(mean=[0.],
                                 std=[1. / 0.225]),
            transforms.Lambda(lambda x: x + 0.45),
            transforms.ToPILImage()
        ])
    elif trans_type == "rgb-no_norm*":
        trans = transforms.Compose([
            transforms.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),  # Ensure 3 channels
            transforms.Resize(num),
            transforms.CenterCrop(num),
            transforms.PILToTensor(),  # Convert PIL to tensor (3, H, W), uint [0,255]
        ])
        inverse_trans = transforms.Compose([transforms.ToPILImage()])
    elif trans_type == "gray-no_norm*":
        trans = transforms.Compose([
            transforms.Lambda(lambda img: img.convert("L") if img.mode != "L" else img),  # Ensure 1 channel
            transforms.Resize(num),
            transforms.CenterCrop(num),
            transforms.PILToTensor(),  # Convert PIL to tensor (1, H, W), uint [0,255]
        ])
        inverse_trans = transforms.Compose([transforms.ToPILImage()])
    elif trans_type == "rgb":
        trans = transforms.Compose([
            transforms.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),  # Ensure 3 channels
            transforms.ToTensor(),  # Convert PIL to tensor (3, H, W), float [0,1]
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])
        inverse_trans = transforms.Compose([
            transforms.Normalize(mean=[0., 0., 0.],
                                 std=[1. / 0.229, 1. / 0.224, 1. / 0.225]),
            transforms.Lambda(lambda x: x + torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)),
            transforms.ToPILImage()
        ])
    elif trans_type == "gray":
        trans = transforms.Compose([
            transforms.Lambda(lambda img: img.convert("L") if img.mode != "L" else img),  # Ensure 1 channel
            transforms.ToTensor(),  # Convert PIL to tensor (1, H, W), float [0,1]
            transforms.Normalize(mean=[0.45],
                                 std=[0.225])
        ])
        inverse_trans = transforms.Compose([
            transforms.Normalize(mean=[0.],
                                 std=[1. / 0.225]),
            transforms.Lambda(lambda x: x + 0.45),
            transforms.ToPILImage()
        ])
    elif trans_type == "rgb-no_norm":
        trans = transforms.Compose([
            transforms.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),  # Ensure 3 channels
            transforms.PILToTensor(),  # Convert PIL to tensor (3, H, W), uint [0,255]
        ])
        inverse_trans = transforms.Compose([transforms.ToPILImage()])
    elif trans_type == "gray-no_norm":
        trans = transforms.Compose([
            transforms.Lambda(lambda img: img.convert("L") if img.mode != "L" else img),  # Ensure 1 channel
            transforms.PILToTensor(),  # Convert PIL to tensor (1, H, W), uint [0,255]
        ])
        inverse_trans = transforms.Compose([transforms.ToPILImage()])
    elif trans_type == "gray_mnist":
        trans = transforms.Compose([
            transforms.Lambda(lambda img: img.convert("L") if img.mode != "L" else img),  # Ensure 1 channel
            transforms.Resize(28),
            transforms.CenterCrop(28),
            transforms.ToTensor(),  # Convert PIL to tensor (1, H, W), float [0,1]
            transforms.Normalize(mean=[0.1307],  # MNIST
                                 std=[0.3081])  # MNIST
        ])
        inverse_trans = transforms.Compose([
            transforms.Normalize(mean=[0.],
                                 std=[1. / 0.3081]),
            transforms.Lambda(lambda x: x + 0.1307),
            transforms.ToPILImage()
        ])

    if add_batch_dim and (trans is not None and inverse_trans is not None):
        trans.transforms.append(transforms.Lambda(lambda x: x.unsqueeze(0)))
        inverse_trans.transforms.insert(0, transforms.Lambda(lambda x: x.squeeze(0)))

    return trans if not return_inverse else inverse_trans

hard_tanh

hard_tanh(x: Tensor) -> Tensor

Return the element-wise hard tanh of the input tensor.

Equivalent to torch.clamp(x, min=-1., max=1.). Unlike the smooth hyperbolic tangent, this activation saturates exactly at -1 and 1 and is linear in between, making it fast and gradient-friendly for bounded representations.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor of any shape and dtype.	required

Returns:

Type	Description
Tensor	A tensor of the same shape as x with all values clamped to the closed
Tensor	interval [-1, 1].

Source code in unaiverse/modules/utils.py

def hard_tanh(x: torch.Tensor) -> torch.Tensor:
    """Return the element-wise hard tanh of the input tensor.

    Equivalent to ``torch.clamp(x, min=-1., max=1.)``. Unlike the smooth hyperbolic
    tangent, this activation saturates exactly at -1 and 1 and is linear in between,
    making it fast and gradient-friendly for bounded representations.

    Args:
        x: Input tensor of any shape and dtype.

    Returns:
        A tensor of the same shape as ``x`` with all values clamped to the closed
        interval ``[-1, 1]``.
    """
    return torch.clamp(x, min=-1., max=1.)

target_shape_fixed_cross_entropy

target_shape_fixed_cross_entropy(output: Tensor, target: Tensor, *args, **kwargs) -> Tensor

Compute cross-entropy loss after squeezing a spurious leading dimension from the target.

When targets arrive from the framework pipeline they may carry an extra batch-size-1 leading dimension (shape (1, N) instead of (N,)). This wrapper calls target.squeeze(0) when target.ndim > 1 before delegating to torch.nn.functional.cross_entropy, so the loss computation always receives a correctly-shaped target regardless of how the batch dimension was added upstream.

Parameters:

Name	Type	Description	Default
output	Tensor	The model output logits tensor (shape (batch, num_classes)).	required
target	Tensor	The target class-index tensor. If it has more than one dimension, the first (leading) dimension is squeezed before computing the loss.	required
*args		Additional positional arguments forwarded to torch.nn.functional.cross_entropy.	()
**kwargs		Additional keyword arguments forwarded to torch.nn.functional.cross_entropy.	{}

Returns:

Type	Description
Tensor	The scalar cross-entropy loss tensor.

Source code in unaiverse/modules/utils.py

def target_shape_fixed_cross_entropy(output: torch.Tensor, target: torch.Tensor,
                                     *args, **kwargs) -> torch.Tensor:
    """Compute cross-entropy loss after squeezing a spurious leading dimension from the target.

    When targets arrive from the framework pipeline they may carry an extra batch-size-1
    leading dimension (shape ``(1, N)`` instead of ``(N,)``). This wrapper calls
    ``target.squeeze(0)`` when ``target.ndim > 1`` before delegating to
    ``torch.nn.functional.cross_entropy``, so the loss computation always receives a
    correctly-shaped target regardless of how the batch dimension was added upstream.

    Args:
        output: The model output logits tensor (shape ``(batch, num_classes)``).
        target: The target class-index tensor. If it has more than one dimension, the
            first (leading) dimension is squeezed before computing the loss.
        *args: Additional positional arguments forwarded to
            ``torch.nn.functional.cross_entropy``.
        **kwargs: Additional keyword arguments forwarded to
            ``torch.nn.functional.cross_entropy``.

    Returns:
        The scalar cross-entropy loss tensor.
    """
    if len(target.shape) > 1:
        target = target.squeeze(0)
    return torch.nn.functional.cross_entropy(output, target, *args, **kwargs)

set_seed

set_seed(seed: int) -> None

Seed all relevant random-number generators for reproducible runs.

Sets the seed for PyTorch (torch.manual_seed), Python's built-in random module, and NumPy (numpy.random.seed). The NumPy seed is always fixed to 0 regardless of the value of seed; this is intentional for the current version. When seed is negative the function is a no-op.

Parameters:

Name	Type	Description	Default
seed	int	The integer seed value. Pass a negative integer to skip seeding entirely (useful when reproducibility is not required).	required

Note

CUDA randomness is not explicitly seeded here. For fully deterministic GPU training, also call torch.cuda.manual_seed_all and set torch.backends.cudnn.deterministic = True in the calling code.

Source code in unaiverse/modules/utils.py

def set_seed(seed: int) -> None:
    """Seed all relevant random-number generators for reproducible runs.

    Sets the seed for PyTorch (``torch.manual_seed``), Python's built-in ``random``
    module, and NumPy (``numpy.random.seed``). The NumPy seed is always fixed to ``0``
    regardless of the value of ``seed``; this is intentional for the current version.
    When ``seed`` is negative the function is a no-op.

    Args:
        seed: The integer seed value. Pass a negative integer to skip seeding entirely
            (useful when reproducibility is not required).

    Note:
        CUDA randomness is not explicitly seeded here. For fully deterministic GPU
        training, also call ``torch.cuda.manual_seed_all`` and set
        ``torch.backends.cudnn.deterministic = True`` in the calling code.
    """
    if seed >= 0:
        torch.manual_seed(seed)
        random.seed(seed)
        np.random.seed(0)

get_proc_inputs_and_proc_outputs_for_rnn

get_proc_inputs_and_proc_outputs_for_rnn(u_shape: Size | tuple, du_dim: int, y_dim: int) -> tuple[list, list]

Build StreamType descriptors for the inputs and output of an RNN-style processor.

Constructs the two-input, one-output stream specification expected by a recurrent processor: the first input is the main observation tensor of shape (batch, *u_shape); the second input is a delta-u tensor of shape (batch, du_dim); the single output is a tensor of shape (batch, y_dim). All tensors use torch.float32 dtype and are non-pubsub, non-private streams.

If u_shape is a torch.Size object it is converted to a plain tuple before building the descriptor.

Parameters:

Name	Type	Description	Default
u_shape	Size | tuple	Shape of the main input tensor, excluding the batch dimension (e.g. (4,) for a 4-dimensional observation).	required
du_dim	int	Dimensionality of the secondary delta-u input tensor.	required
y_dim	int	Dimensionality of the output tensor.	required

Returns:

Type	Description
list	A two-element tuple (proc_inputs, proc_outputs) where proc_inputs is a
list	list of two StreamType objects and proc_outputs is a list containing one
tuple[list, list]	StreamType object.

Source code in unaiverse/modules/utils.py

def get_proc_inputs_and_proc_outputs_for_rnn(u_shape: torch.Size | tuple, du_dim: int,
                                             y_dim: int) -> tuple[list, list]:
    """Build ``StreamType`` descriptors for the inputs and output of an RNN-style processor.

    Constructs the two-input, one-output stream specification expected by a recurrent
    processor: the first input is the main observation tensor of shape
    ``(batch, *u_shape)``; the second input is a delta-u tensor of shape
    ``(batch, du_dim)``; the single output is a tensor of shape ``(batch, y_dim)``.
    All tensors use ``torch.float32`` dtype and are non-pubsub, non-private streams.

    If ``u_shape`` is a ``torch.Size`` object it is converted to a plain ``tuple``
    before building the descriptor.

    Args:
        u_shape: Shape of the main input tensor, *excluding* the batch dimension (e.g.
            ``(4,)`` for a 4-dimensional observation).
        du_dim: Dimensionality of the secondary delta-u input tensor.
        y_dim: Dimensionality of the output tensor.

    Returns:
        A two-element tuple ``(proc_inputs, proc_outputs)`` where ``proc_inputs`` is a
        list of two ``StreamType`` objects and ``proc_outputs`` is a list containing one
        ``StreamType`` object.
    """
    if isinstance(u_shape, torch.Size):
        u_shape = tuple(u_shape)
    proc_inputs = [
        StreamType(data_type="tensor", tensor_shape=(None,) + u_shape, tensor_dtype=torch.float32,
                   pubsub=False, private_only=False),
        StreamType(data_type="tensor", tensor_shape=(None, du_dim,), tensor_dtype=torch.float32,
                   pubsub=False, private_only=False)
    ]
    proc_outputs = [
        StreamType(data_type="tensor", tensor_shape=(None, y_dim), tensor_dtype=torch.float32,
                   pubsub=False, private_only=False)
    ]
    return proc_inputs, proc_outputs

get_proc_inputs_and_proc_outputs_for_image_classification

get_proc_inputs_and_proc_outputs_for_image_classification(y_dim: int = -1, trans_forms=None) -> tuple[list[StreamType], list[StreamType]]

Build StreamType descriptors for the inputs and output of an image-classification processor.

Constructs the canonical single-image-in, single-logit-vector-out specification for an image classification model. The input stream is of type "img" (no fixed tensor shape) and carries the optional stream_to_proc_transforms transform. The output stream is a float tensor of shape (batch, y_dim).

When y_dim=-1, the output dimensionality defaults to 1000, matching the class count of models pre-trained on ImageNet.

Parameters:

Name	Type	Description	Default
y_dim	int	Number of output classes. Pass -1 to default to 1000 (ImageNet). Defaults to -1.	-1
trans_forms		A torchvision transform (or None) applied to the image stream before it is passed to the processor. Defaults to None.	None

Returns:

Type	Description
list[StreamType]	A two-element tuple (proc_inputs, proc_outputs) where proc_inputs is a
list[StreamType]	list containing one StreamType of type "img" and proc_outputs is a
tuple[list[StreamType], list[StreamType]]	list containing one StreamType of type "tensor".

Source code in unaiverse/modules/utils.py

def get_proc_inputs_and_proc_outputs_for_image_classification(y_dim: int = -1, trans_forms=None) \
        -> tuple[list[StreamType], list[StreamType]]:
    """Build ``StreamType`` descriptors for the inputs and output of an image-classification processor.

    Constructs the canonical single-image-in, single-logit-vector-out specification for
    an image classification model. The input stream is of type ``"img"`` (no fixed tensor
    shape) and carries the optional ``stream_to_proc_transforms`` transform. The output
    stream is a float tensor of shape ``(batch, y_dim)``.

    When ``y_dim=-1``, the output dimensionality defaults to 1000, matching the class
    count of models pre-trained on ImageNet.

    Args:
        y_dim: Number of output classes. Pass ``-1`` to default to 1000 (ImageNet).
            Defaults to ``-1``.
        trans_forms: A torchvision transform (or ``None``) applied to the image stream
            before it is passed to the processor. Defaults to ``None``.

    Returns:
        A two-element tuple ``(proc_inputs, proc_outputs)`` where ``proc_inputs`` is a
        list containing one ``StreamType`` of type ``"img"`` and ``proc_outputs`` is a
        list containing one ``StreamType`` of type ``"tensor"``.
    """
    if y_dim == -1:
        y_dim = 1000  # Assuming ImageNet-trained models
    proc_inputs = [StreamType(data_type="img", pubsub=False, private_only=False, stream_to_proc_transforms=trans_forms)]
    proc_outputs = [StreamType(data_type="tensor", tensor_shape=(None, y_dim), tensor_dtype=torch.float32,
                               pubsub=False, private_only=False)]
    return proc_inputs, proc_outputs

isinstance_fcn

isinstance_fcn(obj: object, class_to_check: type | tuple) -> bool

Return whether obj is an instance of the given type or tuple of types.

Thin wrapper around the built-in isinstance that can be passed as a first-class callable in places where a plain isinstance call cannot be used directly (e.g., as a function argument or in a lambda).

Parameters:

Name	Type	Description	Default
obj	object	The object to test.	required
class_to_check	type | tuple	A single type or a tuple of types to test against.	required

Returns:

Type	Description
bool	True if obj is an instance of class_to_check, False otherwise.

Source code in unaiverse/modules/utils.py

def isinstance_fcn(obj: object, class_to_check: type | tuple) -> bool:
    """Return whether ``obj`` is an instance of the given type or tuple of types.

    Thin wrapper around the built-in ``isinstance`` that can be passed as a
    first-class callable in places where a plain ``isinstance`` call cannot be
    used directly (e.g., as a function argument or in a lambda).

    Args:
        obj: The object to test.
        class_to_check: A single type or a tuple of types to test against.

    Returns:
        ``True`` if ``obj`` is an instance of ``class_to_check``, ``False`` otherwise.
    """
    return isinstance(obj, class_to_check)

error_rate_mnist_test_set

error_rate_mnist_test_set(network: Module, mnist_data_save_path: str) -> float

Evaluate and return a network's top-1 error rate on the full MNIST test set.

Downloads the MNIST test split if it is not already present at mnist_data_save_path, then iterates over all 10,000 samples in batches of 200. The network is temporarily switched to eval mode (network.eval()) for the duration of the evaluation; its original training flag is restored afterwards.

All inputs are moved to the device inferred from the first parameter of network, so the method works transparently on both CPU and CUDA models.

Parameters:

Name	Type	Description	Default
network	Module	The torch.nn.Module to evaluate. Its output must be a 2-D logit tensor of shape (batch, num_classes) from which the predicted class is obtained via argmax.	required
mnist_data_save_path	str	Filesystem path where the MNIST dataset is stored or will be downloaded to.	required

Returns:

Type	Description
float	The fraction of misclassified test samples as a float in the closed interval
float	[0, 1] (lower is better).

Note

The evaluation runs inside torch.no_grad() is NOT called explicitly here; callers that require gradient isolation should wrap the call themselves. The network's training flag is restored to its original value after the call.

Source code in unaiverse/modules/utils.py

def error_rate_mnist_test_set(network: torch.nn.Module, mnist_data_save_path: str) -> float:
    """Evaluate and return a network's top-1 error rate on the full MNIST test set.

    Downloads the MNIST test split if it is not already present at
    ``mnist_data_save_path``, then iterates over all 10,000 samples in batches of 200.
    The network is temporarily switched to eval mode (``network.eval()``) for the
    duration of the evaluation; its original training flag is restored afterwards.

    All inputs are moved to the device inferred from the first parameter of ``network``,
    so the method works transparently on both CPU and CUDA models.

    Args:
        network: The ``torch.nn.Module`` to evaluate. Its output must be a 2-D logit
            tensor of shape ``(batch, num_classes)`` from which the predicted class is
            obtained via ``argmax``.
        mnist_data_save_path: Filesystem path where the MNIST dataset is stored or will
            be downloaded to.

    Returns:
        The fraction of misclassified test samples as a ``float`` in the closed interval
        ``[0, 1]`` (lower is better).

    Note:
        The evaluation runs inside ``torch.no_grad()`` is NOT called explicitly here;
        callers that require gradient isolation should wrap the call themselves. The
        network's training flag is restored to its original value after the call.
    """

    # Getting MNIST test set
    mnist_test = datasets.MNIST(root=mnist_data_save_path,
                                train=False, download=True,
                                transform=transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.1307,), (0.3081,))
                                ]))
    mnist_test = DataLoader(mnist_test, batch_size=200, shuffle=False)

    # Checking error rate
    error_rate = 0.
    n = 0
    training_flag_backup = network.training
    network.eval()
    device = next(network.parameters()).device
    for x, y in mnist_test:
        x = x.to(device)
        y = y.to(device)
        o = network(x)
        c = torch.argmax(o, dim=1)
        error_rate += float(torch.sum(c != y).item())
        n += x.shape[0]
    error_rate /= n
    network.training = training_flag_backup

    return error_rate

featherless_call async

featherless_call(sys_prompt: str, prompt: str, sampler: dict | None = None, model: str | None = None) -> str | None

Call the Featherless AI chat-completions API and return the generated text (async).

Sends a chat-completions request to the Featherless endpoint using the process-wide shared AsyncOpenAI client (see _get_featherless_client). Sampler parameters are split into two groups: keys present in OPENAI_NATIVE_SAMPLER_PARAMS are forwarded as top-level fields; all other keys (e.g. top_k, repetition_penalty, min_p) are forwarded via extra_body so that non-standard vLLM/Featherless extensions reach the backend without triggering an SDK validation error.

The call is retried up to three times with exponential backoff (1 s, 2 s, 4 s). An empty or blank response is treated as a transient failure and also triggers a retry. Diagnostic information (model, finish reason, completion token count, wall time, tokens per second) is logged via log.user on each successful API response.

If the MODEL_ID environment variable is not set and model is None, a critical log message is emitted but the call still proceeds with the placeholder value "MOD". Similarly, if FEATHERLESS_API_KEY is not set, a critical message is logged.

Parameters:

Name	Type	Description	Default
sys_prompt	str	The system prompt string. Pass an empty string "" to omit the system message entirely.	required
prompt	str	The user prompt string to send to the model.	required
sampler	dict | None	A dictionary of sampling parameters. Recognized keys include "temperature", "max_tokens", "top_p", "top_k", "frequency_penalty", "presence_penalty", "repetition_penalty", and "min_p". None or an empty dict uses the API defaults. Defaults to None.	None
model	str | None	The Featherless model identifier (e.g. "meta-llama/Meta-Llama-3.1-8B-Instruct"). If None, the MODEL_ID environment variable is used. Defaults to None.	None

Returns:

Type	Description
str | None	The non-empty text of the first valid response choice, or an empty string
str | None	"" if all retries are exhausted (in practice this path is preceded by a
str | None	log.critical call that raises an exception in the configured log handler
str | None	before the empty string is returned).

Raises:

Type	Description
Exception	Any exception raised by the AsyncOpenAI client on the final retry attempt is logged via log.critical (which may itself raise) and the empty-string fallback is returned.

Note

This function is a coroutine and must be awaited. It is the default call argument to APIGatewayServer and serve_api_gateway.

Source code in unaiverse/modules/utils.py

async def featherless_call(sys_prompt: str, prompt: str, sampler: dict | None = None,
                           model: str | None = None) -> str | None:
    """Call the Featherless AI chat-completions API and return the generated text (async).

    Sends a chat-completions request to the Featherless endpoint using the process-wide
    shared ``AsyncOpenAI`` client (see ``_get_featherless_client``). Sampler parameters
    are split into two groups: keys present in ``OPENAI_NATIVE_SAMPLER_PARAMS`` are
    forwarded as top-level fields; all other keys (e.g. ``top_k``,
    ``repetition_penalty``, ``min_p``) are forwarded via ``extra_body`` so that
    non-standard vLLM/Featherless extensions reach the backend without triggering an SDK
    validation error.

    The call is retried up to three times with exponential backoff (1 s, 2 s, 4 s).
    An empty or blank response is treated as a transient failure and also triggers a
    retry. Diagnostic information (model, finish reason, completion token count, wall
    time, tokens per second) is logged via ``log.user`` on each successful API response.

    If the ``MODEL_ID`` environment variable is not set and ``model`` is ``None``, a
    critical log message is emitted but the call still proceeds with the placeholder
    value ``"MOD"``. Similarly, if ``FEATHERLESS_API_KEY`` is not set, a critical
    message is logged.

    Args:
        sys_prompt: The system prompt string. Pass an empty string ``""`` to omit the
            system message entirely.
        prompt: The user prompt string to send to the model.
        sampler: A dictionary of sampling parameters. Recognized keys include
            ``"temperature"``, ``"max_tokens"``, ``"top_p"``, ``"top_k"``,
            ``"frequency_penalty"``, ``"presence_penalty"``, ``"repetition_penalty"``,
            and ``"min_p"``. ``None`` or an empty dict uses the API defaults.
            Defaults to ``None``.
        model: The Featherless model identifier (e.g.
            ``"meta-llama/Meta-Llama-3.1-8B-Instruct"``). If ``None``, the
            ``MODEL_ID`` environment variable is used. Defaults to ``None``.

    Returns:
        The non-empty text of the first valid response choice, or an empty string
        ``""`` if all retries are exhausted (in practice this path is preceded by a
        ``log.critical`` call that raises an exception in the configured log handler
        before the empty string is returned).

    Raises:
        Exception: Any exception raised by the ``AsyncOpenAI`` client on the final
            retry attempt is logged via ``log.critical`` (which may itself raise) and
            the empty-string fallback is returned.

    Note:
        This function is a coroutine and must be awaited. It is the default ``call``
        argument to ``APIGatewayServer`` and ``serve_api_gateway``.
    """
    client = _get_featherless_client()
    model = model or os.environ.get("MODEL_ID", "MOD")

    if model == "MOD":
        log.critical("Calling Featherless API: the model argument was not set; "
                     "the MODEL_ID environment variable was not set!")

    # Split the requested sampler params: OpenAI-native ones go top-level, the rest ride along in extra_body
    sampler = sampler or {}
    extra_body = {k: v for k, v in sampler.items() if k not in OPENAI_NATIVE_SAMPLER_PARAMS}
    native = {k: v for k, v in sampler.items() if k in OPENAI_NATIVE_SAMPLER_PARAMS}

    # Retry with exponential backoff, keeping the last error so a final failure reports its cause
    last_err = None

    for attempt in range(3):
        try:
            # Always use chat completions; include the system message only when a system prompt is given
            msgs = [{"role": "system", "content": sys_prompt}] if sys_prompt else []
            msgs.append({"role": "user", "content": prompt})

            # The dict is typed dict[str, object] so PyCharm accepts the mixed value types (str / list / int / float)
            kwargs: dict[str, object] = {"model": model, "messages": msgs, **native}
            if extra_body:
                kwargs["extra_body"] = extra_body

            t0 = time.perf_counter()
            resp = await client.chat.completions.create(**kwargs)  # type: ignore
            elapsed = time.perf_counter() - t0

            # Diagnostics: wall time, throughput, completion length + stop reason. A large completion_tokens count
            # suggests thinking is still on; a low tokens/sec on a warm model points at a slow/cold API replica.
            usage = getattr(resp, "usage", None)
            c_toks = getattr(usage, "completion_tokens", 0)
            finish_reason = resp.choices[0].finish_reason if resp.choices else "n/a"
            tps = f"{c_toks / elapsed:.1f}" if (c_toks and elapsed > 0) else "n/a"
            log.user(f"[featherless_call] model={model} finish_reason={finish_reason} "
                     f"completion_tokens={c_toks if c_toks is not None else 'n/a'} "
                     f"elapsed={elapsed:.2f}s tokens/sec={tps}")

            # Return the first choice that carries non-empty text (content may be None)
            for choice in resp.choices:
                content = choice.message.content
                if content is not None and content.strip():
                    return content

            # Empty/blank content: treat it as a transient failure so the loop retries instead of returning ""
            finish = resp.choices[0].finish_reason if resp.choices else "n/a"
            last_err = RuntimeError(f"Empty content from model (finish_reason={finish})")
        except Exception as e:
            last_err = e

        # Retry on either an exception or empty content
        log.error(f"Empty response or exception! (prompt={prompt}, attempt={attempt + 1}, last_err={last_err!r}), "
                  f"waiting {2 ** attempt} seconds before trying again...")
        await asyncio.sleep(2 ** attempt)

    log.critical(f"Featherless AI call failed after retries: {last_err!r}")  # This will raise an exception
    return ""  # Never reached

serve_api_gateway

serve_api_gateway() -> None

Process entrypoint for the detached gateway server (see the APIGatewayServer class).

Spawned by the client via python -c "from unaiverse.modules.utils import serve_api_gateway; serve_api_gateway()". We deliberately do NOT use python -m unaiverse.modules.utils here: importing the parent package unaiverse.modules already imports this module, so running it again via runpy would execute a second copy of it under the name __main__ (RuntimeWarning, and two distinct APIGatewayServer classes). Calling this function avoids that.

Source code in unaiverse/modules/utils.py

def serve_api_gateway() -> None:
    """Process entrypoint for the detached gateway server (see the APIGatewayServer class).

    Spawned by the client via
    ``python -c "from unaiverse.modules.utils import serve_api_gateway; serve_api_gateway()"``. We deliberately do
    NOT use ``python -m unaiverse.modules.utils`` here: importing the parent package ``unaiverse.modules`` already
    imports this module, so running it again via ``runpy`` would execute a second copy of it under the name
    ``__main__`` (RuntimeWarning, and two distinct APIGatewayServer classes). Calling this function avoids that.
    """

    async def _main():

        # Construct the server INSIDE the running loop: the scheduler builds an asyncio.Condition, which binds to
        # the current loop at creation. Building it before asyncio.run() would bind it to the wrong loop.
        await APIGatewayServer(call=featherless_call).run()

    try:
        asyncio.run(_main())
    except OSError:
        pass  # Port already bound: another process beat us to it, just exit