gradexp 0.1.0__tar.gz

gradexp-0.1.0/.agent/workflows/run_tests.md

@@ -0,0 +1,11 @@
+---
+description: Run gradexp tests
+---
+
+Run the unit tests for the gradexp client.
+
+// turbo
+1. Run the tests using python unittest
+   ```bash
+   python3 -m unittest discover tests
+   ```

gradexp-0.1.0/.gitignore

@@ -0,0 +1,36 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Virtual Env
+venv/
+env/
+ENV/
+
+# Mac
+.DS_Store
+
+# IDEs
+.idea/
+.vscode/
+
+wandb

gradexp-0.1.0/MANIFEST.in

@@ -0,0 +1 @@
+include README.md

gradexp-0.1.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,74 @@
+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/gradexp/__init__.py

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

gradexp-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-0.1.0/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-0.1.0/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/pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "gradexp"
+version = "0.1.0"
+description = "Gradient Explorer Client Library"
+readme = "README.md"
+requires-python = ">=3.7"
+authors = [
+    { name = "Misha Obu", email = "misha@parallel-ocean.xyz" }
+]
+dependencies = [
+    "click",
+    "requests",
+    "appdirs",
+    "numpy",
+]
+
+[project.scripts]
+gradexp = "gradexp.cli:main"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"

gradexp-0.1.0/schema.ts

@@ -0,0 +1,62 @@
+/**
+ * GradExp Schema Definitions -- Documentation
+ * 
+ * Defines the data structures for the new GradExp API (Projects -> Runs -> Tensors -> Steps).
+ * 
+ * Endpoints:
+ * - POST /api/v1/runs (Create Run, presumably handles Project creation/lookup)
+ * - POST /api/v1/runs/{run_id}/tensors (Create Tensor)
+ * - POST /api/v1/tensors/{tensor_id}/step (Log Tensor Value)
+ */
+
+/**
+ * Payload for creating a run.
+ * POST /api/v1/runs
+ */
+export interface CreateRunRequest {
+    project_name: string;
+    run_name: string;
+}
+
+export interface CreateRunResponse {
+    run_id: string;
+    project_id: string;
+    name: string;
+    status: 'active' | 'complete';
+}
+
+/**
+ * Payload for creating a tensor.
+ * POST /api/v1/runs/{run_id}/tensors
+ */
+export interface CreateTensorRequest {
+    name: string;
+    dtype: string;
+    shape: number[];
+}
+
+export interface CreateTensorResponse {
+    tensor_id: string;
+    run_id: string;
+    name: string;
+    dtype: string;
+    shape: number[];
+}
+
+/**
+ * Payload for logging a step value.
+ * POST /api/v1/tensors/{tensor_id}/step?step={step}
+ * 
+ * The body should be the raw binary data of the tensor.
+ */
+// export interface LogStepRequest { ... } (Removed as it's no longer JSON)
+
+
+/**
+ * Payload for updating run status.
+ * PATCH /api/v1/runs/{run_id}
+ */
+export interface UpdateRunStatusRequest {
+    status: 'complete' | 'stopped' | 'stalled';
+}
+

gradexp-0.1.0/setup.txt

@@ -0,0 +1,6 @@
+local: 
+
+python3 -m venv venv
+source venv/bin/activate
+pip install -e .
+gradexp login

gradexp-0.1.0/test-scripts/init-log.py

@@ -0,0 +1,7 @@
+import gradexp
+import numpy as np
+
+gradexp.init()
+for i in range(10):
+    data = np.random.rand(3, 224, 224).astype(np.float32)
+    gradexp.log(data)

gradexp-0.1.0/test-scripts/mock_verify.py

