Skip to content

V
vidai
Commit d5d30329 authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

Fix client disconnection handling in cluster master 

- Made assign_job_to_worker, _transfer_job_files, _transfer_file_via_websocket, enable_process, disable_process, update_process_weight, restart_client_workers, and restart_client_worker async methods
- Added proper exception handling for websocket send operations
- When websocket send fails due to broken connection, clients are now properly removed from available workers selection
- This ensures that disconnected clients are immediately removed from the worker pool and jobs are re-assigned to available workers
parent a2c308f1
Hide whitespace changes
Inline Side-by-side
Showing with 139 additions and 75 deletions
vidai/cluster_master.py
View file @ d5d30329
......@@ -444,7 +444,7 @@ class ClusterMaster:
return max(1, cuda_count + rocm_count)
def assign_job_to_worker(self, worker_key: str, job_data: dict) -> Optional[str]:
async def assign_job_to_worker(self, worker_key: str, job_data: dict) -> Optional[str]:
"""Assign a job to a worker and handle file/model transfer."""
from .models import estimate_model_vram_requirements
import uuid
......@@ -491,11 +491,15 @@ class ClusterMaster:
'job_data': job_data
}
# Send via websocket (async)
asyncio.create_task(self.client_websockets[client_id].send(json.dumps(job_message)))
# Send via websocket synchronously to catch connection errors
await self.client_websockets[client_id].send(json.dumps(job_message))
print(f"Job {job_id} assigned to worker {worker_key} on client {client_id}")
except Exception as e:
print(f"Failed to send job {job_id} to worker {worker_key}: {e}")
# Connection is broken, remove the client
if client_id in self.clients:
print(f"Removing disconnected client {client_id}")
self._remove_client(client_id)
# Clean up the failed assignment
self.worker_jobs[worker_key].remove(job_id)
self.worker_vram_usage[worker_key] -= vram_required
......@@ -511,7 +515,7 @@ class ClusterMaster:
return job_id
def _transfer_job_files(self, client_id: str, job_data: dict, job_id: str) -> None:
async def _transfer_job_files(self, client_id: str, job_data: dict, job_id: str) -> None:
"""Transfer job files to client."""
media_path = job_data.get('local_path')
if not media_path:
......@@ -529,21 +533,29 @@ class ClusterMaster:
shutil.copy2(media_path, shared_file_path)
# Notify client of shared file location
asyncio.create_task(self.client_websockets[client_id].send(json.dumps({
'type': 'job_file_shared',
'job_id': job_id,
'shared_file_path': shared_file_path,
'original_path': media_path
})))
try:
await self.client_websockets[client_id].send(json.dumps({
'type': 'job_file_shared',
'job_id': job_id,
'shared_file_path': shared_file_path,
'original_path': media_path
}))
except Exception as e:
print(f"Failed to send shared file notification for job {job_id}: {e}")
# Connection is broken, remove the client
if client_id in self.clients:
print(f"Removing disconnected client {client_id}")
self._remove_client(client_id)
return
except Exception as e:
print(f"Failed to use shared directory for job {job_id}: {e}")
# Fall back to websocket transfer
self._transfer_file_via_websocket(client_id, media_path, job_id)
await self._transfer_file_via_websocket(client_id, media_path, job_id)
else:
# Use websocket transfer
self._transfer_file_via_websocket(client_id, media_path, job_id)
await self._transfer_file_via_websocket(client_id, media_path, job_id)
def _transfer_file_via_websocket(self, client_id: str, file_path: str, job_id: str) -> None:
async def _transfer_file_via_websocket(self, client_id: str, file_path: str, job_id: str) -> None:
"""Transfer file via websocket."""
try:
with open(file_path, 'rb') as f:
......@@ -554,29 +566,37 @@ class ClusterMaster:
total_size = len(file_data)
# Send start message
asyncio.create_task(self.client_websockets[client_id].send(json.dumps({
'type': 'job_file_transfer_start',
'job_id': job_id,
'file_path': file_path,
'total_size': total_size
})))
# Send chunks
for i in range(0, total_size, chunk_size):
chunk = file_data[i:i + chunk_size]
asyncio.create_task(self.client_websockets[client_id].send(json.dumps({
'type': 'job_file_chunk',
try:
await self.client_websockets[client_id].send(json.dumps({
'type': 'job_file_transfer_start',
'job_id': job_id,
'offset': i,
'data': chunk.hex()
})))
'file_path': file_path,
'total_size': total_size
}))
# Send completion
asyncio.create_task(self.client_websockets[client_id].send(json.dumps({
'type': 'job_file_transfer_complete',
'job_id': job_id,
'file_path': file_path
})))
# Send chunks
for i in range(0, total_size, chunk_size):
chunk = file_data[i:i + chunk_size]
await self.client_websockets[client_id].send(json.dumps({
'type': 'job_file_chunk',
'job_id': job_id,
'offset': i,
'data': chunk.hex()
}))
# Send completion
await self.client_websockets[client_id].send(json.dumps({
'type': 'job_file_transfer_complete',
'job_id': job_id,
'file_path': file_path
}))
except Exception as e:
print(f"Failed to transfer file {file_path} for job {job_id}: {e}")
# Connection is broken, remove the client
if client_id in self.clients:
print(f"Removing disconnected client {client_id}")
self._remove_client(client_id)
raise # Re-raise to indicate failure
except Exception as e:
print(f"Failed to transfer file {file_path} for job {job_id}: {e}")
......@@ -702,12 +722,20 @@ class ClusterMaster:
f.write(model_data)
# Send file path to client
await self.client_websockets[client_id].send(json.dumps({
'type': 'model_shared_file',
'model_path': model_path,
'shared_file_path': shared_file_path,
'total_size': len(model_data)
}))
try:
await self.client_websockets[client_id].send(json.dumps({
'type': 'model_shared_file',
'model_path': model_path,
'shared_file_path': shared_file_path,
'total_size': len(model_data)
}))
except Exception as e:
print(f"Failed to send model shared file notification to client {client_id}: {e}")
# Connection is broken, remove the client
if client_id in self.clients:
print(f"Removing disconnected client {client_id}")
self._remove_client(client_id)
return False
print(f"Model {model_path} placed in shared directory for client {client_id}: {shared_file_path}")
return True
......@@ -770,19 +798,27 @@ class ClusterMaster:
print("huggingface_hub not available for model download")
return None
def enable_process(self, process_key: str) -> bool:
async def enable_process(self, process_key: str) -> bool:
"""Enable a specific process."""
if process_key in self.processes:
self.processes[process_key]['status'] = 'active'
# Send command to client
client_id = self.processes[process_key]['client_id']
if client_id in self.client_websockets:
asyncio.create_task(self.client_websockets[client_id].send(
json.dumps({
'type': 'enable_process',
'process_name': self.processes[process_key]['name']
})
))
try:
await self.client_websockets[client_id].send(
json.dumps({
'type': 'enable_process',
'process_name': self.processes[process_key]['name']
})
)
except Exception as e:
print(f"Failed to send enable_process command to client {client_id}: {e}")
# Connection is broken, remove the client
if client_id in self.clients:
print(f"Removing disconnected client {client_id}")
self._remove_client(client_id)
return False
return True
return False
......@@ -951,23 +987,31 @@ class ClusterMaster:
return None
def disable_process(self, process_key: str) -> bool:
async def disable_process(self, process_key: str) -> bool:
"""Disable a specific process."""
if process_key in self.processes:
self.processes[process_key]['status'] = 'disabled'
# Send command to client
client_id = self.processes[process_key]['client_id']
if client_id in self.client_websockets:
asyncio.create_task(self.client_websockets[client_id].send(
json.dumps({
'type': 'disable_process',
'process_name': self.processes[process_key]['name']
})
))
try:
await self.client_websockets[client_id].send(
json.dumps({
'type': 'disable_process',
'process_name': self.processes[process_key]['name']
})
)
except Exception as e:
print(f"Failed to send disable_process command to client {client_id}: {e}")
# Connection is broken, remove the client
if client_id in self.clients:
print(f"Removing disconnected client {client_id}")
self._remove_client(client_id)
return False
return True
return False
def update_process_weight(self, process_key: str, weight: int) -> bool:
async def update_process_weight(self, process_key: str, weight: int) -> bool:
"""Update process weight for load balancing."""
if process_key in self.processes:
self.processes[process_key]['weight'] = weight
......@@ -986,17 +1030,25 @@ class ClusterMaster:
# Send command to client
client_id = self.processes[process_key]['client_id']
if client_id in self.client_websockets:
asyncio.create_task(self.client_websockets[client_id].send(
json.dumps({
'type': 'update_weight',
'process_name': self.processes[process_key]['name'],
'weight': weight
})
))
try:
await self.client_websockets[client_id].send(
json.dumps({
'type': 'update_weight',
'process_name': self.processes[process_key]['name'],
'weight': weight
})
)
except Exception as e:
print(f"Failed to send update_weight command to client {client_id}: {e}")
# Connection is broken, remove the client
if client_id in self.clients:
print(f"Removing disconnected client {client_id}")
self._remove_client(client_id)
return False
return True
return False
def restart_client_workers(self, client_id: str, backend: str) -> bool:
async def restart_client_workers(self, client_id: str, backend: str) -> bool:
"""Restart all workers on a client with a different backend."""
if client_id not in self.client_websockets:
return False
......@@ -1006,15 +1058,23 @@ class ClusterMaster:
return False
# Send restart command to client
asyncio.create_task(self.client_websockets[client_id].send(
json.dumps({
'type': 'restart_workers',
'backend': backend
})
))
return True
def restart_client_worker(self, client_id: str, worker_name: str, backend: str) -> bool:
try:
await self.client_websockets[client_id].send(
json.dumps({
'type': 'restart_workers',
'backend': backend
})
)
return True
except Exception as e:
print(f"Failed to send restart_workers command to client {client_id}: {e}")
# Connection is broken, remove the client
if client_id in self.clients:
print(f"Removing disconnected client {client_id}")
self._remove_client(client_id)
return False
async def restart_client_worker(self, client_id: str, worker_name: str, backend: str) -> bool:
"""Restart a specific worker on a client with a different backend."""
if client_id not in self.client_websockets:
return False
......@@ -1025,16 +1085,20 @@ class ClusterMaster:
# Send restart command for specific worker to client
try:
asyncio.create_task(self.client_websockets[client_id].send(
await self.client_websockets[client_id].send(
json.dumps({
'type': 'restart_worker',
'worker_name': worker_name,
'backend': backend
})
))
)
return True
except Exception as e:
print(f"Failed to send restart command for worker {worker_name} to client {client_id}: {e}")
# Connection is broken, remove the client
if client_id in self.clients:
print(f"Removing disconnected client {client_id}")
self._remove_client(client_id)
return False
async def _management_loop(self) -> None:
......@@ -1074,7 +1138,7 @@ class ClusterMaster:
print(f"Job {job['id']} waiting for available workers")
worker_key = self.select_worker_for_job(process_type, job['data'].get('model_path', 'Qwen/Qwen2.5-VL-7B-Instruct'), job['data'])
if worker_key:
job_id = self.assign_job_to_worker(worker_key, job['data'])
job_id = await self.assign_job_to_worker(worker_key, job['data'])
if job_id:
from .database import update_queue_status
update_queue_status(job['id'], 'processing', {'job_id': job_id, 'status': 'Assigned to worker'}, job_id=job_id)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment