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

nemo_evaluator_launcher/exporters/mlflow.py

@@ -16,13 +16,10 @@
 """Evaluation results exporter for MLflow tracking."""
 
 import os
-import shutil
 import tempfile
 from pathlib import Path
 from typing import Any, Dict, List
 
-import yaml
-
 try:
     import mlflow
 
@@ -42,6 +39,7 @@ from nemo_evaluator_launcher.exporters.utils import (
     get_available_artifacts,
     get_benchmark_info,
     get_task_name,
+    mlflow_sanitize,
 )
 
 
@@ -117,11 +115,42 @@ class MLflowExporter(BaseExporter):
                     message="tracking_uri is required (set export.mlflow.tracking_uri or MLFLOW_TRACKING_URI)",
                 )
 
-            # Extract metrics
+            # Stage artifacts locally if remote_ssh (e.g., Slurm), so we can extract metrics
+            staged_base_dir = None
+            try:
+                paths = self.get_job_paths(job_data)
+                if paths.get("storage_type") == "remote_ssh":
+                    tmp_stage = Path(tempfile.mkdtemp(prefix="mlflow_stage_"))
+                    LocalExporter(
+                        {
+                            "output_dir": str(tmp_stage),
+                            "copy_logs": mlflow_config.get(
+                                "log_logs", False
+                            ),  # log_logs -> copy_logs
+                            "only_required": mlflow_config.get("only_required", True),
+                        }
+                    ).export_job(job_data)
+                    staged_base_dir = (
+                        tmp_stage / job_data.invocation_id / job_data.job_id
+                    )
+            except Exception as e:
+                logger.warning(f"Failed staging artifacts for {job_data.job_id}: {e}")
+
+            # Extract metrics (prefer staged if available)
             log_metrics = mlflow_config.get("log_metrics", [])
-            accuracy_metrics = extract_accuracy_metrics(
-                job_data, self.get_job_paths, log_metrics
-            )
+            if staged_base_dir and (staged_base_dir / "artifacts").exists():
+                accuracy_metrics = extract_accuracy_metrics(
+                    job_data,
+                    lambda _: {
+                        "artifacts_dir": staged_base_dir / "artifacts",
+                        "storage_type": "local_filesystem",
+                    },
+                    log_metrics,
+                )
+            else:
+                accuracy_metrics = extract_accuracy_metrics(
+                    job_data, self.get_job_paths, log_metrics
+                )
 
             if not accuracy_metrics:
                 return ExportResult(
@@ -160,10 +189,13 @@ class MLflowExporter(BaseExporter):
                     }
                 )
 
-            # Truncate params
+            # Sanitize params
             safe_params = {
-                str(k)[:250]: str(v)[:250] for k, v in all_params.items() if v
+                mlflow_sanitize(k, "param_key"): mlflow_sanitize(v, "param_value")
+                for k, v in (all_params or {}).items()
+                if v is not None
             }
+
             # Prepare tags
             tags = {}
             if mlflow_config.get("tags"):
@@ -173,7 +205,10 @@ class MLflowExporter(BaseExporter):
             benchmark = bench_info.get("benchmark", get_task_name(job_data))
             harness = bench_info.get("harness", "unknown")
 
-            # Tag the run with invocation_id and task metadata (task_name is benchmark-only)
+            # Tag the run with invocation_id and task metadata
+            exec_type = (job_data.config or {}).get("execution", {}).get(
+                "type"
+            ) or job_data.executor
             tags.update(
                 {
                     "invocation_id": job_data.invocation_id,
@@ -181,11 +216,16 @@ class MLflowExporter(BaseExporter):
                     "task_name": benchmark,
                     "benchmark": benchmark,
                     "harness": harness,
-                    "executor": job_data.executor,
+                    "executor": exec_type,
                 }
             )
