diff --git a/sdks/python/apache_beam/utils/profiler.py b/sdks/python/apache_beam/utils/profiler.py
index 74d6abc600f5..9f9c8cd16296 100644
--- a/sdks/python/apache_beam/utils/profiler.py
+++ b/sdks/python/apache_beam/utils/profiler.py
@@ -21,15 +21,20 @@
 """
 
 import cProfile
-import io
 import logging
 import os
 import pstats
+import sys
 import tempfile
 import time
 import warnings
 from threading import Timer
 
+if sys.version_info[0] < 3:
+  import StringIO
+else:
+  from io import StringIO
+
 
 class Profile(object):
   """cProfile wrapper context for saving and logging profiler results."""
@@ -66,7 +71,7 @@ def __exit__(self, *args):
       os.remove(filename)
 
     if self.log_results:
-      s = io.StringIO()
+      s = StringIO()
       self.stats = pstats.Stats(
           self.profile, stream=s).sort_stats(Profile.SORTBY)
       self.stats.print_stats()