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

nemo_evaluator_launcher/cli/run.py

@@ -16,18 +16,32 @@
 import pathlib
 import time
 from dataclasses import dataclass
+from typing import Literal
 
-import yaml
-from omegaconf import OmegaConf
 from simple_parsing import field
 
-from nemo_evaluator_launcher.api.functional import RunConfig, run_eval
+from nemo_evaluator_launcher.common.logging_utils import logger
+from nemo_evaluator_launcher.common.printing_utils import (
+    bold,
+    cyan,
+    green,
+    magenta,
+    red,
+    yellow,
+)
 
 
 @dataclass
 class Cmd:
     """Run command parameters"""
 
+    config: str | None = field(
+        default=None,
+        alias=["--config"],
+        metadata={
+            "help": "Full path to config file. Uses Hydra by default (--config-mode=hydra). Use --config-mode=raw to load directly (bypasses Hydra)."
+        },
+    )
     config_name: str = field(
         default="default",
         alias=["-c", "--config-name"],
@@ -42,11 +56,11 @@ class Cmd:
             "help": "Path to user config directory. If provided, searches here first, then falls back to internal configs."
         },
     )
-    run_config_file: str | None = field(
-        default=None,
-        alias=["-f", "--run-config-file"],
+    config_mode: Literal["hydra", "raw"] = field(
+        default="hydra",
+        alias=["--config-mode"],
         metadata={
-            "help": "Path to a run config file to load directly (bypasses Hydra config loading)."
+            "help": "Config loading mode: 'hydra' (default) uses Hydra config system, 'raw' loads config file directly bypassing Hydra."
         },
     )
     override: list[str] = field(
@@ -63,40 +77,96 @@ class Cmd:
         alias=["-n", "--dry-run"],
         metadata={"help": "Do not run the evaluation, just print the config."},
     )
+    config_output: str | None = field(
+        default=None,
+        alias=["--config-output"],
+        metadata={
+            "help": "Directory to save the complete run config. Defaults to ~/.nemo-evaluator/run_configs/"
+        },
+    )
 
     def execute(self) -> None:
-        # Load configuration either from Hydra or from a run config file
-        if self.run_config_file:
-            # Validate that run config file is not used with other config options
+        # Import heavy dependencies only when needed
+        import yaml
+        from omegaconf import OmegaConf
+
+        from nemo_evaluator_launcher.api.functional import RunConfig, run_eval
+
+        # Validate config_mode value
+        if self.config_mode not in ["hydra", "raw"]:
+            raise ValueError(
+                f"Invalid --config-mode value: {self.config_mode}. Must be 'hydra' or 'raw'."
+            )
+
+        # Validate that raw mode requires --config
+        if self.config_mode == "raw" and self.config is None:
+            raise ValueError(
+                "--config-mode=raw requires --config to be specified. Raw mode loads config files directly."
+            )
+
+        # Load configuration either from Hydra or directly from a config file
+        if self.config_mode == "raw" and self.config:
+            # Validate that raw config loading is not used with other config options
             if self.config_name != "default":
-                raise ValueError("Cannot use --run-config-file with --config-name")
+                raise ValueError(
+                    "Cannot use --config-mode=raw with --config-name. Raw mode only works with --config."
+                )
             if self.config_dir is not None:
-                raise ValueError("Cannot use --run-config-file with --config-dir")
+                raise ValueError(
+                    "Cannot use --config-mode=raw with --config-dir. Raw mode only works with --config."
+                )
             if self.override:
-                raise ValueError("Cannot use --run-config-file with --override")
+                raise ValueError(
+                    "Cannot use --config-mode=raw with --override. Raw mode only works with --config."
+                )
 
-            # Load from run config file
-            with open(self.run_config_file, "r") as f:
+            # Load from config file directly (bypass Hydra)
+            with open(self.config, "r") as f:
                 config_dict = yaml.safe_load(f)
 
             # Create RunConfig from the loaded data
             config = OmegaConf.create(config_dict)
         else:
+            # Handle --config parameter: split path into config_dir and config_name for Hydra
+            if self.config:
+                if self.config_name != "default":
+                    raise ValueError("Cannot use --config with --config-name")
+                if self.config_dir is not None:
+                    raise ValueError("Cannot use --config with --config-dir")
+                config_path = pathlib.Path(self.config)
+                config_dir = str(config_path.parent)
+                config_name = str(config_path.stem)
+            else:
+                config_dir = self.config_dir
+                config_name = self.config_name
+
             # Load the complete Hydra configuration
             config = RunConfig.from_hydra(
-                config_name=self.config_name,
+                config_dir=config_dir,
+                config_name=config_name,
                 hydra_overrides=self.override,
-                config_dir=self.config_dir,
             )
 
-        invocation_id = run_eval(config, self.dry_run)
+        try:
+            invocation_id = run_eval(config, self.dry_run)
+        except Exception as e:
+            print(red(f"✗ Job submission failed, see logs | Error: {e}"))
+            logger.error("Job submission failed", error=e)
+            raise
 
-        # Save the complete configuration to the raw_configs directory
+        # Save the complete configuration
         if not self.dry_run and invocation_id is not None:
-            # Create ~/.nemo-evaluator/run_configs directory
-            home_dir = pathlib.Path.home()
-            run_configs_dir = home_dir / ".nemo-evaluator" / "run_configs"
-            run_configs_dir.mkdir(parents=True, exist_ok=True)
+            # Determine config output directory
+            if self.config_output:
+                # Use custom directory specified by --config-output
+                config_dir = pathlib.Path(self.config_output)
+            else:
+                # Default to original location: ~/.nemo-evaluator/run_configs
+                home_dir = pathlib.Path.home()
+                config_dir = home_dir / ".nemo-evaluator" / "run_configs"
+
+            # Ensure the directory exists
+            config_dir.mkdir(parents=True, exist_ok=True)
 
             # Convert DictConfig to dict and save as YAML
             config_dict = OmegaConf.to_container(config, resolve=True)
@@ -106,7 +176,7 @@ class Cmd:
 
             # Create config filename with invocation ID
             config_filename = f"{invocation_id}_config.yml"
-            config_path = run_configs_dir / config_filename
+            config_path = config_dir / config_filename
 
             # Save the complete Hydra configuration
             with open(config_path, "w") as f:
@@ -120,16 +190,69 @@ class Cmd:
                 f.write("#\n")
                 f.write("# To rerun this exact configuration:\n")
                 f.write(
-                    f"# nemo-evaluator-launcher run --run-config-file {config_path}\n"
+                    f"# nemo-evaluator-launcher run --config {config_path} --config-mode=raw\n"
                 )
                 f.write("#\n")
                 f.write(config_yaml)
 
-            print(f"Complete run config saved to: {config_path}")
+            print(bold(cyan("Complete run config saved to: ")) + f"\n  {config_path}\n")
+            logger.info("Saved complete config", path=config_path)
 
-        if invocation_id is not None:
-            print(f"to check status: nemo-evaluator-launcher status {invocation_id}")
-            print(f"to kill all jobs: nemo-evaluator-launcher kill {invocation_id}")
+        # Print general success message with invocation ID and helpful commands
+        if invocation_id is not None and not self.dry_run:
             print(
-                f"to kill individual jobs: nemo-evaluator-launcher kill <job_id> (e.g., {invocation_id}.0)"
+                bold(cyan("To check status: "))
+                + f"nemo-evaluator-launcher status {invocation_id}"
+            )
+            print(
+                bold(cyan("To view job info: "))
+                + f"nemo-evaluator-launcher info {invocation_id}"
+            )
+            print(
+                bold(cyan("To kill all jobs: "))
+                + f"nemo-evaluator-launcher kill {invocation_id}"
+            )
+
+            # Show actual job IDs and task names
+            print(bold(cyan("To kill individual jobs:")))
+            # Access tasks - will work after normalization in run_eval
+            tasks = (
+                config.evaluation.tasks
+                if hasattr(config.evaluation, "tasks")
+                else config.evaluation
+            )
+            for idx, task in enumerate(tasks):
+                job_id = f"{invocation_id}.{idx}"
+                print(f"  nemo-evaluator-launcher kill {job_id}  # {task.name}")
+
+            print(
+                magenta(
+                    "(all commands accept shortened IDs as long as there are no conflicts)"
+                )
+            )
+            print(
+                bold(cyan("To print all jobs: ")) + "nemo-evaluator-launcher ls runs"
+                "\n  (--since 1d or --since 6h for time span, see --help)"
+            )
+
+            print(
+                green(
+                    bold(
+                        f"✓ Job submission successful | Invocation ID: {invocation_id}"
+                    )
+                )
+            )
+
+        # Warn if both config_dir and config_name are provided (and config_name is not default)
+        if (
+            self.config is None
+            and self.config_dir is not None
+            and self.config_name != "default"
+        ):
+            joint_path = pathlib.Path(self.config_dir) / f"{self.config_name}.yaml"
+            print(
+                yellow(
+                    f"Warning: Using --config-dir and --config-name together is deprecated. "
+                    f"Please use --config {joint_path} instead."
+                )
             )

nemo_evaluator_launcher/cli/status.py

@@ -13,12 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import json
 from dataclasses import dataclass
 
 from simple_parsing import field
 
-from nemo_evaluator_launcher.api.functional import get_status
+import nemo_evaluator_launcher.common.printing_utils as pu
+from nemo_evaluator_launcher.executors.base import ExecutionState
 
 
 @dataclass
@@ -36,6 +36,11 @@ class Cmd:
     )
 
     def execute(self) -> None:
