补充

tools/tn_contest_runner.py

@@ -199,7 +199,7 @@ def build_parallel_opts(args, tree_file=None, search_only=False):
         "search_workers": args.tn_search_workers or args.torch_threads,
         "max_repeats": args.tn_search_repeats,
         "max_time": args.tn_search_time,
-        "print_stats": not args.no_tn_stats,
+        "print_stats": False,
     }
     if args.tn_search_backend is not None:
         opts["search_backend"] = args.tn_search_backend
@@ -303,7 +303,7 @@ def run_one(args, case_name, obs_name, mode):
             f"failed_trials={search_stats.get('failed_trials', 'na')} "
             f"requested_trials={search_stats.get('requested_trials', 'na')} "
             f"best_score={search_stats.get('best_score', float('nan')):.6g} "
-            f"slices={cost.get('slices')} "
+            f"slices={cost.get('nslices')} "
             f"log10_flops={cost.get('log10_flops', float('nan')):.3f} "
             f"log10_write={cost.get('log10_write', float('nan')):.3f} "
             f"log2_size={cost.get('log2_size', float('nan')):.3f} "
@@ -337,6 +337,11 @@ def apply_case_defaults(args):
 def stop_dask_cluster(args):
     if args.keep_dask or args.tn_search_backend != "dask" or not args.dask_address:
         return
+    if args.mpi:
+        from mpi4py import MPI
+
+        if MPI.COMM_WORLD.Get_rank() != 0:
+            return
     script = ROOT / "tools" / "manage_tn_dask_cluster.sh"
     if not script.exists():
         print(f"dask_stop_skipped reason=missing_script path={script}", flush=True)