Skip to content

Commit 9c0cb83

Browse files
dpj135dpj135
committed
Reduced logics
Signed-off-by: dpj135 <958208521@qq.com>
1 parent a1c5c9b commit 9c0cb83

2 files changed

Lines changed: 90 additions & 161 deletions

File tree

‎transfer_queue/config.yaml‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,12 +48,11 @@ backend:
4848

4949
# For Yuanrong:
5050
Yuanrong:
51-
# Whether to let TQ automatically start etcd and datasystem services (default false)
52-
auto_init: false
51+
# Whether to let TQ automatically start etcd and datasystem services
52+
auto_init: True
5353
# etcd service address (used to start etcd when auto_init=true)
5454
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"
55+
# datasystem worker host and port (used to start dscli when auto_init=true)
5756
# YuanrongStorageClient connection parameters (required)
5857
host: "127.0.0.1"
5958
port: 31501

‎transfer_queue/interface.py‎

Lines changed: 87 additions & 157 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
import math
1818
import os
1919
import shutil
20-
import socket
2120
import subprocess
2221
import tempfile
2322
import time
@@ -174,169 +173,100 @@ def _maybe_create_transferqueue_storage(conf: DictConfig) -> DictConfig:
174173
if conf.backend.Yuanrong.auto_init:
175174
etcd_process = None
176175
etcd_data_dir = None
177-
176+
worker_address = None
177+
if not shutil.which("etcd"):
178+
raise RuntimeError(
179+
"etcd executable not found in PATH. Please install etcd and make sure it's in the PATH."
180+
)
181+
if not shutil.which("dscli"):
182+
raise RuntimeError(
183+
"dscli executable not found in PATH. Please run `pip install openyuanrong-datasystem`."
184+
)
178185
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}"
186+
# ========== Start etcd ==========
187+
etcd_address = "127.0.0.1:2379"
188+
try:
189+
etcd_address = conf.backend.Yuanrong.etcd_address
190+
except Exception:
191+
pass
258192

259-
# Clean up data directory on failure
260-
try:
261-
shutil.rmtree(etcd_data_dir, ignore_errors=True)
262-
except Exception:
263-
pass
193+
# Assume host:port format
194+
parts = etcd_address.split(":")
195+
if len(parts) != 2:
196+
raise ValueError(f"Invalid etcd_address format: {etcd_address}. Expected host:port")
197+
host = parts[0]
198+
port = int(parts[1])
199+
200+
# Create temporary data directory
201+
etcd_data_dir = tempfile.mkdtemp(prefix="tq_etcd_")
202+
logger.info(f"Starting etcd with data directory: {etcd_data_dir}")
203+
204+
cmd = [
205+
"etcd",
206+
f"--data-dir={etcd_data_dir}",
207+
f"--listen-client-urls=http://{host}:{port}",
208+
f"--advertise-client-urls=http://{host}:{port}",
209+
]
210+
211+
etcd_process = subprocess.Popen(
212+
cmd,
213+
stdout=subprocess.DEVNULL,
214+
stderr=subprocess.DEVNULL,
215+
text=True,
216+
bufsize=1,
217+
universal_newlines=True,
218+
start_new_session=True,
219+
)
220+
time.sleep(3) # Wait for etcd to start
221+
# TODO: check if etcd is healthy
222+
etcd_is_healthy = etcd_process.poll() is None
223+
# check
224+
#
225+
#
226+
#
227+
#
228+
if not etcd_is_healthy:
229+
# etcd exited immediately, indicate failure
230+
# Clean up data directory on failure
231+
try:
232+
shutil.rmtree(etcd_data_dir, ignore_errors=True)
233+
except Exception:
234+
pass
235+
raise RuntimeError(f"etcd exited immediately with return code {etcd_process.returncode}")
264236

265-
raise RuntimeError(
266-
f"etcd exited with error. Check {log_file_path} for detailed logs. Output:\n{error_msg}"
267-
)
237+
# Wait a moment for etcd to be ready
238+
time.sleep(2)
268239

269-
# Wait a moment for etcd to be ready if we started it
270-
if etcd_process is not None:
271-
time.sleep(2)
240+
# ========== Start datasystem worker ==========
241+
# Assume host:port format
242+
worker_host = conf.backend.Yuanrong.host
243+
worker_port = conf.backend.Yuanrong.port
244+
worker_address = worker_host + ":" + str(worker_port)
272245

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])
246+
cmd = [
247+
"dscli",
248+
"start",
249+
"-w",
250+
f"--worker_address={worker_address}",
251+
f"--etcd_address={etcd_address}",
252+
]
281253

282254
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
255+
ds_result = subprocess.run(
256+
cmd,
257+
stdout=subprocess.PIPE,
258+
stderr=subprocess.STDOUT,
259+
text=True,
260+
timeout=90,
261+
)
262+
except subprocess.TimeoutExpired as err:
263+
raise RuntimeError("dscli start timed out") from err
264+
# Wait for dscli to start and exit (it starts worker and exits)
265+
if ds_result.returncode == 0 and "[ OK ]" in ds_result.stdout:
266+
logger.info(f"dscli started Yuanrong datasystem worker at {worker_address} successfully.")
267+
268+
else:
269+
raise RuntimeError(f"Failed to start datasystem worker at {worker_address}.")
340270

341271
# Store processes and data directory
342272
_TRANSFER_QUEUE_STORAGE["Yuanrong"] = {

0 commit comments

Comments
 (0)