@@ -413,16 +413,23 @@ capabilities, such as exporting aggregated metrics to TensorBoard, are WIP. Once
413413the functionality is complete, v2 will be replacing the original version.
414414
415415` ` ` python
416- from tunix.perf import metrics as perf_metrics
417- from tunix.perf.experimental import export as perf_export_v2
416+ from tunix import PerfMetricsConfig
417+ from tunix.perf.experimental.export import PerfMetricsExport
418418from tunix.rl import rl_cluster
419419
420- # 1. Create a PerfMetricsConfig object.
421- perf_config = perf_metrics.PerfMetricsConfig()
422420
423- # 2. Create the v2 metrics export function, specifying the trace directory.
424- perf_config.custom_export_fn_v2 = (
425- perf_export_v2.PerfMetricsExport(trace_dir="/tmp/perf_trace").export_metrics
421+ # 1. Define cluster_config (as you would for your RL job)
422+ cluster_config = rl_cluster.ClusterConfig(
423+ training_config=training_config,
424+ # ... other configurations
425+ )
426+
427+ # 2. Create a PerfMetricsConfig with custom_export_fn_v2 defined.
428+ perf_metrics_config = PerfMetricsConfig(
429+ custom_export_fn_v2=PerfMetricsExport.from_cluster_config(
430+ cluster_config=cluster_config,
431+ trace_dir="/tmp/agentic_perf",
432+ ).export_metrics
426433)
427434
428435# 3. Pass the config to the RLCluster.
0 commit comments