nemo-evaluator-launcher 0.1.12__py3-none-any.whl → 0.1.14__py3-none-any.whl

nemo_evaluator_launcher/exporters/wandb.py

@@ -19,7 +19,7 @@ import os
 import shutil
 import tempfile
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 import yaml
 
@@ -38,6 +38,7 @@ from nemo_evaluator_launcher.exporters.registry import register_exporter
 from nemo_evaluator_launcher.exporters.utils import (
     extract_accuracy_metrics,
     extract_exporter_config,
+    get_artifact_root,
     get_available_artifacts,
     get_benchmark_info,
     get_task_name,
@@ -163,29 +164,92 @@ class WandBExporter(BaseExporter):
             return {"success": False, "error": f"W&B export failed: {str(e)}"}
 
     def _log_artifacts(
-        self, job_data: JobData, wandb_config: Dict[str, Any], artifact
+        self,
+        job_data: JobData,
+        wandb_config: Dict[str, Any],
+        artifact,
+        register_staging_dir=None,
     ) -> List[str]:
-        """Log evaluation artifacts to WandB using LocalExporter for transfer."""
+        """Log evaluation artifacts to WandB using LocalExporter for staging."""
         if not wandb_config.get("log_artifacts", True):
             return []
         try:
             temp_dir = tempfile.mkdtemp(prefix="wandb_artifacts_")
-            local_exporter = LocalExporter({"output_dir": temp_dir})
+            if callable(register_staging_dir):
+                register_staging_dir(temp_dir)
+            local_exporter = LocalExporter(
+                {
+                    "output_dir": temp_dir,
+                    "copy_logs": wandb_config.get(
+                        "log_logs", wandb_config.get("copy_logs", False)
+                    ),
+                    "only_required": wandb_config.get("only_required", True),
+                    "format": wandb_config.get("format"),
+                    "log_metrics": wandb_config.get("log_metrics", []),
+                    "output_filename": wandb_config.get("output_filename"),
+                }
+            )
             local_result = local_exporter.export_job(job_data)
 
             if not local_result.success:
                 logger.error(f"Failed to download artifacts: {local_result.message}")
                 return []
 
-            artifacts_dir = Path(local_result.dest) / "artifacts"
-            logged_names = []
-            task_name = get_task_name(job_data)
-            for fname in get_available_artifacts(artifacts_dir):
-                fpath = artifacts_dir / fname
-                if fpath.exists():
-                    artifact.add_file(str(fpath), name=f"{task_name}/{fname}")
-                    logged_names.append(fname)
-            shutil.rmtree(temp_dir)
+            base_dir = Path(local_result.dest)
+            artifacts_dir = base_dir / "artifacts"
+            logs_dir = base_dir / "logs"
+            logged_names: list[str] = []
+
+            artifact_root = get_artifact_root(job_data)  # "<harness>.<benchmark>"
+
+            # Add config file only when artifacts logging is enabled
+            if wandb_config.get("log_artifacts", True):
+                cfg_added = False
+                for fname in ("config.yml", "run_config.yml"):
+                    p = artifacts_dir / fname
+                    if p.exists():
+                        artifact.add_file(str(p), name=f"{artifact_root}/{fname}")
+                        logged_names.append(fname)
+                        cfg_added = True
+                        break
+                if not cfg_added:
+                    with tempfile.NamedTemporaryFile(
+                        "w", suffix=".yaml", delete=False
+                    ) as tmp_cfg:
+                        yaml.dump(
+                            job_data.config or {},
+                            tmp_cfg,
+                            default_flow_style=False,
+                            sort_keys=False,
+                        )
+                        cfg_path = tmp_cfg.name
+                    artifact.add_file(cfg_path, name=f"{artifact_root}/config.yaml")
+                    os.unlink(cfg_path)
+                    logged_names.append("config.yaml")
+
+            files_to_upload: list[Path] = []
+            if wandb_config.get("only_required", True):
+                for fname in get_available_artifacts(artifacts_dir):
+                    p = artifacts_dir / fname
+                    if p.exists():
+                        files_to_upload.append(p)
+            else:
+                for p in artifacts_dir.iterdir():
+                    if p.is_file():
+                        files_to_upload.append(p)
+
+            for fpath in files_to_upload:
+                rel = fpath.relative_to(artifacts_dir).as_posix()
+                artifact.add_file(str(fpath), name=f"{artifact_root}/artifacts/{rel}")
+                logged_names.append(rel)
+
+            if wandb_config.get("log_logs", False) and logs_dir.exists():
+                for p in logs_dir.rglob("*"):
+                    if p.is_file():
+                        rel = p.relative_to(logs_dir).as_posix()
+                        artifact.add_file(str(p), name=f"{artifact_root}/logs/{rel}")
+                        logged_names.append(f"logs/{rel}")
+
             return logged_names
         except Exception as e:
             logger.error(f"Error logging artifacts: {e}")
@@ -193,7 +257,7 @@ class WandBExporter(BaseExporter):
 
     def _check_existing_run(
         self, identifier: str, job_data: JobData, config: Dict[str, Any]
-    ) -> tuple[bool, str]:
+    ) -> tuple[bool, Optional[str]]:
         """Check if run exists based on webhook metadata then name patterns."""
         try:
             import wandb
