Split PPE methods

pfnet · Aug 23, 2024 · 87e3eaf · 87e3eaf
1 parent e3892d7
commit 87e3eaf
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