[docs] uv configuration best practices (#118)

* move managed-python to toml block instead of flag

* include example/docs for customizing uv behavior

* add custom index example to nav

* clean up docslop

Author: OwenPriceSkelly

docs/api/environment_variables.md

@@ -62,3 +62,48 @@ GROUNDHOG_PROXYSTORE_DIR=/scratch/username/proxystore python script.py
 
 !!! warning "Under Construction 👷🚧"
     Proxystore integration is currently `.local`-only, this does not (yet) have any effect on `.remote` or `.submit` calls.
+
+## GROUNDHOG_CACHE_DIR
+
+**Type:** path string
+
+**Default:** Falls back to `$SCRATCH`, then `$TMPDIR`, then `/tmp`
+
+Directory where uv caches packages and Python installations on remote endpoints. This is used to set `UV_CACHE_DIR` and `UV_PYTHON_INSTALL_DIR` in the remote environment if they are not already set.
+
+**Example:**
+```bash
+export GROUNDHOG_CACHE_DIR=/gpfs/shared/uv-cache
+```
+
+**Why this matters:** HPC clusters often have NFS-mounted home directories that can cause file locking issues or have limited quotas. Using fast scratch storage or a shared cache directory improves performance and avoids these issues.
+
+**Precedence:** Existing `UV_CACHE_DIR` and `UV_PYTHON_INSTALL_DIR` environment variables take precedence over `GROUNDHOG_CACHE_DIR`. If none are set, Groundhog uses this fallback chain:
+1. `$GROUNDHOG_CACHE_DIR` (if set)
+2. `$SCRATCH` (HPC scratch space)
+3. `$TMPDIR` (temporary directory)
+4. `/tmp` (system temp)
+
+## `uv` Environment Variables
+
+Groundhog uses `uv` to manage Python environments on remote endpoints. Any `UV_*` environment variable can be used to override `[tool.uv]` configuration in your script.
+
+**Example - Per-endpoint package index:**
+```toml
+[tool.hog.cpu_endpoint]
+endpoint = "..."
+worker_init = """
+export UV_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cpu
+"""
+
+[tool.hog.gpu_endpoint]
+endpoint = "..."
+worker_init = """
+export UV_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu121
+"""
+```
+
+**See also:**
+
+- [`uv` environment variable reference](https://docs.astral.sh/uv/reference/environment/) - Official documentation of `UV_*` env vars
+- [PEP 723 Concepts](../concepts/pep723.md#configuring-uv-via-tooluv) - Configuring uv via `[tool.uv]`

docs/concepts/pep723.md

@@ -90,13 +90,96 @@ PEP 723 defines standard fields:
 
 Tools can add their own sections under `[tool.*]`:
 
-- `[tool.uv]` - uv-specific settings (e.g., `exclude-newer` for reproducibility)
+- `[tool.uv]` - uv package manager configuration (see below)
 - `[tool.hog.*]` - Groundhog endpoint configurations
 
 Standard fields control the Python environment. Tool-specific fields configure behavior.
 
+## Configuring `uv` via `[tool.uv]`
+
+Groundhog uses `uv` to manage Python environments on remote endpoints. You can configure `uv`'s behavior through the `[tool.uv]` section in your PEP 723 metadata.
+
+### Common `[tool.uv]` settings example:
+
+```python
+# /// script
+# requires-python = ">=3.11"
+# dependencies = ["numpy", "torch"]
+#
+# [tool.uv]
+# exclude-newer = "2025-12-19T00:00:00Z"  # Lock packages to a point in time
+# python-preference = "managed"            # Use uv-managed Python
+# extra-index-url = [                      # Additional package indexes
+#     "https://download.pytorch.org/whl/cpu"
+# ]
+# ///
+```
+
+**See also**: Any [uv settings](https://docs.astral.sh/uv/reference/settings/) can be used in `[tool.uv]` - the configuration is passed through to uv when creating the remote environment.
+
+### Custom package sources with `[tool.uv.sources]`
+
+For finer control over where specific packages come from, use `[tool.uv.sources]`:
+
+```python
+# /// script
+# requires-python = ">=3.11"
+# dependencies = ["torch==2.5.1", "my-internal-lib", "my-github-dependency"]
+#
+# [[tool.uv.index]]
+# name = "pytorch-cpu"
+# url = "https://download.pytorch.org/whl/cpu"
+#
+# [[tool.uv.index]]
+# name = "facility-pypi"
+# url = "https://pypi.facility.gov/simple"
+#
+# [tool.uv.sources] (1)
+# torch = { index = "pytorch-cpu" }
+# my-internal-lib = { index = "facility-pypi" }
+# my-github-dependency = { git = "https://github.com/some-org/my-github-dependency", tag = "1.0.0" }
+# ///
+```
+
+1. See also: [`uv` sources documentation](https://docs.astral.sh/uv/concepts/projects/dependencies/#dependency-sources)
+
+This is useful for:
+
+- Installing PyTorch CPU/CUDA variants from PyTorch's custom wheel server
+- Using private package registries for internal packages
+- Pulling specific packages from Git repositories or local paths
+
+See the [PyTorch Custom Index Example](../examples/pytorch_custom_index.md) for a complete example.
+
+### Configuration precedence
+
+`uv` reads configuration with precedence: **Environment variables > `[tool.uv]` in script**
+
+This means:
+
+- Settings in `[tool.uv]` become the baseline for your script
+- Environment variables like `UV_INDEX_URL` can override them (useful for endpoint-specific configuration)
+
+You can use environment variables to override `[tool.uv]` settings per endpoint:
+
+```toml
+[tool.hog.cpu_cluster]
+endpoint = "..."
+worker_init = """
+export UV_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cpu
+"""
+
+[tool.hog.gpu_cluster]
+endpoint = "..."
+worker_init = """
+export UV_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu121
+"""
+```
+
+This lets the same script work on both CPU and GPU clusters without code changes.
+
 ## Next Steps
 
 - **[Dependencies Example](../examples/dependencies.md)** - Add and use packages
-- **[Configuration Example](../examples/configuration.md)** - What else can `[tool.hog]` do?
-- **[`uv` Scripts Guide](https://docs.astral.sh/uv/guides/scripts/)** - Complete `uv` documentation for PEP 723
+- **[Configuration Example](../examples/configuration.md)** - What `[tool.hog.*]` config blocks do
+- **[`uv` Scripts Guide](https://docs.astral.sh/uv/guides/scripts/)** - Official `uv` reference for PEP 723 scripts

docs/examples/index.md

@@ -17,6 +17,7 @@ Examples showing how to handle typical workflows:
 
 - **[Parallel Execution](parallel-execution.md)** - Using `.submit()` for concurrent remote execution
 - **[Endpoint Configuration](configuration.md)** - How the configuration system merges settings from multiple sources (PEP 723, decorators, call-time overrides)
+- **[PyTorch from Custom Sources](pytorch_custom_index.md)** - Configuring uv to install packages from cluster-specific indexes, local paths, or internal mirrors
 - **[Importing Groundhog Functions](imported_function.md)** - Calling Groundhog functions from regular Python scripts, REPLs, and notebooks (includes import safety and `mark_import_safe()`)
 
 ## Running the Examples

docs/examples/pytorch_custom_index.md

@@ -0,0 +1,178 @@
+# PyTorch from Custom Package Sources
+
+This example demonstrates how to configure uv to install PyTorch from cluster-specific package sources, such as internal mirrors, pre-built wheels on shared filesystems, or custom builds optimized for specific hardware.
+
+## Common HPC Use Cases
+
+- **Cluster-optimized builds**: System admins provide PyTorch wheels optimized for cluster hardware
+- **Internal mirrors**: Packages hosted on internal servers for air-gapped or bandwidth-restricted clusters
+- **Shared filesystem wheels**: Pre-built wheels on `/gpfs` or `/scratch` to avoid repeated downloads
+- **Custom PyTorch builds**: Modified PyTorch with cluster-specific patches or optimizations
+
+## Full Example
+
+```python title="pytorch_custom_index.py"
+# /// script
+# requires-python = ">=3.11,<3.13"
+# dependencies = [
+#     "torch==2.5.1",
+#     "torchvision==0.20.1",
+# ]
+#
+# [tool.uv]
+# exclude-newer = "2025-12-19T00:00:00Z"
+# python-preference = "managed"
+#
+# [[tool.uv.index]]  (1)
+# name = "pytorch-cpu"
+# url = "https://download.pytorch.org/whl/cpu"
+#
+# [tool.uv.sources]  (2)
+# torch = { index = "pytorch-cpu" }
+# torchvision = { index = "pytorch-cpu" }
+#
+# [tool.hog.my_endpoint]
+# endpoint = "your-endpoint-uuid"
+# ///
+
+import groundhog_hpc as hog
+
+
+@hog.function(endpoint="my_endpoint")
+def check_pytorch() -> dict[str, str]:
+    """Check PyTorch installation details."""
+    import torch
+
+    return {
+        "version": torch.__version__,
+        "cuda_available": str(torch.cuda.is_available()),
+        "device": str(torch.device("cuda" if torch.cuda.is_available() else "cpu")),
+    }
+
+
+@hog.function(endpoint="my_endpoint")
+def matrix_multiply(size: int = 1000) -> dict[str, float]:
+    """Simple PyTorch matrix multiplication benchmark."""
+    import time
+
+    import torch
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    start = time.time()
+    a = torch.randn(size, size, device=device)
+    b = torch.randn(size, size, device=device)
+    c = torch.mm(a, b)
+    elapsed = time.time() - start
+
+    return {
+        "size": size,
+        "device": str(device),
+        "time_seconds": elapsed,
+        "mean": float(c.mean()),
+    }
+
+
+@hog.harness()
+def main():
+    """Run PyTorch functions remotely."""
+    info = check_pytorch.remote()
+    print(f"PyTorch {info['version']} on {info['device']}")
+
+    result = matrix_multiply.remote(500)
+    print(f"{result['size']}x{result['size']} matmul: {result['time_seconds']:.3f}s")
+```
+
+1. Define a named index pointing to your package source. In this example, PyTorch's public index for CPU wheels. Replace with your cluster's internal index URL.
+
+2. Specify which packages should use which source. This tells uv to fetch `torch` and `torchvision` from the custom index instead of PyPI.
+
+## Configuration Options
+
+### Custom Package Index
+
+For internal PyPI mirrors or cluster-specific package servers:
+
+```toml
+[[tool.uv.index]]
+name = "cluster-pypi"
+url = "https://pypi.internal.mylab.edu/simple"
+
+[tool.uv.sources]
+torch = { index = "cluster-pypi" }
+```
+
+### Local Filesystem Path
+
+For pre-built wheels on shared storage:
+
+```toml
+[tool.uv.sources]
+torch = { path = "/gpfs/shared/wheels/torch-2.5.1+cu121-cp311-linux_x86_64.whl" }
+```
+
+Or for a local package directory:
+
+```toml
+[tool.uv.sources]
+torch = { path = "/gpfs/shared/pytorch-build", editable = true }
+```
+
+### Direct URL
+
+For wheels hosted on a web server:
+
+```toml
+[tool.uv.sources]
+torch = { url = "https://internal.server.edu/wheels/torch-2.5.1-custom-py3-none-any.whl" }
+```
+
+### Git Repository
+
+For custom builds from Git:
+
+```toml
+[tool.uv.sources]
+torch = { git = "https://github.com/myorg/pytorch", tag = "v2.5.1-custom" }
+```
+
+## Per-Endpoint Configuration
+
+Different endpoints may need different PyTorch builds. Use environment variables to override per endpoint:
+
+```toml
+[tool.hog.cluster_a]
+endpoint = "cluster-a-uuid"
+worker_init = """
+# Cluster A has PyTorch wheels on shared storage
+export UV_FIND_LINKS=/gpfs/cluster-a/wheels
+"""
+
+[tool.hog.cluster_b]
+endpoint = "cluster-b-uuid"
+worker_init = """
+# Cluster B uses an internal PyPI mirror
+export UV_INDEX_URL=https://pypi.cluster-b.edu/simple
+"""
+```
+
+See also: [Environment Variables](../api/environment_variables.md#uv-environment-variables)
+
+## Running the Example
+
+```bash
+hog run pytorch_custom_index.py
+```
+
+Output:
+
+```
+PyTorch 2.5.1 on cuda
+500x500 matmul: 0.015s
+```
+
+## Next Steps
+
+- **[PEP 723 Concepts](../concepts/pep723.md#configuring-uv-via-tooluv)** - Complete uv configuration reference
+- **[Environment Variables](../api/environment_variables.md#uv-environment-variables)** - Override uv settings per endpoint
+- **[uv Dependencies](https://docs.astral.sh/uv/concepts/projects/dependencies/)** - Full uv dependency configuration docs

docs/getting-started/quickstart.md

@@ -24,6 +24,7 @@ This creates `hello.py` with the following structure:
 #
 # [tool.uv]
 # exclude-newer = "2025-12-10T00:00:00Z"
+# python-preference = "managed"
 #
 # [tool.hog.my-endpoint]
 # endpoint = "your-endpoint-uuid"