-            # Truncate tags
-            safe_tags = {str(k)[:250]: str(v)[:5000] for k, v in tags.items() if v}
+
+            # Sanitize tags
+            safe_tags = {
+                mlflow_sanitize(k, "tag_key"): mlflow_sanitize(v, "tag_value")
+                for k, v in (tags or {}).items()
+                if v is not None
+            }
 
             # skip run if it already exists
             exists, existing_run_id = self._get_existing_run_info(
@@ -204,26 +244,34 @@ class MLflowExporter(BaseExporter):
                 if safe_tags:
                     mlflow.set_tags(safe_tags)
 
-                # Set run name)
+                # Set run name
                 run_name = (
                     mlflow_config.get("run_name")
                     or f"eval-{job_data.invocation_id}-{benchmark}"
                 )
-                mlflow.set_tag("mlflow.runName", run_name)
+                mlflow.set_tag("mlflow.runName", mlflow_sanitize(run_name, "tag_value"))
 
                 # Set description only if provided
                 description = mlflow_config.get("description")
                 if description:
-                    mlflow.set_tag("mlflow.note.content", str(description)[:5000])
+                    mlflow.set_tag(
+                        "mlflow.note.content", mlflow_sanitize(description, "tag_value")
+                    )
 
                 # Log parameters
                 mlflow.log_params(safe_params)
 
-                # Log metrics
-                mlflow.log_metrics(accuracy_metrics)
+                # Sanitize metric keys before logging
+                safe_metrics = {
+                    mlflow_sanitize(k, "metric"): v
+                    for k, v in (accuracy_metrics or {}).items()
+                }
+                mlflow.log_metrics(safe_metrics)
 
                 # Log artifacts
-                artifacts_logged = self._log_artifacts(job_data, mlflow_config)
+                artifacts_logged = self._log_artifacts(
+                    job_data, mlflow_config, staged_base_dir
+                )
 
                 # Build run URL
                 run_url = None
@@ -253,7 +301,10 @@ class MLflowExporter(BaseExporter):
             )
 
     def _log_artifacts(
-        self, job_data: JobData, mlflow_config: Dict[str, Any]
+        self,
+        job_data: JobData,
+        mlflow_config: Dict[str, Any],
+        pre_staged_dir: Path = None,
     ) -> List[str]:
         """Log evaluation artifacts to MLflow using LocalExporter for transfer."""
 
@@ -262,34 +313,39 @@ class MLflowExporter(BaseExporter):
             return []
 
         try:
-            # Use LocalExporter to get files locally first
-            temp_dir = tempfile.mkdtemp(prefix="mlflow_artifacts_")
-            local_exporter = LocalExporter(
-                {
-                    "output_dir": temp_dir,
-                    "copy_logs": mlflow_config.get(
-                        "log_logs", mlflow_config.get("copy_logs", False)
-                    ),
-                    "only_required": mlflow_config.get("only_required", True),
-                    "format": mlflow_config.get("format", None),
-                    "log_metrics": mlflow_config.get("log_metrics", []),
-                    "output_filename": mlflow_config.get("output_filename", None),
-                }
-            )
-            local_result = local_exporter.export_job(job_data)
-
-            if not local_result.success:
-                logger.error(f"Failed to download artifacts: {local_result.message}")
-                return []
+            should_cleanup = False
+            # Use pre-staged dir if available; otherwise stage now
+            if pre_staged_dir and pre_staged_dir.exists():
+                base_dir = pre_staged_dir
+            else:
+                temp_dir = tempfile.mkdtemp(prefix="mlflow_artifacts_")
+                local_exporter = LocalExporter(
+                    {
+                        "output_dir": str(temp_dir),
+                        "copy_logs": mlflow_config.get(
+                            "log_logs", mlflow_config.get("copy_logs", False)
+                        ),
+                        "only_required": mlflow_config.get("only_required", True),
+                        "format": mlflow_config.get("format", None),
+                        "log_metrics": mlflow_config.get("log_metrics", []),
+                        "output_filename": mlflow_config.get("output_filename", None),
+                    }
+                )
+                local_result = local_exporter.export_job(job_data)
+                if not local_result.success:
+                    logger.error(
+                        f"Failed to download artifacts: {local_result.message}"
+                    )
+                    return []
+                base_dir = Path(local_result.dest)
+                should_cleanup = True
 
-            base_dir = Path(local_result.dest)
             artifacts_dir = base_dir / "artifacts"
             logs_dir = base_dir / "logs"
             logged_names: list[str] = []
