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

nemo_platform/config/models.py

@@ -0,0 +1,377 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Core configuration models for nemo_platform."""
+
+from abc import ABC
+from typing import Annotated, Any, Literal, TypedDict
+
+from pydantic import (
+    BaseModel,
+    Discriminator,
+    Field,
+    HttpUrl,
+    SecretStr,
+    SerializationInfo,
+    Tag,
+    field_serializer,
+    field_validator,
+    model_validator,
+)
+
+from nemo_platform.config.types import OutputFormat, TimestampFormat
+
+DEFAULT_WORKSPACE = "default"
+DEFAULT_CONTEXT = "default"
+DEFAULT_BASE_URL = "http://localhost:8080"
+
+
+class BaseUser(BaseModel, ABC):
+    def get_client_config(self) -> dict[str, object]:
+        return {}
+
+
+class OAuthUser(BaseUser):
+    """User with OAuth token authentication (from device flow, browser flow, etc.)."""
+
+    type: Literal["oauth"] = "oauth"
+
+    name: str = Field(..., min_length=1, description="Unique user name")
+    token: SecretStr = Field(..., min_length=1, description="Access token (JWT)")
+    refresh_token: SecretStr | None = Field(default=None, description="Refresh token for automatic renewal")
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, v: str) -> str:
+        """Validate user name format."""
+        if not v or v.isspace():
+            raise ValueError("User name cannot be empty")
+        return v
+
+    @field_serializer("token", "refresh_token", when_used="json")
+    def serialize_tokens(self, value: SecretStr | None, info: SerializationInfo) -> str | None:
+        """Serialize tokens, revealing secret only when context requests it."""
+        if value is None:
+            return None
+        if info.context and info.context.get("include_secrets"):
+            return value.get_secret_value()
+        return "***REDACTED***"
+
+    def get_client_config(self) -> dict[str, object]:
+        """Return client config with Bearer token authentication."""
+        return {
+            "default_headers": {
+                "Authorization": f"Bearer {self.token.get_secret_value()}",
+            }
+        }
+
+
+class NoAuthUser(BaseUser):
+    """User with no authentication."""
+
+    type: Literal["no-auth"] = "no-auth"
+
+    name: str = Field(..., min_length=1, description="Unique user name")
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, v: str) -> str:
+        """Validate user name format."""
+        if not v or v.isspace():
+            raise ValueError("User name cannot be empty")
+        return v
+
+
+def _get_user_type(data: Any) -> str:
+    """
+    Determine the type of the user for the discriminator with backward-compatibility.
+    """
+    default = "no-auth"
+    if isinstance(data, dict):
+        return data.get("type", default)
+    return getattr(data, "type", default)
+
+
+# Union type for all user types
+User = Annotated[
+    Annotated[OAuthUser, Tag("oauth")] | Annotated[NoAuthUser, Tag("no-auth")],
+    Discriminator(_get_user_type),
+]
+
+
+class Cluster(BaseModel):
+    """Cluster configuration - connection info only, defined in config file."""
+
+    name: str = Field(..., min_length=1, description="Unique cluster name")
+    base_url: HttpUrl = Field(..., description="Base URL for the cluster")
+    metadata: dict[str, Any] = Field(default_factory=dict, description="Additional cluster metadata")
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, v: str) -> str:
+        """Validate cluster name format."""
+        if not v or v.isspace():
+            raise ValueError("Cluster name cannot be empty")
+        if not v.replace("-", "").replace("_", "").isalnum():
+            raise ValueError("Cluster name must be alphanumeric (hyphens and underscores allowed)")
+        return v
+
+
+class Preferences(BaseModel):
+    """User preferences for CLI/SDK behavior - can be overridden at runtime."""
+
+    output_format: OutputFormat = Field(default="table", description="Default output format")
+    timestamp_format: TimestampFormat = Field(default="iso8601", description="Timestamp display format")
+    truncate: bool = Field(default=True, description="Truncate long strings in output")
+    color_output: bool = Field(default=True, description="Enable colored output")
+
+
+class ContextDefinition(BaseModel):
+    """
+    Context definition from the config file (serialization model).
+
+    This represents a context as defined in the YAML config file. It contains
+    string references to cluster and user names that must be resolved at runtime.
+    This is NOT what application code should use - use Context instead.
+    """
+
+    name: str = Field(..., min_length=1, description="Unique context name")
+    cluster: str = Field(..., min_length=1, description="Reference to cluster name (string, not resolved)")
+    user: str = Field(..., min_length=1, description="Reference to user name (string, not resolved)")
+    workspace: str | None = Field(default=None, description="Default workspace")
+    default_model: str | None = Field(default=None, description="Default model entity ID for inference")
+    preferences: Preferences = Field(default_factory=Preferences, description="Context-specific preferences")
+    metadata: dict[str, Any] = Field(default_factory=dict, description="Additional context metadata")
+
+    @field_validator("name", "cluster", "user")
+    @classmethod
+    def validate_identifier(cls, v: str) -> str:
+        """Validate identifier format."""
+        if not v or v.isspace():
+            raise ValueError("Identifier cannot be empty")
+        return v
+
+
+class ConfigParams(TypedDict, total=False):
+    """
+    Type-safe configuration parameters.
+
+    Configuration parameters that can be set via env vars, CLI, or programmatically.
+    These parameters override values within the scope of the current context, not the
+    config file definitions themselves. When no config file exists, cluster_url is required
+    to directly configure the cluster connection.
+
+    All fields are optional - you only need to specify the ones you want to override.
+    """
+
+    current_context: str
+
+    base_url: str
+
+    # OAuth fields (for OAuthUser)
+    access_token: str
+    refresh_token: str
+
+    workspace: str
+    default_model: str
+
+    output_format: OutputFormat
+    timestamp_format: TimestampFormat
+    truncate: bool
+    color_output: bool
+
+
+class LocalServicesConfig(BaseModel):
+    """User-selected paths for ``nemo services run`` local state.
+
+    Persisted by ``nemo setup`` so that subsequent runs reuse the same
+    on-disk locations without re-prompting. ``None`` means "fall back to
+    the service default" (``~/.local/share/nemo/``).
+    """
+
+    data_dir: str | None = Field(
+        default=None,
+        description=(
+            "Directory under which local services persist state (SQLite DB, "
+            "encryption key, files). Passed to the service subprocess as "
+            "``NMP_DATA_DIR``."
+        ),
+    )
+
+
+class ConfigFile(BaseModel):
+    """
+    Root configuration from file.
+
+    This represents the YAML config file content. Contains ContextDefinitions
+    (serialization models) that reference clusters and users by name.
+    These definitions are resolved into Context objects by the Config.resolve() method.
+
+    If no config file exists, an empty ConfigFile can be used with all configuration
+    coming from environment variables or SDK parameters.
+    """
+
+    current_context: str | None = Field(default=None, description="Name of the currently active context")
+    clusters: list[Cluster] = Field(default_factory=list, description="List of configured clusters")
+    users: list[User] = Field(default_factory=list, description="List of configured users")
+    contexts: list[ContextDefinition] = Field(
+        default_factory=list, description="List of context definitions from the config file"
+    )
+    local_services: LocalServicesConfig | None = Field(
+        default=None,
+        description="User-selected paths for local services (set by `nemo setup`).",
+    )
+
+    def ensure_context(
+        self,
+        context_name: str,
+        params: ConfigParams,
+    ) -> tuple[Cluster, User, ContextDefinition]:
+        """
+        Find or create cluster/user/context entities with the given name.
+
+        - If entities exist, updates them with params
+        - If not, creates new entities with descriptive names
+          (e.g. context "local" → cluster "local-cluster", user "local-user")
+
+        Returns (cluster, user, context) tuple.
+        """
+        # First check if context exists to get its cluster/user references
+        existing_context = next((c for c in self.contexts if c.name == context_name), None)
+
+        # If context exists, use its references; otherwise derive descriptive names
+        cluster_name = existing_context.cluster if existing_context else f"{context_name}-cluster"
+        user_name = existing_context.user if existing_context else f"{context_name}-user"
+
+        # Find existing or create cluster
+        cluster = next((c for c in self.clusters if c.name == cluster_name), None)
+        if cluster is None:
+            base_url = params.get("base_url")
+            if not base_url:
+                raise ValueError(f"Cluster '{cluster_name}' does not exist and no base_url provided to create it!")
+
+            cluster = Cluster(name=cluster_name, base_url=HttpUrl(base_url))
+            self.clusters.append(cluster)
+        elif "base_url" in params:
+            cluster.base_url = HttpUrl(params["base_url"])
+
+        # Find existing or create user
+        user: User = next((u for u in self.users if u.name == user_name), None)  # type: ignore[assignment]
+        access_token = params.get("access_token")
+        refresh_token = params.get("refresh_token")
+
+        if user is None:
+            if access_token:
+                # Create OAuthUser for OAuth tokens
+                user = OAuthUser(
+                    name=user_name,
+                    token=SecretStr(access_token),
+                    refresh_token=SecretStr(refresh_token) if refresh_token else None,
+                )
+            else:
+                user = NoAuthUser(name=user_name)
+            self.users.append(user)
+        elif access_token:
+            # Replace existing user with OAuthUser
+            idx = next(i for i, u in enumerate(self.users) if u.name == user_name)
+            user = OAuthUser(
+                name=user_name,
+                token=SecretStr(access_token),
+                refresh_token=SecretStr(refresh_token) if refresh_token else None,
+            )
+            self.users[idx] = user
+        elif isinstance(user, OAuthUser) and refresh_token:
+            # Update existing OAuthUser with new refresh token info
+            idx = next(i for i, u in enumerate(self.users) if u.name == user_name)
+            user = OAuthUser(
+                name=user_name,
+                token=user.token,
+                refresh_token=SecretStr(refresh_token) if refresh_token else user.refresh_token,
+            )
+            self.users[idx] = user
+
+        # Find existing or create context
+        context = existing_context
+        if context is None:
+            context = ContextDefinition(
+                name=context_name,
+                cluster=cluster_name,
+                user=user_name,
+                workspace=params.get("workspace", DEFAULT_WORKSPACE),
+                preferences=Preferences(),
+            )
+            self.contexts.append(context)
+
+        # Apply context-level params
+        if "workspace" in params:
+            context.workspace = params["workspace"]
+        if "default_model" in params:
+            context.default_model = params["default_model"]
+        if "output_format" in params:
+            context.preferences.output_format = params["output_format"]
+        if "timestamp_format" in params:
+            context.preferences.timestamp_format = params["timestamp_format"]
+        if "truncate" in params:
+            context.preferences.truncate = params["truncate"]
+        if "color_output" in params:
+            context.preferences.color_output = params["color_output"]
+
+        return cluster, user, context
+
+    @model_validator(mode="after")
+    def validate_references(self) -> "ConfigFile":
+        """Validate that contexts reference existing clusters and users."""
+        cluster_names = {cluster.name for cluster in self.clusters}
+        user_names = {user.name for user in self.users}
+        context_names = {context.name for context in self.contexts}
+
+        # Validate cluster and user references
+        for context in self.contexts:
+            if context.cluster not in cluster_names:
+                raise ValueError(f"Context '{context.name}' references non-existent cluster '{context.cluster}'")
+            if context.user not in user_names:
+                raise ValueError(f"Context '{context.name}' references non-existent user '{context.user}'")
+
+        # Validate current context exists (only if contexts are defined)
+        if self.current_context and self.contexts and self.current_context not in context_names:
+            raise ValueError(f"Current context '{self.current_context}' does not exist")
+
+        return self
+
+    @model_validator(mode="after")
+    def validate_unique_names(self) -> "ConfigFile":
+        """Ensure cluster, user, and context names are unique."""
+        cluster_names = [cluster.name for cluster in self.clusters]
+        if len(cluster_names) != len(set(cluster_names)):
+            raise ValueError("Cluster names must be unique")
+
+        user_names = [user.name for user in self.users]
+        if len(user_names) != len(set(user_names)):
+            raise ValueError("User names must be unique")
+
+        context_names = [context.name for context in self.contexts]
+        if len(context_names) != len(set(context_names)):
+            raise ValueError("Context names must be unique")
+
+        return self
+
+
+class Context(BaseModel):
+    """
+    Resolved runtime configuration (what application code should use).
+
+    This is the effective configuration after:
+    1. Resolving the cluster and user references (strings -> full objects)
+    2. Applying runtime overrides (env vars, CLI flags, SDK parameters)
+    3. Merging preferences from multiple sources
+
+    Unlike ContextDefinition which is a serialization model from the config file,
+    this contains fully resolved objects ready for immediate use.
+    """
+
+    context_name: str = Field(..., description="Active context name")
+    cluster: Cluster = Field(..., description="Fully resolved cluster configuration with URL")
+    user: User | None = Field(default=None, description="Resolved user with authentication credentials")
+    workspace: str = Field(..., description="Active workspace (required)")
+    default_model: str | None = Field(default=None, description="Default model entity ID for inference")
+    preferences: Preferences = Field(..., description="Effective preferences")

nemo_platform/config/types.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Literal
+
+OutputFormat = Literal["table", "json", "yaml", "markdown", "csv", "raw"]  # Supported output formats
+TimestampFormat = Literal["relative", "iso8601"]  # Supported timestamp formats

nemo_platform/filesets/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Extended filesets module with FilesetFileSystem support.
+
+This module provides high-level file operations (upload, download, etc.) and
+fsspec integration for NeMo Platform filesets via the sdk.files.fsspec property.
+
+Located at: nemo_platform/filesets/ (after vendoring)
+"""
+
+from .filesystem.callbacks import RichFileProgressCallback as RichFileProgressCallback
+from .filesystem.callbacks import RichProgressCallback as RichProgressCallback
+from .filesystem.filesystem import FilesetFileSystem as FilesetFileSystem
+from .filesystem.filesystem import FilesetPathError as FilesetPathError
+from .filesystem.filesystem import build_fileset_ref as build_fileset_ref
+from .filesystem.filesystem import parse_fileset_path as parse_fileset_path
+from .filesystem.filesystem import parse_fileset_ref as parse_fileset_ref
+from .resources import ListFilesResponse as ListFilesResponse

