nemo-evaluator-launcher 0.1.0rc6__py3-none-any.whl → 0.1.41__py3-none-any.whl

nemo_evaluator_launcher/__init__.py

@@ -20,6 +20,7 @@ It automatically initializes logging and conditionally loads internal components
 """
 
 import importlib
+import warnings
 
 from nemo_evaluator_launcher.common.logging_utils import logger
 from nemo_evaluator_launcher.package_info import (
@@ -32,9 +33,22 @@ from nemo_evaluator_launcher.package_info import (
     __version__,
 )
 
-logger.info("Version info", pkg=__package_name__, ver=__version__)
+# Suppress pydantic warnings from third-party libraries (e.g., wandb) that are not
+# compatible with Pydantic 2.x field metadata on Python 3.13+
+warnings.filterwarnings(
+    "ignore",
+    message=r"The 'repr' attribute.*Field\(\).*",
+    category=Warning,
+)
+warnings.filterwarnings(
+    "ignore",
+    message=r"The 'frozen' attribute.*Field\(\).*",
+    category=Warning,
+)
 
 
+logger.info("Version info", pkg=__package_name__, ver=__version__)
+
 try:
     importlib.import_module("nemo_evaluator_launcher_internal")
     logger.debug(

nemo_evaluator_launcher/api/functional.py

@@ -18,8 +18,10 @@
 This module provides the main functional entry points for running evaluations, querying job status, and listing available tasks. These functions are intended to be used by CLI commands and external integrations.
 """
 
-from typing import Any, List, Optional, Union
+from pathlib import Path
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
+import yaml
 from omegaconf import DictConfig, OmegaConf
 
 from nemo_evaluator_launcher.api.types import RunConfig
@@ -97,11 +99,13 @@ def run_eval(cfg: RunConfig, dry_run: bool = False) -> Optional[str]:
     return get_executor(cfg.execution.type).execute_eval(cfg, dry_run)
 
 
