Skip to content

Commit a1c5c9b

Browse files
dpj135dpj135
committed
Added yr init
Signed-off-by: dpj135 <958208521@qq.com>
1 parent f1f041a commit a1c5c9b

2 files changed

Lines changed: 240 additions & 3 deletions

File tree

‎transfer_queue/config.yaml‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,4 +47,15 @@ backend:
4747
RayStore:
4848

4949
# For Yuanrong:
50-
# TODO
50+
Yuanrong:
51+
# Whether to let TQ automatically start etcd and datasystem services (default false)
52+
auto_init: false
53+
# etcd service address (used to start etcd when auto_init=true)
54+
etcd_address: "127.0.0.1:2379"
55+
# datasystem worker address (used to start dscli when auto_init=true)
56+
worker_address: "127.0.0.1:31501"
57+
# YuanrongStorageClient connection parameters (required)
58+
host: "127.0.0.1"
59+
port: 31501
60+
# Client name (optional, default is YuanrongStorageClient)
61+
client_name: "YuanrongStorageClient"

‎transfer_queue/interface.py‎

Lines changed: 228 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,10 @@
1616
import logging
1717
import math
1818
import os
19+
import shutil
20+
import socket
1921
import subprocess
22+
import tempfile
2023
import time
2124
from importlib import resources
2225
from typing import Any, Optional
@@ -102,11 +105,11 @@ def _maybe_create_transferqueue_storage(conf: DictConfig) -> DictConfig:
102105
check = subprocess.run(["pgrep", "-f", "mooncake_master"], stdout=subprocess.PIPE, text=True)
103106
if check.returncode == 0:
104107
pids = check.stdout.strip().replace("\n", ", ")
105-
logging.info(f"Find existing mooncake_master (PID: {pids}), try to kill first...")
108+
logger.info(f"Find existing mooncake_master (PID: {pids}), try to kill first...")
106109

107110
result = os.system('pkill -f "[m]ooncake_master"')
108111
if result == 0:
109-
logging.info("Successfully killed existing mooncake_master processes.")
112+
logger.info("Successfully killed existing mooncake_master processes.")
110113
else:
111114
raise RuntimeError(f"Failed to kill existing mooncake_master processes (exit code: {result}).")
112115

