nemo-platform-plugin 0.1.0__py3-none-any.whl

nemo_platform/beta/safe_synthesizer/job_builder.py

@@ -0,0 +1,355 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+import random
+import string
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import pandas as pd
+from typing_extensions import Self
+
+from .job import SafeSynthesizerJob
+
+if TYPE_CHECKING:
+    from nemo_platform import NeMoPlatform
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Type alias for the config inputs accepted by each with_* method:
+# a Pydantic model, a plain dict, or None (use service-side defaults).
+_ConfigInput = Any
+
+
+def _to_dict(config: _ConfigInput) -> dict:
+    """Coerce a Pydantic model, dict, or None to a plain dict."""
+    if config is None:
+        return {}
+    if hasattr(config, "model_dump"):
+        return config.model_dump()
+    if isinstance(config, dict):
+        return config.copy()
+    raise ValueError(f"Config must be a dict, a Pydantic model, or None; got {type(config)!r}")
+
+
+def _merge_config(config: _ConfigInput, kwargs: dict) -> dict:
+    base = _to_dict(config)
+    base.update(kwargs)
+    return base
+
+
+class SafeSynthesizerJobBuilder:
+    """Builder for Safe Synthesizer Jobs ran with the NeMo Platform.
+
+    This class provides a fluent interface for building Safe Synthesizer configurations.
+    It allows you to configure all the parameters needed to create and run a Safe Synthesizer
+    job. Each method returns the builder instance to allow method chaining.
+
+    Examples:
+        >>> from nemo_platform import NeMoPlatform
+        >>> from nemo_platform.beta.safe_synthesizer.job_builder import SafeSynthesizerJobBuilder
+        >>> client = NeMoPlatform(base_url=..., inference_base_url=...)
+        >>> builder = (
+        ...     SafeSynthesizerJobBuilder(client)
+        ...     .with_data_source(your_dataframe)
+        ...     .with_replace_pii()
+        ...     .synthesize()
+        ...     .with_train(learning_rate=0.0001)
+        ...     .with_generate(num_records=10000)
+        ...     .with_evaluate(enable=False)
+        ... )
+        >>> job = builder.create_job()
+    """
+
+    def __init__(self, client: NeMoPlatform, workspace: str = "default"):
+        self._client = client
+        self._workspace = workspace
+
+        # Job-level fields
+        self._hf_token_secret: str | None = None
+        self._classify_model_provider: str | None = None
+        self._data_source: pd.DataFrame | str | None = None
+        self._data_source_path: str | None = None
+
+        # Enable flags
+        self._enable_synthesis: bool = False
+        self._enable_replace_pii: bool = False
+
+        # Per-section config dicts (populated by with_* methods)
+        self._data_config: dict = {}
+        self._training_config: dict = {}
+        self._generation_config: dict = {}
+        self._evaluation_config: dict = {}
+        self._privacy_config: dict | None = None
+        self._time_series_config: dict = {}
+        self._replace_pii_config: dict = {}
+
+    # ------------------------------------------------------------------
+    # Data source
+    # ------------------------------------------------------------------
+
+    def with_data_source(self, df_source: pd.DataFrame | str) -> Self:
+        """Set the data source for synthetic data generation.
+
+        Args:
+            df_source: Training dataset as a pandas DataFrame or a fetchable URL.
+
+        Returns:
+            The builder instance for method chaining.
+        """
+        self._data_source = df_source
+        self._data_source_path = None  # invalidate any cached upload
+        return self
+
+    # ------------------------------------------------------------------
+    # Mode flags
+    # ------------------------------------------------------------------
+
+    def synthesize(self) -> Self:
+        """Enable data synthesis for the job run."""
+        self._enable_synthesis = True
+        return self
+
+    # ------------------------------------------------------------------
+    # Section configs
+    # ------------------------------------------------------------------
+
+    def with_data(self, config: _ConfigInput = None, **kwargs) -> Self:
+        """Configure data parameters."""
+        self._data_config = _merge_config(config, kwargs)
+        return self
+
+    def with_train(self, config: _ConfigInput = None, **kwargs) -> Self:
+        """Configure training hyperparameters.
+
+        Calling this method also enables synthesis.
+        """
+        self._training_config = _merge_config(config, kwargs)
+        self._enable_synthesis = True
+        return self
+
+    def with_generate(self, config: _ConfigInput = None, **kwargs) -> Self:
+        """Configure generation parameters.
+
+        Calling this method also enables synthesis.
+        """
+        self._generation_config = _merge_config(config, kwargs)
+        self._enable_synthesis = True
+        return self
+
+    def with_evaluate(self, config: _ConfigInput = None, **kwargs) -> Self:
+        """Configure evaluation parameters."""
+        self._evaluation_config = _merge_config(config, kwargs)
+        return self
+
+    def with_differential_privacy(self, config: _ConfigInput = None, **kwargs) -> Self:
+        """Configure differential privacy parameters."""
+        self._privacy_config = _merge_config(config, kwargs)
+        return self
+
+    def with_time_series(self, config: _ConfigInput = None, **kwargs) -> Self:
+        """Configure time-series parameters."""
+        self._time_series_config = _merge_config(config, kwargs)
+        return self
+
+    def with_replace_pii(self, config: _ConfigInput = None, **kwargs) -> Self:
+        """Configure PII replacement.
+
+        Calling this method enables PII replacement for the job. If no config is
+        provided the service will use its own defaults.
+
+        Args:
+            config: PII replacement config as a dict or Pydantic model, or None to
+                use service-side defaults.
+            **kwargs: Individual PII replacement parameters to override.
+
+        Returns:
+            The builder instance for method chaining.
+        """
+        self._replace_pii_config = _merge_config(config, kwargs)
+        self._enable_replace_pii = True
+        return self
+
+    # ------------------------------------------------------------------
+    # Classify model provider
+    # ------------------------------------------------------------------
+
+    def with_classify_model_provider(self, provider_name: str) -> Self:
+        """Configure column classification using an Inference Gateway model provider.
+
+        The model provider should be configured to serve an LLM suitable for column
+        classification tasks.
+
+        Args:
+            provider_name: Name of the model provider. Provide just the name to use a
+                provider in the current workspace (e.g. ``"my-classify-llm"``), or a
+                fully-qualified ``workspace/provider_name`` reference to use a provider
+                from a different workspace.
+
+        Returns:
+            The builder instance for method chaining.
+        """
+        if "/" in provider_name:
+            self._classify_model_provider = provider_name
+        else:
+            self._classify_model_provider = f"{self._workspace}/{provider_name}"
+        logger.info(f"Configured classify model provider: {self._classify_model_provider}")
+        return self
+
+    # ------------------------------------------------------------------
+    # HuggingFace token
+    # ------------------------------------------------------------------
+
+    def with_hf_token_secret(self, secret_name: str) -> Self:
+        """Configure HuggingFace authentication using a platform secret.
+
+        The secret must exist in the same workspace as the job and should contain
+        a valid HuggingFace token.
+
+        Args:
+            secret_name: Name of the platform secret containing the HuggingFace token.
+
+        Returns:
+            The builder instance for method chaining.
+        """
+        self._hf_token_secret = secret_name
+        return self
+
+    # ------------------------------------------------------------------
+    # Internal resolution
+    # ------------------------------------------------------------------
+
+    def _resolve_datasource(self, **kwargs) -> None:
+        if self._data_source_path is not None:
+            return  # already uploaded; reuse the cached result
+        match self._data_source:
+            case pd.DataFrame() as df:
+                pass
+            case str(url):
+                df = pd.read_csv(url, **kwargs)
+            case _:
+                raise ValueError("Data source must be a pandas DataFrame or a URL")
+
+        tmp_path: Path | None = None
+        try:
+            with tempfile.NamedTemporaryFile(mode="w+", suffix=".csv", delete=False) as tmp:
+                tmp_path = Path(tmp.name)
+            df.to_csv(tmp_path, index=False)
+            file_name = f"dataset{self._generate_random_string()}.csv"
+            self._data_source_path = self._upload_to_fileset(
+                dataset_path=tmp_path,
+                filename=file_name,
+                fileset_name="safe-synthesizer-inputs",
+            )
+        finally:
+            if tmp_path is not None:
+                tmp_path.unlink(missing_ok=True)
+
+    def _build_job_spec(self) -> dict:
+        """Assemble the final job spec dict to send to the API."""
+        if not self._enable_replace_pii and not self._enable_synthesis:
+            raise ValueError("Data synthesis and/or replace PII must be enabled")
+        if not self._data_source_path:
+            raise ValueError("No data source path found after uploading dataset")
+
+        pii_config: dict | None = None
+        if self._enable_replace_pii:
+            pii_config = self._replace_pii_config.copy()
+            if self._classify_model_provider:
+                globals_cfg = pii_config.setdefault("globals", {})
+                classify_cfg = globals_cfg.setdefault("classify", {})
+                classify_cfg["classify_model_provider"] = self._classify_model_provider
+                logger.debug(f"Injected classify model provider: {self._classify_model_provider}")
+
+        nss_config: dict = {
+            "enable_synthesis": self._enable_synthesis,
+            "enable_replace_pii": self._enable_replace_pii,
+            "data": self._data_config,
+            "training": self._training_config,
+            "generation": self._generation_config,
+            "evaluation": self._evaluation_config,
+            "time_series": self._time_series_config,
+        }
+        if self._privacy_config is not None:
+            nss_config["privacy"] = self._privacy_config
+        if pii_config is not None:
+            nss_config["replace_pii"] = pii_config
+
+        return {
+            "data_source": self._data_source_path,
+            "config": nss_config,
+            "hf_token_secret": self._hf_token_secret,
+        }
+
+    def resolve_job_config(self) -> Self:
+        """Resolve and validate the final job configuration without submitting.
+
+        Returns:
+            The builder instance for method chaining.
+        """
+        self._resolve_datasource()
+        # _build_job_spec validates; call it just for the side-effect of early error detection.
+        self._build_job_spec()
+        return self
+
+    def create_job(self, **kwargs) -> SafeSynthesizerJob:
+        """Upload the dataset and submit the job.
+
+        Args:
+            **kwargs: Additional job creation parameters passed to the API.
+
+        Returns:
+            A :class:`SafeSynthesizerJob` for monitoring and retrieving results.
+        """
+        self._resolve_datasource()
+        spec = self._build_job_spec()
+        response = self._client.safe_synthesizer.jobs.create(
+            workspace=self._workspace,
+            spec=spec,
+            **kwargs,  # type: ignore  # spec accepts dict at runtime
+        )
+        return SafeSynthesizerJob(response.name, self._client, workspace=self._workspace)
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _generate_random_string(self, length: int = 6) -> str:
+        characters = string.ascii_uppercase + string.digits
+        return "".join(random.choice(characters) for _ in range(length))
+
+    def _upload_to_fileset(
+        self,
+        dataset_path: str | Path,
+        filename: str,
+        fileset_name: str,
+    ) -> str:
+        """Upload a dataset file to a fileset.
+
+        Returns:
+            The fileset:// URL of the uploaded file.
+        """
+        dataset_path = self._validate_dataset_path(dataset_path)
+        self._client.files.upload(
+            local_path=str(dataset_path),
+            remote_path=filename,
+            fileset=fileset_name,
+            workspace=self._workspace,
+            fileset_auto_create=True,
+        )
+        return f"{self._workspace}/{fileset_name}#{filename}"
+
+    @staticmethod
+    def _validate_dataset_path(dataset_path: str | Path) -> Path:
+        if not Path(dataset_path).is_file():
+            raise ValueError("🛑 To upload a dataset, you must provide a valid file path.")
+        if not Path(dataset_path).name.endswith((".parquet", ".csv", ".json", ".jsonl")):
+            raise ValueError(
+                "🛑 Dataset files must be in `parquet`, `csv`, or `json` (orient='records', lines=True) format."
+            )
+        return Path(dataset_path)