-def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
+def get_status(ids_or_prefixes: list[str]) -> list[dict[str, Any]]:
     """Get status of jobs by their IDs or invocation IDs.
 
     Args:
-        job_ids: List of job IDs or invocation IDs to check status for.
+        job_ids: List of job IDs or invocation IDs to check status for. Short ones are allowed,
+                 we would try to match the full ones from prefixes if no collisions are
+                 present.
 
     Returns:
         list[dict[str, Any]]: List of status dictionaries for each job or invocation.
@@ -112,14 +116,15 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
     db = ExecutionDB()
     results: List[dict[str, Any]] = []
 
-    for job_id in job_ids:
-        # If id looks like an invocation_id (8 hex digits, no dot), get all jobs for it
-        if len(job_id) == 8 and "." not in job_id:
-            jobs = db.get_jobs(job_id)
+    # TODO(agronskiy): refactor the `.`-checking job in all the functions.
+    for id_or_prefix in ids_or_prefixes:
+        # If id looks like an invocation_id (no dot), get all jobs for it
+        if "." not in id_or_prefix:
+            jobs = db.get_jobs(id_or_prefix)
             if not jobs:
                 results.append(
                     {
-                        "invocation": job_id,
+                        "invocation": id_or_prefix,
                         "job_id": None,
                         "status": "not_found",
                         "data": {},
@@ -134,7 +139,7 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
             except ValueError as e:
                 results.append(
                     {
-                        "invocation": job_id,
+                        "invocation": id_or_prefix,
                         "job_id": None,
                         "status": "error",
                         "data": {"error": str(e)},
@@ -144,7 +149,7 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
 
             # Get status from the executor for all jobs in the invocation
             try:
-                status_list = executor_cls.get_status(job_id)
+                status_list = executor_cls.get_status(id_or_prefix)
 
                 # Create a result for each job in the invocation
                 for job_id_in_invocation, job_data in jobs.items():
@@ -159,7 +164,7 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
 
                     results.append(
                         {
-                            "invocation": job_id,
+                            "invocation": job_data.invocation_id,
                             "job_id": job_id_in_invocation,
                             "status": (
                                 job_status if job_status is not None else "unknown"
@@ -174,7 +179,7 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
             except Exception as e:
                 results.append(
                     {
-                        "invocation": job_id,
+                        "invocation": id_or_prefix,
                         "job_id": None,
                         "status": "error",
                         "data": {"error": str(e)},
@@ -182,13 +187,13 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
                 )
         else:
             # Otherwise, treat as job_id
-            single_job_data: Optional[JobData] = db.get_job(job_id)
+            single_job_data: Optional[JobData] = db.get_job(id_or_prefix)
 
             if single_job_data is None:
                 results.append(
                     {
                         "invocation": None,
-                        "job_id": job_id,
+                        "job_id": id_or_prefix,
                         "status": "not_found",
                         "data": {},
                     }
@@ -202,7 +207,7 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
                 results.append(
                     {
                         "invocation": None,
-                        "job_id": job_id,
+                        "job_id": id_or_prefix,
                         "status": "error",
                         "data": {"error": str(e)},
                     }
@@ -211,13 +216,13 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
 
             # Get status from the executor
             try:
-                status_list = executor_cls.get_status(job_id)
+                status_list = executor_cls.get_status(id_or_prefix)
 
                 if not status_list:
                     results.append(
                         {
                             "invocation": single_job_data.invocation_id,
-                            "job_id": job_id,
+                            "job_id": single_job_data.job_id,
                             "status": "unknown",
                             "data": single_job_data.data,
                         }
@@ -227,7 +232,7 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
                     results.append(
                         {
                             "invocation": single_job_data.invocation_id,
-                            "job_id": job_id,
+                            "job_id": single_job_data.job_id,
                             "status": (
                                 status_list[0].state.value if status_list else "unknown"
                             ),
@@ -244,7 +249,9 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
                         "invocation": (
                             single_job_data.invocation_id if single_job_data else None
                         ),
-                        "job_id": job_id,
+                        "job_id": (
+                            single_job_data.job_id if single_job_data else id_or_prefix
+                        ),
                         "status": "error",
                         "data": {"error": str(e)},
                     }
@@ -253,6 +260,108 @@ def get_status(job_ids: list[str]) -> list[dict[str, Any]]:
     return results
 
 
+def stream_logs(
+    ids_or_prefixes: Union[str, list[str]],
+) -> Iterator[Tuple[str, str, str]]:
+    """Stream logs from jobs or invocations by their IDs or invocation IDs.
+
+    Args:
+        ids_or_prefixes: Single ID/prefix or list of job IDs or invocation IDs to stream logs from.
+                         Short prefixes are allowed, we would try to match the full ones from
+                         prefixes if no collisions are present.
+
+    Yields:
+        Tuple[str, str, str]: Tuples of (job_id, task_name, log_line) for each log line.
+            Empty lines are yielded as empty strings.
+
+    Raises:
+        ValueError: If the executor doesn't support log streaming.
+    """
+    db = ExecutionDB()
+
+    # Normalize to list for consistent processing
+    if isinstance(ids_or_prefixes, str):
+        ids_or_prefixes = [ids_or_prefixes]
+
+    # Collect all jobs from all IDs, grouped by executor
+    executor_to_jobs: Dict[str, Dict[str, JobData]] = {}
+    executor_to_invocations: Dict[str, List[str]] = {}
+
+    # TODO(agronskiy): refactor the `.`-checking job in all the functions.
+    for id_or_prefix in ids_or_prefixes:
+        # Determine if this is a job ID or invocation ID
+        if "." in id_or_prefix:
+            # This is a job ID
+            job_data = db.get_job(id_or_prefix)
+            if job_data is None:
+                continue
+
+            executor = job_data.executor
+            if executor not in executor_to_jobs:
+                executor_to_jobs[executor] = {}
+            executor_to_jobs[executor][id_or_prefix] = job_data
+        else:
+            # This is an invocation ID
+            jobs = db.get_jobs(id_or_prefix)
+            if not jobs:
+                continue
+
+            # Get the executor class from the first job
+            first_job_data = next(iter(jobs.values()))
+            executor = first_job_data.executor
+            if executor not in executor_to_invocations:
+                executor_to_invocations[executor] = []
+            executor_to_invocations[executor].append(id_or_prefix)
+
+    # Stream logs from each executor simultaneously
+    # For each executor, collect all job IDs and stream them together
+    for executor, jobs_dict in executor_to_jobs.items():
+        try:
+            executor_cls = get_executor(executor)
+        except ValueError:
+            continue
+
+        # For local executor with multiple jobs, pass list to stream simultaneously
+        # For other executors or single jobs, pass individual job IDs
+        if executor == "local" and len(jobs_dict) > 1:
+            # Pass all job IDs as a list to stream simultaneously
+            try:
+                yield from executor_cls.stream_logs(
+                    list(jobs_dict.keys()), executor_name=executor
+                )
+            except NotImplementedError:
+                raise ValueError(
+                    f"Log streaming is not yet implemented for executor '{executor}'"
+                )
+        else:
+            # Single job or non-local executor
+            for job_id in jobs_dict.keys():
+                try:
+                    yield from executor_cls.stream_logs(job_id, executor_name=executor)
+                except NotImplementedError:
+                    raise ValueError(
+                        f"Log streaming is not yet implemented for executor '{executor}'"
+                    )
+
+    # Stream logs from invocation IDs
+    for executor, invocation_ids in executor_to_invocations.items():
+        try:
+            executor_cls = get_executor(executor)
+        except ValueError:
+            continue
+
+        # Stream each invocation (each invocation already handles multiple jobs internally)
+        for invocation_id in invocation_ids:
+            try:
+                yield from executor_cls.stream_logs(
+                    invocation_id, executor_name=executor
+                )
+            except NotImplementedError:
+                raise ValueError(
+                    f"Log streaming is not yet implemented for executor '{executor}'"
+                )
+
+
 def list_all_invocations_summary() -> list[dict[str, Any]]:
     """Return a concise per-invocation summary from the exec DB.
 
