groundhog-hpc 0.5.6__py3-none-any.whl → 0.7.0__py3-none-any.whl

groundhog_hpc/configuration/endpoints.py

@@ -6,7 +6,6 @@ conform to the EndpointConfig/EndpointVariant models for consistency with
 existing configuration parsing logic.
 """
 
-from dataclasses import dataclass
 from typing import Any
 from uuid import UUID
 
@@ -16,6 +15,9 @@ from groundhog_hpc.compute import get_endpoint_metadata, get_endpoint_schema
 KNOWN_ENDPOINTS: dict[str, dict[str, Any]] = {
     "anvil": {
         "uuid": "5aafb4c1-27b2-40d8-a038-a0277611868f",
+        "base": {
+            "requirements": "",
+        },
         "variants": {
             "gpu": {
                 "partition": "gpu-debug",
@@ -26,24 +28,12 @@ KNOWN_ENDPOINTS: dict[str, dict[str, Any]] = {
     },
     "tutorial": {
         "uuid": "4b116d3c-1703-4f8f-9f6f-39921e5864df",
+        "base": {},
         "variants": {},
     },
 }
 
 
-@dataclass
-class FormattedEndpoint:
-    """Formatted endpoint configuration with metadata for template rendering.
-
-    Attributes:
-        name: Endpoint name for use in @hog.function(endpoint="name")
-        toml_block: Formatted TOML configuration block
-    """
-
-    name: str
-    toml_block: str
-
-
 class EndpointSpec:
     """Parsed endpoint specification from --endpoint flag.
 
@@ -51,6 +41,7 @@ class EndpointSpec:
         name: Table name for [tool.hog.{name}]
         variant: Optional variant name for [tool.hog.{name}.{variant}]
         uuid: Globus Compute endpoint UUID
+        base_defaults: Dict of defaults to apply to base endpoint (if known endpoint)
         variant_defaults: Dict of defaults to apply to variant (if known variant)
     """
 
@@ -59,11 +50,13 @@ class EndpointSpec:
         name: str,
         variant: str | None,
         uuid: str,
+        base_defaults: dict[str, Any] | None = None,
         variant_defaults: dict[str, Any] | None = None,
     ):
         self.name = name
         self.variant = variant
         self.uuid = uuid
+        self.base_defaults = base_defaults or {}
         self.variant_defaults = variant_defaults or {}
 
 
@@ -115,29 +108,40 @@ def parse_endpoint_spec(spec: str) -> EndpointSpec:
     if "." in spec:
         base_name, variant = spec.split(".", 1)
         if base_name not in KNOWN_ENDPOINTS:
-            known = ", ".join(KNOWN_ENDPOINTS.keys())
-            raise ValueError(
-                f"Unknown endpoint '{base_name}'. Known endpoints: {known}"
+            # Stub out unknown endpoint with variant with TODO placeholder
+            return EndpointSpec(
+                name=base_name,
+                variant=variant,
+                uuid="TODO: Replace with your endpoint UUID",
             )
 
         endpoint_info = KNOWN_ENDPOINTS[base_name]
         uuid = endpoint_info["uuid"]
+        base_defaults = endpoint_info.get("base", {})
         variant_defaults = endpoint_info["variants"].get(variant, {})
 
         return EndpointSpec(
             name=base_name,
             variant=variant,
             uuid=uuid,
+            base_defaults=base_defaults,
             variant_defaults=variant_defaults,
         )
 
-    # Must be a known endpoint name
+    # Must be a known endpoint name, or stub out unknown ones
     if spec not in KNOWN_ENDPOINTS:
-        known = ", ".join(KNOWN_ENDPOINTS.keys())
-        raise ValueError(f"Unknown endpoint '{spec}'. Known endpoints: {known}")
+        # Stub out unknown endpoint with TODO placeholder
+        return EndpointSpec(
+            name=spec,
+            variant=None,
+            uuid="TODO: Replace with your endpoint UUID",
+        )
 
     endpoint_info = KNOWN_ENDPOINTS[spec]
-    return EndpointSpec(name=spec, variant=None, uuid=endpoint_info["uuid"])
+    base_defaults = endpoint_info.get("base", {})
+    return EndpointSpec(
+        name=spec, variant=None, uuid=endpoint_info["uuid"], base_defaults=base_defaults
+    )
 
 
 def generate_endpoint_config(spec: EndpointSpec) -> dict[str, dict[str, Any]]:
@@ -158,10 +162,22 @@ def generate_endpoint_config(spec: EndpointSpec) -> dict[str, dict[str, Any]]:
     """
     result: dict[str, Any] = {}
 
