From f0e4ae967fa0cbb05f9f654f310726a02536261e Mon Sep 17 00:00:00 2001 From: Lucas Tonon Date: Tue, 12 May 2026 20:28:56 -0300 Subject: [PATCH 001/107] =?UTF-8?q?feat(download):=20bulk=20workspace=20ex?= =?UTF-8?q?port=20=E2=80=94=20session=20+=20project=20zip=20endpoints?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `GET /api/sessions/{sid}/download` and `GET /api/projects/{pid}/download` that stream a self-contained zip of the agent's workspace: every file under `/sessions/{sid}/` plus a synthetic `trainable_local.py` shim, a filtered `requirements.txt`, and a runbook README. The shim lets downloaded scripts that import `from trainable import log, log_image, ...` run on a vanilla Python install — calls land in `./trainable_out/` instead of the Modal volume. The zip is streamed via `StreamingResponse` over an in-memory buffer drained per write, so multi-hundred-MB sessions don't materialize on disk or in RAM. A 2 GB uncompressed cap with a trailing `__truncated.txt` marker keeps a runaway walk bounded. Frontend ships two entry points: - Download icon in the WorkspaceSidebar header → session zip - Download icon on every Project row in the sidebar → project zip Closes #79. Co-Authored-By: Claude Opus 4.7 (1M context) --- backend/main.py | 2 + backend/routers/download.py | 78 ++++++ backend/services/trainable_sdk.py | 364 +++++++++++++++++++++++++ backend/services/workspace_export.py | 283 +++++++++++++++++++ backend/tests/test_workspace_export.py | 270 ++++++++++++++++++ frontend/src/app/page.tsx | 16 ++ frontend/src/components/Sidebar.tsx | 12 + 7 files changed, 1025 insertions(+) create mode 100644 backend/routers/download.py create mode 100644 backend/services/trainable_sdk.py create mode 100644 backend/services/workspace_export.py create mode 100644 backend/tests/test_workspace_export.py diff --git a/backend/main.py b/backend/main.py index 7cf6d27..a5d605f 100644 --- a/backend/main.py +++ b/backend/main.py @@ -13,6 +13,7 @@ from routers import ( compare, data_explorer, + download, experiments, files, lineage, @@ -94,6 +95,7 @@ async def lifespan(app: FastAPI): app.include_router(compare.router, prefix="/api") app.include_router(snapshots.router, prefix="/api") app.include_router(lineage.router, prefix="/api") +app.include_router(download.router, prefix="/api") @app.get("/api/health") diff --git a/backend/routers/download.py b/backend/routers/download.py new file mode 100644 index 0000000..1339b87 --- /dev/null +++ b/backend/routers/download.py @@ -0,0 +1,78 @@ +"""Bulk workspace export endpoints. + +Streams the contents of `/sessions/{sid}` (or every session in a project) +as a single zip the user can `cd` into and run. The zip ships a local +`trainable` shim, a filtered `requirements.txt`, and a runbook README — +so downloaded scripts that `from trainable import log, ...` don't +`NameError` on a vanilla Python install. + +These endpoints inherit the auth posture of `routers/files.py`: they +read the same `/sessions/...` tree under the same caller and add no new +storage. See issue #79 for the runnability contract and rollout plan. +""" + +from __future__ import annotations + +import logging + +from fastapi import APIRouter, Depends, HTTPException +from fastapi.responses import StreamingResponse +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from db import get_db +from models import Project +from models import Session as SessionModel +from services.workspace_export import stream_project_zip, stream_session_zip + +logger = logging.getLogger(__name__) +router = APIRouter() + + +def _attachment_headers(filename: str) -> dict[str, str]: + # `Content-Disposition: attachment` is the part the browser keys on to + # trigger a file save rather than rendering the response inline. + return {"Content-Disposition": f'attachment; filename="{filename}"'} + + +@router.get("/sessions/{session_id}/download") +async def download_session(session_id: str, db: AsyncSession = Depends(get_db)): + result = await db.execute(select(SessionModel).where(SessionModel.id == session_id)) + session = result.scalar_one_or_none() + if not session: + raise HTTPException(status_code=404, detail="Session not found") + + filename = f"session-{session_id[:8]}.zip" + return StreamingResponse( + stream_session_zip(session_id), + media_type="application/zip", + headers=_attachment_headers(filename), + ) + + +@router.get("/projects/{project_id}/download") +async def download_project(project_id: str, db: AsyncSession = Depends(get_db)): + result = await db.execute(select(Project).where(Project.id == project_id)) + project = result.scalar_one_or_none() + if not project: + raise HTTPException(status_code=404, detail="Project not found") + + sessions_result = await db.execute( + select(SessionModel.id, SessionModel.name) + .where(SessionModel.project_id == project_id) + .order_by(SessionModel.created_at) + ) + rows = sessions_result.all() + if not rows: + raise HTTPException( + status_code=404, + detail="Project has no sessions to export", + ) + sessions = [(row[0], row[1]) for row in rows] + + filename = f"project-{project_id[:8]}.zip" + return StreamingResponse( + stream_project_zip(project_id, sessions), + media_type="application/zip", + headers=_attachment_headers(filename), + ) diff --git a/backend/services/trainable_sdk.py b/backend/services/trainable_sdk.py new file mode 100644 index 0000000..b59f81e --- /dev/null +++ b/backend/services/trainable_sdk.py @@ -0,0 +1,364 @@ +"""Synthetic files shipped inside a workspace-export zip. + +The Modal sandbox preamble in :mod:`services.sandbox` injects a `trainable` +module whose helpers write into the volume at well-known paths. When the +user downloads their session's workspace, the agent's scripts still +reference `from trainable import log, log_image, ...` — so the zip ships +a local shim that resolves those imports against the user's filesystem +instead of the volume. + +The shim's *surface* (method names + arity) must match the sandbox +preamble. When you add a helper there, mirror it here, otherwise the +downloaded scripts will `AttributeError` on unfamiliar Python. +""" + +from __future__ import annotations + +from datetime import datetime, timezone + + +# Local shim — written verbatim into the zip as `trainable_local.py`. +# Importing this file (or having it on sys.path under the name `trainable`) +# gives the downloaded scripts the same `trainable.log(...)` API they +# used in the sandbox, but rerouted to `./trainable_out/`. +LOCAL_SHIM = '''\ +"""Local trainable shim (auto-generated by the workspace exporter). + +Drop this file anywhere on sys.path and `import trainable` will resolve +to this shim. Helpers write to ./trainable_out/ instead of the Modal +volume the cloud agent used. Override the output root with the +TRAINABLE_LOCAL_OUT environment variable. + +This is best-effort parity with the cloud SDK — `log`/`log_image`/etc. +all work, but advanced features (live SSE broadcast, server-side chart +configuration) are no-ops. The point is to make downloaded scripts RUN +locally without `NameError`, not to mirror the studio experience. +""" + +from __future__ import annotations + +import json +import os +import pathlib +import re +import sys +import time +import types + + +_OUT = pathlib.Path(os.environ.get("TRAINABLE_LOCAL_OUT", "./trainable_out")).resolve() +_OUT.mkdir(parents=True, exist_ok=True) +_FIG_BASE = _OUT / "figures" +_HTML_BASE = _OUT / "html" +_METRICS_FILE = _OUT / "metrics.jsonl" +_LOG_FILE = _OUT / "log_events.jsonl" +_TABLE_ROW_LIMIT = 1000 + + +def _safe_key(key) -> str: + return re.sub(r"[^A-Za-z0-9_./-]", "_", str(key)).strip("/") or "log" + + +def _append_jsonl(path: pathlib.Path, payload: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "a", encoding="utf-8") as f: + f.write(json.dumps(payload) + "\\n") + + +def _save_image(img, dest_path: pathlib.Path) -> None: + dest_path.parent.mkdir(parents=True, exist_ok=True) + if isinstance(img, (str, bytes, os.PathLike)): + src = os.fspath(img) + if str(src) == str(dest_path): + return + with open(src, "rb") as r, open(dest_path, "wb") as w: + w.write(r.read()) + return + try: + from PIL import Image as _PILImage # type: ignore + + if isinstance(img, _PILImage.Image): + img.convert("RGB").save(dest_path, format="PNG") + return + except Exception: + pass + try: + import numpy as _np # type: ignore + + if isinstance(img, _np.ndarray): + from PIL import Image as _PILImage # type: ignore + + arr = img + if arr.dtype != _np.uint8: + a = arr.astype(_np.float32) + if a.max() <= 1.0 + 1e-6: + a = a * 255.0 + arr = a.clip(0, 255).astype(_np.uint8) + if arr.ndim == 2: + _PILImage.fromarray(arr, mode="L").save(dest_path, format="PNG") + elif arr.ndim == 3 and arr.shape[2] == 4: + _PILImage.fromarray(arr, mode="RGBA").save(dest_path, format="PNG") + else: + _PILImage.fromarray(arr).convert("RGB").save(dest_path, format="PNG") + return + except Exception: + pass + raise TypeError("log_image: unsupported image type %r" % (type(img),)) + + +def log(step, metrics, run=None): + payload = { + "ts": time.time(), + "step": int(step), + "metrics": {k: float(v) for k, v in dict(metrics).items()}, + } + if run: + payload["run"] = str(run) + _append_jsonl(_METRICS_FILE, payload) + print(f"[trainable] step={payload['step']} {payload['metrics']}") + + +def configure_dashboard(charts): + (_OUT / "dashboard.json").write_text(json.dumps({"charts": charts}, indent=2)) + + +def log_image(step, key, image, caption=None, run=None): + safe = _safe_key(key) + dest = _FIG_BASE / safe / f"{int(step)}.png" + _save_image(image, dest) + payload = { + "type": "image", + "step": int(step), + "key": safe, + "path": str(dest.relative_to(_OUT)), + } + if caption: + payload["caption"] = str(caption) + if run: + payload["run"] = str(run) + _append_jsonl(_LOG_FILE, payload) + + +def log_images(step, key, images, captions=None, run=None): + safe = _safe_key(key) + items = [] + for i, img in enumerate(images): + dest = _FIG_BASE / safe / f"{int(step)}_{i}.png" + _save_image(img, dest) + item = {"path": str(dest.relative_to(_OUT))} + if captions and i < len(captions) and captions[i] is not None: + item["caption"] = str(captions[i]) + items.append(item) + payload = {"type": "image_grid", "step": int(step), "key": safe, "items": items} + if run: + payload["run"] = str(run) + _append_jsonl(_LOG_FILE, payload) + + +def log_figure(step, key, fig, run=None): + safe = _safe_key(key) + dest = _FIG_BASE / safe / f"{int(step)}.png" + dest.parent.mkdir(parents=True, exist_ok=True) + try: + fig.savefig(str(dest), format="png", bbox_inches="tight", dpi=120) + except Exception as e: + raise TypeError("log_figure: object is not a matplotlib Figure (%s)" % e) + payload = { + "type": "image", + "step": int(step), + "key": safe, + "path": str(dest.relative_to(_OUT)), + } + if run: + payload["run"] = str(run) + _append_jsonl(_LOG_FILE, payload) + + +def log_table(step, key, columns, rows, run=None): + cols = [str(c) for c in columns] + rs = list(rows)[:_TABLE_ROW_LIMIT] + norm = [] + for r in rs: + row = list(r) if not isinstance(r, dict) else [r.get(c) for c in cols] + norm.append( + [ + None + if v is None + else ( + float(v) + if isinstance(v, bool) is False and isinstance(v, (int, float)) + else str(v) + ) + for v in row + ] + ) + payload = { + "type": "table", + "step": int(step), + "key": _safe_key(key), + "columns": cols, + "rows": norm, + "truncated": len(list(rows)) > _TABLE_ROW_LIMIT, + } + if run: + payload["run"] = str(run) + _append_jsonl(_LOG_FILE, payload) + + +def log_confusion_matrix(step, key, y_true, y_pred, labels=None, run=None): + try: + from sklearn.metrics import confusion_matrix as _cm # type: ignore + + labs = list(labels) if labels is not None else sorted(set(list(y_true) + list(y_pred))) + m = _cm(y_true, y_pred, labels=labs).tolist() + except Exception: + labs = list(labels) if labels is not None else sorted(set(list(y_true) + list(y_pred))) + idx = {lab: i for i, lab in enumerate(labs)} + m = [[0] * len(labs) for _ in labs] + for t, p in zip(y_true, y_pred): + if t in idx and p in idx: + m[idx[t]][idx[p]] += 1 + payload = { + "type": "confusion_matrix", + "step": int(step), + "key": _safe_key(key), + "labels": [str(lab) for lab in labs], + "matrix": m, + } + if run: + payload["run"] = str(run) + _append_jsonl(_LOG_FILE, payload) + + +def show_html(html, *, title=None, key=None): + name = _safe_key(key) if key else "artifact" + dest = _HTML_BASE / f"{name}.html" + dest.parent.mkdir(parents=True, exist_ok=True) + if isinstance(html, str): + dest.write_text(html, encoding="utf-8") + elif hasattr(html, "to_html"): + dest.write_text(html.to_html(), encoding="utf-8") + else: + dest.write_text(str(html), encoding="utf-8") + + +_m = types.ModuleType("trainable") +_m.log = log +_m.configure_dashboard = configure_dashboard +_m.log_image = log_image +_m.log_images = log_images +_m.log_figure = log_figure +_m.log_table = log_table +_m.log_confusion_matrix = log_confusion_matrix +_m.show_html = show_html +sys.modules["trainable"] = _m +''' + + +# Subset of `backend/requirements.txt` actually useful for running the +# agent's downloaded scripts. Excludes server-only deps (FastAPI, modal, +# anthropic, opentelemetry, etc.) so a `pip install -r requirements.txt` +# from the zip stays small and quick. +LOCAL_REQUIREMENTS = """\ +# Generated by the Trainable workspace exporter — minimal set to run +# agent-produced scripts locally. Mirrors what the cloud sandbox installs +# (see backend/services/sandbox.py:get_image). + +pandas>=2.0.0 +numpy>=1.24.0 +scikit-learn>=1.3.0 +matplotlib>=3.7.0 +seaborn>=0.12.0 +plotly>=5.18.0 +xgboost>=2.0.0 +lightgbm>=4.0.0 +pyarrow>=14.0.0 +duckdb>=1.0.0 +openpyxl>=3.1.0 +imbalanced-learn>=0.11.0 +optuna>=3.4.0 +category_encoders>=2.6.0 +pandera>=0.18.0 +shap>=0.43.0 +statsmodels>=0.14.0 +pypdf>=5.0.0 +pillow>=10.0.0 +jupyter>=1.0.0 +nbformat>=5.10.0 +""" + + +def render_readme( + *, scope: str, identifier: str, file_count: int, total_bytes: int +) -> str: + """Generate the README packaged at the zip root. + + `scope` is either "session" or "project" — affects the top-level + description but the run instructions are identical. + """ + ts = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") + short = identifier[:8] if len(identifier) > 8 else identifier + title = f"Trainable {scope} export — {short}" + size_mb = total_bytes / (1024 * 1024) + + if scope == "project": + layout_note = ( + "Each session is nested under `sessions//`. Two sessions\n" + "with identical labels are disambiguated by short-id suffix.\n" + ) + else: + layout_note = ( + "Files preserve the layout the agent wrote in the cloud, so\n" + "imports between `src/` modules keep working unchanged.\n" + ) + + return f"""# {title} + +Generated on {ts} from {scope} `{identifier}` ({file_count} files, {size_mb:.1f} MB). + +## Run it locally + +```bash +python -m venv .venv +source .venv/bin/activate # Windows: .venv\\Scripts\\activate +pip install -r requirements.txt +python -c "import trainable_local" # registers the local `trainable` shim +``` + +Once `trainable_local` is imported, any of the agent's scripts that do +`from trainable import log, log_image, ...` will work against your +filesystem. The simplest way to make the shim available to every script +is to drop a line at the top of your entrypoint: + +```python +import trainable_local # registers ./trainable_out/-backed `trainable` module +``` + +…or set `PYTHONSTARTUP=trainable_local.py` for a session-wide shim. + +## What's inside + +- `src/` — agent-written Python modules (the session was a + proper Python package: `from features import ...` + works the same here). +- `notebooks/` — Jupyter notebooks as-is. +- `figures/` — image artifacts the agent logged. +- `scripts/` — the audit trail of every sandbox call. +- `trainable_local.py` — local shim for the `trainable` SDK. +- `requirements.txt` — pinned subset of the cloud sandbox image. + +{layout_note} + +## What's NOT inside + +- **Raw datasets.** The agent's scripts reference inputs by their original + upload path (e.g. `/sessions/.../data/raw.csv`). Drop your local copy + in place and adjust the path, or set the `DATA_DIR` env var if a script + reads it. +- **The GPU/Modal environment.** Library versions are best-effort, not + byte-identical. Scripts that pin specific GPU kernels or Modal-only + paths may need light edits. +- **Round-trip telemetry.** `trainable.log(...)` writes to + `./trainable_out/metrics.jsonl` here — nothing is sent back to the + studio. +""" diff --git a/backend/services/workspace_export.py b/backend/services/workspace_export.py new file mode 100644 index 0000000..4311a67 --- /dev/null +++ b/backend/services/workspace_export.py @@ -0,0 +1,283 @@ +"""Streaming zip export of a session or project workspace. + +The exporter walks the Modal volume under `/sessions/{sid}` (or the union +of one project's sessions), reads each file via the async volume helpers, +and feeds the bytes into a `zipfile.ZipFile` backed by an in-memory +buffer that we drain after every write — so the response body trickles +out to the client without ever materializing the full archive in RAM or +on disk. + +After the workspace files, three synthetic entries (`README.md`, +`requirements.txt`, `trainable_local.py`) are appended at the zip root. +For a project export they appear once at the top level, and each +session's files are namespaced under `sessions/{slug}/`. + +Size safety +----------- +A per-export uncompressed byte cap (default 2 GB) protects the backend +from a runaway walk. When hit, the walk stops and a trailing +`__truncated.txt` entry lists the omitted paths and total skipped bytes +— the archive itself is still a valid zip the browser will finish +downloading. +""" + +from __future__ import annotations + +import io +import logging +import re +import zipfile +from typing import AsyncIterator, Sequence + +from services.trainable_sdk import LOCAL_REQUIREMENTS, LOCAL_SHIM, render_readme +from services.volume import ( + listdir_async, + read_volume_file_async, + reload_volume_async, + should_ignore_workspace_path, +) + +logger = logging.getLogger(__name__) + + +# Default uncompressed cap for any single export. ~2 GB matches Modal's +# typical session ceiling once `figures/` image grids accumulate; the +# value is overridable per-call so a future opt-in endpoint can raise +# the cap deliberately. +DEFAULT_MAX_UNCOMPRESSED_BYTES = 2 * 1024 * 1024 * 1024 # 2 GB + +# 1 MB read buffer is the sweet spot for Modal volume reads — small +# enough that an aggregator pulling tiny files doesn't stall waiting on +# a large one, big enough that the per-call overhead amortizes. +READ_CHUNK_BYTES = 1024 * 1024 + + +def _slug(name: str, fallback: str) -> str: + """Sanitize a label into a path-safe slug. + + Two sessions with identical labels are disambiguated by the caller + (it suffixes the short id). This function only strips characters + that would break the zip's archive-name layout. + """ + cleaned = re.sub(r"[^A-Za-z0-9._-]+", "-", (name or "").strip()).strip("-._") + return cleaned or fallback + + +async def _iter_files(root: str) -> AsyncIterator[tuple[str, str]]: + """Yield (volume_path, archive_relative_path) for every file under `root`. + + Skips entries that match `should_ignore_workspace_path` so the zip + doesn't ship `__pycache__/`, `.DS_Store`, etc. + """ + root_clean = root.rstrip("/") + try: + entries = await listdir_async(root_clean, recursive=True) + except FileNotFoundError: + return + for entry in entries: + # Modal's listdir returns directories AND files in recursive + # mode; only zip files. The exact enum varies by SDK version, so + # check the `.name` attribute instead of comparing enum objects. + try: + etype = entry.type.name + except AttributeError: + etype = str(entry.type) + if etype != "FILE": + continue + if should_ignore_workspace_path(entry.path): + continue + rel = entry.path + if rel.startswith(root_clean + "/"): + rel = rel[len(root_clean) + 1 :] + elif rel == root_clean: + continue + rel = rel.lstrip("/") + if not rel: + continue + yield entry.path, rel + + +class _StreamingZipBuffer(io.RawIOBase): + """`io.BytesIO`-shaped sink that ZipFile writes into; drained per chunk.""" + + def __init__(self) -> None: + super().__init__() + self._buf = bytearray() + self._pos = 0 + + def writable(self) -> bool: # pragma: no cover — required override + return True + + def write(self, b) -> int: + data = bytes(b) + self._buf.extend(data) + self._pos += len(data) + return len(data) + + def tell(self) -> int: + return self._pos + + def flush(self) -> None: # pragma: no cover — interface stub + return None + + def drain(self) -> bytes: + out = bytes(self._buf) + self._buf.clear() + return out + + +async def _stream_workspace_zip( + *, + scope: str, + identifier: str, + sources: Sequence[tuple[str, str]], + max_bytes: int = DEFAULT_MAX_UNCOMPRESSED_BYTES, +) -> AsyncIterator[bytes]: + """Yield zip-body chunks for `sources`. + + `sources` is a list of `(volume_root, archive_prefix)` — session + exports pass one tuple, project exports pass one per session. Each + prefix is joined onto every file's archive name, so a project export + interleaves `sessions/foo/src/a.py`, `sessions/bar/src/a.py`, etc., + without collisions. + + The terminal `__truncated.txt` entry is only added if the walk + actually hit the cap. + """ + # Refresh once at the top so the export sees writes from a sandbox + # that finished moments ago. A per-source reload would multiply the + # latency without much benefit — sessions in the same project rarely + # diverge by more than a few seconds. + try: + await reload_volume_async() + except Exception as exc: + logger.debug("workspace_export: volume reload skipped: %s", exc) + + buf = _StreamingZipBuffer() + zf = zipfile.ZipFile( + buf, mode="w", compression=zipfile.ZIP_DEFLATED, allowZip64=True + ) + + written_bytes = 0 + file_count = 0 + truncated_paths: list[str] = [] + + try: + for volume_root, archive_prefix in sources: + prefix = archive_prefix.strip("/") + async for volume_path, rel in _iter_files(volume_root): + arcname = f"{prefix}/{rel}" if prefix else rel + try: + payload = await read_volume_file_async(volume_path) + except Exception as exc: + logger.warning( + "workspace_export: skipping unreadable %s: %s", volume_path, exc + ) + continue + + if written_bytes + len(payload) > max_bytes: + truncated_paths.append(arcname) + # Keep walking so the truncated.txt is complete. + continue + + zf.writestr(arcname, payload) + written_bytes += len(payload) + file_count += 1 + chunk = buf.drain() + if chunk: + yield chunk + + # Synthetic entries at the archive root — one set per export, not + # per session, so a project zip doesn't ship N redundant copies. + readme = render_readme( + scope=scope, + identifier=identifier, + file_count=file_count, + total_bytes=written_bytes, + ) + zf.writestr("README.md", readme) + zf.writestr("requirements.txt", LOCAL_REQUIREMENTS) + zf.writestr("trainable_local.py", LOCAL_SHIM) + + if truncated_paths: + truncated_blob = ( + f"# Workspace export hit the {max_bytes:,}-byte cap.\n" + f"# The following {len(truncated_paths)} file(s) were omitted:\n\n" + + "\n".join(truncated_paths) + + "\n" + ) + zf.writestr("__truncated.txt", truncated_blob) + logger.warning( + "workspace_export %s/%s truncated: %d files omitted", + scope, + identifier, + len(truncated_paths), + ) + + chunk = buf.drain() + if chunk: + yield chunk + + finally: + zf.close() + chunk = buf.drain() + if chunk: + yield chunk + + logger.info( + "workspace_export %s/%s done: %d files, %d bytes (%d truncated)", + scope, + identifier, + file_count, + written_bytes, + len(truncated_paths), + ) + + +async def stream_session_zip( + session_id: str, + *, + max_bytes: int = DEFAULT_MAX_UNCOMPRESSED_BYTES, +) -> AsyncIterator[bytes]: + """Stream the zip body for one session.""" + sources = [(f"/sessions/{session_id}", "")] + async for chunk in _stream_workspace_zip( + scope="session", + identifier=session_id, + sources=sources, + max_bytes=max_bytes, + ): + yield chunk + + +async def stream_project_zip( + project_id: str, + sessions: Sequence[tuple[str, str | None]], + *, + max_bytes: int = DEFAULT_MAX_UNCOMPRESSED_BYTES, +) -> AsyncIterator[bytes]: + """Stream the zip body for a project. + + `sessions` is `[(session_id, optional_label), ...]`. Sessions with + duplicate labels are disambiguated by short-id suffix. + """ + if not sessions: + return + + # Resolve slugs up-front so per-file archive names are stable. + used: set[str] = set() + sources: list[tuple[str, str]] = [] + for session_id, label in sessions: + slug = _slug(label or session_id, session_id[:8]) + if slug in used: + slug = f"{slug}-{session_id[:8]}" + used.add(slug) + sources.append((f"/sessions/{session_id}", f"sessions/{slug}")) + + async for chunk in _stream_workspace_zip( + scope="project", + identifier=project_id, + sources=sources, + max_bytes=max_bytes, + ): + yield chunk diff --git a/backend/tests/test_workspace_export.py b/backend/tests/test_workspace_export.py new file mode 100644 index 0000000..1ee7f17 --- /dev/null +++ b/backend/tests/test_workspace_export.py @@ -0,0 +1,270 @@ +"""Tests for the streaming-zip workspace exporter and download endpoints.""" + +from __future__ import annotations + +import io +import sys +import zipfile +from contextlib import ExitStack +from pathlib import Path + +import pytest + +from tests.conftest import MockVolume, mock_volume_patches + + +def _files_for_session(session_id: str) -> dict[str, bytes]: + base = f"/sessions/{session_id}" + return { + f"{base}/src/__init__.py": b"", + f"{base}/src/data.py": b"import pandas as pd\nprint('hi')\n", + f"{base}/src/features.py": b"def build_x(df):\n return df\n", + f"{base}/notebooks/01_eda.ipynb": b'{"cells": []}', + f"{base}/figures/loss/10.png": b"fake-png-bytes", + f"{base}/figures/loss/20.png": b"more-fake-png-bytes", + f"{base}/scripts/step_07_train.py": b"# audit script\n", + # Noise that the exporter must skip. + f"{base}/__pycache__/data.cpython-311.pyc": b"junk", + f"{base}/.DS_Store": b"\x00\x00\x00", + } + + +async def _drain(agen) -> bytes: + chunks = [] + async for c in agen: + chunks.append(c) + return b"".join(chunks) + + +@pytest.mark.asyncio +async def test_session_zip_contains_workspace_and_synthetic_files(): + vol = MockVolume(_files_for_session("sess-aaaa")) + + with ExitStack() as stack: + for p in mock_volume_patches(vol, "services.workspace_export"): + stack.enter_context(p) + + from services.workspace_export import stream_session_zip + + body = await _drain(stream_session_zip("sess-aaaa")) + + zf = zipfile.ZipFile(io.BytesIO(body)) + names = set(zf.namelist()) + + # Workspace files preserved with relative paths (no /sessions// prefix). + assert "src/__init__.py" in names + assert "src/data.py" in names + assert "src/features.py" in names + assert "notebooks/01_eda.ipynb" in names + assert "figures/loss/10.png" in names + assert "figures/loss/20.png" in names + assert "scripts/step_07_train.py" in names + + # Synthetic entries at the zip root. + assert "README.md" in names + assert "requirements.txt" in names + assert "trainable_local.py" in names + + # No noise. + assert not any("__pycache__" in n for n in names) + assert not any(n.endswith(".DS_Store") for n in names) + + # Content sanity. + assert zf.read("src/data.py") == b"import pandas as pd\nprint('hi')\n" + assert b"trainable" in zf.read("trainable_local.py") + assert b"pip install -r requirements.txt" in zf.read("README.md") + + +@pytest.mark.asyncio +async def test_session_zip_is_well_formed_when_workspace_is_empty(): + vol = MockVolume({}) + + with ExitStack() as stack: + for p in mock_volume_patches(vol, "services.workspace_export"): + stack.enter_context(p) + + from services.workspace_export import stream_session_zip + + body = await _drain(stream_session_zip("empty-sess")) + + zf = zipfile.ZipFile(io.BytesIO(body)) + names = set(zf.namelist()) + # Empty session still ships the three synthetic files (issue acceptance). + assert names == {"README.md", "requirements.txt", "trainable_local.py"} + + +@pytest.mark.asyncio +async def test_project_zip_namespaces_each_session(): + files = {} + files.update(_files_for_session("sess-alpha")) + files.update(_files_for_session("sess-bravo")) + vol = MockVolume(files) + + with ExitStack() as stack: + for p in mock_volume_patches(vol, "services.workspace_export"): + stack.enter_context(p) + + from services.workspace_export import stream_project_zip + + body = await _drain( + stream_project_zip( + "proj-x", + [("sess-alpha", "EDA pass"), ("sess-bravo", "Train pass")], + ) + ) + + zf = zipfile.ZipFile(io.BytesIO(body)) + names = set(zf.namelist()) + assert "sessions/EDA-pass/src/data.py" in names + assert "sessions/Train-pass/src/data.py" in names + # Synthetic entries appear exactly once at the project root. + assert "README.md" in names + assert "requirements.txt" in names + assert "trainable_local.py" in names + assert sum(1 for n in names if n == "README.md") == 1 + + +@pytest.mark.asyncio +async def test_project_zip_disambiguates_duplicate_labels(): + files = {} + files.update(_files_for_session("sess-aaaa1111")) + files.update(_files_for_session("sess-bbbb2222")) + vol = MockVolume(files) + + with ExitStack() as stack: + for p in mock_volume_patches(vol, "services.workspace_export"): + stack.enter_context(p) + + from services.workspace_export import stream_project_zip + + body = await _drain( + stream_project_zip( + "proj-dup", + [ + ("sess-aaaa1111", "draft"), + ("sess-bbbb2222", "draft"), + ], + ) + ) + + zf = zipfile.ZipFile(io.BytesIO(body)) + names = set(zf.namelist()) + assert any(n.startswith("sessions/draft/") for n in names) + # Second "draft" gets a short-id suffix (first 8 chars of session id) + # so the two sessions don't collide in the archive. + assert any(n.startswith("sessions/draft-sess-bbb/") for n in names) + + +@pytest.mark.asyncio +async def test_export_caps_uncompressed_bytes_and_emits_truncated_marker(): + big_payload = b"x" * (50 * 1024) + files = {f"/sessions/big/data/file_{i:02d}.bin": big_payload for i in range(10)} + vol = MockVolume(files) + + with ExitStack() as stack: + for p in mock_volume_patches(vol, "services.workspace_export"): + stack.enter_context(p) + + from services.workspace_export import stream_session_zip + + # Cap = 150 KB → only ~3 files fit before truncation kicks in. + body = await _drain(stream_session_zip("big", max_bytes=150 * 1024)) + + zf = zipfile.ZipFile(io.BytesIO(body)) + names = set(zf.namelist()) + assert "__truncated.txt" in names + truncated_blob = zf.read("__truncated.txt").decode("utf-8") + assert "153,600-byte cap" in truncated_blob + # At least one file fit, at least one was skipped. + data_files = [n for n in names if n.startswith("data/")] + assert 0 < len(data_files) < 10 + + +def test_local_shim_imports_and_logs(tmp_path: Path): + """The shipped trainable_local.py must work on a vanilla Python install.""" + from services.trainable_sdk import LOCAL_SHIM + + shim_dir = tmp_path / "pkg" + shim_dir.mkdir() + (shim_dir / "trainable_local.py").write_text(LOCAL_SHIM) + + # Point the shim's output dir at the test's tmp path. + out_dir = tmp_path / "out" + monkey_env = {"TRAINABLE_LOCAL_OUT": str(out_dir)} + + import os + import subprocess + import textwrap + + runner = tmp_path / "runner.py" + runner.write_text( + textwrap.dedent( + """ + import sys + sys.path.insert(0, %r) + import trainable_local # registers ./trainable_out-style module + import trainable + trainable.log(1, {"loss": 0.5}) + trainable.log(2, {"loss": 0.25, "acc": 0.9}) + """ + ) + % str(shim_dir) + ) + + env = os.environ.copy() + env.update(monkey_env) + result = subprocess.run( + [sys.executable, str(runner)], + capture_output=True, + env=env, + text=True, + ) + assert result.returncode == 0, result.stderr + metrics_file = out_dir / "metrics.jsonl" + assert metrics_file.exists(), result.stderr + lines = metrics_file.read_text().strip().splitlines() + assert len(lines) == 2 + assert '"loss": 0.5' in lines[0] + assert '"acc": 0.9' in lines[1] + + +@pytest.mark.asyncio +async def test_session_download_endpoint_returns_zip(client, default_project_id): + # Insert a session via the ORM directly — avoids depending on the + # exact shape of the session-create REST endpoint, which is tested + # elsewhere. + import uuid + + from db import async_session + from models import Session as SessionModel + + sid = str(uuid.uuid4()) + async with async_session() as db: + db.add(SessionModel(id=sid, project_id=default_project_id)) + await db.commit() + + vol = MockVolume(_files_for_session(sid)) + with ExitStack() as stack: + for p in mock_volume_patches(vol, "services.workspace_export"): + stack.enter_context(p) + resp = await client.get(f"/api/sessions/{sid}/download") + + assert resp.status_code == 200, resp.text + assert resp.headers["content-type"] == "application/zip" + assert resp.headers["content-disposition"].startswith("attachment; filename=") + zf = zipfile.ZipFile(io.BytesIO(resp.content)) + names = set(zf.namelist()) + assert "src/data.py" in names + assert "trainable_local.py" in names + + +@pytest.mark.asyncio +async def test_session_download_404_when_session_missing(client): + resp = await client.get("/api/sessions/does-not-exist/download") + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_project_download_404_when_no_sessions(client, default_project_id): + resp = await client.get(f"/api/projects/{default_project_id}/download") + assert resp.status_code == 404 diff --git a/frontend/src/app/page.tsx b/frontend/src/app/page.tsx index 6e90963..2c77f91 100644 --- a/frontend/src/app/page.tsx +++ b/frontend/src/app/page.tsx @@ -66,6 +66,7 @@ import { GitBranch, Globe, ExternalLink, + Download, } from 'lucide-react'; import Sidebar from '@/components/Sidebar'; import Notebook from '@/components/notebook/Notebook'; @@ -3136,6 +3137,21 @@ function WorkspaceSidebar({ > + + + + e.stopPropagation()} + className="p-0.5 rounded opacity-0 group-hover:opacity-100 hover:bg-white/[0.1] transition-all shrink-0" + title="Download project workspace (zip)" + > + +