-
             artifact_path = get_artifact_root(job_data)  # "<harness>.<benchmark>"
 
-            # Log config at root level
+            # Log config at root level (or synthesize)
             cfg_logged = False
             for fname in ("config.yml", "run_config.yml"):
                 p = artifacts_dir / fname
@@ -299,16 +355,19 @@ class MLflowExporter(BaseExporter):
                     break
             if not cfg_logged:
                 with tempfile.TemporaryDirectory() as tmpdir:
+                    from yaml import dump as ydump
+
                     cfg_file = Path(tmpdir) / "config.yaml"
-                    with cfg_file.open("w") as f:
-                        yaml.dump(
+                    cfg_file.write_text(
+                        ydump(
                             job_data.config or {},
-                            f,
                             default_flow_style=False,
                             sort_keys=False,
                         )
+                    )
                     mlflow.log_artifact(str(cfg_file))
 
+            # Choose files to upload
             files_to_upload: list[Path] = []
             if mlflow_config.get("only_required", True):
                 for fname in get_available_artifacts(artifacts_dir):
@@ -316,10 +375,11 @@ class MLflowExporter(BaseExporter):
                     if p.exists():
                         files_to_upload.append(p)
             else:
-                for p in artifacts_dir.iterdir():
+                for p in artifacts_dir.iterdir():  # top-level files only
                     if p.is_file():
                         files_to_upload.append(p)
 
+            # Upload artifacts (with DEBUG per-file)
             for fpath in files_to_upload:
                 rel = fpath.relative_to(artifacts_dir).as_posix()
                 parent = os.path.dirname(rel)
@@ -328,32 +388,28 @@ class MLflowExporter(BaseExporter):
                     artifact_path=f"{artifact_path}/artifacts/{parent}".rstrip("/"),
                 )
                 logged_names.append(rel)
+                logger.debug(f"mlflow upload artifact: {rel}")
 
             # Optionally upload logs under "<harness.task>/logs"
             if mlflow_config.get("log_logs", False) and logs_dir.exists():
-                for p in logs_dir.rglob("*"):
+                for p in logs_dir.iterdir():
                     if p.is_file():
+                        rel = p.name
                         mlflow.log_artifact(
-                            str(p),
-                            artifact_path=f"{artifact_path}/logs",
+                            str(p), artifact_path=f"{artifact_path}/logs"
                         )
-                        logged_names.append(f"logs/{p.name}")
+                        logged_names.append(f"logs/{rel}")
+                        logger.debug(f"mlflow upload log: {rel}")
 
-            # Debug summary of what we uploaded
             logger.info(
                 f"MLflow upload summary: files={len(logged_names)}, only_required={mlflow_config.get('only_required', True)}, log_logs={mlflow_config.get('log_logs', False)}"
             )
-            if logger.isEnabledFor(10):  # DEBUG
-                try:
-                    preview = "\n  - " + "\n  - ".join(sorted(logged_names)[:50])
-                    logger.debug(f"Uploaded files preview (first 50):{preview}")
-                except Exception:
-                    pass
+            if should_cleanup:
+                import shutil
 
-            # cleanup temp
-            shutil.rmtree(temp_dir)
-            return logged_names
+                shutil.rmtree(base_dir, ignore_errors=True)
 
+            return logged_names
         except Exception as e:
             logger.error(f"Error logging artifacts: {e}")
             return []
@@ -391,11 +447,42 @@ class MLflowExporter(BaseExporter):
 
             # Collect metrics from ALL jobs
             all_metrics = {}
+            staged_map: dict[str, Path] = {}
+            for job_id, job_data in jobs.items():
+                try:
+                    paths = self.get_job_paths(job_data)
+                    if paths.get("storage_type") == "remote_ssh":
+                        tmp_stage = Path(tempfile.mkdtemp(prefix="mlflow_inv_stage_"))
+                        LocalExporter(
+                            {
+                                "output_dir": str(tmp_stage),
+                                "copy_logs": mlflow_config.get("log_logs", False),
+                                "only_required": mlflow_config.get(
+                                    "only_required", True
+                                ),
+                            }
+                        ).export_job(job_data)
+                        staged_map[job_id] = (
+                            tmp_stage / job_data.invocation_id / job_data.job_id
+                        )
+                except Exception as e:
+                    logger.warning(f"Staging failed for {job_id}: {e}")
+
             for job_id, job_data in jobs.items():
                 log_metrics = mlflow_config.get("log_metrics", [])
