vajra-streamer 0.0.13__tar.gz

vajra_streamer-0.0.13/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Tvashta Labs
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

vajra_streamer-0.0.13/MANIFEST.in

@@ -0,0 +1,3 @@
+include vajra/*.so
+include vajra/*.dylib
+include vajra/*.dll

vajra_streamer-0.0.13/PKG-INFO

@@ -0,0 +1,43 @@
+Metadata-Version: 2.4
+Name: vajra-streamer
+Version: 0.0.13
+Summary: A fast zero-copy PyTorch tensor streamer
+License-Expression: MIT
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0.0
+Dynamic: license-file
+
+# Vajra
+
+A fast zero-copy PyTorch tensor streamer powered by Dlang.
+
+## Usage
+
+Once installed, you can use Vajra in any Python script or Jupyter Notebook:
+
+```python
+import torch
+from vajra import VajraStreamer, StreamConfig
+
+# Configure the streamer (optional, uses defaults otherwise)
+config = StreamConfig(
+    auth_token="hf_YOUR_TOKEN", # Required for gated models like Llama 3
+    chunk_size_mb=64,
+    chunk_workers=16,
+    gpu_workers=3,
+    disable_cache=False
+)
+
+# Context manager ensures VRAM is freed when done
+with VajraStreamer(config) as streamer:
+    # Pass the URL or a Hugging Face repo ID
+    # (e.g. "meta-llama/Meta-Llama-3-8B")
+    tensors = streamer.load("meta-llama/Meta-Llama-3-8B")
+
+    # 'tensors' is a dictionary mapping tensor names to zero-copy PyTorch tensors
+    # that are backed directly by the downloaded GPU memory.
+    for name, tensor in tensors.items():
+        print(f"Tensor: {name}, Shape: {tensor.shape}, Dtype: {tensor.dtype}")
+```

vajra_streamer-0.0.13/README.md

@@ -0,0 +1,32 @@
+# Vajra
+
+A fast zero-copy PyTorch tensor streamer powered by Dlang.
+
+## Usage
+
+Once installed, you can use Vajra in any Python script or Jupyter Notebook:
+
+```python
+import torch
+from vajra import VajraStreamer, StreamConfig
+
+# Configure the streamer (optional, uses defaults otherwise)
+config = StreamConfig(
+    auth_token="hf_YOUR_TOKEN", # Required for gated models like Llama 3
+    chunk_size_mb=64,
+    chunk_workers=16,
+    gpu_workers=3,
+    disable_cache=False
+)
+
+# Context manager ensures VRAM is freed when done
+with VajraStreamer(config) as streamer:
+    # Pass the URL or a Hugging Face repo ID
+    # (e.g. "meta-llama/Meta-Llama-3-8B")
+    tensors = streamer.load("meta-llama/Meta-Llama-3-8B")
+
+    # 'tensors' is a dictionary mapping tensor names to zero-copy PyTorch tensors
+    # that are backed directly by the downloaded GPU memory.
+    for name, tensor in tensors.items():
+        print(f"Tensor: {name}, Shape: {tensor.shape}, Dtype: {tensor.dtype}")
+```

vajra_streamer-0.0.13/example/test.py

@@ -0,0 +1,52 @@
+import os
+import sys
+
+# Ensure the parent directory (python) is in path to import vajra
+sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+
+from vajra import VajraStreamer, StreamConfig
+
+def main():
+    # URL to a small model for testing
+    url = "https://huggingface.co/meta-llama/Meta-Llama-3-8B/resolve/main/model-00001-of-00004.safetensors"
+
+    config = StreamConfig(
+        auth_token="", # Using public model
+        chunk_size_mb=16, # Small chunk size for the small model
+        chunk_workers=4,
+        gpu_workers=2,
+        disable_cache=True,
+        log_level=4 # 4 = info, 3 = diagnostic, 2 = debug, 1 = debugV, 0 = trace
+    )
+
+    print(f"Streaming {url} into GPU memory (zero-copy)...")
+
+    print("[Python] Entering VajraStreamer context")
+    with VajraStreamer(config) as streamer:
+        # Load the model directly to VRAM
+        print("[Python] Calling streamer.load(url)")
+        tensors = streamer.load(url)
+        print("[Python] streamer.load(url) completed")
+
+        print("\nSuccessfully loaded model!")
+        print(f"Total tensors loaded: {len(tensors)}")
+
+        if len(tensors) > 0:
+            print("\nSample Tensors:")
+            # Print the first few tensors to verify
+            count = 0
+            for name, tensor in tensors.items():
+                if hasattr(tensor, 'dtype'):
+                    print(f" - {name}: shape={tensor.shape}, dtype={tensor.dtype}, device={tensor.device}")
+                else:
+                    print(f" - {name}: shape={tensor.shape}, dtype={tensor.typestr}, device=cuda (raw VajraGPUTensor)")
+                count += 1
+                if count >= 5:
+                    print(" - ...")
+                    break
+
+    print("\nExiting with block. All VRAM will be freed automatically.")
+    print("[Python] Context manager exited successfully")
+
+if __name__ == "__main__":
+    main()

vajra_streamer-0.0.13/pyproject.toml

@@ -0,0 +1,25 @@
+[build-system]
+requires = ["setuptools>=80.0", "setuptools-scm[simple]>=9.2", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "vajra-streamer"
+dynamic = ["version"]
+description = "A fast zero-copy PyTorch tensor streamer"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.8"
+dependencies = [
+    "torch>=2.0.0"
+]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["vajra"]
+
+[tool.setuptools.package-data]
+"vajra" = ["*.so", "*.dylib", "*.dll"]
+
+[tool.setuptools_scm]
+root = ".."

vajra_streamer-0.0.13/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

vajra_streamer-0.0.13/setup.py

@@ -0,0 +1,36 @@
+import platform
+import sys
+
+from setuptools import Distribution, setup
+
+# Guard must run before setuptools does any work. This sdist intentionally
+# bundles a prebuilt Linux x86_64 binary (libvajra.so). It cannot be built
+# on any other platform. Fail loudly here rather than with a confusing
+# linker or import error later.
+if not (sys.platform == "linux" and platform.machine() == "x86_64"):
+    raise RuntimeError(
+        "vajra ships a prebuilt Linux x86_64 shared library (libvajra.so). "
+        "This sdist cannot be built on your platform. "
+        "Install the pre-built wheel instead:\n\n"
+        "    pip install vajra-streamer\n\n"
+        "If no wheel is available for your platform, "
+        "please open an issue at https://github.com/your-org/vibe-streamer"
+    )
+
+
+class BinaryDistribution(Distribution):
+    """Force setuptools to treat this as a platform-specific wheel.
+
+    setuptools has no declarative flag in pyproject.toml for this case —
+    where a pre-compiled .so is injected from outside the Python build
+    system. Overriding has_ext_modules() is the standard escape hatch.
+    Without this, `python -m build` produces a `*-any.whl` (pure-Python
+    tag), which PyPI would serve to Windows/macOS users who can't run .so
+    files.
+    """
+
+    def has_ext_modules(self):
+        return True
+
+
+setup(distclass=BinaryDistribution)

vajra_streamer-0.0.13/vajra/__init__.py

@@ -0,0 +1,249 @@
+import ctypes
+import os
+from dataclasses import dataclass
+from typing import Dict, Optional, Any, Callable
+import torch
+
+# Load the shared library
+import platform
+
+_current_dir = os.path.dirname(os.path.abspath(__file__))
+
+_ext = ".dylib" if platform.system() == "Darwin" else ".so"
+# 1. First, check inside the installed package directory (where the wheel puts it)
+_lib_path = os.path.join(_current_dir, f"libvajra{_ext}")
+
+# 2. Fallbacks for local development if the .so isn't copied yet
+if not os.path.exists(_lib_path):
+    _lib_path = os.path.join(os.path.dirname(_current_dir), f"libvajra{_ext}")
+if not os.path.exists(_lib_path):
+    _lib_path = os.path.join(os.path.dirname(os.path.dirname(_current_dir)), f"libvajra{_ext}")
+
+try:
+    _lib = ctypes.CDLL(_lib_path, mode=ctypes.RTLD_GLOBAL)
+except OSError:
+    # If not in the parent dir, fallback to standard library search paths
+    _lib = ctypes.CDLL(f"libvajra{_ext}", mode=ctypes.RTLD_GLOBAL)
+
+# 1. Define the C-ABI Structs
+class CTensor(ctypes.Structure):
+    _fields_ = [
+        ("name", ctypes.c_char_p),
+        ("gpu_ptr", ctypes.c_void_p),
+        ("shape", ctypes.POINTER(ctypes.c_long)),
+        ("ndim", ctypes.c_size_t),
+        ("dtype", ctypes.c_char_p)
+    ]
+
+class CStreamConfig(ctypes.Structure):
+    _fields_ = [
+        ("auth_token", ctypes.c_char_p),
+        ("chunk_size", ctypes.c_size_t),
+        ("chunk_workers", ctypes.c_uint),
+        ("gpu_workers", ctypes.c_int),
+        ("disable_cache", ctypes.c_bool),
+        ("log_level", ctypes.c_int)
+    ]
+
+# Callback typedef: typedef void (*ProgressCallback)(size_t bytes_downloaded, size_t total_bytes)
+ProgressCallbackType = ctypes.CFUNCTYPE(None, ctypes.c_size_t, ctypes.c_size_t)
+
+# Configure the main entry point: 
+# int stream_model(const char* url, StreamConfig config, CTensor** out_tensors, size_t* num_tensors)
+_lib.stream_model.argtypes = [
+    ctypes.c_char_p,                   # url
+    ctypes.POINTER(CStreamConfig),     # config
+    ctypes.POINTER(ctypes.POINTER(CTensor)), # out_tensors
+    ctypes.POINTER(ctypes.c_size_t)    # num_tensors
+]
+_lib.stream_model.restype = ctypes.c_int
+
+# Tracker control functions
+_lib.vajra_tracker_pause.argtypes = []
+_lib.vajra_tracker_pause.restype = None
+
+_lib.vajra_tracker_resume.argtypes = []
+_lib.vajra_tracker_resume.restype = None
+
+# Configure cleanup function: void free_model_memory(CTensor* tensors_array, size_t num_tensors)
+_lib.free_model_memory.argtypes = [ctypes.POINTER(CTensor), ctypes.c_size_t]
+_lib.free_model_memory.restype = None
+
+# Configure teardown function: void vajra_teardown()
+_lib.vajra_teardown.argtypes = []
+_lib.vajra_teardown.restype = None
+
+import atexit
+import os
+
+# Register os._exit(0) FIRST so it runs LAST.
+# This prevents glibc from running Dlang's .fini_array (rt_term), which segfaults 
+# due to vibe.d daemon threads sleeping in epoll_wait when the process exits.
+atexit.register(os._exit, 0)
+
+
+# 2. Python Config Dataclass
+@dataclass
+class StreamConfig:
+    auth_token: str = ""
+    chunk_size_mb: int = 64
+    chunk_workers: int = 16
+    gpu_workers: int = 3
+    disable_cache: bool = False
+    log_level: int = 4 # default to info (vibe-d log level 4)
+
+    def to_c_struct(self) -> CStreamConfig:
+        return CStreamConfig(
+            auth_token=self.auth_token.encode('utf-8') if self.auth_token else None,
+            chunk_size=self.chunk_size_mb * 1024 * 1024,
+            chunk_workers=self.chunk_workers,
+            gpu_workers=self.gpu_workers,
+            disable_cache=self.disable_cache,
+            log_level=self.log_level
+        )
+
+# 3. Zero-Copy Tensor Wrapper
+class VajraGPUTensor:
+    def __init__(self, c_tensor: CTensor):
+        self._gpu_ptr = c_tensor.gpu_ptr
+        
+        # Safely decode the name 
+        try:
+            self.name = c_tensor.name.decode('utf-8', errors='ignore')
+        except Exception:
+            self.name = "unknown_tensor"
+        
+        # Extract shape
+        self.shape = tuple(c_tensor.shape[i] for i in range(c_tensor.ndim))
+        
+        # Map dtype string to numpy/cuda_array_interface format
+        dtype_str = c_tensor.dtype.decode('utf-8')
+        # Map Safetensors dtypes to __cuda_array_interface__ typestrings
+        self._typestr_map = {
+            "F32": "<f4",
+            "F16": "<f2",
+            "BF16": "<V2", # BF16 doesn't have a standard typestr, often handled specially by torch
+            "I64": "<i8",
+            "I32": "<i4",
+            "I16": "<i2",
+            "I8": "|i1",
+            "U8": "|u1",
+            "BOOL": "|b1",
+        }
+        self.typestr = self._typestr_map.get(dtype_str, "<f4")
+        self.original_dtype = dtype_str
+
+    @property
+    def __cuda_array_interface__(self) -> Dict[str, Any]:
+        """PyTorch zero-copy protocol"""
+        return {
+            "shape": self.shape,
+            "typestr": self.typestr,
+            "data": (self._gpu_ptr, False), # (ptr, read_only=False)
+            "version": 3
+        }
+
+
+# 4. Main Context Manager Streamer
+class VajraStreamer:
+    def __init__(self, config: Optional[StreamConfig] = None):
+        self.config = config or StreamConfig()
+        self._loaded_tensors: list[VajraGPUTensor] = []
+
+    @staticmethod
+    def pause_vram_tracking():
+        """Pause the VRAM tracker to stop logging allocations"""
+        _lib.vajra_tracker_pause()
+
+    @staticmethod
+    def resume_vram_tracking():
+        """Resume the VRAM tracker to start logging allocations again"""
+        _lib.vajra_tracker_resume()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Destroy the VRAM Arenas and the CPU Metadata map at once"""
+        print("[Python] Entering __exit__ block")
+        ptr = getattr(self, '_out_tensors_ptr', None)
+        num = getattr(self, '_num_tensors', None)
+        print(f"[Python] ptr={bool(ptr)}, num={bool(num)}")
+        
+        if ptr is not None and num is not None:
+            print("[Python] Calling _lib.free_model_memory...")
+            _lib.free_model_memory(ptr, num.value)
+            print("[Python] _lib.free_model_memory returned")
+
+            # Prevent double-freeing if __exit__ is called multiple times
+            self._out_tensors_ptr = None
+            self._num_tensors = None
+        else:
+            print("[Python] Did not call free_model_memory because ptr or num is None")
+
+    def _default_progress(self, downloaded: int, total: int):
+        # Basic progress callback if none provided
+        if total > 0:
+            percent = (downloaded / total) * 100
+            print(f"\rStreaming: {downloaded / 1024 / 1024:.2f}MB / {total / 1024 / 1024:.2f}MB ({percent:.1f}%)", end="")
+            if downloaded == total:
+                print() # New line when done
+
+    def load(self, url: str) -> Dict[str, torch.Tensor]:
+        """
+        Stream the model and return a dictionary of zero-copy PyTorch tensors mapping tensor names to VRAM.
+        """
+        url_bytes = url.encode('utf-8')
+        c_config = self.config.to_c_struct()
+
+        # Output pointers
+        out_tensors_ptr = ctypes.POINTER(CTensor)()
+        num_tensors = ctypes.c_size_t(0)
+
+        # Call Dlang
+        error_code = _lib.stream_model(
+            url_bytes,
+            ctypes.byref(c_config),
+            ctypes.byref(out_tensors_ptr),
+            ctypes.byref(num_tensors)
+        )
+
+        # Handle Errors
+        if error_code != 0:
+            if error_code == 1:
+                raise ConnectionError(f"Network error while streaming {url}")
+            elif error_code == 2:
+                raise MemoryError(f"CUDA Out Of Memory while allocating VRAM for {url}")
+            else:
+                raise RuntimeError(f"VajraStreamer failed with generic error code {error_code}")
+
+        # Save the pointers to the class instance so __exit__ can clean them up later
+        self._out_tensors_ptr = out_tensors_ptr
+        self._num_tensors = num_tensors
+
+        # Wrap the returned structs
+        torch_tensors = {}
+        for i in range(num_tensors.value):
+            c_tensor = out_tensors_ptr[i]
+            vibe_tensor = VajraGPUTensor(c_tensor)
+            
+            # Adopt into PyTorch natively (Zero-Copy)
+            # Special handling for bfloat16 as it's not standard in cuda_array_interface
+            if vibe_tensor.original_dtype == "BF16":
+                # PyTorch doesn't officially support bf16 __cuda_array_interface__ via typestr
+                # We do a workaround by creating a view on Int16 data
+                vibe_tensor.typestr = "<i2" 
+                try:
+                    i16_tensor = torch.as_tensor(vibe_tensor, device='cuda')
+                    torch_tensors[vibe_tensor.name] = i16_tensor.view(torch.bfloat16)
+                except Exception as e:
+                    print(f"Skipping torch conversion for {vibe_tensor.name} due to error: {e}")
+                    torch_tensors[vibe_tensor.name] = vibe_tensor
+            else:
+                try:
+                    torch_tensors[vibe_tensor.name] = torch.as_tensor(vibe_tensor, device='cuda')
+                except Exception as e:
+                    print(f"Skipping torch conversion for {vibe_tensor.name} due to error: {e}")
+                    torch_tensors[vibe_tensor.name] = vibe_tensor
+
+        return torch_tensors

vajra_streamer-0.0.13/vajra_streamer.egg-info/PKG-INFO

@@ -0,0 +1,43 @@
+Metadata-Version: 2.4
+Name: vajra-streamer
+Version: 0.0.13
+Summary: A fast zero-copy PyTorch tensor streamer
+License-Expression: MIT
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0.0
+Dynamic: license-file
+
+# Vajra
+
+A fast zero-copy PyTorch tensor streamer powered by Dlang.
+
+## Usage
+
+Once installed, you can use Vajra in any Python script or Jupyter Notebook:
+
+```python
+import torch
+from vajra import VajraStreamer, StreamConfig
+
+# Configure the streamer (optional, uses defaults otherwise)
+config = StreamConfig(
+    auth_token="hf_YOUR_TOKEN", # Required for gated models like Llama 3
+    chunk_size_mb=64,
+    chunk_workers=16,
+    gpu_workers=3,
+    disable_cache=False
+)
+
+# Context manager ensures VRAM is freed when done
+with VajraStreamer(config) as streamer:
+    # Pass the URL or a Hugging Face repo ID
+    # (e.g. "meta-llama/Meta-Llama-3-8B")
+    tensors = streamer.load("meta-llama/Meta-Llama-3-8B")
+
+    # 'tensors' is a dictionary mapping tensor names to zero-copy PyTorch tensors
+    # that are backed directly by the downloaded GPU memory.
+    for name, tensor in tensors.items():
+        print(f"Tensor: {name}, Shape: {tensor.shape}, Dtype: {tensor.dtype}")
+```

vajra_streamer-0.0.13/vajra_streamer.egg-info/SOURCES.txt

@@ -0,0 +1,13 @@
+LICENSE
+MANIFEST.in
+README.md
+pyproject.toml
+setup.py
+example/test.py
+vajra/__init__.py
+vajra/libvajra.so
+vajra_streamer.egg-info/PKG-INFO
+vajra_streamer.egg-info/SOURCES.txt
+vajra_streamer.egg-info/dependency_links.txt
+vajra_streamer.egg-info/requires.txt
+vajra_streamer.egg-info/top_level.txt

vajra_streamer-0.0.13/vajra_streamer.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

vajra_streamer-0.0.13/vajra_streamer.egg-info/requires.txt

@@ -0,0 +1 @@
+torch>=2.0.0

vajra_streamer-0.0.13/vajra_streamer.egg-info/top_level.txt

@@ -0,0 +1 @@
+vajra