Skip to content

Commit

Permalink
attempt to deal with missing Linux-only os.sched_getaffinity
Browse files Browse the repository at this point in the history
Signed-off-by: Luke Nezda <[email protected]>
  • Loading branch information
nezda committed Dec 29, 2023
1 parent e68000f commit c1f89be
Showing 1 changed file with 8 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# 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.
import multiprocessing
import os

import torch
Expand Down Expand Up @@ -132,7 +133,13 @@ def init_from_bin(self, model_type, model_path, **generate_kwargs):
if "threads" not in generate_kwargs:
threads = os.getenv("OMP_NUM_THREADS")
if threads is None:
generate_kwargs["threads"] = len(os.sched_getaffinity(0))
# hack from https://stackoverflow.com/questions/74048135/alternatives-to-os-sched-getaffinity-for-macos
# do we need physical core count instead?
try:
# NOTE: only available on some Unix platforms
generate_kwargs["threads"] = len(os.sched_getaffinity(0))
except AttributeError:
generate_kwargs["threads"] = multiprocessing.cpu_count()
else:
generate_kwargs["threads"] = int(threads)
self.model.init_model(model_path, **generate_kwargs)
Expand Down

0 comments on commit c1f89be

Please sign in to comment.