Skip to content

Commit

Permalink
PERF-#6614: HDK: Use MODIN_CPUS instead of os.cpu_count() for the…
Browse files Browse the repository at this point in the history
… fragment size calculation (#6615)

Signed-off-by: Andrey Pavlenko <[email protected]>
  • Loading branch information
AndreyPavlenko authored Oct 2, 2023
1 parent fc17c48 commit 4531ae1
Showing 1 changed file with 2 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,14 @@
# governing permissions and limitations under the License.

"""Module provides ``HdkWorker`` class."""
import os
from typing import List, Optional, Tuple, Union

import pyarrow as pa
import pyhdk
from packaging import version
from pyhdk.hdk import HDK, ExecutionResult, QueryNode, RelAlgExecutor

from modin.config import HdkFragmentSize, HdkLaunchParameters
from modin.config import CpuCount, HdkFragmentSize, HdkLaunchParameters
from modin.utils import _inherit_docstrings

from .base_worker import BaseDbWorker, DbTable
Expand Down Expand Up @@ -143,7 +142,7 @@ def compute_fragment_size(cls, table):
fragment_size = HdkFragmentSize.get()
if fragment_size is None:
if cls._preferred_device == "CPU":
cpu_count = os.cpu_count()
cpu_count = CpuCount.get()
if cpu_count is not None:
fragment_size = table.num_rows // cpu_count
fragment_size = min(fragment_size, 2**25)
Expand Down

0 comments on commit 4531ae1

Please sign in to comment.