@@ -0,0 +1,91 @@
+import unittest
+from unittest.mock import MagicMock, patch
+import json
+import base64
+import numpy as np
+import time
+import os
+import requests
+
+# Set a dummy API key for the test
+os.environ["GRADIENT_EXPLORER_API_KEY"] = "dummy_key"
+
+import gradexp
+from gradexp import auth, client
+
+class TestGradExpClient(unittest.TestCase):
+    @patch('gradexp.auth.load_token', return_value="dummy_key")
+    @patch('requests.Session.post')
+    @patch('requests.Session.patch')
+    def test_init_and_log(self, mock_patch, mock_post, mock_load_token):
+        # Setup mock responses
+        # 1. init -> runs (409 Conflict)
+        mock_response_conflict = MagicMock()
+        mock_response_conflict.status_code = 409
+        mock_response_conflict.raise_for_status.side_effect = requests.exceptions.HTTPError("Conflict", response=mock_response_conflict)
+        
+        # 2. init -> runs (success with new name)
+        mock_response_init = MagicMock()
+        mock_response_init.json.return_value = {"run_id": "run_123", "project_id": "proj_123"}
+        mock_response_init.status_code = 200
+        
+        # 3. create tensor
+        mock_response_tensor = MagicMock()
+        mock_response_tensor.json.return_value = {"tensor_id": "tensor_123"}
+        mock_response_tensor.status_code = 200
+
+        # 4. log step
+        mock_response_step = MagicMock()
+        mock_response_step.status_code = 200
+        
+        # 5. status update (PATCH)
+        mock_response_patch = MagicMock()
+        mock_response_patch.status_code = 200
+
+        mock_post.side_effect = [mock_response_conflict, mock_response_init, mock_response_tensor, mock_response_step]
+        mock_patch.return_value = mock_response_patch
+
+        # Reset client singleton
+        client._client = client.Client()
+
+        # RUN INIT
+        gradexp.init(project_name="my-project", run_name="my-run")
+        
+        # Verify call 1 (Conflict)
+        call_args_list = mock_post.call_args_list
+        self.assertTrue(len(call_args_list) >= 2)
+        
+        # First attempt: my-run
+        args1, kwargs1 = call_args_list[0]
+        self.assertIn("/api/v1/runs", args1[0])
+        self.assertEqual(kwargs1['json']['run_name'], "my-run")
+        
+        # Second attempt: my-run-2
+        args2, kwargs2 = call_args_list[1]
+        self.assertIn("/api/v1/runs", args2[0])
+        self.assertEqual(kwargs2['json']['run_name'], "my-run-2")
+
+        # RUN LOG
+        data = np.zeros((2, 2), dtype=np.float32)
+        gradexp.log(data, name="my-tensor")
+        
+        # Wait for worker thread
+        time.sleep(1)
+        
+        # Verify LOG calls...
+        # Check create tensor call
+        # note: index 2 because first two were init
+        args_tensor, kwargs_tensor = call_args_list[2]
+        self.assertIn("/api/v1/runs/run_123/tensors", args_tensor[0])
+        
+        # FINISH (triggers status update)
+        gradexp.finish()
+        
+        # Verify PATCH call
+        self.assertTrue(mock_patch.called)
+        patch_args, patch_kwargs = mock_patch.call_args
+        self.assertIn("/api/v1/runs/run_123", patch_args[0])
+        self.assertEqual(patch_kwargs['json']['status'], "complete")
+
+if __name__ == '__main__':
+    unittest.main()

gradexp-0.1.0/test-scripts/with-wandb.py

@@ -0,0 +1,20 @@
+import gradexp
+import numpy as np
+import wandb
+import os
+
+
+
+# 1. Initialize wandb
+wandb.init(project="wandb-gradexp-integration", name="test-run-metadata")
+wandb.log({"test_score": 0.95})
+
+# 2. Initialize gradexp without arguments - it should pull from wandb
+gradexp.init()
+
+
+# 3. Log some data to gradexp
+for i in range(5):
+    data = np.random.rand(10, 10).astype(np.float32)
+    gradexp.log(data, name=f"tensor_{i}")
+

gradexp-0.1.0/tests/test_client_status.py

@@ -0,0 +1,62 @@
+import sys
+import unittest
+from unittest.mock import MagicMock, patch
+import gradexp.client as client
+import gradexp
+
+class TestStatusUpdates(unittest.TestCase):
+    def setUp(self):
+        # Reset client state
+        client._client = client.Client()
+        client._client.api_key = "fake_key" # Mock logged in
+        client._client.data_instance_id = "fake_instance_id"
+        client._client.active = True
+        # Mock requests
+        self.mock_session = MagicMock()
+        client._client._session = self.mock_session
+        
+        # Mock auth.BASE_URL
+        self.auth_patcher = patch('gradexp.auth.BASE_URL', "test.url")
+        self.auth_patcher.start()
+
+    def tearDown(self):
+        self.auth_patcher.stop()
+
+    def test_finish_complete(self):
+        client._client.finish() # Default status="complete"
+        self.mock_session.post.assert_called_with(
+            "https://api.test.url/api/v1/data-instances/fake_instance_id/complete",
+            headers={"Authorization": "Bearer fake_key", "Content-Type": "application/json"}
+        )
+
+    def test_finish_stopped(self):
+        client._client.finish(status="stopped")
+        self.mock_session.post.assert_called_with(
+            "https://api.test.url/api/v1/data-instances/fake_instance_id/stopped",
+            headers={"Authorization": "Bearer fake_key", "Content-Type": "application/json"}
+        )
+        
+    def test_excepthook_installed(self):
+        # Check if excepthook wraps the original
+        original = sys.excepthook
+        
+        # Force reinstall for test (since init might have been called or not)
+        client.init(project_name="test", run_name="test")
+        
+        self.assertNotEqual(sys.excepthook, original)
+        self.assertTrue(hasattr(sys.excepthook, "__code__")) # simple check
+        
+        # Test that calling it triggers finish("stopped")
+        client._client.finish = MagicMock()
+        try:
+            # We don't want to actually crash, just call the hook
+            sys.excepthook(RuntimeError, RuntimeError("test"), None)
+        except:
+            pass
+            
+        client._client.finish.assert_called_with("stopped")
+
+if __name__ == '__main__':
+    # Need to prevent init from running real network requests during import/setup if possible
+    # But init is only called explicitly.
+    unittest.main()

