weco 0.2.19__tar.gz → 0.2.20__tar.gz

{weco-0.2.19 → weco-0.2.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.19
+Version: 0.2.20
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -29,6 +29,9 @@ Dynamic: license-file
 [![docs](https://img.shields.io/website?url=https://docs.weco.ai/&label=docs)](https://docs.weco.ai/)
 [![PyPI version](https://badge.fury.io/py/weco.svg)](https://badge.fury.io/py/weco)
 [![AIDE](https://img.shields.io/badge/AI--Driven_Exploration-arXiv-orange?style=flat-square&logo=arxiv)](https://arxiv.org/abs/2502.13138)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
+
+`pip install weco`
 
 </div>
 

{weco-0.2.19 → weco-0.2.20}/README.md

@@ -6,6 +6,9 @@
 [![docs](https://img.shields.io/website?url=https://docs.weco.ai/&label=docs)](https://docs.weco.ai/)
 [![PyPI version](https://badge.fury.io/py/weco.svg)](https://badge.fury.io/py/weco)
 [![AIDE](https://img.shields.io/badge/AI--Driven_Exploration-arXiv-orange?style=flat-square&logo=arxiv)](https://arxiv.org/abs/2502.13138)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
+
+`pip install weco`
 
 </div>
 

weco-0.2.20/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb

@@ -0,0 +1,310 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Hello Kernel World 🔥"
+      ],
+      "metadata": {
+        "id": "RO1o5fS5W8xc"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "<p align=\"left\">\n",
+        "  <img src=\"https://raw.githubusercontent.com/WecoAI/weco-cli/main/assets/example-optimization.gif\"\n",
+        "       alt=\"Optimization demo\"\n",
+        "       width=\"720\">\n",
+        "</p>\n",
+        "\n",
+        "## 🖥️ Weco CLI Resources\n",
+        "\n",
+        "- 📖 [CLI Reference](https://docs.weco.ai/cli/cli-reference) - Explore our docs for an in-depth look at what the tool can do\n",
+        "- ✨ [Examples](https://docs.weco.ai/examples) - Explore automated R&D across kernel engineering, ML engineering and prompt engineering"
+      ],
+      "metadata": {
+        "id": "yorBWlqGuC7-"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Setup Dependencies"
+      ],
+      "metadata": {
+        "id": "5BQGGqbJW2Eq"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Install requirements\n",
+        "%pip install -q weco ipywidgets numpy torch\n",
+        "\n",
+        "# Enable custom widgets\n",
+        "from google.colab import output\n",
+        "output.enable_custom_widget_manager()"
+      ],
+      "metadata": {
+        "id": "89doT3fbWcGi"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Now we need to determine what `DEVICE` we can run this on, a CPU or GPU..."
+      ],
+      "metadata": {
+        "id": "gYxLwOXzfmiF"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import torch\n",
+        "from rich import print as rprint\n",
+        "\n",
+        "# Check if you're connected to a GPU (it's free!)\n",
+        "if not torch.cuda.is_available():\n",
+        "    DEVICE = \"cpu\"\n",
+        "    rprint(\n",
+        "        \"\"\"\n",
+        "[bold yellow]⚠️  GPU is not enabled.[/bold yellow] The notebook will fall back to [bold]CPU[/bold], but [italic]performance may be lower[/italic].\n",
+        "\n",
+        "[bold]👉 To enable GPU (FREE):[/bold]\n",
+        "• Go to [green]Runtime > Change runtime type[/green]\n",
+        "• Set [bold]'Hardware Accelerator'[/bold] to [bold green]'GPU'[/bold green]\n",
+        "• Click [bold]Save[/bold] and [bold]rerun all cells[/bold]\n",
+        "\n",
+        "[dim]Continuing with CPU for now...[/dim]\n",
+        "\"\"\"\n",
+        "    )\n",
+        "else:\n",
+        "    DEVICE = \"cuda\"\n",
+        "    rprint(\"[bold green]✅ GPU is enabled.[/bold green] Proceeding with [bold green]CUDA[/bold green]...\")"
+      ],
+      "metadata": {
+        "id": "PFAn_bzAXLGO"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Download the example files from CLI repo\n",
+        "!wget https://github.com/WecoAI/weco-cli/archive/refs/heads/main.zip -O repo.zip\n",
+        "!unzip -j repo.zip \"weco-cli-main/examples/hello-kernel-world/*\" -d .\n",
+        "!rm repo.zip"
+      ],
+      "metadata": {
+        "id": "dFGClqxpzwyM"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Google AI Studio has a free API usage quota. Create a key [here](https://aistudio.google.com/apikey) to use `weco` for free!"
+      ],
+      "metadata": {
+        "id": "PGlEsI78bMtN"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import os\n",
+        "\n",
+        "# Pass your API key below\n",
+        "os.environ[\"GEMINI_API_KEY\"] = \"\"\n",
+        "# os.environ[\"OPENAI_API_KEY\"] = \"\"\n",
+        "# os.environ[\"ANTHROPIC_API_KEY\"] = \"\"\n",
+        "\n",
+        "\n",
+        "if not any([os.environ.get(key) for key in [\"GEMINI_API_KEY\", \"OPENAI_API_KEY\", \"ANTHROPIC_API_KEY\"]]):\n",
+        "    rprint(\n",
+        "\"[bold red]❌ No API keys found.[/bold red]\\n\"\n",
+        "\"\\n\"\n",
+        "\"Please set one of the following environment variables:\\n\"\n",
+        "\"  • [cyan]GEMINI_API_KEY[/cyan]\\n\"\n",
+        "\"  • [cyan]OPENAI_API_KEY[/cyan]\\n\"\n",
+        "\"  • [cyan]ANTHROPIC_API_KEY[/cyan]\\n\"\n",
+        "\"\\n\"\n",
+        "\"Setup your [cyan]GEMINI_API_KEY[/cyan] for free - [underline white]https://aistudio.google.com/apikey[/underline white] !\"\n",
+        "    )\n",
+        "else:\n",
+        "    rprint(\"[bold green]✅ API keys found.[/bold green]\\n\\nWe'll only be able to know if they are correct once the optimization starts.\")"
+      ],
+      "metadata": {
+        "id": "b4XuOeNzYTdp"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Let's Start Optimizing!"
+      ],
+      "metadata": {
+        "id": "sbvA8oQceOt5"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Now that we've got our dependecies, GPU and LLM API key sorted out, let's take a look at what code we're optimizing!\n",
+        "\n",
+        "Earlier, we downloaded two files:\n",
+        "1. An evaluation script to help score *how good a solution is* (`evaluate.py`)\n",
+        "2. A snippet of code we'd like to optimize (`optimize.py`)\n",
+        "\n",
+        "Let's take a look at what the code we want to optimize looks like..."
+      ],
+      "metadata": {
+        "id": "4OjXTBkjc4Id"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from IPython.display import display, HTML\n",
+        "from pygments import highlight\n",
+        "from pygments.lexers import PythonLexer\n",
+        "from pygments.formatters import HtmlFormatter\n",
+        "\n",
+        "def view_code_block(path: str):\n",
+        "    with open(path) as f:\n",
+        "        display(HTML(highlight(f.read(), PythonLexer(), HtmlFormatter(full=True, style=\"monokai\"))))\n",
+        "\n",
+        "view_code_block(\"optimize.py\")"
+      ],
+      "metadata": {
+        "id": "rUTxqxWgcC34"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Real-world code is often more complex but this is a good place to start. You can find more advanced examples [here](https://docs.weco.ai/examples), however, we'd recommend starting with this notebook as the optimization setup is the exact same, no matter the complexity!"
+      ],
+      "metadata": {
+        "id": "5C5dvasXdmNw"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "It's simple to start optimizing any piece of code! You just need to set:\n",
+        "1. Path to source code - we can point this to our `optimize.py`\n",
+        "2. Command to run evaluation - notice how we are using the `DEVICE` we setup earlier\n",
+        "3. The metric we are optimizing for - in this case, the evaluation script (`evaluate.py`) prints the `'speedup'` achieved to the terminal\n",
+        "4. Whether you want to maximize or minimize the metric you mentioned above - in our case, we want to make this code faster!\n",
+        "5. Number of steps to optimize for - we'll keep it low to avoid any rate limits being hit on your free Gemini API key\n",
+        "6. Additional context - anything information you think should guide the optimization process"
+      ],
+      "metadata": {
+        "id": "YfDg-pP9fAdC"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Now let's get straight into it. Keep an eye on the `Best Solution` panel!\n",
+        "\n",
+        "Note that you can track the optimization in the logs directory (`.runs/`) and on our dashboard (links shown in the `Summary` panel)."
+      ],
+      "metadata": {
+        "id": "TbG_3nwEhs5G"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import sys, weco.cli as weco_cli\n",
+        "\n",
+        "# When running in a terminal, you can use this instead:\n",
+        "# weco run --source optimize.py \\\n",
+        "#      --eval-command f\"python evaluate.py --solution-path optimize.py --device {DEVICE}\" \\\n",
+        "#      --metric speedup \\\n",
+        "#      --goal maximize \\\n",
+        "#      --steps 10 \\\n",
+        "#      --additional-instructions \"Fuse operations in the forward method while ensuring the max float deviation remains small.\"\n",
+        "\n",
+        "sys.argv = [\n",
+        "    \"weco\", \"run\",\n",
+        "    \"--source\", \"optimize.py\",\n",
+        "    \"--eval-command\", f\"python evaluate.py --solution-path optimize.py --device {DEVICE}\",\n",
+        "    \"--metric\", \"speedup\",\n",
+        "    \"--goal\", \"maximize\",\n",
+        "    \"--steps\", \"10\",\n",
+        "    \"--additional-instructions\", \"Fuse operations in the forward method while ensuring the max float deviation remains small.\"\n",
+        "]\n",
+        "\n",
+        "try: weco_cli.main()\n",
+        "except SystemExit: pass"
+      ],
+      "metadata": {
+        "id": "17YZ2euZplDJ"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Let's take a look at what our optimized code looks like (`optimize.py`)!"
+      ],
+      "metadata": {
+        "id": "990ueX_JsO_1"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "view_code_block(\"optimize.py\")"
+      ],
+      "metadata": {
+        "id": "9dqfzXkajQKs"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Happy Optimizing from the [Weco](https://www.weco.ai/) Team!\n",
+        "\n",
+        "If you'd like to learn more about what Weco can do, here are some spots to check out:\n",
+        "- 📖 [CLI Reference](https://docs.weco.ai/cli/cli-reference) - Explore our docs for an in-depth look at what the tool can do\n",
+        "- ✨ [Examples](https://docs.weco.ai/examples) - Explore automated R&D across kernel engineering, ML engineering and prompt engineering"
+      ],
+      "metadata": {
+        "id": "17oICow9yjn8"
+      }
+    }
+  ]
+}

{weco-0.2.19 → weco-0.2.20}/pyproject.toml

@@ -8,7 +8,7 @@ name = "weco"
 authors = [{ name = "Weco AI Team", email = "contact@weco.ai" }]
 description = "Documentation for `weco`, a CLI for using Weco AI's code optimizer."
 readme = "README.md"
-version = "0.2.19"
+version = "0.2.20"
 license = { text = "MIT" }
 requires-python = ">=3.8"
 dependencies = ["requests", "rich", "packaging"]

{weco-0.2.19 → weco-0.2.20}/weco/api.py

@@ -17,7 +17,7 @@ def handle_api_error(e: requests.exceptions.HTTPError, console: rich.console.Con
     # sys.exit(1)
 
 
-def start_optimization_session(
+def start_optimization_run(
     console: rich.console.Console,
     source_code: str,
     evaluation_command: str,
@@ -29,14 +29,14 @@ def start_optimization_session(
     search_policy_config: Dict[str, Any],
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
-    auth_headers: dict = {},  # Add auth_headers
+    auth_headers: dict = {},
     timeout: int = 800,
 ) -> Dict[str, Any]:
-    """Start the optimization session."""
+    """Start the optimization run."""
     with console.status("[bold green]Starting Optimization..."):
         try:
             response = requests.post(
-                f"{__base_url__}/sessions",  # Path is relative to base_url
+                f"{__base_url__}/runs",
                 json={
                     "source_code": source_code,
                     "additional_instructions": additional_instructions,
@@ -49,37 +49,37 @@ def start_optimization_session(
                     },
                     "metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
                 },
-                headers=auth_headers,  # Add headers
+                headers=auth_headers,
                 timeout=timeout,
             )
             response.raise_for_status()
             return response.json()
         except requests.exceptions.HTTPError as e:
             handle_api_error(e, console)
-            sys.exit(1)  # Exit if starting session fails
+            sys.exit(1)
         except requests.exceptions.RequestException as e:
-            console.print(f"[bold red]Network Error starting session: {e}[/]")
+            console.print(f"[bold red]Network Error starting run: {e}[/]")
             sys.exit(1)
 
 
 def evaluate_feedback_then_suggest_next_solution(
-    session_id: str,
+    run_id: str,
     execution_output: str,
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
-    auth_headers: dict = {},  # Add auth_headers
+    auth_headers: dict = {},
     timeout: int = 800,
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
     try:
         response = requests.post(
-            f"{__base_url__}/sessions/{session_id}/suggest",  # Path is relative to base_url
+            f"{__base_url__}/runs/{run_id}/suggest",
             json={
                 "execution_output": execution_output,
                 "additional_instructions": additional_instructions,
                 "metadata": {**api_keys},
             },
-            headers=auth_headers,  # Add headers
+            headers=auth_headers,
             timeout=timeout,
         )
         response.raise_for_status()
@@ -93,16 +93,13 @@ def evaluate_feedback_then_suggest_next_solution(
         raise  # Re-raise the exception
 
 
-def get_optimization_session_status(
-    session_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
+def get_optimization_run_status(
+    run_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
 ) -> Dict[str, Any]:
-    """Get the current status of the optimization session."""
+    """Get the current status of the optimization run."""
     try:
         response = requests.get(
-            f"{__base_url__}/sessions/{session_id}",  # Path is relative to base_url
-            params={"include_history": include_history},
-            headers=auth_headers,
-            timeout=timeout,
+            f"{__base_url__}/runs/{run_id}", params={"include_history": include_history}, headers=auth_headers, timeout=timeout
         )
         response.raise_for_status()
         return response.json()
@@ -114,42 +111,30 @@ def get_optimization_session_status(
         raise  # Re-raise
 
 
-def send_heartbeat(
-    session_id: str,
-    auth_headers: dict = {},
-    timeout: int = 10,  # Shorter timeout for non-critical heartbeat
-) -> bool:
+def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: int = 10) -> bool:
     """Send a heartbeat signal to the backend."""
     try:
-        response = requests.put(f"{__base_url__}/sessions/{session_id}/heartbeat", headers=auth_headers, timeout=timeout)
-        response.raise_for_status()  # Raises HTTPError for bad responses (4xx or 5xx)
+        response = requests.put(f"{__base_url__}/runs/{run_id}/heartbeat", headers=auth_headers, timeout=timeout)
+        response.raise_for_status()
         return True
     except requests.exceptions.HTTPError as e:
-        # Log non-critical errors like 409 Conflict (session not running)
         if e.response.status_code == 409:
-            print(f"Heartbeat ignored: Session {session_id} is not running.", file=sys.stderr)
+            print(f"Heartbeat ignored: Run {run_id} is not running.", file=sys.stderr)
         else:
-            print(f"Heartbeat failed for session {session_id}: HTTP {e.response.status_code}", file=sys.stderr)
-        # Don't exit, just report failure
+            print(f"Heartbeat failed for run {run_id}: HTTP {e.response.status_code}", file=sys.stderr)
         return False
     except requests.exceptions.RequestException as e:
-        # Network errors are also non-fatal for heartbeats
-        print(f"Heartbeat network error for session {session_id}: {e}", file=sys.stderr)
+        print(f"Heartbeat network error for run {run_id}: {e}", file=sys.stderr)
         return False
 
 
 def report_termination(
-    session_id: str,
-    status_update: str,
-    reason: str,
-    details: Optional[str] = None,
-    auth_headers: dict = {},
-    timeout: int = 30,  # Reasonably longer timeout for important termination message
+    run_id: str, status_update: str, reason: str, details: Optional[str] = None, auth_headers: dict = {}, timeout: int = 30
 ) -> bool:
     """Report the termination reason to the backend."""
     try:
         response = requests.post(
-            f"{__base_url__}/sessions/{session_id}/terminate",
+            f"{__base_url__}/runs/{run_id}/terminate",
             json={"status_update": status_update, "termination_reason": reason, "termination_details": details},
             headers=auth_headers,
             timeout=timeout,
@@ -157,6 +142,5 @@ def report_termination(
         response.raise_for_status()
         return True
     except requests.exceptions.RequestException as e:
-        # Log failure, but don't prevent CLI exit
-        print(f"Warning: Failed to report termination to backend for session {session_id}: {e}", file=sys.stderr)
+        print(f"Warning: Failed to report termination to backend for run {run_id}: {e}", file=sys.stderr)
         return False