from __future__ import annotations
from typing import Union, TextIO, Tuple, List, Optional
import sys
import asyncio
import logging
import time
from pathlib import Path
from itertools import chain
from datetime import datetime
import ruamel.yaml
from palaestrai.core import RuntimeConfig
from palaestrai.experiment import ExperimentRun, Executor, ExecutorState
from palaestrai.util.runtime_profiling import write_runtime_profile_reports
ExperimentRunInputTypes = Union[ExperimentRun, TextIO, str, Path]
LOG = logging.getLogger("palaestrai.runner")
def _maybe_generate_runtime_profile(
run_start_ns: int, run_end_ns: int
) -> None:
config = RuntimeConfig()
if not config.runtime_profiling_enabled:
return
if not config.otel_enabled:
LOG.warning(
"runtime_profiling_enabled is set, but otel_enabled is false; skipping runtime report generation."
)
return
if config.otel_exporter != "jsonl":
LOG.warning(
"runtime_profiling_enabled requires otel_exporter=jsonl; got %s. Skipping runtime report generation.",
config.otel_exporter,
)
return
if not config.otel_jsonl_file:
LOG.warning(
"runtime_profiling_enabled requires otel_jsonl_file to be set. Skipping runtime report generation."
)
return
try:
summary = write_runtime_profile_reports(
trace_file=Path(config.otel_jsonl_file),
output_dir=Path(config.data_path),
run_start_ns=run_start_ns,
run_end_ns=run_end_ns,
include_critical_path=config.runtime_profiling_critical_path_report,
)
LOG.info(
"Runtime profiling summary generated at %s (steps=%s, acts=%s).",
Path(config.data_path),
summary.get("step_stats_ms", {}).get("count"),
summary.get("act_stats_ms", {}).get("count"),
)
except Exception:
LOG.exception("Failed to generate runtime profiling reports.")
[docs]
async def run(
experiment_run_definition: Union[
ExperimentRunInputTypes, List[ExperimentRunInputTypes]
],
runtime_config: Union[str, TextIO, dict, None] = None,
parallel_runs: int = 1,
user: Optional[str] = None,
) -> Tuple[List[str], ExecutorState]:
"""Provides a single-line command to start an experiment
This is the asynchronous variant of :func:`.execute`. Use it if you
have already an event loop running (e.g., in a Jupyter Notebook).
"""
if runtime_config:
RuntimeConfig().load(runtime_config)
# There is an implicit loading of the default config. The object returned
# by RuntimeConfig() has at least the default loaded, and tries to load
# from the search path. So there is no reason to have an explicit load()
# here.
if not isinstance(experiment_run_definition, List):
experiment_run_definition = [experiment_run_definition]
experiment_run_definition = [
Path(i) if isinstance(i, str) else i for i in experiment_run_definition
]
experiment_run_definition = list(
chain.from_iterable(
i.rglob("*.y*ml") if isinstance(i, Path) and i.is_dir() else [i]
for i in experiment_run_definition
)
)
experiment_runs = [
ExperimentRun.load(i) if not isinstance(i, ExperimentRun) else i
for i in experiment_run_definition
]
if user is not None:
for experiment_run in experiment_runs:
experiment_run.user = user
from palaestrai.util.otel import init_tracer
init_tracer("Executor", RuntimeConfig())
executor = Executor(parallel_runs=parallel_runs)
executor.schedule(experiment_runs)
run_start_ns = time.time_ns()
executor_final_state = await executor.execute()
run_end_ns = time.time_ns()
_maybe_generate_runtime_profile(run_start_ns, run_end_ns)
return [e.uid for e in experiment_runs], executor_final_state
[docs]
def execute(
experiment_run_definition: Union[
ExperimentRunInputTypes, List[ExperimentRunInputTypes]
],
runtime_config: Union[str, TextIO, dict, None] = None,
parallel_runs: int = 1,
user: Optional[str] = None,
) -> Tuple[List[str], ExecutorState]:
"""Provides a single-line command to start an experiment
This function is a high-level wrapper that takes a number of experiment
run defintions
and optionally a runtime config,
executes the runs, and returns the results.
Note
----
This is a sync method. If you already have an event loop running, please
use this function's sibling, :func:`.run`
Parameters
----------
experiment_run_definition: 1. Already set ExperimentRun object
2. Any text stream
3. The path to a file
The configuration from which the experiment is loaded.
runtime_config: 1. Any text stream
2. dict
3. None
The Runtime configuration applicable for the run.
Note that even when no additional source is provided, runtime will load
a minimal configuration from build-in defaults.
parallel_runs : int, default: 1
Number of experiment runs that can be executed in parallel
Returns
-------
typing.Tuple[Sequence[str], ExecutorState]
A tuple containing:
1. The list of all experiment run IDs
2. The final state of the executor
"""
loop = None
try:
loop = asyncio.get_running_loop()
except RuntimeError:
pass # No loop - this is good.
if loop is not None and loop.is_running():
raise RuntimeError(
"An event loop is already running. "
"Please use `await palaestrai.run(...)` instead."
)
try:
import uvloop
LOG.debug("Using uvloop.")
if sys.version_info >= (3, 11):
with asyncio.Runner(loop_factory=uvloop.new_event_loop) as runner:
return runner.run(
run(
experiment_run_definition,
runtime_config,
parallel_runs,
user,
)
)
else:
uvloop.install()
return asyncio.run(
run(
experiment_run_definition,
runtime_config,
parallel_runs,
user,
)
)
except ModuleNotFoundError:
LOG.debug("uvloop not available.")
return asyncio.run(
run(experiment_run_definition, runtime_config, parallel_runs, user)
)