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

groundhog_hpc/__init__.py

@@ -32,6 +32,7 @@ import os
 
 from groundhog_hpc.decorators import function, harness, method
 from groundhog_hpc.import_hook import install_import_hook
+from groundhog_hpc.logging import setup_logging
 from groundhog_hpc.utils import mark_import_safe
 
 try:
@@ -41,5 +42,8 @@ except importlib.metadata.PackageNotFoundError:
 
 __all__ = ["function", "harness", "method", "mark_import_safe", "__version__"]
 
+# Configure logging on import
+setup_logging()
+
 if not os.environ.get("GROUNDHOG_NO_IMPORT_HOOK"):
     install_import_hook()

groundhog_hpc/app/add.py

@@ -1,5 +1,6 @@
 """Add command for managing PEP 723 script dependencies."""
 
+import os
 import subprocess
 from pathlib import Path
 
@@ -11,9 +12,23 @@ from groundhog_hpc.app.utils import (
     normalize_python_version_with_uv,
     update_requires_python,
 )
+from groundhog_hpc.configuration.endpoints import (
+    KNOWN_ENDPOINTS,
+    get_endpoint_schema_comments,
+    parse_endpoint_spec,
+)
+from groundhog_hpc.configuration.pep723 import add_endpoint_to_script
+from groundhog_hpc.logging import setup_logging
 
 console = Console()
 
+KNOWN_ENDPOINT_ALIASES = []
+for name in KNOWN_ENDPOINTS.keys():
+    KNOWN_ENDPOINT_ALIASES += [name]
+    KNOWN_ENDPOINT_ALIASES += [
+        f"{name}.{variant}" for variant in KNOWN_ENDPOINTS[name]["variants"].keys()
+    ]
+
 
 def add(
     script: Path = typer.Argument(..., help="Path to the script to modify"),
@@ -24,8 +39,27 @@ def add(
     python: str | None = typer.Option(
         None, "--python", "-p", help="Python version specifier"
     ),
+    endpoints: list[str] = typer.Option(
+        [],
+        "--endpoint",
+        "-e",
+        help=(
+            "Add endpoint configuration (e.g., anvil, anvil.gpu, name:uuid). "
+            f"Known endpoints: {', '.join(KNOWN_ENDPOINT_ALIASES)}. Can specify multiple."
+        ),
+    ),
+    log_level: str = typer.Option(
+        None,
+        "--log-level",
+        help="Set logging level (DEBUG, INFO, WARNING, ERROR)\n\n[env: GROUNDHOG_LOG_LEVEL=]",
+    ),
 ) -> None:
     """Add dependencies or update Python version in a script's PEP 723 metadata."""
+    if log_level:
+        os.environ["GROUNDHOG_LOG_LEVEL"] = log_level.upper()
+        # Reconfigure logging with the new level
+        setup_logging()
+
     if not script.exists():
         console.print(f"[red]Error: Script '{script}' not found[/red]")
         raise typer.Exit(1)
@@ -60,3 +94,43 @@ def add(
         except subprocess.CalledProcessError as e:
             console.print(f"[red]{e.stderr.strip()}[/red]")
             raise typer.Exit(1)
+
+    # handle --endpoint flags
+    if endpoints:
+        content = script.read_text()
+        added_any = False
+
+        for endpoint_spec_str in endpoints:
+            try:
+                spec = parse_endpoint_spec(endpoint_spec_str)
+            except Exception as e:
+                console.print(f"[red]Error: {e}[/red]")
+                raise typer.Exit(1)
+
+            # Build config dict from the spec
+            endpoint_config = {"endpoint": spec.uuid, **spec.base_defaults}
+            variant_config = spec.variant_defaults if spec.variant else None
+
+            # Fetch schema comments if UUID is valid (not a TODO placeholder)
+            schema_comments = None
+            if not spec.uuid.startswith("TODO"):
+                schema_comments = get_endpoint_schema_comments(spec.uuid)
+
+            content, skip_msg = add_endpoint_to_script(
+                content,
+                endpoint_name=spec.name,
+                endpoint_config=endpoint_config,
+                variant_name=spec.variant,
+                variant_config=variant_config,
+                schema_comments=schema_comments,
+            )
+
+            if skip_msg:
+                console.print(f"[yellow]{skip_msg}[/yellow]")
+            else:
+                added_any = True
+
+        script.write_text(content)
+
+        if added_any:
+            console.print(f"[green]Added endpoint configuration to {script}[/green]")

groundhog_hpc/app/init.py

@@ -1,5 +1,6 @@
 """Init command for creating new Groundhog scripts."""
 
+import os
 import subprocess
 from pathlib import Path
 from typing import Optional
@@ -11,9 +12,15 @@ from rich.console import Console
 from groundhog_hpc.app.utils import normalize_python_version_with_uv
 from groundhog_hpc.configuration.endpoints import (
     KNOWN_ENDPOINTS,
-    fetch_and_format_endpoints,
+    get_endpoint_schema_comments,
+    parse_endpoint_spec,
 )
-from groundhog_hpc.configuration.pep723 import Pep723Metadata
+from groundhog_hpc.configuration.pep723 import (
+    Pep723Metadata,
+    add_endpoint_to_script,
+    remove_endpoint_from_script,
+)
+from groundhog_hpc.logging import setup_logging
 
 console = Console()
 
@@ -46,8 +53,18 @@ def init(
             "Can specify multiple."
         ),
     ),
+    log_level: str = typer.Option(
+        None,
+        "--log-level",
+        help="Set logging level (DEBUG, INFO, WARNING, ERROR)\n\n[env: GROUNDHOG_LOG_LEVEL=]",
+    ),
 ) -> None:
     """Create a new groundhog script with PEP 723 metadata and example code."""