gradexp-0.1.0/tests/test_gradexp_mock.py

@@ -0,0 +1,93 @@
+import unittest
+from unittest.mock import patch, MagicMock
+import numpy as np
+import sys
+import os
+
+# Ensure we can import gradexp from current directory
+sys.path.append(os.getcwd())
+
+import gradexp
+
+class TestGradExp(unittest.TestCase):
+    @patch('gradexp.auth.load_token')
+    @patch('requests.Session.post')
+    def test_full_flow(self, mock_post, mock_load_token):
+        # Setup mocks
+        mock_load_token.return_value = "test_api_key"
+        
+        # Mock Init Response (201 Created)
+        init_response = MagicMock()
+        init_response.status_code = 201
+        init_response.json.return_value = {
+            "run_id": "run_uuid_123",
+            "project_id": "project_uuid_456",
+            "name": "test-run",
+            "status": "active"
+        }
+        
+        # Mock Upload Response (200 OK)
+        upload_response = MagicMock()
+        upload_response.status_code = 200
+        upload_response.json.return_value = {"status": "success"}
+        
+        # Configure side_effect to return different responses based on call order or inspection
+        # But simpler: just cycle them if we know the order, or make the mock return init first then upload
+        # Since we use the same session object, requests.Session.post is called
+        mock_post.side_effect = [init_response, upload_response]
+
+        # 1. run gradexp.init()
+        gradexp.init(project_name="test-project", run_name="test-run")
+        
+        # Verify init call
+        self.assertEqual(mock_post.call_count, 1)
+        init_call_args = mock_post.call_args_list[0]
+        self.assertIn('/api/v1/runs', init_call_args[0][0])
+        self.assertEqual(init_call_args[1]['headers']['Authorization'], 'Bearer test_api_key')
+        
+        payload = init_call_args[1]['json']
+        self.assertEqual(payload['project_name'], 'test-project')
+        self.assertEqual(payload['run_name'], 'test-run')
+        
+        # Mock Tensor Creation Response (200 OK)
+        tensor_response = MagicMock()
+        tensor_response.status_code = 200
+        tensor_response.json.return_value = {"tensor_id": "tensor_uuid_123"}
+        
+        # Update side_effect for subsequent calls
+        # 2nd call: _ensure_tensor (POST /runs/.../tensors)
+        # 3rd call: upload (POST /tensors/.../step)
+        mock_post.side_effect = [tensor_response, upload_response]
+
+        # 2. run gradexp.log()
+        data = np.random.rand(3, 224, 224).astype(np.float32)
+        gradexp.log(data)
+        
+        # Wait for potential background thread (though in this mock setup it might be immediate if we're lucky or we might need to wait)
+        # Actually in the current client, it's a background thread.
+        # We might need to join the queue or something.
+        # For simplicity in this mock test, let's just wait a bit if needed or mock the queue.
+        # But wait, mock_post.side_effect will be called by the background thread.
+        
+        import time
+        max_wait = 2.0
+        start = time.time()
+        while mock_post.call_count < 3 and time.time() - start < max_wait:
+            time.sleep(0.1)
+
+        # Verify calls
+        self.assertEqual(mock_post.call_count, 3) 
+        
+        # Verify tensor creation call
+        tensor_call_args = mock_post.call_args_list[1]
+        self.assertIn('/api/v1/runs/run_uuid_123/tensors', tensor_call_args[0][0])
+        
+        # Verify log call
+        log_call_args = mock_post.call_args_list[2]
+        self.assertIn('/api/v1/tensors/tensor_uuid_123/step', log_call_args[0][0])
+        self.assertEqual(log_call_args[1]['params']['step'], 0)
+        
+        print("\nTest passed successfully!")
+
+if __name__ == '__main__':
+    unittest.main()