autooptm 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- autooptm-0.1.0/.gitignore +12 -0
- autooptm-0.1.0/PKG-INFO +43 -0
- autooptm-0.1.0/README.md +33 -0
- autooptm-0.1.0/autooptm/__init__.py +185 -0
- autooptm-0.1.0/autooptm/cli.py +160 -0
- autooptm-0.1.0/autooptm/mcp_server.py +118 -0
- autooptm-0.1.0/pyproject.toml +23 -0
autooptm-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: autooptm
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Client, CLI and MCP server for AutoOptm — measured GPU speedups for your training and inference code
|
|
5
|
+
License: MIT
|
|
6
|
+
Requires-Python: >=3.9
|
|
7
|
+
Provides-Extra: mcp
|
|
8
|
+
Requires-Dist: mcp>=1.0; extra == 'mcp'
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
|
|
11
|
+
# autooptm
|
|
12
|
+
|
|
13
|
+
Measured GPU speedups for your training and inference code — from the
|
|
14
|
+
terminal, an agent, or Python. Diagnosis and measurement are free; credits
|
|
15
|
+
move only when you unlock the patch.
|
|
16
|
+
|
|
17
|
+
```bash
|
|
18
|
+
pip install autooptm
|
|
19
|
+
autooptm login # email code, no browser needed
|
|
20
|
+
autooptm run https://github.com/you/repo --entrypoint train.py
|
|
21
|
+
autooptm unlock <jobId> # asks before spending credits
|
|
22
|
+
git apply autooptm.patch
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
Python:
|
|
26
|
+
|
|
27
|
+
```python
|
|
28
|
+
from autooptm import AutoOptm
|
|
29
|
+
ao = AutoOptm()
|
|
30
|
+
job = ao.submit(".", entrypoint="bench.py", workload="inference")
|
|
31
|
+
done = ao.wait(job["jobId"])
|
|
32
|
+
print(done["speedup"], done["unlock_credits"])
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
MCP (Claude Code, Cursor, any MCP client) — `pip install 'autooptm[mcp]'`:
|
|
36
|
+
|
|
37
|
+
```json
|
|
38
|
+
{ "mcpServers": { "autooptm": { "command": "autooptm-mcp" } } }
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
Tools: `optimize_submit`, `optimize_status`, `optimize_wait`,
|
|
42
|
+
`unlock_patch` (requires explicit user confirmation), `download_patch`,
|
|
43
|
+
`account_balance`.
|
autooptm-0.1.0/README.md
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
# autooptm
|
|
2
|
+
|
|
3
|
+
Measured GPU speedups for your training and inference code — from the
|
|
4
|
+
terminal, an agent, or Python. Diagnosis and measurement are free; credits
|
|
5
|
+
move only when you unlock the patch.
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install autooptm
|
|
9
|
+
autooptm login # email code, no browser needed
|
|
10
|
+
autooptm run https://github.com/you/repo --entrypoint train.py
|
|
11
|
+
autooptm unlock <jobId> # asks before spending credits
|
|
12
|
+
git apply autooptm.patch
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
Python:
|
|
16
|
+
|
|
17
|
+
```python
|
|
18
|
+
from autooptm import AutoOptm
|
|
19
|
+
ao = AutoOptm()
|
|
20
|
+
job = ao.submit(".", entrypoint="bench.py", workload="inference")
|
|
21
|
+
done = ao.wait(job["jobId"])
|
|
22
|
+
print(done["speedup"], done["unlock_credits"])
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
MCP (Claude Code, Cursor, any MCP client) — `pip install 'autooptm[mcp]'`:
|
|
26
|
+
|
|
27
|
+
```json
|
|
28
|
+
{ "mcpServers": { "autooptm": { "command": "autooptm-mcp" } } }
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
Tools: `optimize_submit`, `optimize_status`, `optimize_wait`,
|
|
32
|
+
`unlock_patch` (requires explicit user confirmation), `download_patch`,
|
|
33
|
+
`account_balance`.
|
|
@@ -0,0 +1,185 @@
|
|
|
1
|
+
"""AutoOptm client.
|
|
2
|
+
|
|
3
|
+
Submit a repository (git URL or local directory), get a measured GPU
|
|
4
|
+
speedup back, unlock the patch when the number convinces you.
|
|
5
|
+
|
|
6
|
+
from autooptm import AutoOptm
|
|
7
|
+
ao = AutoOptm() # token from AUTOOPTM_TOKEN
|
|
8
|
+
job = ao.submit("https://github.com/you/repo", entrypoint="train.py")
|
|
9
|
+
done = ao.wait(job["jobId"])
|
|
10
|
+
print(done["speedup"]) # e.g. 1.87
|
|
11
|
+
ao.unlock(done["id"]) # spends credits — ask first
|
|
12
|
+
ao.download_patch(done["id"], "autooptm.patch")
|
|
13
|
+
|
|
14
|
+
The token is the session token the site stores after sign-in
|
|
15
|
+
(localStorage key "ao.token"); it is valid for seven days.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
from __future__ import annotations
|
|
19
|
+
|
|
20
|
+
import io
|
|
21
|
+
import json
|
|
22
|
+
import os
|
|
23
|
+
import time
|
|
24
|
+
import urllib.error
|
|
25
|
+
import urllib.request
|
|
26
|
+
import zipfile
|
|
27
|
+
from pathlib import Path
|
|
28
|
+
from typing import Any, Optional
|
|
29
|
+
|
|
30
|
+
__version__ = "0.1.0"
|
|
31
|
+
|
|
32
|
+
DEFAULT_API = "https://api.autooptm.com"
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def _token_path() -> Path:
|
|
36
|
+
base = os.environ.get("XDG_CONFIG_HOME") or str(Path.home() / ".config")
|
|
37
|
+
return Path(base) / "autooptm" / "token"
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def save_token(token: str) -> Path:
|
|
41
|
+
p = _token_path()
|
|
42
|
+
p.parent.mkdir(parents=True, exist_ok=True)
|
|
43
|
+
p.write_text(token)
|
|
44
|
+
try:
|
|
45
|
+
p.chmod(0o600)
|
|
46
|
+
except OSError:
|
|
47
|
+
pass
|
|
48
|
+
return p
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def clear_token() -> None:
|
|
52
|
+
_token_path().unlink(missing_ok=True)
|
|
53
|
+
|
|
54
|
+
# Never worth uploading: bulky, derived, or secret-bearing.
|
|
55
|
+
_SKIP_DIRS = {".git", "__pycache__", ".venv", "venv", "node_modules",
|
|
56
|
+
".mypy_cache", ".ruff_cache", ".pytest_cache", "wandb",
|
|
57
|
+
".idea", ".vscode"}
|
|
58
|
+
_SKIP_SUFFIXES = {".pt", ".pth", ".ckpt", ".safetensors", ".onnx", ".npz"}
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
class AutoOptmError(RuntimeError):
|
|
62
|
+
def __init__(self, status: int, body: dict):
|
|
63
|
+
self.status = status
|
|
64
|
+
self.body = body
|
|
65
|
+
super().__init__(f"HTTP {status}: {body.get('error') or body}")
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
class AutoOptm:
|
|
69
|
+
def __init__(self, token: Optional[str] = None, api: Optional[str] = None):
|
|
70
|
+
self.api = (api or os.environ.get("AUTOOPTM_API") or DEFAULT_API).rstrip("/")
|
|
71
|
+
stored = ""
|
|
72
|
+
try:
|
|
73
|
+
stored = _token_path().read_text().strip()
|
|
74
|
+
except OSError:
|
|
75
|
+
pass
|
|
76
|
+
self.token = token or os.environ.get("AUTOOPTM_TOKEN") or stored
|
|
77
|
+
if not self.token:
|
|
78
|
+
raise AutoOptmError(0, {
|
|
79
|
+
"error": "no_token",
|
|
80
|
+
"hint": "run `autooptm login` (email code, no browser needed) "
|
|
81
|
+
"or export AUTOOPTM_TOKEN"})
|
|
82
|
+
|
|
83
|
+
# ------------------------------------------------------------- transport
|
|
84
|
+
def _req(self, method: str, path: str, body: Any = None,
|
|
85
|
+
raw: Optional[bytes] = None, content_type: str = "application/json"):
|
|
86
|
+
data = raw if raw is not None else (
|
|
87
|
+
json.dumps(body).encode() if body is not None else None)
|
|
88
|
+
r = urllib.request.Request(
|
|
89
|
+
self.api + path, data=data, method=method,
|
|
90
|
+
headers={"authorization": f"Bearer {self.token}",
|
|
91
|
+
"content-type": content_type,
|
|
92
|
+
# Cloudflare challenges the bare urllib agent (1010).
|
|
93
|
+
"user-agent": f"autooptm-python/{__version__}"})
|
|
94
|
+
try:
|
|
95
|
+
with urllib.request.urlopen(r, timeout=120) as resp:
|
|
96
|
+
return json.loads(resp.read().decode() or "{}")
|
|
97
|
+
except urllib.error.HTTPError as e:
|
|
98
|
+
try:
|
|
99
|
+
payload = json.loads(e.read().decode() or "{}")
|
|
100
|
+
except Exception:
|
|
101
|
+
payload = {"error": f"http_{e.code}"}
|
|
102
|
+
raise AutoOptmError(e.code, payload) from None
|
|
103
|
+
|
|
104
|
+
# ------------------------------------------------------------------ api
|
|
105
|
+
def balance(self) -> dict:
|
|
106
|
+
return self._req("GET", "/api/balance")
|
|
107
|
+
|
|
108
|
+
def upload_dir(self, path: str) -> str:
|
|
109
|
+
"""Zip a local directory (skipping caches, checkpoints, .git) and
|
|
110
|
+
upload it. Returns the storage key to pass to submit()."""
|
|
111
|
+
root = Path(path).resolve()
|
|
112
|
+
buf = io.BytesIO()
|
|
113
|
+
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z:
|
|
114
|
+
for p in sorted(root.rglob("*")):
|
|
115
|
+
rel = p.relative_to(root)
|
|
116
|
+
if any(part in _SKIP_DIRS for part in rel.parts):
|
|
117
|
+
continue
|
|
118
|
+
if p.is_file() and p.suffix.lower() not in _SKIP_SUFFIXES:
|
|
119
|
+
z.write(p, str(rel))
|
|
120
|
+
blob = buf.getvalue()
|
|
121
|
+
out = self._req("POST", "/api/uploads", raw=blob,
|
|
122
|
+
content_type="application/zip")
|
|
123
|
+
return out["key"]
|
|
124
|
+
|
|
125
|
+
def submit(self, source: str, entrypoint: str = "main.py",
|
|
126
|
+
workload: str = "training", gpu: Optional[str] = None,
|
|
127
|
+
args: Optional[str] = None, setup: Optional[str] = None,
|
|
128
|
+
git_ref: Optional[str] = None,
|
|
129
|
+
scale_args: Optional[str] = None) -> dict:
|
|
130
|
+
"""Submit a git URL or a local directory for optimization.
|
|
131
|
+
|
|
132
|
+
Diagnosis and measurement are free of charge; credits move only
|
|
133
|
+
when the patch is unlocked."""
|
|
134
|
+
if "://" in source:
|
|
135
|
+
kind, src = "git", source
|
|
136
|
+
else:
|
|
137
|
+
kind, src = "zip", self.upload_dir(source)
|
|
138
|
+
body = {"sourceKind": kind, "source": src, "entrypoint": entrypoint,
|
|
139
|
+
"workload": workload}
|
|
140
|
+
for k, v in (("gpu", gpu), ("args", args), ("setup", setup),
|
|
141
|
+
("gitRef", git_ref), ("scaleArgs", scale_args)):
|
|
142
|
+
if v:
|
|
143
|
+
body[k] = v
|
|
144
|
+
return self._req("POST", "/api/jobs", body)
|
|
145
|
+
|
|
146
|
+
def job(self, job_id: str) -> dict:
|
|
147
|
+
out = self._req("GET", f"/api/jobs/{job_id}")
|
|
148
|
+
# The detail endpoint wraps the record: { job: {...}, events: [...] }.
|
|
149
|
+
job = out.get("job", out)
|
|
150
|
+
if isinstance(job, dict) and "events" in out:
|
|
151
|
+
job = {**job, "events": out["events"]}
|
|
152
|
+
return job
|
|
153
|
+
|
|
154
|
+
def wait(self, job_id: str, timeout_s: int = 3600, poll_s: int = 10) -> dict:
|
|
155
|
+
"""Poll until the job finishes. Returns the final job record."""
|
|
156
|
+
deadline = time.time() + timeout_s
|
|
157
|
+
while True:
|
|
158
|
+
j = self.job(job_id)
|
|
159
|
+
if j.get("status") in ("succeeded", "failed", "canceled"):
|
|
160
|
+
return j
|
|
161
|
+
if time.time() > deadline:
|
|
162
|
+
raise AutoOptmError(0, {"error": "timeout", "jobId": job_id,
|
|
163
|
+
"lastStatus": j.get("status")})
|
|
164
|
+
time.sleep(poll_s)
|
|
165
|
+
|
|
166
|
+
def cancel(self, job_id: str) -> dict:
|
|
167
|
+
"""Stop a queued or running job. Queued work dies immediately;
|
|
168
|
+
running work is asked to stop on its next heartbeat."""
|
|
169
|
+
return self._req("POST", "/api/jobs/cancel", {"jobId": job_id})
|
|
170
|
+
|
|
171
|
+
def unlock(self, job_id: str) -> dict:
|
|
172
|
+
"""Accept the quoted price and open the patch. THIS SPENDS CREDITS."""
|
|
173
|
+
return self._req("POST", "/api/jobs/unlock", {"jobId": job_id})
|
|
174
|
+
|
|
175
|
+
def download_patch(self, job_id: str, dest: str = "autooptm.patch") -> str:
|
|
176
|
+
j = self.job(job_id)
|
|
177
|
+
url = j.get("patch_url")
|
|
178
|
+
if not url:
|
|
179
|
+
raise AutoOptmError(0, {"error": "patch_locked",
|
|
180
|
+
"hint": "call unlock() first"})
|
|
181
|
+
dl = urllib.request.Request(
|
|
182
|
+
url, headers={"user-agent": f"autooptm-python/{__version__}"})
|
|
183
|
+
with urllib.request.urlopen(dl, timeout=120) as r:
|
|
184
|
+
Path(dest).write_bytes(r.read())
|
|
185
|
+
return dest
|
|
@@ -0,0 +1,160 @@
|
|
|
1
|
+
"""`autooptm` CLI — submit, watch, unlock, download. One command each.
|
|
2
|
+
|
|
3
|
+
autooptm run https://github.com/you/repo --entrypoint train.py
|
|
4
|
+
autooptm run . --entrypoint bench.py --gpu A10
|
|
5
|
+
autooptm status <jobId>
|
|
6
|
+
autooptm unlock <jobId> --yes
|
|
7
|
+
autooptm balance
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import argparse
|
|
13
|
+
import sys
|
|
14
|
+
|
|
15
|
+
import json
|
|
16
|
+
import os
|
|
17
|
+
import urllib.error
|
|
18
|
+
import urllib.request
|
|
19
|
+
|
|
20
|
+
from . import DEFAULT_API, AutoOptm, AutoOptmError, clear_token, save_token
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def _login(api: str, role: str | None) -> int:
|
|
24
|
+
email = input("email: ").strip()
|
|
25
|
+
if not email:
|
|
26
|
+
print("no email"); return 1
|
|
27
|
+
|
|
28
|
+
def post(path: str, body: dict) -> dict:
|
|
29
|
+
r = urllib.request.Request(api + path, method="POST",
|
|
30
|
+
data=json.dumps(body).encode(),
|
|
31
|
+
headers={"content-type": "application/json",
|
|
32
|
+
"user-agent": "autooptm-cli"})
|
|
33
|
+
with urllib.request.urlopen(r, timeout=60) as resp:
|
|
34
|
+
return json.loads(resp.read().decode() or "{}")
|
|
35
|
+
|
|
36
|
+
try:
|
|
37
|
+
started = post("/api/auth/email/start", {"email": email})
|
|
38
|
+
except urllib.error.HTTPError as e:
|
|
39
|
+
print("could not send the code: " + (e.read().decode() or str(e)))
|
|
40
|
+
return 1
|
|
41
|
+
code = started.get("devCode") or input("6-digit code from your inbox: ").strip()
|
|
42
|
+
body = {"email": email, "code": code}
|
|
43
|
+
if role:
|
|
44
|
+
body["role"] = role
|
|
45
|
+
try:
|
|
46
|
+
done = post("/api/auth/email/verify", body)
|
|
47
|
+
except urllib.error.HTTPError as e:
|
|
48
|
+
print("verify failed: " + (e.read().decode() or str(e)))
|
|
49
|
+
return 1
|
|
50
|
+
tok = done.get("token")
|
|
51
|
+
if not tok:
|
|
52
|
+
print("verify failed: " + json.dumps(done)); return 1
|
|
53
|
+
path = save_token(tok)
|
|
54
|
+
print(f"signed in as {email}; token saved to {path} (valid ~7 days)")
|
|
55
|
+
return 0
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def main(argv=None) -> int:
|
|
59
|
+
ap = argparse.ArgumentParser(prog="autooptm")
|
|
60
|
+
sub = ap.add_subparsers(dest="cmd", required=True)
|
|
61
|
+
|
|
62
|
+
run = sub.add_parser("run", help="submit a repo (git URL or local dir) and wait")
|
|
63
|
+
run.add_argument("source")
|
|
64
|
+
run.add_argument("--entrypoint", default="main.py")
|
|
65
|
+
run.add_argument("--workload", choices=["training", "inference"], default="training")
|
|
66
|
+
run.add_argument("--gpu", default=None)
|
|
67
|
+
run.add_argument("--args", dest="run_args", default=None)
|
|
68
|
+
run.add_argument("--setup", default=None)
|
|
69
|
+
run.add_argument("--git-ref", default=None)
|
|
70
|
+
run.add_argument("--no-wait", action="store_true")
|
|
71
|
+
|
|
72
|
+
st = sub.add_parser("status", help="one job's state")
|
|
73
|
+
st.add_argument("job_id")
|
|
74
|
+
|
|
75
|
+
ul = sub.add_parser("unlock", help="pay the quoted credits and download the patch")
|
|
76
|
+
ul.add_argument("job_id")
|
|
77
|
+
ul.add_argument("--yes", action="store_true", help="skip the confirmation prompt")
|
|
78
|
+
ul.add_argument("--out", default="autooptm.patch")
|
|
79
|
+
|
|
80
|
+
cx = sub.add_parser("cancel", help="stop a queued or running job")
|
|
81
|
+
cx.add_argument("job_id")
|
|
82
|
+
|
|
83
|
+
sub.add_parser("balance", help="credits available")
|
|
84
|
+
|
|
85
|
+
lg = sub.add_parser("login", help="sign in with an email code — no browser needed")
|
|
86
|
+
lg.add_argument("--role", default=None,
|
|
87
|
+
choices=["researcher", "student", "engineer",
|
|
88
|
+
"startup", "indie", "other"],
|
|
89
|
+
help="optional: what you do (helps us build the right thing)")
|
|
90
|
+
sub.add_parser("logout", help="forget the saved token")
|
|
91
|
+
|
|
92
|
+
a = ap.parse_args(argv)
|
|
93
|
+
if a.cmd == "login":
|
|
94
|
+
api = (os.environ.get("AUTOOPTM_API") or DEFAULT_API).rstrip("/")
|
|
95
|
+
return _login(api, a.role)
|
|
96
|
+
if a.cmd == "logout":
|
|
97
|
+
clear_token()
|
|
98
|
+
print("token cleared")
|
|
99
|
+
return 0
|
|
100
|
+
try:
|
|
101
|
+
ao = AutoOptm()
|
|
102
|
+
if a.cmd == "balance":
|
|
103
|
+
b = ao.balance()
|
|
104
|
+
print(f"{b.get('total')} credits (≈ ${b.get('usdValue')})")
|
|
105
|
+
elif a.cmd == "run":
|
|
106
|
+
job = ao.submit(a.source, entrypoint=a.entrypoint, workload=a.workload,
|
|
107
|
+
gpu=a.gpu, args=a.run_args, setup=a.setup,
|
|
108
|
+
git_ref=a.git_ref)
|
|
109
|
+
print(f"job {job['jobId']} {job.get('status')}"
|
|
110
|
+
+ (" (free look)" if job.get("freeLook") else ""))
|
|
111
|
+
if not a.no_wait:
|
|
112
|
+
done = ao.wait(job["jobId"])
|
|
113
|
+
_print_result(done)
|
|
114
|
+
elif a.cmd == "cancel":
|
|
115
|
+
r = ao.cancel(a.job_id)
|
|
116
|
+
print(r.get("result", "ok"))
|
|
117
|
+
elif a.cmd == "status":
|
|
118
|
+
_print_result(ao.job(a.job_id))
|
|
119
|
+
elif a.cmd == "unlock":
|
|
120
|
+
j = ao.job(a.job_id)
|
|
121
|
+
price = j.get("unlock_credits")
|
|
122
|
+
if not a.yes:
|
|
123
|
+
ans = input(f"Unlock spends {price} credits. Continue? [y/N] ")
|
|
124
|
+
if ans.strip().lower() not in ("y", "yes"):
|
|
125
|
+
print("aborted"); return 1
|
|
126
|
+
r = ao.unlock(a.job_id)
|
|
127
|
+
path = ao.download_patch(a.job_id, a.out)
|
|
128
|
+
print(f"charged {r.get('charged')} credits; patch -> {path}")
|
|
129
|
+
print("apply with: git apply " + path)
|
|
130
|
+
return 0
|
|
131
|
+
except AutoOptmError as e:
|
|
132
|
+
print(f"error: {e}", file=sys.stderr)
|
|
133
|
+
if e.body.get("hint"):
|
|
134
|
+
print("hint: " + str(e.body["hint"]), file=sys.stderr)
|
|
135
|
+
return 1
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
def _print_result(j: dict) -> None:
|
|
139
|
+
s = j.get("status")
|
|
140
|
+
if s == "succeeded":
|
|
141
|
+
line = f"{j.get('speedup')}x"
|
|
142
|
+
if j.get("baseline_median") is not None:
|
|
143
|
+
line += (f" ({j['baseline_median']}s -> "
|
|
144
|
+
f"{j.get('optimised_median')}s)")
|
|
145
|
+
print(line)
|
|
146
|
+
if j.get("patch_unlocked"):
|
|
147
|
+
print("patch: unlocked")
|
|
148
|
+
elif j.get("unlock_credits") is not None:
|
|
149
|
+
print(f"patch: locked — unlock for {j['unlock_credits']} credits "
|
|
150
|
+
f"(autooptm unlock {j.get('id')})")
|
|
151
|
+
if j.get("report_url"):
|
|
152
|
+
print("report: " + str(j["report_url"]))
|
|
153
|
+
elif s == "failed":
|
|
154
|
+
print(f"failed: {j.get('failure_reason')} — nothing charged")
|
|
155
|
+
else:
|
|
156
|
+
print(s)
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
if __name__ == "__main__":
|
|
160
|
+
raise SystemExit(main())
|
|
@@ -0,0 +1,118 @@
|
|
|
1
|
+
"""MCP server: lets any MCP-capable agent (Claude Code, Cursor, ...) submit
|
|
2
|
+
code to AutoOptm and bring back a measured speedup.
|
|
3
|
+
|
|
4
|
+
Run: AUTOOPTM_TOKEN=... autooptm-mcp (stdio transport)
|
|
5
|
+
|
|
6
|
+
Client config (e.g. .mcp.json):
|
|
7
|
+
{ "mcpServers": { "autooptm": {
|
|
8
|
+
"command": "autooptm-mcp",
|
|
9
|
+
"env": { "AUTOOPTM_TOKEN": "..." } } } }
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
from __future__ import annotations
|
|
13
|
+
|
|
14
|
+
from typing import Optional
|
|
15
|
+
|
|
16
|
+
try:
|
|
17
|
+
from mcp.server.fastmcp import FastMCP
|
|
18
|
+
except ImportError as e: # pragma: no cover
|
|
19
|
+
raise SystemExit("install the extra first: pip install 'autooptm[mcp]'") from e
|
|
20
|
+
|
|
21
|
+
from . import AutoOptm, AutoOptmError
|
|
22
|
+
|
|
23
|
+
mcp = FastMCP("autooptm")
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _client() -> AutoOptm:
|
|
27
|
+
return AutoOptm()
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def _err(e: AutoOptmError) -> dict:
|
|
31
|
+
return {"error": e.body.get("error", "request_failed"),
|
|
32
|
+
"hint": e.body.get("hint"), "status": e.status}
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
@mcp.tool()
|
|
36
|
+
def optimize_submit(source: str, entrypoint: str = "main.py",
|
|
37
|
+
workload: str = "training", gpu: Optional[str] = None,
|
|
38
|
+
args: Optional[str] = None, setup: Optional[str] = None,
|
|
39
|
+
git_ref: Optional[str] = None) -> dict:
|
|
40
|
+
"""Submit code to AutoOptm for GPU optimization. `source` is a public
|
|
41
|
+
git URL (github/gitlab) or a local directory path (zipped and uploaded,
|
|
42
|
+
skipping .git/caches/checkpoints). Diagnosis and measurement are free;
|
|
43
|
+
money moves only at unlock_patch. Returns {jobId, status}."""
|
|
44
|
+
try:
|
|
45
|
+
return _client().submit(source, entrypoint=entrypoint, workload=workload,
|
|
46
|
+
gpu=gpu, args=args, setup=setup, git_ref=git_ref)
|
|
47
|
+
except AutoOptmError as e:
|
|
48
|
+
return _err(e)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
@mcp.tool()
|
|
52
|
+
def optimize_status(job_id: str) -> dict:
|
|
53
|
+
"""Current state of a job: status, speedup, baseline/optimised medians,
|
|
54
|
+
unlock_credits (patch price), patch_unlocked, report_url."""
|
|
55
|
+
try:
|
|
56
|
+
return _client().job(job_id)
|
|
57
|
+
except AutoOptmError as e:
|
|
58
|
+
return _err(e)
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
@mcp.tool()
|
|
62
|
+
def optimize_wait(job_id: str, timeout_s: int = 1800) -> dict:
|
|
63
|
+
"""Block until the job finishes (succeeded/failed), polling every 10s.
|
|
64
|
+
A full optimization run typically takes tens of minutes."""
|
|
65
|
+
try:
|
|
66
|
+
return _client().wait(job_id, timeout_s=timeout_s)
|
|
67
|
+
except AutoOptmError as e:
|
|
68
|
+
return _err(e)
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
@mcp.tool()
|
|
72
|
+
def optimize_cancel(job_id: str) -> dict:
|
|
73
|
+
"""Stop a queued or running optimization. Queued jobs cancel at once;
|
|
74
|
+
running jobs stop at their next heartbeat. Nothing is charged."""
|
|
75
|
+
try:
|
|
76
|
+
return _client().cancel(job_id)
|
|
77
|
+
except AutoOptmError as e:
|
|
78
|
+
return _err(e)
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
@mcp.tool()
|
|
82
|
+
def unlock_patch(job_id: str, confirmed_by_user: bool = False) -> dict:
|
|
83
|
+
"""Pay the quoted credits and unlock the patch. SPENDS REAL CREDITS:
|
|
84
|
+
call only after the human has explicitly approved the price shown in
|
|
85
|
+
optimize_status (unlock_credits). Set confirmed_by_user=true to attest."""
|
|
86
|
+
if not confirmed_by_user:
|
|
87
|
+
return {"error": "confirmation_required",
|
|
88
|
+
"hint": "show the user unlock_credits and ask before paying"}
|
|
89
|
+
try:
|
|
90
|
+
return _client().unlock(job_id)
|
|
91
|
+
except AutoOptmError as e:
|
|
92
|
+
return _err(e)
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
@mcp.tool()
|
|
96
|
+
def download_patch(job_id: str, dest: str = "autooptm.patch") -> dict:
|
|
97
|
+
"""Download an unlocked patch to `dest`. Apply with `git apply`."""
|
|
98
|
+
try:
|
|
99
|
+
return {"path": _client().download_patch(job_id, dest)}
|
|
100
|
+
except AutoOptmError as e:
|
|
101
|
+
return _err(e)
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
@mcp.tool()
|
|
105
|
+
def account_balance() -> dict:
|
|
106
|
+
"""Credits available on the signed-in AutoOptm account."""
|
|
107
|
+
try:
|
|
108
|
+
return _client().balance()
|
|
109
|
+
except AutoOptmError as e:
|
|
110
|
+
return _err(e)
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
def main() -> None:
|
|
114
|
+
mcp.run()
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
if __name__ == "__main__":
|
|
118
|
+
main()
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "autooptm"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Client, CLI and MCP server for AutoOptm — measured GPU speedups for your training and inference code"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
# Zero runtime dependencies on purpose: the client is stdlib urllib.
|
|
13
|
+
dependencies = []
|
|
14
|
+
|
|
15
|
+
[project.optional-dependencies]
|
|
16
|
+
mcp = ["mcp>=1.0"]
|
|
17
|
+
|
|
18
|
+
[project.scripts]
|
|
19
|
+
autooptm = "autooptm.cli:main"
|
|
20
|
+
autooptm-mcp = "autooptm.mcp_server:main"
|
|
21
|
+
|
|
22
|
+
[tool.hatch.build.targets.wheel]
|
|
23
|
+
packages = ["autooptm"]
|