gradexp 0.1.0__py3-none-any.whl

gradexp/__init__.py

@@ -0,0 +1,2 @@
+from .client import init, log, finish
+__version__ = '0.1.0'

gradexp/auth.py

@@ -0,0 +1,103 @@
+import os
+import appdirs
+import json
+import webbrowser
+import click
+import requests
+
+APP_NAME = "gradexp"
+APP_AUTHOR = "GradientExplorer"
+BASE_URL = "gradient-explorer.xyz"
+
+def get_config_dir():
+    return appdirs.user_config_dir(APP_NAME, APP_AUTHOR)
+
+def get_token_path():
+    config_dir = get_config_dir()
+    return os.path.join(config_dir, "secrets.json")
+
+def save_token(token):
+    config_dir = get_config_dir()
+    if not os.path.exists(config_dir):
+        os.makedirs(config_dir)
+    
+    token_path = get_token_path()
+    with open(token_path, 'w') as f:
+        json.dump({"api_key": token}, f)
+    
+    # Set permissions to be readable only by user (0600)
+    os.chmod(token_path, 0o600)
+
+def load_token():
+    token_path = get_token_path()
+    if not os.path.exists(token_path):
+        return None
+    
+    try:
+        with open(token_path, 'r') as f:
+            data = json.load(f)
+            return data.get("api_key")
+    except (json.JSONDecodeError, IOError):
+        return None
+
+def validate_api_key(api_key, debug=False):
+    try:
+        url = f"https://api.{BASE_URL}/api/v1/me"
+        headers = {"Authorization": f"Bearer {api_key}"}
+        
+        if debug:
+            click.echo(f"DEBUG: Request URL: {url}")
+            click.echo(f"DEBUG: Request Headers: Authorization: Bearer {api_key[:4]}...{api_key[-4:]}")
+
+        response = requests.get(
+            url,
+            headers=headers
+        )
+        
+        if debug:
+            click.echo(f"DEBUG: Response Status: {response.status_code}")
+            click.echo(f"DEBUG: Response Body: {response.text}")
+
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        if debug:
+            click.echo(f"DEBUG: Request failed: {e}")
+        return None
+
+def login_flow(debug=False):
+    """
+    Initiates the login flow:
+    1. Opens the browser to the authorization URL.
+    2. Prompts the user to paste the API key.
+    3. Validates the API key.
+    4. Saves the API key securely.
+    """
+    auth_url = f"https://{BASE_URL}/authorize" 
+    
+    click.echo(f"Opening {auth_url} in your default browser...")
+    webbrowser.open(auth_url)
+    
+    api_key = click.prompt("Please paste your API key here", hide_input=True)
+    
+    if not api_key:
+        click.echo("No API key provided. Login failed.")
+        return
+
+    user_info = validate_api_key(api_key, debug=debug)
+
+    if user_info:
+        save_token(api_key)
+        click.secho("Successfully logged in. Your API key is saved.", fg="green")
+    else:
+        # Construct error message
+        click.secho("Error: Invalid API Key", fg="red")
+
+def check_login():
+    token = load_token()
+    if token:
+        click.echo(f"Logged in with token: {token[:4]}..." + "*" * 10)
+        return True
+    else:
+        click.echo("Not logged in. Please run `gradexp login`.")
+        return False

gradexp/cli.py

@@ -0,0 +1,21 @@
+import click
+from . import auth
+
+@click.group()
+@click.version_option()
+@click.option('--debug', is_flag=True, help="Enable debug logging")
+@click.pass_context
+def main(ctx, debug):
+    """Gradient Explorer CLI"""
+    ctx.ensure_object(dict)
+    ctx.obj['DEBUG'] = debug
+
+@main.command()
+@click.option('--debug', is_flag=True, help="Enable debug logging")
+@click.pass_context
+def login(ctx, debug):
+    """Log in to Gradient Explorer"""
+    auth.login_flow(debug=debug or ctx.obj.get('DEBUG', False))
+
+if __name__ == '__main__':
+    main()