@@ -167,6 +170,193 @@ def _maybe_create_transferqueue_storage(conf: DictConfig) -> DictConfig:
167170
f"Output:\n{error_msg}"
168171
)
169172
_TRANSFER_QUEUE_STORAGE["MooncakeStore"] = process
173+
if conf.backend.storage_backend == "Yuanrong":
174+
if conf.backend.Yuanrong.auto_init:
175+
etcd_process = None
176+
etcd_data_dir = None
177+
178+
try:
179+
# ========== Start etcd (inlined from _start_etcd) ==========
180+
etcd_address = conf.backend.Yuanrong.etcd_address
181+
# Parse host and port
182+
if "://" in etcd_address:
183+
# Remove protocol prefix if present
184+
parsed = urlparse(etcd_address)
185+
host = parsed.hostname
186+
port = parsed.port
187+
else:
188+
# Assume host:port format
189+
parts = etcd_address.split(":")
190+
if len(parts) != 2:
191+
raise ValueError(f"Invalid etcd_address format: {etcd_address}. Expected host:port")
192+
host = parts[0]
193+
port = int(parts[1])
194+
195+
# Check if etcd is already running
196+
check = subprocess.run(["pgrep", "-f", "etcd"], stdout=subprocess.PIPE, text=True)
197+
if check.returncode == 0:
198+
pids = check.stdout.strip().replace("\n", ", ")
199+
logger.warning(f"Found existing etcd processes (PID: {pids}). TQ will not start a new one.")
200+
# Try to connect to see if it's listening on our desired port
201+
try:
202+
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
203+
sock.settimeout(2)
204+
result = sock.connect_ex((host, port))
205+
sock.close()
206+
if result == 0:
207+
logger.info(f"etcd is already listening on {host}:{port}")
208+
etcd_process = None
209+
etcd_data_dir = None
210+
else:
211+
logger.warning(
212+
f"etcd process exists but not listening on {host}:{port}, will start new instance"
213+
)
214+
# Continue to start new instance
215+
except Exception as e:
216+
logger.warning(f"Failed to check etcd port: {e}, will start new instance")
217+
# Continue to start new instance
218+
219+
if etcd_process is None:
220+
# Create temporary data directory
221+
etcd_data_dir = tempfile.mkdtemp(prefix="tq_etcd_")
222+
logger.info(f"Starting etcd with data directory: {etcd_data_dir}")
223+
224+
cmd = [
225+
"etcd",
226+
f"--data-dir={etcd_data_dir}",
227+
f"--listen-client-urls=http://0.0.0.0:{port}",
228+
f"--advertise-client-urls=http://{host}:{port}",
229+
"--listen-peer-urls=http://0.0.0.0:2380",
230+
f"--initial-advertise-peer-urls=http://{host}:2380",
231+
"--initial-cluster=default=http://{}:2380".format(host),
232+
]
233+
234+
log_file_path = "/tmp/etcd.log"
235+
with open(log_file_path, "w") as log_file:
236+
etcd_process = subprocess.Popen(
237+
cmd,
238+
stdout=log_file,
239+
stderr=subprocess.STDOUT,
240+
text=True,
241+
bufsize=1,
242+
universal_newlines=True,
243+
start_new_session=True,
244+
)
245+
time.sleep(3) # Wait for etcd to start
246+
247+
if etcd_process.poll() is None:
248+
logger.info(
249+
f"etcd started, PID: {etcd_process.pid}. Logs at: {os.path.abspath(log_file_path)}"
250+
)
251+
else:
252+
error_msg = ""
253+
try:
254+
with open(log_file_path) as f:
255+
error_msg = f.read()
256+
except Exception as e:
257+
error_msg = f"Failed to read log file: {e}"
258+
259+
# Clean up data directory on failure
260+
try:
261+
shutil.rmtree(etcd_data_dir, ignore_errors=True)
262+
except Exception:
263+
pass
264+
265+
raise RuntimeError(
266+
f"etcd exited with error. Check {log_file_path} for detailed logs. Output:\n{error_msg}"
267+
)
268+
269+
# Wait a moment for etcd to be ready if we started it
270+
if etcd_process is not None:
271+
time.sleep(2)
272+
273+
# ========== Start datasystem worker (inlined from _start_dscli) ==========
274+
worker_address = conf.backend.Yuanrong.worker_address
275+
# Check if datasystem worker is already running (by checking worker port)
276+
parts = worker_address.split(":")
277+
if len(parts) != 2:
278+
raise ValueError(f"Invalid worker_address format: {worker_address}. Expected host:port")
279+
worker_host = parts[0]
280+
worker_port = int(parts[1])
281+
282+
try:
283+
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
284+
sock.settimeout(2)
285+
result = sock.connect_ex((worker_host, worker_port))
286+
sock.close()
287+
if result == 0:
288+
logger.info(f"Datasystem worker already listening on {worker_address}")
289+
else:
290+
# Start datasystem worker
291+
cmd = [
292+
"dscli",
293+
"start",
294+
"-w",
295+
f"--worker_address={worker_address}",
296+
f"--etcd_address={etcd_address}",
297+
]
298+
299+
log_file_path = "/tmp/dscli.log"
300+
with open(log_file_path, "w") as log_file:
301+
process = subprocess.Popen(
302+
cmd,
303+
stdout=log_file,
304+
stderr=subprocess.STDOUT,
305+
text=True,
306+
bufsize=1,
307+
universal_newlines=True,
308+
start_new_session=True,
309+
)
310+
# Wait for dscli to start and exit (it starts worker and exits)
311+
time.sleep(3)
312+
313+
if process.poll() is not None:
314+
# dscli exited, check if it succeeded by verifying port
315+
time.sleep(2) # Give worker time to start
316+
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
317+
sock.settimeout(2)
318+
result = sock.connect_ex((worker_host, worker_port))
319+
sock.close()
320+
if result != 0:
321+
error_msg = ""
322+
try:
323+
with open(log_file_path) as f:
324+
error_msg = f.read()
325+
except Exception as e:
326+
error_msg = f"Failed to read log file: {e}"
327+
raise RuntimeError(
328+
f"Failed to start datasystem worker. Check {log_file_path} for logs. "
329+
f"Output:\n{error_msg}"
330+
)
331+
else:
332+
logger.info(f"Datasystem worker started and listening on {worker_address}")
333+
else:
334+
# dscli is still running (unexpected), log warning
335+
logger.warning(
336+
f"dscli process still running (PID: {process.pid}). This may indicate an issue."
337+
)
338+
except Exception as e:
339+
raise RuntimeError(f"Failed to start datasystem worker: {e}") from e
340+
341+
# Store processes and data directory
342+
_TRANSFER_QUEUE_STORAGE["Yuanrong"] = {
343+
"etcd": etcd_process,
344+
"etcd_data_dir": etcd_data_dir,
345+
"worker_address": worker_address,
346+
"etcd_address": etcd_address,
347+
}
348+
logger.info("Yuanrong backend (etcd + datasystem) started successfully.")
349+
350+
except Exception as e:
351+
# Clean up on failure
352+
if etcd_process is not None and etcd_process.poll() is None:
353+
etcd_process.terminate()
354+
if etcd_data_dir is not None:
355+
try:
356+
shutil.rmtree(etcd_data_dir, ignore_errors=True)
357+
except Exception:
358+
pass
359+
raise RuntimeError(f"Failed to start Yuanrong backend: {e}") from e
170360
return conf
171361

