Skip to content

Commit

Permalink
View and Set N Jobs (#1029)
Browse files Browse the repository at this point in the history
* [Style] isort preference applied

* [Style] isort preference applied

* reset n_jobs to default after testing

* simplify imports
  • Loading branch information
DhruvSrikanth committed Aug 3, 2023
1 parent ac413b6 commit f3a6a7c
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 0 deletions.
14 changes: 14 additions & 0 deletions tests/units/utilities/test_distribution.py
Expand Up @@ -16,9 +16,23 @@
LocalDaskDistributor,
MultiprocessingDistributor,
)
from tsfresh.utilities.profiling import get_n_jobs, set_n_jobs


class MultiprocessingDistributorTestCase(TestCase):
def test_n_jobs(self):
curr = get_n_jobs()
self.assertEqual(curr, get_n_jobs())

set_n_jobs(2)
self.assertEqual(2, get_n_jobs())

set_n_jobs(4)
self.assertEqual(4, get_n_jobs())

set_n_jobs(curr)
self.assertEqual(curr, get_n_jobs())

def test_partition(self):

distributor = MultiprocessingDistributor(n_workers=1)
Expand Down
25 changes: 25 additions & 0 deletions tsfresh/utilities/profiling.py
Expand Up @@ -10,6 +10,8 @@
import logging
import pstats

from tsfresh import defaults

_logger = logging.getLogger(__name__)


Expand Down Expand Up @@ -66,3 +68,26 @@ def end_profiling(profiler, filename, sorting=None):
"[calculate_ts_features] Finished profiling of time series feature extraction"
)
f.write(s.getvalue())


def get_n_jobs():
"""
Get the number of jobs to use for parallel processing.
:return: The number of jobs to use for parallel processing.
:rtype: int
"""
return defaults.N_PROCESSES


def set_n_jobs(n_jobs):
"""
Set the number of jobs to use for parallel processing.
:param n_jobs: The number of jobs to use for parallel processing.
:type n_jobs: int
:return: None
:rtype: None
"""
defaults.N_PROCESSES = n_jobs

0 comments on commit f3a6a7c

Please sign in to comment.