-                job_metrics = extract_accuracy_metrics(
-                    job_data, self.get_job_paths, log_metrics
-                )
+                if job_id in staged_map and (staged_map[job_id] / "artifacts").exists():
+                    job_metrics = extract_accuracy_metrics(
+                        job_data,
+                        lambda _: {
+                            "artifacts_dir": staged_map[job_id] / "artifacts",
+                            "storage_type": "local_filesystem",
+                        },
+                        log_metrics,
+                    )
+                else:
+                    job_metrics = extract_accuracy_metrics(
+                        job_data, self.get_job_paths, log_metrics
+                    )
                 all_metrics.update(job_metrics)
 
             if not all_metrics:
@@ -414,9 +501,12 @@ class MLflowExporter(BaseExporter):
             mlflow.set_experiment(experiment_name)
 
             # Prepare parameters for invocation
+            inv_exec_type = (first_job.config or {}).get("execution", {}).get(
+                "type"
+            ) or first_job.executor
             all_params = {
                 "invocation_id": invocation_id,
-                "executor": first_job.executor,
+                "executor": inv_exec_type,
                 "timestamp": str(first_job.timestamp),
                 "jobs_count": str(len(jobs)),
             }
@@ -472,23 +562,31 @@ class MLflowExporter(BaseExporter):
 
                 # Set run name
                 run_name = mlflow_config.get("run_name") or f"eval-{invocation_id}"
-                mlflow.set_tag("mlflow.runName", run_name)
+                mlflow.set_tag("mlflow.runName", mlflow_sanitize(run_name, "tag_value"))
 
                 # Set description
                 description = mlflow_config.get("description")
                 if description:
-                    mlflow.set_tag("mlflow.note.content", str(description)[:5000])
+                    mlflow.set_tag(
+                        "mlflow.note.content", mlflow_sanitize(description, "tag_value")
+                    )
 
                 # Log parameters
                 mlflow.log_params(safe_params)
 
-                # Log ALL metrics
-                mlflow.log_metrics(all_metrics)
+                # Sanitize metric keys
+                safe_all_metrics = {
+                    mlflow_sanitize(k, "metric"): v
+                    for k, v in (all_metrics or {}).items()
+                }
+                mlflow.log_metrics(safe_all_metrics)
 
                 # Log artifacts from all jobs
                 total_artifacts = 0
-                for job_data in jobs.values():
-                    artifacts_logged = self._log_artifacts(job_data, mlflow_config)
+                for job_id, job_data in jobs.items():
+                    artifacts_logged = self._log_artifacts(
+                        job_data, mlflow_config, staged_map.get(job_id)
+                    )
                     total_artifacts += len(artifacts_logged)
 
                 # Build run URL

nemo_evaluator_launcher/exporters/utils.py

@@ -16,6 +16,7 @@
 """Shared utilities for metrics and configuration handling."""
 
 import json
+import re
 import subprocess
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Tuple
@@ -306,21 +307,28 @@ def ssh_setup_masters(jobs: Dict[str, JobData]) -> Dict[Tuple[str, str], str]:
     remote_pairs: set[tuple[str, str]] = set()
     for jd in jobs.values():
         try:
-            paths = jd.data.get("paths") or {}
-            if paths.get("storage_type") == "remote_ssh":
-                remote_pairs.add((paths["username"], paths["hostname"]))
+            # Preferred: explicit 'paths' from job data
+            p = (jd.data or {}).get("paths") or {}
+            if (
+                p.get("storage_type") == "remote_ssh"
+                and p.get("username")
+                and p.get("hostname")
+            ):
+                remote_pairs.add((p["username"], p["hostname"]))
+                continue
+            # Fallback: common slurm fields (works with BaseExporter.get_job_paths)
+            d = jd.data or {}
+            if jd.executor == "slurm" and d.get("username") and d.get("hostname"):
+                remote_pairs.add((d["username"], d["hostname"]))
         except Exception:
             pass
 
     if not remote_pairs:
-        return {}  # no remote jobs
+        return {}
 
-    # Ensure connections directory exists (like execDB does)
     CONNECTIONS_DIR.mkdir(parents=True, exist_ok=True)
-
     control_paths: Dict[Tuple[str, str], str] = {}
     for username, hostname in remote_pairs:
-        # Simple socket name
         socket_path = CONNECTIONS_DIR / f"{username}_{hostname}.sock"
         try:
             cmd = [
@@ -371,9 +379,10 @@ def ssh_download_artifacts(
     config: Dict[str, Any] | None = None,
     control_paths: Dict[Tuple[str, str], str] | None = None,
 ) -> List[str]:
-    """Download artifacts via SSH with optional connection reuse."""
+    """Download artifacts/logs via SSH with optional connection reuse."""
     exported_files: List[str] = []
     copy_logs = bool((config or {}).get("copy_logs", False))
+    copy_artifacts = bool((config or {}).get("copy_artifacts", True))
     only_required = bool((config or {}).get("only_required", True))
 
     control_path = None
@@ -390,44 +399,49 @@ def ssh_download_artifacts(
                 str(local_path),
             ]
         )
-        result = subprocess.run(cmd, capture_output=True)
-        return result.returncode == 0
+        return subprocess.run(cmd, capture_output=True).returncode == 0
 
     export_dir.mkdir(parents=True, exist_ok=True)
-    (export_dir / "artifacts").mkdir(parents=True, exist_ok=True)
-
-    available_local = (
-        get_available_artifacts(paths.get("artifacts_dir", Path()))
-        if not only_required
-        else None
-    )
-    artifact_names = (
-        [a for a in get_relevant_artifacts()]
-        if only_required
-        else (available_local or [])
-    )
-
-    for artifact in artifact_names:
-        remote_file = f"{paths['remote_path']}/artifacts/{artifact}"
-        local_file = export_dir / "artifacts" / artifact
-        if scp_file(remote_file, local_file):
-            exported_files.append(str(local_file))
 
+    # Artifacts
+    if copy_artifacts:
+        art_dir = export_dir / "artifacts"
+        art_dir.mkdir(parents=True, exist_ok=True)
+
+        if only_required:
+            for artifact in get_relevant_artifacts():
+                remote_file = f"{paths['remote_path']}/artifacts/{artifact}"
+                local_file = art_dir / artifact
+                local_file.parent.mkdir(parents=True, exist_ok=True)
+                if scp_file(remote_file, local_file):
+                    exported_files.append(str(local_file))
+        else:
+            # Copy known files individually to avoid subfolders and satisfy tests
+            for artifact in get_available_artifacts(paths.get("artifacts_dir", Path())):
+                remote_file = f"{paths['remote_path']}/artifacts/{artifact}"
+                local_file = art_dir / artifact
+                if scp_file(remote_file, local_file):
+                    exported_files.append(str(local_file))
+
+    # Logs (top-level only)
     if copy_logs:
-        remote_logs = f"{paths['remote_path']}/logs"
         local_logs = export_dir / "logs"
+        remote_logs = f"{paths['remote_path']}/logs"
         cmd = (
             ["scp", "-r"]
             + ssh_opts
             + [
-                f"{paths['username']}@{paths['hostname']}:{remote_logs}",
+                f"{paths['username']}@{paths['hostname']}:{remote_logs}/.",
                 str(local_logs),
             ]
         )
         if subprocess.run(cmd, capture_output=True).returncode == 0:
-            exported_files.extend(
-                [str(f) for f in local_logs.rglob("*") if f.is_file()]
-            )
+            for p in local_logs.iterdir():
+                if p.is_dir():
+                    import shutil
+
+                    shutil.rmtree(p, ignore_errors=True)
+            exported_files.extend([str(f) for f in local_logs.glob("*") if f.is_file()])
 
     return exported_files
 
