wehub-resource-sync
74 lines
2.7 KiB
Python
74 lines
2.7 KiB
Python
import concurrent.futures
|
|
from threading import RLock
|
|
|
|
from mlflow.entities import Metric
|
|
from mlflow.tracking.client import MlflowClient
|
|
|
|
_metrics_queue_lock = RLock()
|
|
_metrics_queue = []
|
|
_thread_pool = concurrent.futures.ThreadPoolExecutor(
|
|
max_workers=1, thread_name_prefix="MlflowMetricsQueue"
|
|
)
|
|
|
|
_MAX_METRIC_QUEUE_SIZE = 500
|
|
|
|
|
|
def _assoc_list_to_map(lst):
|
|
"""
|
|
Convert an association list to a dictionary.
|
|
"""
|
|
d = {}
|
|
for run_id, metric in lst:
|
|
d[run_id] = d[run_id] + [metric] if run_id in d else [metric]
|
|
return d
|
|
|
|
|
|
def flush_metrics_queue():
|
|
"""Flush the metric queue and log contents in batches to MLflow.
|
|
|
|
Queue is divided into batches according to run id.
|
|
"""
|
|
try:
|
|
# Multiple queue flushes may be scheduled simultaneously on different threads
|
|
# (e.g., if the queue is at its flush threshold and several more items
|
|
# are added before a flush occurs). For correctness and efficiency, only one such
|
|
# flush operation should proceed; all others are redundant and should be dropped
|
|
acquired_lock = _metrics_queue_lock.acquire(blocking=False)
|
|
if acquired_lock:
|
|
# For thread safety and to avoid modifying a list while iterating over it, we record a
|
|
# separate list of the items being flushed and remove each one from the metric queue,
|
|
# rather than clearing the metric queue or reassigning it (clearing / reassigning is
|
|
# dangerous because we don't block threads from adding to the queue while a flush is
|
|
# in progress)
|
|
snapshot = _metrics_queue[:]
|
|
for item in snapshot:
|
|
_metrics_queue.remove(item)
|
|
|
|
# Only create MlflowClient if there are metrics to log
|
|
if snapshot:
|
|
client = MlflowClient()
|
|
metrics_by_run = _assoc_list_to_map(snapshot)
|
|
for run_id, metrics in metrics_by_run.items():
|
|
client.log_batch(run_id, metrics=metrics, params=[], tags=[])
|
|
finally:
|
|
if acquired_lock:
|
|
_metrics_queue_lock.release()
|
|
|
|
|
|
def add_to_metrics_queue(key, value, step, time, run_id):
|
|
"""Add a metric to the metric queue.
|
|
|
|
Flush the queue if it exceeds max size.
|
|
|
|
Args:
|
|
key: string, the metrics key,
|
|
value: float, the metrics value.
|
|
step: int, the step of current metric.
|
|
time: int, the timestamp of current metric.
|
|
run_id: string, the run id of the associated mlflow run.
|
|
"""
|
|
met = Metric(key=key, value=value, timestamp=time, step=step)
|
|
_metrics_queue.append((run_id, met))
|
|
if len(_metrics_queue) > _MAX_METRIC_QUEUE_SIZE:
|
|
_thread_pool.submit(flush_metrics_queue)
|