@@ -372,6 +481,7 @@ def kill_job_or_invocation(id: str) -> list[dict[str, Any]]:
                 "data": {"error": f"Unexpected error: {str(e)}"},
             }
 
+    # TODO(agronskiy): refactor the `.`-checking job in all the functions.
     # Determine if this is a job ID or invocation ID
     if "." in id:
         # This is a job ID - kill single job
@@ -434,14 +544,66 @@ def export_results(
             single_id = invocation_ids[0]
 
             if "." in single_id:  # job_id
-                job_data = ExecutionDB().get_job(single_id)
+                # Try reading config from artifacts working dir (auto-export on remote node)
+                cfg_file = None
+                for name in ("run_config.yml", "config.yml"):
+                    p = Path(name)
+                    if p.exists():
+                        cfg_file = p
+                        break
+
+                md_job_data = None
+                if cfg_file:
+                    try:
+                        cfg_yaml = (
+                            yaml.safe_load(cfg_file.read_text(encoding="utf-8")) or {}
+                        )
+
+                        # Merge exporter override file if present
+                        ypath_export = Path("export_config.yml")
+                        if ypath_export.exists():
+                            exp_yaml = (
+                                yaml.safe_load(ypath_export.read_text(encoding="utf-8"))
+                                or {}
+                            )
+                            exec_cfg = cfg_yaml.get("execution") or {}
+                            auto_exp = (exp_yaml.get("execution") or {}).get(
+                                "auto_export"
+                            )
+                            if auto_exp is not None:
+                                exec_cfg["auto_export"] = auto_exp
+                                cfg_yaml["execution"] = exec_cfg
+                            if "export" in exp_yaml:
+                                cfg_yaml["export"] = exp_yaml["export"]
+                            if "evaluation" in exp_yaml and exp_yaml["evaluation"]:
+                                eval_cfg = cfg_yaml.get("evaluation") or {}
+                                eval_cfg.update(exp_yaml["evaluation"])
+                                cfg_yaml["evaluation"] = eval_cfg
+
+                        executor_name = (cfg_yaml.get("execution") or {}).get(
+                            "type", "local"
+                        )
+                        md_job_data = JobData(
+                            invocation_id=single_id.split(".")[0],
+                            job_id=single_id,
+                            timestamp=0.0,
+                            executor=executor_name,  # ensures slurm tag is preserved
+                            data={
+                                "output_dir": str(Path.cwd().parent),
+                                "storage_type": "remote_local",  # no SSH in auto-export path
+                            },
+                            config=cfg_yaml,
+                        )
+                    except Exception:
+                        md_job_data = None
+
+                job_data = md_job_data or ExecutionDB().get_job(single_id)
                 if job_data is None:
                     return {
                         "success": False,
                         "error": f"Job {single_id} not found in ExecutionDB",
                     }
 
-                # Convert single job result to invocation-like structure
                 job_result = exporter.export_job(job_data)
                 return {
                     "success": job_result.success,
@@ -451,14 +613,14 @@ def export_results(
                             "success": job_result.success,
                             "message": job_result.message,
                             "metadata": job_result.metadata or {},
+                            "dest": getattr(job_result, "dest", None),
                         }
                     },
                     "metadata": job_result.metadata or {},
                 }
