Match with image runner
This commit is contained in:
Ryuichi Leo Takashige
committed by
Evan
Evan
parent
f77a672126
commit
69628383c5
@@ -284,22 +284,16 @@ class MpReceiver[T]:
|
||||
|
||||
|
||||
class NonBlockingGenerator[T](Generator[T | None, None, None]):
|
||||
def __init__(self, source: MpReceiver[T] | Generator[T | None, None, None]) -> None:
|
||||
self._inner: MpReceiver[T] | Generator[T | None, None, None] = source
|
||||
def __init__(self, source: MpReceiver[T]) -> None:
|
||||
self._inner: MpReceiver[T] = source
|
||||
|
||||
def send(self, value: None, /) -> T | None:
|
||||
if isinstance(self._inner, Generator):
|
||||
try:
|
||||
return next(self._inner)
|
||||
except (StopIteration, ClosedResourceError):
|
||||
raise StopIteration from None
|
||||
else:
|
||||
try:
|
||||
return self._inner.receive_nowait()
|
||||
except WouldBlock:
|
||||
return None
|
||||
except (EndOfStream, ClosedResourceError):
|
||||
raise StopIteration from None
|
||||
try:
|
||||
return self._inner.receive_nowait()
|
||||
except WouldBlock:
|
||||
return None
|
||||
except (EndOfStream, ClosedResourceError):
|
||||
raise StopIteration from None
|
||||
|
||||
def throw(
|
||||
self,
|
||||
|
||||
@@ -32,9 +32,12 @@ def entrypoint(
|
||||
# Import main after setting global logger - this lets us just import logger from this module
|
||||
try:
|
||||
if bound_instance.is_image_model:
|
||||
from exo.worker.runner.image_models.runner import main
|
||||
from exo.worker.runner.image_models.runner import Runner as ImageRunner
|
||||
|
||||
main(bound_instance, event_sender, task_receiver, cancel_receiver)
|
||||
runner = ImageRunner(
|
||||
bound_instance, event_sender, task_receiver, cancel_receiver
|
||||
)
|
||||
runner.main()
|
||||
else:
|
||||
from exo.worker.runner.llm_inference.runner import Runner
|
||||
|
||||
|
||||
@@ -182,272 +182,266 @@ def _send_image_chunk(
|
||||
)
|
||||
|
||||
|
||||
def main(
|
||||
bound_instance: BoundInstance,
|
||||
event_sender: MpSender[Event],
|
||||
task_receiver: MpReceiver[Task],
|
||||
cancel_receiver: MpReceiver[TaskId],
|
||||
):
|
||||
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
|
||||
resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
|
||||
class Runner:
|
||||
def __init__(
|
||||
self,
|
||||
bound_instance: BoundInstance,
|
||||
event_sender: MpSender[Event],
|
||||
task_receiver: MpReceiver[Task],
|
||||
cancel_receiver: MpReceiver[TaskId],
|
||||
):
|
||||
self.event_sender = event_sender
|
||||
self.task_receiver = task_receiver
|
||||
self.cancel_receiver = cancel_receiver
|
||||
self.bound_instance = bound_instance
|
||||
|
||||
instance, runner_id, shard_metadata = (
|
||||
bound_instance.instance,
|
||||
bound_instance.bound_runner_id,
|
||||
bound_instance.bound_shard,
|
||||
)
|
||||
device_rank = shard_metadata.device_rank
|
||||
logger.info("hello from the runner")
|
||||
if getattr(shard_metadata, "immediate_exception", False):
|
||||
raise Exception("Fake exception - runner failed to spin up.")
|
||||
if timeout := getattr(shard_metadata, "should_timeout", 0):
|
||||
time.sleep(timeout)
|
||||
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
|
||||
resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
|
||||
|
||||
setup_start_time = time.time()
|
||||
cancelled_tasks = set[TaskId]()
|
||||
self.instance, self.runner_id, self.shard_metadata = (
|
||||
bound_instance.instance,
|
||||
bound_instance.bound_runner_id,
|
||||
bound_instance.bound_shard,
|
||||
)
|
||||
self.device_rank = self.shard_metadata.device_rank
|
||||
|
||||
image_model: DistributedImageModel | None = None
|
||||
group = None
|
||||
logger.info("hello from the runner")
|
||||
if getattr(self.shard_metadata, "immediate_exception", False):
|
||||
raise Exception("Fake exception - runner failed to spin up.")
|
||||
if timeout := getattr(self.shard_metadata, "should_timeout", 0):
|
||||
time.sleep(timeout)
|
||||
|
||||
current_status: RunnerStatus = RunnerIdle()
|
||||
logger.info("runner created")
|
||||
event_sender.send(
|
||||
RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
|
||||
)
|
||||
seen = set[TaskId]()
|
||||
with task_receiver as tasks:
|
||||
for task in tasks:
|
||||
if task.task_id in seen:
|
||||
logger.warning("repeat task - potential error")
|
||||
seen.add(task.task_id)
|
||||
cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
|
||||
event_sender.send(
|
||||
TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
|
||||
self.setup_start_time = time.time()
|
||||
self.cancelled_tasks = set[TaskId]()
|
||||
|
||||
self.image_model: DistributedImageModel | None = None
|
||||
self.group = None
|
||||
|
||||
self.current_status: RunnerStatus = RunnerIdle()
|
||||
logger.info("runner created")
|
||||
self.update_status(RunnerIdle())
|
||||
self.seen = set[TaskId]()
|
||||
|
||||
def update_status(self, status: RunnerStatus):
|
||||
self.current_status = status
|
||||
self.event_sender.send(
|
||||
RunnerStatusUpdated(
|
||||
runner_id=self.runner_id, runner_status=self.current_status
|
||||
)
|
||||
match task:
|
||||
case ConnectToGroup() if isinstance(
|
||||
current_status, (RunnerIdle, RunnerFailed)
|
||||
):
|
||||
logger.info("runner connecting")
|
||||
current_status = RunnerConnecting()
|
||||
event_sender.send(
|
||||
RunnerStatusUpdated(
|
||||
runner_id=runner_id, runner_status=current_status
|
||||
)
|
||||
)
|
||||
event_sender.send(TaskAcknowledged(task_id=task.task_id))
|
||||
group = initialize_mlx(bound_instance)
|
||||
)
|
||||
|
||||
logger.info("runner connected")
|
||||
current_status = RunnerConnected()
|
||||
def send_task_status(self, task: Task, status: TaskStatus):
|
||||
self.event_sender.send(
|
||||
TaskStatusUpdated(task_id=task.task_id, task_status=status)
|
||||
)
|
||||
|
||||
# we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
|
||||
case LoadModel() if (
|
||||
isinstance(current_status, RunnerConnected) and group is not None
|
||||
) or (isinstance(current_status, RunnerIdle) and group is None):
|
||||
current_status = RunnerLoading()
|
||||
logger.info("runner loading")
|
||||
event_sender.send(
|
||||
RunnerStatusUpdated(
|
||||
runner_id=runner_id, runner_status=current_status
|
||||
)
|
||||
)
|
||||
event_sender.send(TaskAcknowledged(task_id=task.task_id))
|
||||
def acknowledge_task(self, task: Task):
|
||||
self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
|
||||
|
||||
assert (
|
||||
ModelTask.TextToImage in shard_metadata.model_card.tasks
|
||||
or ModelTask.ImageToImage in shard_metadata.model_card.tasks
|
||||
), f"Incorrect model task(s): {shard_metadata.model_card.tasks}"
|
||||
|
||||
image_model = initialize_image_model(bound_instance)
|
||||
current_status = RunnerLoaded()
|
||||
logger.info("runner loaded")
|
||||
|
||||
case StartWarmup() if isinstance(current_status, RunnerLoaded):
|
||||
current_status = RunnerWarmingUp()
|
||||
logger.info("runner warming up")
|
||||
event_sender.send(
|
||||
RunnerStatusUpdated(
|
||||
runner_id=runner_id, runner_status=current_status
|
||||
)
|
||||
)
|
||||
event_sender.send(TaskAcknowledged(task_id=task.task_id))
|
||||
|
||||
logger.info(f"warming up inference for instance: {instance}")
|
||||
|
||||
assert image_model
|
||||
image = warmup_image_generator(model=image_model)
|
||||
if image is not None:
|
||||
logger.info(f"warmed up by generating {image.size} image")
|
||||
else:
|
||||
logger.info("warmup completed (non-primary node)")
|
||||
|
||||
logger.info(
|
||||
f"runner initialized in {time.time() - setup_start_time} seconds"
|
||||
)
|
||||
|
||||
current_status = RunnerReady()
|
||||
logger.info("runner ready")
|
||||
|
||||
case ImageGeneration(
|
||||
task_params=task_params, command_id=command_id
|
||||
) if isinstance(current_status, RunnerReady):
|
||||
assert image_model
|
||||
logger.info(f"received image generation request: {str(task)[:500]}")
|
||||
current_status = RunnerRunning()
|
||||
logger.info("runner running")
|
||||
event_sender.send(
|
||||
RunnerStatusUpdated(
|
||||
runner_id=runner_id, runner_status=current_status
|
||||
)
|
||||
)
|
||||
event_sender.send(TaskAcknowledged(task_id=task.task_id))
|
||||
|
||||
try:
|
||||
image_index = 0
|
||||
for response in generate_image(
|
||||
model=image_model, task=task_params
|
||||
):
|
||||
is_primary_output = _is_primary_output_node(shard_metadata)
|
||||
|
||||
if is_primary_output:
|
||||
match response:
|
||||
case PartialImageResponse():
|
||||
logger.info(
|
||||
f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
|
||||
)
|
||||
_process_image_response(
|
||||
response,
|
||||
command_id,
|
||||
shard_metadata,
|
||||
event_sender,
|
||||
image_index,
|
||||
)
|
||||
case ImageGenerationResponse():
|
||||
logger.info("sending final ImageChunk")
|
||||
_process_image_response(
|
||||
response,
|
||||
command_id,
|
||||
shard_metadata,
|
||||
event_sender,
|
||||
image_index,
|
||||
)
|
||||
image_index += 1
|
||||
# can we make this more explicit?
|
||||
except Exception as e:
|
||||
if _is_primary_output_node(shard_metadata):
|
||||
event_sender.send(
|
||||
ChunkGenerated(
|
||||
command_id=command_id,
|
||||
chunk=ErrorChunk(
|
||||
model=shard_metadata.model_card.model_id,
|
||||
finish_reason="error",
|
||||
error_message=str(e),
|
||||
),
|
||||
)
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
_send_traces_if_enabled(event_sender, task.task_id, device_rank)
|
||||
|
||||
current_status = RunnerReady()
|
||||
logger.info("runner ready")
|
||||
|
||||
case ImageEdits(task_params=task_params, command_id=command_id) if (
|
||||
isinstance(current_status, RunnerReady)
|
||||
):
|
||||
assert image_model
|
||||
logger.info(f"received image edits request: {str(task)[:500]}")
|
||||
current_status = RunnerRunning()
|
||||
logger.info("runner running")
|
||||
event_sender.send(
|
||||
RunnerStatusUpdated(
|
||||
runner_id=runner_id, runner_status=current_status
|
||||
)
|
||||
)
|
||||
event_sender.send(TaskAcknowledged(task_id=task.task_id))
|
||||
|
||||
try:
|
||||
image_index = 0
|
||||
for response in generate_image(
|
||||
model=image_model, task=task_params
|
||||
):
|
||||
if _is_primary_output_node(shard_metadata):
|
||||
match response:
|
||||
case PartialImageResponse():
|
||||
logger.info(
|
||||
f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
|
||||
)
|
||||
_process_image_response(
|
||||
response,
|
||||
command_id,
|
||||
shard_metadata,
|
||||
event_sender,
|
||||
image_index,
|
||||
)
|
||||
case ImageGenerationResponse():
|
||||
logger.info("sending final ImageChunk")
|
||||
_process_image_response(
|
||||
response,
|
||||
command_id,
|
||||
shard_metadata,
|
||||
event_sender,
|
||||
image_index,
|
||||
)
|
||||
image_index += 1
|
||||
except Exception as e:
|
||||
if _is_primary_output_node(shard_metadata):
|
||||
event_sender.send(
|
||||
ChunkGenerated(
|
||||
command_id=command_id,
|
||||
chunk=ErrorChunk(
|
||||
model=shard_metadata.model_card.model_id,
|
||||
finish_reason="error",
|
||||
error_message=str(e),
|
||||
),
|
||||
)
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
_send_traces_if_enabled(event_sender, task.task_id, device_rank)
|
||||
|
||||
current_status = RunnerReady()
|
||||
logger.info("runner ready")
|
||||
|
||||
case Shutdown():
|
||||
current_status = RunnerShuttingDown()
|
||||
logger.info("runner shutting down")
|
||||
if not TYPE_CHECKING:
|
||||
del image_model, group
|
||||
mx.clear_cache()
|
||||
import gc
|
||||
|
||||
gc.collect()
|
||||
|
||||
event_sender.send(
|
||||
RunnerStatusUpdated(
|
||||
runner_id=runner_id, runner_status=current_status
|
||||
)
|
||||
)
|
||||
event_sender.send(TaskAcknowledged(task_id=task.task_id))
|
||||
|
||||
current_status = RunnerShutdown()
|
||||
case _:
|
||||
raise ValueError(
|
||||
f"Received {task.__class__.__name__} outside of state machine in {current_status=}"
|
||||
)
|
||||
was_cancelled = (task.task_id in cancelled_tasks) or (
|
||||
TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks
|
||||
)
|
||||
if not was_cancelled:
|
||||
event_sender.send(
|
||||
TaskStatusUpdated(
|
||||
task_id=task.task_id, task_status=TaskStatus.Complete
|
||||
)
|
||||
def main(self):
|
||||
with self.task_receiver as tasks:
|
||||
for task in tasks:
|
||||
if task.task_id in self.seen:
|
||||
logger.warning("repeat task - potential error")
|
||||
self.seen.add(task.task_id)
|
||||
self.cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
|
||||
self.send_task_status(task, TaskStatus.Running)
|
||||
self.handle_task(task)
|
||||
was_cancelled = (task.task_id in self.cancelled_tasks) or (
|
||||
TaskId("CANCEL_CURRENT_TASK") in self.cancelled_tasks
|
||||
)
|
||||
event_sender.send(
|
||||
RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
|
||||
)
|
||||
if not was_cancelled:
|
||||
self.send_task_status(task, TaskStatus.Complete)
|
||||
self.update_status(self.current_status)
|
||||
|
||||
if isinstance(current_status, RunnerShutdown):
|
||||
break
|
||||
if isinstance(self.current_status, RunnerShutdown):
|
||||
break
|
||||
|
||||
def handle_task(self, task: Task):
|
||||
match task:
|
||||
case ConnectToGroup() if isinstance(
|
||||
self.current_status, (RunnerIdle, RunnerFailed)
|
||||
):
|
||||
logger.info("runner connecting")
|
||||
self.update_status(RunnerConnecting())
|
||||
self.acknowledge_task(task)
|
||||
self.group = initialize_mlx(self.bound_instance)
|
||||
|
||||
logger.info("runner connected")
|
||||
self.current_status = RunnerConnected()
|
||||
|
||||
# we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
|
||||
case LoadModel() if (
|
||||
isinstance(self.current_status, RunnerConnected)
|
||||
and self.group is not None
|
||||
) or (isinstance(self.current_status, RunnerIdle) and self.group is None):
|
||||
logger.info("runner loading")
|
||||
self.update_status(RunnerLoading())
|
||||
self.acknowledge_task(task)
|
||||
|
||||
assert (
|
||||
ModelTask.TextToImage in self.shard_metadata.model_card.tasks
|
||||
or ModelTask.ImageToImage in self.shard_metadata.model_card.tasks
|
||||
), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}"
|
||||
|
||||
self.image_model = initialize_image_model(self.bound_instance)
|
||||
self.current_status = RunnerLoaded()
|
||||
logger.info("runner loaded")
|
||||
|
||||
case StartWarmup() if isinstance(self.current_status, RunnerLoaded):
|
||||
logger.info("runner warming up")
|
||||
self.update_status(RunnerWarmingUp())
|
||||
self.acknowledge_task(task)
|
||||
|
||||
logger.info(f"warming up inference for instance: {self.instance}")
|
||||
|
||||
assert self.image_model
|
||||
image = warmup_image_generator(model=self.image_model)
|
||||
if image is not None:
|
||||
logger.info(f"warmed up by generating {image.size} image")
|
||||
else:
|
||||
logger.info("warmup completed (non-primary node)")
|
||||
|
||||
logger.info(
|
||||
f"runner initialized in {time.time() - self.setup_start_time} seconds"
|
||||
)
|
||||
|
||||
self.current_status = RunnerReady()
|
||||
logger.info("runner ready")
|
||||
|
||||
case ImageGeneration(task_params=task_params, command_id=command_id) if (
|
||||
isinstance(self.current_status, RunnerReady)
|
||||
):
|
||||
assert self.image_model
|
||||
logger.info(f"received image generation request: {str(task)[:500]}")
|
||||
logger.info("runner running")
|
||||
self.update_status(RunnerRunning())
|
||||
self.acknowledge_task(task)
|
||||
|
||||
try:
|
||||
image_index = 0
|
||||
for response in generate_image(
|
||||
model=self.image_model, task=task_params
|
||||
):
|
||||
is_primary_output = _is_primary_output_node(self.shard_metadata)
|
||||
|
||||
if is_primary_output:
|
||||
match response:
|
||||
case PartialImageResponse():
|
||||
logger.info(
|
||||
f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
|
||||
)
|
||||
_process_image_response(
|
||||
response,
|
||||
command_id,
|
||||
self.shard_metadata,
|
||||
self.event_sender,
|
||||
image_index,
|
||||
)
|
||||
case ImageGenerationResponse():
|
||||
logger.info("sending final ImageChunk")
|
||||
_process_image_response(
|
||||
response,
|
||||
command_id,
|
||||
self.shard_metadata,
|
||||
self.event_sender,
|
||||
image_index,
|
||||
)
|
||||
image_index += 1
|
||||
# can we make this more explicit?
|
||||
except Exception as e:
|
||||
if _is_primary_output_node(self.shard_metadata):
|
||||
self.event_sender.send(
|
||||
ChunkGenerated(
|
||||
command_id=command_id,
|
||||
chunk=ErrorChunk(
|
||||
model=self.shard_metadata.model_card.model_id,
|
||||
finish_reason="error",
|
||||
error_message=str(e),
|
||||
),
|
||||
)
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
_send_traces_if_enabled(
|
||||
self.event_sender, task.task_id, self.device_rank
|
||||
)
|
||||
|
||||
self.current_status = RunnerReady()
|
||||
logger.info("runner ready")
|
||||
|
||||
case ImageEdits(task_params=task_params, command_id=command_id) if (
|
||||
isinstance(self.current_status, RunnerReady)
|
||||
):
|
||||
assert self.image_model
|
||||
logger.info(f"received image edits request: {str(task)[:500]}")
|
||||
logger.info("runner running")
|
||||
self.update_status(RunnerRunning())
|
||||
self.acknowledge_task(task)
|
||||
|
||||
try:
|
||||
image_index = 0
|
||||
for response in generate_image(
|
||||
model=self.image_model, task=task_params
|
||||
):
|
||||
if _is_primary_output_node(self.shard_metadata):
|
||||
match response:
|
||||
case PartialImageResponse():
|
||||
logger.info(
|
||||
f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
|
||||
)
|
||||
_process_image_response(
|
||||
response,
|
||||
command_id,
|
||||
self.shard_metadata,
|
||||
self.event_sender,
|
||||
image_index,
|
||||
)
|
||||
case ImageGenerationResponse():
|
||||
logger.info("sending final ImageChunk")
|
||||
_process_image_response(
|
||||
response,
|
||||
command_id,
|
||||
self.shard_metadata,
|
||||
self.event_sender,
|
||||
image_index,
|
||||
)
|
||||
image_index += 1
|
||||
except Exception as e:
|
||||
if _is_primary_output_node(self.shard_metadata):
|
||||
self.event_sender.send(
|
||||
ChunkGenerated(
|
||||
command_id=command_id,
|
||||
chunk=ErrorChunk(
|
||||
model=self.shard_metadata.model_card.model_id,
|
||||
finish_reason="error",
|
||||
error_message=str(e),
|
||||
),
|
||||
)
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
_send_traces_if_enabled(
|
||||
self.event_sender, task.task_id, self.device_rank
|
||||
)
|
||||
|
||||
self.current_status = RunnerReady()
|
||||
logger.info("runner ready")
|
||||
|
||||
case Shutdown():
|
||||
logger.info("runner shutting down")
|
||||
if not TYPE_CHECKING:
|
||||
del self.image_model, self.group
|
||||
mx.clear_cache()
|
||||
import gc
|
||||
|
||||
gc.collect()
|
||||
|
||||
self.update_status(RunnerShuttingDown())
|
||||
self.acknowledge_task(task)
|
||||
|
||||
self.current_status = RunnerShutdown()
|
||||
case _:
|
||||
raise ValueError(
|
||||
f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
|
||||
)
|
||||
|
||||
@@ -2,6 +2,7 @@ from abc import ABC, abstractmethod
|
||||
from collections import deque
|
||||
from collections.abc import Generator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import cast
|
||||
|
||||
import mlx.core as mx
|
||||
from mlx_lm.tokenizer_utils import TokenizerWrapper
|
||||
@@ -74,6 +75,9 @@ class SequentialGenerator(InferenceGenerator):
|
||||
check_for_cancel_every: int
|
||||
|
||||
_cancelled_tasks: set[TaskId] = field(default_factory=set, init=False)
|
||||
_maybe_queue: list[tuple[TextGeneration, MpSender[GenerationResponse]]] = field(
|
||||
default_factory=list, init=False
|
||||
)
|
||||
_queue: deque[tuple[TextGeneration, MpSender[GenerationResponse]]] = field(
|
||||
default_factory=deque, init=False
|
||||
)
|
||||
@@ -95,9 +99,19 @@ class SequentialGenerator(InferenceGenerator):
|
||||
sender: MpSender[GenerationResponse],
|
||||
) -> None:
|
||||
self._cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
|
||||
self._queue.append((task, sender))
|
||||
if self._active is None:
|
||||
self._start_next()
|
||||
self._maybe_queue.append((task, sender))
|
||||
|
||||
def agree_on_tasks(self) -> None:
|
||||
"""Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
|
||||
agreed, different = mx_all_gather_tasks(
|
||||
[task for task, _ in self._maybe_queue], self.group
|
||||
)
|
||||
self._queue.extend(
|
||||
(task, sender) for task, sender in self._maybe_queue if task in agreed
|
||||
)
|
||||
self._maybe_queue = [
|
||||
(task, sender) for task, sender in self._maybe_queue if task in different
|
||||
]
|
||||
|
||||
def step(self) -> set[TaskId]:
|
||||
if self._pending_close is not None:
|
||||
@@ -105,6 +119,8 @@ class SequentialGenerator(InferenceGenerator):
|
||||
self._pending_close = None
|
||||
|
||||
if self._active is None:
|
||||
self.agree_on_tasks()
|
||||
|
||||
if self._queue:
|
||||
self._start_next()
|
||||
else:
|
||||
@@ -177,6 +193,8 @@ class SequentialGenerator(InferenceGenerator):
|
||||
if mx_any(want_to_cancel, self.group):
|
||||
raise PrefillCancelled()
|
||||
|
||||
self.agree_on_tasks()
|
||||
|
||||
tokens_since_cancel_check = self.check_for_cancel_every
|
||||
|
||||
def on_generation_token() -> None:
|
||||
@@ -191,6 +209,8 @@ class SequentialGenerator(InferenceGenerator):
|
||||
if mx_any(want_to_cancel, self.group):
|
||||
raise PrefillCancelled()
|
||||
|
||||
self.agree_on_tasks()
|
||||
|
||||
return mlx_generate(
|
||||
model=self.model,
|
||||
tokenizer=self.tokenizer,
|
||||
@@ -205,3 +225,55 @@ class SequentialGenerator(InferenceGenerator):
|
||||
|
||||
def close(self) -> None:
|
||||
del self.model, self.tokenizer, self.group
|
||||
|
||||
|
||||
def mx_all_gather_tasks(
|
||||
tasks: list[TextGeneration],
|
||||
group: mx.distributed.Group | None,
|
||||
) -> tuple[list[TextGeneration], list[TextGeneration]]:
|
||||
def encode_task_id(task_id: TaskId) -> list[int]:
|
||||
utf8_task_id = task_id.encode()
|
||||
return [
|
||||
int.from_bytes(utf8_task_id[i : i + 1]) for i in range(len(utf8_task_id))
|
||||
]
|
||||
|
||||
def decode_task_id(encoded_task_id: list[int]) -> TaskId:
|
||||
return TaskId(
|
||||
bytes.decode(b"".join((x).to_bytes(length=1) for x in encoded_task_id))
|
||||
)
|
||||
|
||||
uuid_byte_length = 36
|
||||
|
||||
n_tasks = len(tasks)
|
||||
all_counts = cast(
|
||||
list[int],
|
||||
mx.distributed.all_gather(mx.array([n_tasks]), group=group).tolist(),
|
||||
)
|
||||
max_tasks = max(all_counts)
|
||||
world_size: int = 1 if group is None else group.size()
|
||||
|
||||
if max_tasks == 0:
|
||||
return [], []
|
||||
|
||||
padded = [encode_task_id(task.task_id) for task in tasks] + [
|
||||
[0] * uuid_byte_length
|
||||
] * (max_tasks - n_tasks)
|
||||
gathered = cast(
|
||||
list[list[list[int]]],
|
||||
mx.distributed.all_gather(mx.array(padded), group=group)
|
||||
.reshape(world_size, max_tasks, -1)
|
||||
.tolist(),
|
||||
)
|
||||
all_task_ids: list[list[TaskId]] = [
|
||||
[decode_task_id(encoded_task_id) for encoded_task_id in rank_tasks[:count]]
|
||||
for rank_tasks, count in zip(gathered, all_counts, strict=True)
|
||||
]
|
||||
|
||||
agreed_ids: set[TaskId] = set(all_task_ids[0])
|
||||
for rank_tasks in all_task_ids[1:]:
|
||||
agreed_ids &= set(rank_tasks)
|
||||
|
||||
local_tasks = {task.task_id: task for task in tasks}
|
||||
agreed = [local_tasks[tid] for tid in sorted(agreed_ids)]
|
||||
different = [task for task in tasks if task.task_id not in agreed_ids]
|
||||
return agreed, different
|
||||
|
||||
@@ -158,6 +158,7 @@ class RunnerSupervisor:
|
||||
async def cancel_task(self, task_id: TaskId):
|
||||
if task_id in self.completed:
|
||||
logger.info(f"Unable to cancel {task_id} as it has been completed")
|
||||
self.cancelled.add(task_id)
|
||||
return
|
||||
self.cancelled.add(task_id)
|
||||
with anyio.move_on_after(0.5) as scope:
|
||||
|
||||
Reference in New Issue
Block a user