@@ -204,7 +268,7 @@ class WandBExporter(BaseExporter):
             if not (entity and project):
                 return False, None
 
-            # # Check webhook metadata for run_id first
+            # Check webhook metadata for run_id first
             webhook_meta = job_data.data.get("webhook_metadata", {})
             if (
                 webhook_meta.get("webhook_source") == "wandb"
@@ -306,6 +370,13 @@ class WandBExporter(BaseExporter):
         # Initialize
         run = wandb.init(**{k: v for k, v in run_args.items() if v is not None})
 
+        # Track staging dirs for this run
+        staging_dirs: List[str] = []
+
+        def register_staging_dir(path: str) -> None:
+            if path and os.path.isdir(path):
+                staging_dirs.append(path)
+
         # In multi_task, aggregate lists after init (no overwrite)
         if log_mode == "multi_task":
             try:
@@ -339,34 +410,42 @@ class WandBExporter(BaseExporter):
                 "harness": harness,
             },
         )
-        with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as tmp_cfg:
-            yaml.dump(job_data.config or {}, tmp_cfg, default_flow_style=False)
-            cfg_path = tmp_cfg.name
-        artifact.add_file(cfg_path, name="config.yaml")
-        os.unlink(cfg_path)
 
-        logged_artifacts = self._log_artifacts(job_data, config, artifact)
-        run.log_artifact(artifact)
+        logged_artifacts = self._log_artifacts(
+            job_data, config, artifact, register_staging_dir=register_staging_dir
+        )
 
-        # charts for each logged metric
         try:
-            for k in metrics.keys():
-                run.define_metric(k, summary="last")
-        except Exception:
-            pass
+            run.log_artifact(artifact)
+            # charts for each logged metric
+            try:
+                for k in metrics.keys():
+                    run.define_metric(k, summary="last")
+            except Exception:
+                pass
 
-        # Log metrics with per-task step
-        try:
-            step_idx = int(job_data.job_id.split(".")[-1])
-        except Exception:
-            step_idx = 0
-        run.log(metrics, step=step_idx)
+            # Log metrics with per-task step
+            try:
+                step_idx = int(job_data.job_id.split(".")[-1])
+            except Exception:
+                step_idx = 0
+            run.log(metrics, step=step_idx)
 