@@ -584,3 +598,41 @@ def _safe_update_metrics(
     """Update target from source safely, raising on collisions with detailed values."""
     for k, v in source.items():
         _safe_set_metric(target, k, v, context)
+
+
+# =============================================================================
+# MLFLOW FUNCTIONS
+# =============================================================================
+
+# MLflow constants
+_MLFLOW_KEY_MAX = 250
+_MLFLOW_PARAM_VAL_MAX = 250
+_MLFLOW_TAG_VAL_MAX = 5000
+
+_INVALID_KEY_CHARS = re.compile(r"[^/\w.\- ]")
+_MULTI_UNDERSCORE = re.compile(r"_+")
+
+
+def mlflow_sanitize(s: Any, kind: str = "key") -> str:
+    """
+    Sanitize strings for MLflow logging.
+
+    kind:
+      - "key", "metric", "tag_key", "param_key": apply key rules
+      - "tag_value": apply tag value rules
+      - "param_value": apply param value rules
+    """
+    s = "" if s is None else str(s)
+
+    if kind in ("key", "metric", "tag_key", "param_key"):
+        # common replacements
+        s = s.replace("pass@", "pass_at_")
+        # drop disallowed chars, collapse underscores, trim
+        s = _INVALID_KEY_CHARS.sub("_", s)
+        s = _MULTI_UNDERSCORE.sub("_", s).strip()
+        return s[:_MLFLOW_KEY_MAX] or "key"
+
+    # values: normalize whitespace, enforce length
+    s = s.replace("\n", " ").replace("\r", " ").strip()
+    max_len = _MLFLOW_TAG_VAL_MAX if kind == "tag_value" else _MLFLOW_PARAM_VAL_MAX
+    return s[:max_len]

nemo_evaluator_launcher/exporters/wandb.py

@@ -68,10 +68,41 @@ class WandBExporter(BaseExporter):
                 "log_mode", "per_task"
             )  # Default per_task for immediate export
 
-            # Get metrics
-            metrics = extract_accuracy_metrics(
-                job_data, self.get_job_paths, wandb_config.get("log_metrics", [])
-            )
+            # Stage artifacts locally if remote_ssh (e.g., Slurm), so we can extract metrics
+            staged_base_dir = None
+            try:
+                paths = self.get_job_paths(job_data)
+                if paths.get("storage_type") == "remote_ssh":
+                    tmp_stage = Path(tempfile.mkdtemp(prefix="wandb_stage_"))
+                    LocalExporter(
+                        {
+                            "output_dir": str(tmp_stage),
+                            "copy_logs": wandb_config.get("log_logs", False),
+                            "only_required": wandb_config.get("only_required", True),
+                        }
+                    ).export_job(job_data)
+                    staged_base_dir = (
+                        tmp_stage / job_data.invocation_id / job_data.job_id
+                    )
+            except Exception as e:
+                logger.warning(f"W&B: staging failed for {job_data.job_id}: {e}")
+
+            # Metrics (prefer staged if available)
+            log_metrics = wandb_config.get("log_metrics", [])
+            if staged_base_dir and (staged_base_dir / "artifacts").exists():
+                metrics = extract_accuracy_metrics(
+                    job_data,
+                    lambda _: {
+                        "artifacts_dir": staged_base_dir / "artifacts",
+                        "storage_type": "local_filesystem",
+                    },
+                    log_metrics,
+                )
+            else:
+                metrics = extract_accuracy_metrics(
+                    job_data, self.get_job_paths, log_metrics
+                )
+
             if not metrics:
                 return ExportResult(
                     success=False, dest="wandb", message="No metrics found"
@@ -345,10 +376,14 @@ class WandBExporter(BaseExporter):
             run_args["resume"] = "allow"
 
         # Config metadata
+        exec_type = (job_data.config or {}).get("execution", {}).get(
+            "type"
+        ) or job_data.executor
         run_config = {
             "invocation_id": job_data.invocation_id,
-            "executor": job_data.executor,
+            "executor": exec_type,
         }
+
         if log_mode == "per_task":
             run_config["job_id"] = job_data.job_id
             run_config["harness"] = harness

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 = 14
+PATCH = 15
 PRE_RELEASE = ""
 
 # Use the following formatting: (major, minor, patch, pre-release)

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

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