Update TabArena EC2 runner scripts (run_tabarena.py)

benchmark/tabarena/run_tabarena.py

@@ -29,6 +29,8 @@ def main():
     ap.add_argument("--quickstart", action="store_true", default=True)
     ap.add_argument("--subset", default=None, help="tabarena subset, e.g. 'lite'")
     ap.add_argument("--full", action="store_true")
+    ap.add_argument("--datasets", default=None,
+                    help="comma-separated dataset names (implies lite subset, split 0)")
     ap.add_argument("--run-name", default="felatab_tabarena")
     args = ap.parse_args()
 
@@ -39,15 +41,25 @@ def main():
     models = [(FelaTabAGModel.config_generator(), 0)]  # 0 = default/manual configs only
     if WITH_LGBM:
         models.append(("LightGBM", 0))  # cross-check against a known registry model
-    experiments = TabArenaV0pt1ExperimentBundle(models=models).build_experiments()
+    # AutoGluon pickles one model per CV fold (FelaTab big tier ~1.6 GB fp32); keep
+    # artifacts off tmpfs /tmp. Honors TMPDIR, overridable via TABARENA_ARTIFACTS_DIR.
+    artifacts_dir = os.environ.get("TABARENA_ARTIFACTS_DIR", os.environ.get("TMPDIR"))
+    experiments = TabArenaV0pt1ExperimentBundle(
+        models=models, model_artifacts_base_path=artifacts_dir,
+    ).build_experiments()
 
     build_kwargs = {}
     subset = args.subset
     if args.full:
         subset = None
+    elif args.datasets:
+        subset = "lite"
     elif not subset:
         subset = "lite"
         build_kwargs["dataset_names"] = QUICKSTART_DATASETS
+    if args.datasets:
+        # applies in both lite and full mode; lets us shard the full suite across machines
+        build_kwargs["dataset_names"] = args.datasets.split(",")
 
     context = TabArenaContext()
     context.build_and_run_jobs(