+    # If UUID is a TODO placeholder, skip schema fetching
+    if spec.uuid.startswith("TODO"):
+        filtered_base_defaults = spec.base_defaults.copy()
+    else:
+        # Filter base_defaults to only include fields present in the endpoint schema
+        schema = get_endpoint_schema(spec.uuid)
+        schema_fields = set(schema.get("properties", {}).keys())
+        filtered_base_defaults = {
+            k: v for k, v in spec.base_defaults.items() if k in schema_fields
+        }
+
     # Base configuration
     base_config = {
         "endpoint": spec.uuid,
-        # Other fields will be added by user, we just provide the endpoint UUID
+        **filtered_base_defaults,
+        # Other fields will be added by user, we just provide the endpoint UUID + defaults
     }
     result[spec.name] = base_config
 
@@ -203,152 +219,3 @@ def get_endpoint_schema_comments(endpoint_uuid: str) -> dict[str, str]:
             comments[field_name] = ". ".join(parts)
 
     return comments
-
-
-def get_endpoint_display_name(endpoint_uuid: str) -> str | None:
-    """Fetch endpoint display name from metadata.
-
-    Args:
-        endpoint_uuid: Globus Compute endpoint UUID
-
-    Returns:
-        Display name if available and different from 'name', otherwise None
-    """
-    metadata = get_endpoint_metadata(endpoint_uuid)
-
-    display_name = metadata.get("display_name")
-    name = metadata.get("name")
-
-    # Only return display_name if it's meaningful
-    if display_name and display_name != "None" and display_name != name:
-        return display_name
-
-    return None
-
-
-def format_endpoint_config_to_toml(
-    config_dict: dict[str, dict[str, Any]],
-    endpoint_uuid: str,
-    include_schema_comments: bool = True,
-) -> str:
-    """Format an endpoint configuration dict as TOML with inline documentation.
-
-    Args:
-        config_dict: Dict with structure {"endpoint_name": {"endpoint": "uuid", ...}}
-        endpoint_uuid: UUID for fetching schema documentation
-        include_schema_comments: If True, add commented schema fields with docs
-
-    Returns:
-        Formatted TOML string with comments
-    """
-    lines = []
-
-    # Get display name and schema comments
-    display_name = get_endpoint_display_name(endpoint_uuid)
-
-    # Calculate padding for aligned inline comments
-    # Align to UUID line length (approx 51 chars: "# endpoint = "uuid..."")
-    # For schema comments: "# # field_name = " should align comment to ~column 52
-    alignment_column = 52
-
-    if include_schema_comments:
-        comments = get_endpoint_schema_comments(endpoint_uuid)
-
-    for endpoint_name, config in config_dict.items():
-        # Check if this is a base config or has nested variants
-        has_variants = any(isinstance(v, dict) for v in config.values())
-
-        if has_variants:
-            # Process base and variants
-            base_config = {k: v for k, v in config.items() if not isinstance(v, dict)}
-            variants = {k: v for k, v in config.items() if isinstance(v, dict)}
-
-            # Base config header
-            header = f"[tool.hog.{endpoint_name}]"
-            if display_name:
-                lines.append(f"# {header}  # {display_name}")
-            else:
-                lines.append(f"# {header}")
-
-            # Base config fields (active, so prefix with # for PEP 723)
-            for key, value in base_config.items():
-                if isinstance(value, str):
-                    lines.append(f'# {key} = "{value}"')
-                else:
-                    lines.append(f"# {key} = {value}")
-
-            # Add schema comments if requested (commented out, so prefix with # #)
-            if include_schema_comments:
-                comments = get_endpoint_schema_comments(endpoint_uuid)
-                for field_name, comment in comments.items():
-                    # Pad to align inline comments (left-align, pad to alignment_column)
-                    lines.append(
-                        f"# # {field_name} = {'':<{alignment_column - 7 - len(field_name)}}# {comment}"
-                    )
-
-            # Variant configs (active headers, active fields)
-            for variant_name, variant_config in variants.items():
-                lines.append("#")
-                lines.append(f"# [tool.hog.{endpoint_name}.{variant_name}]")
-                for key, value in variant_config.items():
-                    if isinstance(value, str):
-                        lines.append(f'# {key} = "{value}"')
-                    else:
-                        lines.append(f"# {key} = {value}")
-        else:
-            # Simple config without variants
-            header = f"[tool.hog.{endpoint_name}]"
-            if display_name:
-                lines.append(f"# {header}  # {display_name}")
-            else:
-                lines.append(f"# {header}")
-
-            # Active fields (prefix with # for PEP 723)
-            for key, value in config.items():
-                if isinstance(value, str):
-                    lines.append(f'# {key} = "{value}"')
-                else:
-                    lines.append(f"# {key} = {value}")
-
-            # Add schema comments if requested (commented out, so prefix with # #)
-            if include_schema_comments:
-                comments = get_endpoint_schema_comments(endpoint_uuid)
-                for field_name, comment in comments.items():
-                    # Pad to align inline comments (left-align, pad to alignment_column)
-                    lines.append(
-                        f"# # {field_name} = {'':<{alignment_column - 7 - len(field_name)}}# {comment}"
-                    )
-
-    return "\n".join(lines)
-
-
-def fetch_and_format_endpoints(endpoint_specs: list[str]) -> list[FormattedEndpoint]:
-    """Parse endpoint specifications and generate formatted TOML blocks.
-
-    Args:
-        endpoint_specs: List of endpoint specification strings
-
-    Returns:
-        List of FormattedEndpoint objects with name and TOML block
-
-    Raises:
-        ValueError: If any endpoint spec is invalid
-        RuntimeError: If unable to fetch endpoint metadata
-    """
-    endpoints = []
-
-    for spec_str in endpoint_specs:
-        try:
-            spec = parse_endpoint_spec(spec_str)
-            config_dict = generate_endpoint_config(spec)
-            toml_block = format_endpoint_config_to_toml(
-                config_dict, spec.uuid, include_schema_comments=True
-            )
-            endpoints.append(FormattedEndpoint(name=spec.name, toml_block=toml_block))
-        except Exception as e:
-            # Re-raise with context
-            raise RuntimeError(
-                f"Failed to process endpoint spec '{spec_str}': {e}"
-            ) from e
-
-    return endpoints

