add trace test

setup.py

@@ -52,7 +52,7 @@
     packages=find_packages(),
     package_data={'pfio': package_data},
     extras_require={
-        'test': ['pytest', 'flake8', 'autopep8', 'parameterized', 'isort', 'moto[server]>=5.0.0', 'numpy', 'mypy'],
+        'test': ['pytest', 'flake8', 'autopep8', 'parameterized', 'isort', 'moto[server]>=5.0.0', 'numpy', 'mypy', 'pytorch-pfn-extras'],
         # When updating doc deps, docs/requirements.txt should be updated too
         'doc': ['sphinx', 'sphinx_rtd_theme'],
         'bench': ['numpy>=1.19.5', 'torch>=1.9.0', 'Pillow<=8.2.0'],

tests/v2_tests/test_local.py

@@ -1,11 +1,13 @@
 import io
+import json
 import os
 import pickle
 import tempfile
 import unittest
 from collections.abc import Iterable
 
 import pytest
+import pytorch_pfn_extras as ppe
 
 from pfio.v2 import Local, LocalFileStat, from_url, open_url
 
@@ -291,6 +293,45 @@ def test_from_url_create_option(self):
         from_url(path, create=True)
         assert os.path.exists(path) and os.path.isdir(path)
 
+    def test_ppe_profiling(self):
+        with Local(self.testdir.name, trace=True) as fs:
+            ppe.profiler.clear_tracer()
+
+            with fs.open('foo.txt', 'w') as fp:
+                fp.write('bar')
+
+            dict = ppe.profiler.get_tracer().state_dict()
+            keys = [event["name"] for event in json.loads(dict['_event_list'])]
+
+            assert "pfio.v2.Local:open" in keys
+            assert "pfio.v2.Local:write" in keys
+
+        with Local(self.testdir.name, trace=True) as fs:
+            ppe.profiler.clear_tracer()
+
+            with fs.open('foo.txt', 'r') as fp:
+                tmp = fp.read()
+                assert tmp == 'bar'
+
+            dict = ppe.profiler.get_tracer().state_dict()
+            keys = [event["name"] for event in json.loads(dict['_event_list'])]
+
+            assert "pfio.v2.Local:open" in keys
+            assert "pfio.v2.Local:read" in keys
+
+        with Local(self.testdir.name, trace=False) as fs:
+            ppe.profiler.clear_tracer()
+
+            with fs.open('foo.txt', 'r') as fp:
+                tmp = fp.read()
+                assert tmp == 'bar'
+
+            dict = ppe.profiler.get_tracer().state_dict()
+            keys = [event["name"] for event in json.loads(dict['_event_list'])]
+
+            assert "pfio.v2.Local:open" not in keys
+            assert "pfio.v2.Local:read" not in keys
+
 
 if __name__ == '__main__':
     unittest.main()