+    if log_level:
+        os.environ["GROUNDHOG_LOG_LEVEL"] = log_level.upper()
+        # Reconfigure logging with the new level
+        setup_logging()
+
     if Path(filename).exists():
         console.print(f"[red]Error: {filename} already exists[/red]")
         raise typer.Exit(1)
@@ -67,32 +84,59 @@ def init(
     assert default_meta.tool and default_meta.tool.uv
     exclude_newer = default_meta.tool.uv.exclude_newer
 
-    # Fetch and format endpoint configurations if provided
-    endpoint_blocks = []
+    # Parse endpoint specs if provided
+    endpoint_specs = []
     if endpoints:
         try:
-            endpoint_blocks = fetch_and_format_endpoints(endpoints)
-            for endpoint in endpoint_blocks:
-                console.print(f"[green]✓[/green] Fetched schema for {endpoint.name}")
+            endpoint_specs = [parse_endpoint_spec(spec) for spec in endpoints]
         except Exception as e:
             console.print(f"[red]Error: {e}[/red]")
             raise typer.Exit(1)
 
+    # Determine endpoint name for decorator (first endpoint or placeholder)
+    first_endpoint_name = endpoint_specs[0].name if endpoint_specs else "my_endpoint"
+
+    # Render template (always includes my_endpoint placeholder)
     env = Environment(loader=PackageLoader("groundhog_hpc", "templates"))
     template = env.get_template("init_script.py.jinja")
     content = template.render(
         filename=filename,
         python=python,
         exclude_newer=exclude_newer,
-        endpoint_blocks=endpoint_blocks,
+        endpoint_name=first_endpoint_name,
     )
+
+    # If endpoints provided, replace placeholder with real endpoints
+    if endpoint_specs:
+        # Remove placeholder
+        content = remove_endpoint_from_script(content, "my_endpoint")
+
+        # Add each requested endpoint
+        for spec in endpoint_specs:
+            endpoint_config = {"endpoint": spec.uuid, **spec.base_defaults}
+            variant_config = spec.variant_defaults if spec.variant else None
+
+            # Fetch schema comments if UUID is valid (not a TODO placeholder)
+            schema_comments = None
+            if not spec.uuid.startswith("TODO"):
+                schema_comments = get_endpoint_schema_comments(spec.uuid)
+
+            content, _ = add_endpoint_to_script(
+                content,
+                endpoint_name=spec.name,
+                endpoint_config=endpoint_config,
+                variant_name=spec.variant,
+                variant_config=variant_config,
+                schema_comments=schema_comments,
+            )
+
     Path(filename).write_text(content)
 
     console.print(f"[green]✓[/green] Created {filename}")
-    if endpoint_blocks:
+    if endpoint_specs:
         console.print("\nNext steps:")
         console.print(
-            f"  1. Update fields in the \\[tool.hog.{endpoint_blocks[0].name}] block"
+            f"  1. Update fields in the \\[tool.hog.{endpoint_specs[0].name}] block"
         )
         console.print(f"  2. Run with: [bold]hog run {filename} main[/bold]")
     else:

groundhog_hpc/app/main.py

@@ -16,7 +16,11 @@ from groundhog_hpc.app.run import run
 
 app = typer.Typer(pretty_exceptions_show_locals=False)
 
-app.command(no_args_is_help=True)(run)
+# Enable extra args for run command to capture harness arguments after --
+app.command(
+    no_args_is_help=True,
+    context_settings={"allow_extra_args": True, "allow_interspersed_args": False},
+)(run)
 app.command(no_args_is_help=True)(init)
 app.command(no_args_is_help=True)(add)
 app.command(no_args_is_help=True)(remove)

groundhog_hpc/app/remove.py

@@ -1,5 +1,6 @@
 """Remove command for managing PEP 723 script dependencies."""
 
+import os
 import subprocess
 from pathlib import Path
 
@@ -7,31 +8,105 @@ import typer
 import uv
 from rich.console import Console
 
+from groundhog_hpc.configuration.endpoints import KNOWN_ENDPOINTS
+from groundhog_hpc.configuration.pep723 import remove_endpoint_from_script
+from groundhog_hpc.logging import setup_logging
+
 console = Console()
 
+KNOWN_ENDPOINT_ALIASES = []
+for name in KNOWN_ENDPOINTS.keys():
+    KNOWN_ENDPOINT_ALIASES += [name]
+    KNOWN_ENDPOINT_ALIASES += [
+        f"{name}.{variant}" for variant in KNOWN_ENDPOINTS[name]["variants"].keys()
+    ]
+
 
 def remove(
     script: Path = typer.Argument(..., help="Path to the script to modify"),
-    packages: list[str] = typer.Argument(..., help="Packages to remove"),
+    packages: list[str] | None = typer.Argument(None, help="Packages to remove"),
+    endpoints: list[str] = typer.Option(
+        [],
+        "--endpoint",
+        "-e",
+        help=(
+            "Remove endpoint or variant configuration (e.g., anvil, anvil.gpu, my_endpoint). "
+            f"Known endpoints: {', '.join(KNOWN_ENDPOINT_ALIASES)}. Can specify multiple. "
+            "Note: Removing a base endpoint (e.g., anvil) removes all its variants. "
+            "Removing a specific variant (e.g., anvil.gpu) leaves the base and other variants intact."
+        ),
+    ),
+    log_level: str = typer.Option(
+        None,
+        "--log-level",
+        help="Set logging level (DEBUG, INFO, WARNING, ERROR)\n\n[env: GROUNDHOG_LOG_LEVEL=]",
+    ),
 ) -> None:
     """Remove dependencies from a script's PEP 723 metadata."""
+    if log_level:
+        os.environ["GROUNDHOG_LOG_LEVEL"] = log_level.upper()
+        # Reconfigure logging with the new level
+        setup_logging()
+
     # Validate script exists
     if not script.exists():
         console.print(f"[red]Error: Script '{script}' not found[/red]")
         raise typer.Exit(1)
 
-    # Shell out to uv
-    cmd = [f"{uv.find_uv_bin()}", "remove", "--script", str(script)]
-    cmd.extend(packages)
-
-    try:
-        subprocess.run(
-            cmd,
-            check=True,
-            capture_output=True,
-            text=True,
-        )
-        console.print(f"[green]Removed packages from {script}[/green]")
-    except subprocess.CalledProcessError as e:
-        console.print(f"[red]{e.stderr.strip()}[/red]")
-        raise typer.Exit(1)
+    # Handle package removal
+    packages = packages or []
+    if packages:
+        # Shell out to uv
+        cmd = [f"{uv.find_uv_bin()}", "remove", "--script", str(script)]
+        cmd.extend(packages)
+
+        try:
+            subprocess.run(
+                cmd,
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            console.print(f"[green]Removed packages from {script}[/green]")
+        except subprocess.CalledProcessError as e:
+            console.print(f"[red]{e.stderr.strip()}[/red]")
+            raise typer.Exit(1)
+
+    # Handle endpoint removal
+    if endpoints:
+        content = script.read_text()
+        removed_any = False
+
+        for endpoint_spec in endpoints:
+            # Parse endpoint spec to extract base name and optional variant
+            # Format can be: "name", "name.variant", or "name:uuid"
+            # Split by ':' first to handle "name:uuid" or "name.variant:uuid"
+            name_part = endpoint_spec.split(":")[0]
+
+            # Check if user specified a variant
+            if "." in name_part:
+                base_name, variant_name = name_part.split(".", 1)
+            else:
+                base_name = name_part
+                variant_name = None
+
+            original_content = content
+            content = remove_endpoint_from_script(content, base_name, variant_name)
+
+            if content != original_content:
+                removed_any = True
+            else:
+                if variant_name:
+                    console.print(
+                        f"[yellow]Variant '{base_name}.{variant_name}' not found in {script}[/yellow]"
+                    )
+                else:
+                    console.print(
+                        f"[yellow]Endpoint '{base_name}' not found in {script}[/yellow]"
+                    )
+
+        if removed_any:
+            script.write_text(content)
+            console.print(
+                f"[green]Removed endpoint configuration(s) from {script}[/green]"
+            )

groundhog_hpc/app/run.py

@@ -1,8 +1,10 @@
 """Run command for executing Groundhog scripts on Globus Compute endpoints."""
 
+import inspect
 import os
 import sys
 from pathlib import Path
+from typing import Any
 
 import typer
 
@@ -13,6 +15,7 @@ from groundhog_hpc.app.utils import (
 from groundhog_hpc.configuration.pep723 import read_pep723
 from groundhog_hpc.errors import RemoteExecutionError
 from groundhog_hpc.harness import Harness
+from groundhog_hpc.logging import setup_logging
 from groundhog_hpc.utils import (
     get_groundhog_version_spec,
     import_user_script,
@@ -20,7 +23,38 @@ from groundhog_hpc.utils import (
 )
 
 
+def invoke_harness_with_args(harness: Harness, args: list[str]) -> Any:
+    """Parse CLI args and invoke harness function.
+
+    Reproduces typer.run() logic but with explicit args and standalone_mode=False
+    to capture return values and exceptions instead of sys.exit().
+
+    Args:
+        harness: The harness to invoke
+        args: CLI arguments to parse (e.g., ["arg1", "--count=5"])
+
+    Returns:
+        The return value from the harness function
+
+    Raises:
+        SystemExit: If argument parsing fails (from Click/Typer)
+        Any exception raised by the harness function
+    """
+    original_argv = sys.argv
+    # Use harness name for better help/error messages
+    sys.argv = [harness.func.__name__] + args
+
+    try:
+        app = typer.Typer(add_completion=False)
+        app.command()(harness.func)
+        result = app(standalone_mode=False)
+        return result
+    finally:
+        sys.argv = original_argv
+
+
 def run(
+    ctx: typer.Context,
     script: Path = typer.Argument(
         ..., help="Path to script with PEP 723 dependencies to deploy to the endpoint"
     ),
@@ -32,11 +66,34 @@ def run(
         "--no-fun-allowed",
         help="Suppress emoji output\n\n[env: GROUNDHOG_NO_FUN_ALLOWED=]",
     ),
+    log_level: str = typer.Option(
+        None,
+        "--log-level",
+        help="Set logging level (DEBUG, INFO, WARNING, ERROR)\n\n[env: GROUNDHOG_LOG_LEVEL=]",
+    ),
 ) -> None:
-    """Run a Python script on a Globus Compute endpoint."""
+    """Run a Python script on a Globus Compute endpoint.
+
+    Use -- to pass arguments to parameterized harnesses:
+        hog run script.py harness -- arg1 --option=value
+    """
+    # Handle the -- separator for harness arguments
+    # ctx.args may contain ['--', 'arg1', 'arg2'] - strip the '--' if present
+    harness_args = ctx.args  # List of extra args, or empty list
+    if harness_args and harness_args[0] == "--":
+        harness_args = harness_args[1:]  # Strip leading '--'
+    if harness == "--":
+        # User typed: hog run script.py -- args
+        # Use default harness "main" and shift args
+        harness = "main"
     if no_fun_allowed:
         os.environ["GROUNDHOG_NO_FUN_ALLOWED"] = str(no_fun_allowed)
 
+    if log_level:
+        os.environ["GROUNDHOG_LOG_LEVEL"] = log_level.upper()
+        # Reconfigure logging with the new level
+        setup_logging()
+
     script_path = script.resolve()
     if not script_path.exists():
         typer.echo(f"Error: Script '{script_path}' not found", err=True)
@@ -87,7 +144,18 @@ def run(
             )
             raise typer.Exit(1)
 
-        result = harness_func()
+        # Dispatch based on whether harness arguments were provided
+        if not harness_args:
+            # No extra args: zero-arg call (backward compatible)
+            result = harness_func()
+        else:
+            # Has extra args: parse and invoke parameterized harness
+            sig = inspect.signature(harness_func.func)
+            if len(sig.parameters) == 0:
+                typer.echo(f"Error: Harness '{harness}' takes no arguments", err=True)
+                raise typer.Exit(1)
+            result = invoke_harness_with_args(harness_func, harness_args)
+
         typer.echo(result)
     except RemoteExecutionError as e:
         if e.returncode == 124:

groundhog_hpc/compute.py

@@ -5,6 +5,7 @@ ShellFunctions, registering them, and submitting them for execution on remote
 endpoints.
 """
 
+import logging
 import os
 import warnings
 from functools import lru_cache
@@ -14,6 +15,8 @@ from uuid import UUID
 from groundhog_hpc.future import GroundhogFuture
 from groundhog_hpc.templating import template_shell_command
 
+logger = logging.getLogger(__name__)
+
 warnings.filterwarnings(
     "ignore",
     category=UserWarning,
@@ -92,10 +95,21 @@ def submit_to_executor(
         expected_keys = set(schema.get("properties", {}).keys())
         unexpected_keys = set(config.keys()) - expected_keys
         if unexpected_keys:
+            logger.debug(
+                f"Filtering unexpected config keys for endpoint {endpoint}: {unexpected_keys}"
+            )
             config = {k: v for k, v in config.items() if k not in unexpected_keys}
 
+    logger.debug(f"Creating Globus Compute executor for endpoint {endpoint}")
     with gc.Executor(endpoint, user_endpoint_config=config) as executor:
+        func_name = getattr(
+            shell_function, "__name__", getattr(shell_function, "name", "unknown")
+        )
+        logger.info(f"Submitting function '{func_name}' to endpoint '{endpoint}'")
         future = executor.submit(shell_function)
+        task_id = getattr(future, "task_id", None)
+        if task_id:
+            logger.info(f"Task submitted with ID: {task_id}")
         deserializing_future = GroundhogFuture(future)
         return deserializing_future
 
@@ -114,7 +128,8 @@ def get_task_status(task_id: str | UUID | None) -> dict[str, Any]:
         return {"status": "status pending", "exception": None}
 
     client = _get_compute_client()
-    return client.get_task(task_id)
+    task_status = client.get_task(task_id)
+    return task_status
 
 
 @lru_cache

groundhog_hpc/configuration/defaults.py

@@ -1,6 +1,7 @@
 # Default Globus Compute Executor configuration
 DEFAULT_USER_CONFIG = {
     "worker_init": "",
+    "endpoint_setup": "",
 }
 
 # default maximum execution time for remote functions (in seconds)