# -*- coding: utf-8 -*-
"""The Workflow Runner Module."""
import asyncio
import os
import time
import traceback
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
from trinity.buffer import get_buffer_reader, get_buffer_writer
from trinity.common.config import Config, StorageConfig
from trinity.common.constants import LOG_DIR_ENV_VAR, LOG_LEVEL_ENV_VAR
from trinity.common.experience import Experience
from trinity.common.models.allocator import Allocator
from trinity.common.models.model import ModelWrapper
from trinity.common.workflows import Task, Workflow
from trinity.utils.log import get_logger
from trinity.utils.metrics import aggregate_run_level_metrics
[docs]
@dataclass(frozen=True)
class Status:
"""Status of the task running result."""
completed_runs: int
total_runs: int
metrics: List[Dict[str, float]]
# A list of metric dictionaries, where each dictionary is from a single run.
message: Optional[str] = None
@property
def ok(self) -> bool:
return self.completed_runs == self.total_runs
[docs]
@dataclass(frozen=True)
class RunnerExecutionResult:
"""Execution result for one runner task."""
status: Status
experiences: List[Experience]
[docs]
class WorkflowRunner:
"""A Ray remote actor to run the workflow and generate experiences."""
[docs]
def __init__(
self,
config: Config,
rollout_model_id: int,
auxiliary_model_ids: Optional[List[int]] = None,
runner_id: Optional[int] = None,
) -> None:
self.name = f"{config.explorer.name}_runner_{runner_id}"
self.logger = get_logger(self.name, in_ray_actor=True)
self.config = config
allocator = Allocator(config.explorer)
self.model_wrapper: ModelWrapper = allocator.get_model(
config.explorer.rollout_model, "rollout", rollout_model_id
)
self.auxiliary_model_wrappers: List[ModelWrapper] = [
allocator.get_model(
config.explorer.auxiliary_models[index], f"auxiliary_{index}", auxiliary_model_id
)
for index, auxiliary_model_id in enumerate(auxiliary_model_ids or [])
]
self.workflow_instance: Workflow = None
self.runner_id = runner_id
self.runner_state = {
"workflow_id": None,
"model_version": None,
"begin_time": 0,
"terminate_time": 0,
}
self.concurrent_mode = config.explorer.concurrent_mode
if self.concurrent_mode == "sequential":
self.concurrent_run_fn = self._sequential_run
elif self.concurrent_mode == "asynchronous":
self.concurrent_run_fn = self._asynchronous_run
elif self.concurrent_mode == "multi-threading":
self.concurrent_run_fn = self._multi_threading_run
else:
self.logger.warning(
f"Unknown concurrent_mode {self.concurrent_mode}, defaulting to sequential."
)
self.concurrent_run_fn = self._sequential_run
self.logger.info(
f"WorkflowRunner [{self.name}]({self.concurrent_mode}) initialized:\n"
f" > rollout model: {self.config.explorer.rollout_model.model_path}\n"
f" > auxiliary models: {[aux_model_config.model_path for aux_model_config in self.config.explorer.auxiliary_models]}"
)
[docs]
async def prepare(self) -> None:
"""Prepare the runner."""
await asyncio.gather(
self.model_wrapper.prepare(),
*(aux_model.prepare() for aux_model in self.auxiliary_model_wrappers),
)
self.logger.info(f"WorkflowRunner [{self.name}] is prepared and ready to run tasks.")
[docs]
def is_alive(self):
return True
def _create_workflow_instance(self, task: Task) -> Workflow:
if task.workflow is None:
raise ValueError("Workflow is not set in the task.")
if (
self.workflow_instance is None
or not self.workflow_instance.__class__ == task.workflow
or not self.workflow_instance.resettable
):
# Pass ModelWrapper directly; Workflow.__init__ will get OpenAI clients automatically
self.workflow_instance = task.to_workflow(
self.model_wrapper,
self.auxiliary_model_wrappers,
)
else:
self.workflow_instance.reset(task)
return self.workflow_instance
async def _run_workflow(self, workflow_instance: Workflow) -> List[Experience]:
if workflow_instance.asynchronous:
exps = await workflow_instance.run_async()
else:
exps = workflow_instance.run()
return exps
def _create_isolated_workflow_instance(self, task: Task) -> Workflow:
return task.to_workflow(
(
self.model_wrapper.clone_with_isolated_history()
if self.config.explorer.rollout_model.enable_history
else self.model_wrapper
),
self.auxiliary_model_wrappers,
)
def _build_execution_result(
self,
total_runs: int,
completed_runs: int,
metrics: List[Dict[str, float]],
experiences: List[Experience],
first_error: Optional[str] = None,
) -> RunnerExecutionResult:
if first_error is None:
message = None
elif completed_runs > 0:
message = (
f"{completed_runs}/{total_runs} runs completed successfully. "
f"First error: {first_error}"
)
else:
message = first_error
return RunnerExecutionResult(
status=Status(
completed_runs=completed_runs,
total_runs=total_runs,
metrics=list(metrics),
message=message,
),
experiences=experiences,
)
def _aggregate_run_results(
self,
total_runs: int,
results: List[Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]],
) -> RunnerExecutionResult:
exps = []
run_metrics = []
first_error = None
for ok, new_exps, run_metric, error in results:
if ok:
exps.extend(new_exps)
if run_metric is not None:
run_metrics.append(run_metric)
continue
if first_error is None:
first_error = error
return self._build_execution_result(
total_runs=total_runs,
completed_runs=len(run_metrics),
metrics=run_metrics,
experiences=exps,
first_error=first_error,
)
async def _run_parallel_runs(
self,
task: Task,
repeat_times: int,
run_id_base: int,
collect_partial_runs: bool = True,
use_threads: bool = False,
) -> RunnerExecutionResult:
async def run_single(
i: int,
) -> Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]:
workflow = self._create_isolated_workflow_instance(task)
return await self._execute_single_run(workflow, task, i, run_id_base)
if collect_partial_runs:
if use_threads:
results = await asyncio.gather(
*(
asyncio.to_thread(lambda idx=i: asyncio.run(run_single(idx))) # type: ignore[misc]
for i in range(repeat_times)
)
)
else:
results = await asyncio.gather(*(run_single(i) for i in range(repeat_times)))
return self._aggregate_run_results(repeat_times, results)
future_to_run_index = {}
for i in range(repeat_times):
if use_threads:
future = asyncio.create_task(
asyncio.to_thread(lambda idx=i: asyncio.run(run_single(idx))) # type: ignore[misc]
)
else:
future = asyncio.create_task(run_single(i))
future_to_run_index[future] = i
results = []
while future_to_run_index:
done, pending = await asyncio.wait(
future_to_run_index.keys(),
return_when=asyncio.FIRST_COMPLETED,
)
should_stop = False
for future in done:
future_to_run_index.pop(future)
result = future.result()
results.append(result)
ok, _, _, _ = result
if not ok:
should_stop = True
if should_stop:
for future in pending:
future.cancel()
if pending:
await asyncio.gather(*pending, return_exceptions=True)
break
return self._aggregate_run_results(repeat_times, results)
async def _execute_single_run(
self,
workflow: Workflow,
task: Task,
run_index: int,
run_id_base: int,
) -> Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]:
st = time.time()
workflow.repeat_index = run_id_base + run_index
await self.model_wrapper.clean_workflow_state()
self.runner_state["workflow_id"] = f"{task.batch_id}/{task.task_id}/{run_index}"
self.runner_state["terminate_time"] = None
self.runner_state["begin_time"] = st
try:
new_exps = await self._run_workflow(workflow)
et = time.time()
self.runner_state["terminate_time"] = et
run_metric = aggregate_run_level_metrics(
[exp.metrics for exp in new_exps if exp.metrics]
)
run_metric["time/run_execution"] = et - st
for exp in new_exps:
exp.eid.run = run_id_base + run_index
return True, new_exps, run_metric, None
except Exception as exc:
self.runner_state["terminate_time"] = time.time()
error_trace_back = traceback.format_exc()
self.logger.error(
"WorkflowRunner single run error: " f"{exc}\nTraceback:\n{error_trace_back}"
)
return False, [], None, error_trace_back.rstrip()
async def _run_task(
self,
task: Task,
repeat_times: int,
run_id_base: int,
collect_partial_runs: bool = True,
) -> RunnerExecutionResult:
"""Init workflow from the task and run it."""
if task.workflow.can_repeat:
workflow_instance = self._create_workflow_instance(task)
workflow_instance.set_repeat_times(repeat_times, run_id_base)
st = time.time()
await self.model_wrapper.clean_workflow_state()
self.runner_state["workflow_id"] = f"{task.batch_id}/{task.task_id}/{run_id_base}"
self.runner_state["terminate_time"] = None
self.runner_state["begin_time"] = st
exps = await self._run_workflow(workflow_instance)
et = time.time()
self.runner_state["terminate_time"] = et
# repeatable workflow cannot calculate run level metrics, we use experience level metrics directly
run_metrics = [exp.metrics for exp in exps if exp.metrics]
for metric in run_metrics:
metric["time/run_execution"] = et - st
return self._build_execution_result(
total_runs=repeat_times,
completed_runs=repeat_times,
metrics=run_metrics,
experiences=exps,
)
else:
return await self.concurrent_run_fn(
task,
repeat_times,
run_id_base,
collect_partial_runs=collect_partial_runs,
)
async def _sequential_run(
self,
task: Task,
repeat_times: int,
run_id_base: int,
collect_partial_runs: bool = True,
) -> RunnerExecutionResult:
results = []
for i in range(repeat_times):
workflow = self._create_workflow_instance(task)
result = await self._execute_single_run(workflow, task, i, run_id_base)
results.append(result)
if collect_partial_runs:
continue
ok, _, _, _ = result
if ok:
continue
break
return self._aggregate_run_results(repeat_times, results)
async def _asynchronous_run(
self,
task: Task,
repeat_times: int,
run_id_base: int,
collect_partial_runs: bool = True,
) -> RunnerExecutionResult:
return await self._run_parallel_runs(
task,
repeat_times,
run_id_base,
collect_partial_runs=collect_partial_runs,
)
async def _multi_threading_run(
self,
task: Task,
repeat_times: int,
run_id_base: int,
collect_partial_runs: bool = True,
) -> RunnerExecutionResult:
return await self._run_parallel_runs(
task,
repeat_times,
run_id_base,
collect_partial_runs=collect_partial_runs,
use_threads=True,
)
[docs]
async def get_runner_state(self) -> Dict:
"""Get the runner state."""
runner_state = self.runner_state.copy()
runner_state.update(await self.model_wrapper.get_workflow_state())
return runner_state
[docs]
async def run_task(
self,
task: Task,
batch_id: str,
repeat_times: int = 1,
run_id_base: int = 0,
collect_partial_runs: bool = True,
) -> Tuple[Status, bytes]:
"""Run the task and return the states."""
st = time.time()
try:
model_version = await self.model_wrapper.model_version_async
self.runner_state["model_version"] = model_version
self.logger.info(
f"Starting task: step={batch_id}, model_version={model_version}, repeat_times={repeat_times}, run_id_base={run_id_base}"
)
execution_result = await self._run_task(
task,
repeat_times,
run_id_base,
collect_partial_runs=collect_partial_runs,
)
model_version_after = await self.model_wrapper.model_version_async
exps = execution_result.experiences
if execution_result.status.completed_runs > 0:
assert exps is not None and len(exps) > 0, "An empty experience is generated"
# set eid for each experience
for exp in exps:
exp.eid.batch = task.batch_id
# keep exp.eid.task if it has been set before (e.g., in workflow)
if exp.eid.task == "": # "" is the default value
exp.eid.task = task.task_id
if not hasattr(exp, "info") or exp.info is None:
exp.info = {}
exp.info["model_version"] = model_version
exp.info["model_version_drift"] = model_version_after - model_version
exp.info["use_count"] = 0
exp.info["task_index"] = task.index
if not hasattr(exp, "metrics") or exp.metrics is None:
exp.metrics = {}
status = execution_result.status
if task.is_eval:
# If the task is an evaluation task, we do not record the experiences to the buffer
return status, b""
else:
exp_payload = Experience.serialize_many(exps)
return status, exp_payload
except Exception as e:
error_trace_back = traceback.format_exc()
self.logger.error(f"WorkflowRunner run task error: {e}\nTraceback:\n{error_trace_back}")
return (
Status(
completed_runs=0,
total_runs=repeat_times,
metrics=[{"time/run_execution": time.time() - st}],
message=error_trace_back.rstrip(),
),
b"",
)
[docs]
class DebugWorkflowRunner(WorkflowRunner):
"""A WorkflowRunner for debugging."""
[docs]
def __init__(
self,
config: Config,
output_dir: str = "debug_output",
enable_profiling: bool = False,
disable_overwrite: bool = False,
) -> None:
if disable_overwrite:
# if output dir is not empty, change to a new dir with datetime suffix
if os.path.isdir(output_dir) and os.listdir(output_dir):
suffix = time.strftime("%Y%m%d%H%M%S", time.localtime())
output_dir = f"{output_dir}_{suffix}"
os.environ[LOG_DIR_ENV_VAR] = os.path.join(output_dir, "log")
os.environ[LOG_LEVEL_ENV_VAR] = "DEBUG"
super().__init__(
config=config,
rollout_model_id=0,
auxiliary_model_ids=[0] * len(config.explorer.auxiliary_models),
runner_id=0,
)
self.taskset = get_buffer_reader(config.buffer.explorer_input.tasksets[0])
self.output_dir = output_dir
self.enable_profiling = enable_profiling
self.logger.info(f"Debug output directory: {self.output_dir}")
os.makedirs(self.output_dir, exist_ok=True)
self.output_profiling_file = os.path.join(
self.output_dir,
"profiling.html",
)
self.output_sqlite_file = "sqlite:///" + os.path.join(
self.output_dir,
"experiences.db",
)
self.sqlite_writer = get_buffer_writer(
StorageConfig(
name="debug_buffer",
schema_type="experience",
path=self.output_sqlite_file,
storage_type="sql",
batch_size=1,
wrap_in_ray=False,
)
)
[docs]
async def debug(self) -> None:
"""Run the debug workflow."""
tasks = await self.taskset.read(batch_size=1)
task = tasks[0]
self.logger.info(f"Start debugging task:\n{task.raw_task}")
if not self.enable_profiling:
status, exp_payload = await self.run_task(
task=task, batch_id="debug", repeat_times=1, run_id_base=0
)
else:
from viztracer import VizTracer
with VizTracer(output_file=self.output_profiling_file):
status, exp_payload = await self.run_task(
task=task, batch_id="debug", repeat_times=1, run_id_base=0
)
experiences = Experience.deserialize_many(exp_payload) if exp_payload else []
if not status.ok and not experiences:
experiences = self.model_wrapper.extract_experience_from_history()
self.logger.info(
f"Debugging failed, extracting {len(experiences)} experiences from history."
)
await self.sqlite_writer.write(experiences)
if status.ok:
print(f"Task {task.task_id} completed successfully with metrics:\n{status.metrics}")
else:
self.logger.error(f"Task {task.task_id} failed with message: {status.message}")
self.logger.info("Debugging completed.")