+
             elif single_id.isdigit():  # pipeline_id
-                # Find job by pipeline_id
                 db = ExecutionDB()
-                # Search all jobs for matching pipeline_id
                 for job_id, job_data in db._jobs.items():
                     if job_data.data.get("pipeline_id") == int(single_id):
                         job_result = exporter.export_job(job_data)
@@ -475,14 +637,13 @@ def export_results(
                             "metadata": job_result.metadata or {},
                         }
                 return {"success": False, "error": f"Pipeline {single_id} not found"}
+
             else:  # invocation_id
                 result = exporter.export_invocation(single_id)
-                # Ensure metadata is present in job results to prevent KeyError
                 if "jobs" in result:
                     for job_id, job_result in result["jobs"].items():
-                        if "metadata" not in job_result:
-                            job_result["metadata"] = {}
-                return result  # type: ignore[no-any-return]
+                        job_result.setdefault("metadata", {})
+                return result
         else:
             # Multiple IDs - parse and group
             db = ExecutionDB()

nemo_evaluator_launcher/api/types.py

@@ -19,9 +19,18 @@ This module defines data structures and helpers for configuration and type safet
 """
 
 import os
+import warnings
 from dataclasses import dataclass
 from typing import cast
 
+# ruff: noqa: E402
+# Later when adding optional module to hydra, since the internal package is optional,
+# will generate a hydra warning. We suppress it as distraction and bad UX, before hydra gets invoked.
+warnings.filterwarnings(
+    "ignore",
+    message="provider=hydra.searchpath.*path=nemo_evaluator_launcher_internal.*is not available\\.",
+)
+
 import hydra
 from hydra.core.global_hydra import GlobalHydra
 from omegaconf import DictConfig, OmegaConf

nemo_evaluator_launcher/cli/export.py

@@ -20,8 +20,6 @@ from typing import Any, List, Optional
 
 from simple_parsing import field
 
-from nemo_evaluator_launcher.api.functional import export_results
-
 
 @dataclass
 class ExportCmd:
@@ -29,8 +27,8 @@ class ExportCmd:
 
     # Short usage examples will show up in -h as the class docstring:
     # Examples:
-    #   nemo-evaluator-launcher export 8abcd123 --dest local --format json -o .
-    #   nemo-evaluator-launcher export 8abcd123.0 9ef01234 --dest local --format csv -o results/ -fname processed_results.csv
+    #   nemo-evaluator-launcher export 8abcd123 --dest local --format json --out .
+    #   nemo-evaluator-launcher export 8abcd123.0 9ef01234 --dest local --format csv --out results/ -fname processed_results.csv
     #   nemo-evaluator-launcher export 8abcd123 --dest jet
 
     invocation_ids: List[str] = field(
@@ -43,9 +41,17 @@ class ExportCmd:
         choices=["local", "wandb", "mlflow", "gsheets", "jet"],
         help="Export destination.",
     )
+    # overrides for exporter config; use -o similar to run command
+    override: List[str] = field(
+        default_factory=list,
+        action="append",
+        nargs="?",
+        alias=["-o", "--override"],
+        help="Hydra-style overrides for exporter config. Use `export.<dest>.key=value` (e.g., -o export.wandb.entity=org-name).",
+    )
     output_dir: Optional[str] = field(
         default=".",
-        alias=["--output-dir", "-o"],
+        alias=["--output-dir", "-out"],
         help="Output directory (default: current directory).",
     )
     output_filename: Optional[str] = field(
@@ -69,17 +75,29 @@ class ExportCmd:
         alias=["--log-metrics"],
         help="Filter metrics by name (repeatable). Examples: score, f1, mmlu_score_micro.",
     )
-    only_required: bool = field(
-        default=True,
+    only_required: Optional[bool] = field(
+        default=None,
         alias=["--only-required"],
         help="Copy only required+optional artifacts (default: True). Set to False to copy all available artifacts.",
     )
 
     def execute(self) -> None:
         """Execute export."""
+        # Import heavy dependencies only when needed
+        from omegaconf import OmegaConf
+
+        from nemo_evaluator_launcher.api.functional import export_results
+
+        # Validation: ensure IDs are provided
+        if not self.invocation_ids:
+            print("Error: No IDs provided. Specify one or more invocation or job IDs.")
+            print(
+                "Usage: nemo-evaluator-launcher export <id> [<id>...] --dest <destination>"
+            )
+            return
+
         config: dict[str, Any] = {
             "copy_logs": self.copy_logs,
-            "only_required": self.only_required,
         }
 
         # Output handling
@@ -94,20 +112,90 @@ class ExportCmd:
         if self.log_metrics:
             config["log_metrics"] = self.log_metrics
 
+        # Add only_required if explicitly passed via CLI
+        if self.only_required is not None:
+            config["only_required"] = self.only_required
+
+        # Parse and validate overrides
+        if self.override:
+            # Flatten possible list-of-lists from parser
+            flat_overrides: list[str] = []
+            for item in self.override:
+                if isinstance(item, list):
+                    flat_overrides.extend(str(x) for x in item)
+                else:
+                    flat_overrides.append(str(item))
+
+            try:
+                self._validate_overrides(flat_overrides, self.dest)
+            except ValueError as e:
+                print(f"Error: {e}")
+                return
+
+            # Expand env vars in override vals ($VAR / ${VAR})
+            import os
+
+            from omegaconf import OmegaConf
+
+            expanded_overrides: list[str] = []
+            for ov in flat_overrides:
+                if "=" in ov:
+                    k, v = ov.split("=", 1)
+                    expanded_overrides.append(f"{k}={os.path.expandvars(v)}")
+                else:
+                    expanded_overrides.append(os.path.expandvars(ov))
+
+            dot_cfg = OmegaConf.from_dotlist(expanded_overrides)
+            as_dict = OmegaConf.to_container(dot_cfg, resolve=True) or {}
+            if isinstance(as_dict, dict) and "export" in as_dict:
+                export_map = as_dict.get("export") or {}
+                if isinstance(export_map, dict) and self.dest in export_map:
+                    config.update(export_map[self.dest] or {})
+                else:
+                    config.update(as_dict)
+            else:
+                config.update(as_dict)
+
         if self.format and self.dest != "local":
             print(
                 "Note: --format is only used by --dest local. It will be ignored for other destinations."
             )
 
-        # Execute
+        if "only_required" in config and self.only_required is True:
+            config.pop("only_required", None)
+
         print(
             f"Exporting {len(self.invocation_ids)} {'invocations' if len(self.invocation_ids) > 1 else 'invocation'} to {self.dest}..."
         )
 
         result = export_results(self.invocation_ids, self.dest, config)
 
-        if not result["success"]:
-            print(f"Export failed: {result.get('error', 'Unknown error')}")
+        if not result.get("success", False):
+            err = result.get("error", "Unknown error")
+            print(f"\nExport failed: {err}")
+            # Provide actionable guidance for common configuration issues
+            if self.dest == "mlflow":
+                if "tracking_uri" in str(err).lower():
+                    print("\nMLflow requires 'tracking_uri' to be configured.")
+                    print(
+                        "Set it via: -o export.mlflow.tracking_uri=http://mlflow-server:5000"
+                    )
+                elif "not installed" in str(err).lower():
+                    print("\nMLflow package not installed.")
+                    print("Install via: pip install nemo-evaluator-launcher[mlflow]")
+            elif self.dest == "wandb":
+                if "entity" in str(err).lower() or "project" in str(err).lower():
+                    print("\nW&B requires 'entity' and 'project' to be configured.")
+                    print(
+                        "Set via: -o export.wandb.entity=my-org -o export.wandb.project=my-proj"
+                    )
+                elif "not installed" in str(err).lower():
+                    print("\nW&B package not installed.")
+                    print("Install via: pip install nemo-evaluator-launcher[wandb]")
+            elif self.dest == "gsheets":
+                if "not installed" in str(err).lower():
+                    print("\nGoogle Sheets package not installed.")
+                    print("Install via: pip install nemo-evaluator-launcher[gsheets]")
             return
 
         # Success path
@@ -124,6 +212,9 @@ class ExportCmd:
                         print(f"    URL: {metadata['run_url']}")
                     if metadata.get("summary_path"):
                         print(f"    Summary: {metadata['summary_path']}")
+                    path_hint = job_result.get("dest") or metadata.get("output_dir")
+                    if self.dest == "local" and path_hint:
+                        print(f"    Path: {path_hint}")
                 else:
                     print(f"  {job_id} failed: {job_result.get('message', '')}")
         else:
@@ -136,7 +227,6 @@ class ExportCmd:
             # Show summary path if available
             if metadata.get("summary_path"):
                 print(f"Summary: {metadata['summary_path']}")
-
             # Show per-invocation status
             for invocation_id, inv_result in result["invocations"].items():
                 if inv_result.get("success"):
@@ -146,3 +236,32 @@ class ExportCmd:
                     print(
                         f"  {invocation_id}: failed, {inv_result.get('error', 'Unknown error')}"
                     )
+
+    def _validate_overrides(self, overrides: List[str], dest: str) -> None:
+        """Validate override list for destination consistency.
+
+        Raises:
+            ValueError: If overrides specify wrong destination or have other issues.
+        """
+        if not overrides:
+            return  # nothing to validate
+
+        # Check each override for destination mismatch
+        for override_str in overrides:
+            if override_str.startswith(
+                "export."
+            ):  # check if override starts with export.
+                # Extract destination from override path
+                try:
+                    key_part = override_str.split("=")[0]  # Get left side before =
+                    parts = key_part.split(".")
+                    if len(parts) >= 2:
+                        override_dest = parts[1]
+                        if override_dest != dest:
+                            raise ValueError(
+                                f"Override destination mismatch: override specifies 'export.{override_dest}' but --dest is '{dest}'. "
+                                f"Either change --dest to '{override_dest}' or use 'export.{dest}' in overrides."
+                            )
+                except (IndexError, AttributeError):
+                    # miconstructed override -> OmegaConf handles this
+                    pass