Skip to content

Commit

Permalink
fix: change the color order
Browse files Browse the repository at this point in the history
Signed-off-by: Frost Ming <me@frostming.com>
  • Loading branch information
frostming committed Aug 27, 2024
1 parent 1ebeeb7 commit a99bab6
Showing 1 changed file with 7 additions and 11 deletions.
18 changes: 7 additions & 11 deletions src/bentoml/_internal/cloud/deployment.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
import os
import time
import typing as t
from functools import partial
from os import path
from threading import Event
from threading import Thread
Expand Down Expand Up @@ -829,18 +828,17 @@ def _tail_logs(
cloud_rest_client: RestApiClient = Provide[BentoMLContainer.rest_api_client],
) -> t.Generator[None, None, None]:
import itertools
from collections import defaultdict

pods = cloud_rest_client.v2.list_deployment_pods(self.name, self.cluster)
stop_event = Event()
threads: list[Thread] = []
workers: list[Thread] = []

colors = itertools.cycle(["green", "cyan", "yellow", "blue", "magenta"])
runner_color: dict[str, str] = {}
colors = itertools.cycle(["cyan", "yellow", "blue", "magenta", "green"])
runner_color: dict[str, str] = defaultdict(lambda: next(colors))

def pod_log_worker(pod: KubePodSchema, stop_event: Event) -> None:
current = ""
if pod.runner_name not in runner_color:
runner_color[pod.runner_name] = next(colors)
color = runner_color[pod.runner_name]
for chunk in cloud_rest_client.v2.tail_logs(
cluster_name=self.cluster,
Expand Down Expand Up @@ -868,15 +866,13 @@ def pod_log_worker(pod: KubePodSchema, stop_event: Event) -> None:
for pod in pods:
if pod.labels.get("yatai.ai/is-bento-image-builder") == "true":
continue
thread = Thread(
target=partial(pod_log_worker, pod=pod, stop_event=stop_event)
)
thread = Thread(target=pod_log_worker, args=(pod, stop_event))
thread.start()
threads.append(thread)
workers.append(thread)
yield
finally:
stop_event.set()
for thread in threads:
for thread in workers:
thread.join()


Expand Down

0 comments on commit a99bab6

Please sign in to comment.