Split PPE methods

pfnet · Aug 26, 2024 · e6c128c · e6c128c
1 parent cb533c8
commit e6c128c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 23 deletions.
diff --git a/pfio/v2/_profiler.py b/pfio/v2/_profiler.py
@@ -0,0 +1,21 @@
+try:
+    from pytorch_pfn_extras.profiler import record, record_iterable
+
+except ImportError:
+
+    class _DummyRecord:
+        def __init__(self):
+            pass
+
+        def __enter__(self):
+            pass
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            pass
+
+    # IF PPE is not available, wrap with noop
+    def record(tag, trace, *args):  # type: ignore # NOQA
+        return _DummyRecord()
+
+    def record_iterable(tag, iter, trace, *args):   # type: ignore # NOQA
+        yield from iter
diff --git a/pfio/v2/local.py b/pfio/v2/local.py
@@ -4,29 +4,7 @@
 import shutil
 from typing import Optional
 
-try:
-    from pytorch_pfn_extras.profiler import record, record_iterable
-
-except ImportError:
-
-    class _DummyRecord:
-        def __init__(self):
-            pass
-
-        def __enter__(self):
-            pass
-
-        def __exit__(self, exc_type, exc_value, traceback):
-            pass
-
-    # IF PPE is not available, wrap with noop
-    def record(tag, trace, *args):  # type: ignore # NOQA
-        return _DummyRecord()
-
-    def record_iterable(tag, iter, trace, *args):   # type: ignore # NOQA
-        yield from iter
-
-
+from ._profiler import record, record_iterable
 from .fs import FS, FileStat, format_repr