172362

@@ -328,6 +518,42 @@ def close():
328518
logger.info("Successfully removed all existing keys in mooncake_master.")
329519
except Exception:
330520
pass
521+
elif key == "Yuanrong":
522+
# Stop etcd process and clean up data directory, stop datasystem worker via dscli
523+
if isinstance(value, dict):
524+
etcd_process = value.get("etcd")
525+
etcd_data_dir = value.get("etcd_data_dir")
526+
worker_address = value.get("worker_address")
527+
528+
# Stop etcd if running
529+
if etcd_process is not None and etcd_process.poll() is None:
530+
etcd_process.terminate()
531+
try:
532+
etcd_process.wait(timeout=5)
533+
except subprocess.TimeoutExpired:
534+
etcd_process.kill()
535+
536+
# Clean up etcd data directory
537+
if etcd_data_dir is not None and os.path.exists(etcd_data_dir):
538+
try:
539+
shutil.rmtree(etcd_data_dir, ignore_errors=True)
540+
logger.info(f"Cleaned up etcd data directory: {etcd_data_dir}")
541+
except Exception as e:
542+
logger.warning(f"Failed to clean up etcd data directory {etcd_data_dir}: {e}")
543+
544+
# Stop datasystem worker via dscli command
545+
if worker_address:
546+
try:
547+
subprocess.run(
548+
["dscli", "stop", "--worker_address", worker_address],
549+
timeout=5,
550+
capture_output=True,
551+
)
552+
logger.info(f"Stopped datasystem worker at {worker_address} via dscli stop")
553+
except Exception as e:
554+
logger.warning(f"Failed to stop datasystem worker via dscli: {e}")
555+
else:
556+
logger.warning(f"Unexpected Yuanrong storage value: {value}")
331557
else:
332558
logger.warning(f"close for _TRANSFER_QUEUE_STORAGE with key {key} is not supported for now.")
333559

0 commit comments

Comments
 (0)