Production hardening: progress bars, device awareness, API boundaries

Five targeted improvements for production readiness:
1. CLI progress bars — compress, analyze, merge show loading progress
2. Device-aware orthogonal_init — tensors created on subspace device/dtype
3. VLoRAModel/SubspaceTrainer __repr__ + thread-safety documentation
4. __all__ in all submodules — proper API boundaries for each module
5. Memory estimation logging in from_adapters() before SVD

Author: tveseli

src/vlora/analysis.py

@@ -2,6 +2,14 @@
 
 from __future__ import annotations
 
+__all__ = [
+    "adapter_diff",
+    "compute_similarity_matrix",
+    "find_clusters",
+    "find_outliers",
+    "subspace_coverage",
+]
+
 from typing import TYPE_CHECKING
 
 import torch

src/vlora/cli.py

@@ -92,15 +92,13 @@ def info(subspace_path: str, as_json: bool):
 @click.option("--adaptive-k", is_flag=True, help="Use per-layer adaptive k selection.")
 def compress(adapter_dirs: tuple[str, ...], output: str, num_components: int | None, variance_threshold: float, adaptive_k: bool):
     """Build shared subspace from adapter directories."""
-    click.echo(f"\n  Loading {len(adapter_dirs)} adapters...")
-
     adapters = []
     task_ids = []