groundhog_hpc/configuration/models.py

@@ -25,10 +25,12 @@ class EndpointConfig(BaseModel, extra="allow"):
     Attributes:
         endpoint: Globus Compute endpoint UUID (required for base configs)
         worker_init: Shell commands to run in worker initialization
+        endpoint_setup: Shell commands to run in endpoint setup
     """
 
     endpoint: str | UUID
     worker_init: str | None = None
+    endpoint_setup: str | None = None
 
 
 class EndpointVariant(BaseModel, extra="allow"):
@@ -45,10 +47,12 @@ class EndpointVariant(BaseModel, extra="allow"):
     Attributes:
         endpoint: Always None (variants must inherit endpoint from base)
         worker_init: Additional worker init commands (concatenated with base)
+        endpoint_setup: Additional endpoint setup commands (concatenated with base)
     """
 
     endpoint: None = None
     worker_init: str | None = None
+    endpoint_setup: str | None = None
 
     @model_validator(mode="before")
     @classmethod
@@ -63,9 +67,31 @@ class EndpointVariant(BaseModel, extra="allow"):
 
 
 class UvMetadata(BaseModel, extra="allow", serialize_by_alias=True):
+    """Configuration for uv package manager via [tool.uv].
+
+    Common fields are modeled for validation and defaults. Additional uv settings
+    are supported via extra="allow" - see https://docs.astral.sh/uv/reference/settings/
+
+    Note: Environment variables (UV_*) and CLI flags take precedence over TOML config.
+    See uv documentation for full precedence hierarchy.
+
+    Attributes:
+        exclude_newer: Limit packages to versions uploaded before cutoff (ISO 8601 timestamp)
+        python_preference: Control system vs managed Python ("managed" | "only-managed" | "system" | "only-system")
+        index_url: Primary package index URL (default: PyPI)
+        extra_index_url: Additional package indexes (searched after index_url)
+        python_downloads: Control automatic Python downloads ("automatic" | "manual" | "never")
+        offline: Disable all network access (use only cache and local files)
+    """
+
     exclude_newer: str | None = Field(
         default_factory=_default_exclude_newer, alias="exclude-newer"
     )
