-
Notifications
You must be signed in to change notification settings - Fork 4.4k
[bug-fix] Set number of threads based on allocated CPU count in Docker containers #4471
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
|
Nice catch. I've run into this before and didn't think about it for pytorch. |
| def _read_in_integer_file(filename: str) -> int: | ||
| try: | ||
| with open(filename) as f: | ||
| return int(f.readlines()[0]) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There's some code on https://bugs.python.org/issue36054 - they use the equivalent of int( f.read().rstrip() )
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I like the read().rstrip() better that what I had, changed
| return max(min(num_cpus // 2, 4), 1) if num_cpus is not None else None | ||
|
|
||
|
|
||
| def _get_num_cpus() -> Optional[int]: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: maybe _get_num_available_cpus()?
…r containers (#4471) * Set num threads properly for Docker * Pylint-friendly logic * Use f.read().rstrip() * Change function names
Proposed change(s)
PyTorch by default sets number of threads to 1/2 available cores. However, in Docker containers, the built-in Python methods return the number of cores of the host machine, not the container. This PR reads the allocated CPU count from
cgroupand uses that to set the num threads.Types of change(s)
Checklist
Other comments