gradexp/client.py

@@ -0,0 +1,468 @@
+import os
+import sys
+import signal
+import atexit
+import base64
+import json
+import uuid
+import time
+import threading
+import queue
+import shutil
+import appdirs
+import requests
+import numpy as np
+
+from . import auth
+
+# Calm red: \033[38;5;174m
+# Reset: \033[0m
+GRADEXP_PREFIX = "\033[38;5;174mgradexp\033[0m: "
+
+def _term_log(message, end="\n"):
+    """
+    Prints a message with a colored 'gradexp:' prefix.
+    """
+    # Handle multi-line messages or messages starting with newline
+    if isinstance(message, str) and message.startswith("\n"):
+        print(f"\n{GRADEXP_PREFIX}{message[1:]}", end=end)
+    else:
+        print(f"{GRADEXP_PREFIX}{message}", end=end)
+
+class Client:
+    def __init__(self):
+        self.api_key = None
+        self.run_id = None
+        self.project_id = None
+        self.tensors = {} # {name: {"id": uuid, "step": 0}}
+        self.active = False
+        self._session = requests.Session()
+        self._upload_queue = queue.Queue()
+        self._stop_event = threading.Event()
+        self._worker_thread = None
+        self._cache_dir = None
+        self._interrupted = False
+        self._original_sigint_handler = None
+        self._original_sigterm_handler = None
+        
+        # Progress tracking
+        self._progress_lock = threading.Lock()
+        self._bytes_uploaded = 0
+        self._tensors_uploaded = 0
+        self._total_tensors_queued = 0
+        self._upload_start_time = None
+
+    def init(self, project_name=None, run_name=None):
+        """
+        Initialize a new run session.
+        """
+        # Try to pull from wandb if not provided
+        if project_name is None or run_name is None:
+            if "wandb" in sys.modules:
+                import wandb
+                if wandb.run:
+                    if project_name is None:
+                        project_name = wandb.run.project
+                    if run_name is None:
+                        run_name = wandb.run.name
+                        
+        # Fallback to defaults if still None
+        if project_name is None:
+            project_name = "default"
+        if run_name is None:
+            run_name = "default-run"
+
+        self.api_key = auth.load_token()
+        if not self.api_key:
+            raise RuntimeError("Not logged in. Please run 'gradexp login' first.")
+
+        current_run_name = run_name
+        while True:
+            payload = {
+                "project_name": project_name,
+                "run_name": current_run_name
+            }
+
+            try:
+                url = f"https://api.{auth.BASE_URL}/api/v1/runs"
+                headers = {
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json"
+                }
+                
+                response = self._session.post(url, headers=headers, json=payload)
+                response.raise_for_status()
+                
+                data = response.json()
+                self.run_id = data.get("run_id")
+                self.project_id = data.get("project_id")
+                self.tensors = {} 
+                self.active = True
+                
+                self._cache_dir = appdirs.user_cache_dir(auth.APP_NAME, auth.APP_AUTHOR)
+                os.makedirs(self._cache_dir, exist_ok=True)
+                
+                self._stop_event.clear()
+                self._worker_thread = threading.Thread(target=self._upload_worker, daemon=True)
+                self._worker_thread.start()
+
+                base_url = f"https://{auth.BASE_URL}"
+                run_url = f"{base_url}/?project={self.project_id}&run={self.run_id}"
+                _term_log(f"Initialized. See run progress at {run_url}")
+
+                atexit.register(self.finish)
+
+                try:
+                    self._original_sigint_handler = signal.getsignal(signal.SIGINT)
+                    self._original_sigterm_handler = signal.getsignal(signal.SIGTERM)
+                    signal.signal(signal.SIGINT, self._handle_interrupt)
+                    signal.signal(signal.SIGTERM, self._handle_interrupt)
+                except ValueError:
+                    pass
+
+                break # Success
+
+            except requests.exceptions.HTTPError as e:
+                if e.response.status_code == 409:
+                    try:
+                        err_data = e.response.json()
+                        if "suggested_name" in err_data:
+                            current_run_name = err_data["suggested_name"]
+                            continue
+                    except Exception:
+                        pass
+                    
+                    if "-" in current_run_name and current_run_name.rsplit("-", 1)[1].isdigit():
+                        parts = current_run_name.rsplit("-", 1)
+                        base = parts[0]
+                        num = int(parts[1]) + 1
+                        current_run_name = f"{base}-{num}"
+                    else:
+                        current_run_name = f"{current_run_name}-2"
+                    continue
+                
+                _term_log(f"Failed to initialize: {e}")
+                try:
+                    _term_log(f"Backend error message: {e.response.text}")
+                except:
+                    pass
+
+                if e.response.status_code == 401:
+                    _term_log("Authentication failed: Unauthorized. Please run 'gradexp login' to authenticate.")
+                    sys.exit(1)
+                else:
+                    raise
+            except requests.exceptions.RequestException as e:
+                _term_log(f"Failed to initialize: {e}")
+                raise
+
+    def _ensure_tensor(self, name, array):
+        """
+        Ensures a tensor exists on the backend. Returns its ID.
+        """
+        if name in self.tensors:
+            return self.tensors[name]["id"]
+
+        try:
+            url = f"https://api.{auth.BASE_URL}/api/v1/runs/{self.run_id}/tensors"
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json"
+            }
+            payload = {
+                "name": name,
+                "dtype": str(array.dtype),
+                "shape": list(array.shape)
+            }
+            response = self._session.post(url, headers=headers, json=payload)
+            response.raise_for_status()
+            
+            data = response.json()
+            tensor_id = data.get("tensor_id")
+            self.tensors[name] = {"id": tensor_id, "step": 0}
+            return tensor_id
+        except Exception as e:
+            _term_log(f"Failed to create tensor '{name}': {e}")
+            return None
+
+    def log(self, tensor, name="default"):
+        """
+        Upload data for the tensor.
+        """
+        if not self.active:
+            _term_log("Not initialized. Please call gradexp.init() first.")
+            return
+
+        # Convert to numpy if needed
+        if not isinstance(tensor, np.ndarray):
+            try:
+                tensor = np.array(tensor)
+            except Exception:
+                _term_log("Could not convert input to numpy array.")
+                return
+
+        # Ensure tensor is registered
+        tensor_id = self._ensure_tensor(name, tensor)
+        if not tensor_id:
+            return
+
+        try:
+            # Full consistency: float32 for now as per previous logic, 
+            # but ideally we respect the tensor's own dtype if we sent that to backend.
+            # Assuming backend handles what we told it in _ensure_tensor.
+            # But previous code cast to float32. Let's stick to float32 for consistency for now unless specified.
+            if tensor.dtype != np.float32:
+                tensor = tensor.astype(np.float32)
+
+            tensor_bytes = tensor.tobytes()
+            
+            # Store locally
+            unique_id = str(uuid.uuid4())
+            file_path = os.path.join(self._cache_dir, f"step_{unique_id}.bin")
+            
+            with open(file_path, "wb") as f:
+                f.write(tensor_bytes)
+            
+            # Get and increment step
+            step_index = self.tensors[name]["step"]
+            self.tensors[name]["step"] += 1
+            
+            # Enqueue for background upload
+            with self._progress_lock:
+                self._total_tensors_queued += 1
+            self._upload_queue.put({
+                "file_path": file_path,
+                "tensor_id": tensor_id,
+                "step": step_index
+            })
+            
+        except Exception as e:
+            _term_log(f"Failed to buffer tensor data: {e}")
+
+    def _update_status(self, status):
+        """
+        Updates the status of the run.
+        Uses PATCH /api/v1/runs/{run_id}
+        """
+        if not self.run_id or not self.api_key:
+            return
+
+        try:
+            url = f"https://api.{auth.BASE_URL}/api/v1/runs/{self.run_id}"
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json"
+            }
+            payload = {"status": status}
+            response = self._session.patch(url, headers=headers, json=payload)
+            response.raise_for_status()
+        except Exception as e:
+            _term_log(f"Failed to update session status to {status}: {e}")
+
+    def _print_progress(self):
+        """
+        Prints a progress bar showing upload status and throughput.
+        """
+        with self._progress_lock:
+            uploaded = self._tensors_uploaded
+            total = self._total_tensors_queued
+            bytes_up = self._bytes_uploaded
+            start_time = self._upload_start_time
+        
+        # Calculate throughput
+        if start_time and bytes_up > 0:
+            elapsed = time.time() - start_time
+            if elapsed > 0:
+                mbps = (bytes_up / (1024 * 1024)) / elapsed
+            else:
+                mbps = 0.0
+        else:
+            mbps = 0.0
+        
+        # Build progress bar
+        bar_width = 30
+        if total > 0:
+            filled = int(bar_width * uploaded / total)
+        else:
+            filled = 0
+        bar = "█" * filled + "░" * (bar_width - filled)
+        
+        sys.stdout.write(f"\r{GRADEXP_PREFIX}[{bar}] {uploaded}/{total} tensors | {mbps:.2f} MB/s ")
+        sys.stdout.flush()
+
+    def _handle_interrupt(self, signum, frame):
+        """
+        Handles SIGINT (Ctrl+C) or SIGTERM.
+        """
+        if not self._interrupted:
+            self._interrupted = True
+            q_size = self._upload_queue.qsize()
+            
+            if q_size == 0:
+                _term_log("\nInterrupt received. No pending uploads. Shutting down...")
+                self.finish("stopped")
+                sys.exit(0)
+            else:
+                _term_log(f"\nInterrupt received. Waiting for {q_size} pending uploads to complete.")
+                _term_log("Press Ctrl+C again to force quit.\n")
+                
+                # Show progress while waiting for uploads
+                try:
+                    while not self._upload_queue.empty() or (self._worker_thread and self._worker_thread.is_alive() and self._tensors_uploaded < self._total_tensors_queued):
+                        self._print_progress()
+                        time.sleep(0.2)
+                    self._print_progress()  # Final update
+                    print()  # Newline after progress bar
+                    self.finish("stopped")
+                    sys.exit(0)
+                except KeyboardInterrupt:
+                    # Second Ctrl+C during progress display - force quit
+                    _term_log("\nForce quitting...")
+                    self._update_status("stopped")
+                    os._exit(1)
+        else:
+            _term_log("\nForce quitting...")
+            # Attempt a quick status update if possible, but don't block
+            try:
+                # Direct status update attempt to backend, skipping the queue
+                # This might fail if the session is already half-closed but worth a try
+                self._update_status("stopped")
+            except:
+                pass
+            os._exit(1) # Immediate exit
+
+    def _upload_worker(self):
+        """
+        Background worker that uploads data from the queue.
+        """
+        while not self._stop_event.is_set() or not self._upload_queue.empty():
+            try:
+                # Use a timeout to occasionally check the stop event
+                item = self._upload_queue.get(timeout=1.0)
+            except queue.Empty:
+                continue
+
+            file_path = item["file_path"]
+            tensor_id = item["tensor_id"]
+            step = item["step"]
+
+            try:
+                if not os.path.exists(file_path):
+                    self._upload_queue.task_done()
+                    continue
+
+                with open(file_path, "rb") as f:
+                    content = f.read()
+                
+                content_len = len(content)
+                
+                # Set upload start time on first upload
+                with self._progress_lock:
+                    if self._upload_start_time is None:
+                        self._upload_start_time = time.time()
+                
+                url = f"https://api.{auth.BASE_URL}/api/v1/tensors/{tensor_id}/step"
+                headers = {
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/octet-stream"
+                }
+                params = {
+                    "step": step,
+                    "index": step # Kept for backward compatibility if needed
+                }
+
+                response = self._session.post(url, headers=headers, data=content, params=params)
+                response.raise_for_status()
+                
+                # Success: update progress and delete the local file
+                with self._progress_lock:
+                    self._bytes_uploaded += content_len
+                    self._tensors_uploaded += 1
+                os.remove(file_path)
+
+            except requests.exceptions.RequestException as e:
+                _term_log(f"Failed to upload tensor step in background: {e}")
+                if os.path.exists(file_path):
+                    os.remove(file_path)
+            except Exception as e:
+                _term_log(f"Unexpected error in upload worker: {e}")
+                if os.path.exists(file_path):
+                    os.remove(file_path)
+            finally:
+                self._upload_queue.task_done()
+
+    def finish(self, status="complete"):
+        """
+        Marks the session as finished and flushes the upload queue.
+        """
+        if not self.active:
+            return
+
+        # Signal the worker to stop after processing remaining items
+        self._stop_event.set()
+        
+        if self._worker_thread and self._worker_thread.is_alive():
+            # Check if there are pending uploads
+            with self._progress_lock:
+                pending = self._total_tensors_queued - self._tensors_uploaded
+            
+            if pending > 0:
+                _term_log("finishing... Waiting for background uploads to complete.")
+                # Show progress while waiting
+                while self._worker_thread.is_alive() and self._tensors_uploaded < self._total_tensors_queued:
+                    self._print_progress()
+                    time.sleep(0.2)
+                self._print_progress()  # Final update
+                print()  # Newline after progress bar
+            else:
+                # No pending uploads, just wait for thread to finish
+                self._worker_thread.join()
+
+        # Update status on backend
+        self._update_status(status)
+
+        # Restore signal handlers
+        try:
+            if self._original_sigint_handler:
+                signal.signal(signal.SIGINT, self._original_sigint_handler)
+            if self._original_sigterm_handler:
+                signal.signal(signal.SIGTERM, self._original_sigterm_handler)
+        except ValueError:
+            pass
+
+        self._interrupted = False
+        self.active = False
+        _term_log(f"Session {status}.")
+
+# Singleton instance
+_client = Client()
+
+def _excepthook(type, value, traceback):
+    if _client.active:
+        _client.finish("stopped")
+    sys.__excepthook__(type, value, traceback)
+
+def init(project_name=None, run_name=None):
+    # Install excepthook if not already installed
+    if sys.excepthook is not _excepthook:
+        # We might want to chain if there's already a custom one?
+        # For simplicity, we just use ours which calls the default __excepthook__.
+        # If user has another custom one, this might override it.
+        # A safer way is to store the previous one.
+        global _original_excepthook
+        _original_excepthook = sys.excepthook
+        
+        def tagged_excepthook(type, value, traceback):
+            if _client.active:
+                _client.finish("stopped")
+            _original_excepthook(type, value, traceback)
+            
+        sys.excepthook = tagged_excepthook
+        
+    _client.init(project_name=project_name, run_name=run_name)
+
+def log(tensor, name="default"):
+    _client.log(tensor, name=name)
+
+def finish(status="complete"):
+    _client.finish(status)