-        # metrics summary
-        try:
-            run.summary.update(metrics)
-        except Exception:
-            pass
+            # metrics summary
+            try:
+                run.summary.update(metrics)
+            except Exception:
+                pass
+        finally:
+            for d in staging_dirs:
+                try:
+                    shutil.rmtree(d, ignore_errors=True)
+                except Exception:
+                    pass
+            try:
+                run.finish()
+            except Exception:
+                pass
 
         return {
             "run_id": run.id,

nemo_evaluator_launcher/package_info.py

@@ -16,7 +16,7 @@
 # Below is the _next_ version that will be published, not the currently published one.
 MAJOR = 0
 MINOR = 1
-PATCH = 12
+PATCH = 14
 PRE_RELEASE = ""
 
 # Use the following formatting: (major, minor, patch, pre-release)

{nemo_evaluator_launcher-0.1.12.dist-info → nemo_evaluator_launcher-0.1.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nemo-evaluator-launcher
-Version: 0.1.12
+Version: 0.1.14
 Summary: Launcher for the evaluations provided by NeMo Evaluator containers with different runtime backends
 Author: NVIDIA
 Author-email: nemo-toolkit@nvidia.com

{nemo_evaluator_launcher-0.1.12.dist-info → nemo_evaluator_launcher-0.1.14.dist-info}/RECORD

@@ -1,17 +1,17 @@
 nemo_evaluator_launcher/__init__.py,sha256=2F703fttLaIyMHoVD54rptHMXt4AWnplHDrwWJ3e3PM,1930
-nemo_evaluator_launcher/package_info.py,sha256=iWHbqOvfmglxSklsH5FIsiv5MtRx7ib_9fmE3Wj6hcc,1586
+nemo_evaluator_launcher/package_info.py,sha256=mb9aX6B6nFEfIArIwK-S6yhj2ORaUTmNybkpiR6ltDw,1586
 nemo_evaluator_launcher/api/__init__.py,sha256=U9q_MJK2vRsFaymanhyy0nD1SNAZQZC8oY45RXPX7ac,1024
-nemo_evaluator_launcher/api/functional.py,sha256=6QuCVOkox4SAxtNQhEBzzsMTHaTwHCTLFRqz6eC0gDY,27053
+nemo_evaluator_launcher/api/functional.py,sha256=NyALslqIaFVqZM2eZlEKwp_8huOffd-yzmOLR519cBQ,28448
 nemo_evaluator_launcher/api/types.py,sha256=RXr_QoKdhejj1T9-HybSjd4KTxJmSv0bE0uLUFtF7Zc,3269
 nemo_evaluator_launcher/api/utils.py,sha256=q5HArRj7PKgBfeH3bOX8q1U97yMyQQp72yRRA5JP9PE,818
 nemo_evaluator_launcher/cli/__init__.py,sha256=lNC_skFLYTOt-arnY3ZQnZMWzHlrtD2wAoHvDcHddwM,673
-nemo_evaluator_launcher/cli/export.py,sha256=32DwyG6oM3MZQh0otY2T6BYZTp2aMPOzvUmyEFVKAp8,5733
+nemo_evaluator_launcher/cli/export.py,sha256=GRXxusKDq_1qjMKN6MKOIjZ8x4u5ERgXwHSAGrvsGCY,11211
 nemo_evaluator_launcher/cli/kill.py,sha256=C-4PWmMu8mIITo92o5AHxtq_s-8Cckbp7wAlG0I_ylw,1323
 nemo_evaluator_launcher/cli/ls_runs.py,sha256=5QGUofeKK6tNtQnUF5wJ5YBZ5TL7RRueKlPkvA1wclY,3873
 nemo_evaluator_launcher/cli/ls_tasks.py,sha256=Pd2lBQOQBNHBWrjk4tZg0SQ9Ul9F2Ak-zOyh-G9x-DY,5293
 nemo_evaluator_launcher/cli/main.py,sha256=Wu_AB_6LJL2tvbl-9jw6dPYQ_uvgYivlyM2E9soHvMc,6649
-nemo_evaluator_launcher/cli/run.py,sha256=qNWbwnm0R5YJ00WcHWvSl-Yr3iEQJpMtrEsu2Xi3VDs,5843
-nemo_evaluator_launcher/cli/status.py,sha256=QCZfXQVnQQKNyyVWJtQd1Xe9uhl5YebfhSeT3TmYPmw,6066
+nemo_evaluator_launcher/cli/run.py,sha256=h5DZM3agAV7FSOLt3evPgt5kx9dX0i1cPaeqfxSP-P4,6608
+nemo_evaluator_launcher/cli/status.py,sha256=3kv-VHMHlD2qHwgKzgmTZ5x_D1c9-rW86hBOghHk35M,6108
 nemo_evaluator_launcher/cli/version.py,sha256=puMwIvkmfD3HESjftdTSP6T3Nc8J4cbz8uXWHJcTemY,2030
 nemo_evaluator_launcher/common/__init__.py,sha256=6-xb4KpG8-lZbWBI42c_Gax-Sq0kMSW8UG0Vn8dOBlo,744
 nemo_evaluator_launcher/common/execdb.py,sha256=WPzg5Iu2ojvFpBuYahSt3voP_iEUpoO8NgqMLUBwFxA,9767
@@ -33,26 +33,26 @@ nemo_evaluator_launcher/executors/__init__.py,sha256=mSU1op5r7R_vqOCLDP84z6utfFg
 nemo_evaluator_launcher/executors/base.py,sha256=c_n8LasxAa_3_GbaAy-SkIK9A6HIl-T5Cp-kWZ8V2l4,2939
 nemo_evaluator_launcher/executors/registry.py,sha256=8QXSrsJyHeNi8iSttJ8KWQLXmZve1vxnnCNw_CkeopI,1409
 nemo_evaluator_launcher/executors/lepton/__init__.py,sha256=F_7yuBaYQ6WWTcptADdkL3AIZ_jXJQHGgKag-Hm7BbQ,698
-nemo_evaluator_launcher/executors/lepton/deployment_helpers.py,sha256=m4HLFt9J03PMaPXyrrokBv4Zfpf8hVFsQXYBmPJElA8,20956
-nemo_evaluator_launcher/executors/lepton/executor.py,sha256=37ZjRtu69Mff1o9dXSI_SppQd1yR6yhAX9JKZn_Wegg,37281
+nemo_evaluator_launcher/executors/lepton/deployment_helpers.py,sha256=AAIlHHn-WifevNosug0DlSDLN6NtjkclEu5LHyu1xq8,21799
+nemo_evaluator_launcher/executors/lepton/executor.py,sha256=HsdHhQ4C_pSG7jircGkFXlR8poAzUTttTzd25CzrDXg,37221
 nemo_evaluator_launcher/executors/lepton/job_helpers.py,sha256=6baTxcygfP1oFgAJ7I9EL4xRlcJDWqbqzZoE1CRrwSk,13528
 nemo_evaluator_launcher/executors/local/__init__.py,sha256=lNC_skFLYTOt-arnY3ZQnZMWzHlrtD2wAoHvDcHddwM,673
-nemo_evaluator_launcher/executors/local/executor.py,sha256=z5xeJyWydPtdv95sexvegdqc2_ip0gSED3Hyzd_v78M,17957
+nemo_evaluator_launcher/executors/local/executor.py,sha256=5DasUn_2r3X3ICNHM4-WHcFHR0vnN1jpv267cfdwV9U,18870
 nemo_evaluator_launcher/executors/local/run.template.sh,sha256=oD7QvDqEvd9NpS_5SZhYqIhCjVdcXI-GD98g5j-Ek4I,3344
 nemo_evaluator_launcher/executors/slurm/__init__.py,sha256=lNC_skFLYTOt-arnY3ZQnZMWzHlrtD2wAoHvDcHddwM,673
-nemo_evaluator_launcher/executors/slurm/executor.py,sha256=eG0hR0BuA6KRWMCaqtx2dhKeEZgm1QCC_edaEy32kDs,36274
+nemo_evaluator_launcher/executors/slurm/executor.py,sha256=LdHG4S9gB1pZS92nBkh9U-d2r_PF7j7HDC_fP857rlQ,38363
 nemo_evaluator_launcher/exporters/__init__.py,sha256=mBXG9FG48FeYrs8sF0zA2mgo1eqBmRgoml7zjJrqDso,1323
-nemo_evaluator_launcher/exporters/base.py,sha256=toeitHi-reouJvhRULtsceMlpZat4fHcQIXIbAKury0,3904
-nemo_evaluator_launcher/exporters/gsheets.py,sha256=P2TeHr63qXIGh5nzXhs2l10n5_r2I2C4uWx3pbjX-EY,15084
-nemo_evaluator_launcher/exporters/local.py,sha256=HHDkHxSHtb2IBAQagFwXYXd21XOszPVQSGWMz_66UTM,19185
-nemo_evaluator_launcher/exporters/mlflow.py,sha256=jo0wC1DT7xtaXiTUdo7dZ94IzQSQIQGgWO0b13zYp9Y,16934
+nemo_evaluator_launcher/exporters/base.py,sha256=0BEqS-Zjez-KsrGE9yfo8S5w2uwMW3btBZve3SiiUp0,4307
+nemo_evaluator_launcher/exporters/gsheets.py,sha256=hBOL3vaomCW2fPMDEOQWkZkFCgF4jCoS4U5ZlsNVENs,15911
+nemo_evaluator_launcher/exporters/local.py,sha256=oaHMyFaxihtHJUSmMdp9nlrVe8KfZ8IIq0jqDGbgS6s,19487
+nemo_evaluator_launcher/exporters/mlflow.py,sha256=1yCgB7--lLD8an2IdVjZVfZ190oeYqx-T6bUFG5H0sI,20192
 nemo_evaluator_launcher/exporters/registry.py,sha256=XsPTv_SBAFjcErO6BJ3OHqs3EvXQpLeyKRJuK9Ql4_M,1299
-nemo_evaluator_launcher/exporters/utils.py,sha256=uXH4b-Hk7_FQyLOjMRB0b3zK-Ksb2rGlSdc-oECfGHI,24756
-nemo_evaluator_launcher/exporters/wandb.py,sha256=xdaPNw0QM0jZo20UERbViy_vFT-HgbLYzTgmWaev_kk,13430
+nemo_evaluator_launcher/exporters/utils.py,sha256=5AXO5C6s93lavA2EHHfGGx6o9u_nUQ5WZtXEWxG6Hg0,20999
+nemo_evaluator_launcher/exporters/wandb.py,sha256=FhhrVxtwE0CGiSMXSgAwF6dFD5jojA4qGFEEyKkIrX4,16531
 nemo_evaluator_launcher/resources/mapping.toml,sha256=uOg4Y-gDXXskbbba2vuwJ5FLJ3W0kSZz7Fap_nJnFQc,11322
-nemo_evaluator_launcher-0.1.12.dist-info/licenses/LICENSE,sha256=DyGb0fqHPZAsd_uXHA0DGcOCqsvrNsImuLC0Ts4s1zI,23413
-nemo_evaluator_launcher-0.1.12.dist-info/METADATA,sha256=QGjX9KDLeOgKMio7WWjGfMTDWpz0yxX7bbLVktSMS30,28725
-nemo_evaluator_launcher-0.1.12.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nemo_evaluator_launcher-0.1.12.dist-info/entry_points.txt,sha256=64z1T5GKSB9PW1fCENQuor6X6eqH1rcfg0NQGfKrEy8,130
-nemo_evaluator_launcher-0.1.12.dist-info/top_level.txt,sha256=5PvawNm9TXKqPRjZita1xPOtFiMOipcoRf50FI1iY3s,24
-nemo_evaluator_launcher-0.1.12.dist-info/RECORD,,
+nemo_evaluator_launcher-0.1.14.dist-info/licenses/LICENSE,sha256=DyGb0fqHPZAsd_uXHA0DGcOCqsvrNsImuLC0Ts4s1zI,23413
+nemo_evaluator_launcher-0.1.14.dist-info/METADATA,sha256=Q3OhzGfr5E4fGGvNfyLJgijXvjGFwFZrhrCQG7f8xtU,28725
+nemo_evaluator_launcher-0.1.14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nemo_evaluator_launcher-0.1.14.dist-info/entry_points.txt,sha256=64z1T5GKSB9PW1fCENQuor6X6eqH1rcfg0NQGfKrEy8,130
+nemo_evaluator_launcher-0.1.14.dist-info/top_level.txt,sha256=5PvawNm9TXKqPRjZita1xPOtFiMOipcoRf50FI1iY3s,24
+nemo_evaluator_launcher-0.1.14.dist-info/RECORD,,