Skip to content

Conversation

@ervteng
Copy link
Contributor

@ervteng ervteng commented Sep 11, 2020

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 cgroup and uses that to set the num threads.

Types of change(s)

  • Bug fix
  • New feature
  • Code refactor
  • Breaking change
  • Documentation update
  • Other (please describe)

Checklist

  • Added tests that prove my fix is effective or that my feature works
  • Updated the changelog (if applicable)
  • Updated the documentation (if applicable)
  • Updated the migration guide (if applicable)

Other comments

@chriselion
Copy link
Contributor

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])
Copy link
Contributor

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() )

Copy link
Contributor Author

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]:
Copy link
Contributor

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()?

@ervteng ervteng merged commit e49a0df into master Sep 11, 2020
@delete-merged-branch delete-merged-branch bot deleted the develop-torch-docker-cpu branch September 11, 2020 21:24
chriselion pushed a commit that referenced this pull request Sep 12, 2020
…r containers (#4471)

* Set num threads properly for Docker

* Pylint-friendly logic

* Use f.read().rstrip()

* Change function names
chriselion pushed a commit that referenced this pull request Sep 14, 2020
…r containers (#4471) (#4478)

* Set num threads properly for Docker

* Pylint-friendly logic

* Use f.read().rstrip()

* Change function names

Co-authored-by: Ervin T <ervin@unity3d.com>
@github-actions github-actions bot locked as resolved and limited conversation to collaborators Sep 12, 2021
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants