Skip to content

Instantly share code, notes, and snippets.

@jart
Created October 23, 2016 06:44
Show Gist options
  • Save jart/0a71cde3ca7261f77080a3625a21672b to your computer and use it in GitHub Desktop.
Save jart/0a71cde3ca7261f77080a3625a21672b to your computer and use it in GitHub Desktop.
Multithreaded Python os.walk
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Routines for multi-threaded i/o."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import sys
import threading
def walk(top, threads=60):
"""Multi-threaded version of os.walk().
This routine provides multiple orders of a magnitude performance improvement
when top is mapped to a network filesystem where i/o operations are slow, but
unlimited. For spinning disks it should still run faster regardless of thread
count because it uses a LIFO scheduler that guarantees locality. For SSDs it
will go tolerably slower.
The more exotic coroutine features of os.walk() can not be supported, such as
the ability to selectively inhibit recursion by mutating subdirs.
Args:
top: Path of parent directory to search recursively.
threads: Size of fixed thread pool.
Yields:
A (path, subdirs, files) tuple for each directory within top, including
itself. These tuples come in no particular order; however, the contents of
each tuple itself is sorted.
"""
if not os.path.isdir(top):
return
lock = threading.Lock()
on_input = threading.Condition(lock)
on_output = threading.Condition(lock)
state = {'tasks': 1}
paths = [top]
output = []
def worker():
while True:
with lock:
while True:
if not state['tasks']:
output.append(None)
on_output.notify()
return
if not paths:
on_input.wait()
continue
path = paths.pop()
break
try:
dirs = []
files = []
for item in sorted(os.listdir(path)):
subpath = os.path.join(path, item)
if os.path.isdir(subpath):
dirs.append(item)
with lock:
state['tasks'] += 1
paths.append(subpath)
on_input.notify()
else:
files.append(item)
with lock:
output.append((path, dirs, files))
on_output.notify()
except OSError as e:
print(e, file=sys.stderr)
finally:
with lock:
state['tasks'] -= 1
if not state['tasks']:
on_input.notifyAll()
workers = [threading.Thread(target=worker,
name="fastio.walk %d %s" % (i, top))
for i in range(threads)]
for w in workers:
w.start()
while threads or output: # TODO(jart): Why is 'or output' necessary?
with lock:
while not output:
on_output.wait()
item = output.pop()
if item:
yield item
else:
threads -= 1
@alkuzad
Copy link

alkuzad commented Nov 17, 2017

This is threaded but not super fast as it:

  1. sorts output of os.listdir (not every case needs that)
  2. uses os.path.isdir which is implemented in such way that it uses os.stat. This is especially crucial on NFS.

This code can be much faster when implemented with optional sorting done at the end (or dropped) and leveraging https://pypi.python.org/pypi/scandir (builtin in 3.5, packaged in pypi for 2.7 to 3.6)

Edit: Just note that besides that it's very good code indeed. Well done.

@TheCodeGeek2
Copy link

TheCodeGeek2 commented Jan 10, 2023

Perhaps this just doesn't work when replacing os.listdir() but this does not work in my case.
Is there something that I'm not doing correctly?
for filename in fastio.walk(sourcedir, 60):

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment