some refactoring; add dataset sizes to log messages

Author: johnnygreco

src/data_designer/cli/controllers/download_controller.py

@@ -4,7 +4,7 @@
 import subprocess
 from pathlib import Path
 
-from data_designer.cli.services.download_service import DownloadService
+from data_designer.cli.services.download_service import DATASET_SIZES, DownloadService
 from data_designer.cli.ui import (
     confirm_action,
     console,
@@ -15,6 +15,7 @@
     print_text,
     select_multiple_with_arrows,
 )
+from data_designer.cli.utils import check_ngc_cli_available, get_ngc_version
 
 NGC_URL = "https://catalog.ngc.nvidia.com/"
 NGC_CLI_INSTALL_URL = "https://org.ngc.nvidia.com/setup/installers/cli"
@@ -40,8 +41,8 @@ def run_personas(self, locales: list[str] | None, all_locales: bool, dry_run: bo
         print_info(f"Datasets will be saved to: {self.service.get_managed_assets_directory()}")
         console.print()
 
-        # Check NGC CLI availability (skip in dry run mode)
-        if not dry_run and not self._check_ngc_cli():
+        # Check NGC CLI availability (skip checking in dry run mode)
+        if not dry_run and not check_ngc_cli_with_instructions():
             return
 
         # Determine which locales to download
@@ -57,8 +58,8 @@ def run_personas(self, locales: list[str] | None, all_locales: bool, dry_run: bo
         print_text(f"📦 {action} {len(selected_locales)} Nemotron-Persona dataset(s):")
         for locale in selected_locales:
             already_downloaded = self.service.is_locale_downloaded(locale)
-            status = " (already exists, will update)" if already_downloaded else ""
-            print_text(f"  • {locale}{status}")
+            status = " - already exists, will update" if already_downloaded else ""
+            print_text(f"  • {locale} ({DATASET_SIZES[locale]}){status}")
 
         console.print()
 
@@ -87,29 +88,11 @@ def run_personas(self, locales: list[str] | None, all_locales: bool, dry_run: bo
         console.print()
         if successful:
             print_success(f"Successfully downloaded {len(successful)} dataset(s): {', '.join(successful)}")
-            print_info(f"Location: {self.service.get_managed_assets_directory()}")
+            print_info(f"Saved datasets to: {self.service.get_managed_assets_directory()}")
 
         if failed:
             print_error(f"Failed to download {len(failed)} dataset(s): {', '.join(failed)}")
 
-    def _check_ngc_cli(self) -> bool:
-        """Check if NGC CLI is installed and guide user if not."""
-        if self.service.check_ngc_cli_available():
-            version = self.service.get_ngc_version()
-            if version:
-                print_info(f"NGC CLI: {version}")
-            return True
-
-        print_error("NGC CLI not found!")
-        console.print()
-        print_text("The NGC CLI is required to download the Nemotron-Personas datasets.")
-        console.print()
-        print_text("To download the Nemotron-Personas datasets, follow these steps:")
-        print_text(f"    1. Create an NVIDIA NGC account: {NGC_URL}")
-        print_text(f"    2. Install the NGC CLI: {NGC_CLI_INSTALL_URL}")
-        print_text("    3. Following the install instructions to set up the NGC CLI")
-        print_text("    4. Run 'data-designer download personas'")
-
     def _determine_locales(self, locales: list[str] | None, all_locales: bool) -> list[str]:
         """Determine which locales to download based on user input.
 
@@ -190,3 +173,23 @@ def _download_locale(self, locale: str) -> bool:
             print_error(f"✗ Failed to download Nemotron-Persona dataset for {locale}")
             print_error(f"Unexpected error: {e}")
             return False
+
+
+def check_ngc_cli_with_instructions() -> bool:
+    """Check if NGC CLI is installed and guide user if not."""
+    if check_ngc_cli_available():
+        version = get_ngc_version()
+        if version:
+            print_info(f"NGC CLI: {version}")
+        return True
+
+    print_error("NGC CLI not found!")
+    console.print()
+    print_text("The NGC CLI is required to download the Nemotron-Personas datasets.")
+    console.print()
+    print_text("To download the Nemotron-Personas datasets, follow these steps:")
+    print_text(f"    1. Create an NVIDIA NGC account: {NGC_URL}")
+    print_text(f"    2. Install the NGC CLI: {NGC_CLI_INSTALL_URL}")
+    print_text("    3. Following the install instructions to set up the NGC CLI")
+    print_text("    4. Run 'data-designer download personas'")
+    return False

src/data_designer/cli/services/download_service.py

@@ -7,8 +7,15 @@
 import tempfile
 from pathlib import Path
 
+DATASET_SIZES = {
+    "en_US": "1.24 GB",
+    "en_IN": "2.39 GB",
+    "hi_Deva_IN": "4.14 GB",
+    "hi_Latn_IN": "2.7 GB",
+    "ja_JP": "1.69 GB",
+}
+SUPPORTED_LOCALES = list[str](DATASET_SIZES.keys())
 DATASET_PREFIX = "nemotron-personas-dataset-"
-SUPPORTED_LOCALES = ["en_US", "en_IN", "hi_Deva_IN", "hi_Latn_IN", "ja_JP"]
 
 
 class DownloadService:
@@ -18,31 +25,6 @@ def __init__(self, config_dir: Path):
         self.config_dir = config_dir
         self.managed_assets_dir = config_dir / "managed-assets" / "datasets"
 
-    def check_ngc_cli_available(self) -> bool:
-        """Check if NGC CLI is installed and available.
-
-        Returns:
-            True if NGC CLI is in PATH and executable, False otherwise.
-        """
-        if shutil.which("ngc") is None:
-            return False
-
-        return self.get_ngc_version() is not None
-
-    def get_ngc_version(self) -> str | None:
-        """Get the NGC CLI version if available."""
-        try:
-            result = subprocess.run(
-                ["ngc", "--version"],
-                capture_output=True,
-                text=True,
-                check=True,
-                timeout=5,
-            )
-            return result.stdout.strip()
-        except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
-            return None
-
     def get_available_locales(self) -> dict[str, str]:
         """Get dictionary of available persona locales (locale code -> locale code)."""
         return {locale: locale for locale in SUPPORTED_LOCALES}
@@ -73,14 +55,14 @@ def download_persona_dataset(self, locale: str) -> Path:
                 "registry",
                 "resource",
                 "download-version",
-                f"nvidia/nemotron-personas/{_get_dataset_name(locale)}",
+                f"nvidia/nemotron-personas/{_get_downloaded_dataset_name(locale)}",
                 "--dest",
                 temp_dir,
             ]
 
             subprocess.run(cmd, check=True)
 
-            dataset_pattern = _get_dataset_name(locale)
+            dataset_pattern = _get_downloaded_dataset_name(locale)
             download_pattern = f"{temp_dir}/{dataset_pattern}*/*.parquet"
             parquet_files = glob.glob(download_pattern)
 
@@ -114,16 +96,14 @@ def is_locale_downloaded(self, locale: str) -> bool:
         if not self.managed_assets_dir.exists():
             return False
 
-        # Check for parquet files matching this locale in managed assets
-        dataset_pattern = _get_dataset_name(locale)
         # Look for any parquet files that start with the dataset pattern
-        parquet_files = glob.glob(str(self.managed_assets_dir / f"{dataset_pattern}*.parquet"))
+        parquet_files = glob.glob(str(self.managed_assets_dir / f"{locale}.parquet"))
 
         return len(parquet_files) > 0
 
 
-def _get_dataset_name(locale: str) -> str:
-    """Build dataset name pattern for the given locale.
+def _get_downloaded_dataset_name(locale: str) -> str:
+    """Build the downloaded dataset name pattern for the given locale.
 
     Args:
         locale: Locale code (e.g., 'en_US', 'ja_JP')

src/data_designer/cli/utils.py

@@ -1,6 +1,40 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import shutil
+import subprocess
+
+
+def check_ngc_cli_available() -> bool:
+    """Check if NGC CLI is installed and available.
+
+    Returns:
+        True if NGC CLI is in PATH and executable, False otherwise.
+    """
+    if shutil.which("ngc") is None:
+        return False
+
+    return get_ngc_version() is not None
+
+
+def get_ngc_version() -> str | None:
+    """Get the NGC CLI version if available.
+
+    Returns:
+        NGC CLI version string if available, None otherwise.
+    """
+    try:
+        result = subprocess.run(
+            ["ngc", "--version"],
+            capture_output=True,
+            text=True,
+            check=True,
+            timeout=5,
+        )
+        return result.stdout.strip()
+    except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
+        return None
+
 
 def validate_url(url: str) -> bool:
     """Validate that a string is a valid URL.

tests/cli/controllers/test_download_controller.py

@@ -23,7 +23,7 @@ def controller_with_datasets(tmp_path: Path) -> DownloadController:
     # Create managed assets directory with sample parquet files
     managed_assets_dir = tmp_path / "managed-assets" / "datasets"
     managed_assets_dir.mkdir(parents=True, exist_ok=True)
-    (managed_assets_dir / "nemotron-personas-dataset-en_us_0.parquet").touch()
+    (managed_assets_dir / "en_US.parquet").touch()
     return controller
 
 
@@ -36,7 +36,7 @@ def test_init(tmp_path: Path) -> None:
 
 @patch("data_designer.cli.controllers.download_controller.confirm_action", return_value=False)
 @patch("data_designer.cli.controllers.download_controller.select_multiple_with_arrows", return_value=["en_US"])
-@patch.object(DownloadController, "_check_ngc_cli", return_value=True)
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions", return_value=True)
 def test_run_personas_user_cancels_confirmation(
     mock_check_ngc: MagicMock,
     mock_select: MagicMock,
@@ -58,7 +58,7 @@ def test_run_personas_user_cancels_confirmation(
 
 @patch.object(DownloadController, "_download_locale", return_value=True)
 @patch("data_designer.cli.controllers.download_controller.confirm_action", return_value=True)
-@patch.object(DownloadController, "_check_ngc_cli", return_value=True)
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions", return_value=True)
 def test_run_personas_with_all_flag(
     mock_check_ngc: MagicMock,
     mock_confirm: MagicMock,
@@ -85,7 +85,7 @@ def test_run_personas_with_all_flag(
 
 @patch.object(DownloadController, "_download_locale", return_value=True)
 @patch("data_designer.cli.controllers.download_controller.confirm_action", return_value=True)
-@patch.object(DownloadController, "_check_ngc_cli", return_value=True)
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions", return_value=True)
 def test_run_personas_with_specific_locales(
     mock_check_ngc: MagicMock,
     mock_confirm: MagicMock,
@@ -105,7 +105,7 @@ def test_run_personas_with_specific_locales(
     assert "ja_JP" in downloaded_locales
 
 
-@patch.object(DownloadController, "_check_ngc_cli", return_value=True)
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions", return_value=True)
 def test_run_personas_with_invalid_locales(
     mock_check_ngc: MagicMock,
     controller: DownloadController,
@@ -122,7 +122,7 @@ def test_run_personas_with_invalid_locales(
 @patch.object(DownloadController, "_download_locale", return_value=True)
 @patch("data_designer.cli.controllers.download_controller.confirm_action", return_value=True)
 @patch("data_designer.cli.controllers.download_controller.select_multiple_with_arrows", return_value=["en_US"])
-@patch.object(DownloadController, "_check_ngc_cli", return_value=True)
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions", return_value=True)
 def test_run_personas_interactive_selection(
     mock_check_ngc: MagicMock,
     mock_select: MagicMock,
@@ -147,7 +147,7 @@ def test_run_personas_interactive_selection(
 
 
 @patch("data_designer.cli.controllers.download_controller.select_multiple_with_arrows", return_value=None)
-@patch.object(DownloadController, "_check_ngc_cli", return_value=True)
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions", return_value=True)
 def test_run_personas_interactive_cancelled(
     mock_check_ngc: MagicMock,
     mock_select: MagicMock,
@@ -165,7 +165,7 @@ def test_run_personas_interactive_cancelled(
     # Function should exit early
 
 
-@patch.object(DownloadController, "_check_ngc_cli", return_value=False)
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions", return_value=False)
 def test_run_personas_ngc_cli_not_available(
     mock_check_ngc: MagicMock,
     controller: DownloadController,
@@ -177,20 +177,24 @@ def test_run_personas_ngc_cli_not_available(
     mock_check_ngc.assert_called_once()
 
 
-def test_check_ngc_cli_available_with_version(controller: DownloadController) -> None:
-    """Test _check_ngc_cli displays version when NGC CLI is available."""
-    with patch.object(controller.service, "check_ngc_cli_available", return_value=True):
-        with patch.object(controller.service, "get_ngc_version", return_value="NGC CLI 3.41.4"):
-            result = controller._check_ngc_cli()
+def test_check_ngc_cli_available_with_version() -> None:
+    """Test check_ngc_cli_with_instructions displays version when NGC CLI is available."""
+    from data_designer.cli.controllers.download_controller import check_ngc_cli_with_instructions
+
+    with patch("data_designer.cli.controllers.download_controller.check_ngc_cli_available", return_value=True):
+        with patch("data_designer.cli.controllers.download_controller.get_ngc_version", return_value="NGC CLI 3.41.4"):
+            result = check_ngc_cli_with_instructions()
 
     assert result is True
 
 
-def test_check_ngc_cli_available_without_version(controller: DownloadController) -> None:
-    """Test _check_ngc_cli when version cannot be determined."""
-    with patch.object(controller.service, "check_ngc_cli_available", return_value=True):
-        with patch.object(controller.service, "get_ngc_version", return_value=None):
-            result = controller._check_ngc_cli()
+def test_check_ngc_cli_available_without_version() -> None:
+    """Test check_ngc_cli_with_instructions when version cannot be determined."""
+    from data_designer.cli.controllers.download_controller import check_ngc_cli_with_instructions
+
+    with patch("data_designer.cli.controllers.download_controller.check_ngc_cli_available", return_value=True):
+        with patch("data_designer.cli.controllers.download_controller.get_ngc_version", return_value=None):
+            result = check_ngc_cli_with_instructions()
 
     assert result is True
 
@@ -276,7 +280,7 @@ def test_download_locale_generic_error(controller: DownloadController) -> None:
 
 @patch.object(DownloadController, "_download_locale")
 @patch("data_designer.cli.controllers.download_controller.confirm_action", return_value=True)
-@patch.object(DownloadController, "_check_ngc_cli", return_value=True)
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions", return_value=True)
 def test_run_personas_mixed_success_and_failure(
     mock_check_ngc: MagicMock,
     mock_confirm: MagicMock,
@@ -295,7 +299,7 @@ def test_run_personas_mixed_success_and_failure(
 
 @patch.object(DownloadController, "_download_locale", return_value=True)
 @patch("data_designer.cli.controllers.download_controller.confirm_action", return_value=True)
-@patch.object(DownloadController, "_check_ngc_cli", return_value=True)
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions", return_value=True)
 def test_run_personas_shows_existing_status(
     mock_check_ngc: MagicMock,
     mock_confirm: MagicMock,
@@ -311,7 +315,7 @@ def test_run_personas_shows_existing_status(
 
 @patch.object(DownloadController, "_download_locale")
 @patch("data_designer.cli.controllers.download_controller.confirm_action")
-@patch.object(DownloadController, "_check_ngc_cli")
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions")
 def test_run_personas_with_dry_run_flag(
     mock_check_ngc: MagicMock,
     mock_confirm: MagicMock,
@@ -333,7 +337,7 @@ def test_run_personas_with_dry_run_flag(
 
 @patch.object(DownloadController, "_download_locale")
 @patch("data_designer.cli.controllers.download_controller.confirm_action")
-@patch.object(DownloadController, "_check_ngc_cli")
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions")
 def test_run_personas_with_all_and_dry_run(
     mock_check_ngc: MagicMock,
     mock_confirm: MagicMock,
@@ -356,7 +360,7 @@ def test_run_personas_with_all_and_dry_run(
 @patch.object(DownloadController, "_download_locale")
 @patch("data_designer.cli.controllers.download_controller.confirm_action")
 @patch("data_designer.cli.controllers.download_controller.select_multiple_with_arrows", return_value=["en_US"])
-@patch.object(DownloadController, "_check_ngc_cli")
+@patch("data_designer.cli.controllers.download_controller.check_ngc_cli_with_instructions")
 def test_run_personas_interactive_with_dry_run(
     mock_check_ngc: MagicMock,
     mock_select: MagicMock,