gradexp-0.1.0.dist-info/METADATA

@@ -0,0 +1,86 @@
+Metadata-Version: 2.4
+Name: gradexp
+Version: 0.1.0
+Summary: Gradient Explorer Client Library
+Author-email: Misha Obu <misha@parallel-ocean.xyz>
+Requires-Python: >=3.7
+Requires-Dist: appdirs
+Requires-Dist: click
+Requires-Dist: numpy
+Requires-Dist: requests
+Description-Content-Type: text/markdown
+
+source ./venv/bin/activate 
+
+
+
+This repo is intended to reimplement some key features of wandb for our own purposes. This is the pythonic frontend for our wandb-for-tensors service. 
+
+Repo features:
+pip install gradexp
+gradexp login -> opens webpage -> gets token and stores in in permanent context
+In python file: 
+```
+import gradexp
+gradexp.init("project name")
+
+...
+
+# Automatically starts tracking run id, step, etc
+gradexp.log(TODO ... )
+
+...
+
+```
+
+App-wide features:
+
+1) Local staging + asynchronous streaming
+When you call wandb.log() in your training script, the client SDK:
+
+collects the metric/media payload locally (in memory and also writes to local run files).
+
+hands it off to a separate streaming thread/process that runs alongside your main process.
+
+this uploader process asynchronously batches and sends events over the network. This avoids blocking your training process.
+
+events are queued in memory and written to disk if needed (e.g., offline mode) and then eventually synced.
+
+if you set WANDB_MODE=offline, nothing is sent to the server until a sync is triggered.
+This behavior effectively decouples ingestion from your training loop. It’s not “direct write to bucket” from your app — it goes through this upload pipeline first. 
+Weights & Biases Documentation
+
+2) Upload endpoints and storage target
+On the server side (hosted or self-managed):
+
+incoming streaming events (metrics, history tuples) arrive via W&B’s application backend.
+
+metadata and small event records (like per-step metrics) are stored relationally (MySQL in self-managed reference architecture).
+
+larger blobs (logs, media files, artifacts) are written out to object storage buckets (e.g., S3 or your own BYOB bucket).
+You don’t stream directly into the bucket in small per-log increments — the backend receives the event first and the service layer persists the data. 
+Weights & Biases Documentation
++1
+
+3) How “appending” works
+Object storage (S3/compatible) isn’t a traditional file system — you can’t literally open and append to an existing file like with a local file. Instead:
+
+W&B will write each piece of logged data as a separate object or part of a structured object.
+
+for metric history exports (e.g., after a run completes), W&B creates Parquet exports and pushes those into the bucket as artifacts/history files.
+
+this is a batch write rather than incremental byte-level append.
+The UI and service layer then stitch these pieces together logically for history views.
+(This pattern is common across object-store backed systems — you don’t append bytes to objects at every metric call in practice.) 
+Weights & Biases Documentation
+
+4) Back-end buffering & batching
+Even though metrics are “live” in the UI, there’s batching:
+
+the SDK batches small updates and flushes over HTTP.
+
+on the server side, those are ingested through API layers and persisted into the database or stored as discrete objects.
+
+visualization updates pull from the latest ingested state.
+The “real-time” aspect is achieved by frequent flushes and update propagation — not by writing to a monolithic streaming file.
+W&B’s process model means it tolerates network latency/outages by buffering locally and then syncing when available.

gradexp-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+gradexp/__init__.py,sha256=kZfL0MotOvzx9JGu4rOv_WQkhVcOk4f__8-6M2lWZW4,60
+gradexp/auth.py,sha256=I3idBycevSCYC02znm13cRvqIJdYnEndMZz6JfToHHw,2915
+gradexp/cli.py,sha256=a9Y4b7o09oKcFBB2N9eB_CRB8tKuKcLjTVYAQYR_nGY,540
+gradexp/client.py,sha256=vqzPhzIGjNuIE56k4p-JWVxKlX7m0NtOcrDoGQ1zQ8U,17336
+gradexp-0.1.0.dist-info/METADATA,sha256=_tsPJaRgHM6oz0oAfqkqyE_dQaxNFRFjWgOFMLLlE80,3550
+gradexp-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+gradexp-0.1.0.dist-info/entry_points.txt,sha256=8Azz89F2lTdJNq0x8lItv713gIj6YSk4ePS_C5gncEk,45
+gradexp-0.1.0.dist-info/RECORD,,

gradexp-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

gradexp-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+gradexp = gradexp.cli:main