+    python_preference: str | None = Field(default="managed", alias="python-preference")
+    index_url: str | None = Field(default=None, alias="index-url")
+    extra_index_url: list[str] | None = Field(default=None, alias="extra-index-url")
+    python_downloads: str | None = Field(default=None, alias="python-downloads")
+    offline: bool | None = None
 
 
 class ToolMetadata(BaseModel, extra="allow"):

groundhog_hpc/configuration/pep723.py

@@ -6,8 +6,10 @@ using the PEP 723 inline script metadata format (# /// script ... # ///).
 
 import re
 import sys
+from typing import Any, cast
 
-import tomli_w
+import tomlkit
+import tomlkit.items
 
 from groundhog_hpc.configuration.models import Pep723Metadata
 
@@ -73,7 +75,7 @@ def write_pep723(metadata: Pep723Metadata) -> str:
     metadata_dict = metadata.model_dump(by_alias=True, exclude_none=True)
 
     # Convert dict to TOML format
-    toml_content = tomli_w.dumps(metadata_dict)
+    toml_content = tomlkit.dumps(metadata_dict)
 
     # Format as PEP 723 inline metadata block
     lines = ["# /// script"]
@@ -137,3 +139,277 @@ def insert_or_update_metadata(script_content: str, metadata: Pep723Metadata) ->
             lines.insert(insert_index + 1, "")
 
         return "\n".join(lines)
+
+
+def extract_pep723_toml(
+    script: str,
+) -> tuple[tomlkit.TOMLDocument, re.Match] | tuple[None, None]:
+    """Extract TOML document from PEP 723 block using tomlkit for round-trip preservation.
+
+    Args:
+        script: The full text content of a Python script
+
+    Returns:
+        Tuple of (tomlkit document, regex match) or (None, None) if no block exists.
+    """
+    name = "script"
+    matches = list(
+        filter(
+            lambda m: m.group("type") == name,
+            re.finditer(INLINE_METADATA_REGEX, script),
+        )
+    )
+    if len(matches) > 1:
+        raise ValueError(f"Multiple {name} blocks found")
+    elif len(matches) == 1:
+        match = matches[0]
+        content = "".join(
+            line[2:] if line.startswith("# ") else line[1:]
+            for line in match.group("content").splitlines(keepends=True)
+        )
+        doc = tomlkit.parse(content)
+        return doc, match
+    else:
+        return None, None
+
+
+def embed_pep723_toml(
+    script: str, doc: tomlkit.TOMLDocument, match: re.Match | None
+) -> str:
+    """Replace PEP 723 block with updated TOML document.
+
+    Args:
+        script: The full text content of a Python script
+        doc: tomlkit TOMLDocument to embed
+        match: regex match from extract_pep723_toml, or None to insert new block
+
+    Returns:
+        Updated script content with the new/updated PEP 723 block
+    """
+    # Format TOML document as PEP 723 block
+    toml_content = tomlkit.dumps(doc)
+    lines = ["# /// script"]
+    for line in toml_content.splitlines():
+        if line.strip():
+            lines.append(f"# {line}")
+        else:
+            lines.append("#")
+    lines.append("# ///")
+    metadata_block = "\n".join(lines)
+
+    if match:
+        # Replace existing block
+        return script[: match.start()] + metadata_block + script[match.end() :]
+    else:
+        # Insert at the beginning (after shebang/encoding if present)
+        script_lines = script.split("\n")
+        insert_index = 0
+
+        # Skip shebang line if present
+        if script_lines and script_lines[0].startswith("#!"):
+            insert_index = 1
+
+        # Skip encoding declaration if present
+        if insert_index < len(script_lines) and (
+            script_lines[insert_index].startswith("# -*- coding:")
+            or script_lines[insert_index].startswith("# coding:")
+        ):
+            insert_index += 1
+
+        # Insert metadata block at the appropriate position
+        script_lines.insert(insert_index, metadata_block)
+
+        # Add blank line after metadata if there isn't one
+        if (
+            insert_index + 1 < len(script_lines)
+            and script_lines[insert_index + 1].strip()
+        ):
+            script_lines.insert(insert_index + 1, "")
+
+        return "\n".join(script_lines)
+
+
+def add_endpoint_to_toml(
+    doc: tomlkit.TOMLDocument,
+    endpoint_name: str,
+    endpoint_config: dict[str, Any],
+    variant_name: str | None = None,
+    variant_config: dict[str, Any] | None = None,
+    schema_comments: dict[str, str] | None = None,
+) -> str | None:
+    """Add endpoint config to TOML document in-place.
+
+    Args:
+        doc: tomlkit TOMLDocument to modify
+        endpoint_name: Base endpoint name (e.g., "anvil")
+        endpoint_config: Base endpoint config dict
+        variant_name: Optional variant name (e.g., "gpu")
+        variant_config: Optional variant config dict
+        schema_comments: Optional dict mapping field names to comment strings
+            (e.g., {"account": "Type: string. Your allocation account"})
+
+    Returns:
+        Skip message if endpoint/variant already exists, None on success.
+    """
+    # Ensure tool.hog exists
+    if "tool" not in doc:
+        doc["tool"] = tomlkit.table()
+
+    tool_table = cast(tomlkit.items.Table, doc["tool"])
+    if "hog" not in tool_table:
+        tool_table["hog"] = tomlkit.table()
+
+    hog = cast(tomlkit.items.Table, tool_table["hog"])
+
+    # Check if we're just adding a variant to an existing base
+    if variant_name is not None:
+        if endpoint_name in hog:
+            # Base exists - check if variant exists
+            endpoint_table = cast(tomlkit.items.Table, hog[endpoint_name])
+            if variant_name in endpoint_table:
+                return f"Variant '{endpoint_name}.{variant_name}' already exists"
+            # Add variant to existing base
+            endpoint_table[variant_name] = tomlkit.table()
+            for key, value in (variant_config or {}).items():
+                variant_table = cast(tomlkit.items.Table, endpoint_table[variant_name])
+                variant_table[key] = value
+            return None
+        else:
+            # Base doesn't exist - add base + variant
+            hog[endpoint_name] = tomlkit.table()
+            endpoint_table = cast(tomlkit.items.Table, hog[endpoint_name])
+            for key, value in endpoint_config.items():
+                endpoint_table[key] = value
+            # Add schema comments for fields not already in config
+            _add_schema_comments(endpoint_table, endpoint_config, schema_comments)
+            endpoint_table[variant_name] = tomlkit.table()
+            for key, value in (variant_config or {}).items():
+                variant_table = cast(tomlkit.items.Table, endpoint_table[variant_name])
+                variant_table[key] = value
+            return None
+    else:
+        # Just adding base endpoint
+        if endpoint_name in hog:
+            return f"Endpoint '{endpoint_name}' already exists"
+        hog[endpoint_name] = tomlkit.table()
+        endpoint_table = cast(tomlkit.items.Table, hog[endpoint_name])
+        for key, value in endpoint_config.items():
+            endpoint_table[key] = value
+        # Add schema comments for fields not already in config
+        _add_schema_comments(endpoint_table, endpoint_config, schema_comments)
+        return None
+
+
+def _add_schema_comments(
+    table: tomlkit.items.Table,
+    existing_config: dict[str, Any],
+    schema_comments: dict[str, str] | None,
+) -> None:
+    """Add commented-out schema fields to a tomlkit table.
+
+    Args:
+        table: tomlkit Table to add comments to
+        existing_config: Dict of fields already in the config (to skip)
+        schema_comments: Dict mapping field names to comment strings
+    """
+    if not schema_comments:
+        return
+
+    # Align comments to column 52 (matches format_endpoint_config_to_toml)
+    alignment_column = 52
+
+    for field_name, comment in schema_comments.items():
+        # Skip fields already in the active config
+        if field_name in existing_config:
+            continue
+        # Add as a commented-out field with documentation
+        # Format: # field_name =                  # Type: string. Description
+        # Note: tomlkit.comment adds "# " prefix, embed_pep723_toml adds another "# "
+        # so final output is "# # field_name = {padding}# comment"
+        padding = " " * max(1, alignment_column - 5 - len(field_name))
+        table.add(tomlkit.comment(f"{field_name} ={padding}# {comment}"))
+
+
+def remove_endpoint_from_script(
+    script_content: str, endpoint_name: str, variant_name: str | None = None
+) -> str:
+    """Remove an endpoint or variant from a script's PEP 723 block.
+
+    Args:
+        script_content: Full script file content
+        endpoint_name: Name of the endpoint (e.g., "my_endpoint")
+        variant_name: Optional variant name. If provided, only removes that variant.
+                     If None, removes the entire endpoint and all its variants.
+
+    Returns:
+        Updated script content with the endpoint/variant removed
+    """
+    doc, match = extract_pep723_toml(script_content)
+    if doc is None or match is None:
+        return script_content
+
+    if "tool" in doc:
+        tool_table = cast(tomlkit.items.Table, doc["tool"])
+        if "hog" in tool_table:
+            hog = cast(tomlkit.items.Table, tool_table["hog"])
+
+            if variant_name is not None:
+                # Remove only the specific variant
+                if endpoint_name in hog:
+                    endpoint_table = cast(tomlkit.items.Table, hog[endpoint_name])
+                    if variant_name in endpoint_table:
+                        del endpoint_table[variant_name]
+            else:
+                # Remove the entire endpoint (and all its variants)
+                if endpoint_name in hog:
+                    del hog[endpoint_name]
+
+    return embed_pep723_toml(script_content, doc, match)
+
+
+def add_endpoint_to_script(
+    script_content: str,
+    endpoint_name: str,
+    endpoint_config: dict[str, Any],
+    variant_name: str | None = None,
+    variant_config: dict[str, Any] | None = None,
+    schema_comments: dict[str, str] | None = None,
+) -> tuple[str, str | None]:
+    """Add endpoint config to existing script, preserving formatting.
+
+    Convenience wrapper that combines extract_pep723_toml, add_endpoint_to_toml,
+    and embed_pep723_toml into a single call.
+
+    Args:
+        script_content: Full script file content
+        endpoint_name: Base endpoint name (e.g., "anvil")
+        endpoint_config: Base endpoint config dict
+        variant_name: Optional variant name (e.g., "gpu")
+        variant_config: Optional variant config dict
+        schema_comments: Optional dict mapping field names to comment strings
+            (e.g., {"account": "Type: string. Your allocation account"})
+
+    Returns:
+        Tuple of (updated_content, skip_message)
+        skip_message is None if endpoint was added, or info string if skipped
+    """
+    doc, match = extract_pep723_toml(script_content)
+
+    if doc is None:
+        # No PEP 723 block exists - create minimal one with defaults
+        doc = tomlkit.document()
+        metadata = Pep723Metadata()
+        doc["requires-python"] = metadata.requires_python
+        doc["dependencies"] = []
+
+    skip_msg = add_endpoint_to_toml(
+        doc,
+        endpoint_name,
+        endpoint_config,
+        variant_name,
+        variant_config,
+        schema_comments,
+    )
+
+    updated_content = embed_pep723_toml(script_content, doc, match)
+    return updated_content, skip_msg