-    for d in adapter_dirs:
-        path = Path(d)
-        adapters.append(load_adapter(path))
-        task_ids.append(path.name)
-        click.echo(f"    Loaded: {path.name}")
+    with click.progressbar(adapter_dirs, label="  Loading adapters") as bar:
+        for d in bar:
+            path = Path(d)
+            adapters.append(load_adapter(path))
+            task_ids.append(path.name)
 
     click.echo("  Building subspace...")
     sub = SharedSubspace.from_adapters(
@@ -182,14 +180,16 @@ def analyze(adapter_dirs: tuple[str, ...], threshold: float, as_json: bool):
 
     adapters = []
     names = []
-    for d in adapter_dirs:
-        path = Path(d)
-        adapters.append(load_adapter(path))
-        names.append(path.name)
+    with click.progressbar(adapter_dirs, label="  Loading adapters") as bar:
+        for d in bar:
+            path = Path(d)
+            adapters.append(load_adapter(path))
+            names.append(path.name)
 
     if len(adapters) < 2:
         raise click.ClickException("Need at least 2 adapters for analysis.")
 
+    click.echo("  Computing similarity matrix...")
     sim = compute_similarity_matrix(adapters)
     clusters = find_clusters(sim, threshold=threshold)
 
@@ -206,10 +206,6 @@ def analyze(adapter_dirs: tuple[str, ...], threshold: float, as_json: bool):
         click.echo(json_mod.dumps(output, indent=2))
         return
 
-    click.echo(f"\n  Loading {len(adapter_dirs)} adapters...")
-    for n in names:
-        click.echo(f"    Loaded: {n}")
-
     click.echo("\n  Pairwise Cosine Similarity:")
     header = "  " + " " * 20 + "  ".join(f"{n[:8]:>8}" for n in names)
     click.echo(header)
@@ -399,12 +395,11 @@ def merge(adapter_dirs: tuple[str, ...], output: str, method: str, weights: str
     """Merge multiple adapters into one using task arithmetic, TIES, or DARE."""
     from vlora.merge import MERGE_METHODS
 
-    click.echo(f"\n  Loading {len(adapter_dirs)} adapters...")
     adapters = []
-    for d in adapter_dirs:
-        path = Path(d)
-        adapters.append(load_adapter(path))
-        click.echo(f"    Loaded: {path.name}")
+    with click.progressbar(adapter_dirs, label="  Loading adapters") as bar:
+        for d in bar:
+            path = Path(d)
+            adapters.append(load_adapter(path))
 
     if len(adapters) < 2:
         raise click.ClickException("Need at least 2 adapters to merge.")

src/vlora/io.py

@@ -6,6 +6,15 @@
 
 from __future__ import annotations
 
+__all__ = [
+    "LoRAWeights",
+    "load_adapter",
+    "load_adapter_from_hub",
+    "save_adapter",
+    "parse_state_dict",
+    "stack_lora_weights",
+]
+
 import json
 import logging
 import re

src/vlora/merge.py

@@ -12,6 +12,8 @@
 
 from __future__ import annotations
 
+__all__ = ["task_arithmetic", "ties_merge", "dare_merge", "MERGE_METHODS"]
+
 import logging
 
 import torch

src/vlora/model.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+__all__ = ["VLoRAModel"]
+
 from typing import Any
 
 import torch
@@ -38,6 +40,11 @@ class VLoRAModel(nn.Module):
     ``compute_dtype`` to match the model's compute precision (typically
     ``torch.bfloat16``).
 
+    Note: This class is **not thread-safe**. Concurrent calls to
+    ``set_task()``, ``merge()``, or ``forward()`` from multiple threads
+    may produce incorrect results. Use a lock or separate model instances
+    for multi-threaded serving.
+
     Usage:
         subspace = SharedSubspace.load("shared_subspace/")
         base_model = AutoModelForCausalLM.from_pretrained("model-name")
@@ -312,6 +319,15 @@ def is_merged(self) -> bool:
         """Whether LoRA deltas are currently baked into base weights."""
         return self._merged
 
+    def __repr__(self) -> str:
+        task = self._active_task or "none"
+        merged = " merged" if self._merged else ""
+        return (
+            f"VLoRAModel(tasks={len(self.subspace.tasks)}, "
+            f"active={task!r}{merged}, "
+            f"layers={len(self._target_modules)})"
+        )
+
     def compile(self, **kwargs) -> VLoRAModel:
         """Compile the base model with torch.compile for faster inference.
 

src/vlora/ops.py

@@ -6,6 +6,20 @@
 
 from __future__ import annotations
 
+__all__ = [
+    "NF4_QUANT_TABLE",
+    "compute_svd",
+    "explained_variance_ratio",
+    "gram_schmidt",
+    "incremental_svd_update",
+    "nf4_pack",
+    "nf4_quantize_dequantize",
+    "nf4_unpack",
+    "project_onto_subspace",
+    "reconstruct_from_subspace",
+    "select_num_components",
+]
+
 import torch
 from torch import Tensor
 

src/vlora/pipeline.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+__all__ = ["init_subspace", "absorb_task", "extract_adapter"]
+
 from pathlib import Path
 
 from vlora.io import LoRAWeights, load_adapter, save_adapter

src/vlora/router.py

@@ -5,6 +5,7 @@
 shared subspace, blending is a cheap linear combination rather than
 reconstructing and merging full LoRA matrices.
 
+
 Usage:
     subspace = SharedSubspace.load("shared_subspace/")
     router = TaskRouter.from_subspace(subspace, hidden_dim=64)
@@ -17,6 +18,8 @@
 
 from __future__ import annotations
 
+__all__ = ["TaskRouter"]
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F

src/vlora/subspace.py

@@ -7,6 +7,8 @@
 
 from __future__ import annotations
 
+__all__ = ["SharedSubspace", "TaskProjection"]
+
 import logging
 from dataclasses import dataclass
 from pathlib import Path
@@ -133,7 +135,20 @@ def from_adapters(
                 the threshold. Overrides num_components.
         """
         check_adapters_compatible(adapters)
-        logger.info("Building subspace from %d adapters", len(adapters))
+
+        # Log memory estimate to help users anticipate resource needs
+        n_adapters = len(adapters)
+        sample_layer = adapters[0].layer_names[0]
+        dim_a = adapters[0].lora_a[sample_layer].numel()
+        dim_b = adapters[0].lora_b[sample_layer].numel()
+        n_layers = len(adapters[0].layer_names)
+        # SVD working memory: ~2 * (N * D * 4 bytes) per layer for A and B
+        svd_bytes = 2 * n_adapters * (dim_a + dim_b) * 4 * n_layers
+        svd_mb = svd_bytes / (1024 * 1024)
+        logger.info(
+            "Building subspace from %d adapters (%d layers, ~%.0f MB estimated)",
+            n_adapters, n_layers, svd_mb,
+        )
 
         if task_ids is None:
             task_ids = [f"task_{i}" for i in range(len(adapters))]

src/vlora/training.py

@@ -19,6 +19,8 @@
 
 from __future__ import annotations
 
+__all__ = ["SubspaceTrainer", "orthogonal_init"]
+
 import torch
 from torch import Tensor
 
@@ -49,9 +51,11 @@ def orthogonal_init(
 
     for layer in subspace.layer_names:
         actual_k = subspace.components_a[layer].shape[0]
-        loadings_a[layer] = torch.randn(actual_k) * scale
+        device = subspace.components_a[layer].device
+        dtype = subspace.components_a[layer].dtype
+        loadings_a[layer] = torch.randn(actual_k, device=device, dtype=dtype) * scale
         # Initialize B-side to zero (like standard LoRA) so initial delta is zero
-        loadings_b[layer] = torch.zeros(actual_k)
+        loadings_b[layer] = torch.zeros(actual_k, device=device, dtype=dtype)
 
     proj = TaskProjection(task_id=task_id, loadings_a=loadings_a, loadings_b=loadings_b)
     subspace.tasks[task_id] = proj
@@ -145,3 +149,10 @@ def num_trainable_params(self) -> int:
     def step_count(self) -> int:
         """Number of optimizer steps taken."""
         return self._step_count
+
+    def __repr__(self) -> str:
+        return (
+            f"SubspaceTrainer(task={self.task_id!r}, "
+            f"params={self.num_trainable_params}, "
+            f"steps={self._step_count})"
+        )