Avoid multiprocessing data loader for shape export

export_weibull_death_parameter_stats.py

@@ -18,6 +18,7 @@ import numpy as np
 import pandas as pd
 import torch
 import torch.nn.functional as F
+import torch.multiprocessing as torch_mp
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
@@ -40,6 +41,11 @@ from evaluate_auc_v2 import (
     validate_dataset_metadata,
 )
 
+try:
+    torch_mp.set_sharing_strategy("file_system")
+except RuntimeError:
+    pass
+
 
 def quantile_summary(df: pd.DataFrame, group_cols: List[str], value_cols: List[str]) -> pd.DataFrame:
     probs = [0.01, 0.05, 0.25, 0.50, 0.75, 0.95, 0.99]
@@ -293,7 +299,15 @@ def main() -> None:
     parser.add_argument("--landmark_step", type=float, default=None)
     parser.add_argument("--horizons", type=str, default=None)
     parser.add_argument("--batch_size", type=int, default=None)
-    parser.add_argument("--num_workers", type=int, default=None)
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=0,
+        help=(
+            "DataLoader workers. Default 0 avoids Linux multiprocessing "
+            "'received 0 items of ancdata' failures on shared filesystems."
+        ),
+    )
     parser.add_argument("--device", type=str, default=None)
     parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None)
     parser.add_argument("--hidden_cache_dtype", type=str, default="float32", choices=["float16", "float32"])
@@ -386,17 +400,18 @@ def main() -> None:
     )
 
     batch_size = int(cfg_get(args, cfg, "batch_size", 128))
-    num_workers = int(cfg_get(args, cfg, "num_workers", 4))
-    loader = DataLoader(
-        landmark_dataset,
-        batch_size=batch_size,
-        shuffle=False,
-        collate_fn=collate_landmark_fn,
-        num_workers=num_workers,
-        pin_memory=device.type == "cuda",
-        persistent_workers=num_workers > 0,
-        prefetch_factor=2 if num_workers > 0 else None,
-    )
+    num_workers = int(cfg_get(args, cfg, "num_workers", 0))
+    loader_kwargs = {
+        "batch_size": batch_size,
+        "shuffle": False,
+        "collate_fn": collate_landmark_fn,
+        "num_workers": num_workers,
+        "pin_memory": device.type == "cuda",
+    }
+    if num_workers > 0:
+        loader_kwargs["persistent_workers"] = True
+        loader_kwargs["prefetch_factor"] = 2
+    loader = DataLoader(landmark_dataset, **loader_kwargs)
 
     use_amp = bool(cfg_get(args, cfg, "use_amp", False))
     hidden_all, row_arrays = infer_landmark_hidden_local(