lockbot 2.1.2__tar.gz → 2.2.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.
- {lockbot-2.1.2/python/lockbot.egg-info → lockbot-2.2.0}/PKG-INFO +6 -3
- {lockbot-2.1.2 → lockbot-2.2.0}/README.md +5 -2
- {lockbot-2.1.2 → lockbot-2.2.0}/pyproject.toml +1 -1
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/admin/router.py +71 -9
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/auth/schemas.py +7 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/bots/models.py +5 -1
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/bots/router.py +64 -36
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/config.py +3 -1
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/main.py +28 -1
- {lockbot-2.1.2 → lockbot-2.2.0/python/lockbot.egg-info}/PKG-INFO +6 -3
- {lockbot-2.1.2 → lockbot-2.2.0}/LICENSE +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/MANIFEST.in +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/admin/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/auth/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/auth/dependencies.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/auth/models.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/auth/router.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/bots/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/bots/encryption.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/bots/manager.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/bots/schemas.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/bots/webhook_handler.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/database.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/logs/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/settings/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/settings/models.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/backend/app/settings/router.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/base_bot.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/bot_instance.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/config.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/device_bot.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/device_usage_alert.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/device_usage_utils.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/entry.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/env.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/handler.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/i18n/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/i18n/en.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/i18n/zh.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/io.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/message_adapter.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/msg_utils.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/node_bot.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/platforms/__init__.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/platforms/infoflow.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/queue_bot.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/request.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot/core/utils.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot.egg-info/SOURCES.txt +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot.egg-info/dependency_links.txt +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot.egg-info/requires.txt +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/python/lockbot.egg-info/top_level.txt +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/setup.cfg +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/tools/create_super_admin.py +0 -0
- {lockbot-2.1.2 → lockbot-2.2.0}/tools/gen_keys.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: lockbot
|
|
3
|
-
Version: 2.
|
|
3
|
+
Version: 2.2.0
|
|
4
4
|
Summary: Cluster resource management bot for IM platforms
|
|
5
5
|
Author-email: Jianbang Yang <yangjianbang112@gmail.com>
|
|
6
6
|
License: MIT
|
|
@@ -47,6 +47,9 @@ Supports both standalone Flask deployment and a full platform mode with FastAPI
|
|
|
47
47
|
|
|
48
48
|
[中文文档](README_CN.md) | [Live Demo](https://dynamicheart.github.io/lockbot/)
|
|
49
49
|
|
|
50
|
+
[](https://pypi.org/project/lockbot/)
|
|
51
|
+
[](https://github.com/DynamicHeart/lockbot/pkgs/container/lockbot)
|
|
52
|
+
|
|
50
53
|
## Features
|
|
51
54
|
|
|
52
55
|
- **Device Lock Bot** — Lock/unlock individual GPUs or devices on a cluster
|
|
@@ -100,14 +103,14 @@ docker pull ghcr.io/dynamicheart/lockbot:latest
|
|
|
100
103
|
docker run -d --name lockbot -p 8000:8000 \
|
|
101
104
|
-e JWT_SECRET=your-secret \
|
|
102
105
|
-e ENCRYPTION_KEY=your-fernet-key \
|
|
103
|
-
-v lockbot-data:/
|
|
106
|
+
-v lockbot-data:/data \
|
|
104
107
|
ghcr.io/dynamicheart/lockbot:latest
|
|
105
108
|
|
|
106
109
|
# 4. Create super_admin (password auto-generated and printed)
|
|
107
110
|
docker exec -it lockbot python tools/create_super_admin.py --username admin --email admin@example.com
|
|
108
111
|
```
|
|
109
112
|
|
|
110
|
-
> **
|
|
113
|
+
> **Data persistence**: All data (SQLite DB, bot state files) stored under `/data`. Override with `DATA_DIR` env var.
|
|
111
114
|
|
|
112
115
|
## Bot Configuration
|
|
113
116
|
|
|
@@ -7,6 +7,9 @@ Supports both standalone Flask deployment and a full platform mode with FastAPI
|
|
|
7
7
|
|
|
8
8
|
[中文文档](README_CN.md) | [Live Demo](https://dynamicheart.github.io/lockbot/)
|
|
9
9
|
|
|
10
|
+
[](https://pypi.org/project/lockbot/)
|
|
11
|
+
[](https://github.com/DynamicHeart/lockbot/pkgs/container/lockbot)
|
|
12
|
+
|
|
10
13
|
## Features
|
|
11
14
|
|
|
12
15
|
- **Device Lock Bot** — Lock/unlock individual GPUs or devices on a cluster
|
|
@@ -60,14 +63,14 @@ docker pull ghcr.io/dynamicheart/lockbot:latest
|
|
|
60
63
|
docker run -d --name lockbot -p 8000:8000 \
|
|
61
64
|
-e JWT_SECRET=your-secret \
|
|
62
65
|
-e ENCRYPTION_KEY=your-fernet-key \
|
|
63
|
-
-v lockbot-data:/
|
|
66
|
+
-v lockbot-data:/data \
|
|
64
67
|
ghcr.io/dynamicheart/lockbot:latest
|
|
65
68
|
|
|
66
69
|
# 4. Create super_admin (password auto-generated and printed)
|
|
67
70
|
docker exec -it lockbot python tools/create_super_admin.py --username admin --email admin@example.com
|
|
68
71
|
```
|
|
69
72
|
|
|
70
|
-
> **
|
|
73
|
+
> **Data persistence**: All data (SQLite DB, bot state files) stored under `/data`. Override with `DATA_DIR` env var.
|
|
71
74
|
|
|
72
75
|
## Bot Configuration
|
|
73
76
|
|
|
@@ -3,13 +3,16 @@ Admin API routes.
|
|
|
3
3
|
"""
|
|
4
4
|
|
|
5
5
|
import contextlib
|
|
6
|
+
import io
|
|
7
|
+
import json
|
|
6
8
|
import os
|
|
7
9
|
import shutil
|
|
8
10
|
import tempfile
|
|
11
|
+
import zipfile
|
|
9
12
|
from datetime import datetime
|
|
10
13
|
|
|
11
14
|
from fastapi import APIRouter, Depends, HTTPException
|
|
12
|
-
from fastapi.responses import FileResponse
|
|
15
|
+
from fastapi.responses import FileResponse, Response
|
|
13
16
|
from sqlalchemy.orm import Session
|
|
14
17
|
from starlette.background import BackgroundTask
|
|
15
18
|
|
|
@@ -21,6 +24,7 @@ from lockbot.backend.app.auth.schemas import (
|
|
|
21
24
|
AdminEditUser,
|
|
22
25
|
PasswordResetOut,
|
|
23
26
|
UserOut,
|
|
27
|
+
UserOutWithStats,
|
|
24
28
|
)
|
|
25
29
|
from lockbot.backend.app.bots.models import Bot
|
|
26
30
|
from lockbot.backend.app.database import get_db
|
|
@@ -68,18 +72,36 @@ def admin_create_user(
|
|
|
68
72
|
return PasswordResetOut(id=user.id, username=user.username, new_password=raw_password)
|
|
69
73
|
|
|
70
74
|
|
|
71
|
-
@router.get("/users", response_model=list[
|
|
75
|
+
@router.get("/users", response_model=list[UserOutWithStats])
|
|
72
76
|
def list_users(
|
|
73
77
|
operator: User = Depends(require_admin),
|
|
74
78
|
db: Session = Depends(get_db),
|
|
75
79
|
):
|
|
76
80
|
"""List users visible to the operator.
|
|
77
|
-
Super_admin sees all; admin sees only users (not admins/super_admins).
|
|
81
|
+
Super_admin sees all; admin sees only users (not admins/super_admins).
|
|
82
|
+
Includes bot_count and running_count for each user."""
|
|
78
83
|
if operator.role == "super_admin":
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
84
|
+
users = db.query(User).all()
|
|
85
|
+
else:
|
|
86
|
+
# admin: see regular users + self
|
|
87
|
+
users = db.query(User).filter((User.role == "user") | (User.id == operator.id)).all()
|
|
88
|
+
|
|
89
|
+
# Attach bot stats for each user
|
|
90
|
+
result = []
|
|
91
|
+
for u in users:
|
|
92
|
+
data = UserOut.model_validate(u)
|
|
93
|
+
bot_count = db.query(Bot).filter(Bot.user_id == u.id, Bot.is_deleted.is_(False)).count()
|
|
94
|
+
running_count = (
|
|
95
|
+
db.query(Bot)
|
|
96
|
+
.filter(
|
|
97
|
+
Bot.user_id == u.id,
|
|
98
|
+
Bot.status == "running",
|
|
99
|
+
Bot.is_deleted.is_(False),
|
|
100
|
+
)
|
|
101
|
+
.count()
|
|
102
|
+
)
|
|
103
|
+
result.append({**data.model_dump(), "bot_count": bot_count, "running_count": running_count})
|
|
104
|
+
return result
|
|
83
105
|
|
|
84
106
|
|
|
85
107
|
@router.put("/users/{user_id}", response_model=UserOut)
|
|
@@ -169,7 +191,7 @@ def list_all_bots(
|
|
|
169
191
|
_admin: User = Depends(require_admin),
|
|
170
192
|
db: Session = Depends(get_db),
|
|
171
193
|
):
|
|
172
|
-
rows = db.query(Bot, User.username).join(User, Bot.user_id == User.id).all()
|
|
194
|
+
rows = db.query(Bot, User.username).join(User, Bot.user_id == User.id).filter(Bot.is_deleted.is_(False)).all()
|
|
173
195
|
return [
|
|
174
196
|
{c.name: getattr(bot, c.name) for c in bot.__table__.columns} | {"owner": username} for bot, username in rows
|
|
175
197
|
]
|
|
@@ -181,7 +203,7 @@ def platform_stats(
|
|
|
181
203
|
db: Session = Depends(get_db),
|
|
182
204
|
):
|
|
183
205
|
total_users = db.query(User).count()
|
|
184
|
-
bots = db.query(Bot).all()
|
|
206
|
+
bots = db.query(Bot).filter(Bot.is_deleted.is_(False)).all()
|
|
185
207
|
return {
|
|
186
208
|
"totalUsers": total_users,
|
|
187
209
|
"totalBots": len(bots),
|
|
@@ -222,3 +244,43 @@ def download_backup(
|
|
|
222
244
|
media_type="application/x-sqlite3",
|
|
223
245
|
background=BackgroundTask(os.unlink, tmp_path),
|
|
224
246
|
)
|
|
247
|
+
|
|
248
|
+
|
|
249
|
+
_DEFAULT_DATA_DIR = os.environ.get("DATA_DIR", "/data")
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
@router.get("/bot-states")
|
|
253
|
+
def download_all_bot_states(
|
|
254
|
+
_admin: User = Depends(require_super_admin),
|
|
255
|
+
db: Session = Depends(get_db),
|
|
256
|
+
):
|
|
257
|
+
"""Download all bot state files as a zip archive (super_admin only)."""
|
|
258
|
+
bots = db.query(Bot).filter(Bot.is_deleted.is_(False)).all()
|
|
259
|
+
|
|
260
|
+
buf = io.BytesIO()
|
|
261
|
+
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
|
262
|
+
for bot in bots:
|
|
263
|
+
# Determine data dir from config_overrides or default
|
|
264
|
+
data_dir = _DEFAULT_DATA_DIR
|
|
265
|
+
try:
|
|
266
|
+
overrides = json.loads(bot.config_overrides or "{}")
|
|
267
|
+
data_dir = overrides.get("DATA_DIR") or _DEFAULT_DATA_DIR
|
|
268
|
+
except (json.JSONDecodeError, AttributeError):
|
|
269
|
+
pass
|
|
270
|
+
state_file = os.path.join(data_dir, "bots", str(bot.id), "bot_state.json")
|
|
271
|
+
folder_name = f"bot_{bot.id}_{bot.name}"
|
|
272
|
+
if os.path.exists(state_file):
|
|
273
|
+
zf.write(state_file, f"{folder_name}/bot_state.json")
|
|
274
|
+
else:
|
|
275
|
+
# Create an empty placeholder so the admin knows the bot has no state
|
|
276
|
+
zf.writestr(f"{folder_name}/bot_state.json", "{}")
|
|
277
|
+
|
|
278
|
+
buf.seek(0)
|
|
279
|
+
content = buf.read()
|
|
280
|
+
|
|
281
|
+
filename = f"lockbot_states_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
|
|
282
|
+
return Response(
|
|
283
|
+
content=content,
|
|
284
|
+
media_type="application/zip",
|
|
285
|
+
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
|
|
286
|
+
)
|
|
@@ -30,6 +30,13 @@ class UserOut(BaseModel):
|
|
|
30
30
|
model_config = {"from_attributes": True}
|
|
31
31
|
|
|
32
32
|
|
|
33
|
+
class UserOutWithStats(UserOut):
|
|
34
|
+
"""User info with bot statistics."""
|
|
35
|
+
|
|
36
|
+
bot_count: int = 0
|
|
37
|
+
running_count: int = 0
|
|
38
|
+
|
|
39
|
+
|
|
33
40
|
class TokenOut(BaseModel):
|
|
34
41
|
access_token: str
|
|
35
42
|
token_type: str = "bearer"
|
|
@@ -6,7 +6,7 @@ from __future__ import annotations
|
|
|
6
6
|
|
|
7
7
|
from datetime import datetime
|
|
8
8
|
|
|
9
|
-
from sqlalchemy import DateTime, ForeignKey, Integer, String, Text, func
|
|
9
|
+
from sqlalchemy import Boolean, DateTime, ForeignKey, Integer, String, Text, func
|
|
10
10
|
from sqlalchemy.orm import Mapped, mapped_column
|
|
11
11
|
|
|
12
12
|
from lockbot.backend.app.database import Base
|
|
@@ -35,6 +35,10 @@ class Bot(Base):
|
|
|
35
35
|
last_request_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
|
|
36
36
|
config_overrides: Mapped[str] = mapped_column(Text, default="{}") # JSON
|
|
37
37
|
|
|
38
|
+
# Soft delete
|
|
39
|
+
is_deleted: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
|
|
40
|
+
deleted_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
|
|
41
|
+
|
|
38
42
|
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now())
|
|
39
43
|
updated_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), onupdate=func.now())
|
|
40
44
|
|
|
@@ -15,7 +15,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request, status
|
|
|
15
15
|
from fastapi.responses import PlainTextResponse
|
|
16
16
|
from sqlalchemy.orm import Session
|
|
17
17
|
|
|
18
|
-
from lockbot.backend.app.auth.dependencies import get_current_user, require_admin
|
|
18
|
+
from lockbot.backend.app.auth.dependencies import get_current_user, require_admin, require_super_admin
|
|
19
19
|
from lockbot.backend.app.auth.models import User
|
|
20
20
|
from lockbot.backend.app.bots import encryption
|
|
21
21
|
from lockbot.backend.app.bots.manager import bot_manager
|
|
@@ -57,7 +57,7 @@ def _write_log(bot_id: int, message: str, level: str = "INFO", category: str = "
|
|
|
57
57
|
"category": category,
|
|
58
58
|
"level": level,
|
|
59
59
|
"message": message,
|
|
60
|
-
"created_at": datetime.utcnow().isoformat(),
|
|
60
|
+
"created_at": datetime.utcnow().isoformat() + "Z",
|
|
61
61
|
}
|
|
62
62
|
try:
|
|
63
63
|
with open(log_path, "a", encoding="utf-8") as f:
|
|
@@ -101,7 +101,7 @@ def _mask_bot(bot: Bot, db: Session | None = None) -> dict:
|
|
|
101
101
|
def _get_user_bot(bot_id: int, user: User, db: Session) -> Bot:
|
|
102
102
|
"""Fetch a bot owned by the user (or any bot if admin), or raise 404."""
|
|
103
103
|
bot = db.get(Bot, bot_id)
|
|
104
|
-
if not bot or (bot.user_id != user.id and user.role not in ("admin", "super_admin")):
|
|
104
|
+
if not bot or bot.is_deleted or (bot.user_id != user.id and user.role not in ("admin", "super_admin")):
|
|
105
105
|
raise HTTPException(status_code=404, detail="Bot not found")
|
|
106
106
|
return bot
|
|
107
107
|
|
|
@@ -167,7 +167,7 @@ def list_bots(
|
|
|
167
167
|
user: User = Depends(get_current_user),
|
|
168
168
|
db: Session = Depends(get_db),
|
|
169
169
|
):
|
|
170
|
-
return db.query(Bot).filter(Bot.user_id == user.id).all()
|
|
170
|
+
return db.query(Bot).filter(Bot.user_id == user.id, Bot.is_deleted.is_(False)).all()
|
|
171
171
|
|
|
172
172
|
|
|
173
173
|
@router.post("", response_model=BotOut, status_code=status.HTTP_201_CREATED)
|
|
@@ -179,7 +179,15 @@ def create_bot(
|
|
|
179
179
|
if body.bot_type.upper() not in VALID_BOT_TYPES:
|
|
180
180
|
raise HTTPException(status_code=422, detail=f"Invalid bot_type, must be one of {VALID_BOT_TYPES}")
|
|
181
181
|
|
|
182
|
-
exists =
|
|
182
|
+
exists = (
|
|
183
|
+
db.query(Bot)
|
|
184
|
+
.filter(
|
|
185
|
+
Bot.user_id == user.id,
|
|
186
|
+
Bot.name == body.name,
|
|
187
|
+
Bot.is_deleted.is_(False),
|
|
188
|
+
)
|
|
189
|
+
.first()
|
|
190
|
+
)
|
|
183
191
|
if exists:
|
|
184
192
|
raise HTTPException(status_code=409, detail="Bot name already exists")
|
|
185
193
|
|
|
@@ -202,7 +210,15 @@ def create_bot(
|
|
|
202
210
|
# Auto-start the bot if within quota
|
|
203
211
|
auto_started = False
|
|
204
212
|
if user.role not in ("admin", "super_admin"):
|
|
205
|
-
running_count =
|
|
213
|
+
running_count = (
|
|
214
|
+
db.query(Bot)
|
|
215
|
+
.filter(
|
|
216
|
+
Bot.user_id == user.id,
|
|
217
|
+
Bot.status == "running",
|
|
218
|
+
Bot.is_deleted.is_(False),
|
|
219
|
+
)
|
|
220
|
+
.count()
|
|
221
|
+
)
|
|
206
222
|
if running_count >= user.max_running_bots:
|
|
207
223
|
pass # Quota reached — bot created but not started
|
|
208
224
|
else:
|
|
@@ -237,7 +253,7 @@ def get_running_states(
|
|
|
237
253
|
db: Session = Depends(get_db),
|
|
238
254
|
):
|
|
239
255
|
"""Return state for all bots owned by the user (running from memory, stopped from file)."""
|
|
240
|
-
user_bots = db.query(Bot).filter(Bot.user_id == user.id).all()
|
|
256
|
+
user_bots = db.query(Bot).filter(Bot.user_id == user.id, Bot.is_deleted.is_(False)).all()
|
|
241
257
|
result = {}
|
|
242
258
|
for bot in user_bots:
|
|
243
259
|
instance = bot_manager.get_instance(bot.id)
|
|
@@ -277,16 +293,27 @@ def update_bot(
|
|
|
277
293
|
):
|
|
278
294
|
bot = _get_user_bot(bot_id, user, db)
|
|
279
295
|
|
|
280
|
-
#
|
|
281
|
-
if user.role
|
|
282
|
-
|
|
283
|
-
if owner and owner.role in ("admin", "super_admin"):
|
|
284
|
-
raise HTTPException(status_code=403, detail="Cannot edit another admin's bot")
|
|
296
|
+
# Only owner or super_admin can edit
|
|
297
|
+
if user.role != "super_admin" and bot.user_id != user.id:
|
|
298
|
+
raise HTTPException(status_code=403, detail="Cannot edit another user's bot")
|
|
285
299
|
|
|
286
300
|
if bot.status == "running":
|
|
287
301
|
raise HTTPException(status_code=409, detail="Cannot update a running bot. Stop it first.")
|
|
288
302
|
|
|
289
303
|
if body.name is not None:
|
|
304
|
+
# Check for duplicate name (excluding current bot and deleted bots)
|
|
305
|
+
dup = (
|
|
306
|
+
db.query(Bot)
|
|
307
|
+
.filter(
|
|
308
|
+
Bot.user_id == bot.user_id,
|
|
309
|
+
Bot.name == body.name,
|
|
310
|
+
Bot.id != bot_id,
|
|
311
|
+
Bot.is_deleted.is_(False),
|
|
312
|
+
)
|
|
313
|
+
.first()
|
|
314
|
+
)
|
|
315
|
+
if dup:
|
|
316
|
+
raise HTTPException(status_code=409, detail="Bot name already exists")
|
|
290
317
|
bot.name = body.name
|
|
291
318
|
if body.group_id is not None:
|
|
292
319
|
bot.group_id = body.group_id
|
|
@@ -314,16 +341,19 @@ def delete_bot(
|
|
|
314
341
|
):
|
|
315
342
|
bot = _get_user_bot(bot_id, user, db)
|
|
316
343
|
|
|
317
|
-
#
|
|
318
|
-
if user.role
|
|
319
|
-
|
|
320
|
-
if owner and owner.role in ("admin", "super_admin"):
|
|
321
|
-
raise HTTPException(status_code=403, detail="Cannot delete another admin's bot")
|
|
344
|
+
# Only owner or super_admin can delete
|
|
345
|
+
if user.role != "super_admin" and bot.user_id != user.id:
|
|
346
|
+
raise HTTPException(status_code=403, detail="Cannot delete another user's bot")
|
|
322
347
|
|
|
323
|
-
|
|
324
|
-
|
|
348
|
+
# Stop the bot first if running
|
|
349
|
+
if bot.status in ("running", "error"):
|
|
350
|
+
with contextlib.suppress(RuntimeError):
|
|
351
|
+
bot_manager.stop_bot(bot_id)
|
|
352
|
+
bot.status = "stopped"
|
|
353
|
+
bot.pid = None
|
|
325
354
|
|
|
326
|
-
|
|
355
|
+
bot.is_deleted = True
|
|
356
|
+
bot.deleted_at = datetime.utcnow()
|
|
327
357
|
db.commit()
|
|
328
358
|
|
|
329
359
|
|
|
@@ -331,20 +361,14 @@ def delete_bot(
|
|
|
331
361
|
def transfer_bot_owner(
|
|
332
362
|
bot_id: int,
|
|
333
363
|
body: dict,
|
|
334
|
-
user: User = Depends(
|
|
364
|
+
user: User = Depends(require_super_admin),
|
|
335
365
|
db: Session = Depends(get_db),
|
|
336
366
|
):
|
|
337
|
-
"""Transfer bot ownership.
|
|
367
|
+
"""Transfer bot ownership. Super_admin can transfer all bots."""
|
|
338
368
|
bot = db.get(Bot, bot_id)
|
|
339
|
-
if not bot:
|
|
369
|
+
if not bot or bot.is_deleted:
|
|
340
370
|
raise HTTPException(status_code=404, detail="Bot not found")
|
|
341
371
|
|
|
342
|
-
# Permission: admin cannot transfer other admins' bots
|
|
343
|
-
if user.role == "admin" and bot.user_id != user.id:
|
|
344
|
-
owner = db.get(User, bot.user_id)
|
|
345
|
-
if owner and owner.role in ("admin", "super_admin"):
|
|
346
|
-
raise HTTPException(status_code=403, detail="Cannot transfer another admin's bot")
|
|
347
|
-
|
|
348
372
|
target_username = (body.get("username") or "").strip()
|
|
349
373
|
if not target_username:
|
|
350
374
|
raise HTTPException(status_code=400, detail="username is required")
|
|
@@ -353,10 +377,6 @@ def transfer_bot_owner(
|
|
|
353
377
|
if not target_user:
|
|
354
378
|
raise HTTPException(status_code=404, detail="Target user not found")
|
|
355
379
|
|
|
356
|
-
# Admin cannot transfer bot to another admin
|
|
357
|
-
if user.role == "admin" and target_user.role in ("admin", "super_admin"):
|
|
358
|
-
raise HTTPException(status_code=403, detail="Cannot transfer bot to an admin")
|
|
359
|
-
|
|
360
380
|
old_owner = db.get(User, bot.user_id)
|
|
361
381
|
old_name = old_owner.username if old_owner else "unknown"
|
|
362
382
|
bot.user_id = target_user.id
|
|
@@ -391,7 +411,15 @@ def start_bot(
|
|
|
391
411
|
|
|
392
412
|
# Enforce max_running_bots quota (admins are exempt)
|
|
393
413
|
if user.role not in ("admin", "super_admin"):
|
|
394
|
-
running_count =
|
|
414
|
+
running_count = (
|
|
415
|
+
db.query(Bot)
|
|
416
|
+
.filter(
|
|
417
|
+
Bot.user_id == user.id,
|
|
418
|
+
Bot.status == "running",
|
|
419
|
+
Bot.is_deleted.is_(False),
|
|
420
|
+
)
|
|
421
|
+
.count()
|
|
422
|
+
)
|
|
395
423
|
if running_count >= user.max_running_bots:
|
|
396
424
|
raise HTTPException(
|
|
397
425
|
status_code=403,
|
|
@@ -889,7 +917,7 @@ async def webhook(bot_id: int, request: Request, db: Session = Depends(get_db)):
|
|
|
889
917
|
|
|
890
918
|
# Update last_request_at
|
|
891
919
|
bot = db.get(Bot, bot_id)
|
|
892
|
-
if bot:
|
|
920
|
+
if bot and not bot.is_deleted:
|
|
893
921
|
bot.last_request_at = datetime.utcnow()
|
|
894
922
|
db.commit()
|
|
895
923
|
|
|
@@ -921,7 +949,7 @@ async def _reply_bot_not_running(
|
|
|
921
949
|
) -> PlainTextResponse:
|
|
922
950
|
"""When a bot is not running, POST a 'not started' message back to the IM platform."""
|
|
923
951
|
bot = db.get(Bot, bot_id)
|
|
924
|
-
if not bot:
|
|
952
|
+
if not bot or bot.is_deleted:
|
|
925
953
|
raise HTTPException(status_code=404, detail="Bot not found")
|
|
926
954
|
|
|
927
955
|
# URL verification must still work even when the bot is stopped
|
|
@@ -8,10 +8,12 @@ from pathlib import Path
|
|
|
8
8
|
# Project root directory (lockbot/)
|
|
9
9
|
BASE_DIR = Path(__file__).resolve().parent.parent.parent
|
|
10
10
|
|
|
11
|
+
DATA_DIR = os.environ.get("DATA_DIR", "/data")
|
|
12
|
+
|
|
11
13
|
# Database
|
|
12
14
|
DATABASE_URL = os.environ.get(
|
|
13
15
|
"DATABASE_URL",
|
|
14
|
-
f"sqlite:///{
|
|
16
|
+
f"sqlite:///{os.path.join(DATA_DIR, 'lockbot.db')}",
|
|
15
17
|
)
|
|
16
18
|
|
|
17
19
|
# JWT
|
|
@@ -72,6 +72,25 @@ def _migrate_users_must_change_password():
|
|
|
72
72
|
logger.info("Migrated users: added 'must_change_password' column")
|
|
73
73
|
|
|
74
74
|
|
|
75
|
+
def _migrate_bot_soft_delete():
|
|
76
|
+
"""Add 'is_deleted' and 'deleted_at' columns to bots if they don't exist (SQLite migration)."""
|
|
77
|
+
from sqlalchemy import inspect as sa_inspect
|
|
78
|
+
from sqlalchemy import text
|
|
79
|
+
|
|
80
|
+
insp = sa_inspect(engine)
|
|
81
|
+
if "bots" not in insp.get_table_names():
|
|
82
|
+
return
|
|
83
|
+
columns = [c["name"] for c in insp.get_columns("bots")]
|
|
84
|
+
if "is_deleted" not in columns:
|
|
85
|
+
with engine.begin() as conn:
|
|
86
|
+
conn.execute(text("ALTER TABLE bots ADD COLUMN is_deleted BOOLEAN NOT NULL DEFAULT 0"))
|
|
87
|
+
logger.info("Migrated bots: added 'is_deleted' column")
|
|
88
|
+
if "deleted_at" not in columns:
|
|
89
|
+
with engine.begin() as conn:
|
|
90
|
+
conn.execute(text("ALTER TABLE bots ADD COLUMN deleted_at DATETIME"))
|
|
91
|
+
logger.info("Migrated bots: added 'deleted_at' column")
|
|
92
|
+
|
|
93
|
+
|
|
75
94
|
def _seed_dev_admin():
|
|
76
95
|
"""Create admin user in dev mode if it doesn't exist."""
|
|
77
96
|
from lockbot.backend.app.config import (
|
|
@@ -153,6 +172,7 @@ async def lifespan(app: FastAPI):
|
|
|
153
172
|
_migrate_bot_logs_category()
|
|
154
173
|
_migrate_bot_consecutive_failures()
|
|
155
174
|
_migrate_users_must_change_password()
|
|
175
|
+
_migrate_bot_soft_delete()
|
|
156
176
|
_seed_dev_admin()
|
|
157
177
|
_seed_dev_users()
|
|
158
178
|
_reset_running_bots()
|
|
@@ -196,7 +216,14 @@ def _reset_running_bots():
|
|
|
196
216
|
db = SessionLocal()
|
|
197
217
|
try:
|
|
198
218
|
# Collect all bots that need recovery (running + error)
|
|
199
|
-
recover_bots =
|
|
219
|
+
recover_bots = (
|
|
220
|
+
db.query(Bot)
|
|
221
|
+
.filter(
|
|
222
|
+
Bot.status.in_(["running", "error"]),
|
|
223
|
+
Bot.is_deleted.is_(False),
|
|
224
|
+
)
|
|
225
|
+
.all()
|
|
226
|
+
)
|
|
200
227
|
if not recover_bots:
|
|
201
228
|
return
|
|
202
229
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: lockbot
|
|
3
|
-
Version: 2.
|
|
3
|
+
Version: 2.2.0
|
|
4
4
|
Summary: Cluster resource management bot for IM platforms
|
|
5
5
|
Author-email: Jianbang Yang <yangjianbang112@gmail.com>
|
|
6
6
|
License: MIT
|
|
@@ -47,6 +47,9 @@ Supports both standalone Flask deployment and a full platform mode with FastAPI
|
|
|
47
47
|
|
|
48
48
|
[中文文档](README_CN.md) | [Live Demo](https://dynamicheart.github.io/lockbot/)
|
|
49
49
|
|
|
50
|
+
[](https://pypi.org/project/lockbot/)
|
|
51
|
+
[](https://github.com/DynamicHeart/lockbot/pkgs/container/lockbot)
|
|
52
|
+
|
|
50
53
|
## Features
|
|
51
54
|
|
|
52
55
|
- **Device Lock Bot** — Lock/unlock individual GPUs or devices on a cluster
|
|
@@ -100,14 +103,14 @@ docker pull ghcr.io/dynamicheart/lockbot:latest
|
|
|
100
103
|
docker run -d --name lockbot -p 8000:8000 \
|
|
101
104
|
-e JWT_SECRET=your-secret \
|
|
102
105
|
-e ENCRYPTION_KEY=your-fernet-key \
|
|
103
|
-
-v lockbot-data:/
|
|
106
|
+
-v lockbot-data:/data \
|
|
104
107
|
ghcr.io/dynamicheart/lockbot:latest
|
|
105
108
|
|
|
106
109
|
# 4. Create super_admin (password auto-generated and printed)
|
|
107
110
|
docker exec -it lockbot python tools/create_super_admin.py --username admin --email admin@example.com
|
|
108
111
|
```
|
|
109
112
|
|
|
110
|
-
> **
|
|
113
|
+
> **Data persistence**: All data (SQLite DB, bot state files) stored under `/data`. Override with `DATA_DIR` env var.
|
|
111
114
|
|
|
112
115
|
## Bot Configuration
|
|
113
116
|
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|