Skip to content
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.

Sign up for GitHub

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

Receive a customized parameter to specify the number of concurrent running threads for the code_* operations. #195

Merged
merged 2 commits into from
Nov 19, 2024
Merged

Receive a customized parameter to specify the number of concurrent running threads for the code_* operations. #195

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
241af17
chore: added concurrent threads count parameter
lapsule Nov 19, 2024
cbec912
Merge branch 'release/concurrent_count_config'
lapsule Nov 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions docetl/operations/code_operations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import os
from typing import Any, Dict, List, Optional, Tuple
from concurrent.futures import ThreadPoolExecutor
from docetl.operations.base import BaseOperation
Expand All @@ -7,6 +8,7 @@ class CodeMapOperation(BaseOperation):
class schema(BaseOperation.schema):
type: str = "code_map"
code: str
concurrent_thread_count: int = os.cpu_count()
drop_keys: Optional[List[str]] = None

def syntax_check(self) -> None:
Expand All @@ -27,7 +29,7 @@ def execute(self, input_data: List[Dict]) -> Tuple[List[Dict], float]:
transform_fn = namespace["transform"]

results = []
with ThreadPoolExecutor() as executor:
with ThreadPoolExecutor(max_workers=self.config.get('concurrent_thread_count', os.cpu_count())) as executor:
futures = [executor.submit(transform_fn, doc) for doc in input_data]
pbar = RichLoopBar(
range(len(futures)),
Expand All @@ -51,6 +53,7 @@ class CodeReduceOperation(BaseOperation):
class schema(BaseOperation.schema):
type: str = "code_reduce"
code: str
concurrent_thread_count: int = os.cpu_count()

def syntax_check(self) -> None:
config = self.schema(**self.config)
Expand Down Expand Up @@ -89,7 +92,7 @@ def get_group_key(item):
grouped_data = list(grouped_data.items())

results = []
with ThreadPoolExecutor() as executor:
with ThreadPoolExecutor(max_workers=self.config.get('concurrent_thread_count', os.cpu_count())) as executor:
futures = [executor.submit(reduce_fn, group) for _, group in grouped_data]
pbar = RichLoopBar(
range(len(futures)),
Expand All @@ -113,6 +116,7 @@ class CodeFilterOperation(BaseOperation):
class schema(BaseOperation.schema):
type: str = "code_filter"
code: str
concurrent_thread_count: int = os.cpu_count()

def syntax_check(self) -> None:
config = self.schema(**self.config)
Expand All @@ -132,7 +136,7 @@ def execute(self, input_data: List[Dict]) -> Tuple[List[Dict], float]:
filter_fn = namespace["transform"]

results = []
with ThreadPoolExecutor() as executor:
with ThreadPoolExecutor(max_workers=self.config.get('concurrent_thread_count', os.cpu_count())) as executor:
futures = [executor.submit(filter_fn, doc) for doc in input_data]
pbar = RichLoopBar(
range(len(futures)),
Expand Down
Loading