wehub-resource-sync
58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
import threading
|
|
|
|
from mlflow.entities.metric import Metric
|
|
from mlflow.entities.param import Param
|
|
from mlflow.entities.run_tag import RunTag
|
|
|
|
|
|
class RunBatch:
|
|
def __init__(
|
|
self,
|
|
run_id: str,
|
|
params: list["Param"] | None = None,
|
|
tags: list["RunTag"] | None = None,
|
|
metrics: list["Metric"] | None = None,
|
|
completion_event: threading.Event | None = None,
|
|
):
|
|
"""Initializes an instance of `RunBatch`.
|
|
|
|
Args:
|
|
run_id: The ID of the run.
|
|
params: A list of parameters. Default is None.
|
|
tags: A list of tags. Default is None.
|
|
metrics: A list of metrics. Default is None.
|
|
completion_event: A threading.Event object. Default is None.
|
|
"""
|
|
self.run_id = run_id
|
|
self.params = params or []
|
|
self.tags = tags or []
|
|
self.metrics = metrics or []
|
|
self.completion_event = completion_event
|
|
self._exception = None
|
|
self.child_batches = []
|
|
|
|
@property
|
|
def exception(self):
|
|
"""Exception raised during logging the batch."""
|
|
return self._exception
|
|
|
|
@exception.setter
|
|
def exception(self, exception):
|
|
self._exception = exception
|
|
|
|
def add_child_batch(self, child_batch):
|
|
"""Add a child batch to the current batch.
|
|
|
|
This is useful when merging child batches into a parent batch. Child batches are kept so
|
|
that we can properly notify the system when child batches have been processed.
|
|
"""
|
|
self.child_batches.append(child_batch)
|
|
|
|
def complete(self):
|
|
"""Mark the batch as completed."""
|
|
if self.completion_event:
|
|
self.completion_event.set()
|
|
|
|
for child_batch in self.child_batches:
|
|
child_batch.complete()
|