+        # Import heavy dependencies only when needed
+        import json
+
+        from nemo_evaluator_launcher.api.functional import get_status
+
         res = get_status(self.job_ids)
         if self.json:
             # Remove progress field from JSON output as it's a WIP feature
@@ -94,10 +99,16 @@ class Cmd:
             else:
                 location = ""
 
+            # Format status with visual indicators and colors
+            status = job.get("status", "")
+            formatted_status = self._format_status_with_indicators(status)
+
+            # Extract task name
+
             rows.append(
                 [
                     job.get("job_id", ""),
-                    job.get("status", ""),
+                    formatted_status,
                     # job.get("progress", ""), temporarily disabled as this is a WIP feature
                     executor_info,
                     location,
@@ -106,7 +117,10 @@ class Cmd:
 
         # Calculate column widths and print
         widths = [
-            max(len(str(headers[i])), max(len(str(row[i])) for row in rows))
+            max(
+                len(str(headers[i])),
+                max(len(self._strip_ansi_codes(str(row[i]))) for row in rows),
+            )
             for i in range(len(headers))
         ]
 
@@ -117,4 +131,34 @@ class Cmd:
         print("-" * len(header_row))
 
         for row in rows:
-            print(" | ".join(str(row[i]).ljust(widths[i]) for i in range(len(row))))
+            # Adjust padding for ANSI color codes
+            formatted_row = []
+            for i in range(len(row)):
+                content = str(row[i])
+                visible_length = len(self._strip_ansi_codes(content))
+                padding = widths[i] - visible_length
+                formatted_row.append(content + " " * padding)
+            print(" | ".join(formatted_row))
+
+    def _format_status_with_indicators(self, status: str) -> str:
+        """Format status with Unicode visual indicators only."""
+        # Status mapping based on ExecutionState enum
+        status_formats = {
+            ExecutionState.SUCCESS.value: pu.green("✓ SUCCESS"),
+            ExecutionState.FAILED.value: pu.red("✗ FAILED"),
+            ExecutionState.RUNNING.value: pu.yellow("▶ RUNNING"),
+            ExecutionState.PENDING.value: pu.cyan("⧗ PENDING"),
+            ExecutionState.KILLED.value: pu.magenta("✗ KILLED"),
+            # Additional states for error handling
+            "not_found": pu.grey("? NOT FOUND"),
+            "error": pu.red("✗ ERROR"),
+        }
+
+        return status_formats.get(status.lower(), pu.grey(status.upper()))
+
+    def _strip_ansi_codes(self, text: str) -> str:
+        """Remove ANSI color codes from text for length calculation."""
+        import re
+
+        ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+        return ansi_escape.sub("", text)

nemo_evaluator_launcher/cli/version.py

@@ -19,6 +19,29 @@ import importlib
 from dataclasses import dataclass
 
 from nemo_evaluator_launcher import __package_name__, __version__
+from nemo_evaluator_launcher.common.logging_utils import logger
+
+
+def get_versions() -> dict:
+    internal_module_name = "nemo_evaluator_launcher_internal"
+    res = {__package_name__: __version__}
+    # Check for internal package
+    try:
+        internal_module = importlib.import_module(internal_module_name)
+        # Try to get version from internal package
+        internal_version = getattr(internal_module, "__version__", None)
+        if internal_version:
+            res[internal_module_name] = internal_version
+        else:
+            res[internal_module_name] = "available (version unknown)"
+    except ImportError:
+        # Internal package not available - this is expected in many cases
+        pass
+    except Exception as e:
+        logger.error(f"nemo_evaluator_launcher_internal: error loading ({e})")
+        raise
+
+    return res
 
 
 @dataclass
@@ -27,26 +50,6 @@ class Cmd:
 
     def execute(self) -> None:
         """Execute the version command."""
-        print(f"{__package_name__}: {__version__}")
-
-        # Check for internal package
-        try:
-            internal_module = importlib.import_module(
-                "nemo_evaluator_launcher_internal"
-            )
-            # Try to get version from internal package
-            try:
-                internal_version = getattr(internal_module, "__version__", None)
-                if internal_version:
-                    print(f"nemo-evaluator-launcher-internal: {internal_version}")
-                else:
-                    print(
-                        "nemo-evaluator-launcher-internal: available (version unknown)"
-                    )
-            except Exception:
-                print("nemo-evaluator-launcher-internal: available (version unknown)")
-        except ImportError:
-            # Internal package not available - this is expected in many cases
-            pass
-        except Exception as e:
-            print(f"nemo-evaluator-launcher-internal: error loading ({e})")
+        res = get_versions()
+        for package, version in res.items():
+            print(f"{package}: {version}")

nemo_evaluator_launcher/common/execdb.py

@@ -29,15 +29,15 @@ EXEC_DB_FILE = EXEC_DB_DIR / "exec.v1.jsonl"
 
 
 def generate_invocation_id() -> str:
-    """Generate a unique invocation ID as an 8-digit hex string."""
-    return secrets.token_hex(4)
+    """Generate a unique invocation ID as an 16-digit hex string."""
+    return secrets.token_hex(8)
 
 
 def generate_job_id(invocation_id: str, index: int) -> str:
     """Generate a job ID as <invocation_id>.<n>.
 
     Args:
-        invocation_id: The invocation group ID (8-digit hex).
+        invocation_id: The invocation group ID (16-digit hex).
         index: The job index (0-based integer).
     Returns:
         The job ID string.
@@ -50,7 +50,7 @@ class JobData:
     """Data structure for job execution information.
 
     Attributes:
-        invocation_id: 8-digit hex string.
+        invocation_id: 16-digit hex string.
         job_id: <invocation_id>.<n> string.
         timestamp: Unix timestamp when the job was created.
         executor: Name of the executor that handled this job.
@@ -148,41 +148,135 @@ class ExecutionDB:
             )
             raise
 
+    def _resolve_invocation_id(self, short_id: str) -> Optional[str]:
+        """Resolve a short invocation ID to the full one.
+
+        Args:
+            short_id: Partial or full invocation ID.
+
+        Returns:
+            Full invocation ID if found uniquely, None if not found.
+
+        Raises:
+            ValueError: If the short_id matches multiple invocation IDs.
+        """
+        if not short_id:
+            return None
+
+        short_id = short_id.lower()
+
+        # NOTE(agronskiy): this is a non-optimized implementation that assumes small amount
+        # of jobs in ExecDB(), a typical scenario. Speeding up would involve building a
+        # prefix tree when loading invocations/jobs.
+        matches = [
+            inv_id
+            for inv_id in self._invocations.keys()
+            if inv_id.lower().startswith(short_id)
+        ]
+
+        if len(matches) == 1:
+            return matches[0]
+        elif len(matches) > 1:
+            raise ValueError(f"Ambiguous invocation ID '{short_id}': matches {matches}")
+        else:
+            return None
+
+    def _resolve_job_id(self, short_job_id: str) -> Optional[str]:
+        """Resolve a short job ID to the full one.
+
+        Args:
+            short_job_id: Partial or full job ID.
+
+        Returns:
+            Full job ID if found uniquely, None if not found.
+
+        Raises:
+            ValueError: If the short_job_id matches multiple job IDs.
+        """
+        if not short_job_id:
+            return None
+
+        # Normalize to lowercase for case-insensitive matching
+        short_job_id = short_job_id.lower()
+
+        if "." in short_job_id:
+            parts = short_job_id.split(".", 1)
+            short_inv_id, job_index = parts[0], parts[1]
+
+            # Resolve the invocation part
+            full_inv_id = self._resolve_invocation_id(short_inv_id)
+            if full_inv_id:
+                candidate_job_id = f"{full_inv_id}.{job_index}"
+                if candidate_job_id in self._jobs:
+                    return candidate_job_id
+
+        # NOTE(agronskiy): unfortunately, due to legacy, there exist usecases where
+        # job_id is the same format as invocation_id
+        candidate_job_id = self._resolve_invocation_id(short_job_id)
+        if candidate_job_id and candidate_job_id in self._jobs:
+            return candidate_job_id
+
+        return None
+
     def get_job(self, job_id: str) -> Optional[JobData]:
-        return self._jobs.get(job_id)
+        """Get job by full or partial job ID.
 
-    def get_jobs(self, invocation_id: str) -> Dict[str, JobData]:
-        job_ids = self._invocations.get(invocation_id, [])
-        return {
-            job_id: self._jobs[job_id] for job_id in job_ids if job_id in self._jobs
-        }
+        Args:
+            job_id: Full or partial job ID.
 
-    def get_invocation_jobs(self, invocation_id: str) -> List[str]:
-        return self._invocations.get(invocation_id, [])
+        Returns:
+            JobData if found, None otherwise.
 
-    def get_all_jobs(self) -> Dict[str, JobData]:
-        """Return a copy of all jobs in the execution DB."""
-        return dict(self._jobs)
+        Raises:
+            ValueError: If the job_id matches multiple jobs.
+        """
+        resolved_id = self._resolve_job_id(job_id)
+        if resolved_id:
+            return self._jobs.get(resolved_id)
 
+        return None
 
-def write_job(job: JobData) -> None:
-    db = ExecutionDB()
-    db.write_job(job)
+    def get_jobs(self, invocation_id: str) -> Dict[str, JobData]:
+        """Get all jobs for a full or partial invocation ID.
 
+        Args:
+            invocation_id: Full or partial invocation ID.
 
-def get_job(job_id: str) -> Optional[JobData]:
-    db = ExecutionDB()
-    return db.get_job(job_id)
+        Returns:
+            Dictionary mapping job_id to JobData for all jobs in the invocation.
 
+        Raises:
+            ValueError: If the invocation_id matches multiple invocations.
+        """
+        resolved_inv_id = self._resolve_invocation_id(invocation_id)
+        if not resolved_inv_id:
+            return {}
 
-def get_jobs(invocation_id: str) -> Dict[str, JobData]:
-    db = ExecutionDB()
-    return db.get_jobs(invocation_id)
+        job_ids = self._invocations.get(resolved_inv_id, [])
+        return {
+            job_id: self._jobs[job_id] for job_id in job_ids if job_id in self._jobs
+        }
+
+    def get_invocation_jobs(self, invocation_id: str) -> List[str]:
+        """Get job IDs for a full or partial invocation ID.
+
+        Args:
+            invocation_id: Full or partial invocation ID.
 
+        Returns:
+            List of job IDs for the invocation.
 
-def get_all_jobs() -> Dict[str, JobData]:
-    db = ExecutionDB()
-    return db.get_all_jobs()
+        Raises:
+            ValueError: If the invocation_id matches multiple invocations.
+        """
+        resolved_inv_id = self._resolve_invocation_id(invocation_id)
+        if not resolved_inv_id:
+            return []
+        return self._invocations.get(resolved_inv_id, [])
+
+    def get_all_jobs(self) -> Dict[str, JobData]:
+        """Return a copy of all jobs in the execution DB."""
+        return dict(self._jobs)
 
 
 # Ensure all the paths