weco 0.2.24__tar.gz → 0.2.26__tar.gz

{weco-0.2.24 → weco-0.2.26}/.gitignore

@@ -75,4 +75,7 @@ CLAUDE.md
 
 # Repomix
 .repomixignore
-repomix-output.*
+repomix-output.*
+
+# Claude config
+.claude/

{weco-0.2.24 → weco-0.2.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.24
+Version: 0.2.26
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -268,6 +268,7 @@ Weco supports the following LLM models:
 - `gpt-4o-mini`
 
 ### Anthropic Models
+- `claude-opus-4-1`
 - `claude-opus-4-0`
 - `claude-sonnet-4-0`
 - `claude-3-7-sonnet-latest`

{weco-0.2.24 → weco-0.2.26}/README.md

@@ -242,6 +242,7 @@ Weco supports the following LLM models:
 - `gpt-4o-mini`
 
 ### Anthropic Models
+- `claude-opus-4-1`
 - `claude-opus-4-0`
 - `claude-sonnet-4-0`
 - `claude-3-7-sonnet-latest`

{weco-0.2.24 → weco-0.2.26}/examples/cuda/evaluate.py

@@ -154,7 +154,7 @@ if __name__ == "__main__":
     max_diff_avg /= n_correctness_trials
     print(f"max float diff between values of baseline and optimized model: {max_diff_avg}")
     if max_diff_avg > correctness_tolerance:
-        print("invalid solution: max float diff is too high")
+        print("Incorrect solution: max float diff is too high")
 
     # measure performance
     inputs = get_inputs(batch_size=batch_size, seq_len=seq_len, n_embd=n_embd, device="cuda")

weco-0.2.26/examples/hello-kernel-world/README.md

@@ -0,0 +1,59 @@
+# Hello Kernel World
+
+This example demonstrates the basics of using Weco to optimize a simple PyTorch model. The model performs a series of basic operations: matrix multiplication, division, summation, and scaling. It's designed as an introductory tutorial to help you understand how Weco works before moving on to more advanced optimization tasks.
+
+## Setup
+
+Install the CLI using `pip`:
+```bash
+pip install weco
+```
+
+Create your API key from one of the supported providers:
+- **OpenAI:** Create your API key [here](https://platform.openai.com/api-keys), then run: `export OPENAI_API_KEY="your_key_here"`
+- **Anthropic:** Create your API key [here](https://console.anthropic.com/settings/keys), then run: `export ANTHROPIC_API_KEY="your_key_here"`  
+- **Google:** Create your API key [here](https://aistudio.google.com/apikey), then run: `export GEMINI_API_KEY="your_key_here"`
+
+Install the required dependencies:
+```bash
+pip install torch
+```
+
+## Run Weco
+
+Now run Weco to optimize your code:
+```bash
+weco run --source optimize.py \
+     --eval-command "python evaluate.py --solution-path optimize.py --device cpu" \
+     --metric speedup \
+     --goal maximize \
+     --steps 15 \
+     --additional-instructions "Fuse operations in the forward method while ensuring the max float deviation remains small. Maintain the same format of the code."
+```
+
+**Note:** If you have an NVIDIA GPU, change the device in the `--eval-command` to `cuda`. If you are running this on Apple Silicon, set it to `mps`.
+
+### Explanation
+
+*   `--source optimize.py`: The simple PyTorch model to be optimized.
+*   `--eval-command "python evaluate.py --solution-path optimize.py --device cpu"`: Runs the evaluation script, which benchmarks the optimized code against a baseline and prints the `speedup`.
+*   `--metric speedup`: The optimization target metric.
+*   `--goal maximize`: To increase the speedup.
+*   `--steps 15`: The number of optimization iterations.
+*   `--additional-instructions "..."`: Provides specific guidance to focus on operation fusion while maintaining correctness.
+
+Weco will iteratively modify `optimize.py`, attempting to fuse and optimize the operations in the forward method, guided by the performance feedback from the evaluation script.
+
+## Interactive Tutorial
+****
+For a hands-on walkthrough of this example, check out the [Colab notebook](colab_notebook_walkthrough.ipynb) that provides step-by-step guidance through the optimization process.
+
+## Next Steps
+
+Once you've mastered the basics with this example, explore more advanced optimization techniques:
+- [Triton Optimization](/examples/triton/README.md) for GPU kernel programming
+- [CUDA Optimization](/examples/cuda/README.md) for low-level GPU optimization
+- [Model Development](/examples/spaceship-titanic/README.md) for ML model optimization
+- [Prompt Engineering](/examples/prompt/README.md) for LLM prompt optimization
+
+You can also check out our [CLI Reference](https://docs.weco.ai/cli/cli-reference) to learn more about what you can do with the tool.

{weco-0.2.24 → weco-0.2.26}/examples/hello-kernel-world/evaluate.py

@@ -132,7 +132,7 @@ if __name__ == "__main__":
     max_diff_avg /= n_correctness_trials
     print(f"max float diff between values of baseline and optimized model: {max_diff_avg}")
     if max_diff_avg > correctness_tolerance:
-        print("invalid solution: max float diff is too high")
+        print("Incorrect solution: max float diff is too high")
 
     # measure performance
     inputs = get_inputs(batch_size, input_size, args.device)

{weco-0.2.24 → weco-0.2.26}/examples/spaceship-titanic/evaluate.py

@@ -5,7 +5,7 @@ from sklearn.metrics import accuracy_score
 from sklearn.model_selection import train_test_split
 
 
-class InvalidSubmissionError(Exception):
+class IncorrectSubmissionError(Exception):
     pass
 
 
@@ -18,18 +18,18 @@ def evaluate_for_accuracy(
 
     # Submission checks
     if len(submission_df) != len(answers_df):
-        raise InvalidSubmissionError("Submission must have the same length as the answers.")
+        raise IncorrectSubmissionError("Submission must have the same length as the answers.")
     if target_column not in submission_df.columns:
-        raise InvalidSubmissionError(f"Submission must have a `{target_column}` column")
+        raise IncorrectSubmissionError(f"Submission must have a `{target_column}` column")
     if id_column not in submission_df.columns:
-        raise InvalidSubmissionError(f"Submission must have a `{id_column}` column")
+        raise IncorrectSubmissionError(f"Submission must have a `{id_column}` column")
 
     # Sort on id to ensure correct ordering
     submission_df = submission_df.sort_values(by=id_column)
     answers_df = answers_df.sort_values(by=id_column)
 
     if (submission_df[id_column].values != answers_df[id_column].values).any():
-        raise InvalidSubmissionError(f"Submission and Answers `{id_column}`'s do not match")
+        raise IncorrectSubmissionError(f"Submission and Answers `{id_column}`'s do not match")
 
     return accuracy_score(submission_df[target_column], answers_df[target_column])
 

{weco-0.2.24 → weco-0.2.26}/examples/triton/evaluate.py

@@ -149,7 +149,7 @@ if __name__ == "__main__":
     max_diff_avg /= n_correctness_trials
     print(f"max float diff between values of baseline and optimized model: {max_diff_avg}")
     if max_diff_avg > correctness_tolerance:
-        print("invalid solution: max float diff is too high")
+        print("Incorrect solution: max float diff is too high")
 
     # measure performance
     inputs = get_inputs(batch_size=batch_size, seq_len=seq_len, n_embd=n_embd, device="cuda")

{weco-0.2.24 → weco-0.2.26}/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.24"
+version = "0.2.26"
 license = { text = "MIT" }
 requires-python = ">=3.8"
 dependencies = [

{weco-0.2.24 → weco-0.2.26}/weco/api.py

@@ -14,8 +14,6 @@ def handle_api_error(e: requests.exceptions.HTTPError, console: Console) -> None
     except (ValueError, KeyError):  # Handle cases where response is not JSON or detail key is missing
         detail = f"HTTP {e.response.status_code} Error: {e.response.text}"
     console.print(f"[bold red]{detail}[/]")
-    # Avoid exiting here, let the caller decide if the error is fatal
-    # sys.exit(1)
 
 
 def start_optimization_run(
@@ -32,7 +30,7 @@ def start_optimization_run(
     api_keys: Dict[str, Any] = {},
     auth_headers: dict = {},
     timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
-) -> Dict[str, Any]:
+) -> Optional[Dict[str, Any]]:
     """Start the optimization run."""
     with console.status("[bold green]Starting Optimization..."):
         try:
@@ -61,12 +59,12 @@ def start_optimization_run(
             if result.get("code") is None:
                 result["code"] = ""
             return result
-        except Exception as e:
+        except requests.exceptions.HTTPError as e:
             handle_api_error(e, console)
-            sys.exit(1)
+            return None
         except Exception as e:
             console.print(f"[bold red]Error starting run: {e}[/]")
-            sys.exit(1)
+            return None
 
 
 def evaluate_feedback_then_suggest_next_solution(
@@ -101,11 +99,11 @@ def evaluate_feedback_then_suggest_next_solution(
         return result
     except requests.exceptions.HTTPError as e:
         # Allow caller to handle suggest errors, maybe retry or terminate
-        handle_api_error(e, console)  # Use default console if none passed
-        raise  # Re-raise the exception
+        handle_api_error(e, console)
+        raise
     except Exception as e:
-        print(f"Error: {e}")  # Use print as console might not be available
-        raise  # Re-raise the exception
+        console.print(f"[bold red]Error: {e}[/]")
+        raise
 
 
 def get_optimization_run_status(
@@ -137,11 +135,11 @@ def get_optimization_run_status(
                     result["nodes"][i]["code"] = ""
         return result
     except requests.exceptions.HTTPError as e:
-        handle_api_error(e, console)  # Use default console
-        raise  # Re-raise
+        handle_api_error(e, console)
+        raise
     except Exception as e:
-        print(f"Error getting run status: {e}")
-        raise  # Re-raise
+        console.print(f"[bold red]Error getting run status: {e}[/]")
+        raise
 
 
 def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = (10, 10)) -> bool:
@@ -152,7 +150,7 @@ def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tup
         return True
     except requests.exceptions.HTTPError as e:
         if e.response.status_code == 409:
-            print("Polling ignore: Run {run_id} is not running.", file=sys.stderr)
+            print(f"Polling ignore: Run {run_id} is not running.", file=sys.stderr)
         else:
             print(f"Polling failed for run {run_id}: HTTP {e.response.status_code}", file=sys.stderr)
         return False
@@ -221,7 +219,6 @@ def get_optimization_suggestions_from_codebase(
     """Analyze codebase and get optimization suggestions using the model-agnostic backend API."""
     model, api_key_dict = _determine_model_and_api_key()
     try:
-        model, api_key_dict = _determine_model_and_api_key()
         response = requests.post(
             f"{__base_url__}/onboard/analyze-codebase",
             json={
@@ -238,7 +235,7 @@ def get_optimization_suggestions_from_codebase(
         result = response.json()
         return [option for option in result.get("options", [])]
 
-    except Exception as e:
+    except requests.exceptions.HTTPError as e:
         handle_api_error(e, console)
         return None
     except Exception as e:
@@ -257,7 +254,6 @@ def generate_evaluation_script_and_metrics(
     """Generate evaluation script and determine metrics using the model-agnostic backend API."""
     model, api_key_dict = _determine_model_and_api_key()
     try:
-        model, api_key_dict = _determine_model_and_api_key()
         response = requests.post(
             f"{__base_url__}/onboard/generate-script",
             json={
@@ -294,7 +290,6 @@ def analyze_evaluation_environment(
     """Analyze existing evaluation scripts and environment using the model-agnostic backend API."""
     model, api_key_dict = _determine_model_and_api_key()
     try:
-        model, api_key_dict = _determine_model_and_api_key()
         response = requests.post(
             f"{__base_url__}/onboard/analyze-environment",
             json={
@@ -312,7 +307,7 @@ def analyze_evaluation_environment(
         response.raise_for_status()
         return response.json()
 
-    except Exception as e:
+    except requests.exceptions.HTTPError as e:
         handle_api_error(e, console)
         return None
     except Exception as e:
@@ -331,7 +326,6 @@ def analyze_script_execution_requirements(
     """Analyze script to determine proper execution command using the model-agnostic backend API."""
     model, api_key_dict = _determine_model_and_api_key()
     try:
-        model, api_key_dict = _determine_model_and_api_key()
         response = requests.post(
             f"{__base_url__}/onboard/analyze-script",
             json={
@@ -348,7 +342,7 @@ def analyze_script_execution_requirements(
         result = response.json()
         return result.get("command", f"python {script_path}")
 
-    except Exception as e:
+    except requests.exceptions.HTTPError as e:
         handle_api_error(e, console)
         return f"python {script_path}"
     except Exception as e:

{weco-0.2.24 → weco-0.2.26}/weco/chatbot.py

@@ -2,6 +2,7 @@ import pathlib
 import shlex
 import argparse
 from typing import List, Optional, Dict, Any, Tuple
+import sys
 
 from rich.console import Console
 from rich.prompt import Prompt
@@ -682,9 +683,9 @@ class Chatbot:
 
         # Import and execute the actual optimization function
         # (Import here to avoid circular imports)
-        from .optimizer import execute_optimization as actual_execute_optimization
+        from .optimizer import execute_optimization as execute_optimization_run
 
-        success = actual_execute_optimization(
+        success = execute_optimization_run(
             source=target_file,
             eval_command=eval_config["eval_command"],
             metric=eval_config["metric_name"],
@@ -702,6 +703,9 @@ class Chatbot:
         else:
             self.console.print("\n[bold yellow]⚠️  Optimization ended early or encountered issues.[/]")
 
+        exit_code = 0 if success else 1
+        sys.exit(exit_code)
+
     def show_and_copy_command(self, command: str) -> None:
         """Show the command and copy it to clipboard."""
         import subprocess

{weco-0.2.24 → weco-0.2.26}/weco/cli.py

@@ -71,7 +71,7 @@ def configure_run_parser(run_parser: argparse.ArgumentParser) -> None:
 
 def execute_run_command(args: argparse.Namespace) -> None:
     """Execute the 'weco run' command with all its logic."""
-    from .optimizer import execute_optimization  # Moved import inside
+    from .optimizer import execute_optimization
 
     success = execute_optimization(
         source=args.source,

{weco-0.2.24 → weco-0.2.26}/weco/optimizer.py

@@ -183,6 +183,10 @@ def execute_optimization(
             auth_headers=auth_headers,
             timeout=api_timeout,
         )
+        # Indicate the endpoint failed to return a response and the optimization was unsuccessful
+        if run_response is None:
+            return False
+
         run_id = run_response["run_id"]
         run_name = run_response["run_name"]
         current_run_id_for_heartbeat = run_id

{weco-0.2.24 → weco-0.2.26}/weco.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.24
+Version: 0.2.26
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -268,6 +268,7 @@ Weco supports the following LLM models:
 - `gpt-4o-mini`
 
 ### Anthropic Models
+- `claude-opus-4-1`
 - `claude-opus-4-0`
 - `claude-sonnet-4-0`
 - `claude-3-7-sonnet-latest`

{weco-0.2.24 → weco-0.2.26}/weco.egg-info/SOURCES.txt

@@ -11,6 +11,7 @@ examples/cuda/README.md
 examples/cuda/evaluate.py
 examples/cuda/guide.md
 examples/cuda/optimize.py
+examples/hello-kernel-world/README.md
 examples/hello-kernel-world/colab_notebook_walkthrough.ipynb
 examples/hello-kernel-world/evaluate.py
 examples/hello-kernel-world/optimize.py