121 lines
3.9 KiB
Python
121 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
gatekeeper-client — 成员服务器接入铸渊算力调度池
|
|
|
|
用法：
|
|
python3 gatekeeper-client.py --register --name YOUR_NAME --master 43.139.217.141:3910
|
|
python3 gatekeeper-client.py --status
|
|
"""
|
|
import json, os, sys, time, socket, uuid, hashlib, subprocess
|
|
from datetime import datetime
|
|
|
|
CONFIG_DIR = os.path.expanduser('~/.gatekeeper')
|
|
CONFIG_FILE = os.path.join(CONFIG_DIR, 'config.json')
|
|
|
|
def log(m):
|
|
print(f'[{datetime.now().strftime("%H:%M:%S")}] {m}', flush=True)
|
|
|
|
def ensure_dir():
|
|
os.makedirs(CONFIG_DIR, exist_ok=True)
|
|
|
|
def load_config():
|
|
try:
|
|
with open(CONFIG_FILE) as f:
|
|
return json.load(f)
|
|
except:
|
|
return {}
|
|
|
|
def save_config(cfg):
|
|
ensure_dir()
|
|
with open(CONFIG_FILE, 'w') as f:
|
|
json.dump(cfg, f, indent=2)
|
|
|
|
def register(name, master):
|
|
"""向主控闸门注册"""
|
|
log(f'注册到 {master}')
|
|
host, port = master.split(':')
|
|
port = int(port)
|
|
|
|
# 生成本机身份
|
|
node_id = str(uuid.uuid4())[:8]
|
|
key = hashlib.sha256((node_id + name + str(time.time())).encode()).hexdigest()[:32]
|
|
|
|
# 通过 socket 发送注册请求
|
|
try:
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
s.settimeout(10)
|
|
s.connect((host, port))
|
|
req = json.dumps({
|
|
'action': 'register',
|
|
'name': name,
|
|
'node_id': node_id,
|
|
'key': key
|
|
})
|
|
s.sendall((req + '\n').encode())
|
|
resp = s.recv(4096).decode()
|
|
s.close()
|
|
|
|
data = json.loads(resp)
|
|
if data.get('status') == 'ok':
|
|
save_config({'name': name, 'node_id': node_id, 'key': key, 'master': master, 'status': 'registered'})
|
|
log('注册成功！身份密钥已保存在本地')
|
|
log(f'名称: {name}')
|
|
log(f'节点ID: {node_id}')
|
|
return True
|
|
else:
|
|
log(f'注册失败: {data.get("message", "?")}')
|
|
return False
|
|
except Exception as e:
|
|
log(f'连接失败: {e}')
|
|
log('请检查服务器是否能访问 ' + master)
|
|
return False
|
|
|
|
def show_status():
|
|
"""显示当前连接状态"""
|
|
cfg = load_config()
|
|
if not cfg:
|
|
log('未注册。请先运行: python3 gatekeeper-client.py --register --name YOUR_NAME --master IP:PORT')
|
|
return
|
|
log(f'名称: {cfg.get("name", "?")}')
|
|
log(f'节点ID: {cfg.get("node_id", "?")}')
|
|
log(f'主控闸门: {cfg.get("master", "?")}')
|
|
log(f'状态: {cfg.get("status", "?")}')
|
|
|
|
# 测试连接
|
|
master = cfg.get('master', '')
|
|
if master:
|
|
try:
|
|
host, port = master.split(':')
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
s.settimeout(5)
|
|
s.connect((host, int(port)))
|
|
s.sendall(json.dumps({'action': 'ping', 'node_id': cfg.get('node_id')})['\n'].encode())
|
|
resp = s.recv(4096).decode()
|
|
s.close()
|
|
data = json.loads(resp)
|
|
if data.get('status') == 'ok':
|
|
log('连接: ✅ connected')
|
|
else:
|
|
log('连接: ❌ connection error')
|
|
except Exception as e:
|
|
log(f'连接: ❌ {e}')
|
|
|
|
if __name__ == '__main__':
|
|
import argparse
|
|
parser = argparse.ArgumentParser(description='gatekeeper-client: 接入铸渊算力调度池')
|
|
parser.add_argument('--register', action='store_true', help='注册到主控闸门')
|
|
parser.add_argument('--name', help='你的名称 (如 awen)')
|
|
parser.add_argument('--master', default='43.139.217.141:3910', help='主控闸门地址')
|
|
parser.add_argument('--status', action='store_true', help='查看连接状态')
|
|
args = parser.parse_args()
|
|
|
|
if args.status:
|
|
show_status()
|
|
elif args.register:
|
|
if not args.name:
|
|
print('请提供 --name 参数')
|
|
sys.exit(1)
|
|
register(args.name, args.master)
|
|
else:
|
|
parser.print_help()
|