diff --git a/backend/openmlr/app.py b/backend/openmlr/app.py index a6b80ea..9a0c890 100644 --- a/backend/openmlr/app.py +++ b/backend/openmlr/app.py @@ -109,13 +109,20 @@ async def lifespan(app: FastAPI): from .auth.router import router as auth_router from .routes.agent import router as agent_router from .routes.compute import router as compute_router +from .routes.datasets import router as datasets_router from .routes.eval import router as eval_router +from .routes.experiments import router as experiments_router +from .routes.figures import router as figures_router from .routes.health import router as health_router from .routes.keys import router as keys_router from .routes.mcp import router as mcp_router +from .routes.models import router as models_router from .routes.projects import router as projects_router +from .routes.reproducibility import router as reproducibility_router +from .routes.research import router as research_router from .routes.review import router as review_router from .routes.settings import router as settings_router +from .routes.sweeps import router as sweeps_router from .routes.terminal import router as terminal_router app.include_router(auth_router) @@ -126,8 +133,15 @@ async def lifespan(app: FastAPI): app.include_router(compute_router) app.include_router(mcp_router) app.include_router(projects_router) +app.include_router(research_router) app.include_router(review_router) app.include_router(eval_router) +app.include_router(experiments_router) +app.include_router(datasets_router) +app.include_router(sweeps_router) +app.include_router(models_router) +app.include_router(figures_router) +app.include_router(reproducibility_router) app.include_router(terminal_router) diff --git a/backend/openmlr/routes/datasets.py b/backend/openmlr/routes/datasets.py new file mode 100644 index 0000000..5c01456 --- /dev/null +++ b/backend/openmlr/routes/datasets.py @@ -0,0 +1,201 @@ +"""Dataset management, profiling, validation, and split REST API routes. + +Provides endpoints for inspecting data files, analyzing column statistics, validating schemas, +and generating partitioned splits for ML training workflows. +""" + +from __future__ import annotations + +import logging +import os +import re +from pathlib import Path +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException, status +from pydantic import BaseModel, Field + +from ..db.models import User +from ..dependencies import get_current_user_optional +from ..services.dataset_profiler import DatasetProfiler + +router = APIRouter(prefix="/api/datasets", tags=["datasets"]) +logger = logging.getLogger(__name__) + +DATASET_PATH_DESC = "Path to the dataset file" +SAFE_PATH_PATTERN = re.compile(r"^[a-zA-Z0-9_\-./ ]+$") + + +def _safe_dataset_path(path_str: str) -> Path: + """Safely validate and resolve a dataset file path, mitigating path injection risks.""" + clean = os.path.normpath(str(path_str).strip()) + if not clean or not SAFE_PATH_PATTERN.match(clean) or ".." in clean or clean.startswith(("/etc", "/proc", "/sys")): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Invalid or disallowed dataset path parameter", + ) + resolved = Path(clean).resolve() + if not resolved.exists() or not resolved.is_file(): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Dataset file not found at '{path_str}'", + ) + return resolved + + +def _safe_output_dir(dir_str: str) -> Path: + """Safely validate and resolve an output directory path.""" + clean = os.path.normpath(str(dir_str).strip()) + if not clean or not SAFE_PATH_PATTERN.match(clean) or ".." in clean or clean.startswith(("/etc", "/proc", "/sys")): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Invalid or disallowed output directory path parameter", + ) + return Path(clean).resolve() + + +class ProfileDatasetRequest(BaseModel): + """Payload for profiling a dataset file.""" + + path: str = Field(..., min_length=1, description=DATASET_PATH_DESC) + sample_size: int = Field(default=5000, ge=10, le=100000, description="Max rows to sample") + + +class InspectSamplesRequest(BaseModel): + """Payload for sampling rows from a dataset.""" + + path: str = Field(..., min_length=1, description=DATASET_PATH_DESC) + n: int = Field(default=5, ge=1, le=100, description="Number of rows to return") + offset: int = Field(default=0, ge=0, description="Row offset") + strategy: str = Field(default="head", description="Sampling strategy: head, random, stratified") + label_column: str | None = Field(default=None, description="Column for stratification") + + +class ValidateDatasetRequest(BaseModel): + """Payload for validating a dataset.""" + + path: str = Field(..., min_length=1, description=DATASET_PATH_DESC) + expected_columns: list[str] | None = Field(default=None, description="Required column names") + max_null_pct: float = Field(default=20.0, ge=0.0, le=100.0, description="Max allowed null percentage") + max_token_length: int | None = Field(default=None, ge=1, description="Max allowed text tokens") + + +class SplitDatasetRequest(BaseModel): + """Payload for splitting a dataset into train/val/test partitions.""" + + path: str = Field(..., min_length=1, description="Path to the source dataset file") + output_dir: str = Field(..., min_length=1, description="Target directory for output partitions") + train_ratio: float = Field(default=0.8, gt=0.0, lt=1.0, description="Train ratio") + val_ratio: float = Field(default=0.1, ge=0.0, lt=1.0, description="Validation ratio") + test_ratio: float = Field(default=0.1, ge=0.0, lt=1.0, description="Test ratio") + stratify_column: str | None = Field(default=None, description="Column to stratify on") + seed: int = Field(default=42, description="Random seed") + + +@router.post("/profile") +async def profile_dataset( + req: ProfileDatasetRequest, + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Compute comprehensive statistical profile and diagnostics for a dataset file.""" + path = _safe_dataset_path(req.path) + + try: + profile = DatasetProfiler.profile(path, sample_size=req.sample_size) + return { + "success": True, + "profile": profile.to_dict(), + } + except Exception as e: + logger.exception("Failed profiling dataset '%s': %s", req.path, e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Error profiling dataset: {e}", + ) + + +@router.post("/inspect") +async def inspect_samples( + req: InspectSamplesRequest, + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Retrieve sample records from a dataset.""" + path = _safe_dataset_path(req.path) + + try: + samples = DatasetProfiler.sample_records( + path, + n=req.n, + offset=req.offset, + strategy=req.strategy, + label_column=req.label_column, + ) + return { + "success": True, + "total_sampled": len(samples), + "samples": samples, + } + except Exception as e: + logger.exception("Failed inspecting dataset '%s': %s", req.path, e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Error inspecting dataset: {e}", + ) + + +@router.post("/validate") +async def validate_dataset( + req: ValidateDatasetRequest, + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Validate dataset structure, required schema, and constraints.""" + path = _safe_dataset_path(req.path) + + try: + result = DatasetProfiler.validate_dataset( + path, + expected_columns=req.expected_columns, + max_null_pct=req.max_null_pct, + max_token_length=req.max_token_length, + ) + return { + "success": True, + "validation": result, + } + except Exception as e: + logger.exception("Failed validating dataset '%s': %s", req.path, e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Error validating dataset: {e}", + ) + + +@router.post("/split") +async def split_dataset( + req: SplitDatasetRequest, + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Partition a dataset file into train/val/test splits.""" + path = _safe_dataset_path(req.path) + out_dir = _safe_output_dir(req.output_dir) + + try: + manifest = DatasetProfiler.split_dataset( + path, + output_dir=out_dir, + train_ratio=req.train_ratio, + val_ratio=req.val_ratio, + test_ratio=req.test_ratio, + stratify_column=req.stratify_column, + seed=req.seed, + ) + return { + "success": True, + "manifest": manifest, + } + except Exception as e: + logger.exception("Failed splitting dataset '%s': %s", req.path, e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Error splitting dataset: {e}", + ) diff --git a/backend/openmlr/routes/experiments.py b/backend/openmlr/routes/experiments.py new file mode 100644 index 0000000..7597a3f --- /dev/null +++ b/backend/openmlr/routes/experiments.py @@ -0,0 +1,325 @@ +"""Machine Learning Experiments and Run Tracking API routes. + +Provides endpoints for creating, monitoring, updating, and comparing ML experiment runs, +metric trajectories (train/val loss, throughput, GPU stats), checkpoints, and terminal logs. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException, Query, Request, status +from pydantic import BaseModel, Field + +from ..db.models import User +from ..dependencies import get_current_user_optional +from ..services.experiment_tracker import ExperimentTracker +from .projects import WORKSPACES_ROOT + +router = APIRouter(tags=["experiments"]) +logger = logging.getLogger(__name__) + +# Fallback in-memory tracker +_global_tracker = ExperimentTracker() + + +def _get_tracker(request: Request, project_uuid: str | None = None) -> ExperimentTracker: + """Get the active ExperimentTracker, with workspace storage if available.""" + if hasattr(request.app.state, "experiment_tracker") and request.app.state.experiment_tracker: + return request.app.state.experiment_tracker + + if project_uuid: + storage_dir = WORKSPACES_ROOT / project_uuid / ".project-meta" / "experiments" + event_bus = getattr(request.app.state, "event_bus", None) + return ExperimentTracker(event_bus=event_bus, storage_dir=storage_dir) + + return _global_tracker + + +class CreateRunRequest(BaseModel): + """Payload for initiating a new experiment run.""" + + name: str = Field(..., min_length=1, max_length=255, description="Descriptive name of the run") + description: str = Field(default="", description="Hypothesis or goal of this experiment run") + hyperparameters: dict[str, Any] = Field(default_factory=dict, description="Model and training hyperparameters") + compute_target: str = Field(default="Local GPU", description="Target hardware (e.g. Local H100, Modal A100)") + tags: list[str] = Field(default_factory=list, description="Categorization tags") + total_steps: int = Field(default=100, ge=1, description="Expected total optimization steps") + total_epochs: int = Field(default=1, ge=1, description="Expected total epochs") + project_uuid: str | None = Field(default=None, description="Associated project UUID") + + +class UpdateRunStatusRequest(BaseModel): + """Payload for updating experiment run status.""" + + status: str = Field(..., description="New status: running, paused, completed, failed, idle") + reason: str | None = Field(default=None, description="Optional explanation or failure error message") + + +class LogMetricsRequest(BaseModel): + """Payload for logging metric points during training.""" + + step: int = Field(..., ge=0, description="Global training step") + epoch: int = Field(default=1, ge=1, description="Current epoch") + metrics: dict[str, float] = Field(..., description="Key-value metrics (e.g. train_loss, val_loss, lr)") + timestamp: float | None = Field(default=None, description="Optional epoch millisecond timestamp") + + +class LogEntriesRequest(BaseModel): + """Payload for appending stdout/stderr log lines.""" + + lines: list[str] = Field(..., description="Array of log lines") + + +class RegisterCheckpointRequest(BaseModel): + """Payload for saving an experiment checkpoint.""" + + name: str = Field(..., description="Checkpoint name (e.g. step_500_best.pt)") + step: int = Field(..., ge=0, description="Step at which checkpoint was saved") + epoch: int = Field(default=1, ge=1, description="Epoch at which checkpoint was saved") + path: str = Field(default="", description="Relative or absolute file path to checkpoint") + file_size_mb: float = Field(default=0.0, ge=0, description="Checkpoint size in megabytes") + metrics: dict[str, float] = Field(default_factory=dict, description="Metric snapshot at checkpoint") + download_url: str = Field(default="", description="Direct download URL if uploaded to object store") + + +@router.get("/api/experiments/runs") +async def list_experiment_runs( + request: Request, + project_uuid: str | None = Query(None, description="Filter by project UUID"), + status: str | None = Query(None, description="Filter by status (running, completed, failed, paused, all)"), + search: str | None = Query(None, description="Search term across name, description, tags"), + limit: int = Query(50, ge=1, le=200, description="Page limit"), + offset: int = Query(0, ge=0, description="Page offset"), + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """List all tracked experiment runs with summary metrics and pagination.""" + tracker = _get_tracker(request, project_uuid) + runs, total = tracker.list_runs( + project_uuid=project_uuid, + status=status, + search=search, + limit=limit, + offset=offset, + ) + return { + "runs": [r.to_dict() for r in runs], + "total": total, + "limit": limit, + "offset": offset, + } + + +@router.post("/api/experiments/runs", status_code=status.HTTP_201_CREATED) +async def create_experiment_run( + req: CreateRunRequest, + request: Request, + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Start and register a new experiment run.""" + tracker = _get_tracker(request, req.project_uuid) + run = tracker.create_run( + name=req.name, + description=req.description, + hyperparameters=req.hyperparameters, + compute_target=req.compute_target, + tags=req.tags, + total_steps=req.total_steps, + total_epochs=req.total_epochs, + project_uuid=req.project_uuid, + ) + + event_bus = getattr(request.app.state, "event_bus", None) + if event_bus: + try: + await event_bus.broadcast({ + "type": "experiment_run_created", + "run_id": run.id, + "name": run.name, + "project_uuid": req.project_uuid, + }) + except Exception as ex: + logger.warning("Failed to broadcast experiment run created event: %s", ex) + + return {"status": "created", "run": run.to_dict()} + + +@router.get("/api/experiments/runs/{run_id}") +async def get_experiment_run( + run_id: str, + request: Request, + project_uuid: str | None = Query(None), + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Retrieve detailed metrics, parameters, checkpoints, and logs for a specific run.""" + tracker = _get_tracker(request, project_uuid) + run = tracker.get_run(run_id) + if not run: + raise HTTPException(status_code=404, detail=f"Experiment run '{run_id}' not found") + return {"run": run.to_dict()} + + +@router.post("/api/experiments/runs/{run_id}/metrics") +async def log_run_metrics( + run_id: str, + req: LogMetricsRequest, + request: Request, + project_uuid: str | None = Query(None), + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Log one or more training/evaluation metric points for a run.""" + tracker = _get_tracker(request, project_uuid) + try: + run = tracker.log_metrics( + run_id=run_id, + step=req.step, + epoch=req.epoch, + metrics=req.metrics, + timestamp=req.timestamp, + ) + except KeyError: + raise HTTPException(status_code=404, detail=f"Experiment run '{run_id}' not found") + + event_bus = getattr(request.app.state, "event_bus", None) + if event_bus: + try: + await event_bus.broadcast({ + "type": "experiment_metric_logged", + "run_id": run_id, + "step": req.step, + "epoch": req.epoch, + "metrics": req.metrics, + }) + except Exception as ex: + logger.warning("Failed to broadcast metric event: %s", ex) + + return { + "status": "success", + "run_id": run.id, + "current_step": run.current_step, + "best_val_loss": run.best_val_loss, + } + + +@router.post("/api/experiments/runs/{run_id}/status") +async def update_run_status( + run_id: str, + req: UpdateRunStatusRequest, + request: Request, + project_uuid: str | None = Query(None), + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Update execution status of an experiment run (running, paused, completed, failed).""" + tracker = _get_tracker(request, project_uuid) + try: + run = tracker.update_status(run_id=run_id, status=req.status, reason=req.reason) + except KeyError: + raise HTTPException(status_code=404, detail=f"Experiment run '{run_id}' not found") + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + + event_bus = getattr(request.app.state, "event_bus", None) + if event_bus: + try: + await event_bus.broadcast({ + "type": "experiment_status_changed", + "run_id": run_id, + "status": req.status, + "reason": req.reason, + }) + except Exception as ex: + logger.warning("Failed to broadcast status event: %s", ex) + + return {"status": "success", "run": run.to_dict()} + + +@router.post("/api/experiments/runs/{run_id}/logs") +async def append_run_logs( + run_id: str, + req: LogEntriesRequest, + request: Request, + project_uuid: str | None = Query(None), + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Append stdout/stderr lines to the experiment run log buffer.""" + tracker = _get_tracker(request, project_uuid) + try: + logs = tracker.append_logs(run_id=run_id, lines=req.lines) + except KeyError: + raise HTTPException(status_code=404, detail=f"Experiment run '{run_id}' not found") + return {"status": "success", "total_lines": len(logs)} + + +@router.get("/api/experiments/runs/{run_id}/logs") +async def get_run_logs( + run_id: str, + request: Request, + limit: int = Query(200, ge=1, le=2000), + project_uuid: str | None = Query(None), + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Retrieve the recent stdout/stderr log buffer for a run.""" + tracker = _get_tracker(request, project_uuid) + run = tracker.get_run(run_id) + if not run: + raise HTTPException(status_code=404, detail=f"Experiment run '{run_id}' not found") + return {"run_id": run_id, "logs": run.logs[-limit:]} + + +@router.post("/api/experiments/runs/{run_id}/checkpoints") +async def register_run_checkpoint( + run_id: str, + req: RegisterCheckpointRequest, + request: Request, + project_uuid: str | None = Query(None), + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Register a trained model checkpoint for an experiment run.""" + tracker = _get_tracker(request, project_uuid) + try: + cp = tracker.register_checkpoint( + run_id=run_id, + name=req.name, + step=req.step, + epoch=req.epoch, + path=req.path, + file_size_mb=req.file_size_mb, + metrics=req.metrics, + download_url=req.download_url, + ) + except KeyError: + raise HTTPException(status_code=404, detail=f"Experiment run '{run_id}' not found") + return {"status": "registered", "checkpoint": cp.to_dict()} + + +@router.get("/api/experiments/compare") +async def compare_experiment_runs( + request: Request, + run_ids: str = Query(..., description="Comma-separated run IDs to compare"), + project_uuid: str | None = Query(None), + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Compare metrics trajectories, hyperparameters, and best losses across runs.""" + ids = [rid.strip() for rid in run_ids.split(",") if rid.strip()] + if not ids: + raise HTTPException(status_code=400, detail="At least one run_id must be provided") + + tracker = _get_tracker(request, project_uuid) + comparison = tracker.compare_runs(ids) + return comparison + + +@router.delete("/api/experiments/runs/{run_id}") +async def delete_experiment_run( + run_id: str, + request: Request, + project_uuid: str | None = Query(None), + user: User | None = Depends(get_current_user_optional), +) -> dict[str, Any]: + """Delete an experiment run from storage.""" + tracker = _get_tracker(request, project_uuid) + deleted = tracker.delete_run(run_id) + if not deleted: + raise HTTPException(status_code=404, detail=f"Experiment run '{run_id}' not found") + return {"status": "deleted", "run_id": run_id} diff --git a/backend/openmlr/routes/figures.py b/backend/openmlr/routes/figures.py new file mode 100644 index 0000000..e5204cf --- /dev/null +++ b/backend/openmlr/routes/figures.py @@ -0,0 +1,92 @@ +"""REST API routes for Publication Figure Studio and LaTeX Diagram Generation.""" + +from __future__ import annotations + +import logging +from typing import Any + +from fastapi import APIRouter, HTTPException, Query, status + +from ..services.figure_generator import FigureGeneratorService +from ..services.figure_types import ( + GenerateFigureRequest, + MultiPanelLayoutRequest, +) + +router = APIRouter(prefix="/api/figures", tags=["figures"]) +logger = logging.getLogger("openmlr.routes.figures") + +PROJECT_ID_DESC = "Project ID" + + +@router.get("") +async def list_figures( + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """List all figure artifacts for a project.""" + figures = FigureGeneratorService.list_figures(project_id) + return { + "figures": [f.to_dict() for f in figures], + "total_count": len(figures), + } + + +@router.post("", status_code=status.HTTP_201_CREATED) +async def generate_figure( + request: GenerateFigureRequest, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Generate a new publication figure artifact.""" + artifact = FigureGeneratorService.generate_figure(project_id, request) + return { + "figure": artifact.to_dict(), + "message": f"Figure '{artifact.title}' generated successfully.", + } + + +@router.get("/{figure_id}") +async def get_figure( + figure_id: str, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Get details of a single figure artifact.""" + artifact = FigureGeneratorService.get_figure(project_id, figure_id) + if not artifact: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Figure '{figure_id}' not found.", + ) + return {"figure": artifact.to_dict()} + + +@router.delete("/{figure_id}") +async def delete_figure( + figure_id: str, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Delete a figure artifact.""" + success = FigureGeneratorService.delete_figure(project_id, figure_id) + if not success: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Figure '{figure_id}' not found.", + ) + return { + "success": True, + "message": f"Figure '{figure_id}' deleted successfully.", + } + + +@router.post("/multi-panel") +async def create_multipanel_layout( + request: MultiPanelLayoutRequest, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Combine multiple figures into a multi-panel LaTeX subfigure layout.""" + result = FigureGeneratorService.create_multipanel_layout(project_id, request) + if "error" in result: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=result["error"], + ) + return result diff --git a/backend/openmlr/routes/models.py b/backend/openmlr/routes/models.py new file mode 100644 index 0000000..af462b3 --- /dev/null +++ b/backend/openmlr/routes/models.py @@ -0,0 +1,204 @@ +"""REST API routes for Model Registry, Checkpoint Governance, and Model Cards.""" + +from __future__ import annotations + +import logging +from typing import Any + +from fastapi import APIRouter, HTTPException, Query, status + +from ..services.model_registry import ModelRegistryService +from ..services.model_types import ( + CompareModelsRequest, + GenerateModelCardRequest, + InspectCheckpointRequest, + PlanQuantizationRequest, + RegisterModelRequest, + UpdateModelRequest, +) + +router = APIRouter(prefix="/api/model-registry", tags=["model-registry"]) +logger = logging.getLogger("openmlr.routes.models") + + +PROJECT_ID_DESC = "Project ID" + + +@router.get("") +async def list_models( + project_id: str = Query("default", description=PROJECT_ID_DESC), + task_type: str | None = Query(None, description="Filter by task type"), + framework: str | None = Query(None, description="Filter by framework"), + status: str | None = Query(None, description="Filter by status"), + tag: str | None = Query(None, description="Filter by tag"), +) -> dict[str, Any]: + """List all registered model artifacts for a project.""" + models = ModelRegistryService.list_models( + project_id=project_id, + task_type=task_type, + framework=framework, + status=status, + tag=tag, + ) + return { + "models": [m.to_dict() for m in models], + "total_count": len(models), + } + + +@router.post("", status_code=status.HTTP_201_CREATED) +async def register_model( + request: RegisterModelRequest, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Register a new model artifact.""" + artifact = ModelRegistryService.register_model(project_id, request) + return { + "model": artifact.to_dict(), + "message": f"Model artifact '{artifact.name}' (v{artifact.version}) registered successfully.", + } + + +@router.get("/{model_id}") +async def get_model( + model_id: str, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Get full details of a registered model artifact.""" + artifact = ModelRegistryService.get_model(project_id, model_id) + if not artifact: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Model artifact '{model_id}' not found.", + ) + return {"model": artifact.to_dict()} + + +@router.put("/{model_id}") +async def update_model( + model_id: str, + request: UpdateModelRequest, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Update metadata and properties of a model artifact.""" + artifact = ModelRegistryService.update_model(project_id, model_id, request) + if not artifact: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Model artifact '{model_id}' not found.", + ) + return { + "model": artifact.to_dict(), + "message": f"Model artifact '{artifact.name}' updated successfully.", + } + + +@router.delete("/{model_id}") +async def delete_model( + model_id: str, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Delete a model artifact from the registry.""" + success = ModelRegistryService.delete_model(project_id, model_id) + if not success: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Model artifact '{model_id}' not found.", + ) + return { + "success": True, + "message": f"Model artifact '{model_id}' deleted successfully.", + } + + +@router.post("/{model_id}/card") +async def generate_model_card( + model_id: str, + request: GenerateModelCardRequest, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Generate a multi-format Model Card (Markdown, LaTeX, BibTeX, Carbon).""" + card = ModelRegistryService.generate_model_card(project_id, model_id, request) + if not card: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Model artifact '{model_id}' not found.", + ) + return { + "model_name": card.model_name, + "version": card.version, + "markdown": card.markdown, + "latex": card.latex, + "bibtex": card.bibtex, + "co2_emissions_kg": card.co2_emissions_kg, + "summary": card.summary, + } + + +@router.post("/{model_id}/quantization") +async def plan_quantization( + model_id: str, + request: PlanQuantizationRequest, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Calculate quantization trade-offs and memory savings for target precisions.""" + artifact = ModelRegistryService.get_model(project_id, model_id) + if not artifact: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Model artifact '{model_id}' not found.", + ) + estimates = ModelRegistryService.plan_quantization(artifact, request.target_precisions) + return { + "model_id": artifact.id, + "model_name": artifact.name, + "base_parameters": artifact.parameters_count, + "estimates": [ + { + "target_precision": e.target_precision, + "estimated_size_mb": e.estimated_size_mb, + "estimated_vram_mb": e.estimated_vram_mb, + "compression_ratio": e.compression_ratio, + "expected_latency_speedup": e.expected_latency_speedup, + "suggested_engine": e.suggested_engine, + "loss_tolerance_level": e.loss_tolerance_level, + } + for e in estimates + ], + } + + +@router.post("/inspect") +async def inspect_checkpoint(request: InspectCheckpointRequest) -> dict[str, Any]: + """Inspect checkpoint structure, layer breakdown, and memory requirements.""" + inspection = ModelRegistryService.inspect_checkpoint(request) + return { + "file_format": inspection.file_format, + "total_parameters": inspection.total_parameters, + "trainable_parameters": inspection.trainable_parameters, + "total_size_mb": inspection.total_size_mb, + "estimated_vram_fp32_mb": inspection.estimated_vram_fp32_mb, + "estimated_vram_fp16_mb": inspection.estimated_vram_fp16_mb, + "estimated_vram_int8_mb": inspection.estimated_vram_int8_mb, + "estimated_vram_int4_mb": inspection.estimated_vram_int4_mb, + "dtype_breakdown": inspection.dtype_breakdown, + "layers_count": inspection.layers_count, + "top_layers": inspection.top_layers, + "has_optimizer_state": inspection.has_optimizer_state, + "metadata": inspection.metadata, + } + + +@router.post("/compare") +async def compare_models( + request: CompareModelsRequest, + project_id: str = Query("default", description=PROJECT_ID_DESC), +) -> dict[str, Any]: + """Compare multiple model artifacts side-by-side.""" + comparison = ModelRegistryService.compare_models(project_id, request.model_ids) + if "error" in comparison: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=comparison["error"], + ) + return comparison diff --git a/backend/openmlr/routes/reproducibility.py b/backend/openmlr/routes/reproducibility.py new file mode 100644 index 0000000..3c2bb98 --- /dev/null +++ b/backend/openmlr/routes/reproducibility.py @@ -0,0 +1,109 @@ +"""REST API routes for Reproducibility Studio and Artifact Verification.""" + +from __future__ import annotations + +import logging +from typing import Any + +from fastapi import APIRouter, HTTPException, Query, status + +from ..services.reproducibility_auditor import ReproducibilityAuditorService +from ..services.reproducibility_types import ( + AuditCodebaseRequest, + FixDeterminismRequest, + GenerateAppendixRequest, + GenerateDockerfileRequest, + ReproducibilityAuditReport, +) + +logger = logging.getLogger("openmlr.routes.reproducibility") +router = APIRouter(prefix="/api/reproducibility", tags=["reproducibility"]) + + +@router.get("/reports", response_model=list[ReproducibilityAuditReport]) +async def list_reports( + project_id: str | None = Query(None, description="Filter reports by project ID"), +) -> Any: + """List all reproducibility audit reports for a project.""" + return ReproducibilityAuditorService.list_reports(project_id) + + +@router.get("/reports/{report_id}", response_model=ReproducibilityAuditReport) +async def get_report( + report_id: str, + project_id: str | None = Query(None, description="Project ID"), +) -> Any: + """Retrieve a specific reproducibility audit report.""" + report = ReproducibilityAuditorService.get_report(report_id, project_id) + if not report: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Report '{report_id}' not found.", + ) + return report + + +@router.delete("/reports/{report_id}") +async def delete_report( + report_id: str, + project_id: str | None = Query(None, description="Project ID"), +) -> dict[str, str]: + """Delete a reproducibility audit report.""" + success = ReproducibilityAuditorService.delete_report(report_id, project_id) + if not success: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Report '{report_id}' not found.", + ) + return {"status": "deleted", "report_id": report_id} + + +@router.post("/audit", response_model=ReproducibilityAuditReport) +async def audit_codebase( + request: AuditCodebaseRequest, + project_id: str | None = Query(None, description="Project ID"), +) -> Any: + """Run an automated reproducibility audit on codebase or provided snippets.""" + try: + return ReproducibilityAuditorService.audit_codebase(request, project_id) + except Exception as e: + logger.exception("Failed to run reproducibility audit: %s", e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Audit failed: {e}", + ) from e + + +@router.post("/dockerfile") +async def generate_dockerfile( + request: GenerateDockerfileRequest, +) -> dict[str, str]: + """Generate a reproducible Dockerfile.""" + dockerfile = ReproducibilityAuditorService.generate_dockerfile(request) + return {"dockerfile": dockerfile} + + +@router.post("/appendix") +async def generate_appendix( + request: GenerateAppendixRequest, + project_id: str | None = Query(None, description="Project ID"), +) -> dict[str, str]: + """Generate a LaTeX Reproducibility Statement appendix.""" + report = None + if request.report_id: + report = ReproducibilityAuditorService.get_report(request.report_id, project_id) + appendix = ReproducibilityAuditorService.generate_latex_appendix(request, report) + return {"latex_appendix": appendix} + + +@router.post("/fix-determinism") +async def fix_determinism( + request: FixDeterminismRequest, +) -> dict[str, str]: + """Generate boilerplate determinism code.""" + snippet = ReproducibilityAuditorService.generate_determinism_snippet( + framework=request.framework, + seed=request.seed, + strict_mode=request.strict_mode, + ) + return {"determinism_snippet": snippet} diff --git a/backend/openmlr/routes/research.py b/backend/openmlr/routes/research.py new file mode 100644 index 0000000..dadb3b1 --- /dev/null +++ b/backend/openmlr/routes/research.py @@ -0,0 +1,408 @@ +"""Autonomous ML Research Workflow API routes. + +Provides endpoints for tracking, driving, and interacting with the 5-phase +scientific research state machine (Reconnaissance -> Hypothesis -> Experimentation -> Analysis -> Paper Drafting). +""" + +from __future__ import annotations + +import logging +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException, Request +from pydantic import BaseModel, Field +from sqlalchemy.ext.asyncio import AsyncSession + +from ..agent.research_orchestrator import PHASE_GUIDELINES, ResearchOrchestrator +from ..agent.states import MilestoneStatus, ResearchPhase +from ..db import operations as ops +from ..db.engine import get_db +from ..db.models import User +from ..dependencies import get_current_user +from .projects import WORKSPACES_ROOT + +router = APIRouter(tags=["research"]) +logger = logging.getLogger(__name__) + +# Standard default milestones for new research projects +DEFAULT_PHASE_MILESTONES = [ + { + "phase": ResearchPhase.RECONNAISSANCE, + "title": "Literature Reconnaissance", + "description": "Search academic databases (arXiv, OpenAlex, Semantic Scholar) and catalog relevant papers.", + "criteria": ["Identify at least 5 foundational papers", "Map baseline benchmark metrics"], + }, + { + "phase": ResearchPhase.HYPOTHESIS, + "title": "Hypothesis & Experimental Design", + "description": "Formulate testable scientific claims, architectural modifications, and ablation matrices.", + "criteria": ["Define falsifiable hypothesis", "Specify baseline vs proposed comparison criteria"], + }, + { + "phase": ResearchPhase.EXPERIMENTATION, + "title": "Code Implementation & Model Training", + "description": "Implement model code and execute training/benchmarks on local or cloud compute.", + "criteria": ["Verify train/eval loss trajectories", "Save experiment checkpoints and logs"], + }, + { + "phase": ResearchPhase.ANALYSIS, + "title": "Empirical Analysis & Self-Correction", + "description": "Evaluate metrics against baselines, run ablations, and resolve training anomalies.", + "criteria": ["Generate comparison tables", "Quantify statistical significance"], + }, + { + "phase": ResearchPhase.PAPER_DRAFTING, + "title": "LaTeX Manuscript & Bibliography", + "description": "Author standard conference paper sections (Abstract, Intro, Method, Results) and compile PDF.", + "criteria": ["Ensure LaTeX compiles without error", "Validate BibTeX citation keys"], + }, +] + + +class ResearchStartRequest(BaseModel): + """Request payload to initiate a research workflow.""" + + goal: str = Field(..., min_length=5, description="Scientific objective or research question") + initial_phase: str = Field( + default="reconnaissance", + description="Starting phase: idle, reconnaissance, hypothesis, experimentation, analysis, paper_drafting", + ) + generate_default_milestones: bool = Field( + default=True, + description="Whether to prepopulate standard research milestones for each phase", + ) + + +class ResearchTransitionRequest(BaseModel): + """Request payload to advance or switch research phases.""" + + next_phase: str = Field(..., description="Target phase name") + reason: str = Field(..., min_length=3, description="Rationale for phase transition") + artifacts_produced: list[str] = Field(default_factory=list, description="Artifact keys or identifiers generated") + milestone_id: str | None = Field(default=None, description="Optional milestone ID triggering this transition") + + +class MilestoneCreateRequest(BaseModel): + """Request payload to add a milestone.""" + + title: str = Field(..., min_length=2, description="Short title of milestone") + description: str = Field(default="", description="Detailed milestone criteria or instructions") + phase: str | None = Field(default=None, description="Target phase; defaults to active phase") + criteria: list[str] = Field(default_factory=list, description="Verification acceptance criteria") + + +class MilestoneUpdateRequest(BaseModel): + """Request payload to update or complete a milestone.""" + + status: str | None = Field(default=None, description="pending, in_progress, completed, failed, skipped") + output_artifacts: list[str] = Field(default_factory=list, description="Artifact identifiers produced by milestone") + + +class ArtifactCreateRequest(BaseModel): + """Request payload to register a research artifact.""" + + type: str = Field(..., description="Artifact type: paper, hypothesis, experiment, metrics, manuscript_section, bibtex") + data: Any = Field(..., description="Artifact payload content or metadata dictionary") + section_name: str | None = Field(default=None, description="Section name if type is manuscript_section") + + +def _get_project_workspace(user_id: int, project_slug: str): + ws_dir = WORKSPACES_ROOT / str(user_id) / project_slug + ws_dir.mkdir(parents=True, exist_ok=True) + return ws_dir + + +def _load_orchestrator(user_id: int, project_slug: str) -> ResearchOrchestrator: + ws_dir = _get_project_workspace(user_id, project_slug) + orch = ResearchOrchestrator(workspace_path=ws_dir) + orch.load_state() + return orch + + +@router.get("/api/research/phases") +async def list_research_phases() -> dict[str, Any]: + """List all supported research phases and their metadata.""" + phases = [ + { + "id": p.value, + "name": p.value.replace("_", " ").title(), + "description": PHASE_GUIDELINES.get(p, ""), + } + for p in ResearchPhase + ] + return {"phases": phases} + + +@router.get("/api/research/guidelines") +async def get_all_phase_guidelines() -> dict[str, str]: + """Get prompt guidelines for all research phases.""" + return {p.value: guidelines for p, guidelines in PHASE_GUIDELINES.items()} + + +@router.get("/api/projects/{project_id}/research/state") +async def get_project_research_state( + project_id: int, + db: AsyncSession = Depends(get_db), + user: User = Depends(get_current_user), +) -> dict[str, Any]: + """Retrieve current research state, milestones, artifacts, and transition history for a project.""" + project = await ops.get_project_by_id(db, project_id) + if not project or project.user_id != user.id: + raise HTTPException(status_code=404, detail="Project not found") + + orch = _load_orchestrator(user.id, project.slug) + return { + "project_id": project.id, + "project_name": project.name, + "state": orch.state.to_dict(), + "guidelines": orch.get_phase_guidelines(), + "context_prompt": orch.format_research_context(), + } + + +@router.post("/api/projects/{project_id}/research/start") +async def start_project_research( + project_id: int, + req: ResearchStartRequest, + request: Request, + db: AsyncSession = Depends(get_db), + user: User = Depends(get_current_user), +) -> dict[str, Any]: + """Initialize or reset the structured scientific research state machine for a project.""" + project = await ops.get_project_by_id(db, project_id) + if not project or project.user_id != user.id: + raise HTTPException(status_code=404, detail="Project not found") + + try: + init_phase = ResearchPhase(req.initial_phase) + except ValueError: + raise HTTPException( + status_code=400, + detail=f"Invalid phase '{req.initial_phase}'. Valid phases: {[p.value for p in ResearchPhase]}", + ) + + orch = _load_orchestrator(user.id, project.slug) + transition = orch.start_research(goal=req.goal, initial_phase=init_phase) + + if req.generate_default_milestones and not orch.state.milestones: + for m in DEFAULT_PHASE_MILESTONES: + orch.add_milestone( + title=m["title"], + description=m["description"], + phase=m["phase"], + criteria=m["criteria"], + ) + + orch.save_state() + + # Broadcast event if event bus is available + if hasattr(request.app.state, "event_bus") and request.app.state.event_bus: + try: + request.app.state.event_bus.publish( + "research_started", + { + "project_id": project.id, + "goal": req.goal, + "current_phase": orch.current_phase.value, + "transition": transition.to_dict(), + }, + project_id=str(project.id), + ) + except Exception as ex: + logger.warning("Failed to publish research_started event: %s", ex) + + return { + "status": "started", + "project_id": project.id, + "state": orch.state.to_dict(), + "transition": transition.to_dict(), + "guidelines": orch.get_phase_guidelines(), + } + + +@router.post("/api/projects/{project_id}/research/transition") +async def transition_project_research_phase( + project_id: int, + req: ResearchTransitionRequest, + request: Request, + db: AsyncSession = Depends(get_db), + user: User = Depends(get_current_user), +) -> dict[str, Any]: + """Transition research phase to next stage with recorded rationale and artifacts.""" + project = await ops.get_project_by_id(db, project_id) + if not project or project.user_id != user.id: + raise HTTPException(status_code=404, detail="Project not found") + + try: + next_phase = ResearchPhase(req.next_phase) + except ValueError: + raise HTTPException( + status_code=400, + detail=f"Invalid phase '{req.next_phase}'. Valid phases: {[p.value for p in ResearchPhase]}", + ) + + orch = _load_orchestrator(user.id, project.slug) + transition = orch.transition_to( + next_phase=next_phase, + reason=req.reason, + artifacts_produced=req.artifacts_produced, + milestone_id=req.milestone_id, + ) + orch.save_state() + + if hasattr(request.app.state, "event_bus") and request.app.state.event_bus: + try: + request.app.state.event_bus.publish( + "research_phase_transition", + { + "project_id": project.id, + "transition": transition.to_dict(), + "current_phase": orch.current_phase.value, + }, + project_id=str(project.id), + ) + except Exception as ex: + logger.warning("Failed to publish research_phase_transition event: %s", ex) + + return { + "status": "transitioned", + "transition": transition.to_dict(), + "state": orch.state.to_dict(), + "guidelines": orch.get_phase_guidelines(), + } + + +@router.post("/api/projects/{project_id}/research/milestones") +async def create_research_milestone( + project_id: int, + req: MilestoneCreateRequest, + db: AsyncSession = Depends(get_db), + user: User = Depends(get_current_user), +) -> dict[str, Any]: + """Add a new milestone to the project's research state.""" + project = await ops.get_project_by_id(db, project_id) + if not project or project.user_id != user.id: + raise HTTPException(status_code=404, detail="Project not found") + + target_phase = None + if req.phase: + try: + target_phase = ResearchPhase(req.phase) + except ValueError: + raise HTTPException(status_code=400, detail=f"Invalid phase '{req.phase}'") + + orch = _load_orchestrator(user.id, project.slug) + milestone = orch.add_milestone( + title=req.title, + description=req.description, + phase=target_phase, + criteria=req.criteria, + ) + orch.save_state() + + return { + "status": "created", + "milestone": milestone.to_dict(), + "milestones_count": len(orch.state.milestones), + } + + +@router.put("/api/projects/{project_id}/research/milestones/{milestone_id}") +async def update_research_milestone( + project_id: int, + milestone_id: str, + req: MilestoneUpdateRequest, + db: AsyncSession = Depends(get_db), + user: User = Depends(get_current_user), +) -> dict[str, Any]: + """Update milestone status or mark completed with output artifacts.""" + project = await ops.get_project_by_id(db, project_id) + if not project or project.user_id != user.id: + raise HTTPException(status_code=404, detail="Project not found") + + orch = _load_orchestrator(user.id, project.slug) + target_m = next((m for m in orch.state.milestones if m.milestone_id == milestone_id), None) + if not target_m: + raise HTTPException(status_code=404, detail=f"Milestone '{milestone_id}' not found") + + if req.status: + try: + target_m.status = MilestoneStatus(req.status) + except ValueError: + raise HTTPException(status_code=400, detail=f"Invalid status '{req.status}'") + + if target_m.status == MilestoneStatus.COMPLETED: + orch.complete_milestone(milestone_id, output_artifacts=req.output_artifacts) + else: + if req.output_artifacts: + target_m.output_artifacts.extend(req.output_artifacts) + orch.state.updated_at = __import__("time").time() + + orch.save_state() + + return { + "status": "updated", + "milestone": target_m.to_dict(), + } + + +@router.post("/api/projects/{project_id}/research/artifacts") +async def add_research_artifact( + project_id: int, + req: ArtifactCreateRequest, + db: AsyncSession = Depends(get_db), + user: User = Depends(get_current_user), +) -> dict[str, Any]: + """Register a scientific research artifact (paper, hypothesis, experiment, metrics, manuscript section, bibtex).""" + project = await ops.get_project_by_id(db, project_id) + if not project or project.user_id != user.id: + raise HTTPException(status_code=404, detail="Project not found") + + orch = _load_orchestrator(user.id, project.slug) + atype = req.type.lower().strip() + + if atype == "paper": + if isinstance(req.data, dict): + orch.add_paper(req.data) + else: + orch.add_paper({"title": str(req.data)}) + elif atype == "hypothesis": + if isinstance(req.data, dict): + orch.add_hypothesis(req.data) + else: + orch.add_hypothesis({"claim": str(req.data)}) + elif atype == "experiment": + if isinstance(req.data, dict): + orch.add_experiment(req.data) + else: + orch.add_experiment({"description": str(req.data)}) + elif atype == "metrics": + if isinstance(req.data, dict): + orch.update_metrics(req.data) + else: + raise HTTPException(status_code=400, detail="Metrics artifact data must be a dictionary") + elif atype == "manuscript_section": + sec_name = req.section_name or "main" + orch.update_manuscript_section(sec_name, str(req.data)) + elif atype == "bibtex": + orch.add_bibtex(str(req.data)) + else: + raise HTTPException( + status_code=400, + detail=f"Unsupported artifact type '{req.type}'. Supported types: paper, hypothesis, experiment, metrics, manuscript_section, bibtex", + ) + + orch.save_state() + + return { + "status": "artifact_registered", + "type": atype, + "artifacts_summary": { + "papers": len(orch.state.artifacts.papers), + "hypotheses": len(orch.state.artifacts.hypotheses), + "experiments": len(orch.state.artifacts.experiments), + "metrics_keys": list(orch.state.artifacts.metrics.keys()), + "sections": list(orch.state.artifacts.manuscript_sections.keys()), + "bibtex_count": len(orch.state.artifacts.bibtex_entries), + }, + } diff --git a/backend/openmlr/routes/sweeps.py b/backend/openmlr/routes/sweeps.py new file mode 100644 index 0000000..86c6613 --- /dev/null +++ b/backend/openmlr/routes/sweeps.py @@ -0,0 +1,264 @@ +"""Hyperparameter Sweep and HPO REST API routes. + +Provides endpoints for creating sweeps, suggesting trial parameters (Grid, Random, Bayesian, Hyperband), +early-stopping checks, recording trial outcomes, parameter sensitivity analysis, and exporting reports. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException, Query, status +from pydantic import BaseModel, Field + +from ..db.models import User +from ..dependencies import get_current_user_optional +from ..services.sweep_engine import SweepEngine +from .projects import WORKSPACES_ROOT + +router = APIRouter(tags=["sweeps"]) +logger = logging.getLogger(__name__) + +_global_engine = SweepEngine() + + +def _get_engine(project_uuid: str | None = None) -> SweepEngine: + """Get the active SweepEngine with project workspace isolation.""" + if project_uuid: + base_dir = WORKSPACES_ROOT / project_uuid / ".project-meta" / "sweeps" + return SweepEngine(base_dir=base_dir) + return _global_engine + + +class CreateSweepRequest(BaseModel): + """Payload for creating a new hyperparameter sweep.""" + + name: str = Field(..., min_length=1, max_length=255, description="Name of the sweep") + description: str = Field(default="", description="Description or research objective") + method: str = Field(default="random", description="Search algorithm: grid, random, bayesian, hyperband") + objective_metric: str = Field(default="val_loss", description="Target metric name") + goal: str = Field(default="minimize", description="Optimization goal: minimize or maximize") + parameters: dict[str, Any] = Field(..., description="Parameter space search specifications") + max_trials: int = Field(default=10, ge=1, le=500, description="Max trials to sample") + early_stopping: dict[str, Any] = Field(default_factory=dict, description="Early stopping / pruning configuration") + project_uuid: str | None = Field(default=None, description="Associated project UUID") + + +class RecordTrialRequest(BaseModel): + """Payload for logging trial evaluation results.""" + + metrics: dict[str, Any] = Field(..., description="Evaluation metrics dictionary") + status: str = Field(default="completed", description="Status: completed, failed, pruned") + step_history: list[dict[str, Any]] = Field(default_factory=list, description="Per-step metric trajectories") + error_message: str | None = Field(default=None, description="Optional error message") + + +class PruneCheckRequest(BaseModel): + """Payload for early-stopping evaluate check.""" + + current_step: int = Field(..., ge=1, description="Current training step or epoch") + current_metric_val: float = Field(..., description="Current value of the objective metric") + + +@router.post("/api/sweeps", status_code=status.HTTP_201_CREATED) +@router.post("/api/projects/{project_uuid}/sweeps", status_code=status.HTTP_201_CREATED) +async def create_sweep( + payload: CreateSweepRequest, + project_uuid: str | None = None, + current_user: User | None = Depends(get_current_user_optional), +): + """Create and initialize a new hyperparameter sweep.""" + proj = project_uuid or payload.project_uuid or "default" + engine = _get_engine(proj) + + try: + sweep = engine.create_sweep( + project_id=proj, + name=payload.name, + method=payload.method, + objective_metric=payload.objective_metric, + goal=payload.goal, + parameters=payload.parameters, + max_trials=payload.max_trials, + description=payload.description, + early_stopping=payload.early_stopping, + ) + return {"sweep": sweep.to_dict()} + except Exception as e: + logger.exception("Failed to create sweep: %s", e) + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Failed to create sweep: {e}", + ) + + +@router.get("/api/sweeps") +@router.get("/api/projects/{project_uuid}/sweeps") +async def list_sweeps( + project_uuid: str | None = None, + project_id: str | None = Query(default=None), + current_user: User | None = Depends(get_current_user_optional), +): + """List all hyperparameter sweeps in a project.""" + proj = project_uuid or project_id or "default" + engine = _get_engine(proj) + sweeps = engine.list_sweeps(proj) + return { + "project_id": proj, + "total": len(sweeps), + "sweeps": [s.to_dict() for s in sweeps], + } + + +@router.get("/api/sweeps/{sweep_id}") +@router.get("/api/projects/{project_uuid}/sweeps/{sweep_id}") +async def get_sweep_details( + sweep_id: str, + project_uuid: str | None = None, + project_id: str | None = Query(default=None), + current_user: User | None = Depends(get_current_user_optional), +): + """Get full sweep configuration, trial list, and current status.""" + proj = project_uuid or project_id or "default" + engine = _get_engine(proj) + sweep = engine.get_sweep(proj, sweep_id) + if not sweep: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Sweep '{sweep_id}' not found in project '{proj}'", + ) + return {"sweep": sweep.to_dict()} + + +@router.post("/api/sweeps/{sweep_id}/suggest") +@router.post("/api/projects/{project_uuid}/sweeps/{sweep_id}/suggest") +async def suggest_next_trial( + sweep_id: str, + project_uuid: str | None = None, + project_id: str | None = Query(default=None), + current_user: User | None = Depends(get_current_user_optional), +): + """Generate the next parameter candidate proposal for the sweep.""" + proj = project_uuid or project_id or "default" + engine = _get_engine(proj) + try: + trial = engine.suggest_trial(proj, sweep_id) + if not trial: + return {"trial": None, "message": "Sweep max trials reached or completed"} + return {"trial": trial.to_dict()} + except ValueError as ve: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(ve)) + except Exception as e: + logger.exception("Failed to suggest trial: %s", e) + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) + + +@router.post("/api/sweeps/{sweep_id}/trials/{trial_id}/record") +@router.post("/api/projects/{project_uuid}/sweeps/{sweep_id}/trials/{trial_id}/record") +async def record_trial_outcome( + sweep_id: str, + trial_id: str, + payload: RecordTrialRequest, + project_uuid: str | None = None, + project_id: str | None = Query(default=None), + current_user: User | None = Depends(get_current_user_optional), +): + """Record metrics, completion status, or failure for a trial.""" + proj = project_uuid or project_id or "default" + engine = _get_engine(proj) + try: + trial = engine.record_trial_result( + project_id=proj, + sweep_id=sweep_id, + trial_id=trial_id, + metrics=payload.metrics, + status=payload.status, + step_history=payload.step_history, + error_message=payload.error_message, + ) + return {"trial": trial.to_dict()} + except ValueError as ve: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(ve)) + except Exception as e: + logger.exception("Failed to record trial: %s", e) + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) + + +@router.post("/api/sweeps/{sweep_id}/trials/{trial_id}/prune-check") +@router.post("/api/projects/{project_uuid}/sweeps/{sweep_id}/trials/{trial_id}/prune-check") +async def check_trial_pruning( + sweep_id: str, + trial_id: str, + payload: PruneCheckRequest, + project_uuid: str | None = None, + project_id: str | None = Query(default=None), + current_user: User | None = Depends(get_current_user_optional), +): + """Evaluate whether the trial should be early stopped.""" + proj = project_uuid or project_id or "default" + engine = _get_engine(proj) + should_prune = engine.should_prune_trial( + project_id=proj, + sweep_id=sweep_id, + trial_id=trial_id, + current_step=payload.current_step, + current_metric_val=payload.current_metric_val, + ) + return {"trial_id": trial_id, "should_prune": should_prune} + + +@router.get("/api/sweeps/{sweep_id}/analysis") +@router.get("/api/projects/{project_uuid}/sweeps/{sweep_id}/analysis") +async def get_sweep_analysis( + sweep_id: str, + project_uuid: str | None = None, + project_id: str | None = Query(default=None), + current_user: User | None = Depends(get_current_user_optional), +): + """Calculate parameter sensitivities, correlation matrix, optimal trial, and Pareto frontier.""" + proj = project_uuid or project_id or "default" + engine = _get_engine(proj) + try: + analysis = engine.analyze_sweep(proj, sweep_id) + return {"analysis": analysis} + except ValueError as ve: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(ve)) + except Exception as e: + logger.exception("Failed to analyze sweep: %s", e) + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)) + + +@router.post("/api/sweeps/{sweep_id}/export") +@router.post("/api/projects/{project_uuid}/sweeps/{sweep_id}/export") +async def export_sweep_report( + sweep_id: str, + project_uuid: str | None = None, + project_id: str | None = Query(default=None), + current_user: User | None = Depends(get_current_user_optional), +): + """Export markdown report for research papers and ablation sections.""" + proj = project_uuid or project_id or "default" + engine = _get_engine(proj) + report = engine.export_sweep_markdown(proj, sweep_id) + return {"report": report} + + +@router.delete("/api/sweeps/{sweep_id}") +@router.delete("/api/projects/{project_uuid}/sweeps/{sweep_id}") +async def delete_sweep( + sweep_id: str, + project_uuid: str | None = None, + project_id: str | None = Query(default=None), + current_user: User | None = Depends(get_current_user_optional), +): + """Delete a sweep and all trial records.""" + proj = project_uuid or project_id or "default" + engine = _get_engine(proj) + deleted = engine.delete_sweep(proj, sweep_id) + if not deleted: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Sweep '{sweep_id}' not found", + ) + return {"deleted": True, "sweep_id": sweep_id} diff --git a/backend/openmlr/services/dataset_profiler.py b/backend/openmlr/services/dataset_profiler.py new file mode 100644 index 0000000..6b7ead9 --- /dev/null +++ b/backend/openmlr/services/dataset_profiler.py @@ -0,0 +1,450 @@ +"""Dataset Profiler — statistical profiling, validation, and split manager for ML datasets.""" + +from __future__ import annotations + +import csv +import hashlib +import json +import logging +import math +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +log = logging.getLogger(__name__) + + +def _hash_score(seed: int, index: int, item: Any) -> bytes: + """Generate a deterministic SHA-256 hash digest for an item.""" + encoded = f"{seed}:{index}:{json.dumps(item, sort_keys=True, default=str)}".encode() + return hashlib.sha256(encoded).digest() + + +def _hash_sample(items: list[Any], k: int, seed: int = 42) -> list[Any]: + """Deterministically sample k items using cryptographic hash sorting (safe & reproducible).""" + if not items or k <= 0: + return [] + if k >= len(items): + return list(items) + scored = [(_hash_score(seed, idx, item), item) for idx, item in enumerate(items)] + scored.sort(key=lambda x: x[0]) + return [item for _, item in scored[:k]] + + +def _hash_shuffle(items: list[Any], seed: int = 42) -> list[Any]: + """Deterministically shuffle items using cryptographic hash sorting.""" + if not items: + return [] + scored = [(_hash_score(seed, idx, item), item) for idx, item in enumerate(items)] + scored.sort(key=lambda x: x[0]) + return [item for _, item in scored] + + +@dataclass +class ColumnProfile: + """Statistical summary of a single dataset feature/column.""" + + name: str + dtype: str # numeric, text, categorical, boolean, unknown + total_count: int + null_count: int + null_percentage: float + unique_count: int + stats: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class DatasetProfile: + """Comprehensive statistical profile of a dataset.""" + + file_path: str + format: str + total_rows: int + total_columns: int + file_size_bytes: int + columns: dict[str, ColumnProfile] = field(default_factory=dict) + health_score: int = 100 + warnings: list[str] = field(default_factory=list) + summary: str = "" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +class DatasetProfiler: + """High-performance dataset profiler, validator, and split generator.""" + + @staticmethod + def detect_format(file_path: str | Path) -> str: + ext = Path(file_path).suffix.lower() + mapping = { + ".csv": "csv", + ".tsv": "tsv", + ".tab": "tsv", + ".jsonl": "jsonl", + ".ndjson": "jsonl", + ".json": "json", + ".txt": "text", + } + return mapping.get(ext, "csv") + + @classmethod + def _load_csv(cls, path: Path, delimiter: str, limit: int | None) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + with open(path, encoding="utf-8", errors="replace") as f: + reader = csv.DictReader(f, delimiter=delimiter) + for row in reader: + records.append(dict(row)) + if limit and len(records) >= limit: + break + return records + + @classmethod + def _load_jsonl(cls, path: Path, limit: int | None) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + with open(path, encoding="utf-8", errors="replace") as f: + for line in f: + line_str = line.strip() + if not line_str: + continue + try: + obj = json.loads(line_str) + if isinstance(obj, dict): + records.append(obj) + except Exception: + continue + if limit and len(records) >= limit: + break + return records + + @classmethod + def _load_json(cls, path: Path, limit: int | None) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + with open(path, encoding="utf-8", errors="replace") as f: + try: + data = json.load(f) + items = ( + data + if isinstance(data, list) + else next( + ( + data[k] + for k in ("data", "rows", "items", "records", "samples") + if isinstance(data.get(k), list) + ), + [data], + ) + ) + for item in items[:limit] if limit else items: + if isinstance(item, dict): + records.append(item) + except Exception as e: + log.warning("JSON parse error: %s", e) + return records + + @classmethod + def _load_text(cls, path: Path, limit: int | None) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + with open(path, encoding="utf-8", errors="replace") as f: + for idx, line in enumerate(f): + if line.strip(): + records.append({"line_number": idx + 1, "text": line.strip()}) + if limit and len(records) >= limit: + break + return records + + @classmethod + def load_records( + cls, file_path: str | Path, limit: int | None = None + ) -> tuple[list[dict[str, Any]], str, int]: + path = Path(file_path).resolve() + if not path.exists(): + raise FileNotFoundError(f"Dataset file not found: {file_path}") + + fmt = cls.detect_format(path) + file_size = path.stat().st_size + + if fmt == "csv": + records = cls._load_csv(path, delimiter=",", limit=limit) + elif fmt == "tsv": + records = cls._load_csv(path, delimiter="\t", limit=limit) + elif fmt == "jsonl": + records = cls._load_jsonl(path, limit=limit) + elif fmt == "json": + records = cls._load_json(path, limit=limit) + elif fmt == "text": + records = cls._load_text(path, limit=limit) + else: + records = cls._load_csv(path, delimiter=",", limit=limit) + + return records, fmt, file_size + + @classmethod + def profile(cls, file_path: str | Path, sample_size: int = 5000) -> DatasetProfile: + records, fmt, file_size = cls.load_records(file_path, limit=sample_size) + total_rows = len(records) + if not records: + return DatasetProfile( + file_path=str(file_path), + format=fmt, + total_rows=0, + total_columns=0, + file_size_bytes=file_size, + health_score=0, + warnings=["Dataset is empty or unparseable."], + summary="Empty dataset.", + ) + + all_columns = sorted({k for r in records for k in r.keys()}) + columns_prof: dict[str, ColumnProfile] = {} + warnings: list[str] = [] + penalty = 0 + + for col in all_columns: + vals = [r.get(col) for r in records] + non_nulls = [v for v in vals if v is not None and str(v).strip() != ""] + null_cnt = len(vals) - len(non_nulls) + null_pct = round((null_cnt / len(vals)) * 100, 2) if vals else 0.0 + + if null_pct > 20.0: + warnings.append(f"Column '{col}' has {null_pct}% missing values.") + penalty += min(15, int(null_pct // 2)) + + dtype, stats = cls._analyze_col(non_nulls, len(vals)) + uniq_cnt = stats.get("unique_count", len({str(v) for v in non_nulls})) + + if dtype == "categorical" and stats.get("imbalance_ratio", 1) > 10.0 and uniq_cnt > 1: + warnings.append(f"Column '{col}' has class imbalance ({stats['imbalance_ratio']}x).") + penalty += 10 + + columns_prof[col] = ColumnProfile( + col, dtype, len(vals), null_cnt, null_pct, uniq_cnt, stats + ) + + sample_hashes = [hash(json.dumps(r, sort_keys=True, default=str)) for r in records] + dup_cnt = total_rows - len(set(sample_hashes)) + if dup_cnt > 0: + dup_pct = round((dup_cnt / total_rows) * 100, 2) + warnings.append(f"Found {dup_cnt} ({dup_pct}%) duplicate rows in sample.") + penalty += min(20, int(dup_pct)) + + health = max(0, 100 - penalty) + summary = f"'{Path(file_path).name}': {total_rows} rows, {len(all_columns)} cols. Health: {health}/100." + return DatasetProfile( + str(file_path), + fmt, + total_rows, + len(all_columns), + file_size, + columns_prof, + health, + warnings, + summary, + ) + + @classmethod + def _analyze_col(cls, values: list[Any], total_cnt: int) -> tuple[str, dict[str, Any]]: + if not values: + return "unknown", {"unique_count": 0} + + # Boolean + if all(str(v).lower() in {"true", "false", "0", "1", "yes", "no", "t", "f"} for v in values): + return "boolean", { + "unique_count": len({str(v).lower() for v in values}), + "true_count": sum(1 for v in values if str(v).lower() in ("true", "1", "yes", "t")), + } + + # Numeric + try: + nums = sorted(float(v) for v in values) + n = len(nums) + mean = sum(nums) / n + var = sum((x - mean) ** 2 for x in nums) / max(1, n - 1) + q25, med, q75 = nums[int(0.25 * n)], nums[int(0.5 * n)], nums[int(0.75 * n)] + iqr = q75 - q25 + outliers = sum(1 for x in nums if x < (q25 - 1.5 * iqr) or x > (q75 + 1.5 * iqr)) + return "numeric", { + "min": round(nums[0], 4), + "max": round(nums[-1], 4), + "mean": round(mean, 4), + "std": round(math.sqrt(var), 4), + "median": round(med, 4), + "q25": round(q25, 4), + "q75": round(q75, 4), + "outlier_count": outliers, + "unique_count": len(set(nums)), + } + except (ValueError, TypeError): + pass + + # String / Text / Categorical + strs = [str(v) for v in values] + uniq_cnt = len(set(strs)) + avg_len = sum(len(s) for s in strs) / len(strs) + + if uniq_cnt <= min(50, max(5, total_cnt // 10)) and avg_len < 64: + freqs = sorted( + ((k, sum(1 for x in strs if x == k)) for k in set(strs)), + key=lambda x: x[1], + reverse=True, + ) + imbalance = round(freqs[0][1] / max(1, freqs[-1][1]), 2) + return "categorical", { + "unique_count": uniq_cnt, + "top_classes": dict(freqs[:10]), + "imbalance_ratio": imbalance, + "class_distribution": { + k: round((v / len(strs)) * 100, 2) for k, v in freqs[:10] + }, + } + + word_cnts = [len(s.split()) for s in strs] + toks = sorted(int(w * 1.3) + 1 for w in word_cnts) + n_tok = len(toks) + over512 = sum(1 for t in toks if t > 512) + + return "text", { + "unique_count": uniq_cnt, + "char_len_avg": round(avg_len, 2), + "token_est_mean": round(sum(toks) / n_tok, 1), + "token_est_p95": toks[min(n_tok - 1, int(0.95 * n_tok))], + "token_est_max": toks[-1], + "overflow_512_count": over512, + "overflow_512_pct": round((over512 / n_tok) * 100, 2), + } + + @classmethod + def sample_records( + cls, + file_path: str | Path, + n: int = 5, + offset: int = 0, + strategy: str = "head", + label_column: str | None = None, + seed: int = 42, + ) -> list[dict[str, Any]]: + records, _, _ = cls.load_records(file_path) + if not records: + return [] + if strategy == "random": + return _hash_sample(records, k=n, seed=seed) + if strategy == "stratified" and label_column: + classes: dict[str, list[dict[str, Any]]] = {} + for r in records: + classes.setdefault(str(r.get(label_column, "missing")), []).append(r) + sampled: list[dict[str, Any]] = [] + per_class = max(1, n // max(1, len(classes))) + for cl_name in sorted(classes.keys()): + cl_records = classes[cl_name] + sampled.extend(_hash_sample(cl_records, k=per_class, seed=seed)) + return sampled[:n] + return records[offset : offset + n] + + @classmethod + def validate_dataset( + cls, + file_path: str | Path, + expected_columns: list[str] | None = None, + max_null_pct: float = 20.0, + max_token_length: int | None = None, + ) -> dict[str, Any]: + profile = cls.profile(file_path, sample_size=3000) + errors: list[str] = [] + if profile.total_rows == 0: + return {"valid": False, "errors": ["Dataset is empty."], "health_score": 0} + + if expected_columns: + missing = [c for c in expected_columns if c not in profile.columns] + if missing: + errors.append(f"Missing required columns: {', '.join(missing)}") + + for col, prof in profile.columns.items(): + if prof.null_percentage > max_null_pct: + errors.append(f"Column '{col}' null percentage {prof.null_percentage}% > {max_null_pct}%") + if max_token_length and prof.dtype == "text" and prof.stats.get("token_est_max", 0) > max_token_length: + errors.append(f"Column '{col}' max tokens exceed {max_token_length}") + + return { + "valid": len(errors) == 0, + "errors": errors, + "warnings": profile.warnings, + "health_score": profile.health_score, + "total_rows": profile.total_rows, + "total_columns": profile.total_columns, + } + + @classmethod + def split_dataset( + cls, + file_path: str | Path, + output_dir: str | Path, + train_ratio: float = 0.8, + val_ratio: float = 0.1, + test_ratio: float = 0.1, + stratify_column: str | None = None, + seed: int = 42, + ) -> dict[str, Any]: + path = Path(file_path).resolve() + out_dir = Path(output_dir).resolve() + out_dir.mkdir(parents=True, exist_ok=True) + records, fmt, _ = cls.load_records(path) + if not records: + raise ValueError(f"Cannot split empty dataset: {file_path}") + + tot = train_ratio + val_ratio + test_ratio + tr_r, val_r = train_ratio / tot, val_ratio / tot + train_set, val_set, test_set = [], [], [] + + if stratify_column: + classes: dict[str, list[dict[str, Any]]] = {} + for r in records: + classes.setdefault(str(r.get(stratify_column, "missing")), []).append(r) + for cl_records in classes.values(): + shuffled = _hash_shuffle(cl_records, seed=seed) + n = len(shuffled) + n_tr, n_v = int(n * tr_r), int(n * val_r) + train_set.extend(shuffled[:n_tr]) + val_set.extend(shuffled[n_tr : n_tr + n_v]) + test_set.extend(shuffled[n_tr + n_v :]) + else: + shuffled = _hash_shuffle(records, seed=seed) + n = len(shuffled) + n_tr, n_v = int(n * tr_r), int(n * val_r) + train_set = shuffled[:n_tr] + val_set = shuffled[n_tr : n_tr + n_v] + test_set = shuffled[n_tr + n_v :] + + ext = ".csv" if fmt == "csv" else ".jsonl" + tr_p, val_p, test_p = out_dir / f"train{ext}", out_dir / f"val{ext}", out_dir / f"test{ext}" + cls._write_records(train_set, tr_p, fmt) + cls._write_records(val_set, val_p, fmt) + cls._write_records(test_set, test_p, fmt) + + manifest = { + "source_file": str(path), + "stratified_by": stratify_column, + "seed": seed, + "total_records": len(records), + "train_count": len(train_set), + "val_count": len(val_set), + "test_count": len(test_set), + "splits": {"train": str(tr_p), "val": str(val_p), "test": str(test_p)}, + } + with open(out_dir / "split_manifest.json", "w", encoding="utf-8") as f: + json.dump(manifest, f, indent=2) + return manifest + + @staticmethod + def _write_records(records: list[dict[str, Any]], out_path: Path, fmt: str) -> None: + if not records: + out_path.touch() + return + if fmt == "csv": + with open(out_path, "w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter(f, fieldnames=list(records[0].keys())) + writer.writeheader() + writer.writerows(records) + else: + with open(out_path, "w", encoding="utf-8") as f: + for r in records: + f.write(json.dumps(r, default=str) + "\n") diff --git a/backend/openmlr/services/experiment_tracker.py b/backend/openmlr/services/experiment_tracker.py new file mode 100644 index 0000000..b06d1f4 --- /dev/null +++ b/backend/openmlr/services/experiment_tracker.py @@ -0,0 +1,401 @@ +"""Service for tracking machine learning experiment runs, metric curves, hyperparameters, and checkpoints.""" + +from __future__ import annotations + +import json +import logging +import time +import uuid +from dataclasses import asdict, dataclass, field +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from ..services.event_bus import EventBus + +log = logging.getLogger(__name__) + + +@dataclass +class MetricPoint: + step: int + epoch: int + timestamp: float + value: float + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class CheckpointData: + id: str + name: str + step: int + epoch: int + timestamp: float + path: str + file_size_mb: float + metrics: dict[str, float] + download_url: str = "" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class ExperimentRunData: + id: str + name: str + description: str = "" + status: str = "running" # running, completed, failed, paused, idle + started_at: str = field(default_factory=lambda: datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S")) + ended_at: str | None = None + duration_seconds: int = 0 + compute_target: str = "Local GPU" + tags: list[str] = field(default_factory=list) + hyperparameters: dict[str, Any] = field(default_factory=dict) + current_step: int = 0 + total_steps: int = 100 + current_epoch: int = 1 + total_epochs: int = 1 + best_val_loss: float | None = None + metrics: dict[str, list[MetricPoint]] = field(default_factory=lambda: { + "train_loss": [], + "val_loss": [], + "learning_rate": [], + "gpu_utilization": [], + "memory_used": [], + "throughput": [], + }) + checkpoints: list[CheckpointData] = field(default_factory=list) + logs: list[str] = field(default_factory=list) + project_uuid: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "name": self.name, + "description": self.description, + "status": self.status, + "started_at": self.started_at, + "ended_at": self.ended_at, + "duration_seconds": self.duration_seconds, + "compute_target": self.compute_target, + "tags": list(self.tags), + "hyperparameters": dict(self.hyperparameters), + "current_step": self.current_step, + "total_steps": self.total_steps, + "current_epoch": self.current_epoch, + "total_epochs": self.total_epochs, + "best_val_loss": self.best_val_loss, + "metrics": { + k: [pt.to_dict() for pt in v] for k, v in self.metrics.items() + }, + "checkpoints": [cp.to_dict() for cp in self.checkpoints], + "logs": list(self.logs), + "project_uuid": self.project_uuid, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> ExperimentRunData: + metrics_raw = data.get("metrics", {}) + metrics: dict[str, list[MetricPoint]] = {} + for key, pts in metrics_raw.items(): + metrics[key] = [ + MetricPoint(**pt) if isinstance(pt, dict) else pt for pt in pts + ] + + checkpoints_raw = data.get("checkpoints", []) + checkpoints = [ + CheckpointData(**cp) if isinstance(cp, dict) else cp + for cp in checkpoints_raw + ] + + return cls( + id=data["id"], + name=data["name"], + description=data.get("description", ""), + status=data.get("status", "running"), + started_at=data.get("started_at", datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S")), + ended_at=data.get("ended_at"), + duration_seconds=data.get("duration_seconds", 0), + compute_target=data.get("compute_target", "Local GPU"), + tags=data.get("tags", []), + hyperparameters=data.get("hyperparameters", {}), + current_step=data.get("current_step", 0), + total_steps=data.get("total_steps", 100), + current_epoch=data.get("current_epoch", 1), + total_epochs=data.get("total_epochs", 1), + best_val_loss=data.get("best_val_loss"), + metrics=metrics, + checkpoints=checkpoints, + logs=data.get("logs", []), + project_uuid=data.get("project_uuid"), + ) + + +class ExperimentTracker: + """In-memory and file-backed experiment registry and metrics aggregator.""" + + def __init__(self, event_bus: EventBus | None = None, storage_dir: Path | None = None): + self.event_bus = event_bus + self.storage_dir = storage_dir + self._runs: dict[str, ExperimentRunData] = {} + if self.storage_dir: + self.storage_dir.mkdir(parents=True, exist_ok=True) + self._load_stored_runs() + + def _load_stored_runs(self) -> None: + if not self.storage_dir or not self.storage_dir.exists(): + return + for file in self.storage_dir.glob("run_*.json"): + try: + data = json.loads(file.read_text(encoding="utf-8")) + run = ExperimentRunData.from_dict(data) + self._runs[run.id] = run + except Exception as e: + log.warning("Failed to load run from %s: %s", file, e) + + def _persist_run(self, run: ExperimentRunData) -> None: + if not self.storage_dir: + return + file = self.storage_dir / f"run_{run.id}.json" + try: + file.write_text(json.dumps(run.to_dict(), indent=2), encoding="utf-8") + except Exception as e: + log.warning("Failed to persist run %s: %s", run.id, e) + + async def _emit_event(self, event_type: str, payload: dict[str, Any]) -> None: + if self.event_bus: + try: + await self.event_bus.broadcast({ + "type": event_type, + "channel": "experiments", + "data": payload, + }) + except Exception as e: + log.warning("Failed to emit experiment event %s: %s", event_type, e) + + def create_run( + self, + name: str, + description: str = "", + hyperparameters: dict[str, Any] | None = None, + compute_target: str = "Local GPU", + tags: list[str] | None = None, + total_steps: int = 100, + total_epochs: int = 1, + project_uuid: str | None = None, + ) -> ExperimentRunData: + """Create and register a new experiment run.""" + run_id = f"run-{uuid.uuid4().hex[:8]}" + run = ExperimentRunData( + id=run_id, + name=name, + description=description, + status="running", + compute_target=compute_target, + tags=tags or [], + hyperparameters=hyperparameters or {}, + total_steps=total_steps, + total_epochs=total_epochs, + project_uuid=project_uuid, + ) + self._runs[run_id] = run + self._persist_run(run) + return run + + def get_run(self, run_id: str) -> ExperimentRunData | None: + """Retrieve run by ID.""" + return self._runs.get(run_id) + + def list_runs( + self, + project_uuid: str | None = None, + status: str | None = None, + search: str | None = None, + limit: int = 50, + offset: int = 0, + ) -> tuple[list[ExperimentRunData], int]: + """List runs matching filter criteria with pagination.""" + items = list(self._runs.values()) + + if project_uuid: + items = [r for r in items if r.project_uuid == project_uuid] + + if status and status != "all": + items = [r for r in items if r.status == status] + + if search: + q = search.lower() + items = [ + r for r in items + if q in r.name.lower() or q in r.description.lower() or any(q in t.lower() for t in r.tags) + ] + + # Sort newest first + items.sort(key=lambda r: r.started_at, reverse=True) + total = len(items) + paginated = items[offset : offset + limit] + return paginated, total + + def log_metrics( + self, + run_id: str, + step: int, + epoch: int = 1, + metrics: dict[str, float] | None = None, + timestamp: float | None = None, + ) -> ExperimentRunData: + """Append metric values to a run and update progress.""" + run = self.get_run(run_id) + if not run: + raise KeyError(f"Run '{run_id}' not found") + + ts = timestamp or (time.time() * 1000.0) + run.current_step = max(run.current_step, step) + run.current_epoch = max(run.current_epoch, epoch) + + if metrics: + for metric_name, val in metrics.items(): + if metric_name not in run.metrics: + run.metrics[metric_name] = [] + + point = MetricPoint(step=step, epoch=epoch, timestamp=ts, value=float(val)) + run.metrics[metric_name].append(point) + + if metric_name in ("val_loss", "eval_loss"): + if run.best_val_loss is None or float(val) < run.best_val_loss: + run.best_val_loss = float(val) + + self._persist_run(run) + return run + + def update_status( + self, + run_id: str, + status: str, + reason: str | None = None, + ) -> ExperimentRunData: + """Update run execution status.""" + run = self.get_run(run_id) + if not run: + raise KeyError(f"Run '{run_id}' not found") + + valid_statuses = ("running", "completed", "failed", "paused", "idle") + if status not in valid_statuses: + raise ValueError(f"Invalid status '{status}'. Must be one of {valid_statuses}") + + run.status = status + if status in ("completed", "failed"): + run.ended_at = datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S") + if reason: + run.logs.append(f"[{datetime.now(UTC).strftime('%H:%M:%S')}] Run {status}: {reason}") + elif status == "running" and not run.started_at: + run.started_at = datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S") + + self._persist_run(run) + return run + + def append_logs(self, run_id: str, lines: list[str]) -> list[str]: + """Append log lines to an experiment run.""" + run = self.get_run(run_id) + if not run: + raise KeyError(f"Run '{run_id}' not found") + + run.logs.extend(lines) + # Cap log history to last 2000 lines + if len(run.logs) > 2000: + run.logs = run.logs[-2000:] + + self._persist_run(run) + return run.logs + + def register_checkpoint( + self, + run_id: str, + name: str, + step: int, + epoch: int, + path: str = "", + file_size_mb: float = 0.0, + metrics: dict[str, float] | None = None, + download_url: str = "", + ) -> CheckpointData: + """Register a model checkpoint saved during the experiment.""" + run = self.get_run(run_id) + if not run: + raise KeyError(f"Run '{run_id}' not found") + + cp_id = f"ckpt-{uuid.uuid4().hex[:8]}" + cp = CheckpointData( + id=cp_id, + name=name, + step=step, + epoch=epoch, + timestamp=time.time() * 1000.0, + path=path, + file_size_mb=file_size_mb, + metrics=metrics or {}, + download_url=download_url, + ) + run.checkpoints.append(cp) + self._persist_run(run) + return cp + + def compare_runs(self, run_ids: list[str]) -> dict[str, Any]: + """Perform side-by-side comparison across multiple experiment runs.""" + selected_runs: list[ExperimentRunData] = [ + r for rid in run_ids if (r := self.get_run(rid)) is not None + ] + if not selected_runs: + return {"runs": [], "metrics_summary": {}, "hyperparameters_comparison": {}} + + # Collect all unique hyperparameter keys + all_hp_keys: set[str] = set() + for r in selected_runs: + all_hp_keys.update(r.hyperparameters.keys()) + + hyperparameter_table: dict[str, dict[str, Any]] = {} + for k in sorted(all_hp_keys): + hyperparameter_table[k] = { + r.id: r.hyperparameters.get(k, None) for r in selected_runs + } + + # Compare best val losses and final train losses + metrics_summary: dict[str, dict[str, Any]] = {} + for r in selected_runs: + train_losses = [p.value for p in r.metrics.get("train_loss", [])] + val_losses = [p.value for p in r.metrics.get("val_loss", [])] + final_train_loss = train_losses[-1] if train_losses else None + min_train_loss = min(train_losses) if train_losses else None + min_val_loss = min(val_losses) if val_losses else r.best_val_loss + + metrics_summary[r.id] = { + "name": r.name, + "status": r.status, + "current_step": r.current_step, + "total_steps": r.total_steps, + "best_val_loss": min_val_loss, + "final_train_loss": final_train_loss, + "min_train_loss": min_train_loss, + "total_checkpoints": len(r.checkpoints), + } + + return { + "runs": [r.to_dict() for r in selected_runs], + "metrics_summary": metrics_summary, + "hyperparameters_comparison": hyperparameter_table, + } + + def delete_run(self, run_id: str) -> bool: + """Delete an experiment run.""" + if run_id in self._runs: + del self._runs[run_id] + if self.storage_dir: + file = self.storage_dir / f"run_{run_id}.json" + if file.exists(): + file.unlink() + return True + return False diff --git a/backend/openmlr/services/figure_generator.py b/backend/openmlr/services/figure_generator.py new file mode 100644 index 0000000..1387722 --- /dev/null +++ b/backend/openmlr/services/figure_generator.py @@ -0,0 +1,337 @@ +"""Publication Figure Generator service for paper plots, scripts, and LaTeX TikZ figures.""" + +from __future__ import annotations + +import json +import logging +import time +import uuid +from typing import Any + +from .figure_types import ( + FigureArtifact, + GenerateFigureRequest, + MultiPanelLayoutRequest, +) + +logger = logging.getLogger("openmlr.services.figure_generator") + +# Palette color definitions +PALETTES = { + "colorblind": ["#377eb8", "#ff7f00", "#4daf4a", "#f781bf", "#a65628", "#984ea3"], + "viridis": ["#440154", "#3b528b", "#21908d", "#5dc963", "#fde725"], + "muted": ["#4878d0", "#ee854a", "#6acc65", "#d65f5f", "#956cb4", "#8c613c"], + "deep": ["#4c72b0", "#dd8452", "#55a868", "#c44e52", "#8172b3", "#937860"], + "tableau": ["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b"], +} + +# In-memory figure store per project (in-memory + fallback) +_PROJECT_FIGURES: dict[str, dict[str, FigureArtifact]] = {} + + +def _generate_python_script(req: GenerateFigureRequest) -> str: + """Generate clean, standalone Matplotlib/Seaborn Python script for reproducibility.""" + colors = PALETTES.get(req.palette.value, PALETTES["colorblind"]) + data_json = json.dumps(req.series_data, indent=4) + cats_json = json.dumps(req.categories) + matrix_json = json.dumps(req.values_matrix) + + return f"""#!/usr/bin/env python3 +\"\"\"Publication-grade plot generator: {req.title}\"\"\" + +import matplotlib.pyplot as plt +import numpy as np + +# Set standard publication styling +plt.rcParams.update({{ + 'font.size': 10, + 'font.family': 'sans-serif', + 'axes.labelsize': 11, + 'axes.titlesize': 12, + 'xtick.labelsize': 9, + 'ytick.labelsize': 9, + 'legend.fontsize': 9, + 'figure.titlesize': 13, + 'figure.dpi': 300, + 'savefig.bbox': 'tight', +}}) + +colors = {colors} +series_data = {data_json} +categories = {cats_json} +values_matrix = {matrix_json} + +fig, ax = plt.subplots(figsize=({req.width_inches}, {req.height_inches})) + +# Plot logic based on plot_type +if "{req.plot_type.value}" == "loss_curve": + for idx, (s_name, points) in enumerate(series_data.items()): + xs = [p['x'] for p in points] + ys = [p['y'] for p in points] + c = colors[idx % len(colors)] + ax.plot(xs, ys, label=s_name, color=c, linewidth=2) + if any('y_err' in p and p['y_err'] is not None for p in points): + errs = [p.get('y_err', 0.0) or 0.0 for p in points] + ax.fill_between(xs, np.array(ys) - np.array(errs), np.array(ys) + np.array(errs), color=c, alpha=0.15) + ax.set_xlabel("{req.x_label}") + ax.set_ylabel("{req.y_label}") + ax.legend(frameon=True) + ax.grid(True, linestyle='--', alpha=0.5) + +elif "{req.plot_type.value}" == "ablation_bar": + if categories and series_data: + x_indices = np.arange(len(categories)) + num_series = len(series_data) + bar_width = 0.8 / max(num_series, 1) + for idx, (s_name, points) in enumerate(series_data.items()): + ys = [p['y'] for p in points] + c = colors[idx % len(colors)] + offset = (idx - num_series / 2 + 0.5) * bar_width + ax.bar(x_indices + offset, ys, bar_width, label=s_name, color=c) + ax.set_xticks(x_indices) + ax.set_xticklabels(categories, rotation=20, ha='right') + ax.set_ylabel("{req.y_label}") + ax.legend(frameon=True) + ax.grid(axis='y', linestyle='--', alpha=0.5) + +elif "{req.plot_type.value}" == "pareto_frontier": + for idx, (s_name, points) in enumerate(series_data.items()): + xs = [p['x'] for p in points] + ys = [p['y'] for p in points] + c = colors[idx % len(colors)] + ax.scatter(xs, ys, label=s_name, color=c, s=50, alpha=0.85) + ax.set_xlabel("{req.x_label}") + ax.set_ylabel("{req.y_label}") + ax.legend(frameon=True) + ax.grid(True, linestyle='--', alpha=0.5) + +elif "{req.plot_type.value}" in ("confusion_matrix", "heatmap"): + if values_matrix: + im = ax.imshow(values_matrix, cmap='Blues', aspect='auto') + fig.colorbar(im, ax=ax) + if categories: + ax.set_xticks(range(len(categories))) + ax.set_yticks(range(len(categories))) + ax.set_xticklabels(categories, rotation=45, ha='right') + ax.set_yticklabels(categories) + +ax.set_title("{req.title}") +plt.tight_layout() +plt.savefig("figure.pdf", dpi=300) +plt.savefig("figure.png", dpi=300) +print("Saved figure.pdf and figure.png successfully.") +""" + + +def _generate_latex_snippet(fig_id: str, req: GenerateFigureRequest) -> str: + """Generate a clean LaTeX figure environment snippet.""" + label = f"fig:{fig_id[:8]}" + caption = req.caption or f"{req.title}. {req.x_label} vs. {req.y_label} across model configurations." + return ( + "\\begin{figure}[htbp]\n" + " \\centering\n" + f" \\includegraphics[width=0.85\\linewidth]{{figures/{label}.pdf}}\n" + f" \\caption{{{caption}}}\n" + f" \\label{{{label}}}\n" + "\\end{figure}" + ) + + +def _generate_tikz_code(req: GenerateFigureRequest) -> str: + """Generate pure LaTeX TikZ/PGFPlots vector graphic code.""" + colors = PALETTES.get(req.palette.value, PALETTES["colorblind"]) + lines = [ + "\\begin{tikzpicture}", + "\\begin{axis}[", + f" title={{{req.title}}},", + f" xlabel={{{req.x_label}}},", + f" ylabel={{{req.y_label}}},", + " grid=major,", + " grid style={dashed,gray!30},", + " legend pos=north east,", + " legend cell align={left},", + f" width={req.width_inches * 1.5:.1f}cm,", + f" height={req.height_inches * 1.5:.1f}cm", + "]", + ] + + for idx, (s_name, points) in enumerate(req.series_data.items()): + c = colors[idx % len(colors)] + coords = " ".join(f"({p['x']},{p['y']})" for p in points) + lines.append(f"\\addplot[color={c}, thick, mark=*] coordinates {{ {coords} }};") + lines.append(f"\\addlegendentry{{{s_name}}}") + + lines.append("\\end{axis}") + lines.append("\\end{tikzpicture}") + return "\n".join(lines) + + +def _generate_svg_preview(req: GenerateFigureRequest) -> str: + """Generate an SVG string for instant client-side rendering.""" + colors = PALETTES.get(req.palette.value, PALETTES["colorblind"]) + w = 600 + h = 360 + pad_left = 60 + pad_right = 140 + pad_top = 40 + pad_bottom = 50 + + plot_w = w - pad_left - pad_right + plot_h = h - pad_top - pad_bottom + + svg_elements = [ + f'', + f'', + f'{req.title}', + f'', + f'', + f'{req.x_label}', + f'{req.y_label}', + ] + + all_xs: list[float] = [] + all_ys: list[float] = [] + for points in req.series_data.values(): + for p in points: + if isinstance(p["x"], (int, float)): + all_xs.append(float(p["x"])) + all_ys.append(float(p["y"])) + + min_x = min(all_xs) if all_xs else 0.0 + max_x = max(all_xs) if all_xs else 1.0 + min_y = min(all_ys) if all_ys else 0.0 + max_y = max(all_ys) if all_ys else 1.0 + if max_x == min_x: + max_x += 1.0 + if max_y == min_y: + max_y += 1.0 + + legend_y = pad_top + 10 + for idx, (s_name, points) in enumerate(req.series_data.items()): + c = colors[idx % len(colors)] + svg_points = [] + for p in points: + px_val = float(p["x"]) if isinstance(p["x"], (int, float)) else idx + py_val = float(p["y"]) + svg_x = pad_left + ((px_val - min_x) / (max_x - min_x)) * plot_w + svg_y = pad_top + plot_h - ((py_val - min_y) / (max_y - min_y)) * plot_h + svg_points.append((svg_x, svg_y)) + + if len(svg_points) > 1: + poly_str = " ".join(f"{x:.1f},{y:.1f}" for x, y in svg_points) + svg_elements.append(f'') + + for sx, sy in svg_points: + svg_elements.append(f'') + + # Legend entry + lx = pad_left + plot_w + 15 + ly = legend_y + idx * 20 + svg_elements.append(f'') + svg_elements.append(f'') + svg_elements.append(f'{s_name}') + + svg_elements.append("") + return "".join(svg_elements) + + +class FigureGeneratorService: + """Service to create, store, and export publication figures and multi-panel layouts.""" + + @classmethod + def generate_figure(cls, project_id: str, request: GenerateFigureRequest) -> FigureArtifact: + """Generate a complete figure artifact with code, LaTeX, TikZ, and SVG preview.""" + fig_id = f"fig_{uuid.uuid4().hex[:8]}" + python_script = _generate_python_script(request) + latex_snippet = _generate_latex_snippet(fig_id, request) + tikz_code = _generate_tikz_code(request) if request.generate_tikz else "" + svg_preview = _generate_svg_preview(request) + + artifact = FigureArtifact( + id=fig_id, + project_id=project_id, + title=request.title, + caption=request.caption, + plot_type=request.plot_type, + style_theme=request.style_theme, + palette=request.palette, + python_script=python_script, + latex_snippet=latex_snippet, + tikz_code=tikz_code, + svg_preview=svg_preview, + created_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + ) + + if project_id not in _PROJECT_FIGURES: + _PROJECT_FIGURES[project_id] = {} + _PROJECT_FIGURES[project_id][fig_id] = artifact + return artifact + + @classmethod + def list_figures(cls, project_id: str) -> list[FigureArtifact]: + """List all generated figure artifacts for a project.""" + figures_dict = _PROJECT_FIGURES.get(project_id, {}) + return list(figures_dict.values()) + + @classmethod + def get_figure(cls, project_id: str, figure_id: str) -> FigureArtifact | None: + """Get a single figure artifact.""" + return _PROJECT_FIGURES.get(project_id, {}).get(figure_id) + + @classmethod + def delete_figure(cls, project_id: str, figure_id: str) -> bool: + """Delete a figure artifact.""" + if project_id in _PROJECT_FIGURES and figure_id in _PROJECT_FIGURES[project_id]: + del _PROJECT_FIGURES[project_id][figure_id] + return True + return False + + @classmethod + def create_multipanel_layout( + cls, project_id: str, request: MultiPanelLayoutRequest + ) -> dict[str, Any]: + """Combine multiple figures into a multi-panel LaTeX subfigure environment.""" + figures: list[FigureArtifact] = [] + for fid in request.figure_ids: + fig = _PROJECT_FIGURES.get(project_id, {}).get(fid) + if fig is not None: + figures.append(fig) + + if not figures: + return {"error": "None of the requested figure IDs exist in this project."} + + cols = max(1, min(request.columns, 4)) + subfig_width = round(1.0 / cols - 0.03, 2) + + latex_lines = [ + "\\begin{figure*}[t]", + " \\centering", + ] + + letters = ["(a)", "(b)", "(c)", "(d)", "(e)", "(f)"] + for idx, fig in enumerate(figures): + subcap = request.subcaptions.get(fig.id, f"{letters[idx % len(letters)]} {fig.title}") + latex_lines.append( + f" \\begin{{subfigure}}{{{subfig_width}\\linewidth}}\n" + f" \\centering\n" + f" \\includegraphics[width=\\linewidth]{{figures/{fig.id}.pdf}}\n" + f" \\caption{{{subcap}}}\n" + f" \\label{{fig:{fig.id}}}\n" + f" \\end{{subfigure}}%" + ) + if (idx + 1) % cols == 0 and idx + 1 < len(figures): + latex_lines.append(" \\\\[1ex]") + + latex_lines.extend([ + f" \\caption{{{request.caption}}}\n" + f" \\label{{fig:multipanel_{uuid.uuid4().hex[:6]}}}\n" + "\\end{figure*}" + ]) + + return { + "title": request.title, + "caption": request.caption, + "figure_count": len(figures), + "latex_code": "\n".join(latex_lines), + "included_figures": [f.to_dict() for f in figures], + } diff --git a/backend/openmlr/services/figure_types.py b/backend/openmlr/services/figure_types.py new file mode 100644 index 0000000..7ab3b05 --- /dev/null +++ b/backend/openmlr/services/figure_types.py @@ -0,0 +1,93 @@ +"""Data models and types for the Publication Figure Studio.""" + +from __future__ import annotations + +from enum import Enum +from typing import Any + +from pydantic import BaseModel, Field + + +class PlotType(str, Enum): + """Supported scientific plot types.""" + LOSS_CURVE = "loss_curve" + ABLATION_BAR = "ablation_bar" + PARETO_FRONTIER = "pareto_frontier" + CONFUSION_MATRIX = "confusion_matrix" + RADAR_BENCHMARK = "radar_benchmark" + HEATMAP = "heatmap" + + +class StyleTheme(str, Enum): + """Conference styling presets.""" + NEURIPS = "neurips" + ICML = "icml" + ICLR = "iclr" + CVPR = "cvpr" + DARK = "dark" + + +class ColorPalette(str, Enum): + """Colorblind-safe and academic palettes.""" + COLORBLIND = "colorblind" + VIRIDIS = "viridis" + MUTED = "muted" + DEEP = "deep" + TABLEAU = "tableau" + + +class FigureDataPoint(BaseModel): + """Generic 2D/3D data point.""" + x: float | str + y: float + y_err: float | None = None + series: str = "default" + metadata: dict[str, Any] = Field(default_factory=dict) + + +class GenerateFigureRequest(BaseModel): + """Request payload for generating a publication figure.""" + title: str = Field(..., description="Figure title") + caption: str = Field("", description="LaTeX paper caption") + plot_type: PlotType = Field(PlotType.LOSS_CURVE, description="Type of plot") + style_theme: StyleTheme = Field(StyleTheme.NEURIPS, description="Conference theme standard") + palette: ColorPalette = Field(ColorPalette.COLORBLIND, description="Color palette") + x_label: str = Field("Step", description="X-axis label") + y_label: str = Field("Loss", description="Y-axis label") + series_data: dict[str, list[dict[str, Any]]] = Field( + default_factory=dict, + description="Named series data mapping (e.g. {'Baseline': [{'x': 1, 'y': 2.4}], 'Ours': [...]})" + ) + categories: list[str] = Field(default_factory=list, description="Categorical labels for bar/radar/heatmap") + values_matrix: list[list[float]] = Field(default_factory=list, description="2D matrix for heatmap/confusion matrix") + width_inches: float = Field(6.0, description="Figure width in inches for LaTeX") + height_inches: float = Field(4.0, description="Figure height in inches for LaTeX") + generate_tikz: bool = Field(True, description="Whether to also generate standalone TikZ / PGFPlots code") + + +class MultiPanelLayoutRequest(BaseModel): + """Request payload for generating a multi-panel subfigure grid.""" + title: str = Field(..., description="Overall figure title") + caption: str = Field(..., description="Combined multi-panel LaTeX caption") + figure_ids: list[str] = Field(..., description="IDs of figures to include in subfigure grid") + columns: int = Field(2, description="Number of columns in subfigure grid (1, 2, 3)") + subcaptions: dict[str, str] = Field(default_factory=dict, description="Subcaption per figure ID") + + +class FigureArtifact(BaseModel): + """Stored figure artifact with scripts and LaTeX snippets.""" + id: str + project_id: str + title: str + caption: str + plot_type: PlotType + style_theme: StyleTheme + palette: ColorPalette + python_script: str + latex_snippet: str + tikz_code: str + svg_preview: str + created_at: str + + def to_dict(self) -> dict[str, Any]: + return self.model_dump() diff --git a/backend/openmlr/services/model_card_generator.py b/backend/openmlr/services/model_card_generator.py new file mode 100644 index 0000000..42fa9de --- /dev/null +++ b/backend/openmlr/services/model_card_generator.py @@ -0,0 +1,192 @@ +"""Model Card and Artifact documentation generator adhering to NeurIPS/HuggingFace standards.""" + +from __future__ import annotations + +from .model_types import ModelArtifact, ModelCardContent + +# GPU TDP in Watts for carbon estimation +GPU_TDP_MAP = { + "nvidia a100": 400.0, + "nvidia h100": 700.0, + "nvidia l40s": 350.0, + "nvidia rtx 4090": 450.0, + "nvidia rtx 3090": 350.0, + "nvidia v100": 300.0, + "nvidia t4": 70.0, + "google tpu v4": 250.0, + "apple m3 max": 60.0, +} +DEFAULT_PUE = 1.2 # Datacenter Power Usage Effectiveness +CARBON_INTENSITY_KG_PER_KWH = 0.385 # Global grid average carbon intensity in kg CO2eq / kWh + + +def estimate_carbon_footprint(gpu_type: str, gpu_hours: float) -> float: + """Estimate CO2 equivalent emissions in kilograms.""" + lower_gpu = gpu_type.lower() + power_watts = 350.0 + for key, val in GPU_TDP_MAP.items(): + if key in lower_gpu: + power_watts = val + break + total_kwh = (power_watts * gpu_hours * DEFAULT_PUE) / 1000.0 + return round(total_kwh * CARBON_INTENSITY_KG_PER_KWH, 2) + + +def generate_bibtex_entry(model: ModelArtifact, author: str) -> str: + """Generate a clean BibTeX citation for the model artifact.""" + tag = model.name.lower().replace(" ", "_").replace("-", "_") + return ( + f"@misc{{{tag}_{model.version},\n" + f" title = {{{model.name} (v{model.version}): Autonomous Research Model Artifact}},\n" + f" author = {{{author}}},\n" + f" year = {{2026}},\n" + f" publisher = {{OpenMLR Autonomous Research Platform}},\n" + f" howpublished = {{\\url{{https://openmlr.local/models/{model.id}}}}}\n" + f"}}" + ) + + +def generate_latex_card(model: ModelArtifact, author: str, co2_kg: float) -> str: + """Generate a LaTeX table/section snippet for academic manuscripts.""" + metrics_str = " & ".join(f"{k} = {v:.4f}" if isinstance(v, float) else f"{k} = {v}" for k, v in model.metrics.items()) or "N/A" + return ( + "\\begin{table}[h]\n" + "\\centering\n" + "\\small\n" + "\\begin{tabular}{ll}\n" + "\\toprule\n" + "\\textbf{Model Property} & \\textbf{Specification} \\\\\n" + "\\midrule\n" + f"Model Name & {model.name} (v{model.version}) \\\\\n" + f"Author & {author} \\\\\n" + f"Architecture & {model.architecture} \\\\\n" + f"Framework & {model.framework.capitalize()} \\\\\n" + f"Task Type & {model.task_type.replace('_', ' ').capitalize()} \\\\\n" + f"Parameter Count & {model.parameters_count:,} parameters \\\\\n" + f"Artifact Size & {model.model_size_mb:.1f} MB \\\\\n" + f"Primary Metrics & {metrics_str} \\\\\n" + f"Estimated Carbon & {co2_kg:.2f} kg $\\text{{CO}}_2\\text{{eq}}$ \\\\\n" + "\\bottomrule\n" + "\\end{tabular}\n" + f"\\caption{{Model Card and Artifact Specifications for {model.name}.}}\n" + f"\\label{{tab:model_card_{model.id}}}\n" + "\\end{table}" + ) + + +def generate_markdown_card( + model: ModelArtifact, + author: str, + license_str: str, + intended_use: str, + limitations: str, + evaluation_notes: str, + gpu_type: str, + gpu_hours: float, + co2_kg: float, +) -> str: + """Generate standard Markdown model card.""" + metrics_rows = "\n".join( + f"| `{k}` | `{v:.4f}` |" if isinstance(v, float) else f"| `{k}` | `{v}` |" + for k, v in model.metrics.items() + ) or "| Metric | None recorded |" + + hparams_rows = "\n".join( + f"| `{k}` | `{v}` |" for k, v in model.hyperparameters.items() + ) or "| Hyperparameter | None specified |" + + tags_str = ", ".join(f"`{t}`" for t in model.tags) or "`research`" + + return f"""# Model Card: {model.name} (v{model.version}) + +## Model Details +- **Model Name:** {model.name} +- **Version:** {model.version} +- **Architecture:** {model.architecture} +- **Framework:** {model.framework} +- **Task Type:** {model.task_type} +- **Status:** {model.status} +- **Author / Lead:** {author} +- **License:** {license_str} +- **Tags:** {tags_str} +- **Created Date:** {model.created_at} + +## Description & Summary +{model.description or "Autonomous model artifact trained and evaluated via OpenMLR research pipeline."} + +## Intended Use +{intended_use or "Scientific research, academic benchmarking, ablation verification, and reproducible autonomous machine learning."} + +## Model Architecture & Capacity +- **Total Parameters:** {model.parameters_count:,} +- **Artifact Disk Size:** {model.model_size_mb:.2f} MB +- **Checkpoint Reference:** `{model.checkpoint_path or "N/A"}` +- **Base Model Foundation:** `{model.base_model or "Trained from scratch"}` + +## Hyperparameters & Training Configuration +| Parameter | Value | +| :--- | :--- | +{hparams_rows} + +## Quantitative Evaluation Results +| Metric Name | Value | +| :--- | :--- | +{metrics_rows} + +{evaluation_notes} + +## Limitations & Ethical Considerations +{limitations or "This model is intended strictly for research purposes. Out-of-distribution inputs may degrade performance. Verify downstream safety alignments before real-world deployment."} + +## Environmental Impact & Carbon Estimation +- **Hardware Utilized:** {gpu_type} +- **Compute Time:** {gpu_hours:.1f} GPU hours +- **Estimated Carbon Emissions:** **{co2_kg:.2f} kg CO2eq** (PUE: {DEFAULT_PUE}, Grid Intensity: {CARBON_INTENSITY_KG_PER_KWH} kg/kWh) + +## Citation +```bibtex +{generate_bibtex_entry(model, author)} +``` +""" + + +def build_model_card( + model: ModelArtifact, + author: str = "OpenMLR Research Agent", + license_str: str = "Apache-2.0", + intended_use: str = "", + limitations: str = "", + evaluation_notes: str = "", + gpu_type: str = "NVIDIA A100-SXM4-80GB", + gpu_hours: float = 24.0, +) -> ModelCardContent: + """Build complete multi-format model card artifact.""" + co2_kg = estimate_carbon_footprint(gpu_type, gpu_hours) + bibtex = generate_bibtex_entry(model, author) + latex = generate_latex_card(model, author, co2_kg) + md = generate_markdown_card( + model=model, + author=author, + license_str=license_str, + intended_use=intended_use, + limitations=limitations, + evaluation_notes=evaluation_notes, + gpu_type=gpu_type, + gpu_hours=gpu_hours, + co2_kg=co2_kg, + ) + return ModelCardContent( + model_name=model.name, + version=model.version, + markdown=md, + latex=latex, + bibtex=bibtex, + co2_emissions_kg=co2_kg, + summary={ + "parameters": model.parameters_count, + "size_mb": model.model_size_mb, + "architecture": model.architecture, + "framework": model.framework, + "co2_kg": co2_kg, + }, + ) diff --git a/backend/openmlr/services/model_registry.py b/backend/openmlr/services/model_registry.py new file mode 100644 index 0000000..6e4fe42 --- /dev/null +++ b/backend/openmlr/services/model_registry.py @@ -0,0 +1,336 @@ +"""Model Registry & Governance Service — artifact lifecycle, checkpoint inspection, quantization planning.""" + +from __future__ import annotations + +import os +import uuid +from datetime import UTC, datetime +from typing import Any + +from .model_card_generator import build_model_card +from .model_types import ( + CheckpointInspection, + GenerateModelCardRequest, + InspectCheckpointRequest, + ModelArtifact, + ModelCardContent, + QuantizationEstimate, + RegisterModelRequest, + UpdateModelRequest, +) + + +class ModelRegistryService: + """Service for managing model artifacts, checkpoint inspections, and quantization planning.""" + + _models_store: dict[str, dict[str, ModelArtifact]] = {} + + @classmethod + def _get_project_store(cls, project_id: str) -> dict[str, ModelArtifact]: + if project_id not in cls._models_store: + cls._models_store[project_id] = {} + return cls._models_store[project_id] + + @classmethod + def register_model(cls, project_id: str, request: RegisterModelRequest) -> ModelArtifact: + """Register a new model artifact in the registry.""" + store = cls._get_project_store(project_id) + model_id = f"model_{uuid.uuid4().hex[:12]}" + now = datetime.now(UTC).isoformat() + + # If parameters_count or model_size_mb is not explicitly provided, estimate from checkpoint + params = request.parameters_count + size_mb = request.model_size_mb + if params == 0 and size_mb > 0: + params = int((size_mb * 1024 * 1024) / 4) # Assume FP32 base + elif size_mb == 0 and params > 0: + size_mb = round((params * 4) / (1024 * 1024), 2) + + artifact = ModelArtifact( + id=model_id, + project_id=project_id, + name=request.name, + version=request.version, + architecture=request.architecture, + framework=request.framework, + task_type=request.task_type, + status=request.status, + created_at=now, + updated_at=now, + description=request.description, + parameters_count=params, + model_size_mb=size_mb, + checkpoint_path=request.checkpoint_path, + base_model=request.base_model, + tags=list(request.tags), + metrics=dict(request.metrics), + hyperparameters=dict(request.hyperparameters), + lineage=dict(request.lineage), + metadata=dict(request.metadata), + ) + store[model_id] = artifact + return artifact + + @classmethod + def list_models( + cls, + project_id: str, + task_type: str | None = None, + framework: str | None = None, + status: str | None = None, + tag: str | None = None, + ) -> list[ModelArtifact]: + """List all model artifacts matching optional filter criteria.""" + store = cls._get_project_store(project_id) + results = list(store.values()) + + if task_type: + results = [m for m in results if m.task_type == task_type] + if framework: + results = [m for m in results if m.framework == framework] + if status: + results = [m for m in results if m.status == status] + if tag: + results = [m for m in results if tag in m.tags] + + results.sort(key=lambda m: m.created_at, reverse=True) + return results + + @classmethod + def get_model(cls, project_id: str, model_id: str) -> ModelArtifact | None: + """Fetch a specific model artifact by id.""" + store = cls._get_project_store(project_id) + return store.get(model_id) + + @classmethod + def update_model(cls, project_id: str, model_id: str, request: UpdateModelRequest) -> ModelArtifact | None: + """Update an existing model artifact.""" + store = cls._get_project_store(project_id) + artifact = store.get(model_id) + if not artifact: + return None + + if request.name is not None: + artifact.name = request.name + if request.version is not None: + artifact.version = request.version + if request.status is not None: + artifact.status = request.status + if request.description is not None: + artifact.description = request.description + if request.parameters_count is not None: + artifact.parameters_count = request.parameters_count + if request.model_size_mb is not None: + artifact.model_size_mb = request.model_size_mb + if request.checkpoint_path is not None: + artifact.checkpoint_path = request.checkpoint_path + if request.tags is not None: + artifact.tags = list(request.tags) + if request.metrics is not None: + artifact.metrics.update(request.metrics) + if request.hyperparameters is not None: + artifact.hyperparameters.update(request.hyperparameters) + if request.metadata is not None: + artifact.metadata.update(request.metadata) + + artifact.updated_at = datetime.now(UTC).isoformat() + return artifact + + @classmethod + def delete_model(cls, project_id: str, model_id: str) -> bool: + """Delete a model artifact from the registry.""" + store = cls._get_project_store(project_id) + if model_id in store: + del store[model_id] + return True + return False + + @classmethod + def inspect_checkpoint(cls, req: InspectCheckpointRequest) -> CheckpointInspection: + """Inspect checkpoint metadata, calculate parameter distribution, and estimate VRAM requirements.""" + total_params = req.parameters_count + size_mb = req.model_size_mb + path = req.checkpoint_path + fmt = "pytorch (.pt/.pth)" + + if path.endswith(".safetensors"): + fmt = "safetensors" + elif path.endswith(".onnx"): + fmt = "onnx" + elif path.endswith(".gguf"): + fmt = "gguf" + elif path.endswith(".bin"): + fmt = "pytorch_bin" + + # Check real file if accessible + if path and os.path.exists(path) and size_mb == 0: + size_mb = round(os.path.getsize(path) / (1024 * 1024), 2) + + if total_params == 0 and size_mb > 0: + total_params = int((size_mb * 1024 * 1024) / 4) + elif total_params > 0 and size_mb == 0: + size_mb = round((total_params * 4) / (1024 * 1024), 2) + elif total_params == 0 and size_mb == 0: + total_params = 125_000_000 # Default 125M parameter reference + size_mb = 500.0 + + # Memory footprint calculations (weights + 20% runtime activation overhead) + vram_fp32 = round((total_params * 4 * 1.2) / (1024 * 1024), 1) + vram_fp16 = round((total_params * 2 * 1.2) / (1024 * 1024), 1) + vram_int8 = round((total_params * 1 * 1.2) / (1024 * 1024), 1) + vram_int4 = round((total_params * 0.5 * 1.2) / (1024 * 1024), 1) + + dtype_breakdown = { + "torch.float32": int(total_params * 0.95), + "torch.int64": int(total_params * 0.05), + } + + top_layers = req.layer_samples or [ + {"name": "transformer.encoder.layers.0.self_attn.q_proj.weight", "params": int(total_params * 0.05), "dtype": "float32"}, + {"name": "transformer.encoder.layers.0.mlp.gate_proj.weight", "params": int(total_params * 0.12), "dtype": "float32"}, + {"name": "transformer.output_projection.weight", "params": int(total_params * 0.08), "dtype": "float32"}, + ] + + return CheckpointInspection( + file_format=fmt, + total_parameters=total_params, + trainable_parameters=total_params, + total_size_mb=size_mb, + estimated_vram_fp32_mb=vram_fp32, + estimated_vram_fp16_mb=vram_fp16, + estimated_vram_int8_mb=vram_int8, + estimated_vram_int4_mb=vram_int4, + dtype_breakdown=dtype_breakdown, + layers_count=len(top_layers) + 24, + top_layers=top_layers, + has_optimizer_state=False, + metadata={"framework": req.framework, "path": path}, + ) + + @classmethod + def plan_quantization(cls, model: ModelArtifact, target_precisions: list[str]) -> list[QuantizationEstimate]: + """Generate precision quantization trade-off estimates for model deployment.""" + params = model.parameters_count if model.parameters_count > 0 else 125_000_000 + fp32_size_mb = (params * 4) / (1024 * 1024) + estimates: list[QuantizationEstimate] = [] + + precision_specs = { + "fp16": { + "bytes_per_param": 2.0, + "speedup": 1.7, + "engine": "vLLM / HuggingFace Transformers (native half-precision)", + "loss": "Negligible (<0.1% accuracy drop)", + }, + "bf16": { + "bytes_per_param": 2.0, + "speedup": 1.7, + "engine": "FlashAttention-2 / PyTorch AMP", + "loss": "Negligible (<0.05% accuracy drop, higher dynamic range)", + }, + "fp8": { + "bytes_per_param": 1.0, + "speedup": 2.4, + "engine": "TensorRT-LLM / vLLM FP8 (Ada/Hopper GPUs)", + "loss": "Minimal (<0.5% accuracy drop)", + }, + "int8": { + "bytes_per_param": 1.0, + "speedup": 2.1, + "engine": "bitsandbytes LLM.int8() / SmoothQuant", + "loss": "Low (<1.0% accuracy drop)", + }, + "int4": { + "bytes_per_param": 0.55, # includes group scale & zero-point overhead + "speedup": 3.2, + "engine": "AutoAWQ / GPTQ / llama.cpp GGUF Q4_K_M", + "loss": "Moderate (1.5-2.5% perplexity delta, 4x memory savings)", + }, + } + + for prec in target_precisions: + norm_prec = prec.lower() + spec = precision_specs.get(norm_prec, { + "bytes_per_param": 2.0, + "speedup": 1.5, + "engine": "Custom Quantizer", + "loss": "Variable", + }) + size_mb = round((params * spec["bytes_per_param"]) / (1024 * 1024), 2) + vram_mb = round(size_mb * 1.2, 2) + comp_ratio = round(fp32_size_mb / max(size_mb, 0.001), 2) + + estimates.append( + QuantizationEstimate( + target_precision=norm_prec.upper(), + estimated_size_mb=size_mb, + estimated_vram_mb=vram_mb, + compression_ratio=comp_ratio, + expected_latency_speedup=spec["speedup"], + suggested_engine=spec["engine"], + loss_tolerance_level=spec["loss"], + ) + ) + + return estimates + + @classmethod + def generate_model_card( + cls, + project_id: str, + model_id: str, + req: GenerateModelCardRequest, + ) -> ModelCardContent | None: + """Generate a complete multi-format model card.""" + artifact = cls.get_model(project_id, model_id) + if not artifact: + return None + + return build_model_card( + model=artifact, + author=req.author, + license_str=req.license, + intended_use=req.intended_use, + limitations=req.limitations, + evaluation_notes=req.evaluation_notes, + gpu_type=req.gpu_type, + gpu_hours=req.gpu_hours, + ) + + @classmethod + def compare_models(cls, project_id: str, model_ids: list[str]) -> dict[str, Any]: + """Compare multiple model artifacts side-by-side.""" + store = cls._get_project_store(project_id) + models = [store[mid] for mid in model_ids if mid in store] + if len(models) < 2: + return {"error": "At least 2 valid models are required for comparison"} + + # Collect all metric keys + all_metrics: set[str] = set() + for m in models: + all_metrics.update(m.metrics.keys()) + + metric_matrix: dict[str, dict[str, float | None]] = {} + for metric_name in sorted(all_metrics): + metric_matrix[metric_name] = {m.id: m.metrics.get(metric_name) for m in models} + + # Find best model on primary metrics (val_loss minimum or accuracy/f1 maximum) + best_candidate = models[0] + for m in models[1:]: + if "accuracy" in m.metrics and "accuracy" in best_candidate.metrics: + if m.metrics["accuracy"] > best_candidate.metrics["accuracy"]: + best_candidate = m + elif "val_loss" in m.metrics and "val_loss" in best_candidate.metrics: + if m.metrics["val_loss"] < best_candidate.metrics["val_loss"]: + best_candidate = m + + return { + "compared_models": [m.to_dict() for m in models], + "metric_matrix": metric_matrix, + "parameter_comparison": {m.id: m.parameters_count for m in models}, + "size_comparison_mb": {m.id: m.model_size_mb for m in models}, + "recommended_model_id": best_candidate.id, + "recommendation_reason": ( + f"Model '{best_candidate.name}' (v{best_candidate.version}) demonstrates the strongest " + f"empirical performance across primary benchmark metrics with {best_candidate.parameters_count:,} parameters." + ), + } diff --git a/backend/openmlr/services/model_types.py b/backend/openmlr/services/model_types.py new file mode 100644 index 0000000..f5cc36c --- /dev/null +++ b/backend/openmlr/services/model_types.py @@ -0,0 +1,181 @@ +"""Model registry and governance domain types.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal + +from pydantic import BaseModel, Field + +FrameworkType = Literal["pytorch", "safetensors", "jax", "onnx", "gguf", "huggingface", "tensorrt"] +TaskType = Literal[ + "causal_lm", + "seq2seq", + "classification", + "object_detection", + "segmentation", + "diffusion", + "embedding", + "reinforcement_learning", + "custom", +] +ModelStatus = Literal["draft", "training", "evaluated", "production", "archived"] +PrecisionType = Literal["fp32", "fp16", "bf16", "int8", "int4", "fp8", "mixed"] + + +@dataclass +class LayerSummary: + name: str + param_count: int + dtype: str + trainable: bool = True + shape: list[int] = field(default_factory=list) + + +@dataclass +class CheckpointInspection: + file_format: str + total_parameters: int + trainable_parameters: int + total_size_mb: float + estimated_vram_fp32_mb: float + estimated_vram_fp16_mb: float + estimated_vram_int8_mb: float + estimated_vram_int4_mb: float + dtype_breakdown: dict[str, int] = field(default_factory=dict) + layers_count: int = 0 + top_layers: list[dict[str, Any]] = field(default_factory=list) + has_optimizer_state: bool = False + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class QuantizationEstimate: + target_precision: str + estimated_size_mb: float + estimated_vram_mb: float + compression_ratio: float + expected_latency_speedup: float + suggested_engine: str + loss_tolerance_level: str + + +@dataclass +class ModelCardContent: + model_name: str + version: str + markdown: str + latex: str + bibtex: str + co2_emissions_kg: float + summary: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class ModelArtifact: + id: str + project_id: str + name: str + version: str + architecture: str + framework: FrameworkType + task_type: TaskType + status: ModelStatus + created_at: str + updated_at: str + description: str = "" + parameters_count: int = 0 + model_size_mb: float = 0.0 + checkpoint_path: str = "" + base_model: str = "" + tags: list[str] = field(default_factory=list) + metrics: dict[str, float] = field(default_factory=dict) + hyperparameters: dict[str, Any] = field(default_factory=dict) + lineage: dict[str, Any] = field(default_factory=dict) + metadata: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "project_id": self.project_id, + "name": self.name, + "version": self.version, + "architecture": self.architecture, + "framework": self.framework, + "task_type": self.task_type, + "status": self.status, + "created_at": self.created_at, + "updated_at": self.updated_at, + "description": self.description, + "parameters_count": self.parameters_count, + "model_size_mb": self.model_size_mb, + "checkpoint_path": self.checkpoint_path, + "base_model": self.base_model, + "tags": list(self.tags), + "metrics": dict(self.metrics), + "hyperparameters": dict(self.hyperparameters), + "lineage": dict(self.lineage), + "metadata": dict(self.metadata), + } + + +# Pydantic Schemas for API Requests & Responses + +class RegisterModelRequest(BaseModel): + name: str + version: str = "1.0.0" + architecture: str = "Transformer" + framework: FrameworkType = "pytorch" + task_type: TaskType = "causal_lm" + status: ModelStatus = "evaluated" + description: str = "" + parameters_count: int = 0 + model_size_mb: float = 0.0 + checkpoint_path: str = "" + base_model: str = "" + tags: list[str] = Field(default_factory=list) + metrics: dict[str, float] = Field(default_factory=dict) + hyperparameters: dict[str, Any] = Field(default_factory=dict) + lineage: dict[str, Any] = Field(default_factory=dict) + metadata: dict[str, Any] = Field(default_factory=dict) + + +class UpdateModelRequest(BaseModel): + name: str | None = None + version: str | None = None + status: ModelStatus | None = None + description: str | None = None + parameters_count: int | None = None + model_size_mb: float | None = None + checkpoint_path: str | None = None + tags: list[str] | None = None + metrics: dict[str, float] | None = None + hyperparameters: dict[str, Any] | None = None + metadata: dict[str, Any] | None = None + + +class InspectCheckpointRequest(BaseModel): + checkpoint_path: str = Field("", max_length=500) + parameters_count: int = Field(0, ge=0) + model_size_mb: float = Field(0.0, ge=0.0) + framework: str = "pytorch" + layer_samples: list[dict[str, Any]] = Field(default_factory=list) + + +class GenerateModelCardRequest(BaseModel): + include_carbon_estimate: bool = True + author: str = "OpenMLR Research Agent" + license: str = "Apache-2.0" + intended_use: str = "" + limitations: str = "" + evaluation_notes: str = "" + gpu_hours: float = Field(24.0, ge=0.0) + gpu_type: str = "NVIDIA A100-SXM4-80GB" + + +class PlanQuantizationRequest(BaseModel): + target_precisions: list[str] = Field(default_factory=lambda: ["fp16", "int8", "int4", "fp8"]) + + +class CompareModelsRequest(BaseModel): + model_ids: list[str] = Field(..., min_length=2, max_length=10) diff --git a/backend/openmlr/services/reproducibility_auditor.py b/backend/openmlr/services/reproducibility_auditor.py new file mode 100644 index 0000000..2404345 --- /dev/null +++ b/backend/openmlr/services/reproducibility_auditor.py @@ -0,0 +1,427 @@ +"""Reproducibility Auditor service for ML research artifacts, determinism, and conference checklists.""" + +from __future__ import annotations + +import logging +import os +import re +import uuid +from datetime import UTC, datetime + +from .reproducibility_templates import ( + generate_badge_markdown, + generate_badge_svg, + generate_latex_appendix, +) +from .reproducibility_types import ( + AuditCodebaseRequest, + CategoryScore, + CheckCategory, + CheckItem, + CheckSeverity, + CheckStatus, + GenerateAppendixRequest, + GenerateDockerfileRequest, + ReproducibilityAuditReport, +) + +logger = logging.getLogger("openmlr.services.reproducibility_auditor") + + +class ReproducibilityAuditorService: + """Audits ML codebases for determinism, environment pinning, hardware, and conference reproducibility.""" + + _reports_store: dict[str, dict[str, ReproducibilityAuditReport]] = {} + + @classmethod + def _get_project_store(cls, project_id: str | None) -> dict[str, ReproducibilityAuditReport]: + pid = project_id or "default" + if pid not in cls._reports_store: + cls._reports_store[pid] = {} + return cls._reports_store[pid] + + @classmethod + def list_reports(cls, project_id: str | None = None) -> list[ReproducibilityAuditReport]: + store = cls._get_project_store(project_id) + return sorted(store.values(), key=lambda r: r.created_at, reverse=True) + + @classmethod + def get_report(cls, report_id: str, project_id: str | None = None) -> ReproducibilityAuditReport | None: + store = cls._get_project_store(project_id) + return store.get(report_id) + + @classmethod + def delete_report(cls, report_id: str, project_id: str | None = None) -> bool: + store = cls._get_project_store(project_id) + if report_id in store: + del store[report_id] + return True + return False + + @classmethod + def generate_determinism_snippet(cls, framework: str = "pytorch", seed: int = 42, strict_mode: bool = True) -> str: + """Generate boilerplate Python snippet for 100% deterministic experiment execution.""" + fw = framework.lower() + if fw == "jax": + return ( + f"# JAX Deterministic Random State\n" + f"import jax\n" + f"import jax.numpy as jnp\n" + f"import numpy as np\n" + f"import os\n" + f"import random\n\n" + f"os.environ['PYTHONHASHSEED'] = str({seed})\n" + f"random.seed({seed})\n" + f"np.random.seed({seed})\n" + f"rng_key = jax.random.PRNGKey({seed})\n" + ) + elif fw in ("tensorflow", "tf"): + return ( + f"# TensorFlow Deterministic Setup\n" + f"import os\n" + f"import random\n" + f"import numpy as np\n" + f"import tensorflow as tf\n\n" + f"os.environ['PYTHONHASHSEED'] = str({seed})\n" + f"os.environ['TF_DETERMINISTIC_OPS'] = '1'\n" + f"random.seed({seed})\n" + f"np.random.seed({seed})\n" + f"tf.random.set_seed({seed})\n" + ) + + strict_call = "torch.use_deterministic_algorithms(True)" if strict_mode else "# torch.use_deterministic_algorithms(True)" + return ( + f"# PyTorch Reproducibility & Determinism Boilerplate\n" + f"import os\n" + f"import random\n" + f"import numpy as np\n" + f"import torch\n\n" + f"def set_seed(seed: int = {seed}) -> None:\n" + f" os.environ['PYTHONHASHSEED'] = str(seed)\n" + f" os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'\n" + f" random.seed(seed)\n" + f" np.random.seed(seed)\n" + f" torch.manual_seed(seed)\n" + f" torch.cuda.manual_seed_all(seed)\n" + f" torch.backends.cudnn.deterministic = True\n" + f" torch.backends.cudnn.benchmark = False\n" + f" {strict_call}\n\n" + f"set_seed({seed})\n" + ) + + @classmethod + def generate_dockerfile(cls, req: GenerateDockerfileRequest) -> str: + """Generate a production-ready, reproducible Docker container definition.""" + pkgs = "\n".join([f" {p} \\" for p in req.requirements]) if req.requirements else " torch torchvision --index-url https://download.pytorch.org/whl/cu121 \\" + return ( + f"# Generated Reproducible ML Research Container\n" + f"FROM nvidia/cuda:{req.cuda_version}-runtime-ubuntu22.04\n\n" + f"ENV DEBIAN_FRONTEND=noninteractive \\\n" + f" PYTHONUNBUFFERED=1 \\\n" + f" PYTHONHASHSEED=0 \\\n" + f" CUBLAS_WORKSPACE_CONFIG=:4096:8\n\n" + f"RUN apt-get update && apt-get install -y --no-install-recommends \\\n" + f" python{req.python_version} \\\n" + f" python3-pip \\\n" + f" git \\\n" + f" curl \\\n" + f" && rm -rf /var/lib/apt/lists/*\n\n" + f"WORKDIR /workspace\n" + f"COPY . /workspace\n\n" + f"RUN pip install --no-cache-dir --upgrade pip && \\\n" + f" pip install --no-cache-dir \\\n" + f"{pkgs}\n" + f" pyyaml\n\n" + f"CMD [\"{req.entrypoint_cmd}\"]\n" + ) + + @classmethod + def generate_conda_env(cls, env_name: str = "openmlr-reproduce", python_version: str = "3.11", dependencies: list[str] | None = None) -> str: + """Generate an environment.yml for Conda reproducibility.""" + deps = dependencies or ["pytorch", "torchvision", "pytorch-cuda=12.1", "numpy", "scipy", "pyyaml"] + dep_lines = "\n".join([f" - {d}" for d in deps]) + return ( + f"name: {env_name}\n" + f"channels:\n" + f" - pytorch\n" + f" - nvidia\n" + f" - conda-forge\n" + f"dependencies:\n" + f" - python={python_version}\n" + f"{dep_lines}\n" + f" - pip:\n" + f" - openmlr\n" + ) + + @classmethod + def generate_badge_markdown(cls, score: float, grade: str) -> str: + return generate_badge_markdown(score, grade) + + @classmethod + def generate_badge_svg(cls, score: float, grade: str) -> str: + return generate_badge_svg(score, grade) + + @classmethod + def generate_latex_appendix(cls, req: GenerateAppendixRequest, report: ReproducibilityAuditReport | None = None) -> str: + return generate_latex_appendix(req, report) + + @classmethod + def audit_codebase( + cls, + request: AuditCodebaseRequest, + project_id: str | None = None, + ) -> ReproducibilityAuditReport: + """Perform comprehensive static analysis and reproducibility auditing on target code.""" + code_map = request.code_snippets or {} + target_path = request.target_path + + if not code_map and os.path.exists(target_path): + try: + for root, _, files in os.walk(target_path): + for file in files: + if file.endswith((".py", ".txt", ".toml", ".yaml", ".yml", ".json", ".md", "Dockerfile")): + full_path = os.path.join(root, file) + rel_path = os.path.relpath(full_path, target_path) + if len(code_map) < 25: + try: + with open(full_path, encoding="utf-8", errors="ignore") as f: + code_map[rel_path] = f.read(50000) + except Exception: + pass + except Exception as e: + logger.warning("Error reading files in %s: %s", target_path, e) + + combined_text = "\n".join(code_map.values()) + filenames = list(code_map.keys()) + + # 1. Determinism Audit + seeds_found: dict[str, int | str] = {} + checklist: list[CheckItem] = [] + + has_torch_seed = bool(re.search(r"torch\.manual_seed\s*\(\s*([^)]+)\s*\)", combined_text)) + has_np_seed = bool(re.search(r"np\.random\.seed\s*\(\s*([^)]+)\s*\)", combined_text)) + has_random_seed = bool(re.search(r"random\.seed\s*\(\s*([^)]+)\s*\)", combined_text)) + has_cudnn_det = "cudnn.deterministic" in combined_text + has_det_algo = "use_deterministic_algorithms" in combined_text or "CUBLAS_WORKSPACE_CONFIG" in combined_text + + seed_match = re.search(r"(?:manual_seed|seed)\s*\(\s*([^)]+)\s*\)", combined_text) + if seed_match: + val_str = seed_match.group(1).strip() + if val_str.isdigit(): + seeds_found["main_seed"] = int(val_str) + else: + seeds_found["main_seed"] = val_str + + if has_torch_seed or has_np_seed or has_random_seed: + checklist.append( + CheckItem( + id="det_seed_init", + category=CheckCategory.DETERMINISM, + title="Random Seed Initialization", + description="RNG generator is explicitly initialized with deterministic seeds.", + status=CheckStatus.PASS, + severity=CheckSeverity.CRITICAL, + details=f"Found explicit random seed calls ({', '.join(str(v) for v in seeds_found.values()) or 'present'}).", + ) + ) + else: + checklist.append( + CheckItem( + id="det_seed_init", + category=CheckCategory.DETERMINISM, + title="Random Seed Initialization", + description="RNG generator is explicitly initialized with deterministic seeds.", + status=CheckStatus.FAIL, + severity=CheckSeverity.CRITICAL, + details="No explicit torch.manual_seed(), np.random.seed(), or random.seed() call detected.", + remediation="Add `set_seed(42)` at the beginning of training scripts.", + ) + ) + + checklist.append( + CheckItem( + id="det_cudnn", + category=CheckCategory.DETERMINISM, + title="cuDNN Determinism Flags", + description="torch.backends.cudnn.deterministic is configured.", + status=CheckStatus.PASS if has_cudnn_det else CheckStatus.WARN, + severity=CheckSeverity.HIGH, + details="torch.backends.cudnn.deterministic is enabled." if has_cudnn_det else "cuDNN algorithm selection can lead to variance.", + remediation="Set `torch.backends.cudnn.deterministic = True` and `torch.backends.cudnn.benchmark = False`." if not has_cudnn_det else "", + ) + ) + + # 2. Environment & Dependencies Audit + has_req_file = any("requirements" in fn or "pyproject" in fn or "environment.yml" in fn or "Pipfile" in fn for fn in filenames) + has_pinned_deps = bool(re.search(r"[a-zA-Z0-9_\-]+==\d+\.\d+", combined_text)) + has_dockerfile = any("Dockerfile" in fn or "docker-compose" in fn for fn in filenames) + + checklist.append( + CheckItem( + id="env_deps_manifest", + category=CheckCategory.ENVIRONMENT, + title="Dependency Manifest File", + description="Project includes requirements.txt, pyproject.toml, or environment.yml.", + status=CheckStatus.PASS if has_req_file else CheckStatus.FAIL, + severity=CheckSeverity.CRITICAL, + details="Found dependency configuration file." if has_req_file else "No requirements.txt or pyproject.toml detected.", + remediation="Generate locked dependencies using `pip freeze > requirements.txt` or `uv export`." if not has_req_file else "", + ) + ) + + checklist.append( + CheckItem( + id="env_pinned_versions", + category=CheckCategory.ENVIRONMENT, + title="Exact Package Pinning", + description="Package versions are strictly pinned with `==` to avoid breaking changes.", + status=CheckStatus.PASS if has_pinned_deps else CheckStatus.WARN, + severity=CheckSeverity.HIGH, + details="Found pinned package dependencies (e.g. torch==2.x)." if has_pinned_deps else "Dependencies appear unpinned.", + remediation="Pin exact versions for all dependencies using `==`." if not has_pinned_deps else "", + ) + ) + + checklist.append( + CheckItem( + id="env_container_recipe", + category=CheckCategory.ENVIRONMENT, + title="Containerization / Docker Recipe", + description="Dockerfile is provided for isolated reproduction.", + status=CheckStatus.PASS if has_dockerfile else CheckStatus.WARN, + severity=CheckSeverity.MEDIUM, + details="Found Dockerfile in workspace." if has_dockerfile else "No container recipe detected.", + remediation="Use OpenMLR's generate_dockerfile action to build a reproducible container image." if not has_dockerfile else "", + ) + ) + + # 3. Hardware & Compute + has_cuda_check = "cuda.is_available()" in combined_text or "cuda:0" in combined_text or "to(device)" in combined_text + checklist.append( + CheckItem( + id="hw_device_agnostic", + category=CheckCategory.HARDWARE, + title="Device Selection & Hardware Portability", + description="Code dynamically identifies compute device (CUDA / MPS / CPU).", + status=CheckStatus.PASS if has_cuda_check else CheckStatus.WARN, + severity=CheckSeverity.HIGH, + details="Dynamic device checking identified." if has_cuda_check else "No dynamic device selection found.", + remediation="Use `device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')`." if not has_cuda_check else "", + ) + ) + + # 4. Dataset & Splits + has_splits = any(k in combined_text.lower() for k in ["train_test_split", "train_split", "val_dataset", "test_loader", "random_split"]) + checklist.append( + CheckItem( + id="data_splits_spec", + category=CheckCategory.DATASET, + title="Evaluation Partition Separation", + description="Training, validation, and test datasets are partitioned without leakage.", + status=CheckStatus.PASS if has_splits else CheckStatus.WARN, + severity=CheckSeverity.CRITICAL, + details="Train/val/test split logic present." if has_splits else "Explicit partition splitting not detected.", + remediation="Ensure train, validation, and test splits are strictly separated." if not has_splits else "", + ) + ) + + # 5. Hyperparameters & Logging + has_hparams = any(k in combined_text.lower() for k in ["argparse", "hydra", "click", "learning_rate", "batch_size", "config.yaml"]) + checklist.append( + CheckItem( + id="hp_config_logging", + category=CheckCategory.HYPERPARAMETERS, + title="Hyperparameter Specification", + description="Training hyperparameters are configurable via CLI args or config files.", + status=CheckStatus.PASS if has_hparams else CheckStatus.WARN, + severity=CheckSeverity.HIGH, + details="Hyperparameter configuration pattern found." if has_hparams else "Hyperparameters might be hardcoded.", + remediation="Expose learning rate, batch size, and epoch count via argparse or config YAML." if not has_hparams else "", + ) + ) + + # 6. Checkpoints & Model Artifacts + has_checkpointing = "torch.save" in combined_text or "save_pretrained" in combined_text or "checkpoint" in combined_text.lower() + checklist.append( + CheckItem( + id="ckpt_state_save", + category=CheckCategory.CHECKPOINTS, + title="Model & Optimizer State Checkpointing", + description="Saves model weights and training state for post-training validation.", + status=CheckStatus.PASS if has_checkpointing else CheckStatus.WARN, + severity=CheckSeverity.HIGH, + details="Checkpoint persistence logic detected." if has_checkpointing else "No checkpoint saving calls found.", + remediation="Save model and optimizer state dictionaries periodically using `torch.save()`." if not has_checkpointing else "", + ) + ) + + # Calculate category scores + categories_dict: dict[CheckCategory, list[CheckItem]] = {} + for item in checklist: + categories_dict.setdefault(item.category, []).append(item) + + categories_scores: list[CategoryScore] = [] + total_score_sum = 0.0 + + for cat in CheckCategory: + items = categories_dict.get(cat, []) + if not items: + categories_scores.append(CategoryScore(category=cat, score=100.0, passed_checks=1, total_checks=1, status=CheckStatus.PASS)) + total_score_sum += 100.0 + continue + + passed = sum(1 for i in items if i.status == CheckStatus.PASS) + warns = sum(1 for i in items if i.status == CheckStatus.WARN) + score = max(0.0, min(100.0, ((passed * 1.0) + (warns * 0.6)) / len(items) * 100.0)) + status = CheckStatus.PASS if score >= 80 else CheckStatus.WARN if score >= 60 else CheckStatus.FAIL + categories_scores.append( + CategoryScore( + category=cat, + score=round(score, 1), + passed_checks=passed, + total_checks=len(items), + status=status, + ) + ) + total_score_sum += score + + overall_score = round(total_score_sum / len(CheckCategory), 1) + grade = "A+" if overall_score >= 95 else "A" if overall_score >= 85 else "B" if overall_score >= 70 else "C" if overall_score >= 50 else "F" + + detected_frameworks = [] + if "torch" in combined_text: + detected_frameworks.append("PyTorch") + if "jax" in combined_text: + detected_frameworks.append("JAX") + if "tensorflow" in combined_text or "keras" in combined_text: + detected_frameworks.append("TensorFlow") + if "transformers" in combined_text or "datasets" in combined_text: + detected_frameworks.append("HuggingFace") + if not detected_frameworks: + detected_frameworks.append(request.framework_hint or "Python ML") + + report_id = f"rep_{uuid.uuid4().hex[:12]}" + now = datetime.now(UTC).isoformat() + + report = ReproducibilityAuditReport( + id=report_id, + project_id=project_id, + created_at=now, + overall_score=overall_score, + grade=grade, + venue=request.venue, + categories=categories_scores, + checklist=checklist, + detected_frameworks=detected_frameworks, + seeds_detected=seeds_found, + cuda_requirements={"cuda_version": "12.1", "memory_mb": 8192, "deterministic_required": has_det_algo}, + dockerfile_recipe=cls.generate_dockerfile(GenerateDockerfileRequest()), + conda_recipe=cls.generate_conda_env(), + badge_markdown=cls.generate_badge_markdown(overall_score, grade), + badge_svg=cls.generate_badge_svg(overall_score, grade), + ) + report.latex_appendix = cls.generate_latex_appendix(GenerateAppendixRequest(report_id=report_id), report) + + store = cls._get_project_store(project_id) + store[report_id] = report + return report diff --git a/backend/openmlr/services/reproducibility_templates.py b/backend/openmlr/services/reproducibility_templates.py new file mode 100644 index 0000000..fc320a2 --- /dev/null +++ b/backend/openmlr/services/reproducibility_templates.py @@ -0,0 +1,49 @@ +"""Template and snippet generators for Reproducibility Artifacts and Badges.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .reproducibility_types import GenerateAppendixRequest, ReproducibilityAuditReport + + +def generate_badge_markdown(score: float, grade: str) -> str: + color = "brightgreen" if score >= 90 else "green" if score >= 80 else "yellow" if score >= 70 else "red" + return f"[![OpenMLR Reproducibility](https://img.shields.io/badge/reproducibility-{grade}%20({score:.0f}%25)-{color}.svg)](#reproducibility)" + + +def generate_badge_svg(score: float, grade: str) -> str: + bg_color = "#10b981" if score >= 85 else "#f59e0b" if score >= 70 else "#ef4444" + return ( + f'\n' + f' \n' + f' \n' + f' reproducibility\n' + f' {grade} ({score:.0f}%)\n' + f'' + ) + + +def generate_latex_appendix(req: GenerateAppendixRequest, report: ReproducibilityAuditReport | None = None) -> str: + """Generate LaTeX Reproducibility Statement section adhering to conference guidelines.""" + seeds_str = ", ".join(str(s) for s in req.random_seeds) + score_val = f"{report.overall_score:.0f}" if report else "95" + grade_val = report.grade if report else "A+" + return ( + f"\\section{{Reproducibility Statement}}\n" + f"\\label{{sec:reproducibility}}\n\n" + f"To ensure full scientific reproducibility, this work adheres strictly to conference reproducibility guidelines. " + f"The codebase achieves an automated reproducibility index of {score_val}/100 (Grade {grade_val}).\n\n" + f"\\subsection{{Hardware and Execution Environment}}\n" + f"All experiments were conducted on: {req.hardware_specs}. " + f"CUDA deterministic flags (\\texttt{{CUBLAS\\_WORKSPACE\\_CONFIG=:4096:8}}) and seed initializations were strictly applied across all runs.\n\n" + f"\\subsection{{Random Seeds and Determinism}}\n" + f"Evaluations were repeated over {len(req.random_seeds)} distinct random seeds: \\{{{seeds_str}\\}}. " + f"Both PyTorch and NumPy random generators were explicitly seeded, and standard deviations are reported across all benchmark tables.\n\n" + f"\\subsection{{Dataset & Code Availability}}\n" + f"The benchmark datasets used in our evaluations are publicly accessible at \\url{{{req.dataset_url}}}. " + f"Full source code, hyperparameter configuration files, and standalone Docker recipes are available at \\url{{{req.code_url}}}.\n\n" + f"\\subsection{{Hyperparameters & Optimization}}\n" + f"All learning rates, optimizer states, batch sizes, and learning rate schedules are documented in the main experimental tables and supplementary YAML config files.\n" + ) diff --git a/backend/openmlr/services/reproducibility_types.py b/backend/openmlr/services/reproducibility_types.py new file mode 100644 index 0000000..45d1cfb --- /dev/null +++ b/backend/openmlr/services/reproducibility_types.py @@ -0,0 +1,108 @@ +"""Data models and types for the Reproducibility Auditor & Artifact Governance Service.""" + +from __future__ import annotations + +from enum import Enum +from typing import Any + +from pydantic import BaseModel, Field + + +class ChecklistVenue(str, Enum): + NEURIPS = "neurips" + ICML = "icml" + ICLR = "iclr" + CVPR = "cvpr" + GENERAL = "general" + + +class CheckStatus(str, Enum): + PASS = "pass" + WARN = "warn" + FAIL = "fail" + SKIP = "skip" + + +class CheckSeverity(str, Enum): + CRITICAL = "critical" + HIGH = "high" + MEDIUM = "medium" + LOW = "low" + + +class CheckCategory(str, Enum): + DETERMINISM = "determinism" + ENVIRONMENT = "environment" + HARDWARE = "hardware" + DATASET = "dataset" + HYPERPARAMETERS = "hyperparameters" + CHECKPOINTS = "checkpoints" + + +class CheckItem(BaseModel): + id: str = Field(..., description="Unique check identifier") + category: CheckCategory = Field(..., description="Audit category") + title: str = Field(..., description="Brief title of check") + description: str = Field(..., description="Detailed description") + status: CheckStatus = Field(..., description="Status result") + severity: CheckSeverity = Field(default=CheckSeverity.MEDIUM, description="Severity if failing") + details: str = Field(default="", description="Findings and context") + remediation: str = Field(default="", description="Suggested fix or code snippet") + + +class CategoryScore(BaseModel): + category: CheckCategory = Field(..., description="Category") + score: float = Field(..., description="Score 0 to 100") + passed_checks: int = Field(..., description="Number of passing checks") + total_checks: int = Field(..., description="Total checks evaluated") + status: CheckStatus = Field(default=CheckStatus.PASS, description="Overall category status") + + +class ReproducibilityAuditReport(BaseModel): + id: str = Field(..., description="Unique audit report identifier") + project_id: str | None = Field(default=None, description="Associated project ID") + created_at: str = Field(..., description="ISO timestamp") + overall_score: float = Field(..., description="Overall reproducibility score (0-100)") + grade: str = Field(..., description="Letter grade (A+, A, B, C, F)") + venue: ChecklistVenue = Field(default=ChecklistVenue.NEURIPS, description="Evaluation rubric venue") + categories: list[CategoryScore] = Field(default_factory=list, description="Category scores") + checklist: list[CheckItem] = Field(default_factory=list, description="Detailed checklist items") + detected_frameworks: list[str] = Field(default_factory=list, description="Detected ML frameworks") + seeds_detected: dict[str, int | str] = Field(default_factory=dict, description="Discovered seeds") + cuda_requirements: dict[str, Any] = Field(default_factory=dict, description="Hardware & CUDA requirements") + dockerfile_recipe: str = Field(default="", description="Reproducible Dockerfile") + conda_recipe: str = Field(default="", description="Reproducible environment.yml") + latex_appendix: str = Field(default="", description="LaTeX Reproducibility Statement") + badge_markdown: str = Field(default="", description="Markdown badge string") + badge_svg: str = Field(default="", description="SVG badge markup") + + +class AuditCodebaseRequest(BaseModel): + target_path: str = Field(default=".", description="Target workspace or directory path to audit") + venue: ChecklistVenue = Field(default=ChecklistVenue.NEURIPS, description="Conference standard to evaluate") + framework_hint: str | None = Field(default=None, description="Optional framework hint (pytorch, jax, etc.)") + code_snippets: dict[str, str] | None = Field(default=None, description="Optional in-memory code snippets") + + +class GenerateDockerfileRequest(BaseModel): + framework: str = Field(default="pytorch", description="ML framework: pytorch, jax, tensorflow, scikit-learn") + cuda_version: str = Field(default="12.1.0", description="CUDA base image version") + python_version: str = Field(default="3.11", description="Python runtime version") + entrypoint_cmd: str = Field(default="python train.py", description="Main training entrypoint") + requirements: list[str] = Field(default_factory=list, description="Pinned Python packages") + + +class GenerateAppendixRequest(BaseModel): + report_id: str | None = Field(default=None, description="Optional report ID to base appendix on") + paper_title: str = Field(default="Reproducible Machine Learning Study", description="Paper title") + authors: str = Field(default="Autonomous Research Agent", description="Authors") + hardware_specs: str = Field(default="NVIDIA A100-SXM4-80GB (1 GPU), 8 CPU cores, 64GB RAM", description="Hardware") + random_seeds: list[int] = Field(default_factory=lambda: [42, 1337, 2026], description="Seeds tested") + dataset_url: str = Field(default="https://huggingface.co/datasets/...", description="Dataset repository") + code_url: str = Field(default="https://github.com/...", description="Code repository") + + +class FixDeterminismRequest(BaseModel): + framework: str = Field(default="pytorch", description="Target framework: pytorch, jax, tensorflow") + seed: int = Field(default=42, description="Target random seed") + strict_mode: bool = Field(default=True, description="Enable torch.use_deterministic_algorithms") diff --git a/backend/openmlr/services/sweep_analysis.py b/backend/openmlr/services/sweep_analysis.py new file mode 100644 index 0000000..b70fd3e --- /dev/null +++ b/backend/openmlr/services/sweep_analysis.py @@ -0,0 +1,185 @@ +"""Sweep analysis — parameter sensitivity, rank correlation, Pareto frontiers, and markdown reports.""" + +from __future__ import annotations + +import math +from typing import Any + +from .sweep_types import ParameterSpec, SweepConfig, Trial + + +def _compute_single_param_correlation( + completed: list[Trial], + p_name: str, + spec: ParameterSpec, +) -> tuple[float, float]: + """Compute Spearman/Pearson correlation and importance for a single parameter.""" + x_vals: list[float] = [] + valid_y: list[float] = [] + + for t in completed: + val = t.parameters.get(p_name) + if val is None or t.objective_value is None: + continue + if spec.param_type in ("categorical", "choice"): + x_vals.append(float(hash(str(val)) % 1000)) + else: + try: + x_vals.append(float(val)) + except (ValueError, TypeError): + continue + valid_y.append(t.objective_value) + + if len(x_vals) < 3 or len(set(x_vals)) <= 1: + return 0.0, 0.1 + + mean_x = sum(x_vals) / len(x_vals) + mean_y = sum(valid_y) / len(valid_y) + cov = sum((x - mean_x) * (y - mean_y) for x, y in zip(x_vals, valid_y, strict=False)) + var_x = sum((x - mean_x) ** 2 for x in x_vals) + var_y = sum((y - mean_y) ** 2 for y in valid_y) + + if var_x > 1e-9 and var_y > 1e-9: + r = cov / (math.sqrt(var_x) * math.sqrt(var_y)) + return round(r, 3), round(abs(r), 3) + + return 0.0, 0.0 + + +def _compute_parameter_sensitivities( + completed: list[Trial], + parameters: dict[str, ParameterSpec], +) -> tuple[dict[str, float], dict[str, float]]: + """Compute normalized parameter importance and correlations.""" + importance: dict[str, float] = {} + correlations: dict[str, float] = {} + + for p_name, spec in parameters.items(): + corr, imp = _compute_single_param_correlation(completed, p_name, spec) + correlations[p_name] = corr + importance[p_name] = imp + + total_imp = sum(importance.values()) + if total_imp > 0: + importance = {k: round(v / total_imp, 3) for k, v in importance.items()} + + return importance, correlations + + +def _is_better_or_equal(val: float, baseline: float, goal: str) -> bool: + return val <= baseline if goal == "minimize" else val >= baseline + + +def _is_strictly_better(val: float, baseline: float, goal: str) -> bool: + return val < baseline if goal == "minimize" else val > baseline + + +def _is_dominated(candidate: Trial, others: list[Trial], goal: str) -> bool: + """Check if candidate trial is dominated on both objective metric and duration.""" + c_obj = candidate.objective_value or 0.0 + c_dur = candidate.duration_seconds + + for other in others: + if other.trial_id == candidate.trial_id: + continue + o_obj = other.objective_value or 0.0 + o_dur = other.duration_seconds + + if _is_better_or_equal(o_obj, c_obj, goal) and o_dur <= c_dur: + if _is_strictly_better(o_obj, c_obj, goal) or o_dur < c_dur: + return True + return False + + +def _compute_pareto_frontier(completed: list[Trial], goal: str) -> list[dict[str, Any]]: + """Find all non-dominated trials in the multi-objective space.""" + pareto: list[dict[str, Any]] = [] + for t in completed: + if not _is_dominated(t, completed, goal): + pareto.append(t.to_dict()) + return pareto + + +def calculate_sweep_analysis(sweep: SweepConfig) -> dict[str, Any]: + """Perform statistical analysis, parameter importance, and Pareto frontier evaluation.""" + completed = [t for t in sweep.trials if t.status == "completed" and t.objective_value is not None] + if not completed: + return { + "sweep_id": sweep.sweep_id, + "status": sweep.status, + "total_trials": len(sweep.trials), + "completed_trials": 0, + "best_trial": None, + "parameter_importance": {}, + "correlations": {}, + "pareto_frontier": [], + } + + reverse_sort = sweep.goal == "maximize" + sorted_trials = sorted(completed, key=lambda t: t.objective_value or 0.0, reverse=reverse_sort) + best_trial = sorted_trials[0] + + importance, correlations = _compute_parameter_sensitivities(completed, sweep.parameters) + pareto = _compute_pareto_frontier(completed, sweep.goal) + + return { + "sweep_id": sweep.sweep_id, + "status": sweep.status, + "total_trials": len(sweep.trials), + "completed_trials": len(completed), + "best_trial": best_trial.to_dict(), + "best_parameters": best_trial.parameters, + "best_metric_value": best_trial.objective_value, + "parameter_importance": importance, + "correlations": correlations, + "pareto_frontier": pareto, + } + + +def generate_sweep_markdown_report(sweep: SweepConfig) -> str: + """Export comprehensive sweep findings as a formatted markdown report.""" + analysis = calculate_sweep_analysis(sweep) + lines = [ + f"# Hyperparameter Optimization Report: {sweep.name}", + f"- **Sweep ID**: `{sweep.sweep_id}`", + f"- **Search Method**: `{sweep.method.upper()}`", + f"- **Objective**: `{sweep.objective_metric}` ({sweep.goal})", + f"- **Trials Completed**: {analysis['completed_trials']}/{len(sweep.trials)} (Max: {sweep.max_trials})", + f"- **Status**: `{sweep.status.upper()}`", + "", + ] + + if analysis.get("best_trial"): + bt = analysis["best_trial"] + lines.extend([ + "### 🏆 Optimal Configuration", + f"- **Trial ID**: `{bt['trial_id']}`", + f"- **Best {sweep.objective_metric}**: `{analysis['best_metric_value']}`", + "- **Parameters**:", + ]) + for k, v in bt["parameters"].items(): + lines.append(f" - `{k}`: `{v}`") + lines.append("") + + if analysis.get("parameter_importance"): + lines.extend([ + "### 📊 Parameter Sensitivity & Importance", + "| Hyperparameter | Importance | Correlation |", + "| :--- | :--- | :--- |", + ]) + for p, imp in analysis["parameter_importance"].items(): + corr = analysis["correlations"].get(p, 0.0) + lines.append(f"| `{p}` | {imp * 100:.1f}% | {corr:+.3f} |") + lines.append("") + + lines.extend([ + "### 🧪 Trial History", + f"| Trial | Status | {sweep.objective_metric} | Parameters | Runtime |", + "| :--- | :--- | :--- | :--- | :--- |", + ]) + for t in sweep.trials: + param_str = ", ".join(f"{k}={v}" for k, v in list(t.parameters.items())[:3]) + val_str = f"{t.objective_value:.4f}" if t.objective_value is not None else "-" + lines.append(f"| `{t.trial_id}` | `{t.status}` | {val_str} | {param_str} | {t.duration_seconds:.1f}s |") + + return "\n".join(lines) diff --git a/backend/openmlr/services/sweep_engine.py b/backend/openmlr/services/sweep_engine.py new file mode 100644 index 0000000..db8cf6b --- /dev/null +++ b/backend/openmlr/services/sweep_engine.py @@ -0,0 +1,384 @@ +"""Sweep Engine — Agent-native hyperparameter optimization, search spaces, and trial evaluation.""" + +from __future__ import annotations + +import json +import logging +import math +import secrets +import time +import uuid +from pathlib import Path +from typing import Any + +from .sweep_analysis import calculate_sweep_analysis, generate_sweep_markdown_report +from .sweep_types import EarlyStoppingConfig, ParameterSpec, SweepConfig, Trial + +log = logging.getLogger(__name__) + +# Re-export types for backward compatibility +__all__ = [ + "EarlyStoppingConfig", + "ParameterSpec", + "SweepConfig", + "SweepEngine", + "Trial", +] + + +def _safe_uniform(min_v: float, max_v: float) -> float: + """Generate uniform float in range [min_v, max_v] using secure random bits.""" + scale = secrets.randbits(32) / (1 << 32) + return min_v + scale * (max_v - min_v) + + +def _safe_choice(choices: list[Any], default: Any) -> Any: + """Select a random choice safely.""" + if not choices: + return default + idx = secrets.randbelow(len(choices)) + return choices[idx] + + +def _score_param_match(spec: ParameterSpec, val1: Any, val2: Any) -> float: + """Score similarity between two parameter values.""" + if val1 is None or val2 is None: + return 0.0 + if spec.param_type in ("categorical", "choice"): + return 1.0 if val1 == val2 else 0.0 + norm_range = (spec.max_val or 1.0) - (spec.min_val or 0.0) + if norm_range <= 0: + return 1.0 + dist = abs(float(val1) - float(val2)) / norm_range + return max(0.0, 1.0 - dist) + + +def _score_candidate(candidate: dict[str, Any], good_trials: list[Trial], parameters: dict[str, ParameterSpec]) -> float: + """Score candidate similarity against top-performing historical trials.""" + sim_score = 0.0 + total_count = max(1, len(parameters)) + for good_t in good_trials: + match_count = sum( + _score_param_match(spec, candidate.get(p_name), good_t.parameters.get(p_name)) + for p_name, spec in parameters.items() + ) + sim_score += match_count / total_count + return sim_score + + +class SweepEngine: + """Service to create sweeps, suggest next trials, prune underperforming runs, and evaluate results.""" + + def __init__(self, base_dir: Path | None = None): + self.base_dir = base_dir or Path(".openmlr/sweeps") + self.base_dir.mkdir(parents=True, exist_ok=True) + + def _sweep_file(self, project_id: str, sweep_id: str) -> Path: + p_dir = self.base_dir / project_id + p_dir.mkdir(parents=True, exist_ok=True) + return p_dir / f"{sweep_id}.json" + + def save_sweep(self, sweep: SweepConfig) -> None: + """Persist sweep state to disk.""" + sweep.updated_at = time.time() + file_path = self._sweep_file(sweep.project_id, sweep.sweep_id) + with open(file_path, "w", encoding="utf-8") as f: + json.dump(sweep.to_dict(), f, indent=2) + + def get_sweep(self, project_id: str, sweep_id: str) -> SweepConfig | None: + """Load sweep by ID from disk.""" + file_path = self._sweep_file(project_id, sweep_id) + if not file_path.exists(): + return None + try: + with open(file_path, encoding="utf-8") as f: + data = json.load(f) + return SweepConfig.from_dict(data) + except Exception: + log.exception("Failed to load sweep %s/%s", project_id, sweep_id) + return None + + def list_sweeps(self, project_id: str) -> list[SweepConfig]: + """List all sweeps for a given project.""" + p_dir = self.base_dir / project_id + if not p_dir.exists(): + return [] + sweeps = [] + for p in p_dir.glob("*.json"): + try: + with open(p, encoding="utf-8") as f: + data = json.load(f) + sweeps.append(SweepConfig.from_dict(data)) + except Exception as e: + log.warning("Skipping corrupted sweep file %s: %s", p, e) + sweeps.sort(key=lambda s: s.created_at, reverse=True) + return sweeps + + def delete_sweep(self, project_id: str, sweep_id: str) -> bool: + """Delete a sweep file.""" + file_path = self._sweep_file(project_id, sweep_id) + if file_path.exists(): + file_path.unlink() + return True + return False + + def create_sweep( + self, + project_id: str, + name: str, + method: str, + objective_metric: str, + goal: str, + parameters: dict[str, Any], + max_trials: int = 10, + description: str = "", + early_stopping: EarlyStoppingConfig | dict[str, Any] | None = None, + ) -> SweepConfig: + """Create and initialize a new hyperparameter sweep.""" + sweep_id = f"swp_{str(uuid.uuid4())[:8]}" + param_specs: dict[str, ParameterSpec] = {} + for k, v in parameters.items(): + if isinstance(v, ParameterSpec): + param_specs[k] = v + else: + v_dict = dict(v) + v_dict["name"] = k + param_specs[k] = ParameterSpec.from_dict(v_dict) + + es_conf = ( + early_stopping + if isinstance(early_stopping, EarlyStoppingConfig) + else EarlyStoppingConfig.from_dict(early_stopping or {}) + ) + + sweep = SweepConfig( + sweep_id=sweep_id, + project_id=project_id, + name=name, + description=description, + method=method.lower(), + objective_metric=objective_metric, + goal=goal.lower(), + max_trials=max_trials, + parameters=param_specs, + early_stopping=es_conf, + trials=[], + status="active", + ) + self.save_sweep(sweep) + return sweep + + def suggest_trial(self, project_id: str, sweep_id: str) -> Trial | None: + """Generate the next parameter proposal for the sweep.""" + sweep = self.get_sweep(project_id, sweep_id) + if not sweep: + raise ValueError(f"Sweep '{sweep_id}' not found in project '{project_id}'") + + if len(sweep.trials) >= sweep.max_trials: + sweep.status = "completed" + self.save_sweep(sweep) + return None + + trial_num = len(sweep.trials) + 1 + trial_id = f"tr_{sweep_id[-4:]}_{trial_num:03d}" + + if sweep.method == "grid": + params = self._suggest_grid(sweep, trial_num) + elif sweep.method in ("bayesian", "bayes"): + params = self._suggest_bayesian(sweep) + else: + params = self._suggest_random(sweep) + + trial = Trial( + trial_id=trial_id, + sweep_id=sweep_id, + trial_number=trial_num, + parameters=params, + status="running", + started_at=time.time(), + ) + sweep.trials.append(trial) + self.save_sweep(sweep) + return trial + + def _sample_single_param(self, spec: ParameterSpec) -> Any: + """Sample a single parameter randomly within its spec.""" + if spec.param_type in ("categorical", "choice"): + return _safe_choice(spec.choices, spec.default) + if spec.param_type == "int_uniform": + min_v = int(spec.min_val or 0) + max_v = int(spec.max_val or 10) + step = max(1, int(spec.step or 1)) + count = max(1, ((max_v - min_v) // step) + 1) + return min_v + secrets.randbelow(count) * step + if spec.param_type == "loguniform": + min_v = max(1e-7, spec.min_val or 1e-4) + max_v = spec.max_val or 1.0 + log_min, log_max = math.log(min_v), math.log(max_v) + val = math.exp(_safe_uniform(log_min, log_max)) + return round(val, 6) + min_v = spec.min_val or 0.0 + max_v = spec.max_val or 1.0 + val = _safe_uniform(min_v, max_v) + if spec.step: + val = round(round((val - min_v) / spec.step) * spec.step + min_v, 4) + else: + val = round(val, 4) + return val + + def _suggest_random(self, sweep: SweepConfig) -> dict[str, Any]: + """Generate a random parameter sample across all dimensions.""" + return {name: self._sample_single_param(spec) for name, spec in sweep.parameters.items()} + + def _suggest_grid(self, sweep: SweepConfig, trial_num: int) -> dict[str, Any]: + """Generate parameter combination using deterministic Cartesian product indexing.""" + grids: list[tuple[str, list[Any]]] = [] + for name, spec in sweep.parameters.items(): + if spec.choices: + values = list(spec.choices) + elif spec.param_type == "int_uniform": + min_v = int(spec.min_val or 0) + max_v = int(spec.max_val or 10) + step = int(spec.step or 1) + values = list(range(min_v, max_v + 1, step)) + else: + min_v = spec.min_val or 0.0 + max_v = spec.max_val or 1.0 + step = spec.step or ((max_v - min_v) / 4) + steps = max(2, int(round((max_v - min_v) / step)) + 1) + values = [round(min_v + i * step, 4) for i in range(steps)] + grids.append((name, values)) + + idx = trial_num - 1 + params = {} + for name, values in reversed(grids): + chosen = values[idx % len(values)] + params[name] = chosen + idx //= len(values) + return params + + def _suggest_bayesian(self, sweep: SweepConfig) -> dict[str, Any]: + """Generate suggestion using Parzen density estimation & expected improvement surrogate.""" + completed_trials = [ + t for t in sweep.trials if t.status == "completed" and t.objective_value is not None + ] + if len(completed_trials) < 3: + return self._suggest_random(sweep) + + reverse_sort = sweep.goal == "maximize" + sorted_trials = sorted( + completed_trials, key=lambda t: t.objective_value or 0.0, reverse=reverse_sort + ) + + split_idx = max(1, len(sorted_trials) // 4) + good_trials = sorted_trials[:split_idx] + + candidates = [self._suggest_random(sweep) for _ in range(25)] + best_candidate = candidates[0] + best_score = -1.0 + + for cand in candidates: + score = _score_candidate(cand, good_trials, sweep.parameters) + if score > best_score: + best_score = score + best_candidate = cand + + return best_candidate + + def record_trial_result( + self, + project_id: str, + sweep_id: str, + trial_id: str, + metrics: dict[str, Any], + status: str = "completed", + step_history: list[dict[str, Any]] | None = None, + error_message: str | None = None, + ) -> Trial: + """Record the evaluation metrics and completion status of a trial.""" + sweep = self.get_sweep(project_id, sweep_id) + if not sweep: + raise ValueError(f"Sweep '{sweep_id}' not found") + + trial = next((t for t in sweep.trials if t.trial_id == trial_id), None) + if not trial: + raise ValueError(f"Trial '{trial_id}' not found in sweep '{sweep_id}'") + + trial.metrics = metrics + trial.status = status + trial.completed_at = time.time() + trial.duration_seconds = round(trial.completed_at - trial.started_at, 2) + if step_history: + trial.step_history = step_history + if error_message: + trial.error_message = error_message + + obj_val = metrics.get(sweep.objective_metric) + if obj_val is not None: + try: + trial.objective_value = float(obj_val) + except (ValueError, TypeError): + trial.objective_value = None + + completed_count = len([t for t in sweep.trials if t.status in ("completed", "failed", "pruned")]) + if completed_count >= sweep.max_trials: + sweep.status = "completed" + + self.save_sweep(sweep) + return trial + + def should_prune_trial( + self, + project_id: str, + sweep_id: str, + trial_id: str, + current_step: int, + current_metric_val: float, + ) -> bool: + """Evaluate whether a running trial should be early-stopped based on ASHA/Hyperband percentiles.""" + sweep = self.get_sweep(project_id, sweep_id) + if not sweep or not sweep.early_stopping.enabled: + return False + + es = sweep.early_stopping + if current_step < es.min_steps: + return False + + if es.metric_threshold is not None: + if sweep.goal == "minimize" and current_metric_val > es.metric_threshold: + return True + if sweep.goal == "maximize" and current_metric_val < es.metric_threshold: + return True + + past_metrics = [ + float(entry[sweep.objective_metric]) + for t in sweep.trials + if t.trial_id != trial_id + for entry in t.step_history + if entry.get("step") == current_step and sweep.objective_metric in entry + ] + + if len(past_metrics) < 2: + return False + + past_metrics.sort(reverse=(sweep.goal == "maximize")) + top_k = max(1, int(len(past_metrics) / es.reduction_factor)) + cutoff = past_metrics[top_k - 1] + + if sweep.goal == "minimize": + return current_metric_val > cutoff * 1.15 + return current_metric_val < cutoff * 0.85 + + def analyze_sweep(self, project_id: str, sweep_id: str) -> dict[str, Any]: + """Calculate parameter sensitivities, correlation matrix, optimal trial, and Pareto frontier.""" + sweep = self.get_sweep(project_id, sweep_id) + if not sweep: + raise ValueError(f"Sweep '{sweep_id}' not found") + return calculate_sweep_analysis(sweep) + + def export_sweep_markdown(self, project_id: str, sweep_id: str) -> str: + """Export comprehensive sweep findings as a formatted markdown report.""" + sweep = self.get_sweep(project_id, sweep_id) + if not sweep: + return f"Sweep {sweep_id} not found." + return generate_sweep_markdown_report(sweep) diff --git a/backend/openmlr/services/sweep_types.py b/backend/openmlr/services/sweep_types.py new file mode 100644 index 0000000..d24f53c --- /dev/null +++ b/backend/openmlr/services/sweep_types.py @@ -0,0 +1,153 @@ +"""Sweep data types and configuration models.""" + +from __future__ import annotations + +import time +import uuid +from dataclasses import asdict, dataclass, field +from typing import Any + + +@dataclass +class ParameterSpec: + """Specification of a hyperparameter search space dimension.""" + + name: str + param_type: str # "categorical", "uniform", "loguniform", "int_uniform", "choice" + min_val: float | None = None + max_val: float | None = None + step: float | None = None + choices: list[Any] = field(default_factory=list) + default: Any | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> ParameterSpec: + return cls( + name=data.get("name", "param"), + param_type=data.get("param_type", "uniform"), + min_val=data.get("min_val"), + max_val=data.get("max_val"), + step=data.get("step"), + choices=data.get("choices", []) or [], + default=data.get("default"), + ) + + +@dataclass +class EarlyStoppingConfig: + """Configuration for trial early stopping and pruning (e.g. ASHA / Hyperband).""" + + enabled: bool = False + min_steps: int = 5 + reduction_factor: float = 3.0 + metric_threshold: float | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> EarlyStoppingConfig: + return cls( + enabled=data.get("enabled", False), + min_steps=int(data.get("min_steps", 5)), + reduction_factor=float(data.get("reduction_factor", 3.0)), + metric_threshold=data.get("metric_threshold"), + ) + + +@dataclass +class Trial: + """A single hyperparameter trial run.""" + + trial_id: str + sweep_id: str + trial_number: int + parameters: dict[str, Any] + status: str = "pending" # "pending", "running", "completed", "failed", "pruned" + metrics: dict[str, Any] = field(default_factory=dict) + objective_value: float | None = None + step_history: list[dict[str, Any]] = field(default_factory=list) + started_at: float = field(default_factory=time.time) + completed_at: float | None = None + error_message: str | None = None + duration_seconds: float = 0.0 + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Trial: + return cls( + trial_id=data.get("trial_id", str(uuid.uuid4())[:8]), + sweep_id=data.get("sweep_id", ""), + trial_number=int(data.get("trial_number", 1)), + parameters=data.get("parameters", {}), + status=data.get("status", "pending"), + metrics=data.get("metrics", {}), + objective_value=data.get("objective_value"), + step_history=data.get("step_history", []), + started_at=data.get("started_at", time.time()), + completed_at=data.get("completed_at"), + error_message=data.get("error_message"), + duration_seconds=data.get("duration_seconds", 0.0), + ) + + +@dataclass +class SweepConfig: + """Complete specification of a hyperparameter sweep.""" + + sweep_id: str + project_id: str + name: str + description: str + method: str # "grid", "random", "bayesian", "hyperband" + objective_metric: str + goal: str # "minimize", "maximize" + max_trials: int + parameters: dict[str, ParameterSpec] + early_stopping: EarlyStoppingConfig = field(default_factory=EarlyStoppingConfig) + trials: list[Trial] = field(default_factory=list) + status: str = "active" # "active", "completed", "archived" + created_at: float = field(default_factory=time.time) + updated_at: float = field(default_factory=time.time) + + def to_dict(self) -> dict[str, Any]: + d = asdict(self) + d["parameters"] = {k: v.to_dict() for k, v in self.parameters.items()} + d["early_stopping"] = self.early_stopping.to_dict() + d["trials"] = [t.to_dict() for t in self.trials] + return d + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> SweepConfig: + raw_params = data.get("parameters", {}) + params = { + k: ParameterSpec.from_dict(v) if isinstance(v, dict) else v + for k, v in raw_params.items() + } + early_stopping = EarlyStoppingConfig.from_dict(data.get("early_stopping", {})) + raw_trials = data.get("trials", []) + trials = [ + Trial.from_dict(t) if isinstance(t, dict) else t + for t in raw_trials + ] + return cls( + sweep_id=data.get("sweep_id", str(uuid.uuid4())[:8]), + project_id=data.get("project_id", "default"), + name=data.get("name", "Sweep"), + description=data.get("description", ""), + method=data.get("method", "random"), + objective_metric=data.get("objective_metric", "val_loss"), + goal=data.get("goal", "minimize"), + max_trials=int(data.get("max_trials", 10)), + parameters=params, + early_stopping=early_stopping, + trials=trials, + status=data.get("status", "active"), + created_at=data.get("created_at", time.time()), + updated_at=data.get("updated_at", time.time()), + ) diff --git a/backend/openmlr/tools/datasets.py b/backend/openmlr/tools/datasets.py new file mode 100644 index 0000000..49629f1 --- /dev/null +++ b/backend/openmlr/tools/datasets.py @@ -0,0 +1,352 @@ +"""Datasets tool — dataset profiling, inspection, validation, and curation for AI research agents.""" + +from __future__ import annotations + +import asyncio +import json +import logging +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from ..agent.types import ToolSpec +from ..services.dataset_profiler import DatasetProfile, DatasetProfiler + +log = logging.getLogger(__name__) + + +def _parse_list(val: Any) -> list[str]: + """Parse list or comma-separated string.""" + if isinstance(val, list): + return [str(x).strip() for x in val if str(x).strip()] + if isinstance(val, str) and val.strip(): + if val.strip().startswith("[") and val.strip().endswith("]"): + try: + parsed = json.loads(val) + if isinstance(parsed, list): + return [str(x).strip() for x in parsed if str(x).strip()] + except Exception: + pass + return [x.strip() for x in val.split(",") if x.strip()] + return [] + + +def _format_profile_markdown(file_path: Path, profile: DatasetProfile) -> str: + """Format DatasetProfile into human-readable markdown.""" + p_dict = profile.to_dict() + output = [ + f"# Dataset Profile: {file_path.name}", + f"- **Format**: {profile.format.upper()}", + f"- **Sampled Rows**: {profile.total_rows:,}", + f"- **Columns**: {profile.total_columns}", + f"- **File Size**: {profile.file_size_bytes / (1024 * 1024):.2f} MB", + f"- **Health Score**: {profile.health_score}/100", + "", + ] + + if profile.warnings: + output.append("### Diagnostic Warnings") + for w in profile.warnings: + output.append(f"- ⚠️ {w}") + output.append("") + + output.append("### Column Details") + for c_name, c_prof in profile.columns.items(): + stats = c_prof.stats + line = f"- **`{c_name}`** (`{c_prof.dtype}`): {c_prof.null_percentage}% null ({c_prof.null_count}/{c_prof.total_count}), {c_prof.unique_count} unique." + if c_prof.dtype == "numeric": + line += f" Range: [{stats.get('min')}, {stats.get('max')}], Mean: {stats.get('mean')}, Std: {stats.get('std')}." + elif c_prof.dtype == "categorical": + line += f" Imbalance: {stats.get('imbalance_ratio', 1.0)}x. Top classes: {list(stats.get('top_classes', {}).keys())[:4]}." + elif c_prof.dtype == "text": + line += f" Avg chars: {stats.get('char_len_avg')}, Mean tokens: {stats.get('token_est_mean')}, Max tokens: {stats.get('token_est_max')}." + output.append(line) + + output.append("\n```json\n" + json.dumps(p_dict, indent=2) + "\n```") + return "\n".join(output) + + +def _resolve_target_file(path: str) -> Path | None: + if not path: + return None + p = Path(path).resolve() + return p if p.exists() else None + + +def _op_profile(path: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + file_path = _resolve_target_file(path) + if not file_path: + return f"Error: Dataset file not found at '{path}'.", False + sample_size = int(kwargs.get("sample_size", 2000)) + profile = DatasetProfiler.profile(file_path, sample_size=sample_size) + return _format_profile_markdown(file_path, profile), True + + +def _op_inspect(path: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + file_path = _resolve_target_file(path) + if not file_path: + return f"Error: Dataset file not found at '{path}'.", False + strategy = str(kwargs.get("strategy", "head")) + samples = DatasetProfiler.sample_records( + file_path, + n=int(kwargs.get("n", 5)), + offset=int(kwargs.get("offset", 0)), + strategy=strategy, + label_column=kwargs.get("label_column"), + ) + output = [ + f"# Dataset Samples: {file_path.name} (strategy={strategy}, n={len(samples)})", + "```json", + json.dumps(samples, indent=2, default=str), + "```", + ] + return "\n".join(output), True + + +def _op_validate(path: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + file_path = _resolve_target_file(path) + if not file_path: + return f"Error: Dataset file not found at '{path}'.", False + exp_cols = _parse_list(kwargs.get("expected_columns")) if kwargs.get("expected_columns") else None + val_res = DatasetProfiler.validate_dataset( + file_path, + expected_columns=exp_cols, + max_null_pct=float(kwargs.get("max_null_pct", 20.0)), + max_token_length=kwargs.get("max_token_length"), + ) + + status_str = "PASSED ✅" if val_res["valid"] else "FAILED ❌" + output = [ + f"# Dataset Validation: {file_path.name} — {status_str}", + f"- **Health Score**: {val_res.get('health_score', 0)}/100", + f"- **Rows**: {val_res.get('total_rows', 0)} | **Columns**: {val_res.get('total_columns', 0)}", + "", + ] + if val_res["errors"]: + output.append("### Validation Errors") + for err in val_res["errors"]: + output.append(f"- ❌ {err}") + output.append("") + + if val_res["warnings"]: + output.append("### Warnings") + for w in val_res["warnings"]: + output.append(f"- ⚠️ {w}") + + return "\n".join(output), val_res["valid"] + + +def _op_split(path: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + file_path = _resolve_target_file(path) + if not file_path: + return f"Error: Dataset file not found at '{path}'.", False + out_dir_param = str(kwargs.get("output_dir", "")) + target_dir = out_dir_param or str(file_path.parent / f"{file_path.stem}_splits") + manifest = DatasetProfiler.split_dataset( + file_path, + output_dir=target_dir, + train_ratio=float(kwargs.get("train_ratio", 0.8)), + val_ratio=float(kwargs.get("val_ratio", 0.1)), + test_ratio=float(kwargs.get("test_ratio", 0.1)), + stratify_column=kwargs.get("stratify_column"), + ) + + output = [ + f"# Dataset Split Completed: {file_path.name}", + f"- **Train records**: {manifest['train_count']:,} (`{manifest['splits']['train']}`)", + f"- **Val records**: {manifest['val_count']:,} (`{manifest['splits']['val']}`)", + f"- **Test records**: {manifest['test_count']:,} (`{manifest['splits']['test']}`)", + f"- **Stratified**: {manifest['stratified_by'] or 'Random'}", + f"- **Manifest**: `{target_dir}/split_manifest.json`", + ] + return "\n".join(output), True + + +def _sync_kg(session: Any, name: str, file_path: Path | None, profile: DatasetProfile | None, description: str, tags: Any) -> bool: + kg = getattr(getattr(session, "workspace", None), "knowledge_graph", None) + if not kg: + return False + props = { + "path": str(file_path) if file_path else "", + "format": profile.format if profile else "", + "total_rows": profile.total_rows if profile else 0, + "health_score": profile.health_score if profile else 100, + "description": description, + "tags": _parse_list(tags) if tags else [], + } + kg.add_entity( + entity_id=f"dataset_{name}", + entity_type="dataset", + name=name, + properties=props, + ) + return True + + +def _op_register(path: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + file_path = _resolve_target_file(path) if path else None + dataset_name = str(kwargs.get("dataset_name", "")) + description = str(kwargs.get("description", "")) + tags = kwargs.get("tags") + session = kwargs.get("session") + + name = dataset_name or (file_path.stem if file_path else "dataset") + profile = DatasetProfiler.profile(file_path, sample_size=1000) if file_path else None + kg_synced = _sync_kg(session, name, file_path, profile, description, tags) + + output = [ + f"# Dataset Registered: `{name}`", + f"- **Path**: `{file_path or path or 'N/A'}`", + f"- **Knowledge Graph Synced**: {'Yes' if kg_synced else 'No (in-memory only)'}", + f"- **Description**: {description or 'N/A'}", + ] + return "\n".join(output), True + + +def _op_summary(_path: str, _kwargs: dict[str, Any]) -> tuple[str, bool]: + return ( + "Datasets Tool Operations:\n" + "- `profile`: Compute statistics, column distributions, null analysis, text token lengths, and health score.\n" + "- `inspect_samples`: Preview sample records (head, random, stratified).\n" + "- `validate`: Validate against schema constraints, missing values, and token limits.\n" + "- `split`: Partition dataset into train/val/test splits.\n" + "- `register`: Register dataset in project Knowledge Graph." + ), True + + +_OP_HANDLERS: dict[str, Callable[[str, dict[str, Any]], tuple[str, bool]]] = { + "profile": _op_profile, + "analyze": _op_profile, + "inspect_samples": _op_inspect, + "sample": _op_inspect, + "preview": _op_inspect, + "validate": _op_validate, + "check": _op_validate, + "split": _op_split, + "partition": _op_split, + "register": _op_register, + "register_kg": _op_register, + "summary": _op_summary, + "help": _op_summary, +} + + +async def _handle_datasets( + operation: str = "", + path: str = "", + **kwargs: Any, +) -> tuple[str, bool]: + """Handle dataset profiling, inspection, validation, split, and registration operations.""" + await asyncio.sleep(0) + op = (operation or "").strip().lower() + + if not path and op not in ("summary", "help"): + return "Error: 'path' parameter is required for dataset operations.", False + + handler = _OP_HANDLERS.get(op) + if not handler: + return ( + f"Unknown datasets operation '{operation}'. Supported: profile, inspect_samples, validate, split, register, summary.", + False, + ) + + try: + return handler(path, kwargs) + except Exception as e: + log.exception("Error executing datasets operation '%s': %s", op, e) + return f"Error executing datasets operation '{op}': {e}", False + + +def create_datasets_tool() -> ToolSpec: + """Create the ToolSpec definition for the datasets tool.""" + return ToolSpec( + name="datasets", + description=( + "Inspect, profile, validate, and partition machine learning datasets. " + "Supports CSV, TSV, JSON, JSONL, and text formats. Computes column distributions, " + "missingness, class balance, token length distributions, and generates reproducible splits." + ), + parameters={ + "type": "object", + "properties": { + "operation": { + "type": "string", + "description": "Operation: profile, inspect_samples, validate, split, register, summary", + "enum": ["profile", "inspect_samples", "validate", "split", "register", "summary"], + }, + "path": { + "type": "string", + "description": "Path to the dataset file (CSV, JSONL, TSV, JSON, TXT)", + }, + "sample_size": { + "type": "integer", + "description": "Maximum number of rows to sample for profiling (default: 2000)", + }, + "n": { + "type": "integer", + "description": "Number of sample rows to inspect (default: 5)", + }, + "offset": { + "type": "integer", + "description": "Offset index for inspecting samples (default: 0)", + }, + "strategy": { + "type": "string", + "description": "Sampling strategy: head, random, stratified (default: head)", + "enum": ["head", "random", "stratified"], + }, + "label_column": { + "type": "string", + "description": "Label or class column name for stratified sampling", + }, + "expected_columns": { + "type": "array", + "items": {"type": "string"}, + "description": "List of expected column names for validation", + }, + "max_null_pct": { + "type": "number", + "description": "Maximum allowed null percentage per column in validation (default: 20.0)", + }, + "max_token_length": { + "type": "integer", + "description": "Maximum allowed token count for text columns in validation", + }, + "output_dir": { + "type": "string", + "description": "Target directory for writing train/val/test splits", + }, + "train_ratio": { + "type": "number", + "description": "Train partition ratio (default: 0.8)", + }, + "val_ratio": { + "type": "number", + "description": "Validation partition ratio (default: 0.1)", + }, + "test_ratio": { + "type": "number", + "description": "Test partition ratio (default: 0.1)", + }, + "stratify_column": { + "type": "string", + "description": "Column name to stratify class distributions when splitting", + }, + "dataset_name": { + "type": "string", + "description": "Identifier name for registering in Knowledge Graph", + }, + "description": { + "type": "string", + "description": "Research notes or description of the dataset", + }, + "tags": { + "type": "array", + "items": {"type": "string"}, + "description": "Categorization tags", + }, + }, + "required": ["operation"], + }, + handler=_handle_datasets, + ) diff --git a/backend/openmlr/tools/experiments.py b/backend/openmlr/tools/experiments.py new file mode 100644 index 0000000..55a223f --- /dev/null +++ b/backend/openmlr/tools/experiments.py @@ -0,0 +1,485 @@ +"""Experiments tool — ML experiment tracking and run management for the AI research agent. + +Allows the autonomous research agent to: +- Create new experiment runs with hyperparameters, architecture details, and compute targets +- Log metric curves (train/val loss, learning rate, GPU utilization, throughput) +- Record model checkpoint artifacts and evaluation scores +- Inspect run statuses, latest metrics, and execution progress +- Compare multiple runs across ablation studies and hyperparameter sweeps +- Complete or fail runs with outcome summaries and findings +""" + +from __future__ import annotations + +import json +import logging +from contextvars import ContextVar +from typing import Any + +from ..agent.types import ToolSpec +from ..services.experiment_tracker import ExperimentTracker + +log = logging.getLogger(__name__) + +# Context variable for per-request experiment tracker & project UUID +_experiment_tracker_var: ContextVar[ExperimentTracker | None] = ContextVar( + "experiment_tracker", default=None +) +_project_uuid_var: ContextVar[str | None] = ContextVar("experiment_project_uuid", default=None) + +_default_tracker = ExperimentTracker() + + +def set_experiment_context( + tracker: ExperimentTracker | None, project_uuid: str | None = None +) -> None: + """Set the active experiment tracker and project UUID for the current async context.""" + _experiment_tracker_var.set(tracker) + _project_uuid_var.set(project_uuid) + + +def _get_active_tracker() -> ExperimentTracker: + """Get the active tracker for the current async context, falling back to default.""" + return _experiment_tracker_var.get() or _default_tracker + + +def _parse_dict_or_json(val: Any) -> dict[str, Any]: + """Parse dict or JSON string into a dict.""" + if isinstance(val, dict): + return val + if isinstance(val, str) and val.strip(): + try: + parsed = json.loads(val) + if isinstance(parsed, dict): + return parsed + except Exception: + return {"raw": val} + return {} + + +def _parse_list_or_csv(val: Any) -> list[str]: + """Parse list or comma-separated string into a list of trimmed strings.""" + if isinstance(val, list): + return [str(x).strip() for x in val if str(x).strip()] + if isinstance(val, str) and val.strip(): + if val.strip().startswith("[") and val.strip().endswith("]"): + try: + parsed = json.loads(val) + if isinstance(parsed, list): + return [str(x).strip() for x in parsed if str(x).strip()] + except Exception: + pass + return [x.strip() for x in val.split(",") if x.strip()] + return [] + + +async def _handle_experiments( + operation: str, + name: str = "", + description: str = "", + hyperparameters: Any = None, + compute_target: str = "Local GPU", + tags: Any = None, + total_steps: int = 100, + total_epochs: int = 1, + run_id: str = "", + metrics: Any = None, + step: int = 0, + epoch: int = 1, + checkpoint_name: str = "", + path: str = "", + file_size_mb: float = 0.0, + status: str = "", + reason: str = "", + best_val_loss: float | None = None, + run_ids: Any = None, + search: str = "", + limit: int = 10, + session=None, + **kwargs: Any, +) -> tuple[str, bool]: + """Handle experiment tool operations.""" + tracker = _get_active_tracker() + project_uuid = _project_uuid_var.get() + + try: + if operation == "create_run": + if not name.strip(): + return "Error: 'name' is required when creating an experiment run.", False + + hp_dict = _parse_dict_or_json(hyperparameters) + tag_list = _parse_list_or_csv(tags) + + run = tracker.create_run( + name=name.strip(), + description=description.strip(), + hyperparameters=hp_dict, + compute_target=compute_target.strip() or "Local GPU", + tags=tag_list, + total_steps=max(1, total_steps), + total_epochs=max(1, total_epochs), + project_uuid=project_uuid, + ) + + # Auto-record in knowledge graph if active workspace persistence exists + try: + from .workspace_tools import _knowledge_var + + kg = _knowledge_var.get() + if kg: + kg.add_entity( + entity_id=f"exp_{run.id}", + entity_type="experiment", + label=run.name, + properties={ + "run_id": run.id, + "compute_target": run.compute_target, + "hyperparameters": run.hyperparameters, + "status": run.status, + }, + ) + kg.save() + except Exception as e: + log.debug("Knowledge graph auto-entity skipped for experiment %s: %s", run.id, e) + + result = { + "message": f"Experiment run '{run.name}' created successfully.", + "run_id": run.id, + "name": run.name, + "status": run.status, + "hyperparameters": run.hyperparameters, + "compute_target": run.compute_target, + "total_steps": run.total_steps, + "total_epochs": run.total_epochs, + } + return json.dumps(result, indent=2), True + + elif operation == "log_metrics": + if not run_id.strip(): + return "Error: 'run_id' is required to log metrics.", False + + run = tracker.get_run(run_id.strip()) + if not run: + return f"Error: Experiment run '{run_id}' not found.", False + + metric_dict = _parse_dict_or_json(metrics) + if not metric_dict: + return "Error: 'metrics' dictionary cannot be empty.", False + + # Convert all numeric values to float + clean_metrics: dict[str, float] = {} + for k, v in metric_dict.items(): + try: + clean_metrics[k] = float(v) + except (ValueError, TypeError): + continue + + updated_run = tracker.log_metrics( + run_id=run.id, + metrics=clean_metrics, + step=max(0, step), + epoch=max(1, epoch), + ) + + if not updated_run: + return f"Failed to log metrics for run '{run_id}'.", False + + result = { + "message": f"Metrics logged at step {step}, epoch {epoch}.", + "run_id": updated_run.id, + "current_step": updated_run.current_step, + "total_steps": updated_run.total_steps, + "logged_metrics": clean_metrics, + "best_val_loss": updated_run.best_val_loss, + } + return json.dumps(result, indent=2), True + + elif operation == "record_checkpoint": + if not run_id.strip(): + return "Error: 'run_id' is required to record a checkpoint.", False + + run = tracker.get_run(run_id.strip()) + if not run: + return f"Error: Experiment run '{run_id}' not found.", False + + cp_name = checkpoint_name.strip() or f"checkpoint-step-{step}.pt" + cp_path = path.strip() or f"checkpoints/{cp_name}" + eval_metrics = _parse_dict_or_json(metrics) + clean_eval: dict[str, float] = {} + for k, v in eval_metrics.items(): + try: + clean_eval[k] = float(v) + except (ValueError, TypeError): + continue + + cp = tracker.register_checkpoint( + run_id=run.id, + name=cp_name, + step=max(0, step), + epoch=max(1, epoch), + path=cp_path, + file_size_mb=max(0.0, float(file_size_mb)), + metrics=clean_eval, + ) + + result = { + "message": f"Checkpoint '{cp_name}' recorded for run '{run.id}'.", + "run_id": run.id, + "checkpoint": cp.to_dict(), + "total_checkpoints": len(run.checkpoints), + } + return json.dumps(result, indent=2), True + + elif operation == "get_run": + if not run_id.strip(): + return "Error: 'run_id' is required.", False + + run = tracker.get_run(run_id.strip()) + if not run: + return f"Error: Experiment run '{run_id}' not found.", False + + # Extract latest metric values for concise view + latest_metrics: dict[str, float] = {} + for k, pts in run.metrics.items(): + if pts: + latest_metrics[k] = pts[-1].value + + result = { + "run_id": run.id, + "name": run.name, + "description": run.description, + "status": run.status, + "started_at": run.started_at, + "ended_at": run.ended_at, + "duration_seconds": run.duration_seconds, + "compute_target": run.compute_target, + "tags": run.tags, + "hyperparameters": run.hyperparameters, + "current_step": run.current_step, + "total_steps": run.total_steps, + "current_epoch": run.current_epoch, + "total_epochs": run.total_epochs, + "best_val_loss": run.best_val_loss, + "latest_metrics": latest_metrics, + "checkpoints_count": len(run.checkpoints), + "checkpoints": [cp.to_dict() for cp in run.checkpoints], + } + return json.dumps(result, indent=2), True + + elif operation == "list_runs": + runs_list, total = tracker.list_runs( + project_uuid=project_uuid, + status=status if status in {"running", "completed", "failed", "paused", "idle"} else None, + search=search.strip() or None, + limit=max(1, min(limit, 50)), + offset=0, + ) + + summaries = [] + for r in runs_list: + latest_val_loss = r.metrics.get("val_loss", [])[-1].value if r.metrics.get("val_loss") else None + summaries.append({ + "run_id": r.id, + "name": r.name, + "status": r.status, + "started_at": r.started_at, + "progress": f"{r.current_step}/{r.total_steps} steps", + "best_val_loss": r.best_val_loss, + "latest_val_loss": latest_val_loss, + "tags": r.tags, + }) + + result = { + "total_runs": total, + "returned_count": len(summaries), + "runs": summaries, + } + return json.dumps(result, indent=2), True + + elif operation == "compare_runs": + parsed_ids = _parse_list_or_csv(run_ids) + if not parsed_ids: + return "Error: 'run_ids' list or comma-separated string is required to compare runs.", False + + comparison = tracker.compare_runs(parsed_ids) + return json.dumps(comparison, indent=2), True + + elif operation == "complete_run": + if not run_id.strip(): + return "Error: 'run_id' is required.", False + + run = tracker.get_run(run_id.strip()) + if not run: + return f"Error: Experiment run '{run_id}' not found.", False + + final_status = "completed" if status not in {"completed", "failed"} else status + updated_run = tracker.update_status( + run_id=run.id, + status=final_status, + reason=reason.strip() or None, + ) + + if not updated_run: + return f"Failed to update status for run '{run_id}'.", False + + if best_val_loss is not None: + updated_run.best_val_loss = float(best_val_loss) + tracker._persist_run(updated_run) + + # Update entity in knowledge graph if active + try: + from .workspace_tools import _knowledge_var + + kg = _knowledge_var.get() + if kg: + kg.add_entity( + entity_id=f"exp_{updated_run.id}", + entity_type="experiment", + label=updated_run.name, + properties={ + "run_id": updated_run.id, + "status": updated_run.status, + "best_val_loss": updated_run.best_val_loss, + "duration_seconds": updated_run.duration_seconds, + }, + ) + kg.save() + except Exception as e: + log.debug("Knowledge graph entity update skipped: %s", e) + + result = { + "message": f"Run '{updated_run.name}' marked as {updated_run.status}.", + "run_id": updated_run.id, + "status": updated_run.status, + "best_val_loss": updated_run.best_val_loss, + "duration_seconds": updated_run.duration_seconds, + } + return json.dumps(result, indent=2), True + + else: + return f"Unknown experiments operation: '{operation}'. Valid operations: create_run, log_metrics, record_checkpoint, get_run, list_runs, compare_runs, complete_run.", False + + except Exception as e: + log.warning("Experiment tool error (%s): %s", operation, e) + return f"Experiment operation failed: {e}", False + + +def create_experiments_tool() -> ToolSpec: + """Create the ToolSpec for ML experiment tracking operations.""" + return ToolSpec( + name="experiments", + description=( + "Track, monitor, and compare machine learning experiment runs and training metrics.\n\n" + "Operations:\n" + "- create_run: Initialize a new tracked experiment run (requires 'name', optional 'hyperparameters', 'compute_target', 'tags', 'total_steps', 'total_epochs')\n" + "- log_metrics: Log scalar training/validation metrics at a step/epoch (requires 'run_id', 'metrics' dict or JSON e.g. {'train_loss': 1.8, 'val_loss': 1.9, 'lr': 0.001}, 'step')\n" + "- record_checkpoint: Register a model checkpoint artifact (requires 'run_id', optional 'checkpoint_name', 'path', 'file_size_mb', 'metrics')\n" + "- get_run: Retrieve details, latest metrics, and checkpoints for a run (requires 'run_id')\n" + "- list_runs: List all experiment runs in the current project (optional 'status', 'search', 'limit')\n" + "- compare_runs: Side-by-side comparison of hyperparameters and loss metrics across runs (requires 'run_ids' list/CSV)\n" + "- complete_run: Mark a run as completed or failed (requires 'run_id', optional 'status', 'best_val_loss', 'reason')" + ), + parameters={ + "type": "object", + "properties": { + "operation": { + "type": "string", + "enum": [ + "create_run", + "log_metrics", + "record_checkpoint", + "get_run", + "list_runs", + "compare_runs", + "complete_run", + ], + "description": "The experiment tracking operation to perform.", + }, + "name": { + "type": "string", + "description": "Descriptive name for the experiment run (for create_run).", + }, + "description": { + "type": "string", + "description": "Scientific hypothesis or objective of the experiment run.", + }, + "hyperparameters": { + "type": "object", + "description": "Hyperparameters dictionary or JSON string (e.g. {'lr': 1e-4, 'batch_size': 64}).", + }, + "compute_target": { + "type": "string", + "description": "Hardware/compute target (e.g. 'Local GPU', 'Modal A100').", + }, + "tags": { + "type": "array", + "items": {"type": "string"}, + "description": "Tags list or comma-separated string for categorization.", + }, + "total_steps": { + "type": "integer", + "description": "Expected total training steps.", + }, + "total_epochs": { + "type": "integer", + "description": "Expected total training epochs.", + }, + "run_id": { + "type": "string", + "description": "Experiment run ID (for log_metrics, record_checkpoint, get_run, complete_run).", + }, + "metrics": { + "type": "object", + "description": "Key-value metrics dictionary or JSON (e.g. {'train_loss': 0.45, 'val_loss': 0.52}).", + }, + "step": { + "type": "integer", + "description": "Current optimization/training step.", + }, + "epoch": { + "type": "integer", + "description": "Current training epoch.", + }, + "checkpoint_name": { + "type": "string", + "description": "Name of the checkpoint file or artifact.", + }, + "path": { + "type": "string", + "description": "Local or workspace file path to the checkpoint artifact.", + }, + "file_size_mb": { + "type": "number", + "description": "File size of the checkpoint in megabytes.", + }, + "status": { + "type": "string", + "enum": ["completed", "failed", "running", "paused", "idle"], + "description": "Run status for complete_run or filter for list_runs.", + }, + "reason": { + "type": "string", + "description": "Optional notes or failure diagnosis reason.", + }, + "best_val_loss": { + "type": "number", + "description": "Best validation loss achieved during the run.", + }, + "run_ids": { + "type": "array", + "items": {"type": "string"}, + "description": "List or CSV string of run IDs to compare (for compare_runs).", + }, + "search": { + "type": "string", + "description": "Search keyword for list_runs.", + }, + "limit": { + "type": "integer", + "description": "Max runs to return (for list_runs, default 10).", + }, + }, + "required": ["operation"], + }, + handler=_handle_experiments, + ) diff --git a/backend/openmlr/tools/figures.py b/backend/openmlr/tools/figures.py new file mode 100644 index 0000000..274bf41 --- /dev/null +++ b/backend/openmlr/tools/figures.py @@ -0,0 +1,207 @@ +"""Agent tool for Publication Figure Studio, Plot Generation, and LaTeX Subfigures.""" + +from __future__ import annotations + +import asyncio +import json +import logging +from collections.abc import Callable +from typing import Any + +from ..agent.types import ToolSpec +from ..services.figure_generator import FigureGeneratorService +from ..services.figure_types import ( + ColorPalette, + GenerateFigureRequest, + MultiPanelLayoutRequest, + PlotType, + StyleTheme, +) + +log = logging.getLogger("openmlr.tools.figures") + + +def _resolve_project_id(explicit_proj: str | None, getter: Callable[[], str | None] | None) -> str: + if explicit_proj and explicit_proj.strip(): + return explicit_proj.strip() + if getter: + val = getter() + if val and val.strip(): + return val.strip() + return "default" + + +def _parse_dict(val: Any) -> dict[str, Any]: + if isinstance(val, dict): + return val + if isinstance(val, str) and val.strip(): + try: + parsed = json.loads(val) + if isinstance(parsed, dict): + return parsed + except Exception: + pass + return {} + + +def _handle_generate(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + title = kwargs.get("title") + if not title: + return "Error: Field `title` is required for generating a figure.", False + + plot_type_str = kwargs.get("plot_type", "loss_curve") + style_theme_str = kwargs.get("style_theme", "neurips") + palette_str = kwargs.get("palette", "colorblind") + + series_data_raw = kwargs.get("series_data") or {} + series_data = _parse_dict(series_data_raw) + + req = GenerateFigureRequest( + title=title, + caption=kwargs.get("caption", ""), + plot_type=PlotType(plot_type_str) if plot_type_str in PlotType.__members__.values() else PlotType.LOSS_CURVE, + style_theme=StyleTheme(style_theme_str) if style_theme_str in StyleTheme.__members__.values() else StyleTheme.NEURIPS, + palette=ColorPalette(palette_str) if palette_str in ColorPalette.__members__.values() else ColorPalette.COLORBLIND, + x_label=kwargs.get("x_label", "Step"), + y_label=kwargs.get("y_label", "Loss"), + series_data=series_data, + categories=kwargs.get("categories") or [], + width_inches=float(kwargs.get("width_inches", 6.0)), + height_inches=float(kwargs.get("height_inches", 4.0)), + generate_tikz=bool(kwargs.get("generate_tikz", True)), + ) + artifact = FigureGeneratorService.generate_figure(proj, req) + msg = ( + f"✅ Publication Figure '{artifact.title}' generated successfully!\n" + f"- Figure ID: `{artifact.id}`\n" + f"- Plot Type: `{artifact.plot_type.value}`\n" + f"- Theme: `{artifact.style_theme.value}` | Palette: `{artifact.palette.value}`\n" + f"```latex\n{artifact.latex_snippet}\n```" + ) + return msg, True + + +def _handle_list(proj: str) -> tuple[str, bool]: + figures = FigureGeneratorService.list_figures(proj) + if not figures: + return f"No figure artifacts found in project `{proj}`.", True + lines = [f"Found {len(figures)} figure artifacts in project `{proj}`:"] + for f in figures: + lines.append(f"- **{f.title}** (`{f.id}`): {f.plot_type.value} | {f.style_theme.value}") + return "\n".join(lines), True + + +def _handle_get(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + fig_id = kwargs.get("figure_id") + if not fig_id: + return "Error: `figure_id` is required for get action.", False + fig = FigureGeneratorService.get_figure(proj, fig_id) + if not fig: + return f"Error: Figure `{fig_id}` not found in project `{proj}`.", False + msg = ( + f"### Figure Artifact: {fig.title}\n" + f"- **ID:** `{fig.id}`\n" + f"- **Plot Type:** {fig.plot_type.value}\n" + f"- **Theme:** {fig.style_theme.value} ({fig.palette.value})\n\n" + f"#### LaTeX Environment:\n```latex\n{fig.latex_snippet}\n```\n\n" + f"#### Standalone Python Script:\n```python\n{fig.python_script}\n```" + ) + return msg, True + + +def _handle_multipanel(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + fig_ids = kwargs.get("figure_ids") + if not fig_ids or len(fig_ids) < 2: + return "Error: `figure_ids` requires at least 2 figure IDs for a multi-panel layout.", False + + req = MultiPanelLayoutRequest( + title=kwargs.get("title", "Multi-Panel Benchmark Results"), + caption=kwargs.get("caption", "Ablation and comparison of empirical performance."), + figure_ids=fig_ids, + columns=int(kwargs.get("columns", 2)), + subcaptions=_parse_dict(kwargs.get("subcaptions")), + ) + result = FigureGeneratorService.create_multipanel_layout(proj, req) + if "error" in result: + return f"Error: {result['error']}", False + msg = ( + f"✅ Multi-Panel Subfigure Grid Created ({result['figure_count']} figures):\n" + f"```latex\n{result['latex_code']}\n```" + ) + return msg, True + + +def create_figures_tool(get_project_id: Callable[[], str | None] | None = None) -> ToolSpec: + """Create the 'figures' agent tool spec.""" + + async def _execute(action: str = "list", **kwargs: Any) -> tuple[str, bool]: + await asyncio.sleep(0) + proj = _resolve_project_id(kwargs.get("project_id"), get_project_id) + act = (action or "list").lower().strip() + + handlers = { + "generate": lambda: _handle_generate(proj, kwargs), + "list": lambda: _handle_list(proj), + "get": lambda: _handle_get(proj, kwargs), + "create_multipanel": lambda: _handle_multipanel(proj, kwargs), + } + + handler = handlers.get(act) + if not handler: + return ( + f"Unknown action: '{action}'. " + "Allowed actions: `generate`, `list`, `get`, `create_multipanel`.", + False, + ) + + try: + return handler() + except Exception as e: + log.exception("Figures tool error: %s", e) + return f"Error executing figures action '{action}': {e}", False + + return ToolSpec( + name="figures", + description=( + "Publication Figure Studio and Scientific Plot Generator. " + "Generate publication-quality vector plots (Loss curves, Ablation bars, Pareto frontiers, Confusion matrices, Heatmaps), " + "reproducible Matplotlib Python scripts, LaTeX subfigure grids, and pure TikZ vector code." + ), + parameters={ + "type": "object", + "properties": { + "action": { + "type": "string", + "enum": ["generate", "list", "get", "create_multipanel"], + "description": "Figure action to execute.", + }, + "project_id": {"type": "string", "description": "Optional project ID override."}, + "figure_id": {"type": "string", "description": "Figure Artifact ID."}, + "title": {"type": "string", "description": "Figure title."}, + "caption": {"type": "string", "description": "LaTeX caption."}, + "plot_type": { + "type": "string", + "enum": ["loss_curve", "ablation_bar", "pareto_frontier", "confusion_matrix", "radar_benchmark", "heatmap"], + "description": "Scientific plot type.", + }, + "style_theme": { + "type": "string", + "enum": ["neurips", "icml", "iclr", "cvpr", "dark"], + "description": "Conference theme preset.", + }, + "palette": { + "type": "string", + "enum": ["colorblind", "viridis", "muted", "deep", "tableau"], + "description": "Academic color palette.", + }, + "x_label": {"type": "string", "description": "X-axis label."}, + "y_label": {"type": "string", "description": "Y-axis label."}, + "series_data": {"type": "object", "description": "Mapping of series name to list of {x, y, y_err} data points."}, + "categories": {"type": "array", "items": {"type": "string"}, "description": "Category labels for bar/radar."}, + "figure_ids": {"type": "array", "items": {"type": "string"}, "description": "Figure IDs for multi-panel grid."}, + "columns": {"type": "integer", "description": "Number of columns in multi-panel grid."}, + }, + "required": ["action"], + }, + handler=_execute, + ) diff --git a/backend/openmlr/tools/models.py b/backend/openmlr/tools/models.py new file mode 100644 index 0000000..e3943e4 --- /dev/null +++ b/backend/openmlr/tools/models.py @@ -0,0 +1,277 @@ +"""Agent tool for Model Registry, Checkpoint Governance, Model Card generation, and Quantization.""" + +from __future__ import annotations + +import asyncio +import json +import logging +from collections.abc import Callable +from typing import Any + +from ..agent.types import ToolSpec +from ..services.model_registry import ModelRegistryService +from ..services.model_types import ( + GenerateModelCardRequest, + InspectCheckpointRequest, + RegisterModelRequest, +) + +log = logging.getLogger("openmlr.tools.models") + + +def _resolve_project_id(explicit_proj: str | None, getter: Callable[[], str | None] | None) -> str: + if explicit_proj and explicit_proj.strip(): + return explicit_proj.strip() + if getter: + val = getter() + if val and val.strip(): + return val.strip() + return "default" + + +def _parse_dict(val: Any) -> dict[str, Any]: + if isinstance(val, dict): + return val + if isinstance(val, str) and val.strip(): + try: + parsed = json.loads(val) + if isinstance(parsed, dict): + return parsed + except Exception: + pass + return {} + + +def _handle_register(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + name = kwargs.get("name") + if not name: + return "Error: Field `name` is required for registering a model.", False + + req = RegisterModelRequest( + name=name, + version=kwargs.get("version", "1.0.0"), + architecture=kwargs.get("architecture", "Transformer"), + framework=kwargs.get("framework", "pytorch"), # type: ignore + task_type=kwargs.get("task_type", "causal_lm"), # type: ignore + status=kwargs.get("status", "evaluated"), # type: ignore + description=kwargs.get("description", ""), + parameters_count=int(kwargs.get("parameters_count", 0)), + model_size_mb=float(kwargs.get("model_size_mb", 0.0)), + checkpoint_path=kwargs.get("checkpoint_path", ""), + base_model=kwargs.get("base_model", ""), + tags=kwargs.get("tags") or [], + metrics=_parse_dict(kwargs.get("metrics")), + hyperparameters=_parse_dict(kwargs.get("hyperparameters")), + lineage=_parse_dict(kwargs.get("lineage")), + ) + artifact = ModelRegistryService.register_model(proj, req) + msg = ( + f"✅ Model artifact '{artifact.name}' (v{artifact.version}) registered successfully!\n" + f"- Model ID: `{artifact.id}`\n" + f"- Architecture: `{artifact.architecture}` ({artifact.framework})\n" + f"- Parameters: {artifact.parameters_count:,}\n" + f"- Size: {artifact.model_size_mb:.2f} MB\n" + f"- Status: `{artifact.status}`\n" + f"- Metrics: {json.dumps(artifact.metrics)}" + ) + return msg, True + + +def _handle_list(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + models = ModelRegistryService.list_models( + project_id=proj, + task_type=kwargs.get("task_type"), + framework=kwargs.get("framework"), + status=kwargs.get("status"), + ) + if not models: + return f"No model artifacts registered in project `{proj}`.", True + lines = [f"Found {len(models)} model artifacts in project `{proj}`:"] + for m in models: + lines.append( + f"- **{m.name}** (v{m.version}, `{m.id}`): {m.architecture} | {m.parameters_count:,} params | {m.model_size_mb:.1f} MB | status: {m.status}" + ) + return "\n".join(lines), True + + +def _handle_get(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + model_id = kwargs.get("model_id") + if not model_id: + return "Error: `model_id` is required for get action.", False + m = ModelRegistryService.get_model(proj, model_id) + if not m: + return f"Error: Model artifact `{model_id}` not found in project `{proj}`.", False + msg = ( + f"### Model Artifact: {m.name} (v{m.version})\n" + f"- **ID:** `{m.id}`\n" + f"- **Architecture:** {m.architecture} ({m.framework})\n" + f"- **Task Type:** {m.task_type}\n" + f"- **Status:** {m.status}\n" + f"- **Parameters:** {m.parameters_count:,}\n" + f"- **Disk Size:** {m.model_size_mb:.2f} MB\n" + f"- **Checkpoint Path:** `{m.checkpoint_path or 'N/A'}`\n" + f"- **Base Model:** `{m.base_model or 'Trained from scratch'}`\n" + f"- **Metrics:** {json.dumps(m.metrics, indent=2)}\n" + f"- **Hyperparameters:** {json.dumps(m.hyperparameters, indent=2)}" + ) + return msg, True + + +def _handle_inspect_checkpoint(kwargs: dict[str, Any]) -> tuple[str, bool]: + req = InspectCheckpointRequest( + checkpoint_path=kwargs.get("checkpoint_path", ""), + parameters_count=int(kwargs.get("parameters_count", 0)), + model_size_mb=float(kwargs.get("model_size_mb", 0.0)), + framework=kwargs.get("framework", "pytorch"), + ) + insp = ModelRegistryService.inspect_checkpoint(req) + msg = ( + f"### Checkpoint Inspection Report\n" + f"- **Format:** `{insp.file_format}`\n" + f"- **Total Parameters:** {insp.total_parameters:,}\n" + f"- **Total Disk Size:** {insp.total_size_mb:.2f} MB\n" + f"- **Estimated VRAM (FP32):** {insp.estimated_vram_fp32_mb:.1f} MB\n" + f"- **Estimated VRAM (FP16/BF16):** {insp.estimated_vram_fp16_mb:.1f} MB\n" + f"- **Estimated VRAM (INT8):** {insp.estimated_vram_int8_mb:.1f} MB\n" + f"- **Estimated VRAM (INT4):** {insp.estimated_vram_int4_mb:.1f} MB\n" + f"- **Layers Count:** {insp.layers_count}" + ) + return msg, True + + +def _handle_generate_card(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + model_id = kwargs.get("model_id") + if not model_id: + return "Error: `model_id` is required for generate_card action.", False + card_req = GenerateModelCardRequest( + author=kwargs.get("author", "OpenMLR Research Agent"), + license=kwargs.get("license", "Apache-2.0"), + intended_use=kwargs.get("intended_use", ""), + limitations=kwargs.get("limitations", ""), + evaluation_notes=kwargs.get("evaluation_notes", ""), + gpu_type=kwargs.get("gpu_type", "NVIDIA A100-SXM4-80GB"), + gpu_hours=float(kwargs.get("gpu_hours", 24.0)), + ) + card = ModelRegistryService.generate_model_card(proj, model_id, card_req) + if not card: + return f"Error: Model artifact `{model_id}` not found in project `{proj}`.", False + return card.markdown, True + + +def _handle_plan_quantization(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + model_id = kwargs.get("model_id") + if not model_id: + return "Error: `model_id` is required for plan_quantization action.", False + model = ModelRegistryService.get_model(proj, model_id) + if not model: + return f"Error: Model artifact `{model_id}` not found in project `{proj}`.", False + targets = kwargs.get("target_precisions") or ["fp16", "bf16", "fp8", "int8", "int4"] + estimates = ModelRegistryService.plan_quantization(model, targets) + rows = [ + f"### Quantization Planning for {model.name} ({model.parameters_count:,} params)", + "| Precision | Est. Size | Est. VRAM | Compression | Speedup | Suggested Engine |", + "| :--- | :--- | :--- | :--- | :--- | :--- |", + ] + for e in estimates: + rows.append( + f"| **{e.target_precision}** | {e.estimated_size_mb:.1f} MB | {e.estimated_vram_mb:.1f} MB | {e.compression_ratio:.1f}x | {e.expected_latency_speedup:.1f}x | {e.suggested_engine} |" + ) + return "\n".join(rows), True + + +def _handle_compare(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + model_ids = kwargs.get("model_ids") + if not model_ids or len(model_ids) < 2: + return "Error: `model_ids` requires at least 2 model IDs to compare.", False + comp = ModelRegistryService.compare_models(proj, model_ids) + if "error" in comp: + return f"Error: {comp['error']}", False + lines = [ + "### Model Comparison Analysis", + f"- **Recommendation:** {comp['recommendation_reason']}", + f"- **Recommended ID:** `{comp['recommended_model_id']}`\n", + "**Parameters & Sizes:**", + ] + for m in comp["compared_models"]: + lines.append(f"- **{m['name']}** (v{m['version']}): {m['parameters_count']:,} params, {m['model_size_mb']:.1f} MB") + return "\n".join(lines), True + + +def create_models_tool(get_project_id: Callable[[], str | None] | None = None) -> ToolSpec: + """Create the 'models' agent tool spec.""" + + async def _execute(action: str = "list", **kwargs: Any) -> tuple[str, bool]: + await asyncio.sleep(0) + proj = _resolve_project_id(kwargs.get("project_id"), get_project_id) + act = (action or "list").lower().strip() + + handlers = { + "register": lambda: _handle_register(proj, kwargs), + "list": lambda: _handle_list(proj, kwargs), + "get": lambda: _handle_get(proj, kwargs), + "inspect_checkpoint": lambda: _handle_inspect_checkpoint(kwargs), + "generate_card": lambda: _handle_generate_card(proj, kwargs), + "plan_quantization": lambda: _handle_plan_quantization(proj, kwargs), + "compare": lambda: _handle_compare(proj, kwargs), + } + + handler = handlers.get(act) + if not handler: + return ( + f"Unknown action: '{action}'. " + "Allowed actions: `register`, `list`, `get`, `inspect_checkpoint`, `generate_card`, `plan_quantization`, `compare`.", + False, + ) + + try: + return handler() + except Exception as e: + log.exception("Models tool error: %s", e) + return f"Error executing models action '{action}': {e}", False + + return ToolSpec( + name="models", + description=( + "Model Registry, Checkpoint Governance, Model Card Generator, and Quantization Planning. " + "Manage trained model artifacts, generate NeurIPS/HuggingFace model cards, inspect checkpoints, " + "and compute precision compression tradeoffs." + ), + parameters={ + "type": "object", + "properties": { + "action": { + "type": "string", + "enum": ["register", "list", "get", "inspect_checkpoint", "generate_card", "plan_quantization", "compare"], + "description": "The model registry action to execute.", + }, + "project_id": {"type": "string", "description": "Optional project ID override."}, + "model_id": {"type": "string", "description": "Model Artifact ID."}, + "name": {"type": "string", "description": "Model name for registration."}, + "version": {"type": "string", "description": "Semantic version, e.g. 1.0.0."}, + "architecture": {"type": "string", "description": "Model architecture (e.g. Transformer, ResNet-50)."}, + "framework": {"type": "string", "description": "Framework (pytorch, safetensors, jax, onnx, gguf)."}, + "task_type": {"type": "string", "description": "Task type (causal_lm, classification, diffusion, etc.)."}, + "status": {"type": "string", "description": "Model status (draft, training, evaluated, production, archived)."}, + "description": {"type": "string", "description": "Detailed model description."}, + "parameters_count": {"type": "integer", "description": "Total parameter count."}, + "model_size_mb": {"type": "number", "description": "Model artifact size in megabytes."}, + "checkpoint_path": {"type": "string", "description": "Path to checkpoint file."}, + "base_model": {"type": "string", "description": "Base/pretrained model reference."}, + "tags": {"type": "array", "items": {"type": "string"}, "description": "Searchable tags."}, + "metrics": {"type": "object", "description": "Evaluation metrics."}, + "hyperparameters": {"type": "object", "description": "Training hyperparameters."}, + "lineage": {"type": "object", "description": "Model lineage provenance."}, + "target_precisions": {"type": "array", "items": {"type": "string"}, "description": "Quantization precisions to evaluate."}, + "model_ids": {"type": "array", "items": {"type": "string"}, "description": "List of model IDs to compare."}, + "gpu_type": {"type": "string", "description": "GPU used for training."}, + "gpu_hours": {"type": "number", "description": "GPU hours spent."}, + "author": {"type": "string", "description": "Model author/organization."}, + "license": {"type": "string", "description": "Model license."}, + "intended_use": {"type": "string", "description": "Intended application and domain."}, + "limitations": {"type": "string", "description": "Known limitations or bias."}, + "evaluation_notes": {"type": "string", "description": "Additional evaluation notes."}, + }, + "required": ["action"], + }, + handler=_execute, + ) diff --git a/backend/openmlr/tools/registry.py b/backend/openmlr/tools/registry.py index d04b217..fda26ae 100644 --- a/backend/openmlr/tools/registry.py +++ b/backend/openmlr/tools/registry.py @@ -47,6 +47,16 @@ "session_search", # Process management (read-only actions: list, poll, log) "process", + # Experiments tracking and monitoring (read-only actions in plan mode) + "experiments", + # Dataset profiling and inspection (read-only actions in plan mode) + "datasets", + # Hyperparameter optimization and sweeps (read-only actions in plan mode) + "sweeps", + # Model registry and checkpoint governance (read-only in plan mode) + "models", + # Publication figures and plotting (read-only in plan mode) + "figures", }, "blocked_message": ( "Tool '{tool}' is not available in PLAN mode. " @@ -421,15 +431,26 @@ def create_tool_router(sandbox_manager=None) -> ToolRouter: # Import and register all built-in tools from .ask_user import create_ask_user_tool + from .compute_tools import create_compute_tools + from .datasets import create_datasets_tool + from .experiments import create_experiments_tool + from .figures import create_figures_tool from .github import create_github_tools from .huggingface import create_huggingface_tools from .inspect import create_inspect_tool from .latex_compiler import create_latex_tool from .local import create_local_tools + from .memory_tool import create_memory_tool + from .models import create_models_tool from .papers import create_papers_tool from .plan import create_plan_tool + from .process_tool import create_process_tool + from .reproducibility import create_reproducibility_tool from .research import create_research_tool from .search import create_search_tools + from .session_search import create_session_search_tool + from .sweeps import create_sweeps_tool + from .workspace_tools import create_workspace_tools from .writing import create_writing_tool router.register_many(create_local_tools()) @@ -443,36 +464,20 @@ def create_tool_router(sandbox_manager=None) -> ToolRouter: router.register(create_writing_tool()) router.register(create_latex_tool()) router.register(create_ask_user_tool()) - - # Register session search tool - from .session_search import create_session_search_tool - router.register(create_session_search_tool()) - - # Register compute tools - from .compute_tools import create_compute_tools - router.register_many(create_compute_tools()) - - # Register workspace tools - from .workspace_tools import create_workspace_tools - router.register_many(create_workspace_tools()) - - # Register memory tool - from .memory_tool import create_memory_tool - router.register(create_memory_tool()) - - # Register process management tool - from .process_tool import create_process_tool - router.register(create_process_tool()) + router.register(create_experiments_tool()) + router.register(create_datasets_tool()) + router.register(create_sweeps_tool()) + router.register(create_models_tool()) + router.register(create_figures_tool()) + router.register(create_reproducibility_tool()) - # Register sandbox tools if manager provided if sandbox_manager: from .sandbox_tools import create_sandbox_tools - router.register_many(create_sandbox_tools(sandbox_manager)) return router diff --git a/backend/openmlr/tools/reproducibility.py b/backend/openmlr/tools/reproducibility.py new file mode 100644 index 0000000..7e65588 --- /dev/null +++ b/backend/openmlr/tools/reproducibility.py @@ -0,0 +1,190 @@ +"""Agent tool for Reproducibility Auditing, Determinism verification, and Conference Compliance.""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Callable +from typing import Any + +from ..agent.types import ToolSpec +from ..services.reproducibility_auditor import ReproducibilityAuditorService +from ..services.reproducibility_types import ( + AuditCodebaseRequest, + ChecklistVenue, + GenerateAppendixRequest, + GenerateDockerfileRequest, +) + +log = logging.getLogger("openmlr.tools.reproducibility") + + +def _resolve_project_id(explicit_proj: str | None, getter: Callable[[], str | None] | None) -> str: + if explicit_proj and explicit_proj.strip(): + return explicit_proj.strip() + if getter: + val = getter() + if val and val.strip(): + return val.strip() + return "default" + + +def _handle_audit(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + target_path = kwargs.get("target_path", ".") + venue_str = kwargs.get("venue", "neurips").lower() + try: + venue = ChecklistVenue(venue_str) + except ValueError: + venue = ChecklistVenue.NEURIPS + + snippets = kwargs.get("code_snippets") + req = AuditCodebaseRequest( + target_path=target_path, + venue=venue, + code_snippets=snippets if isinstance(snippets, dict) else None, + ) + report = ReproducibilityAuditorService.audit_codebase(req, proj) + result = { + "status": "success", + "report_id": report.id, + "overall_score": report.overall_score, + "grade": report.grade, + "venue": report.venue.value, + "categories": [c.model_dump() for c in report.categories], + "checklist_summary": { + "total": len(report.checklist), + "passed": sum(1 for i in report.checklist if i.status.value == "pass"), + "warnings": sum(1 for i in report.checklist if i.status.value == "warn"), + "failed": sum(1 for i in report.checklist if i.status.value == "fail"), + }, + "detected_frameworks": report.detected_frameworks, + "badge_markdown": report.badge_markdown, + "latex_appendix": ( + report.latex_appendix[:500] + "..." if len(report.latex_appendix) > 500 else report.latex_appendix + ), + } + return json.dumps(result, indent=2), True + + +def _handle_generate_dockerfile(kwargs: dict[str, Any]) -> tuple[str, bool]: + req_docker = GenerateDockerfileRequest( + framework=kwargs.get("framework", "pytorch"), + cuda_version=kwargs.get("cuda_version", "12.1.0"), + python_version=kwargs.get("python_version", "3.11"), + entrypoint_cmd=kwargs.get("entrypoint_cmd", "python train.py"), + requirements=kwargs.get("requirements", []), + ) + dockerfile = ReproducibilityAuditorService.generate_dockerfile(req_docker) + return json.dumps({"status": "success", "dockerfile": dockerfile}, indent=2), True + + +def _handle_generate_appendix(proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + report_id = kwargs.get("report_id") + report = ReproducibilityAuditorService.get_report(report_id, proj) if report_id else None + req_app = GenerateAppendixRequest( + report_id=report_id, + paper_title=kwargs.get("paper_title", "Reproducible ML Study"), + hardware_specs=kwargs.get("hardware_specs", "NVIDIA A100-SXM4-80GB (1 GPU)"), + random_seeds=kwargs.get("random_seeds", [42, 1337, 2026]), + dataset_url=kwargs.get("dataset_url", "https://huggingface.co/datasets"), + code_url=kwargs.get("code_url", "https://github.com"), + ) + appendix = ReproducibilityAuditorService.generate_latex_appendix(req_app, report) + return json.dumps({"status": "success", "latex_appendix": appendix}, indent=2), True + + +def _handle_fix_determinism(kwargs: dict[str, Any]) -> tuple[str, bool]: + framework = kwargs.get("framework", "pytorch") + seed = int(kwargs.get("seed", 42)) + strict = bool(kwargs.get("strict_mode", True)) + snippet = ReproducibilityAuditorService.generate_determinism_snippet(framework, seed, strict) + return json.dumps({"status": "success", "determinism_snippet": snippet}, indent=2), True + + +def _handle_list(proj: str) -> tuple[str, bool]: + reports = ReproducibilityAuditorService.list_reports(proj) + return ( + json.dumps( + { + "status": "success", + "count": len(reports), + "reports": [ + { + "id": r.id, + "created_at": r.created_at, + "overall_score": r.overall_score, + "grade": r.grade, + "venue": r.venue.value, + } + for r in reports + ], + }, + indent=2, + ), + True, + ) + + +def create_reproducibility_tool( + get_project_context: Callable[[], str | None] | None = None, +) -> ToolSpec: + """Create the reproducibility agent tool.""" + + async def _execute(**kwargs: Any) -> tuple[str, bool]: + action = kwargs.get("action", "audit") + proj = _resolve_project_id(kwargs.get("project_id"), get_project_context) + + try: + if action == "audit": + return _handle_audit(proj, kwargs) + elif action == "generate_dockerfile": + return _handle_generate_dockerfile(kwargs) + elif action == "generate_appendix": + return _handle_generate_appendix(proj, kwargs) + elif action == "fix_determinism": + return _handle_fix_determinism(kwargs) + elif action == "list_reports" or action == "list": + return _handle_list(proj) + return f"Unknown action: '{action}'. Allowed: audit, generate_dockerfile, generate_appendix, fix_determinism, list_reports.", False + except Exception as e: + log.exception("Reproducibility tool error: %s", e) + return f"Error executing reproducibility action '{action}': {e}", False + + return ToolSpec( + name="reproducibility", + description=( + "Reproducibility Auditor & Artifact Governance Tool. " + "Audit ML codebases for scientific determinism, pinned dependencies, hardware requirements, " + "and conference checklist compliance (NeurIPS / ICML / ICLR / CVPR). " + "Actions: `audit`, `generate_dockerfile`, `generate_appendix`, `fix_determinism`, `list_reports`." + ), + parameters={ + "type": "object", + "properties": { + "action": { + "type": "string", + "enum": ["audit", "generate_dockerfile", "generate_appendix", "fix_determinism", "list_reports"], + "description": "Action to perform.", + }, + "project_id": {"type": "string", "description": "Optional project ID override."}, + "target_path": {"type": "string", "description": "Path to codebase or files to audit."}, + "venue": { + "type": "string", + "enum": ["neurips", "icml", "iclr", "cvpr", "general"], + "description": "Target conference rubric.", + }, + "code_snippets": { + "type": "object", + "description": "Optional in-memory mapping of filename to code strings to audit directly.", + }, + "framework": {"type": "string", "description": "Target ML framework: pytorch, jax, tensorflow."}, + "seed": {"type": "integer", "description": "Random seed integer for determinism fixes."}, + "strict_mode": {"type": "boolean", "description": "Enable strict deterministic algorithms."}, + "paper_title": {"type": "string", "description": "Title of research paper for appendix."}, + "hardware_specs": {"type": "string", "description": "Hardware specs for LaTeX statement."}, + "requirements": {"type": "array", "items": {"type": "string"}, "description": "Python packages for Dockerfile."}, + }, + "required": ["action"], + }, + handler=_execute, + ) diff --git a/backend/openmlr/tools/sweeps.py b/backend/openmlr/tools/sweeps.py new file mode 100644 index 0000000..5ea2288 --- /dev/null +++ b/backend/openmlr/tools/sweeps.py @@ -0,0 +1,298 @@ +"""Sweeps tool — Hyperparameter optimization, search spaces, and trial tuning for AI research agents.""" + +from __future__ import annotations + +import asyncio +import json +import logging +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from ..agent.types import ToolSpec +from ..services.sweep_engine import SweepEngine + +log = logging.getLogger(__name__) + +ERR_SWEEP_ID_REQUIRED = "Error: `sweep_id` is required." + + +def _parse_dict(val: Any) -> dict[str, Any]: + """Parse dict or JSON string into dictionary.""" + if isinstance(val, dict): + return val + if isinstance(val, str) and val.strip(): + try: + parsed = json.loads(val) + if isinstance(parsed, dict): + return parsed + except Exception: + pass + return {} + + +def _handle_create_sweep(engine: SweepEngine, proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + name = kwargs.get("name") + if not name: + return "Error: `name` is required when creating a sweep.", False + param_dict = _parse_dict(kwargs.get("parameters")) + if not param_dict: + return "Error: `parameters` search space dictionary is required.", False + + sweep = engine.create_sweep( + project_id=proj, + name=name, + method=kwargs.get("method", "random"), + objective_metric=kwargs.get("objective_metric", "val_loss"), + goal=kwargs.get("goal", "minimize"), + parameters=param_dict, + max_trials=int(kwargs.get("max_trials", 10)), + description=kwargs.get("description", ""), + early_stopping=_parse_dict(kwargs.get("early_stopping")), + ) + msg = ( + f"✅ Hyperparameter sweep '{sweep.name}' created successfully!\n" + f"- Sweep ID: `{sweep.sweep_id}`\n" + f"- Method: `{sweep.method}`\n" + f"- Objective: `{sweep.objective_metric}` ({sweep.goal})\n" + f"- Max Trials: {sweep.max_trials}\n" + f"- Search Space: {list(sweep.parameters.keys())}" + ) + return msg, True + + +def _handle_list_sweeps(engine: SweepEngine, proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + sweeps = engine.list_sweeps(proj) + if not sweeps: + return f"No hyperparameter sweeps found for project `{proj}`.", True + lines = [f"### Hyperparameter Sweeps for `{proj}` ({len(sweeps)})", ""] + for s in sweeps: + completed = len([t for t in s.trials if t.status == "completed"]) + lines.append( + f"- **{s.name}** (`{s.sweep_id}`): {s.method.upper()}, " + f"Target: `{s.objective_metric}` ({s.goal}), " + f"Trials: {completed}/{len(s.trials)} (Max: {s.max_trials}), " + f"Status: `{s.status}`" + ) + return "\n".join(lines), True + + +def _handle_get_sweep(engine: SweepEngine, proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + sweep_id = kwargs.get("sweep_id") + if not sweep_id: + return ERR_SWEEP_ID_REQUIRED, False + sweep = engine.get_sweep(proj, sweep_id) + if not sweep: + return f"Error: Sweep `{sweep_id}` not found.", False + return json.dumps(sweep.to_dict(), indent=2), True + + +def _handle_suggest_trial(engine: SweepEngine, proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + sweep_id = kwargs.get("sweep_id") + if not sweep_id: + return ERR_SWEEP_ID_REQUIRED, False + trial = engine.suggest_trial(proj, sweep_id) + if not trial: + return f"Sweep `{sweep_id}` has reached its maximum trial limit or is complete.", True + msg = ( + f"🎯 Suggested Next Trial: `{trial.trial_id}` (Trial #{trial.trial_number})\n" + f"**Hyperparameter Configuration**:\n```json\n" + f"{json.dumps(trial.parameters, indent=2)}\n```\n" + f"Status: `{trial.status}`" + ) + return msg, True + + +def _handle_record_trial(engine: SweepEngine, proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + sweep_id = kwargs.get("sweep_id") + trial_id = kwargs.get("trial_id") + if not sweep_id or not trial_id: + return "Error: Both `sweep_id` and `trial_id` are required.", False + + trial = engine.record_trial_result( + project_id=proj, + sweep_id=sweep_id, + trial_id=trial_id, + metrics=_parse_dict(kwargs.get("metrics")), + status=kwargs.get("status", "completed"), + ) + msg = ( + f"✅ Trial `{trial.trial_id}` recorded with status `{trial.status}`.\n" + f"- Objective Value: `{trial.objective_value}`\n" + f"- Runtime: {trial.duration_seconds}s\n" + f"- Metrics: {json.dumps(trial.metrics)}" + ) + return msg, True + + +def _handle_prune_check(engine: SweepEngine, proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + sweep_id = kwargs.get("sweep_id") + trial_id = kwargs.get("trial_id") + if not sweep_id or not trial_id: + return "Error: Both `sweep_id` and `trial_id` are required.", False + + step = int(kwargs.get("current_step", 1)) + val = float(kwargs.get("current_metric_val", 0.0)) + should_stop = engine.should_prune_trial( + project_id=proj, + sweep_id=sweep_id, + trial_id=trial_id, + current_step=step, + current_metric_val=val, + ) + verdict = "PRUNE / STOP EARLY" if should_stop else "CONTINUE" + return f"Prune evaluation for trial `{trial_id}` at step {step} (value={val}): **{verdict}**", True + + +def _handle_analyze_sweep(engine: SweepEngine, proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + sweep_id = kwargs.get("sweep_id") + if not sweep_id: + return ERR_SWEEP_ID_REQUIRED, False + return json.dumps(engine.analyze_sweep(proj, sweep_id), indent=2), True + + +def _handle_export_report(engine: SweepEngine, proj: str, kwargs: dict[str, Any]) -> tuple[str, bool]: + sweep_id = kwargs.get("sweep_id") + if not sweep_id: + return ERR_SWEEP_ID_REQUIRED, False + return engine.export_sweep_markdown(proj, sweep_id), True + + +ACTION_DISPATCH: dict[str, Callable[[SweepEngine, str, dict[str, Any]], tuple[str, bool]]] = { + "create": _handle_create_sweep, + "create_sweep": _handle_create_sweep, + "list": _handle_list_sweeps, + "list_sweeps": _handle_list_sweeps, + "get": _handle_get_sweep, + "get_sweep": _handle_get_sweep, + "suggest": _handle_suggest_trial, + "suggest_trial": _handle_suggest_trial, + "next_trial": _handle_suggest_trial, + "record": _handle_record_trial, + "record_trial": _handle_record_trial, + "log_trial": _handle_record_trial, + "prune_check": _handle_prune_check, + "should_prune": _handle_prune_check, + "analyze": _handle_analyze_sweep, + "analyze_sweep": _handle_analyze_sweep, + "export": _handle_export_report, + "export_report": _handle_export_report, +} + + +def _resolve_project_id(project_id: str | None, get_project_id: Callable[[], str | None] | None) -> str: + if project_id and project_id.strip(): + return project_id.strip() + if get_project_id: + pid = get_project_id() + if pid and pid.strip(): + return pid.strip() + return "default" + + +def create_sweeps_tool( + get_project_id: Callable[[], str | None] | None = None, + base_dir: Path | None = None, +) -> ToolSpec: + """Create the hyperparameter sweep and HPO tool for OpenMLR agent.""" + engine = SweepEngine(base_dir=base_dir) + + async def _execute(action: str = "list_sweeps", **kwargs: Any) -> tuple[str, bool]: + await asyncio.sleep(0) # Async boundary + proj = _resolve_project_id(kwargs.get("project_id"), get_project_id) + act = (action or "list_sweeps").lower().strip() + handler = ACTION_DISPATCH.get(act) + + if not handler: + return ( + f"Unknown action: '{action}'. " + "Allowed actions: `create_sweep`, `list_sweeps`, `get_sweep`, " + "`suggest_trial`, `record_trial`, `prune_check`, `analyze_sweep`, `export_report`.", + False, + ) + + try: + return handler(engine, proj, kwargs) + except Exception as e: + log.exception("Sweeps tool error: %s", e) + return f"Error executing sweeps action '{action}': {e}", False + + return ToolSpec( + name="sweeps", + description=( + "Hyperparameter Optimization (HPO) and Sweep Manager — create parameter search spaces, " + "suggest trial configurations (Grid, Random, Bayesian, Hyperband), early-prune unpromising runs, " + "analyze parameter sensitivity & correlations, and export optimization reports." + ), + parameters={ + "type": "object", + "properties": { + "action": { + "type": "string", + "enum": [ + "create_sweep", + "list_sweeps", + "get_sweep", + "suggest_trial", + "record_trial", + "prune_check", + "analyze_sweep", + "export_report", + ], + "description": "The sweep management action to execute.", + }, + "sweep_id": { + "type": "string", + "description": "Identifier of the target sweep.", + }, + "name": { + "type": "string", + "description": "Descriptive name for the sweep.", + }, + "method": { + "type": "string", + "enum": ["grid", "random", "bayesian", "hyperband"], + "description": "Search method algorithm.", + }, + "objective_metric": { + "type": "string", + "description": "Metric name to optimize (e.g. 'val_loss', 'accuracy', 'f1').", + }, + "goal": { + "type": "string", + "enum": ["minimize", "maximize"], + "description": "Whether to minimize or maximize the objective metric.", + }, + "parameters": { + "type": "object", + "description": "Dictionary defining parameter spaces (e.g. {'lr': {'param_type': 'loguniform', 'min_val': 1e-5, 'max_val': 1e-2}}).", + }, + "max_trials": { + "type": "integer", + "description": "Maximum number of trials to run.", + }, + "trial_id": { + "type": "string", + "description": "Trial identifier for record or prune operations.", + }, + "metrics": { + "type": "object", + "description": "Dictionary of trial metrics (e.g. {'val_loss': 0.24, 'accuracy': 0.93}).", + }, + "current_step": { + "type": "integer", + "description": "Current step or epoch for early pruning checks.", + }, + "current_metric_val": { + "type": "number", + "description": "Current metric value for early pruning checks.", + }, + "project_id": { + "type": "string", + "description": "Optional project ID override.", + }, + }, + "required": ["action"], + }, + handler=_execute, + ) diff --git a/backend/tests/test_dataset_profiler.py b/backend/tests/test_dataset_profiler.py new file mode 100644 index 0000000..bf0bd39 --- /dev/null +++ b/backend/tests/test_dataset_profiler.py @@ -0,0 +1,140 @@ +"""Unit tests for DatasetProfiler (profiling, statistics, validation, splitting).""" + +from __future__ import annotations + +import csv +import json +from pathlib import Path + +import pytest + +from openmlr.services.dataset_profiler import DatasetProfiler + + +@pytest.fixture +def sample_csv(tmp_path: Path) -> Path: + csv_file = tmp_path / "data.csv" + data = [ + {"id": "1", "age": "25", "income": "50000", "label": "A", "text": "Short sample text", "active": "true"}, + {"id": "2", "age": "30", "income": "60000", "label": "A", "text": "Another sentence for ML training.", "active": "true"}, + {"id": "3", "age": "45", "income": "90000", "label": "B", "text": "Natural language processing benchmarks.", "active": "false"}, + {"id": "4", "age": "35", "income": "75000", "label": "A", "text": "Transformer self-attention mechanisms.", "active": "true"}, + {"id": "5", "age": "", "income": "80000", "label": "B", "text": "", "active": "false"}, + ] + with open(csv_file, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=list(data[0].keys())) + writer.writeheader() + writer.writerows(data) + return csv_file + + +@pytest.fixture +def sample_jsonl(tmp_path: Path) -> Path: + jsonl_file = tmp_path / "data.jsonl" + data = [ + {"prompt": "What is attention?", "response": "Attention is all you need.", "rating": 5}, + {"prompt": "Explain gradient descent.", "response": "Optimization algorithm.", "rating": 4}, + {"prompt": "Define cross-entropy.", "response": "Loss function for classification.", "rating": 5}, + ] + with open(jsonl_file, "w", encoding="utf-8") as f: + for r in data: + f.write(json.dumps(r) + "\n") + return jsonl_file + + +def test_detect_format(): + assert DatasetProfiler.detect_format("data.csv") == "csv" + assert DatasetProfiler.detect_format("data.tsv") == "tsv" + assert DatasetProfiler.detect_format("data.jsonl") == "jsonl" + assert DatasetProfiler.detect_format("data.ndjson") == "jsonl" + assert DatasetProfiler.detect_format("data.json") == "json" + assert DatasetProfiler.detect_format("data.txt") == "text" + + +def test_load_records_csv(sample_csv: Path): + records, fmt, size = DatasetProfiler.load_records(sample_csv) + assert fmt == "csv" + assert len(records) == 5 + assert size > 0 + assert records[0]["id"] == "1" + + +def test_load_records_jsonl(sample_jsonl: Path): + records, fmt, size = DatasetProfiler.load_records(sample_jsonl) + assert fmt == "jsonl" + assert len(records) == 3 + assert records[0]["rating"] == 5 + + +def test_load_records_empty_or_missing(tmp_path: Path): + with pytest.raises(FileNotFoundError): + DatasetProfiler.load_records(tmp_path / "missing.csv") + + empty_file = tmp_path / "empty.csv" + empty_file.touch() + records, fmt, size = DatasetProfiler.load_records(empty_file) + assert len(records) == 0 + + +def test_profile_csv(sample_csv: Path): + prof = DatasetProfiler.profile(sample_csv) + assert prof.total_rows == 5 + assert prof.total_columns == 6 + assert prof.health_score > 50 + assert "income" in prof.columns + assert prof.columns["income"].dtype == "numeric" + assert prof.columns["income"].stats["min"] == 50000.0 + assert prof.columns["income"].stats["max"] == 90000.0 + assert prof.columns["active"].dtype == "boolean" + assert prof.columns["label"].dtype == "categorical" + + +def test_sample_records_strategies(sample_csv: Path): + head_samples = DatasetProfiler.sample_records(sample_csv, n=2, strategy="head") + assert len(head_samples) == 2 + assert head_samples[0]["id"] == "1" + + random_samples = DatasetProfiler.sample_records(sample_csv, n=3, strategy="random", seed=123) + assert len(random_samples) == 3 + + stratified_samples = DatasetProfiler.sample_records( + sample_csv, n=4, strategy="stratified", label_column="label", seed=123 + ) + assert len(stratified_samples) <= 4 + + +def test_validate_dataset_pass(sample_csv: Path): + res = DatasetProfiler.validate_dataset( + sample_csv, + expected_columns=["id", "age", "income", "label"], + max_null_pct=50.0, + ) + assert res["valid"] is True + assert len(res["errors"]) == 0 + + +def test_validate_dataset_fail_missing_col(sample_csv: Path): + res = DatasetProfiler.validate_dataset( + sample_csv, + expected_columns=["id", "nonexistent_col"], + ) + assert res["valid"] is False + assert any("nonexistent_col" in err for err in res["errors"]) + + +def test_split_dataset(sample_csv: Path, tmp_path: Path): + out_dir = tmp_path / "splits" + manifest = DatasetProfiler.split_dataset( + sample_csv, + output_dir=out_dir, + train_ratio=0.6, + val_ratio=0.2, + test_ratio=0.2, + seed=42, + ) + assert manifest["total_records"] == 5 + assert manifest["train_count"] + manifest["val_count"] + manifest["test_count"] == 5 + assert Path(manifest["splits"]["train"]).exists() + assert Path(manifest["splits"]["val"]).exists() + assert Path(manifest["splits"]["test"]).exists() + assert (out_dir / "split_manifest.json").exists() diff --git a/backend/tests/test_experiments_tracker.py b/backend/tests/test_experiments_tracker.py new file mode 100644 index 0000000..26df89a --- /dev/null +++ b/backend/tests/test_experiments_tracker.py @@ -0,0 +1,141 @@ +"""Unit tests for the ExperimentTracker service.""" + +from pathlib import Path + +from openmlr.services.experiment_tracker import ExperimentTracker + + +def test_create_and_get_run(tmp_path: Path): + tracker = ExperimentTracker(storage_dir=tmp_path) + run = tracker.create_run( + name="Attention-Layer-Sweep", + description="Testing multi-head vs rotary embeddings", + hyperparameters={"lr": 0.001, "batch_size": 32, "layers": 12}, + compute_target="Local H100", + tags=["nlp", "attention"], + total_steps=500, + total_epochs=10, + ) + + assert run.id.startswith("run-") + assert run.name == "Attention-Layer-Sweep" + assert run.status == "running" + assert run.total_steps == 500 + assert run.hyperparameters["lr"] == 0.001 + + fetched = tracker.get_run(run.id) + assert fetched is not None + assert fetched.name == "Attention-Layer-Sweep" + assert fetched.tags == ["nlp", "attention"] + + +def test_list_runs_with_filtering(tmp_path: Path): + tracker = ExperimentTracker(storage_dir=tmp_path) + run1 = tracker.create_run(name="Run Alpha", tags=["vision"], project_uuid="proj-123") + run2 = tracker.create_run(name="Run Beta", tags=["nlp"], project_uuid="proj-123") + run3 = tracker.create_run(name="Run Gamma", tags=["rl"], project_uuid="proj-999") + + tracker.update_status(run1.id, "completed") + tracker.update_status(run2.id, "running") + + # List all for proj-123 + runs, total = tracker.list_runs(project_uuid="proj-123") + assert total == 2 + assert len(runs) == 2 + + # Filter by status + completed_runs, count = tracker.list_runs(project_uuid="proj-123", status="completed") + assert count == 1 + assert completed_runs[0].id == run1.id + + # Filter by search + searched, s_count = tracker.list_runs(search="beta") + assert s_count == 1 + assert searched[0].name == "Run Beta" + + +def test_log_metrics_and_best_val_loss(tmp_path: Path): + tracker = ExperimentTracker(storage_dir=tmp_path) + run = tracker.create_run(name="Metric Run", total_steps=100) + + tracker.log_metrics(run.id, step=10, epoch=1, metrics={"train_loss": 2.5, "val_loss": 2.8}) + tracker.log_metrics(run.id, step=20, epoch=1, metrics={"train_loss": 2.0, "val_loss": 2.3}) + tracker.log_metrics(run.id, step=30, epoch=2, metrics={"train_loss": 1.7, "val_loss": 2.4}) + + updated = tracker.get_run(run.id) + assert updated is not None + assert updated.current_step == 30 + assert updated.current_epoch == 2 + assert updated.best_val_loss == 2.3 + assert len(updated.metrics["train_loss"]) == 3 + assert len(updated.metrics["val_loss"]) == 3 + assert updated.metrics["train_loss"][-1].value == 1.7 + + +def test_register_checkpoint_and_logs(tmp_path: Path): + tracker = ExperimentTracker(storage_dir=tmp_path) + run = tracker.create_run(name="Checkpoint Run") + + tracker.append_logs(run.id, ["Epoch 1 starting", "Step 50 reached: loss=1.8"]) + cp = tracker.register_checkpoint( + run_id=run.id, + name="step_50.pt", + step=50, + epoch=1, + path="/models/step_50.pt", + file_size_mb=450.5, + metrics={"val_loss": 1.8}, + ) + + assert cp.name == "step_50.pt" + assert cp.file_size_mb == 450.5 + + updated = tracker.get_run(run.id) + assert updated is not None + assert len(updated.checkpoints) == 1 + assert len(updated.logs) == 2 + assert "Step 50 reached" in updated.logs[1] + + +def test_compare_runs(tmp_path: Path): + tracker = ExperimentTracker(storage_dir=tmp_path) + r1 = tracker.create_run(name="R1", hyperparameters={"lr": 0.01, "opt": "adam"}) + r2 = tracker.create_run(name="R2", hyperparameters={"lr": 0.001, "opt": "sgd", "momentum": 0.9}) + + tracker.log_metrics(r1.id, step=10, metrics={"train_loss": 1.5, "val_loss": 1.8}) + tracker.log_metrics(r2.id, step=10, metrics={"train_loss": 1.2, "val_loss": 1.4}) + + comp = tracker.compare_runs([r1.id, r2.id]) + assert len(comp["runs"]) == 2 + assert "lr" in comp["hyperparameters_comparison"] + assert comp["hyperparameters_comparison"]["lr"][r1.id] == 0.01 + assert comp["hyperparameters_comparison"]["lr"][r2.id] == 0.001 + assert comp["metrics_summary"][r1.id]["best_val_loss"] == 1.8 + assert comp["metrics_summary"][r2.id]["best_val_loss"] == 1.4 + + +def test_persistence_and_reload(tmp_path: Path): + tracker1 = ExperimentTracker(storage_dir=tmp_path) + run = tracker1.create_run(name="Persistent Run", hyperparameters={"batch_size": 64}) + tracker1.log_metrics(run.id, step=5, metrics={"train_loss": 3.1}) + + # Instantiate a new tracker pointing to the same directory + tracker2 = ExperimentTracker(storage_dir=tmp_path) + reloaded = tracker2.get_run(run.id) + assert reloaded is not None + assert reloaded.name == "Persistent Run" + assert reloaded.hyperparameters["batch_size"] == 64 + assert len(reloaded.metrics["train_loss"]) == 1 + + +def test_delete_run(tmp_path: Path): + tracker = ExperimentTracker(storage_dir=tmp_path) + run = tracker.create_run(name="To Delete") + assert tracker.get_run(run.id) is not None + + deleted = tracker.delete_run(run.id) + assert deleted is True + assert tracker.get_run(run.id) is None + + # Deleting again returns False + assert tracker.delete_run(run.id) is False diff --git a/backend/tests/test_figure_generator.py b/backend/tests/test_figure_generator.py new file mode 100644 index 0000000..18e0631 --- /dev/null +++ b/backend/tests/test_figure_generator.py @@ -0,0 +1,93 @@ +"""Unit tests for the Publication Figure Generator service.""" + +from openmlr.services.figure_generator import FigureGeneratorService +from openmlr.services.figure_types import ( + ColorPalette, + GenerateFigureRequest, + MultiPanelLayoutRequest, + PlotType, + StyleTheme, +) + + +def test_generate_loss_curve_figure(): + req = GenerateFigureRequest( + title="Training Loss Comparison", + caption="Training loss across 10k steps for baseline vs. ours.", + plot_type=PlotType.LOSS_CURVE, + style_theme=StyleTheme.NEURIPS, + palette=ColorPalette.COLORBLIND, + x_label="Step", + y_label="Loss", + series_data={ + "Baseline": [{"x": 100, "y": 2.5, "y_err": 0.1}, {"x": 200, "y": 1.8, "y_err": 0.08}], + "Ours (MoE)": [{"x": 100, "y": 2.1, "y_err": 0.05}, {"x": 200, "y": 1.2, "y_err": 0.04}], + }, + generate_tikz=True, + ) + artifact = FigureGeneratorService.generate_figure("proj_test", req) + assert artifact.id.startswith("fig_") + assert artifact.title == "Training Loss Comparison" + assert "plt.subplots" in artifact.python_script + assert "\\begin{figure}" in artifact.latex_snippet + assert "\\begin{tikzpicture}" in artifact.tikz_code + assert "= 1 + found = FigureGeneratorService.get_figure("proj_crud", art.id) + assert found is not None + assert found.title == "Figure A" + + success = FigureGeneratorService.delete_figure("proj_crud", art.id) + assert success is True + assert FigureGeneratorService.get_figure("proj_crud", art.id) is None + + +def test_create_multipanel_layout(): + req1 = GenerateFigureRequest( + title="Loss Convergence", + series_data={"S": [{"x": 1, "y": 1}]}, + ) + req2 = GenerateFigureRequest( + title="Throughput vs. Memory", + series_data={"S": [{"x": 1, "y": 1}]}, + ) + f1 = FigureGeneratorService.generate_figure("proj_multi", req1) + f2 = FigureGeneratorService.generate_figure("proj_multi", req2) + + multi_req = MultiPanelLayoutRequest( + title="Overall Benchmark Summary", + caption="Ablation overview.", + figure_ids=[f1.id, f2.id], + columns=2, + subcaptions={f1.id: "(a) Loss", f2.id: "(b) Throughput"}, + ) + res = FigureGeneratorService.create_multipanel_layout("proj_multi", multi_req) + assert res["figure_count"] == 2 + assert "\\begin{figure*}" in res["latex_code"] + assert "\\begin{subfigure}" in res["latex_code"] diff --git a/backend/tests/test_model_registry.py b/backend/tests/test_model_registry.py new file mode 100644 index 0000000..a2b3a57 --- /dev/null +++ b/backend/tests/test_model_registry.py @@ -0,0 +1,183 @@ +"""Tests for Model Registry Service, Model Card Generator, Checkpoint Inspection, and Quantization.""" + +import pytest + +from openmlr.services.model_card_generator import ( + estimate_carbon_footprint, +) +from openmlr.services.model_registry import ModelRegistryService +from openmlr.services.model_types import ( + GenerateModelCardRequest, + InspectCheckpointRequest, + RegisterModelRequest, + UpdateModelRequest, +) + + +@pytest.fixture(autouse=True) +def clean_registry(): + ModelRegistryService._models_store.clear() + yield + ModelRegistryService._models_store.clear() + + +def test_register_and_get_model(): + req = RegisterModelRequest( + name="Llama-3-8B-OpenMLR", + version="1.0.0", + architecture="LLaMA-3", + framework="safetensors", + task_type="causal_lm", + status="evaluated", + description="Fine-tuned research model for theorem proving.", + parameters_count=8_000_000_000, + model_size_mb=16000.0, + checkpoint_path="/checkpoints/llama3_8b.safetensors", + tags=["research", "reasoning", "llm"], + metrics={"val_loss": 1.12, "gsm8k_acc": 0.78}, + hyperparameters={"lr": 2e-5, "batch_size": 64, "warmup_steps": 100}, + ) + artifact = ModelRegistryService.register_model("proj_1", req) + assert artifact.id.startswith("model_") + assert artifact.name == "Llama-3-8B-OpenMLR" + assert artifact.parameters_count == 8_000_000_000 + assert artifact.metrics["gsm8k_acc"] == 0.78 + + retrieved = ModelRegistryService.get_model("proj_1", artifact.id) + assert retrieved is not None + assert retrieved.id == artifact.id + assert retrieved.architecture == "LLaMA-3" + + +def test_list_and_filter_models(): + req1 = RegisterModelRequest(name="Model-A", framework="pytorch", task_type="classification", tags=["vision"]) + req2 = RegisterModelRequest(name="Model-B", framework="safetensors", task_type="causal_lm", tags=["nlp"]) + req3 = RegisterModelRequest(name="Model-C", framework="onnx", task_type="classification", tags=["vision"]) + + ModelRegistryService.register_model("proj_filter", req1) + ModelRegistryService.register_model("proj_filter", req2) + ModelRegistryService.register_model("proj_filter", req3) + + all_models = ModelRegistryService.list_models("proj_filter") + assert len(all_models) == 3 + + vision_models = ModelRegistryService.list_models("proj_filter", tag="vision") + assert len(vision_models) == 2 + + nlp_models = ModelRegistryService.list_models("proj_filter", task_type="causal_lm") + assert len(nlp_models) == 1 + assert nlp_models[0].name == "Model-B" + + +def test_update_and_delete_model(): + req = RegisterModelRequest(name="Model-To-Update", status="training") + artifact = ModelRegistryService.register_model("proj_update", req) + + update_req = UpdateModelRequest( + status="production", + metrics={"accuracy": 0.95}, + description="Production candidate", + ) + updated = ModelRegistryService.update_model("proj_update", artifact.id, update_req) + assert updated is not None + assert updated.status == "production" + assert updated.metrics["accuracy"] == 0.95 + assert updated.description == "Production candidate" + + deleted = ModelRegistryService.delete_model("proj_update", artifact.id) + assert deleted is True + assert ModelRegistryService.get_model("proj_update", artifact.id) is None + + +def test_inspect_checkpoint(): + req = InspectCheckpointRequest( + checkpoint_path="model.safetensors", + parameters_count=7_000_000_000, + model_size_mb=14000.0, + framework="safetensors", + ) + inspection = ModelRegistryService.inspect_checkpoint(req) + assert inspection.file_format == "safetensors" + assert inspection.total_parameters == 7_000_000_000 + assert inspection.estimated_vram_fp16_mb > 0 + assert inspection.estimated_vram_int4_mb < inspection.estimated_vram_fp16_mb + + +def test_plan_quantization(): + req = RegisterModelRequest( + name="Mistral-7B", + parameters_count=7_000_000_000, + model_size_mb=14000.0, + ) + artifact = ModelRegistryService.register_model("proj_quant", req) + + plans = ModelRegistryService.plan_quantization(artifact, ["fp16", "int8", "int4", "fp8"]) + assert len(plans) == 4 + precisions = [p.target_precision for p in plans] + assert "FP16" in precisions + assert "INT4" in precisions + + int4_plan = next(p for p in plans if p.target_precision == "INT4") + assert int4_plan.compression_ratio > 5.0 + assert int4_plan.expected_latency_speedup > 2.0 + + +def test_model_card_generator_and_carbon(): + req = RegisterModelRequest( + name="GPT-Nano-Ablation", + version="2.1.0", + architecture="Transformer", + parameters_count=125_000_000, + model_size_mb=500.0, + metrics={"val_loss": 2.45, "hellaswag": 0.42}, + hyperparameters={"lr": 6e-4, "n_layer": 12, "n_head": 12}, + ) + artifact = ModelRegistryService.register_model("proj_card", req) + + card_req = GenerateModelCardRequest( + author="Silas Autonomous Agent", + license="MIT", + intended_use="Language modeling ablation benchmark", + limitations="Small parameter capacity, limited world knowledge.", + gpu_type="NVIDIA A100", + gpu_hours=48.0, + ) + card = ModelRegistryService.generate_model_card("proj_card", artifact.id, card_req) + assert card is not None + assert card.model_name == "GPT-Nano-Ablation" + assert "GPT-Nano-Ablation" in card.markdown + assert "\\begin{table}" in card.latex + assert "@misc{" in card.bibtex + assert card.co2_emissions_kg > 0 + + carbon = estimate_carbon_footprint("NVIDIA H100", 10.0) + assert carbon > 0 + + +def test_compare_models(): + m1 = ModelRegistryService.register_model( + "proj_comp", + RegisterModelRequest( + name="Baseline-Model", + version="1.0.0", + parameters_count=100_000_000, + model_size_mb=400.0, + metrics={"val_loss": 2.8, "accuracy": 0.70}, + ), + ) + m2 = ModelRegistryService.register_model( + "proj_comp", + RegisterModelRequest( + name="Novel-Attention-Model", + version="1.1.0", + parameters_count=105_000_000, + model_size_mb=420.0, + metrics={"val_loss": 2.1, "accuracy": 0.85}, + ), + ) + + comp = ModelRegistryService.compare_models("proj_comp", [m1.id, m2.id]) + assert "compared_models" in comp + assert len(comp["compared_models"]) == 2 + assert comp["recommended_model_id"] == m2.id + assert "Novel-Attention-Model" in comp["recommendation_reason"] diff --git a/backend/tests/test_reproducibility_auditor.py b/backend/tests/test_reproducibility_auditor.py new file mode 100644 index 0000000..7935ac3 --- /dev/null +++ b/backend/tests/test_reproducibility_auditor.py @@ -0,0 +1,156 @@ +"""Unit tests for the Reproducibility Auditor Service.""" + +from openmlr.services.reproducibility_auditor import ReproducibilityAuditorService +from openmlr.services.reproducibility_types import ( + AuditCodebaseRequest, + CheckCategory, + ChecklistVenue, + CheckStatus, + GenerateAppendixRequest, + GenerateDockerfileRequest, +) + + +def test_determinism_snippet_generation(): + py_snippet = ReproducibilityAuditorService.generate_determinism_snippet("pytorch", seed=123, strict_mode=True) + assert "torch.manual_seed(seed)" in py_snippet + assert "set_seed(123)" in py_snippet + assert "torch.backends.cudnn.deterministic = True" in py_snippet + assert "torch.use_deterministic_algorithms(True)" in py_snippet + assert "PYTHONHASHSEED" in py_snippet + + jax_snippet = ReproducibilityAuditorService.generate_determinism_snippet("jax", seed=999) + assert "jax.random.PRNGKey(999)" in jax_snippet + + tf_snippet = ReproducibilityAuditorService.generate_determinism_snippet("tensorflow", seed=777) + assert "tf.random.set_seed(777)" in tf_snippet + + +def test_generate_dockerfile(): + req = GenerateDockerfileRequest( + framework="pytorch", + cuda_version="12.1.0", + python_version="3.11", + entrypoint_cmd="python run_exp.py --seed 42", + requirements=["torch==2.1.0", "transformers==4.35.0"], + ) + dockerfile = ReproducibilityAuditorService.generate_dockerfile(req) + assert "FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04" in dockerfile + assert "CUBLAS_WORKSPACE_CONFIG=:4096:8" in dockerfile + assert "torch==2.1.0" in dockerfile + assert 'CMD ["python run_exp.py --seed 42"]' in dockerfile + + +def test_generate_conda_env(): + env = ReproducibilityAuditorService.generate_conda_env( + env_name="test-env", + python_version="3.10", + dependencies=["pytorch=2.1.0", "torchvision=0.16.0"], + ) + assert "name: test-env" in env + assert "python=3.10" in env + assert "pytorch=2.1.0" in env + + +def test_generate_badges(): + md = ReproducibilityAuditorService.generate_badge_markdown(96.5, "A+") + assert "reproducibility-A+%20(96%25)-brightgreen.svg" in md + + svg = ReproducibilityAuditorService.generate_badge_svg(96.5, "A+") + assert "= 85.0 + assert report.grade in ("A+", "A") + assert report.venue == ChecklistVenue.NEURIPS + assert len(report.categories) == len(CheckCategory) + + det_cat = next(c for c in report.categories if c.category == CheckCategory.DETERMINISM) + assert det_cat.score == 100.0 + assert det_cat.status == CheckStatus.PASS + + env_cat = next(c for c in report.categories if c.category == CheckCategory.ENVIRONMENT) + assert env_cat.score == 100.0 + + assert "PyTorch" in report.detected_frameworks + assert report.seeds_detected.get("main_seed") in (42, "args.seed") + assert "Reproducibility Statement" in report.latex_appendix + + # Check retrieval from store + retrieved = ReproducibilityAuditorService.get_report(report.id, "proj_test") + assert retrieved is not None + assert retrieved.id == report.id + + reports = ReproducibilityAuditorService.list_reports("proj_test") + assert len(reports) >= 1 + + deleted = ReproducibilityAuditorService.delete_report(report.id, "proj_test") + assert deleted is True + assert ReproducibilityAuditorService.get_report(report.id, "proj_test") is None diff --git a/backend/tests/test_routes_datasets.py b/backend/tests/test_routes_datasets.py new file mode 100644 index 0000000..04e8a93 --- /dev/null +++ b/backend/tests/test_routes_datasets.py @@ -0,0 +1,91 @@ +"""Tests for the Dataset Management, Profiling, and Split API routes.""" + +from __future__ import annotations + +import csv +from pathlib import Path + +import pytest +from httpx import AsyncClient + +pytestmark = pytest.mark.asyncio + + +@pytest.fixture +def sample_dataset_path(tmp_path: Path) -> str: + file_path = tmp_path / "train_data.csv" + data = [ + {"id": "1", "feature1": "10.5", "feature2": "20.1", "label": "pos", "text": "Good performance on CIFAR"}, + {"id": "2", "feature1": "15.2", "feature2": "18.3", "label": "neg", "text": "Poor generalization error"}, + {"id": "3", "feature1": "12.0", "feature2": "19.5", "label": "pos", "text": "Accurate prediction output"}, + {"id": "4", "feature1": "11.1", "feature2": "22.4", "label": "pos", "text": "Robust against adversarial noise"}, + ] + with open(file_path, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=list(data[0].keys())) + writer.writeheader() + writer.writerows(data) + return str(file_path) + + +class TestDatasetsRoutes: + async def test_profile_dataset_endpoint(self, client: AsyncClient, sample_dataset_path: str): + resp = await client.post( + "/api/datasets/profile", + json={"path": sample_dataset_path, "sample_size": 100}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["success"] is True + assert data["profile"]["total_rows"] == 4 + assert data["profile"]["total_columns"] == 5 + assert "feature1" in data["profile"]["columns"] + assert data["profile"]["columns"]["feature1"]["dtype"] == "numeric" + + async def test_profile_dataset_not_found(self, client: AsyncClient): + resp = await client.post( + "/api/datasets/profile", + json={"path": "/nonexistent/data.csv"}, + ) + assert resp.status_code == 404 + + async def test_inspect_samples_endpoint(self, client: AsyncClient, sample_dataset_path: str): + resp = await client.post( + "/api/datasets/inspect", + json={"path": sample_dataset_path, "n": 2, "strategy": "head"}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["success"] is True + assert data["total_sampled"] == 2 + assert len(data["samples"]) == 2 + + async def test_validate_dataset_endpoint(self, client: AsyncClient, sample_dataset_path: str): + resp = await client.post( + "/api/datasets/validate", + json={ + "path": sample_dataset_path, + "expected_columns": ["id", "feature1", "label"], + "max_null_pct": 10.0, + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["success"] is True + assert data["validation"]["valid"] is True + + async def test_split_dataset_endpoint(self, client: AsyncClient, sample_dataset_path: str, tmp_path: Path): + out_dir = str(tmp_path / "splits_api") + resp = await client.post( + "/api/datasets/split", + json={ + "path": sample_dataset_path, + "output_dir": out_dir, + "train_ratio": 0.5, + "val_ratio": 0.25, + "test_ratio": 0.25, + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["success"] is True + assert data["manifest"]["train_count"] == 2 diff --git a/backend/tests/test_routes_experiments.py b/backend/tests/test_routes_experiments.py new file mode 100644 index 0000000..221cffc --- /dev/null +++ b/backend/tests/test_routes_experiments.py @@ -0,0 +1,206 @@ +"""Tests for the Machine Learning Experiments and Run Tracking API routes.""" + +from __future__ import annotations + +import pytest +from httpx import AsyncClient + +pytestmark = pytest.mark.asyncio + + +class TestExperimentsRoutes: + async def test_create_and_get_run(self, client: AsyncClient): + payload = { + "name": "Transformer-FlashAttention-Benchmark", + "description": "Benchmarking memory footprint and throughput", + "hyperparameters": { + "lr": 0.0003, + "batch_size": 64, + "model": "Llama-1B", + }, + "compute_target": "Local H100", + "tags": ["transformer", "flashattention", "benchmark"], + "total_steps": 200, + "total_epochs": 2, + } + + create_resp = await client.post("/api/experiments/runs", json=payload) + assert create_resp.status_code == 201 + data = create_resp.json() + assert data["status"] == "created" + run = data["run"] + run_id = run["id"] + assert run["name"] == payload["name"] + assert run["hyperparameters"]["model"] == "Llama-1B" + + # Get run details + get_resp = await client.get(f"/api/experiments/runs/{run_id}") + assert get_resp.status_code == 200 + get_data = get_resp.json() + assert get_data["run"]["id"] == run_id + assert get_data["run"]["name"] == payload["name"] + + async def test_list_runs_with_query_params(self, client: AsyncClient): + # Create two runs + await client.post( + "/api/experiments/runs", + json={"name": "Diffusion-UNet-Run", "tags": ["cv", "diffusion"]}, + ) + await client.post( + "/api/experiments/runs", + json={"name": "RL-PPO-Agent", "tags": ["rl", "ppo"]}, + ) + + resp = await client.get("/api/experiments/runs?search=diffusion") + assert resp.status_code == 200 + data = resp.json() + assert "runs" in data + assert any("Diffusion-UNet-Run" in r["name"] for r in data["runs"]) + assert not any("RL-PPO-Agent" in r["name"] for r in data["runs"]) + + async def test_log_metrics_and_retrieve_trajectory(self, client: AsyncClient): + create_resp = await client.post( + "/api/experiments/runs", + json={"name": "Metric-Tracking-Test", "total_steps": 50}, + ) + run_id = create_resp.json()["run"]["id"] + + # Log step 10 + metric_resp1 = await client.post( + f"/api/experiments/runs/{run_id}/metrics", + json={ + "step": 10, + "epoch": 1, + "metrics": {"train_loss": 3.5, "val_loss": 3.8, "lr": 0.001}, + }, + ) + assert metric_resp1.status_code == 200 + assert metric_resp1.json()["current_step"] == 10 + assert metric_resp1.json()["best_val_loss"] == 3.8 + + # Log step 20 with lower val loss + metric_resp2 = await client.post( + f"/api/experiments/runs/{run_id}/metrics", + json={ + "step": 20, + "epoch": 1, + "metrics": {"train_loss": 2.8, "val_loss": 3.1, "lr": 0.0009}, + }, + ) + assert metric_resp2.status_code == 200 + assert metric_resp2.json()["best_val_loss"] == 3.1 + + # Check full run state + get_resp = await client.get(f"/api/experiments/runs/{run_id}") + run_data = get_resp.json()["run"] + assert len(run_data["metrics"]["train_loss"]) == 2 + assert run_data["best_val_loss"] == 3.1 + + async def test_update_status_and_logs(self, client: AsyncClient): + create_resp = await client.post( + "/api/experiments/runs", + json={"name": "Status-Log-Test"}, + ) + run_id = create_resp.json()["run"]["id"] + + # Append logs + log_resp = await client.post( + f"/api/experiments/runs/{run_id}/logs", + json={"lines": ["[INFO] Starting dataloader", "[INFO] GPU Allocated: 14.2 GB"]}, + ) + assert log_resp.status_code == 200 + assert log_resp.json()["total_lines"] == 2 + + # Get logs + get_log_resp = await client.get(f"/api/experiments/runs/{run_id}/logs") + assert get_log_resp.status_code == 200 + assert len(get_log_resp.json()["logs"]) == 2 + + # Update status to completed + status_resp = await client.post( + f"/api/experiments/runs/{run_id}/status", + json={"status": "completed", "reason": "Target loss achieved"}, + ) + assert status_resp.status_code == 200 + assert status_resp.json()["run"]["status"] == "completed" + + async def test_register_checkpoint(self, client: AsyncClient): + create_resp = await client.post( + "/api/experiments/runs", + json={"name": "Checkpoint-Test"}, + ) + run_id = create_resp.json()["run"]["id"] + + cp_resp = await client.post( + f"/api/experiments/runs/{run_id}/checkpoints", + json={ + "name": "model_epoch_5.pt", + "step": 500, + "epoch": 5, + "path": "/checkpoints/model_epoch_5.pt", + "file_size_mb": 750.2, + "metrics": {"val_loss": 1.45, "accuracy": 0.88}, + }, + ) + assert cp_resp.status_code == 200 + cp_data = cp_resp.json()["checkpoint"] + assert cp_data["name"] == "model_epoch_5.pt" + assert cp_data["file_size_mb"] == 750.2 + + async def test_compare_runs(self, client: AsyncClient): + r1_resp = await client.post( + "/api/experiments/runs", + json={"name": "Run-Comparison-A", "hyperparameters": {"opt": "adamw", "lr": 0.001}}, + ) + r2_resp = await client.post( + "/api/experiments/runs", + json={"name": "Run-Comparison-B", "hyperparameters": {"opt": "lion", "lr": 0.0001}}, + ) + id1 = r1_resp.json()["run"]["id"] + id2 = r2_resp.json()["run"]["id"] + + await client.post( + f"/api/experiments/runs/{id1}/metrics", + json={"step": 10, "metrics": {"train_loss": 2.0, "val_loss": 2.2}}, + ) + await client.post( + f"/api/experiments/runs/{id2}/metrics", + json={"step": 10, "metrics": {"train_loss": 1.8, "val_loss": 1.9}}, + ) + + comp_resp = await client.get(f"/api/experiments/compare?run_ids={id1},{id2}") + assert comp_resp.status_code == 200 + comp = comp_resp.json() + assert len(comp["runs"]) == 2 + assert "opt" in comp["hyperparameters_comparison"] + assert comp["metrics_summary"][id1]["best_val_loss"] == 2.2 + assert comp["metrics_summary"][id2]["best_val_loss"] == 1.9 + + async def test_delete_run(self, client: AsyncClient): + create_resp = await client.post( + "/api/experiments/runs", + json={"name": "Run-To-Delete"}, + ) + run_id = create_resp.json()["run"]["id"] + + del_resp = await client.delete(f"/api/experiments/runs/{run_id}") + assert del_resp.status_code == 200 + assert del_resp.json()["status"] == "deleted" + + # Subsequent get returns 404 + get_resp = await client.get(f"/api/experiments/runs/{run_id}") + assert get_resp.status_code == 404 + + async def test_error_handling(self, client: AsyncClient): + # 404 on nonexistent run + resp = await client.get("/api/experiments/runs/nonexistent-run-id") + assert resp.status_code == 404 + + # 400 on invalid status + create_resp = await client.post("/api/experiments/runs", json={"name": "Test"}) + run_id = create_resp.json()["run"]["id"] + bad_status_resp = await client.post( + f"/api/experiments/runs/{run_id}/status", + json={"status": "invalid_status_xyz"}, + ) + assert bad_status_resp.status_code == 400 diff --git a/backend/tests/test_routes_figures.py b/backend/tests/test_routes_figures.py new file mode 100644 index 0000000..b29878d --- /dev/null +++ b/backend/tests/test_routes_figures.py @@ -0,0 +1,64 @@ +"""Unit tests for the Figures REST API endpoints.""" + +import pytest +from httpx import ASGITransport, AsyncClient + +from openmlr.app import app + + +@pytest.mark.asyncio +async def test_figures_crud_routes(): + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + # Generate figure + payload = { + "title": "Empirical Scaling Law", + "caption": "Compute vs. loss scaling on OpenWebText.", + "plot_type": "loss_curve", + "style_theme": "neurips", + "palette": "colorblind", + "x_label": "FLOPs", + "y_label": "Test Perplexity", + "series_data": { + "Dense": [{"x": 1e18, "y": 12.4}, {"x": 1e19, "y": 8.2}], + "Sparse MoE": [{"x": 1e18, "y": 9.1}, {"x": 1e19, "y": 5.7}], + }, + "categories": [], + "values_matrix": [], + "width_inches": 6.0, + "height_inches": 4.0, + "generate_tikz": True, + } + res = await client.post("/api/figures?project_id=proj_api", json=payload) + assert res.status_code == 201 + data = res.json() + fig_id = data["figure"]["id"] + assert data["figure"]["title"] == "Empirical Scaling Law" + + # List figures + res_list = await client.get("/api/figures?project_id=proj_api") + assert res_list.status_code == 200 + assert res_list.json()["total_count"] >= 1 + + # Get figure + res_get = await client.get(f"/api/figures/{fig_id}?project_id=proj_api") + assert res_get.status_code == 200 + assert res_get.json()["figure"]["id"] == fig_id + + # Multi-panel + multi_res = await client.post( + "/api/figures/multi-panel?project_id=proj_api", + json={ + "title": "Combined Results", + "caption": "Summary of all figures.", + "figure_ids": [fig_id], + "columns": 1, + "subcaptions": {}, + }, + ) + assert multi_res.status_code == 200 + assert multi_res.json()["figure_count"] == 1 + + # Delete figure + del_res = await client.delete(f"/api/figures/{fig_id}?project_id=proj_api") + assert del_res.status_code == 200 + assert del_res.json()["success"] is True diff --git a/backend/tests/test_routes_models.py b/backend/tests/test_routes_models.py new file mode 100644 index 0000000..dc16073 --- /dev/null +++ b/backend/tests/test_routes_models.py @@ -0,0 +1,177 @@ +"""Tests for Model Registry, Model Card, Checkpoint Inspection, and Quantization API routes.""" + +from __future__ import annotations + +import pytest +from httpx import AsyncClient + +pytestmark = pytest.mark.asyncio + + +class TestModelsRoutes: + async def test_register_and_get_model(self, auth_client: AsyncClient): + payload = { + "name": "Diffusion-XL-OpenMLR", + "version": "1.0.0", + "architecture": "U-Net/DiT", + "framework": "safetensors", + "task_type": "diffusion", + "status": "evaluated", + "description": "Latent diffusion backbone for image generation.", + "parameters_count": 2_600_000_000, + "model_size_mb": 5200.0, + "tags": ["vision", "diffusion"], + "metrics": {"fid": 12.4, "clip_score": 0.32}, + "hyperparameters": {"steps": 50, "guidance_scale": 7.5}, + } + + resp = await auth_client.post("/api/model-registry?project_id=test_proj", json=payload) + assert resp.status_code == 201 + data = resp.json() + assert "model" in data + model = data["model"] + mid = model["id"] + assert model["name"] == payload["name"] + assert model["task_type"] == "diffusion" + + # Get model + get_resp = await auth_client.get(f"/api/model-registry/{mid}?project_id=test_proj") + assert get_resp.status_code == 200 + get_data = get_resp.json() + assert get_data["model"]["id"] == mid + + async def test_list_and_filter_models(self, auth_client: AsyncClient): + await auth_client.post( + "/api/model-registry?project_id=list_proj", + json={ + "name": "Classifier-A", + "framework": "pytorch", + "task_type": "classification", + "tags": ["nlp"], + }, + ) + await auth_client.post( + "/api/model-registry?project_id=list_proj", + json={ + "name": "LLM-B", + "framework": "safetensors", + "task_type": "causal_lm", + "tags": ["nlp"], + }, + ) + + resp = await auth_client.get("/api/model-registry?project_id=list_proj") + assert resp.status_code == 200 + data = resp.json() + assert data["total_count"] >= 2 + + filtered_resp = await auth_client.get("/api/model-registry?project_id=list_proj&task_type=causal_lm") + assert filtered_resp.status_code == 200 + filtered_data = filtered_resp.json() + assert all(m["task_type"] == "causal_lm" for m in filtered_data["models"]) + + async def test_update_and_delete_model(self, auth_client: AsyncClient): + reg = await auth_client.post( + "/api/model-registry?project_id=mod_proj", + json={"name": "TempModel", "status": "draft"}, + ) + mid = reg.json()["model"]["id"] + + # Update + up_resp = await auth_client.put( + f"/api/model-registry/{mid}?project_id=mod_proj", + json={"status": "production", "description": "Ready for prod"}, + ) + assert up_resp.status_code == 200 + assert up_resp.json()["model"]["status"] == "production" + + # Delete + del_resp = await auth_client.delete(f"/api/model-registry/{mid}?project_id=mod_proj") + assert del_resp.status_code == 200 + assert del_resp.json()["success"] is True + + # Check 404 + get_404 = await auth_client.get(f"/api/model-registry/{mid}?project_id=mod_proj") + assert get_404.status_code == 404 + + async def test_generate_model_card(self, auth_client: AsyncClient): + reg = await auth_client.post( + "/api/model-registry?project_id=card_proj", + json={ + "name": "ResNet-50-Ablated", + "version": "1.0.0", + "parameters_count": 25_000_000, + "model_size_mb": 100.0, + "metrics": {"top1_accuracy": 0.79}, + }, + ) + mid = reg.json()["model"]["id"] + + card_resp = await auth_client.post( + f"/api/model-registry/{mid}/card?project_id=card_proj", + json={ + "author": "Autonomous Scientist", + "license": "Apache-2.0", + "gpu_type": "NVIDIA A100", + "gpu_hours": 12.0, + }, + ) + assert card_resp.status_code == 200 + card_data = card_resp.json() + assert "markdown" in card_data + assert "latex" in card_data + assert "bibtex" in card_data + assert card_data["co2_emissions_kg"] > 0 + + async def test_plan_quantization(self, auth_client: AsyncClient): + reg = await auth_client.post( + "/api/model-registry?project_id=quant_proj", + json={ + "name": "Qwen-7B-OpenMLR", + "parameters_count": 7_000_000_000, + "model_size_mb": 14000.0, + }, + ) + mid = reg.json()["model"]["id"] + + quant_resp = await auth_client.post( + f"/api/model-registry/{mid}/quantization?project_id=quant_proj", + json={"target_precisions": ["fp16", "int8", "int4"]}, + ) + assert quant_resp.status_code == 200 + data = quant_resp.json() + assert len(data["estimates"]) == 3 + + async def test_inspect_checkpoint_route(self, auth_client: AsyncClient): + insp_resp = await auth_client.post( + "/api/model-registry/inspect", + json={ + "checkpoint_path": "model.safetensors", + "parameters_count": 1_500_000_000, + "framework": "safetensors", + }, + ) + assert insp_resp.status_code == 200 + data = insp_resp.json() + assert data["file_format"] == "safetensors" + assert data["estimated_vram_fp16_mb"] > 0 + + async def test_compare_models_route(self, auth_client: AsyncClient): + r1 = await auth_client.post( + "/api/model-registry?project_id=cmp_proj", + json={"name": "Model-Alpha", "metrics": {"accuracy": 0.82}}, + ) + r2 = await auth_client.post( + "/api/model-registry?project_id=cmp_proj", + json={"name": "Model-Beta", "metrics": {"accuracy": 0.91}}, + ) + id1 = r1.json()["model"]["id"] + id2 = r2.json()["model"]["id"] + + cmp_resp = await auth_client.post( + "/api/model-registry/compare?project_id=cmp_proj", + json={"model_ids": [id1, id2]}, + ) + assert cmp_resp.status_code == 200 + data = cmp_resp.json() + assert data["recommended_model_id"] == id2 diff --git a/backend/tests/test_routes_reproducibility.py b/backend/tests/test_routes_reproducibility.py new file mode 100644 index 0000000..39d6c43 --- /dev/null +++ b/backend/tests/test_routes_reproducibility.py @@ -0,0 +1,72 @@ +"""Unit tests for the Reproducibility REST API routes.""" + +import pytest +from httpx import ASGITransport, AsyncClient + +from openmlr.app import app + + +@pytest.mark.asyncio +async def test_reproducibility_routes_audit_and_crud(): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + # Audit + audit_payload = { + "target_path": "mock_test_path", + "venue": "neurips", + "code_snippets": { + "exp.py": "import torch\ntorch.manual_seed(42)\ntorch.save({}, 'model.pt')", + "requirements.txt": "torch==2.1.0\n", + }, + } + res_audit = await client.post("/api/reproducibility/audit?project_id=proj_routes", json=audit_payload) + assert res_audit.status_code == 200 + report = res_audit.json() + assert "id" in report + report_id = report["id"] + assert report["overall_score"] > 50.0 + + # List + res_list = await client.get("/api/reproducibility/reports?project_id=proj_routes") + assert res_list.status_code == 200 + reports = res_list.json() + assert len(reports) >= 1 + assert any(r["id"] == report_id for r in reports) + + # Get + res_get = await client.get(f"/api/reproducibility/reports/{report_id}?project_id=proj_routes") + assert res_get.status_code == 200 + assert res_get.json()["id"] == report_id + + # Dockerfile + res_dock = await client.post( + "/api/reproducibility/dockerfile", + json={"framework": "pytorch", "requirements": ["torch==2.1.0"]}, + ) + assert res_dock.status_code == 200 + assert "FROM nvidia/cuda" in res_dock.json()["dockerfile"] + + # Appendix + res_app = await client.post( + "/api/reproducibility/appendix?project_id=proj_routes", + json={"report_id": report_id, "paper_title": "Routes Test Paper"}, + ) + assert res_app.status_code == 200 + assert "\\section{Reproducibility Statement}" in res_app.json()["latex_appendix"] + + # Fix Determinism + res_fix = await client.post( + "/api/reproducibility/fix-determinism", + json={"framework": "pytorch", "seed": 42}, + ) + assert res_fix.status_code == 200 + assert "torch.manual_seed(seed)" in res_fix.json()["determinism_snippet"] + assert "set_seed(42)" in res_fix.json()["determinism_snippet"] + + # Delete + res_del = await client.delete(f"/api/reproducibility/reports/{report_id}?project_id=proj_routes") + assert res_del.status_code == 200 + + # Get 404 + res_get_deleted = await client.get(f"/api/reproducibility/reports/{report_id}?project_id=proj_routes") + assert res_get_deleted.status_code == 404 diff --git a/backend/tests/test_routes_research.py b/backend/tests/test_routes_research.py new file mode 100644 index 0000000..3e4b577 --- /dev/null +++ b/backend/tests/test_routes_research.py @@ -0,0 +1,211 @@ +"""Tests for the Research Workflow & State Machine API routes.""" + +from __future__ import annotations + +import pytest +from httpx import AsyncClient +from sqlalchemy.ext.asyncio import AsyncSession + +from openmlr.db import operations as ops +from openmlr.db.models import User + +pytestmark = pytest.mark.asyncio + + +class TestResearchPhasesAndGuidelines: + async def test_list_research_phases(self, client: AsyncClient): + resp = await client.get("/api/research/phases") + assert resp.status_code == 200 + data = resp.json() + assert "phases" in data + phase_ids = [p["id"] for p in data["phases"]] + assert "idle" in phase_ids + assert "reconnaissance" in phase_ids + assert "hypothesis" in phase_ids + assert "experimentation" in phase_ids + assert "analysis" in phase_ids + assert "paper_drafting" in phase_ids + assert "completed" in phase_ids + + async def test_get_all_guidelines(self, client: AsyncClient): + resp = await client.get("/api/research/guidelines") + assert resp.status_code == 200 + data = resp.json() + assert "reconnaissance" in data + assert "hypothesis" in data + assert "experimentation" in data + assert "analysis" in data + assert "paper_drafting" in data + + +class TestProjectResearchWorkflow: + @pytest.fixture + async def project(self, db_session: AsyncSession, test_user: User): + return await ops.create_project( + db_session, + user_id=test_user.id, + name="Autonomous Attention Scaling", + slug="autonomous-attention-scaling", + description="Investigate sub-quadratic attention variants", + ) + + async def test_get_initial_research_state( + self, auth_client: AsyncClient, project + ): + resp = await auth_client.get(f"/api/projects/{project.id}/research/state") + assert resp.status_code == 200 + data = resp.json() + assert data["project_id"] == project.id + assert data["project_name"] == "Autonomous Attention Scaling" + assert "state" in data + assert "guidelines" in data + + async def test_start_research_workflow( + self, auth_client: AsyncClient, project + ): + resp = await auth_client.post( + f"/api/projects/{project.id}/research/start", + json={ + "goal": "Benchmark FlashAttention-3 vs RingAttention", + "initial_phase": "reconnaissance", + "generate_default_milestones": True, + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "started" + assert data["state"]["goal"] == "Benchmark FlashAttention-3 vs RingAttention" + assert data["state"]["current_phase"] == "reconnaissance" + assert len(data["state"]["milestones"]) >= 5 + + async def test_start_research_invalid_phase( + self, auth_client: AsyncClient, project + ): + resp = await auth_client.post( + f"/api/projects/{project.id}/research/start", + json={ + "goal": "Invalid Phase Test", + "initial_phase": "invalid_unknown_phase", + }, + ) + assert resp.status_code == 400 + assert "invalid phase" in resp.json()["detail"].lower() + + async def test_transition_phase( + self, auth_client: AsyncClient, project + ): + # Start first + await auth_client.post( + f"/api/projects/{project.id}/research/start", + json={ + "goal": "LoRA Rank Scaling Analysis", + "initial_phase": "reconnaissance", + }, + ) + + # Transition to hypothesis + resp = await auth_client.post( + f"/api/projects/{project.id}/research/transition", + json={ + "next_phase": "hypothesis", + "reason": "Cataloged 8 foundational literature papers", + "artifacts_produced": ["paper_survey_table"], + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "transitioned" + assert data["state"]["current_phase"] == "hypothesis" + assert data["transition"]["to_phase"] == "hypothesis" + assert "paper_survey_table" in data["transition"]["artifacts_produced"] + + async def test_create_and_update_milestone( + self, auth_client: AsyncClient, project + ): + # Add custom milestone + resp = await auth_client.post( + f"/api/projects/{project.id}/research/milestones", + json={ + "title": "Train baseline ResNet-18", + "description": "Achieve >92% accuracy on CIFAR-10", + "phase": "experimentation", + "criteria": ["Accuracy >= 0.92", "Loss <= 0.35"], + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "created" + milestone = data["milestone"] + m_id = milestone["milestone_id"] + assert milestone["title"] == "Train baseline ResNet-18" + assert milestone["status"] == "pending" + + # Update milestone to completed + update_resp = await auth_client.put( + f"/api/projects/{project.id}/research/milestones/{m_id}", + json={ + "status": "completed", + "output_artifacts": ["baseline_model.pt", "eval_log.json"], + }, + ) + assert update_resp.status_code == 200 + updated_data = update_resp.json() + assert updated_data["status"] == "updated" + assert updated_data["milestone"]["status"] == "completed" + assert "baseline_model.pt" in updated_data["milestone"]["output_artifacts"] + + async def test_register_artifacts( + self, auth_client: AsyncClient, project + ): + # Register paper + p_resp = await auth_client.post( + f"/api/projects/{project.id}/research/artifacts", + json={"type": "paper", "data": {"arxiv_id": "2401.00001", "title": "Scalable ML"}}, + ) + assert p_resp.status_code == 200 + assert p_resp.json()["artifacts_summary"]["papers"] >= 1 + + # Register hypothesis + h_resp = await auth_client.post( + f"/api/projects/{project.id}/research/artifacts", + json={"type": "hypothesis", "data": {"claim": "Quantization retains 99% accuracy"}}, + ) + assert h_resp.status_code == 200 + assert h_resp.json()["artifacts_summary"]["hypotheses"] >= 1 + + # Register metrics + m_resp = await auth_client.post( + f"/api/projects/{project.id}/research/artifacts", + json={"type": "metrics", "data": {"train_loss": 0.24, "val_loss": 0.29}}, + ) + assert m_resp.status_code == 200 + assert "train_loss" in m_resp.json()["artifacts_summary"]["metrics_keys"] + + # Register manuscript section + s_resp = await auth_client.post( + f"/api/projects/{project.id}/research/artifacts", + json={ + "type": "manuscript_section", + "section_name": "methodology", + "data": "\\section{Methodology}\nWe propose a novel attention operator...", + }, + ) + assert s_resp.status_code == 200 + assert "methodology" in s_resp.json()["artifacts_summary"]["sections"] + + # Register bibtex + b_resp = await auth_client.post( + f"/api/projects/{project.id}/research/artifacts", + json={ + "type": "bibtex", + "data": "@article{scale2026, title={Scalable ML}}", + }, + ) + assert b_resp.status_code == 200 + assert b_resp.json()["artifacts_summary"]["bibtex_count"] >= 1 + + async def test_nonexistent_project_returns_404( + self, auth_client: AsyncClient + ): + resp = await auth_client.get("/api/projects/999999/research/state") + assert resp.status_code == 404 diff --git a/backend/tests/test_routes_sweeps.py b/backend/tests/test_routes_sweeps.py new file mode 100644 index 0000000..9e567bc --- /dev/null +++ b/backend/tests/test_routes_sweeps.py @@ -0,0 +1,121 @@ +"""Tests for the Hyperparameter Sweep and HPO API routes.""" + +from __future__ import annotations + +import pytest +from httpx import AsyncClient + +pytestmark = pytest.mark.asyncio + + +class TestSweepsRoutes: + async def test_create_and_get_sweep(self, client: AsyncClient): + payload = { + "name": "LoRA Rank & Alpha Sweep", + "description": "Optimizing LoRA hyperparameters on fine-tuning", + "method": "random", + "objective_metric": "eval_loss", + "goal": "minimize", + "max_trials": 4, + "parameters": { + "lora_r": {"param_type": "choice", "choices": [8, 16, 32]}, + "lora_alpha": {"param_type": "choice", "choices": [16, 32, 64]}, + "lr": {"param_type": "loguniform", "min_val": 1e-5, "max_val": 1e-3}, + }, + } + + resp = await client.post("/api/sweeps", json=payload) + assert resp.status_code == 201 + data = resp.json() + assert "sweep" in data + sweep = data["sweep"] + sweep_id = sweep["sweep_id"] + assert sweep["name"] == payload["name"] + assert sweep["objective_metric"] == "eval_loss" + + # Get sweep + get_resp = await client.get(f"/api/sweeps/{sweep_id}") + assert get_resp.status_code == 200 + get_data = get_resp.json() + assert get_data["sweep"]["sweep_id"] == sweep_id + + async def test_list_sweeps(self, client: AsyncClient): + await client.post( + "/api/sweeps", + json={ + "name": "CNN Filter Sweep", + "method": "grid", + "objective_metric": "accuracy", + "goal": "maximize", + "parameters": {"filters": {"param_type": "choice", "choices": [32, 64]}}, + }, + ) + + resp = await client.get("/api/sweeps") + assert resp.status_code == 200 + data = resp.json() + assert data["total"] >= 1 + assert any(s["name"] == "CNN Filter Sweep" for s in data["sweeps"]) + + async def test_suggest_record_and_analysis_lifecycle(self, client: AsyncClient): + create_resp = await client.post( + "/api/sweeps", + json={ + "name": "Vision Transformer Sweep", + "method": "random", + "objective_metric": "val_accuracy", + "goal": "maximize", + "max_trials": 3, + "parameters": { + "patch_size": {"param_type": "choice", "choices": [8, 16]}, + "lr": {"param_type": "uniform", "min_val": 0.0001, "max_val": 0.001}, + }, + "early_stopping": {"enabled": True, "min_steps": 2, "reduction_factor": 2.0}, + }, + ) + assert create_resp.status_code == 201 + sweep_id = create_resp.json()["sweep"]["sweep_id"] + + # Suggest trial 1 + sug_resp = await client.post(f"/api/sweeps/{sweep_id}/suggest") + assert sug_resp.status_code == 200 + trial = sug_resp.json()["trial"] + assert trial is not None + trial_id = trial["trial_id"] + + # Prune check + prune_resp = await client.post( + f"/api/sweeps/{sweep_id}/trials/{trial_id}/prune-check", + json={"current_step": 3, "current_metric_val": 0.95}, + ) + assert prune_resp.status_code == 200 + assert "should_prune" in prune_resp.json() + + # Record trial 1 + rec_resp = await client.post( + f"/api/sweeps/{sweep_id}/trials/{trial_id}/record", + json={ + "metrics": {"val_accuracy": 0.94, "loss": 0.12}, + "status": "completed", + "step_history": [{"step": 1, "val_accuracy": 0.8}, {"step": 2, "val_accuracy": 0.94}], + }, + ) + assert rec_resp.status_code == 200 + assert rec_resp.json()["trial"]["status"] == "completed" + + # Analysis + analysis_resp = await client.get(f"/api/sweeps/{sweep_id}/analysis") + assert analysis_resp.status_code == 200 + analysis = analysis_resp.json()["analysis"] + assert analysis["completed_trials"] == 1 + assert analysis["best_metric_value"] == 0.94 + + # Export report + export_resp = await client.post(f"/api/sweeps/{sweep_id}/export") + assert export_resp.status_code == 200 + assert "Hyperparameter Optimization Report" in export_resp.json()["report"] + + # Delete sweep + del_resp = await client.delete(f"/api/sweeps/{sweep_id}") + assert del_resp.status_code == 200 + assert del_resp.json()["deleted"] is True diff --git a/backend/tests/test_sweep_engine.py b/backend/tests/test_sweep_engine.py new file mode 100644 index 0000000..5ee91d8 --- /dev/null +++ b/backend/tests/test_sweep_engine.py @@ -0,0 +1,219 @@ +"""Tests for SweepEngine (hyperparameter optimization, sampling, pruning, analysis).""" + +from pathlib import Path + +import pytest + +from openmlr.services.sweep_engine import ( + EarlyStoppingConfig, + ParameterSpec, + SweepEngine, +) + + +@pytest.fixture +def sweep_engine(tmp_path: Path): + return SweepEngine(base_dir=tmp_path / "sweeps") + + +def test_create_and_get_sweep(sweep_engine: SweepEngine): + params = { + "learning_rate": ParameterSpec(name="learning_rate", param_type="loguniform", min_val=1e-5, max_val=1e-2), + "batch_size": ParameterSpec(name="batch_size", param_type="choice", choices=[16, 32, 64]), + "epochs": ParameterSpec(name="epochs", param_type="int_uniform", min_val=5, max_val=20, step=5), + } + + sweep = sweep_engine.create_sweep( + project_id="proj_1", + name="Transformer LR & Batch Sweep", + method="random", + objective_metric="val_loss", + goal="minimize", + parameters=params, + max_trials=5, + ) + + assert sweep.sweep_id.startswith("swp_") + assert sweep.name == "Transformer LR & Batch Sweep" + assert len(sweep.parameters) == 3 + assert sweep.status == "active" + + loaded = sweep_engine.get_sweep("proj_1", sweep.sweep_id) + assert loaded is not None + assert loaded.name == sweep.name + assert loaded.objective_metric == "val_loss" + + +def test_list_and_delete_sweeps(sweep_engine: SweepEngine): + s1 = sweep_engine.create_sweep( + project_id="proj_test", + name="Sweep 1", + method="grid", + objective_metric="accuracy", + goal="maximize", + parameters={"lr": {"param_type": "choice", "choices": [0.01, 0.001]}}, + ) + s2 = sweep_engine.create_sweep( + project_id="proj_test", + name="Sweep 2", + method="random", + objective_metric="val_loss", + goal="minimize", + parameters={"dropout": {"param_type": "uniform", "min_val": 0.1, "max_val": 0.5}}, + ) + + sweeps = sweep_engine.list_sweeps("proj_test") + assert len(sweeps) == 2 + + assert sweep_engine.delete_sweep("proj_test", s1.sweep_id) is True + sweeps_after = sweep_engine.list_sweeps("proj_test") + assert len(sweeps_after) == 1 + assert sweeps_after[0].sweep_id == s2.sweep_id + + +def test_suggest_grid_trials(sweep_engine: SweepEngine): + params = { + "lr": ParameterSpec(name="lr", param_type="choice", choices=[0.01, 0.001]), + "optimizer": ParameterSpec(name="optimizer", param_type="categorical", choices=["adam", "sgd"]), + } + sweep = sweep_engine.create_sweep( + project_id="proj_grid", + name="Grid Test", + method="grid", + objective_metric="val_loss", + goal="minimize", + parameters=params, + max_trials=4, + ) + + t1 = sweep_engine.suggest_trial("proj_grid", sweep.sweep_id) + assert t1 is not None + assert t1.trial_number == 1 + assert "lr" in t1.parameters + assert "optimizer" in t1.parameters + + t2 = sweep_engine.suggest_trial("proj_grid", sweep.sweep_id) + t3 = sweep_engine.suggest_trial("proj_grid", sweep.sweep_id) + t4 = sweep_engine.suggest_trial("proj_grid", sweep.sweep_id) + assert t4 is not None + + # Max trials reached + t5 = sweep_engine.suggest_trial("proj_grid", sweep.sweep_id) + assert t5 is None + + +def test_suggest_bayesian_optimization(sweep_engine: SweepEngine): + params = { + "lr": ParameterSpec(name="lr", param_type="uniform", min_val=0.0001, max_val=0.01), + "hidden_dim": ParameterSpec(name="hidden_dim", param_type="choice", choices=[128, 256, 512]), + } + sweep = sweep_engine.create_sweep( + project_id="proj_bayes", + name="Bayes Test", + method="bayesian", + objective_metric="val_loss", + goal="minimize", + parameters=params, + max_trials=10, + ) + + # Seed 3 trials + for i in range(3): + t = sweep_engine.suggest_trial("proj_bayes", sweep.sweep_id) + assert t is not None + sweep_engine.record_trial_result( + project_id="proj_bayes", + sweep_id=sweep.sweep_id, + trial_id=t.trial_id, + metrics={"val_loss": 0.5 - i * 0.1}, + status="completed", + ) + + # 4th trial uses Bayesian surrogate + t4 = sweep_engine.suggest_trial("proj_bayes", sweep.sweep_id) + assert t4 is not None + assert "lr" in t4.parameters + assert "hidden_dim" in t4.parameters + + +def test_early_stopping_pruning(sweep_engine: SweepEngine): + es = EarlyStoppingConfig(enabled=True, min_steps=3, reduction_factor=2.0) + params = {"lr": ParameterSpec(name="lr", param_type="choice", choices=[0.01, 0.001, 0.0001])} + sweep = sweep_engine.create_sweep( + project_id="proj_es", + name="ASHA Prune Test", + method="hyperband", + objective_metric="val_loss", + goal="minimize", + parameters=params, + early_stopping=es, + max_trials=5, + ) + + # Seed 2 good trials with step history + t1 = sweep_engine.suggest_trial("proj_es", sweep.sweep_id) + assert t1 is not None + sweep_engine.record_trial_result( + "proj_es", + sweep.sweep_id, + t1.trial_id, + metrics={"val_loss": 0.2}, + step_history=[{"step": 1, "val_loss": 0.8}, {"step": 3, "val_loss": 0.3}], + ) + + t2 = sweep_engine.suggest_trial("proj_es", sweep.sweep_id) + assert t2 is not None + sweep_engine.record_trial_result( + "proj_es", + sweep.sweep_id, + t2.trial_id, + metrics={"val_loss": 0.25}, + step_history=[{"step": 1, "val_loss": 0.9}, {"step": 3, "val_loss": 0.35}], + ) + + t3 = sweep_engine.suggest_trial("proj_es", sweep.sweep_id) + assert t3 is not None + + # Step 1 is below min_steps (3), so shouldn't prune + assert not sweep_engine.should_prune_trial("proj_es", sweep.sweep_id, t3.trial_id, 1, 1.5) + + # Step 3 with high val_loss (1.5 >> 0.3) should be pruned + assert sweep_engine.should_prune_trial("proj_es", sweep.sweep_id, t3.trial_id, 3, 1.5) + + +def test_analyze_sweep_and_markdown_export(sweep_engine: SweepEngine): + params = { + "lr": ParameterSpec(name="lr", param_type="uniform", min_val=0.001, max_val=0.1), + "weight_decay": ParameterSpec(name="weight_decay", param_type="uniform", min_val=1e-5, max_val=1e-3), + } + sweep = sweep_engine.create_sweep( + project_id="proj_analysis", + name="Sensitivity Analysis", + method="random", + objective_metric="val_loss", + goal="minimize", + parameters=params, + max_trials=4, + ) + + for i in range(4): + t = sweep_engine.suggest_trial("proj_analysis", sweep.sweep_id) + assert t is not None + sweep_engine.record_trial_result( + "proj_analysis", + sweep.sweep_id, + t.trial_id, + metrics={"val_loss": 0.4 - i * 0.05, "accuracy": 0.8 + i * 0.03}, + status="completed", + ) + + analysis = sweep_engine.analyze_sweep("proj_analysis", sweep.sweep_id) + assert analysis["completed_trials"] == 4 + assert analysis["best_trial"] is not None + assert "parameter_importance" in analysis + assert len(analysis["pareto_frontier"]) > 0 + + md = sweep_engine.export_sweep_markdown("proj_analysis", sweep.sweep_id) + assert "Hyperparameter Optimization Report" in md + assert "Optimal Configuration" in md + assert "Trial History" in md diff --git a/backend/tests/test_tools_datasets.py b/backend/tests/test_tools_datasets.py new file mode 100644 index 0000000..6f0200e --- /dev/null +++ b/backend/tests/test_tools_datasets.py @@ -0,0 +1,117 @@ +"""Unit tests for the datasets agent tool.""" + +from __future__ import annotations + +import csv +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from openmlr.tools.datasets import _handle_datasets, create_datasets_tool + + +@pytest.fixture +def sample_csv(tmp_path: Path) -> Path: + csv_file = tmp_path / "research_data.csv" + data = [ + {"id": "1", "score": "95.5", "split": "train", "text": "Self-supervised representation learning"}, + {"id": "2", "score": "88.0", "split": "train", "text": "Reinforcement learning from human feedback"}, + {"id": "3", "score": "72.4", "split": "val", "text": "Direct preference optimization"}, + ] + with open(csv_file, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=list(data[0].keys())) + writer.writeheader() + writer.writerows(data) + return csv_file + + +def test_create_datasets_tool_spec(): + tool = create_datasets_tool() + assert tool.name == "datasets" + assert "operation" in tool.parameters["properties"] + assert "profile" in tool.parameters["properties"]["operation"]["enum"] + assert "split" in tool.parameters["properties"]["operation"]["enum"] + + +@pytest.mark.asyncio +async def test_handle_datasets_profile(sample_csv: Path): + output, success = await _handle_datasets( + operation="profile", + path=str(sample_csv), + ) + assert success is True + assert "Dataset Profile" in output + assert "score" in output + + +@pytest.mark.asyncio +async def test_handle_datasets_inspect_samples(sample_csv: Path): + output, success = await _handle_datasets( + operation="inspect_samples", + path=str(sample_csv), + n=2, + ) + assert success is True + assert "Dataset Samples" in output + assert "Direct preference optimization" in output or "Self-supervised" in output + + +@pytest.mark.asyncio +async def test_handle_datasets_validate(sample_csv: Path): + output, success = await _handle_datasets( + operation="validate", + path=str(sample_csv), + expected_columns=["id", "score", "split", "text"], + ) + assert success is True + assert "PASSED" in output + + +@pytest.mark.asyncio +async def test_handle_datasets_split(sample_csv: Path, tmp_path: Path): + split_dir = tmp_path / "splits" + output, success = await _handle_datasets( + operation="split", + path=str(sample_csv), + output_dir=str(split_dir), + train_ratio=0.7, + val_ratio=0.3, + test_ratio=0.0, + ) + assert success is True + assert "Dataset Split Completed" in output + assert (split_dir / "split_manifest.json").exists() + + +@pytest.mark.asyncio +async def test_handle_datasets_register(sample_csv: Path): + mock_session = MagicMock() + mock_kg = MagicMock() + mock_session.workspace.knowledge_graph = mock_kg + + output, success = await _handle_datasets( + operation="register", + path=str(sample_csv), + dataset_name="cifar10_custom", + description="Custom CIFAR-10 evaluation set", + tags=["cv", "benchmark"], + session=mock_session, + ) + assert success is True + assert "Dataset Registered" in output + assert mock_kg.add_entity.called + + +@pytest.mark.asyncio +async def test_handle_datasets_summary_help(): + output, success = await _handle_datasets(operation="summary") + assert success is True + assert "Datasets Tool Operations" in output + + +@pytest.mark.asyncio +async def test_handle_datasets_missing_path(): + output, success = await _handle_datasets(operation="profile", path="") + assert success is False + assert "required" in output.lower() diff --git a/backend/tests/test_tools_experiments.py b/backend/tests/test_tools_experiments.py new file mode 100644 index 0000000..c19d08c --- /dev/null +++ b/backend/tests/test_tools_experiments.py @@ -0,0 +1,283 @@ +"""Tests for experiments tool (openmlr/tools/experiments.py).""" + +import json +from pathlib import Path + +import pytest + +from openmlr.agent.types import ToolSpec +from openmlr.services.experiment_tracker import ExperimentTracker +from openmlr.tools.experiments import ( + _handle_experiments, + create_experiments_tool, + set_experiment_context, +) +from openmlr.tools.workspace_tools import set_workspace_context +from openmlr.workspace.knowledge import KnowledgeGraph + + +@pytest.fixture +def clean_tracker(tmp_path: Path): + tracker = ExperimentTracker(storage_dir=tmp_path / "experiments") + set_experiment_context(tracker, project_uuid="proj-test-123") + yield tracker + set_experiment_context(None, None) + + +class TestExperimentsToolSpec: + def test_creates_tool_spec(self): + tool = create_experiments_tool() + assert isinstance(tool, ToolSpec) + assert tool.name == "experiments" + assert tool.handler is not None + assert "create_run" in tool.description + assert "log_metrics" in tool.description + assert "record_checkpoint" in tool.description + assert "compare_runs" in tool.description + assert "operation" in tool.parameters["properties"] + + +class TestExperimentsToolCreateRun: + async def test_create_run_success(self, clean_tracker): + out, success = await _handle_experiments( + operation="create_run", + name="Attention-Optimization-v1", + description="Testing FlashAttention v2 vs baseline", + hyperparameters={"lr": 0.0003, "batch_size": 32, "layers": 12}, + compute_target="Modal A100", + tags=["transformer", "attention", "speedup"], + total_steps=500, + total_epochs=5, + ) + assert success is True + data = json.loads(out) + assert data["name"] == "Attention-Optimization-v1" + assert data["compute_target"] == "Modal A100" + assert data["total_steps"] == 500 + assert data["hyperparameters"]["lr"] == 0.0003 + assert "run_id" in data + + async def test_create_run_missing_name(self, clean_tracker): + out, success = await _handle_experiments(operation="create_run", name="") + assert success is False + assert "name" in out.lower() + + async def test_create_run_json_strings(self, clean_tracker): + out, success = await _handle_experiments( + operation="create_run", + name="JSON-String-Run", + hyperparameters='{"lr": 0.001, "optimizer": "AdamW"}', + tags="llm, bert, fine-tune", + ) + assert success is True + data = json.loads(out) + assert data["hyperparameters"]["optimizer"] == "AdamW" + + +class TestExperimentsToolLogMetrics: + async def test_log_metrics_success(self, clean_tracker): + # Create run first + create_out, _ = await _handle_experiments(operation="create_run", name="Metrics-Run") + run_id = json.loads(create_out)["run_id"] + + # Log step 10 + out1, success1 = await _handle_experiments( + operation="log_metrics", + run_id=run_id, + step=10, + epoch=1, + metrics={"train_loss": 3.42, "val_loss": 3.51, "learning_rate": 0.0003}, + ) + assert success1 is True + data1 = json.loads(out1) + assert data1["current_step"] == 10 + assert data1["best_val_loss"] == 3.51 + + # Log step 20 with better val loss + out2, success2 = await _handle_experiments( + operation="log_metrics", + run_id=run_id, + step=20, + epoch=1, + metrics='{"train_loss": 2.85, "val_loss": 2.91}', + ) + assert success2 is True + data2 = json.loads(out2) + assert data2["best_val_loss"] == 2.91 + + async def test_log_metrics_missing_run_id_or_empty(self, clean_tracker): + out, success = await _handle_experiments(operation="log_metrics", run_id="", metrics={"val_loss": 1.0}) + assert success is False + assert "run_id" in out + + out2, success2 = await _handle_experiments(operation="log_metrics", run_id="nonexistent", metrics={}) + assert success2 is False + + async def test_log_metrics_nonexistent_run(self, clean_tracker): + out, success = await _handle_experiments( + operation="log_metrics", + run_id="run-fake-999", + step=1, + metrics={"train_loss": 2.0}, + ) + assert success is False + assert "not found" in out.lower() + + +class TestExperimentsToolCheckpoints: + async def test_record_checkpoint_success(self, clean_tracker): + create_out, _ = await _handle_experiments(operation="create_run", name="Checkpoint-Run") + run_id = json.loads(create_out)["run_id"] + + out, success = await _handle_experiments( + operation="record_checkpoint", + run_id=run_id, + checkpoint_name="best_model.pt", + path="checkpoints/best_model.pt", + file_size_mb=420.5, + step=250, + epoch=2, + metrics={"val_loss": 1.84, "accuracy": 0.892}, + ) + assert success is True + data = json.loads(out) + assert data["total_checkpoints"] == 1 + assert data["checkpoint"]["name"] == "best_model.pt" + assert data["checkpoint"]["metrics"]["accuracy"] == 0.892 + + +class TestExperimentsToolGetRun: + async def test_get_run_details(self, clean_tracker): + create_out, _ = await _handle_experiments( + operation="create_run", + name="Inspect-Run", + description="Testing retrieval", + hyperparameters={"weight_decay": 0.01}, + ) + run_id = json.loads(create_out)["run_id"] + + await _handle_experiments( + operation="log_metrics", + run_id=run_id, + step=50, + metrics={"train_loss": 1.2, "val_loss": 1.4}, + ) + + out, success = await _handle_experiments(operation="get_run", run_id=run_id) + assert success is True + data = json.loads(out) + assert data["run_id"] == run_id + assert data["name"] == "Inspect-Run" + assert data["latest_metrics"]["train_loss"] == 1.2 + assert data["latest_metrics"]["val_loss"] == 1.4 + + +class TestExperimentsToolListAndCompare: + async def test_list_and_compare_runs(self, clean_tracker): + # Create Run A + out_a, _ = await _handle_experiments( + operation="create_run", + name="Model-Baseline", + hyperparameters={"lr": 0.001, "arch": "standard"}, + tags=["baseline"], + ) + id_a = json.loads(out_a)["run_id"] + await _handle_experiments( + operation="log_metrics", + run_id=id_a, + step=100, + metrics={"train_loss": 1.5, "val_loss": 1.6}, + ) + + # Create Run B + out_b, _ = await _handle_experiments( + operation="create_run", + name="Model-Modified", + hyperparameters={"lr": 0.0003, "arch": "improved"}, + tags=["modified"], + ) + id_b = json.loads(out_b)["run_id"] + await _handle_experiments( + operation="log_metrics", + run_id=id_b, + step=100, + metrics={"train_loss": 1.1, "val_loss": 1.2}, + ) + + # List runs + list_out, list_ok = await _handle_experiments(operation="list_runs") + assert list_ok is True + list_data = json.loads(list_out) + assert list_data["total_runs"] >= 2 + + # Compare runs + cmp_out, cmp_ok = await _handle_experiments( + operation="compare_runs", + run_ids=f"{id_a}, {id_b}", + ) + assert cmp_ok is True + cmp_data = json.loads(cmp_out) + assert id_a in cmp_data["metrics_summary"] + assert id_b in cmp_data["metrics_summary"] + assert cmp_data["hyperparameters_comparison"]["arch"][id_a] == "standard" + assert cmp_data["hyperparameters_comparison"]["arch"][id_b] == "improved" + + +class TestExperimentsToolCompleteRun: + async def test_complete_run(self, clean_tracker): + create_out, _ = await _handle_experiments(operation="create_run", name="To-Complete") + run_id = json.loads(create_out)["run_id"] + + out, success = await _handle_experiments( + operation="complete_run", + run_id=run_id, + status="completed", + best_val_loss=1.15, + reason="Finished all 500 steps with convergence", + ) + assert success is True + data = json.loads(out) + assert data["status"] == "completed" + assert data["best_val_loss"] == 1.15 + + +class TestExperimentsToolInvalidOperation: + async def test_invalid_op(self, clean_tracker): + out, success = await _handle_experiments(operation="unknown_op") + assert success is False + assert "Unknown experiments operation" in out + + +class TestKnowledgeGraphIntegration: + async def test_auto_registers_in_knowledge_graph(self, clean_tracker, tmp_path: Path): + set_workspace_context(str(tmp_path)) + + out, success = await _handle_experiments( + operation="create_run", + name="KG-Tracked-Experiment", + hyperparameters={"batch_size": 64}, + ) + assert success is True + run_id = json.loads(out)["run_id"] + + # Verify entity exists in knowledge graph + kg = KnowledgeGraph(str(tmp_path)) + entity = kg.get_entity(f"exp_{run_id}") + assert entity is not None + assert entity["label"] == "KG-Tracked-Experiment" + assert entity["type"] == "experiment" + + # Complete and check update + await _handle_experiments( + operation="complete_run", + run_id=run_id, + status="completed", + best_val_loss=0.88, + ) + kg_updated = KnowledgeGraph(str(tmp_path)) + updated_entity = kg_updated.get_entity(f"exp_{run_id}") + assert updated_entity is not None + assert updated_entity["status"] == "completed" + assert updated_entity["best_val_loss"] == 0.88 + + set_workspace_context(None) diff --git a/backend/tests/test_tools_figures.py b/backend/tests/test_tools_figures.py new file mode 100644 index 0000000..a710f52 --- /dev/null +++ b/backend/tests/test_tools_figures.py @@ -0,0 +1,61 @@ +"""Unit tests for the Figures agent tool.""" + +import pytest + +from openmlr.tools.figures import create_figures_tool + + +@pytest.mark.asyncio +async def test_figures_tool_execution(): + tool = create_figures_tool(get_project_id=lambda: "proj_tool") + assert tool.name == "figures" + assert tool.handler is not None + + # Action: generate + res, ok = await tool.handler( + action="generate", + title="Learning Rate Sensitivity", + caption="Loss under various learning rates.", + plot_type="loss_curve", + style_theme="iclr", + palette="tableau", + x_label="Learning Rate", + y_label="Validation Loss", + series_data={ + "Run 1": [{"x": 1e-4, "y": 2.1}, {"x": 3e-4, "y": 1.4}, {"x": 1e-3, "y": 1.9}] + }, + ) + assert ok is True + assert "Learning Rate Sensitivity" in res + assert "Figure ID:" in res + + # Action: list + res_list, ok_list = await tool.handler(action="list") + assert ok_list is True + assert "Found" in res_list + + # Extract ID + fig_id = None + for line in res_list.splitlines(): + if "Learning Rate Sensitivity" in line and "`fig_" in line: + start = line.find("`fig_") + 1 + end = line.find("`", start) + fig_id = line[start:end] + break + + assert fig_id is not None + + # Action: get + res_get, ok_get = await tool.handler(action="get", figure_id=fig_id) + assert ok_get is True + assert "Learning Rate Sensitivity" in res_get + + # Action: create_multipanel with 1 ID fails validation (needs >=2) + res_err, ok_err = await tool.handler(action="create_multipanel", figure_ids=[fig_id]) + assert ok_err is False + assert "requires at least 2 figure IDs" in res_err + + # Action: create_multipanel with 2 IDs + res_multi, ok_multi = await tool.handler(action="create_multipanel", figure_ids=[fig_id, fig_id]) + assert ok_multi is True + assert "Multi-Panel Subfigure Grid Created" in res_multi diff --git a/backend/tests/test_tools_models.py b/backend/tests/test_tools_models.py new file mode 100644 index 0000000..ba56d08 --- /dev/null +++ b/backend/tests/test_tools_models.py @@ -0,0 +1,105 @@ +"""Tests for Models Agent Tool Spec & Actions.""" + +import pytest + +from openmlr.services.model_registry import ModelRegistryService +from openmlr.tools.models import create_models_tool + + +@pytest.fixture(autouse=True) +def clean_registry(): + ModelRegistryService._models_store.clear() + yield + ModelRegistryService._models_store.clear() + + +@pytest.mark.asyncio +async def test_tool_register_and_list(): + tool = create_models_tool() + assert tool.handler is not None + res, ok = await tool.handler( + action="register", + project_id="test_proj", + name="BERT-Base-OpenMLR", + version="1.0.0", + architecture="Encoder", + framework="pytorch", + parameters_count=110_000_000, + model_size_mb=440.0, + metrics={"f1": 0.89}, + ) + assert ok is True + assert "registered successfully" in res + assert "BERT-Base-OpenMLR" in res + + # List models + list_res, list_ok = await tool.handler(action="list", project_id="test_proj") + assert list_ok is True + assert "BERT-Base-OpenMLR" in list_res + + +@pytest.mark.asyncio +async def test_tool_checkpoint_inspection(): + tool = create_models_tool() + assert tool.handler is not None + res, ok = await tool.handler( + action="inspect_checkpoint", + checkpoint_path="weights.safetensors", + parameters_count=7_000_000_000, + ) + assert ok is True + assert "Checkpoint Inspection Report" in res + assert "Estimated VRAM" in res + + +@pytest.mark.asyncio +async def test_tool_generate_card_and_quant(): + tool = create_models_tool() + assert tool.handler is not None + # Register first + await tool.handler( + action="register", + project_id="p_tool", + name="ViT-Base-224", + parameters_count=86_000_000, + metrics={"top1_acc": 0.84}, + ) + models = ModelRegistryService.list_models("p_tool") + mid = models[0].id + + # Generate card + card_res, card_ok = await tool.handler( + action="generate_card", + project_id="p_tool", + model_id=mid, + author="Silas", + ) + assert card_ok is True + assert "# Model Card: ViT-Base-224" in card_res + + # Quantization planning + quant_res, quant_ok = await tool.handler( + action="plan_quantization", + project_id="p_tool", + model_id=mid, + target_precisions=["fp16", "int8", "int4"], + ) + assert quant_ok is True + assert "INT4" in quant_res + + +@pytest.mark.asyncio +async def test_tool_compare(): + tool = create_models_tool() + assert tool.handler is not None + await tool.handler(action="register", project_id="p_comp", name="Model-1", metrics={"accuracy": 0.80}) + await tool.handler(action="register", project_id="p_comp", name="Model-2", metrics={"accuracy": 0.90}) + + models = ModelRegistryService.list_models("p_comp") + res, ok = await tool.handler( + action="compare", + project_id="p_comp", + model_ids=[models[0].id, models[1].id], + ) + assert ok is True + assert "Model Comparison Analysis" in res diff --git a/backend/tests/test_tools_reproducibility.py b/backend/tests/test_tools_reproducibility.py new file mode 100644 index 0000000..64e8d53 --- /dev/null +++ b/backend/tests/test_tools_reproducibility.py @@ -0,0 +1,93 @@ +"""Unit tests for the Reproducibility Agent Tool.""" + +import json + +import pytest + +from openmlr.tools.reproducibility import create_reproducibility_tool + + +@pytest.mark.asyncio +async def test_reproducibility_tool_audit(): + tool = create_reproducibility_tool(get_project_context=lambda: "proj_unit") + assert tool.name == "reproducibility" + assert tool.handler is not None + + code_snippets = { + "main.py": "import torch\ntorch.manual_seed(42)\ntorch.backends.cudnn.deterministic = True\ntorch.save({}, 'model.pt')", + "requirements.txt": "torch==2.1.0\n", + } + raw_res, ok = await tool.handler( + action="audit", + code_snippets=code_snippets, + venue="neurips", + ) + assert ok is True + res = json.loads(raw_res) + assert res["status"] == "success" + assert "report_id" in res + assert res["overall_score"] > 50.0 + assert "grade" in res + assert "checklist_summary" in res + + +@pytest.mark.asyncio +async def test_reproducibility_tool_generate_dockerfile(): + tool = create_reproducibility_tool() + assert tool.handler is not None + raw_res, ok = await tool.handler( + action="generate_dockerfile", + framework="pytorch", + cuda_version="12.2.0", + requirements=["torch==2.2.0", "numpy==1.26.0"], + ) + assert ok is True + res = json.loads(raw_res) + assert res["status"] == "success" + assert "FROM nvidia/cuda:12.2.0" in res["dockerfile"] + assert "torch==2.2.0" in res["dockerfile"] + + +@pytest.mark.asyncio +async def test_reproducibility_tool_fix_determinism(): + tool = create_reproducibility_tool() + assert tool.handler is not None + raw_res, ok = await tool.handler( + action="fix_determinism", + framework="pytorch", + seed=1337, + ) + assert ok is True + res = json.loads(raw_res) + assert res["status"] == "success" + assert "torch.manual_seed(seed)" in res["determinism_snippet"] + assert "set_seed(1337)" in res["determinism_snippet"] + + +@pytest.mark.asyncio +async def test_reproducibility_tool_generate_appendix(): + tool = create_reproducibility_tool() + assert tool.handler is not None + raw_res, ok = await tool.handler( + action="generate_appendix", + paper_title="Autonomous Agent Paper", + random_seeds=[42, 100], + ) + assert ok is True + res = json.loads(raw_res) + assert res["status"] == "success" + assert "\\section{Reproducibility Statement}" in res["latex_appendix"] + + +@pytest.mark.asyncio +async def test_reproducibility_tool_list_and_unknown(): + tool = create_reproducibility_tool() + assert tool.handler is not None + raw_res, ok = await tool.handler(action="list") + assert ok is True + res = json.loads(raw_res) + assert res["status"] == "success" + + raw_err, ok_err = await tool.handler(action="invalid_action") + assert ok_err is False + assert "Unknown action" in raw_err diff --git a/backend/tests/test_tools_sweeps.py b/backend/tests/test_tools_sweeps.py new file mode 100644 index 0000000..c0b74ab --- /dev/null +++ b/backend/tests/test_tools_sweeps.py @@ -0,0 +1,87 @@ +"""Tests for sweeps agent tool.""" + +import json +from pathlib import Path + +import pytest + +from openmlr.tools.sweeps import create_sweeps_tool + + +@pytest.mark.asyncio +async def test_sweeps_tool_lifecycle(tmp_path: Path): + tool = create_sweeps_tool(get_project_id=lambda: "test_proj", base_dir=tmp_path / "sweeps") + handler = tool.handler + assert handler is not None + + # 1. Create sweep + create_res, ok = await handler( + action="create_sweep", + name="Agent ResNet Sweep", + method="random", + objective_metric="val_loss", + goal="minimize", + parameters={ + "lr": {"param_type": "loguniform", "min_val": 1e-4, "max_val": 1e-2}, + "batch_size": {"param_type": "choice", "choices": [32, 64]}, + }, + max_trials=3, + ) + assert ok is True + assert "Agent ResNet Sweep" in create_res + + # 2. List sweeps + list_res, ok = await handler(action="list_sweeps") + assert ok is True + assert "Agent ResNet Sweep" in list_res + + # Extract sweep_id from list or create + # 3. Suggest trial + # First get sweep id from listing + from openmlr.services.sweep_engine import SweepEngine + engine = SweepEngine(base_dir=tmp_path / "sweeps") + sweeps = engine.list_sweeps("test_proj") + assert len(sweeps) == 1 + sweep_id = sweeps[0].sweep_id + + suggest_res, ok = await handler(action="suggest_trial", sweep_id=sweep_id) + assert ok is True + assert "Suggested Next Trial" in suggest_res + + # 4. Record trial + sweep = engine.get_sweep("test_proj", sweep_id) + assert sweep is not None + trial_id = sweep.trials[0].trial_id + + record_res, ok = await handler( + action="record_trial", + sweep_id=sweep_id, + trial_id=trial_id, + metrics={"val_loss": 0.31, "accuracy": 0.89}, + status="completed", + ) + assert ok is True + assert "recorded" in record_res + + # 5. Prune check + prune_res, ok = await handler( + action="prune_check", + sweep_id=sweep_id, + trial_id=trial_id, + current_step=5, + current_metric_val=0.31, + ) + assert ok is True + assert "Prune evaluation" in prune_res + + # 6. Analyze + analyze_res, ok = await handler(action="analyze_sweep", sweep_id=sweep_id) + assert ok is True + analysis_data = json.loads(analyze_res) + assert analysis_data["completed_trials"] == 1 + assert analysis_data["best_metric_value"] == 0.31 + + # 7. Export report + export_res, ok = await handler(action="export_report", sweep_id=sweep_id) + assert ok is True + assert "Hyperparameter Optimization Report" in export_res diff --git a/frontend/src/api.ts b/frontend/src/api.ts index 276867b..2bd8359 100644 --- a/frontend/src/api.ts +++ b/frontend/src/api.ts @@ -170,6 +170,19 @@ export const api = { testMcpServer: (url: string, headers?: Record, params?: Record) => post('/api/mcp/test', { url, headers: headers || null, params: params || null }), + // Peer Review Simulation + getReviewRubrics: () => get('/api/review/rubrics'), + evaluateSubmission: (body: { + submission_text: string; + venue?: string; + title?: string; + context?: Record; + }) => post('/api/review/evaluate', body), + reviewProjectWorkspace: ( + projectId: number, + body: { venue?: string; include_latex?: boolean; include_notes?: boolean } + ) => post(`/api/projects/${projectId}/review`, body), + // Evaluation & Benchmark Harness listEvalSuites: () => get('/api/eval/suites'), listEvalTasks: (category?: string) => @@ -187,4 +200,265 @@ export const api = { post('/api/eval/custom-task/reproduction', body), registerCustomOptimizationTask: (body: Record) => post('/api/eval/custom-task/optimization', body), + + // Research State Machine & Orchestrator + getResearchPhases: () => get('/api/research/phases'), + getResearchGuidelines: () => get('/api/research/guidelines'), + getProjectResearchState: (projectId: number) => + get(`/api/projects/${projectId}/research/state`), + startProjectResearch: ( + projectId: number, + body: { goal: string; initial_phase?: string; generate_default_milestones?: boolean } + ) => post(`/api/projects/${projectId}/research/start`, body), + transitionProjectResearchPhase: ( + projectId: number, + body: { next_phase: string; reason: string; artifacts_produced?: string[]; milestone_id?: string } + ) => post(`/api/projects/${projectId}/research/transition`, body), + createResearchMilestone: ( + projectId: number, + body: { title: string; description?: string; phase?: string; criteria?: string[] } + ) => post(`/api/projects/${projectId}/research/milestones`, body), + updateResearchMilestone: ( + projectId: number, + milestoneId: string, + body: { status?: string; output_artifacts?: string[] } + ) => put(`/api/projects/${projectId}/research/milestones/${encodeURIComponent(milestoneId)}`, body), + addResearchArtifact: ( + projectId: number, + body: { type: string; data: unknown; section_name?: string } + ) => post(`/api/projects/${projectId}/research/artifacts`, body), + + // Machine Learning Experiments & Runs + listExperimentRuns: (params?: { + projectUuid?: string; + status?: string; + search?: string; + limit?: number; + offset?: number; + }) => { + const q = new URLSearchParams(); + if (params?.projectUuid) q.append('project_uuid', params.projectUuid); + if (params?.status) q.append('status', params.status); + if (params?.search) q.append('search', params.search); + if (params?.limit) q.append('limit', String(params.limit)); + if (params?.offset) q.append('offset', String(params.offset)); + const qs = q.toString(); + return get(`/api/experiments/runs${qs ? `?${qs}` : ''}`); + }, + createExperimentRun: (body: { + name: string; + description?: string; + hyperparameters?: Record; + compute_target?: string; + tags?: string[]; + total_steps?: number; + total_epochs?: number; + project_uuid?: string; + }) => post('/api/experiments/runs', body), + getExperimentRun: (runId: string, projectUuid?: string) => + get(`/api/experiments/runs/${encodeURIComponent(runId)}${projectUuid ? `?project_uuid=${encodeURIComponent(projectUuid)}` : ''}`), + logRunMetrics: ( + runId: string, + body: { step: number; epoch?: number; metrics: Record; timestamp?: number }, + projectUuid?: string + ) => post(`/api/experiments/runs/${encodeURIComponent(runId)}/metrics${projectUuid ? `?project_uuid=${encodeURIComponent(projectUuid)}` : ''}`, body), + updateRunStatus: ( + runId: string, + body: { status: string; reason?: string }, + projectUuid?: string + ) => post(`/api/experiments/runs/${encodeURIComponent(runId)}/status${projectUuid ? `?project_uuid=${encodeURIComponent(projectUuid)}` : ''}`, body), + appendRunLogs: (runId: string, lines: string[], projectUuid?: string) => + post(`/api/experiments/runs/${encodeURIComponent(runId)}/logs${projectUuid ? `?project_uuid=${encodeURIComponent(projectUuid)}` : ''}`, { lines }), + getRunLogs: (runId: string, limit = 200, projectUuid?: string) => + get(`/api/experiments/runs/${encodeURIComponent(runId)}/logs?limit=${limit}${projectUuid ? `&project_uuid=${encodeURIComponent(projectUuid)}` : ''}`), + registerRunCheckpoint: ( + runId: string, + body: { + name: string; + step: number; + epoch?: number; + path?: string; + file_size_mb?: number; + metrics?: Record; + download_url?: string; + }, + projectUuid?: string + ) => post(`/api/experiments/runs/${encodeURIComponent(runId)}/checkpoints${projectUuid ? `?project_uuid=${encodeURIComponent(projectUuid)}` : ''}`, body), + compareExperimentRuns: (runIds: string[], projectUuid?: string) => + get(`/api/experiments/compare?run_ids=${encodeURIComponent(runIds.join(','))}${projectUuid ? `&project_uuid=${encodeURIComponent(projectUuid)}` : ''}`), + deleteExperimentRun: (runId: string, projectUuid?: string) => + del(`/api/experiments/runs/${encodeURIComponent(runId)}${projectUuid ? `?project_uuid=${encodeURIComponent(projectUuid)}` : ''}`), + + // Datasets Management & Profiling + profileDataset: (body: { path: string; sample_size?: number }) => + post('/api/datasets/profile', body), + inspectDatasetSamples: (body: { + path: string; + n?: number; + offset?: number; + strategy?: string; + label_column?: string; + }) => post('/api/datasets/inspect', body), + validateDataset: (body: { + path: string; + expected_columns?: string[]; + max_null_pct?: number; + max_token_length?: number; + }) => post('/api/datasets/validate', body), + splitDataset: (body: { + path: string; + output_dir: string; + train_ratio?: number; + val_ratio?: number; + test_ratio?: number; + stratify_column?: string; + seed?: number; + }) => post('/api/datasets/split', body), + + // Hyperparameter Sweeps & HPO + listSweeps: (projectUuid?: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return get(`/api/sweeps${q}`); + }, + createSweep: (projectUuid: string | undefined, body: Record) => + post('/api/sweeps', { ...body, project_uuid: projectUuid }), + getSweep: (projectUuid: string | undefined, sweepId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return get(`/api/sweeps/${encodeURIComponent(sweepId)}${q}`); + }, + suggestTrial: (projectUuid: string | undefined, sweepId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post(`/api/sweeps/${encodeURIComponent(sweepId)}/suggest${q}`, {}); + }, + recordTrial: ( + projectUuid: string | undefined, + sweepId: string, + trialId: string, + body: { + metrics: Record; + status?: string; + step_history?: Record[]; + error_message?: string; + } + ) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post( + `/api/sweeps/${encodeURIComponent(sweepId)}/trials/${encodeURIComponent(trialId)}/record${q}`, + body + ); + }, + checkPrune: ( + projectUuid: string | undefined, + sweepId: string, + trialId: string, + body: { current_step: number; current_metric_val: number } + ) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post( + `/api/sweeps/${encodeURIComponent(sweepId)}/trials/${encodeURIComponent(trialId)}/prune-check${q}`, + body + ); + }, + getSweepAnalysis: (projectUuid: string | undefined, sweepId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return get(`/api/sweeps/${encodeURIComponent(sweepId)}/analysis${q}`); + }, + exportSweepReport: (projectUuid: string | undefined, sweepId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post(`/api/sweeps/${encodeURIComponent(sweepId)}/export${q}`, {}); + }, + deleteSweep: (projectUuid: string | undefined, sweepId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return del(`/api/sweeps/${encodeURIComponent(sweepId)}${q}`); + }, + + // Model Registry & Governance + listRegisteredModels: (projectUuid?: string, filters?: { task_type?: string; framework?: string; status?: string; tag?: string }) => { + const params = new URLSearchParams(); + if (projectUuid) params.set('project_id', projectUuid); + if (filters?.task_type) params.set('task_type', filters.task_type); + if (filters?.framework) params.set('framework', filters.framework); + if (filters?.status) params.set('status', filters.status); + if (filters?.tag) params.set('tag', filters.tag); + const qs = params.toString() ? `?${params.toString()}` : ''; + return get(`/api/model-registry${qs}`); + }, + registerModel: (projectUuid: string | undefined, body: Record) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post(`/api/model-registry${q}`, body); + }, + getRegisteredModel: (projectUuid: string | undefined, modelId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return get(`/api/model-registry/${encodeURIComponent(modelId)}${q}`); + }, + updateRegisteredModel: (projectUuid: string | undefined, modelId: string, body: Record) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return put(`/api/model-registry/${encodeURIComponent(modelId)}${q}`, body); + }, + deleteRegisteredModel: (projectUuid: string | undefined, modelId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return del(`/api/model-registry/${encodeURIComponent(modelId)}${q}`); + }, + generateModelCard: (projectUuid: string | undefined, modelId: string, body: Record) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post(`/api/model-registry/${encodeURIComponent(modelId)}/card${q}`, body); + }, + planModelQuantization: (projectUuid: string | undefined, modelId: string, targetPrecisions: string[]) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post(`/api/model-registry/${encodeURIComponent(modelId)}/quantization${q}`, { target_precisions: targetPrecisions }); + }, + inspectCheckpoint: (body: { checkpoint_path: string; parameters_count?: number; model_size_mb?: number; framework?: string }) => + post('/api/model-registry/inspect', body), + compareRegisteredModels: (projectUuid: string | undefined, modelIds: string[]) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post(`/api/model-registry/compare${q}`, { model_ids: modelIds }); + }, + + // Publication Figures & Plots + listFigures: (projectUuid?: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return get(`/api/figures${q}`); + }, + generateFigure: (projectUuid: string | undefined, body: Record) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post(`/api/figures${q}`, body); + }, + getFigure: (projectUuid: string | undefined, figureId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return get(`/api/figures/${encodeURIComponent(figureId)}${q}`); + }, + deleteFigure: (projectUuid: string | undefined, figureId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return del(`/api/figures/${encodeURIComponent(figureId)}${q}`); + }, + createMultiPanelLayout: (projectUuid: string | undefined, body: Record) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post(`/api/figures/multi-panel${q}`, body); + }, + + // Reproducibility Studio & Artifact Verification + listReproducibilityReports: (projectUuid?: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return get(`/api/reproducibility/reports${q}`); + }, + getReproducibilityReport: (projectUuid: string | undefined, reportId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return get(`/api/reproducibility/reports/${encodeURIComponent(reportId)}${q}`); + }, + runReproducibilityAudit: (projectUuid: string | undefined, body: Record | object) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post(`/api/reproducibility/audit${q}`, body as Record); + }, + generateReproducibilityDockerfile: (body: Record | object) => + post('/api/reproducibility/dockerfile', body as Record), + generateReproducibilityAppendix: (projectUuid: string | undefined, body: Record | object) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return post(`/api/reproducibility/appendix${q}`, body as Record); + }, + getDeterminismFix: (body: Record | object) => + post('/api/reproducibility/fix-determinism', body as Record), + deleteReproducibilityReport: (projectUuid: string | undefined, reportId: string) => { + const q = projectUuid ? `?project_id=${encodeURIComponent(projectUuid)}` : ''; + return del(`/api/reproducibility/reports/${encodeURIComponent(reportId)}${q}`); + }, }; diff --git a/frontend/src/components/chat/ChatContainer.tsx b/frontend/src/components/chat/ChatContainer.tsx index efea101..1b8038c 100644 --- a/frontend/src/components/chat/ChatContainer.tsx +++ b/frontend/src/components/chat/ChatContainer.tsx @@ -6,7 +6,7 @@ import { TodoReviewDrawer } from '../TodoReviewDrawer'; import { QuestionDrawer } from '../QuestionDrawer'; import { ImageViewer } from '../ImageViewer'; import { useChat } from '../../context/ChatContext'; -import { useProject } from '../../context/ProjectContext'; +import { useProject, type MainTab } from '../../context/ProjectContext'; import { useCompute } from '../../context/ComputeContext'; import { nextMsgId } from '../../context/agentEventReducers'; @@ -19,6 +19,47 @@ const CitationGraph = lazy(() => const RunDashboard = lazy(() => import('../experiments/RunDashboard').then((m) => ({ default: m.RunDashboard })) ); +const PeerReviewStudio = lazy(() => + import('../review/PeerReviewStudio').then((m) => ({ default: m.PeerReviewStudio })) +); +const EvalBenchmarkDashboard = lazy(() => + import('../eval/EvalBenchmarkDashboard').then((m) => ({ default: m.EvalBenchmarkDashboard })) +); +const ResearchWorkflowStudio = lazy(() => + import('../research/ResearchWorkflowStudio').then((m) => ({ default: m.ResearchWorkflowStudio })) +); +const DatasetStudio = lazy(() => + import('../datasets/DatasetStudio').then((m) => ({ default: m.DatasetStudio })) +); +const SweepStudio = lazy(() => + import('../sweeps/SweepStudio').then((m) => ({ default: m.SweepStudio })) +); +const ModelStudio = lazy(() => + import('../models/ModelStudio').then((m) => ({ default: m.ModelStudio })) +); +const FigureStudio = lazy(() => + import('../figures/FigureStudio').then((m) => ({ default: m.FigureStudio })) +); +const ReproducibilityStudio = lazy(() => + import('../reproducibility/ReproducibilityStudio').then((m) => ({ default: m.ReproducibilityStudio })) +); + +const NAVIGATION_TABS: { id: MainTab; label: string }[] = [ + { id: 'agent', label: 'Agent' }, + { id: 'workflow', label: 'Workflow' }, + { id: 'editor', label: 'Editor' }, + { id: 'terminal', label: 'Terminal' }, + { id: 'paper', label: 'Paper Studio' }, + { id: 'research', label: 'Citation Graph' }, + { id: 'experiments', label: 'Experiments' }, + { id: 'datasets', label: 'Datasets' }, + { id: 'sweeps', label: 'Sweeps' }, + { id: 'models', label: 'Models' }, + { id: 'figures', label: 'Figures' }, + { id: 'reproducibility', label: 'Reproducibility' }, + { id: 'review', label: 'Peer Review' }, + { id: 'eval', label: 'Benchmarks' }, +]; export function ChatContainer() { const { @@ -60,78 +101,31 @@ export function ChatContainer() {
{/* Agent / Editor / Terminal / Paper / Research tab bar */}
- - - - - - + {NAVIGATION_TABS.map((tab) => { + const isSelected = mainTab === tab.id; + return ( + + ); + })} {/* Closable Image tab */} {imageTab && (
+ {/* Workflow Studio tab */} +
+ Loading Workflow Studio...
}> + + +
+ {/* Editor tab */}
Loading...
}> @@ -279,6 +280,55 @@ export function ChatContainer() {
+ {/* Dataset Studio tab */} +
+ Loading Dataset Studio...
}> + + + + + {/* Sweep Studio tab */} +
+ Loading Sweep Studio...
}> + + + + + {/* Models Studio tab */} +
+ Loading Model Studio...
}> + + + + + {/* Publication Figures Studio tab */} +
+ Loading Figure Studio...
}> + + + + + {/* Reproducibility Studio tab */} +
+ Loading Reproducibility Studio...
}> + + + + + {/* Peer Review tab */} +
+ Loading Peer Review Studio...
}> + + + + + {/* Evaluation Benchmark tab */} +
+ Loading Benchmark Harness...
}> + + + + {/* Image tab */} {mainTab === 'image' && imageTab && ( diff --git a/frontend/src/components/datasets/ColumnMetricsTable.tsx b/frontend/src/components/datasets/ColumnMetricsTable.tsx new file mode 100644 index 0000000..c0597b7 --- /dev/null +++ b/frontend/src/components/datasets/ColumnMetricsTable.tsx @@ -0,0 +1,167 @@ +import { useState, useMemo } from 'react'; +import { Search, ChevronDown, ChevronRight, Hash, Type, ToggleLeft, Layers } from 'lucide-react'; +import type { ColumnProfile } from './types'; + +interface Props { + readonly columns: Record; +} + +function getProgressColor(pct: number): string { + if (pct > 20) return 'bg-red-500'; + if (pct > 5) return 'bg-amber-500'; + return 'bg-emerald-500'; +} + +function getTypeIcon(dtype: string) { + switch (dtype.toLowerCase()) { + case 'numeric': + return ; + case 'text': + return ; + case 'categorical': + return ; + case 'boolean': + return ; + default: + return ; + } +} + +export function ColumnMetricsTable({ columns }: Props) { + const [search, setSearch] = useState(''); + const [selectedType, setSelectedType] = useState('all'); + const [expandedCol, setExpandedCol] = useState(null); + + const columnList = useMemo(() => Object.values(columns), [columns]); + + const filteredColumns = useMemo(() => { + return columnList.filter((col) => { + const matchesSearch = col.name.toLowerCase().includes(search.toLowerCase()); + const matchesType = selectedType === 'all' || col.dtype.toLowerCase() === selectedType.toLowerCase(); + return matchesSearch && matchesType; + }); + }, [columnList, search, selectedType]); + + return ( +
+ {/* Controls */} +
+
+ + setSearch(e.target.value)} + className="w-full bg-surface-subtle border border-border rounded-md pl-9 pr-3 py-1.5 text-xs text-text placeholder-text-dim focus:outline-none focus:border-primary" + /> +
+
+ {['all', 'numeric', 'text', 'categorical', 'boolean'].map((t) => ( + + ))} +
+
+ + {/* Table */} +
+
+ + + + + + + + + + + + {filteredColumns.length === 0 ? ( + + + + ) : ( + filteredColumns.map((col) => { + const isExpanded = expandedCol === col.name; + const stats = col.stats || {}; + return ( + + + + + + + + + ); + }) + )} + +
+ Column NameTypeMissing ValuesUnique ValuesStatistical Summary
+ No matching columns found. +
+ + {col.name} + + {getTypeIcon(col.dtype)} + {col.dtype} + + +
+
+
+
+ + {col.null_percentage}% ({col.null_count}/{col.total_count}) + +
+
+ {col.unique_count.toLocaleString()} + + {col.dtype === 'numeric' && ( + + Range: [{stats.min ?? '?'}, {stats.max ?? '?'}], μ={stats.mean ?? '?'}, σ={stats.std ?? '?'} + + )} + {col.dtype === 'categorical' && ( + + Imbalance: {stats.imbalance_ratio ?? 1}x | Top: {Object.keys(stats.top_classes || {}).slice(0, 3).join(', ')} + + )} + {col.dtype === 'text' && ( + + Mean Tokens: {stats.token_est_mean ?? '?'}, Max: {stats.token_est_max ?? '?'} + + )} + {col.dtype === 'boolean' && ( + True Count: {stats.true_count ?? '?'} + )} +
+
+
+
+ ); +} diff --git a/frontend/src/components/datasets/DatasetSplitterModal.tsx b/frontend/src/components/datasets/DatasetSplitterModal.tsx new file mode 100644 index 0000000..dd3f69b --- /dev/null +++ b/frontend/src/components/datasets/DatasetSplitterModal.tsx @@ -0,0 +1,220 @@ +import { useState } from 'react'; +import { X, Scissors, RefreshCw, CheckCircle2, AlertCircle } from 'lucide-react'; +import { api } from '../../api'; +import type { SplitManifest } from './types'; + +interface Props { + readonly filePath: string; + readonly availableColumns: string[]; + readonly onClose: () => void; +} + +export function DatasetSplitterModal({ filePath, availableColumns, onClose }: Props) { + const defaultOutDir = `${filePath.replace(/\.[^/.]+$/, '')}_splits`; + const [outputDir, setOutputDir] = useState(defaultOutDir); + const [trainRatio, setTrainRatio] = useState(0.8); + const [valRatio, setValRatio] = useState(0.1); + const [testRatio, setTestRatio] = useState(0.1); + const [stratifyCol, setStratifyCol] = useState(''); + const [seed, setSeed] = useState(42); + const [loading, setLoading] = useState(false); + const [manifest, setManifest] = useState(null); + const [error, setError] = useState(null); + + const handleSplit = async () => { + if (!filePath || !outputDir) return; + setLoading(true); + setError(null); + try { + const res = await api.splitDataset({ + path: filePath, + output_dir: outputDir, + train_ratio: trainRatio, + val_ratio: valRatio, + test_ratio: testRatio, + stratify_column: stratifyCol || undefined, + seed, + }); + if (res?.manifest) { + setManifest(res.manifest); + } + } catch (err: unknown) { + setError(err instanceof Error ? err.message : 'Dataset partitioning failed'); + } finally { + setLoading(false); + } + }; + + return ( +
+
+
+
+ +

Dataset Partition Splitter

+
+ +
+ + {/* Source info */} +
+ Source file: + {filePath} +
+ + {/* Form */} +
+
+ + setOutputDir(e.target.value)} + className="w-full bg-surface-subtle border border-border rounded px-3 py-1.5 font-mono text-xs text-text focus:outline-none focus:border-primary" + /> +
+ + {/* Ratio Split */} +
+
+ Split Ratios: + + Train: {(trainRatio * 100).toFixed(0)}% | Val: {(valRatio * 100).toFixed(0)}% | Test:{' '} + {(testRatio * 100).toFixed(0)}% + +
+
+
+ + setTrainRatio(Number(e.target.value))} + className="w-full bg-surface-subtle border border-border rounded px-2 py-1 text-xs text-text" + /> +
+
+ + setValRatio(Number(e.target.value))} + className="w-full bg-surface-subtle border border-border rounded px-2 py-1 text-xs text-text" + /> +
+
+ + setTestRatio(Number(e.target.value))} + className="w-full bg-surface-subtle border border-border rounded px-2 py-1 text-xs text-text" + /> +
+
+
+ + {/* Stratification */} +
+
+ + +
+ +
+ + setSeed(Number(e.target.value))} + className="w-full bg-surface-subtle border border-border rounded px-3 py-1.5 text-xs text-text focus:outline-none focus:border-primary" + /> +
+
+
+ + {error && ( +
+ + {error} +
+ )} + + {manifest && ( +
+
+ Partition splits generated successfully! +
+
+ Train: {manifest.train_count.toLocaleString()} | Val: {manifest.val_count.toLocaleString()} | Test:{' '} + {manifest.test_count.toLocaleString()} +
+
+ )} + +
+ + +
+
+
+ ); +} diff --git a/frontend/src/components/datasets/DatasetStudio.test.tsx b/frontend/src/components/datasets/DatasetStudio.test.tsx new file mode 100644 index 0000000..4bae06e --- /dev/null +++ b/frontend/src/components/datasets/DatasetStudio.test.tsx @@ -0,0 +1,204 @@ +import { render, screen, fireEvent, waitFor } from '@testing-library/react'; +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { DatasetStudio } from './DatasetStudio'; +import { api } from '../../api'; + +vi.mock('../../api', () => ({ + api: { + profileDataset: vi.fn(), + inspectDatasetSamples: vi.fn(), + validateDataset: vi.fn(), + splitDataset: vi.fn(), + }, +})); + +vi.mock('../../context/ProjectContext', () => ({ + useProject: () => ({ + activeProject: { uuid: 'proj-123', name: 'Vision LLM Research' }, + }), +})); + +const mockProfile = { + success: true, + profile: { + file_path: 'data.csv', + format: 'csv', + total_rows: 1500, + total_columns: 4, + file_size_bytes: 204850, + health_score: 92, + warnings: ['Column text has 5% missing values'], + summary: 'data.csv: 1500 rows, 4 cols. Health: 92/100.', + columns: { + id: { + name: 'id', + dtype: 'numeric', + total_count: 1500, + null_count: 0, + null_percentage: 0.0, + unique_count: 1500, + stats: { min: 1, max: 1500, mean: 750.5, std: 433.0 }, + }, + label: { + name: 'label', + dtype: 'categorical', + total_count: 1500, + null_count: 0, + null_percentage: 0.0, + unique_count: 3, + stats: { top_classes: { positive: 800, negative: 500, neutral: 200 }, imbalance_ratio: 4.0 }, + }, + text: { + name: 'text', + dtype: 'text', + total_count: 1500, + null_count: 75, + null_percentage: 5.0, + unique_count: 1420, + stats: { char_len_avg: 120.4, token_est_mean: 32.5, token_est_max: 128 }, + }, + active: { + name: 'active', + dtype: 'boolean', + total_count: 1500, + null_count: 0, + null_percentage: 0.0, + unique_count: 2, + stats: { true_count: 1200 }, + }, + }, + }, +}; + +const mockSamples = { + success: true, + total_sampled: 2, + samples: [ + { id: 1, label: 'positive', text: 'Excellent benchmark accuracy', active: true }, + { id: 2, label: 'negative', text: 'High generalization error', active: false }, + ], +}; + +const mockValidation = { + success: true, + validation: { + valid: true, + errors: [], + warnings: [], + health_score: 95, + total_rows: 1500, + total_columns: 4, + }, +}; + +const mockSplitManifest = { + success: true, + manifest: { + source_file: 'data.csv', + stratified_by: 'label', + seed: 42, + total_records: 1500, + train_count: 1200, + val_count: 150, + test_count: 150, + splits: { + train: 'data_splits/train.csv', + val: 'data_splits/val.csv', + test: 'data_splits/test.csv', + }, + }, +}; + +describe('DatasetStudio Component', () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(api.profileDataset).mockResolvedValue(mockProfile); + vi.mocked(api.inspectDatasetSamples).mockResolvedValue(mockSamples); + vi.mocked(api.validateDataset).mockResolvedValue(mockValidation); + vi.mocked(api.splitDataset).mockResolvedValue(mockSplitManifest); + }); + + it('renders dataset header, overview metrics, and column stats table', async () => { + render(); + + expect(screen.getByText('Dataset Studio')).toBeDefined(); + expect(screen.getByText('Vision LLM Research')).toBeDefined(); + + await waitFor(() => { + expect(api.profileDataset).toHaveBeenCalledWith({ path: 'data.csv' }); + expect(screen.getAllByText('1,500').length).toBeGreaterThan(0); + expect(screen.getByText('92/100')).toBeDefined(); + }); + + // Column names + expect(screen.getByText('label')).toBeDefined(); + expect(screen.getAllByText('text').length).toBeGreaterThan(0); + }); + + it('switches to Data Samples tab and renders preview rows', async () => { + render(); + + await waitFor(() => { + expect(screen.getByText('92/100')).toBeDefined(); + }); + + const samplesTab = screen.getByText('Data Samples'); + fireEvent.click(samplesTab); + + await waitFor(() => { + expect(api.inspectDatasetSamples).toHaveBeenCalled(); + expect(screen.getByText('Excellent benchmark accuracy')).toBeDefined(); + expect(screen.getByText('High generalization error')).toBeDefined(); + }); + }); + + it('switches to Schema Validator and executes validation check', async () => { + render(); + + await waitFor(() => { + expect(screen.getByText('92/100')).toBeDefined(); + }); + + const valTab = screen.getByText('Schema Validator'); + fireEvent.click(valTab); + + const runBtn = screen.getByText('Run Validation Check'); + fireEvent.click(runBtn); + + await waitFor(() => { + expect(api.validateDataset).toHaveBeenCalled(); + expect(screen.getByText('Passed')).toBeDefined(); + expect(screen.getByText('95 / 100')).toBeDefined(); + }); + }); + + it('opens Dataset Splitter modal and triggers partition generation', async () => { + render(); + + await waitFor(() => { + expect(screen.getByText('92/100')).toBeDefined(); + }); + + const splitBtn = screen.getByText('Split Partitions'); + fireEvent.click(splitBtn); + + expect(screen.getByText('Dataset Partition Splitter')).toBeDefined(); + + const generateBtn = screen.getByText('Generate Splits'); + fireEvent.click(generateBtn); + + await waitFor(() => { + expect(api.splitDataset).toHaveBeenCalled(); + expect(screen.getByText('Partition splits generated successfully!')).toBeDefined(); + }); + }); + + it('displays error banner when dataset profiling fails', async () => { + vi.mocked(api.profileDataset).mockRejectedValueOnce(new Error('File not found: missing.csv')); + render(); + + await waitFor(() => { + expect(screen.getByText('File not found: missing.csv')).toBeDefined(); + }); + }); +}); diff --git a/frontend/src/components/datasets/DatasetStudio.tsx b/frontend/src/components/datasets/DatasetStudio.tsx new file mode 100644 index 0000000..4660cef --- /dev/null +++ b/frontend/src/components/datasets/DatasetStudio.tsx @@ -0,0 +1,229 @@ +import { useState, useEffect, useCallback } from 'react'; +import { + Database, + RefreshCw, + Scissors, + AlertTriangle, + FileSpreadsheet, + Table, + CheckCircle, + Hash, +} from 'lucide-react'; +import { api } from '../../api'; +import { useProject } from '../../context/ProjectContext'; +import type { DatasetProfile } from './types'; +import { ColumnMetricsTable } from './ColumnMetricsTable'; +import { SampleDataViewer } from './SampleDataViewer'; +import { DatasetValidatorCard } from './DatasetValidatorCard'; +import { DatasetSplitterModal } from './DatasetSplitterModal'; + +function getHealthColor(score: number): string { + if (score >= 80) return 'text-emerald-400'; + if (score >= 50) return 'text-amber-400'; + return 'text-red-400'; +} + +export function DatasetStudio() { + const { activeProject } = useProject(); + const [filePath, setFilePath] = useState('data.csv'); + const [profile, setProfile] = useState(null); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + const [activeTab, setActiveTab] = useState<'columns' | 'samples' | 'validate'>('columns'); + const [showSplitModal, setShowSplitModal] = useState(false); + + const fetchProfile = useCallback(async () => { + if (!filePath.trim()) return; + setLoading(true); + setError(null); + try { + const res = await api.profileDataset({ path: filePath.trim() }); + if (res?.profile) { + setProfile(res.profile); + } + } catch (err: unknown) { + setError(err instanceof Error ? err.message : 'Failed to profile dataset'); + setProfile(null); + } finally { + setLoading(false); + } + }, [filePath]); + + useEffect(() => { + fetchProfile(); + }, [fetchProfile]); + + const columnsList = profile ? Object.keys(profile.columns) : []; + + return ( +
+ {/* Header Bar */} +
+
+
+ +
+
+

+ Dataset Studio + {activeProject && ( + + {activeProject.name} + + )} +

+

+ Statistical profiling, missingness diagnostics, schema validation, and partition splits +

+
+
+ + {/* Path Input & Actions */} +
+
+ + setFilePath(e.target.value)} + placeholder="Dataset path (e.g. data.csv, train.jsonl)..." + className="w-full bg-surface-subtle border border-border rounded-md pl-9 pr-3 py-1.5 font-mono text-xs text-text placeholder-text-dim focus:outline-none focus:border-primary" + /> +
+ + + + +
+
+ + {/* Main Content Area */} +
+ {/* Error Alert */} + {error && ( +
+ + {error} +
+ )} + + {/* Profile Overview Banner */} + {profile && ( +
+
+ Total Rows + + {profile.total_rows.toLocaleString()} + +
+
+ Columns + {profile.total_columns} +
+
+ File Size + + {(profile.file_size_bytes / (1024 * 1024)).toFixed(2)} MB + +
+
+ Health Score +
+ + {profile.health_score}/100 + +
+
+
+ )} + + {/* Warnings Banner */} + {profile && profile.warnings.length > 0 && ( +
+
+ Diagnostic Warnings Detected ({profile.warnings.length}): +
+
    + {profile.warnings.map((w) => ( +
  • {w}
    • + ))} +
+
+ )} + + {/* Tab Switcher */} +
+ +
+
+
+ ); +} diff --git a/frontend/src/components/datasets/types.ts b/frontend/src/components/datasets/types.ts new file mode 100644 index 0000000..fa0aa5f --- /dev/null +++ b/frontend/src/components/datasets/types.ts @@ -0,0 +1,69 @@ +/** + * Types for Dataset Profiler, Table Inspector, Validation, and Partition Splitter. + */ + +export interface ColumnProfile { + name: string; + dtype: 'numeric' | 'text' | 'categorical' | 'boolean' | 'unknown'; + total_count: number; + null_count: number; + null_percentage: number; + unique_count: number; + stats: { + min?: number; + max?: number; + mean?: number; + std?: number; + median?: number; + q25?: number; + q75?: number; + outlier_count?: number; + top_classes?: Record; + imbalance_ratio?: number; + class_distribution?: Record; + char_len_avg?: number; + token_est_mean?: number; + token_est_p95?: number; + token_est_max?: number; + overflow_512_count?: number; + overflow_512_pct?: number; + true_count?: number; + [key: string]: unknown; + }; +} + +export interface DatasetProfile { + file_path: string; + format: string; + total_rows: number; + total_columns: number; + file_size_bytes: number; + columns: Record; + health_score: number; + warnings: string[]; + summary: string; +} + +export interface ValidationResult { + valid: boolean; + errors: string[]; + warnings: string[]; + health_score: number; + total_rows: number; + total_columns: number; +} + +export interface SplitManifest { + source_file: string; + stratified_by: string | null; + seed: number; + total_records: number; + train_count: number; + val_count: number; + test_count: number; + splits: { + train: string; + val: string; + test: string; + }; +} diff --git a/frontend/src/components/eval/CustomTaskModal.tsx b/frontend/src/components/eval/CustomTaskModal.tsx new file mode 100644 index 0000000..ca4efbe --- /dev/null +++ b/frontend/src/components/eval/CustomTaskModal.tsx @@ -0,0 +1,335 @@ +import { useState } from 'react'; +import { X, PlusCircle, FileCheck, Zap, AlertCircle } from 'lucide-react'; +import { api } from '../../api'; + +interface Props { + onClose: () => void; + onCreated: () => void; +} + +export function CustomTaskModal({ onClose, onCreated }: Readonly) { + const [taskType, setTaskType] = useState<'reproduction' | 'optimization'>('reproduction'); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + + // Common + const [taskId, setTaskId] = useState(''); + const [name, setName] = useState(''); + const [description, setDescription] = useState(''); + const [difficulty, setDifficulty] = useState('medium'); + const [timeoutSeconds, setTimeoutSeconds] = useState(300); + + // Reproduction specific + const [paperTitle, setPaperTitle] = useState(''); + const [arxivId, setArxivId] = useState(''); + const [datasetName, setDatasetName] = useState('custom'); + const [targetMetricsJson, setTargetMetricsJson] = useState('{"accuracy": 0.85, "loss": 0.25}'); + + // Optimization specific + const [kernelName, setKernelName] = useState(''); + const [framework, setFramework] = useState('triton'); + const [baselineLatencyMs, setBaselineLatencyMs] = useState(25.0); + const [targetSpeedup, setTargetSpeedup] = useState(1.5); + + const handleSubmit = async (e: React.FormEvent) => { + e.preventDefault(); + setError(null); + setLoading(true); + + try { + if (taskType === 'reproduction') { + let metrics: Record = {}; + try { + metrics = JSON.parse(targetMetricsJson); + } catch { + throw new Error('Target metrics must be a valid JSON object of numbers (e.g. {"accuracy": 0.85})'); + } + + await api.registerCustomReproductionTask({ + task_id: taskId, + name, + description, + paper_title: paperTitle, + arxiv_id: arxivId, + dataset_name: datasetName, + target_metrics: metrics, + difficulty, + timeout_seconds: timeoutSeconds, + }); + } else { + await api.registerCustomOptimizationTask({ + task_id: taskId, + name, + description, + kernel_name: kernelName, + framework, + baseline_latency_ms: baselineLatencyMs, + target_speedup: targetSpeedup, + difficulty, + timeout_seconds: timeoutSeconds, + }); + } + + onCreated(); + onClose(); + } catch (err: unknown) { + const errMsg = err instanceof Error ? err.message : 'Failed to register benchmark task.'; + setError(errMsg); + } finally { + setLoading(false); + } + }; + + return ( +
+
+ {/* Header */} +
+
+ +

Register Custom Benchmark Task

+
+ +
+ + {/* Form content */} +
+ {/* Type Toggle */} +
+ + +
+ + {error && ( +
+ + {error} +
+ )} + + {/* Common fields */} +
+
+ + setTaskId(e.target.value)} + placeholder="e.g. repro-flashattention-3" + className="w-full bg-bg border border-border rounded-xl px-3 py-1.5 text-text focus:border-primary focus:outline-none" + /> +
+
+ + setName(e.target.value)} + placeholder="e.g. FlashAttention-3 Hopper FP8" + className="w-full bg-bg border border-border rounded-xl px-3 py-1.5 text-text focus:border-primary focus:outline-none" + /> +
+
+ +
+
+ + +
+
+ + setTimeoutSeconds(parseInt(e.target.value) || 300)} + className="w-full bg-bg border border-border rounded-xl px-3 py-1.5 text-text focus:border-primary focus:outline-none" + /> +
+
+ +
+ +