From 4e26ffa1addfcbd6ef072394c6ef19ecc2a07f68 Mon Sep 17 00:00:00 2001
From: xRAIKIRI <692817779@qq.com>
Date: Thu, 20 Aug 2026 22:41:29 +0800
Subject: [PATCH] Import data_loader only under TYPE_CHECKING (annotation-only
 import)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Related to #1019.

`openpi/training/checkpoints.py` imports `openpi.training.data_loader` at module top level, but the module already uses `from __future__ import annotations`, so the import is never evaluated at runtime — it exists only for two type annotations and unnecessarily pulls TF-related native libraries at import time.

This PR moves it under `if TYPE_CHECKING:`, a common practice for annotation-only imports.

## Verification
- Import of `openpi.policies.policy` + `openpi.policies.policy_config` passes 20/20 on a clean environment (no workaround).
- `ruff check` / `ruff format` pass on the changed file.

Note: the segfault reported in #1019 could not be reliably reproduced later on clean 535/580 environments; this PR is a defensive hygiene improvement rather than a claimed fix.
---
 src/openpi/training/checkpoints.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/openpi/training/checkpoints.py b/src/openpi/training/checkpoints.py
index f32a831397..d4273447d8 100644
--- a/src/openpi/training/checkpoints.py
+++ b/src/openpi/training/checkpoints.py
@@ -4,7 +4,7 @@
 import concurrent.futures as futures
 import dataclasses
 import logging
-from typing import Protocol
+from typing import TYPE_CHECKING, Protocol
 
 from etils import epath
 import jax
@@ -13,7 +13,9 @@
 
 from openpi.shared import array_typing as at
 import openpi.shared.normalize as _normalize
-import openpi.training.data_loader as _data_loader
+
+if TYPE_CHECKING:
+    import openpi.training.data_loader as _data_loader
 import openpi.training.utils as training_utils