From 238a4b8f8f5d5a550820de2dcd0da160136af16d Mon Sep 17 00:00:00 2001
From: Mick <mickjagger19@icloud.com>
Date: Thu, 26 Mar 2026 16:33:22 +0800
Subject: [PATCH] [diffusion] CI: fix breaking import path in nightly (#21449)

---
 .../diffusion/publish_comparison_results.py   | 40 ++++++++++++------
 .../utils/diffusion/publish_diffusion_gt.py   | 42 ++++++++++++-------
 2 files changed, 55 insertions(+), 27 deletions(-)

diff --git a/scripts/ci/utils/diffusion/publish_comparison_results.py b/scripts/ci/utils/diffusion/publish_comparison_results.py
index e7cfbb8c1..03a40612c 100644
--- a/scripts/ci/utils/diffusion/publish_comparison_results.py
+++ b/scripts/ci/utils/diffusion/publish_comparison_results.py
@@ -20,19 +20,33 @@ import time
 from datetime import datetime, timezone
 from pathlib import Path
 
-# Reuse GitHub API helpers from publish_traces
-sys.path.insert(0, str(Path(__file__).parent))
-from ..publish_traces import (
-    create_blobs,
-    create_commit,
-    create_tree,
-    get_branch_sha,
-    get_tree_sha,
-    is_permission_error,
-    is_rate_limit_error,
-    update_branch_ref,
-    verify_token_permissions,
-)
+# Reuse GitHub API helpers from publish_traces.
+# Support both direct script execution and package-style imports.
+if __package__:
+    from ..publish_traces import (
+        create_blobs,
+        create_commit,
+        create_tree,
+        get_branch_sha,
+        get_tree_sha,
+        is_permission_error,
+        is_rate_limit_error,
+        update_branch_ref,
+        verify_token_permissions,
+    )
+else:
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+    from publish_traces import (
+        create_blobs,
+        create_commit,
+        create_tree,
+        get_branch_sha,
+        get_tree_sha,
+        is_permission_error,
+        is_rate_limit_error,
+        update_branch_ref,
+        verify_token_permissions,
+    )
 
 # Repository configuration
 REPO_OWNER = "sglang-bot"
diff --git a/scripts/ci/utils/diffusion/publish_diffusion_gt.py b/scripts/ci/utils/diffusion/publish_diffusion_gt.py
index eda5de09f..00b5de0c2 100644
--- a/scripts/ci/utils/diffusion/publish_diffusion_gt.py
+++ b/scripts/ci/utils/diffusion/publish_diffusion_gt.py
@@ -6,21 +6,35 @@ via the GitHub API (same pattern as publish_traces.py).
 import argparse
 import os
 import sys
+from pathlib import Path
 
-# Allow importing from the same directory (scripts/ci/utils/)
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-from ..publish_traces import (
-    create_blobs,
-    create_commit,
-    create_tree,
-    get_branch_sha,
-    get_tree_sha,
-    is_permission_error,
-    is_rate_limit_error,
-    update_branch_ref,
-    verify_token_permissions,
-)
+# Reuse GitHub API helpers from publish_traces.
+# Support both direct script execution and package-style imports.
+if __package__:
+    from ..publish_traces import (
+        create_blobs,
+        create_commit,
+        create_tree,
+        get_branch_sha,
+        get_tree_sha,
+        is_permission_error,
+        is_rate_limit_error,
+        update_branch_ref,
+        verify_token_permissions,
+    )
+else:
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+    from publish_traces import (
+        create_blobs,
+        create_commit,
+        create_tree,
+        get_branch_sha,
+        get_tree_sha,
+        is_permission_error,
+        is_rate_limit_error,
+        update_branch_ref,
+        verify_token_permissions,
+    )
 
 REPO_OWNER = "sglang-bot"
 REPO_NAME = "sglang-ci-data"