Use fork for Linux burden workers
This commit is contained in:
@@ -228,6 +228,50 @@ def _eligible_landmark_rows(
|
||||
return rows
|
||||
|
||||
|
||||
def _row_to_worker_spec(row: dict[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
"patient_id": int(row["patient_id"]),
|
||||
"dataset_index": int(row["dataset_index"]),
|
||||
"landmark_age": float(row["landmark_age"]),
|
||||
"t_query": float(row["t_query"]),
|
||||
"followup_end_time": float(row["followup_end_time"]),
|
||||
}
|
||||
|
||||
|
||||
def _materialize_worker_rows(
|
||||
dataset: Any,
|
||||
row_specs: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
rows: list[dict[str, Any]] = []
|
||||
for spec in row_specs:
|
||||
sample = dataset.samples[int(spec["dataset_index"])]
|
||||
seq_event = np.asarray(sample["event_seq"], dtype=np.int64)
|
||||
seq_time = np.asarray(sample["time_seq"], dtype=np.float32)
|
||||
tgt_event = np.asarray(sample["target_event_seq"], dtype=np.int64)
|
||||
tgt_time = np.asarray(sample["target_time_seq"], dtype=np.float32)
|
||||
full_event = np.concatenate([seq_event, tgt_event[-1:]])
|
||||
full_time = np.concatenate([seq_time, tgt_time[-1:]])
|
||||
t_query = np.float32(float(spec["t_query"]))
|
||||
prefix_mask = full_time <= t_query
|
||||
rows.append(
|
||||
{
|
||||
"patient_id": int(spec["patient_id"]),
|
||||
"dataset_index": int(spec["dataset_index"]),
|
||||
"sex": int(sample["sex"]),
|
||||
"landmark_age": np.float32(float(spec["landmark_age"])),
|
||||
"t_query": t_query,
|
||||
"followup_end_time": np.float32(float(spec["followup_end_time"])),
|
||||
"event_seq": full_event[prefix_mask].astype(np.int64, copy=False),
|
||||
"time_seq": full_time[prefix_mask].astype(np.float32, copy=False),
|
||||
"other_type": np.asarray(sample["other_type"], dtype=np.int64),
|
||||
"other_value": np.asarray(sample["other_value"], dtype=np.float32),
|
||||
"other_value_kind": np.asarray(sample["other_value_kind"], dtype=np.int64),
|
||||
"other_time": np.asarray(sample["other_time"], dtype=np.float32),
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def _config_split_indices(
|
||||
n: int,
|
||||
cfg: dict[str, Any],
|
||||
@@ -620,9 +664,25 @@ def _compute_bi_from_readout_table(
|
||||
def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
device = payload["device"]
|
||||
run_path = Path(payload["run_path"])
|
||||
print(
|
||||
f"[BI worker {device}] starting with {len(payload['row_specs'])} rows",
|
||||
flush=True,
|
||||
)
|
||||
load_start = time.perf_counter()
|
||||
ctx = load_burden_context(run_path, device=device)
|
||||
print(
|
||||
f"[BI worker {device}] context loaded in {time.perf_counter() - load_start:.2f}s",
|
||||
flush=True,
|
||||
)
|
||||
materialize_start = time.perf_counter()
|
||||
rows = _materialize_worker_rows(ctx.dataset, payload["row_specs"])
|
||||
print(
|
||||
f"[BI worker {device}] rows materialized in "
|
||||
f"{time.perf_counter() - materialize_start:.2f}s",
|
||||
flush=True,
|
||||
)
|
||||
out, readout_jobs, timings = _compute_bi_from_readout_table(
|
||||
rows=payload["rows"],
|
||||
rows=rows,
|
||||
horizon=payload["horizon"],
|
||||
matrices=payload["matrices"],
|
||||
union_disease_ids=payload["union_disease_ids"],
|
||||
@@ -630,6 +690,13 @@ def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
readout_batch_size=int(payload["readout_batch_size"]),
|
||||
ctx=ctx,
|
||||
)
|
||||
print(
|
||||
f"[BI worker {device}] done: readout_jobs={readout_jobs}, "
|
||||
f"build={timings['build_readout_sec']:.2f}s, "
|
||||
f"forward={timings['forward_sec']:.2f}s, "
|
||||
f"reduce={timings['reduce_sec']:.2f}s",
|
||||
flush=True,
|
||||
)
|
||||
return {"rows": out, "readout_jobs": readout_jobs, "timings": timings}
|
||||
|
||||
|
||||
@@ -774,6 +841,13 @@ def main() -> None:
|
||||
"'cuda:0,cuda:1'. Overrides --device when provided."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mp_start_method",
|
||||
type=str,
|
||||
default="fork",
|
||||
choices=["fork", "forkserver"],
|
||||
help="Multiprocessing start method for Linux multi-GPU runs.",
|
||||
)
|
||||
parser.add_argument("--functional_weight_col", type=str,
|
||||
default="hfrm_normalized_weight")
|
||||
args = parser.parse_args()
|
||||
@@ -845,6 +919,19 @@ def main() -> None:
|
||||
formed_mode=args.formed_mode,
|
||||
horizon=horizon,
|
||||
)
|
||||
for device, chunk_rows in row_chunks:
|
||||
estimated_jobs = sum(
|
||||
_estimate_readout_jobs_for_row(
|
||||
row,
|
||||
formed_mode=args.formed_mode,
|
||||
horizon=horizon,
|
||||
)
|
||||
for row in chunk_rows
|
||||
)
|
||||
print(
|
||||
f"Assigned {len(chunk_rows)} rows / ~{estimated_jobs} readout jobs to {device}",
|
||||
flush=True,
|
||||
)
|
||||
if len(row_chunks) == 1:
|
||||
all_rows, total_readout_jobs, timings = _compute_bi_from_readout_table(
|
||||
rows=rows,
|
||||
@@ -865,7 +952,7 @@ def main() -> None:
|
||||
{
|
||||
"device": device,
|
||||
"run_path": str(run_path),
|
||||
"rows": chunk_rows,
|
||||
"row_specs": [_row_to_worker_spec(row) for row in chunk_rows],
|
||||
"horizon": horizon,
|
||||
"matrices": matrices,
|
||||
"union_disease_ids": union_disease_ids,
|
||||
@@ -876,7 +963,7 @@ def main() -> None:
|
||||
]
|
||||
with ProcessPoolExecutor(
|
||||
max_workers=len(payloads),
|
||||
mp_context=mp.get_context("spawn"),
|
||||
mp_context=mp.get_context(args.mp_start_method),
|
||||
) as executor:
|
||||
futures = [executor.submit(_compute_chunk_worker, p) for p in payloads]
|
||||
for future in tqdm(
|
||||
@@ -901,6 +988,7 @@ def main() -> None:
|
||||
print(f"Landmark rows: {len(rows)}")
|
||||
print(f"Readout jobs: {total_readout_jobs}")
|
||||
print(f"Readout batch size per worker: {int(args.readout_batch_size)}")
|
||||
print(f"Multiprocessing start method: {args.mp_start_method}")
|
||||
print(
|
||||
"Timing seconds: "
|
||||
f"build_readout={total_timings['build_readout_sec']:.2f}, "
|
||||
|
||||
Reference in New Issue
Block a user