Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
4 changes: 3 additions & 1 deletion oc_meta/core/curator.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,9 +176,11 @@ def _collect_identifiers_with_progress(self, task_id=None):
for i in range(0, total_rows, self.min_rows_parallel):
chunks.append(self.data[i : i + self.min_rows_parallel])

mp_method = "spawn" if os.name == "nt" else "forkserver"

with ProcessPoolExecutor(
max_workers=self.workers,
mp_context=multiprocessing.get_context("forkserver"),
mp_context=multiprocessing.get_context(mp_method),
) as executor:
for chunk_metavals, chunk_ids, chunk_vvis in executor.map(
_extract_ids_from_chunk, chunks
Expand Down
10 changes: 7 additions & 3 deletions oc_meta/lib/finder.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from __future__ import annotations

import multiprocessing
import os
from concurrent.futures import ProcessPoolExecutor
from functools import partial
from typing import TYPE_CHECKING, Dict, List, Tuple, TypedDict
Expand Down Expand Up @@ -678,9 +679,10 @@ def process_batch_parallel(subjects, cur_depth, visited_subjects):
for i in range(0, len(batch_queries), queries_per_worker)
]
worker = partial(execute_sparql_queries, ts_url)
mp_method = "spawn" if os.name == "nt" else "forkserver"
with ProcessPoolExecutor(
max_workers=min(len(query_groups), MAX_WORKERS),
mp_context=multiprocessing.get_context("forkserver"),
mp_context=multiprocessing.get_context(mp_method),
) as executor:
grouped_results = list(executor.map(worker, query_groups))
results = [item for sublist in grouped_results for item in sublist]
Expand Down Expand Up @@ -795,9 +797,10 @@ def get_initial_subjects_from_identifiers(identifiers, progress_task=None):
sum(batch_sizes[i : i + QLEVER_QUERIES_PER_GROUP])
)
worker = partial(execute_sparql_queries, ts_url)
mp_method = "spawn" if os.name == "nt" else "forkserver"
with ProcessPoolExecutor(
max_workers=MAX_WORKERS,
mp_context=multiprocessing.get_context("forkserver"),
mp_context=multiprocessing.get_context(mp_method),
) as executor:
results = []
for idx, grouped_result in enumerate(
Expand Down Expand Up @@ -973,9 +976,10 @@ def get_initial_subjects_from_vvis(vvis, progress_task=None):
vvi_count = int(total_vvis * len(group) / len(vvi_queries))
grouped_vvi_counts.append(max(1, vvi_count))
worker = partial(execute_sparql_queries, ts_url)
mp_method = "spawn" if os.name == "nt" else "forkserver"
with ProcessPoolExecutor(
max_workers=MAX_WORKERS,
mp_context=multiprocessing.get_context("forkserver"),
mp_context=multiprocessing.get_context(mp_method),
) as executor:
results = []
for idx, grouped_result in enumerate(
Expand Down
4 changes: 2 additions & 2 deletions oc_meta/lib/rdf_patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
from oc_meta.lib.file_manager import find_rdf_file
from oc_meta.run.meta.generate_csv import load_json_from_file

_forkserver_context = multiprocessing.get_context("forkserver")
mp_method = multiprocessing.get_context("forkserver") if os.name != "nt" else multiprocessing.get_context("spawn")

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah this was the issue causing tests to fail, I didn't notice that there's multiprocessing.get_context() outside of the ProcessPoolExecutor instead of inside like everywhere else


HAS_IDENTIFIER = "http://purl.org/spar/datacite/hasIdentifier"
USES_IDENTIFIER_SCHEME = "http://purl.org/spar/datacite/usesIdentifierScheme"
Expand Down Expand Up @@ -201,7 +201,7 @@ def load_available_entities(
result = {}
with ProcessPoolExecutor(
max_workers=workers,
mp_context=_forkserver_context,
mp_context=mp_method,
) as executor:
for partial in executor.map(_load_target_batch, batches(tasks, 24)):
result.update(partial)
Expand Down
5 changes: 4 additions & 1 deletion oc_meta/lib/sparql.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from __future__ import annotations

import multiprocessing
import os
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from typing import Callable
Expand Down Expand Up @@ -131,9 +132,11 @@ def run_queries_parallel(
query_groups.append(batch_queries[i : i + QLEVER_QUERIES_PER_GROUP])
grouped_sizes.append(sum(batch_sizes[i : i + QLEVER_QUERIES_PER_GROUP]))

mp_method = "spawn" if os.name == "nt" else "forkserver"

with ProcessPoolExecutor(
max_workers=min(len(query_groups), workers),
mp_context=multiprocessing.get_context("forkserver"),
mp_context=multiprocessing.get_context(mp_method),
) as executor:
future_to_size = {
executor.submit(
Expand Down
3 changes: 2 additions & 1 deletion oc_meta/run/count/meta_entities.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,10 +130,11 @@ def count_venues_from_csv(self) -> int:
task = progress.add_task(
"Counting venues from CSV files...", total=len(filepaths)
)
mp_method = "spawn" if os.name == "nt" else "forkserver"

# Use forkserver to avoid deadlocks when forking in a multi-threaded environment
with ProcessPoolExecutor(
mp_context=multiprocessing.get_context("forkserver")
mp_context=multiprocessing.get_context(mp_method)
) as executor:
futures = {
executor.submit(_count_venues_in_file, fp): fp for fp in filepaths
Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/count/triples.py
Original file line number Diff line number Diff line change
Expand Up @@ -268,7 +268,7 @@ def process_files(
task = progress.add_task(f"Counting {unit_name}", total=len(files))

# Use forkserver to avoid deadlocks when forking in a multi-threaded environment
ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")
with ctx.Pool(processes=workers) as pool:
for file_path, count, error in pool.imap_unordered(
worker_fn, files, chunksize=chunksize
Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/find/duplicates.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ def process_identifier_chunk(
entity_info: defaultdict[IdentifierKey, set[str]] = defaultdict(set)

# Use forkserver to avoid deadlocks when forking in a multi-threaded environment.
ctx = mp.get_context("forkserver")
ctx = mp.get_context("spawn") if os.name == "nt" else mp.get_context("forkserver")
with ctx.Pool(processes=mp.cpu_count()) as pool:
results = pool.map(process_identifier_zip_file, zip_files_chunk)

Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/find/hasnext_anomalies.py
Original file line number Diff line number Diff line change
Expand Up @@ -309,7 +309,7 @@ def main() -> None:
total_brs = 0
all_anomalies: List[dict] = []

ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")
with ctx.Pool(
args.workers,
_init_worker,
Expand Down
3 changes: 2 additions & 1 deletion oc_meta/run/find/merged_entities.py
Original file line number Diff line number Diff line change
Expand Up @@ -137,8 +137,9 @@ def main():
all_results: list[tuple[str, str]] = []

# Use forkserver to avoid deadlocks when forking in a multi-threaded environment
mp_method = "spawn" if os.name == "nt" else "forkserver"
with ProcessPoolExecutor(
max_workers=args.workers, mp_context=multiprocessing.get_context("forkserver")
max_workers=args.workers, mp_context=multiprocessing.get_context(mp_method)
) as executor:
futures = {executor.submit(process_prov_file, f): f for f in prov_files}

Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/infodir/_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -164,7 +164,7 @@ def bounded_process_map(
if workers <= 0:
raise ValueError("workers must be greater than zero")
path_iterator = iter(paths)
context = multiprocessing.get_context("forkserver")
context = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")
pending: dict[Future[ResultType], str] = {}
with ProcessPoolExecutor(max_workers=workers, mp_context=context) as executor:
for path in islice(path_iterator, workers * 2):
Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/merge/check_merged_brs_results.py
Original file line number Diff line number Diff line change
Expand Up @@ -546,7 +546,7 @@ def main():
csv_files = [f for f in os.listdir(args.csv_folder) if f.endswith(".csv")]

# Use forkserver to avoid deadlocks when forking in a multi-threaded environment
ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")

# Process CSV files to gather tasks
with ctx.Pool(processes=multiprocessing.cpu_count()) as pool:
Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/merge/check_merged_ids_results.py
Original file line number Diff line number Diff line change
Expand Up @@ -475,7 +475,7 @@ def main():
csv_files = [f for f in os.listdir(args.csv_folder) if f.endswith(".csv")]

# Use forkserver to avoid deadlocks when forking in a multi-threaded environment
ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")

# Process CSV files in parallel
with ctx.Pool(processes=multiprocessing.cpu_count()) as pool:
Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/merge/check_merged_ras_results.py
Original file line number Diff line number Diff line change
Expand Up @@ -515,7 +515,7 @@ def main():
csv_files = [f for f in os.listdir(args.csv_folder) if f.endswith(".csv")]

# Use forkserver to avoid deadlocks when forking in a multi-threaded environment
ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")

# Process CSV files to gather tasks
with ctx.Pool(processes=multiprocessing.cpu_count()) as pool:
Expand Down
3 changes: 2 additions & 1 deletion oc_meta/run/meta/check_rdf_files.py
Original file line number Diff line number Diff line change
Expand Up @@ -627,11 +627,12 @@ def main() -> None:
index: Optional[dict[str, str]] = {} if args.input_csv else None

console.print(f"Checking every row of {os.path.basename(args.csv)}...")
mp_method = "spawn" if os.name == "nt" else "forkserver"
with ProcessPoolExecutor(
max_workers=args.workers,
initializer=_init_worker,
initargs=config,
mp_context=multiprocessing.get_context("forkserver"),
mp_context=multiprocessing.get_context(mp_method),
) as pool:
aggregated = _drive(
pool,
Expand Down
6 changes: 4 additions & 2 deletions oc_meta/run/meta/check_results.py
Original file line number Diff line number Diff line change
Expand Up @@ -480,8 +480,9 @@ def on_id_batch(batch_size: int):
prov_future = None
prov_executor = None
if total_omids > 0:
mp_method = "spawn" if os.name == "nt" else "forkserver"
prov_executor = ProcessPoolExecutor(
max_workers=1, mp_context=multiprocessing.get_context("forkserver")
max_workers=1, mp_context=multiprocessing.get_context(mp_method)
)
prov_future = prov_executor.submit(
check_provenance_existence,
Expand Down Expand Up @@ -510,9 +511,10 @@ def _apply_zip_results(zip_results: dict[str, tuple[bool, bool]]) -> None:
result.prov_graphs_missing += 1

if zip_args and workers > 1:
mp_method = "spawn" if os.name == "nt" else "forkserver"
with ProcessPoolExecutor(
max_workers=min(len(zip_args), workers),
mp_context=multiprocessing.get_context("forkserver"),
mp_context=multiprocessing.get_context(mp_method),
) as executor:
for future in as_completed(
{executor.submit(_check_zip_file, a): a for a in zip_args}
Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/meta/generate_csv.py
Original file line number Diff line number Diff line change
Expand Up @@ -636,7 +636,7 @@ def generate_csv(
result_buffer = ResultBuffer(output_dir)

# Use forkserver to avoid deadlocks when forking in a multi-threaded environment
ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")
with ctx.Pool(
workers,
_init_worker,
Expand Down
6 changes: 4 additions & 2 deletions oc_meta/run/meta/preprocess_input.py
Original file line number Diff line number Diff line change
Expand Up @@ -364,9 +364,10 @@ def main(): # pragma: no cover
results: list[FileResult] = []
with create_progress() as progress:
task = progress.add_task("Filtering existing IDs", total=len(csv_files))
mp_method = "spawn" if os.name == "nt" else "forkserver"
with ProcessPoolExecutor(
max_workers=args.workers,
mp_context=multiprocessing.get_context("forkserver"),
mp_context=multiprocessing.get_context(mp_method),
) as executor:
futures = {
executor.submit(
Expand All @@ -385,9 +386,10 @@ def main(): # pragma: no cover
results = []
with create_progress() as progress:
task = progress.add_task("Reading CSV files", total=len(csv_files))
mp_method = "spawn" if os.name == "nt" else "forkserver"
with ProcessPoolExecutor(
max_workers=args.workers,
mp_context=multiprocessing.get_context("forkserver"),
mp_context=multiprocessing.get_context(mp_method),
) as executor:
futures = {
executor.submit(collect_rows_from_file, f): f for f in csv_files
Expand Down
4 changes: 2 additions & 2 deletions oc_meta/run/meta_process.py
Original file line number Diff line number Diff line change
Expand Up @@ -351,7 +351,7 @@ def _upload_sparql_queries(self) -> None:
# Use forkserver to avoid deadlocks when forking from a multi-threaded process.
# Libraries like Redis and rdflib create background threads, and fork() would
# copy locked mutexes into the child process, causing hangs.
ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")

data_process = ctx.Process(
target=_upload_to_triplestore,
Expand Down Expand Up @@ -410,7 +410,7 @@ def _store_and_upload(
# Use forkserver to avoid deadlocks when forking from a multi-threaded process.
# Libraries like rdflib create background threads, and fork() would
# copy locked mutexes into the child process, causing hangs.
ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")

data_store_process = ctx.Process(
target=_store_rdf_worker,
Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/migration/rdf_to_nquads.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,7 @@ def main() -> None: # pragma: no cover
)

# Use forkserver to avoid deadlocks when forking in a multi-threaded environment
ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")
with ProcessPoolExecutor(max_workers=num_workers, mp_context=ctx) as executor:
iterator = executor.map(task_func, zip_files)

Expand Down
3 changes: 2 additions & 1 deletion oc_meta/run/migration/stream_nquads.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import argparse
import gzip
import multiprocessing
import os
import sys
import zipfile
from collections.abc import Iterable, Iterator
Expand Down Expand Up @@ -208,7 +209,7 @@ def main() -> None: # pragma: no cover
only_prov=args.mode == "prov",
)

ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")
with ctx.Pool(processes=num_workers) as pool:
if output_dir:
output_dir.mkdir(parents=True, exist_ok=True)
Expand Down
6 changes: 3 additions & 3 deletions oc_meta/run/patches/fix_dangling_ars.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,8 @@
_stop_requested = False
_existing_roles: frozenset[str] = frozenset()
_target_roles: frozenset[str] = frozenset()
_fork_context = multiprocessing.get_context("fork")
_forkserver_context = multiprocessing.get_context("forkserver")
_fork_context = multiprocessing.get_context("fork") if os.name != "nt" else multiprocessing.get_context("spawn")
_context = multiprocessing.get_context("forkserver") if os.name != "nt" else multiprocessing.get_context("spawn")


@dataclass(frozen=True, slots=True)
Expand Down Expand Up @@ -311,7 +311,7 @@ def load_provenance_statuses(
tasks = [(path, frozenset(targets)) for path, targets in targets_by_path.items()]
statuses = {}
with ProcessPoolExecutor(
max_workers=workers, mp_context=_forkserver_context
max_workers=workers, mp_context=_context
) as executor:
for partial in executor.map(_provenance_status_batch, _batches(tasks, 24)):
statuses.update(partial)
Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/patches/fix_duplicate_part_of.py
Original file line number Diff line number Diff line change
Expand Up @@ -473,7 +473,7 @@ def check_orphans(
br_files[i : i + batch_size] for i in range(0, len(br_files), batch_size)
]

ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")
with create_progress() as progress:
task = progress.add_task("Checking orphans", total=len(br_files))
executor = ProcessPoolExecutor(
Expand Down
4 changes: 3 additions & 1 deletion oc_meta/run/patches/fix_literal_datatypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from __future__ import annotations

import argparse
import os
import signal
import shutil
import zipfile
Expand Down Expand Up @@ -195,10 +196,11 @@ def main() -> None: # pragma: no cover
]

# Use forkserver to avoid deadlocks when forking in a multi-threaded environment
mp_method = "spawn" if os.name == "nt" else "forkserver"
executor = ProcessPoolExecutor(
max_workers=args.workers,
initializer=_worker_init,
mp_context=multiprocessing.get_context("forkserver"),
mp_context=multiprocessing.get_context(mp_method),
)
try:
with create_progress() as progress:
Expand Down
2 changes: 1 addition & 1 deletion oc_meta/run/patches/fix_misplaced_editor_ars.py
Original file line number Diff line number Diff line change
Expand Up @@ -286,7 +286,7 @@ def find_misplaced_editor_ars(
frbr_part_of: dict[str, list[str]] = {}
content_ars: dict[str, set[str]] = {}

ctx = multiprocessing.get_context("forkserver")
ctx = multiprocessing.get_context("spawn") if os.name == "nt" else multiprocessing.get_context("forkserver")

with create_progress() as progress:
br_task = progress.add_task("Scanning BR files", total=len(br_files))
Expand Down
Loading