nemo_platform/filesets/filesystem/callbacks.py

@@ -0,0 +1,229 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from fsspec.callbacks import Callback
+
+
+def _create_rich_progress(transient: bool = False):
+    """Create a Rich Progress instance with adaptive columns.
+
+    The progress bar uses custom columns that render differently based on
+    whether the task is the main (file count) task or a file (bytes) task.
+    Tasks should set `is_main=True` in their fields for file count display.
+    """
+    from rich.progress import (
+        BarColumn,
+        DownloadColumn,
+        MofNCompleteColumn,
+        Progress,
+        ProgressColumn,
+        SpinnerColumn,
+        Task,
+        TextColumn,
+        TimeElapsedColumn,
+        TransferSpeedColumn,
+    )
+    from rich.table import Column
+    from rich.text import Text
+
+    class FilesOrBytesColumn(ProgressColumn):
+        """Shows 'N/M files' for main task, bytes for file tasks."""
+
+        def __init__(self):
+            super().__init__(table_column=Column(min_width=14))
+            self._files_col = MofNCompleteColumn()
+            self._bytes_col = DownloadColumn()
+
+        def render(self, task: Task) -> Text:
+            if task.fields.get("is_main"):
+                result = self._files_col.render(task)
+                return Text(f"{result.plain} files", style="progress.download")
+            return self._bytes_col.render(task)
+
+    class AdaptiveSpeedColumn(ProgressColumn):
+        """Shows '-' for main task, transfer speed for file tasks."""
+
+        def __init__(self):
+            super().__init__(table_column=Column(min_width=10))
+            self._speed_col = TransferSpeedColumn()
+
+        def render(self, task: Task) -> Text:
+            if task.fields.get("is_main"):
+                return Text("-", style="progress.data.speed")
+            result = self._speed_col.render(task)
+            if result.plain == "?":
+                return Text("-", style="progress.data.speed")
+            return result
+
+    return Progress(
+        SpinnerColumn(),
+        TextColumn("[bold blue]{task.description}", table_column=Column(min_width=32)),
+        BarColumn(),
+        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+        TextColumn("•"),
+        FilesOrBytesColumn(),
+        TextColumn("•"),
+        AdaptiveSpeedColumn(),
+        TextColumn("•"),
+        TimeElapsedColumn(),
+        transient=transient,
+    )
+
+
+class RichProgressCallback(Callback):
+    """fsspec Callback that displays progress using Rich progress bars.
+
+    Requires the `rich` package to be installed. If not installed,
+    raises ImportError when instantiated.
+
+    Shows:
+    - Overall progress bar at top (file count)
+    - Individual progress bar per file being transferred (bytes)
+
+    Example:
+        >>> with RichProgressCallback(description="Uploading") as callback:
+        ...     fs.put(local_path, remote_path, callback=callback)
+    """
+
+    def __init__(
+        self,
+        description: str = "Transferring",
+        transient: bool = False,
+    ):
+        """Initialize the Rich progress callback.
+
+        Args:
+            description: Text to display next to the progress bar.
+            transient: If True, remove the progress bar after completion.
+        """
+        try:
+            import rich  # noqa: F401
+            from rich.progress import TaskID
+        except ImportError as e:
+            raise ImportError(
+                "RichProgressCallback requires the 'rich' package. Install it with: uv pip install rich"
+            ) from e
+
+        super().__init__()
+        self._description = description
+        self._progress = _create_rich_progress(transient=transient)
+        self._main_task_id: TaskID | None = None
+
+    def set_size(self, size: int | None) -> None:
+        """Set the total number of files."""
+        if self._main_task_id is None or size is None:
+            return
+        self._progress.update(self._main_task_id, total=size)
+
+    def absolute_update(self, value: int) -> None:
+        """Set file progress to an absolute value."""
+        if self._main_task_id is None:
+            return
+        self._progress.update(self._main_task_id, completed=value)
+
+    def relative_update(self, inc: int = 1) -> None:
+        """Increment file progress (called when a file completes)."""
+        if self._main_task_id is None:
+            return
+        self._progress.advance(self._main_task_id, advance=inc)
+
+    def branched(self, path_1: str, _path_2: str, **_kwargs) -> RichFileProgressCallback:
+        """Create a child callback for an individual file transfer."""
+        filename = path_1.rsplit("/", 1)[-1] if "/" in path_1 else path_1
+        return RichFileProgressCallback(filename, progress=self._progress)
+
+    def __enter__(self):
+        """Start the progress bar display."""
+        self._progress.start()
+        self._main_task_id = self._progress.add_task(self._description, total=None, is_main=True)
+        return self
+
+    def __exit__(self, *_exc_args):
+        """Stop the progress bar display."""
+        self._progress.stop()
+
+
+class RichFileProgressCallback(Callback):
+    """Callback that shows progress for a single file transfer.
+
+    Can be used standalone with get_file() or as a child of RichProgressCallback
+    for tracking individual files in a batch transfer.
+
+    Example (standalone):
+        >>> with RichFileProgressCallback() as callback:
+        ...     fs.get_file(remote_path, local_path, callback=callback)
+
+    Example (as child of RichProgressCallback):
+        >>> # Typically created automatically via RichProgressCallback.branched()
+        >>> callback = RichFileProgressCallback("myfile.txt", progress=parent._progress)
+    """
+
+    def __init__(
+        self,
+        filename: str = "Transferring",
+        *,
+        progress=None,
+        transient: bool = False,
+    ):
+        """Initialize the file progress callback.
+
+        Args:
+            filename: Name/description to display. Defaults to "Transferring".
+            progress: Optional Rich Progress instance. If not provided, creates its own.
+            transient: If True and creating own progress bar, remove it after completion.
+        """
+        try:
+            import rich  # noqa: F401
+        except ImportError as e:
+            raise ImportError(
+                "RichFileProgressCallback requires the 'rich' package. Install it with: uv pip install rich"
+            ) from e
+
+        from rich.progress import TaskID
+
+        super().__init__()
+        self._filename = filename
+        self._owns_progress = progress is None
+        self._progress = progress if progress is not None else _create_rich_progress(transient=transient)
+        self._task_id: TaskID | None = None
+
+    def set_size(self, size: int | None) -> None:
+        """Create the task with the file size."""
+        if size is None:
+            return
+        # Truncate long filenames
+        display_name = self._filename
+        if len(display_name) > 30:
+            display_name = "..." + display_name[-27:]
+        # Use different formatting for standalone vs child mode
+        if self._owns_progress:
+            task_description = display_name
+        else:
+            task_description = f"[dim]{display_name}[/dim]"
+        self._task_id = self._progress.add_task(task_description, total=size, is_main=False)
+
+    def relative_update(self, inc: int = 1) -> None:
+        """Update bytes transferred for this file."""
+        if self._task_id is None:
+            return
+        self._progress.advance(self._task_id, advance=inc)
+
+    def absolute_update(self, value: int) -> None:
+        """Set bytes transferred to absolute value."""
+        if self._task_id is None:
+            return
+        self._progress.update(self._task_id, completed=value)
+
+    def __enter__(self):
+        """Start the progress bar display (only when used standalone)."""
+        if self._owns_progress:
+            self._progress.start()
+        return self
+
+    def __exit__(self, *_exc_args):
+        """Stop the progress bar display (only when used standalone)."""
+        if self._owns_progress:
+            self._progress.stop()
+        return False