nemo_platform/cli/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""NeMo CLI - Command-line interface for NeMo Platform."""

nemo_platform/cli/app.py

@@ -0,0 +1,321 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""NeMo CLI - Command-line interface for NeMo Platform."""
+
+from __future__ import annotations
+
+import logging
+from importlib.metadata import EntryPoint
+from typing import TYPE_CHECKING, Annotated
+
+import typer
+
+from nemo_platform.cli.commands.api import API_TOP_LEVEL_ENTRIES
+from nemo_platform.cli.commands.manifest_registry import TOP_LEVEL_ENTRIES
+from nemo_platform.cli.core.help_formatter import HELP_OPTION_NAMES
+from nemo_platform.cli.core.lazy_load import (
+    ManifestBackedNmpGroup,
+    attach_lazy_entries,
+)
+from nemo_platform.cli.core.logging import configure_logging
+from nemo_platform.cli.core.types import ListOutputFormat, TimestampFormat
+from nemo_platform.cli.manifest import (
+    TopLevelEntry,
+    build_top_level_entries,
+)
+
+if TYPE_CHECKING:
+    from nemo_platform_plugin.cli import NemoCLI
+    from nemo_platform_plugin.function import NemoFunction
+    from nemo_platform_plugin.job import NemoJob
+
+    from nemo_platform.config.models import ConfigParams
+
+logger = logging.getLogger(__name__)
+
+_SKIP_AUTH_CHECK_SUBCOMMANDS = frozenset(
+    {"agent", "auth", "config", "setup", "quickstart", "cluster-info", "skills", "docs", "services", "plugins"}
+)
+# Create the main CLI app with custom help formatting
+app = typer.Typer(
+    name="nemo",
+    no_args_is_help=True,
+    add_completion=True,
+    pretty_exceptions_enable=False,
+    rich_markup_mode=None,
+    cls=ManifestBackedNmpGroup,
+    context_settings=dict(help_option_names=list(HELP_OPTION_NAMES)),
+)
+
+
+def _build_top_level_lazy_entries() -> tuple[TopLevelEntry, ...]:
+    return build_top_level_entries(
+        (*TOP_LEVEL_ENTRIES, *API_TOP_LEVEL_ENTRIES),
+        _installed_plugin_command_entry_points(),
+        include_hidden=True,
+    )
+
+
+def _installed_plugin_command_entry_points() -> dict[str, EntryPoint]:
+    """Return installed plugin CLI entry points without importing plugin code."""
+    try:
+        from nemo_platform_plugin.discovery import discover_entry_points
+    except ImportError:
+        return {}
+    try:
+        return discover_entry_points("nemo.cli")
+    except Exception:  # noqa: BLE001
+        logger.warning("Failed to discover CLI plugin entry points", exc_info=True)
+        return {}
+
+
+def _add_plugin_job_commands(
+    plugin_app: typer.Typer,
+    plugin_jobs: dict[str, type[NemoJob]],
+    *,
+    cli: NemoCLI | None = None,
+) -> None:
+    # TODO: nemo-platform-plugin is temporarily optional while it is being published
+    # to the nightly PyPI feed. Once available, it should become an
+    # unconditional dependency and these guards can be removed.
+    try:
+        from nemo_platform_plugin.commands import add_job_commands
+    except ImportError:
+        logger.warning(
+            "nemo_platform_plugin.commands unavailable; skipping plugin job command injection", exc_info=True
+        )
+        return
+
+    add_job_commands(plugin_app, plugin_jobs, cli=cli)
+
+
+def _discover_plugin_job_entry_points() -> dict[str, EntryPoint] | None:
+    # TODO: nemo-platform-plugin is temporarily optional while it is being published
+    # to the nightly PyPI feed. Once available, it should become an
+    # unconditional dependency and these guards can be removed.
+    try:
+        from nemo_platform_plugin.discovery import discover_entry_points
+    except ImportError:
+        return None
+
+    return discover_entry_points("nemo.jobs")
+
+
+def _add_plugin_function_commands(
+    plugin_app: typer.Typer,
+    plugin_functions: dict[str, type[NemoFunction]],
+    *,
+    cli: NemoCLI | None = None,
+) -> None:
+    # Same nemo-platform-plugin optionality guard as the jobs path. When the
+    # package isn't installed we silently skip — the plugin's bare CLI
+    # surface (whatever it shipped via `nemo.cli`) still loads.
+    try:
+        from nemo_platform_plugin.commands import add_function_commands
+    except ImportError:
+        logger.warning(
+            "nemo_platform_plugin.commands unavailable; skipping plugin function command injection",
+            exc_info=True,
+        )
+        return
+
+    add_function_commands(plugin_app, plugin_functions, cli=cli)
+
+
+def _discover_plugin_function_entry_points() -> dict[str, EntryPoint] | None:
+    try:
+        from nemo_platform_plugin.discovery import discover_entry_points
+    except ImportError:
+        return None
+
+    return discover_entry_points("nemo.functions")
+
+
+# Global options
+@app.callback(options_metavar="[GLOBAL OPTIONS]")
+def main(
+    ctx: typer.Context,
+    version: Annotated[
+        bool | None,
+        typer.Option(
+            "--version",
+            "-V",
+            help="Show version information and exit.",
+            callback=_version_callback,
+            is_eager=True,
+            rich_help_panel="Help",
+        ),
+    ] = None,
+    context_name: Annotated[
+        str | None,
+        typer.Option(
+            "--context",
+            "-c",
+            help="The name of the context to use. Overrides the current context in the config file.",
+            rich_help_panel="Global Options",
+            hidden=True,
+        ),
+    ] = None,
+    base_url: Annotated[
+        str | None,
+        typer.Option(
+            "--base-url",
+            help="Base URL for the NeMo Platform API",
+            rich_help_panel="Global Options",
+        ),
+    ] = None,
+    output_format: Annotated[
+        ListOutputFormat | None,
+        typer.Option(
+            "--output-format",
+            "-f",
+            help="Output format for how results are printed.",
+            rich_help_panel="Global Options",
+        ),
+    ] = None,
+    no_truncate: Annotated[
+        bool | None,
+        typer.Option(
+            "--no-truncate",
+            help="Don't truncate long values in table/markdown/csv output",
+            rich_help_panel="Global Options",
+        ),
+    ] = None,
+    timestamp_format: Annotated[
+        TimestampFormat | None,
+        typer.Option(
+            help="Timestamp format for table/markdown/csv output",
+            rich_help_panel="Global Options",
+        ),
+    ] = None,
+    verbose: Annotated[
+        bool | None,
+        typer.Option(
+            "--verbose",
+            "-v",
+            help="Enable verbose messaging. This only impacts logs that are visible, it doesn't change any data outputs.",
+            rich_help_panel="Global Options",
+        ),
+    ] = None,
+    agent_mode: Annotated[
+        bool | None,
+        typer.Option(
+            "--agent-mode",
+            "-A",
+            help="Enable agent-friendly output mode with extra context for coding agents.",
+            rich_help_panel="Global Options",
+        ),
+    ] = None,
+    no_auto_refresh: Annotated[
+        bool,
+        typer.Option(
+            "--no-auto-refresh",
+            help="Disable automatic token refresh when token is about to expire.",
+            rich_help_panel="Global Options",
+            hidden=True,
+        ),
+    ] = False,
+) -> None:
+    """
+    Command-line interface for NeMo Platform.
+
+    :books: Documentation: https://docs.nvidia.com/nemo/microservices/latest/about/index.html
+
+    [green]Getting started:[/]
+    - Browse documentation with [cyan]`nemo docs --list`[/]
+    - Run local platform services with [cyan]`nemo services run --help`[/]
+
+    [green]Examples:[/]
+    nemo workspaces list --output-format markdown
+    nemo workspaces get default -f json
+    """
+    # Lazy imports for performance (avoid loading pydantic_settings for --help)
+    from nemo_platform.cli.core.context import CLIContext
+    from nemo_platform.quickstart import QuickstartConfig
+
+    # Configure logging (always call to silence httpx in non-verbose mode)
+    configure_logging(1 if verbose else 0)
+
+    if ctx.obj is None:
+        ctx.obj = CLIContext()
+
+    # Resolve agent mode: explicit flag > env var > default False
+    import os
+
+    if agent_mode is None:
+        env_val = os.environ.get("NMP_AGENT_MODE", "").lower()
+        agent_mode = env_val in ("1", "true", "yes")
+    ctx.obj.agent_mode = agent_mode
+
+    # Build ConfigParams from CLI args
+    overrides: ConfigParams = {}
+    if context_name is not None:
+        overrides["current_context"] = context_name
+    if base_url is not None:
+        overrides["base_url"] = base_url
+    if output_format is not None:
+        overrides["output_format"] = output_format
+    elif agent_mode:
+        overrides["output_format"] = "markdown"
+    if timestamp_format is not None:
+        overrides["timestamp_format"] = timestamp_format
+    if no_truncate is not None:
+        overrides["truncate"] = not no_truncate
+
+    # Update CLIContext overrides
+    ctx.obj.overrides.update(overrides)
+    ctx.obj.verbosity = 1 if verbose else 0
+    ctx.obj.quickstart_config = QuickstartConfig.load()
+
+    # Non-quickstart contexts always require auth. Some quickstart contexts require auth (opt-in)
+    context_requires_auth = ctx.obj.quickstart_config is None or ctx.obj.quickstart_config.auth_enabled
+    command_requires_auth = ctx.invoked_subcommand not in _SKIP_AUTH_CHECK_SUBCOMMANDS
+    auto_refresh_enabled = not no_auto_refresh
+
+    if context_requires_auth and command_requires_auth and auto_refresh_enabled:
+        from nemo_platform.cli.commands.auth import AuthError, ensure_valid_token
+
+        try:
+            token_valid = ensure_valid_token(ctx.obj.get_sdk_context())
+            if not token_valid:
+                typer.echo(
+                    "Error: Your access token has expired and could not be refreshed.\n"
+                    "Hint: Run 'nemo auth login' to re-authenticate.",
+                    err=True,
+                )
+                raise typer.Exit(code=1)
+
+            # reset the context, so we load the refreshed auth info
+            ctx.obj.reset_sdk_context()
+        except typer.Exit:
+            raise  # Re-raise Exit to stop execution
+        except AuthError as e:
+            typer.echo(f"Error: {e}", err=True)
+            raise typer.Exit(code=1)
+        except Exception as e:
+            typer.echo(f"Warning: Failed to check/refresh token: {e}", err=True)
+
+    if ctx.invoked_subcommand is None:
+        typer.echo(ctx.get_help())
+
+
+attach_lazy_entries(main, _build_top_level_lazy_entries())
+
+
+def _version_callback(value: bool) -> None:
+    """Print version information and exit."""
+    if value:
+        import nemo_platform
+
+        typer.echo(f"nemo version {nemo_platform.__version__}")
+        raise typer.Exit()
+
+
+def cli() -> None:
+    """Main entry point for the CLI."""
+    app()
+
+
+if __name__ == "__main__":
+    cli()

nemo_platform/cli/commands/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+