Auto-tag Vast worker images

README.md

@@ -108,13 +108,14 @@ export VAST_API_KEY='...'
 export COMFY_MODAL_VAST_IMAGE='ghcr.io/owner/comfy-modal-worker@sha256:...'
 ```
 
-Build the shared SSH/Vast worker image locally with an explicit version tag, then push it to a registry and use the digest printed by the script:
+Build the shared SSH/Vast worker image locally, push it to a registry, and use the digest printed by the script (the image must be readable by the rented instance; the current implementation supports public images):
 
 ```bash
-uv run python scripts/build_vast_worker_image.py \
-  --tag ghcr.io/owner/comfy-modal-worker:v0.3.6 --push
+uv run python scripts/build_vast_worker_image.py --push
 ```
 
+By default, the builder reads `project.version` and the GitHub repository owner from `pyproject.toml`, producing `ghcr.io/<owner>/comfy-modal-worker:v<version>`. Use `--owner` or `--tag-template` to change those defaults, or `--tag` to supply the complete image tag explicitly.
+
 When running the builder outside an installed ComfyUI `custom_nodes` directory, it checks `COMFYUI_ROOT`, `~/git/ComfyUI`, and `~/git/Latest_ComfyUI`. Pass `--comfyui-root /path/to/ComfyUI` when the source checkout lives elsewhere.
 
 The image must be readable by the rented instance. The first implementation supports public images; private-registry credential plumbing is intentionally not stored in workflows. `COMFY_MODAL_VAST_SSH_IDENTITY_FILE` may name an absolute private-key path when the Vast account does not use the default SSH identity. `COMFY_MODAL_VAST_API_BASE_URL` exists for the loopback simulator and otherwise accepts only HTTPS.

scripts/build_vast_worker_image.py

@@ -6,9 +6,13 @@ import argparse
 import logging
 import subprocess
 import sys
+import tomllib
+from collections.abc import Mapping
 from dataclasses import replace
 from pathlib import Path
-from typing import Sequence
+from string import Formatter
+from typing import Any, Sequence
+from urllib.parse import urlparse
 
 REPO_ROOT = Path(__file__).resolve().parents[1]
 if str(REPO_ROOT) not in sys.path:
@@ -20,14 +24,33 @@ from ssh_runtime import export_worker_image_context
 
 logger = logging.getLogger(__name__)
 
+DEFAULT_TAG_TEMPLATE = "ghcr.io/{owner}/comfy-modal-worker:v{version}"
+
 
 def _parser() -> argparse.ArgumentParser:
     """Return the command-line parser."""
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
         "--tag",
-        required=True,
-        help="Versioned registry tag, for example ghcr.io/owner/comfy-worker:v0.4.0.",
+        help=(
+            "Explicit versioned registry tag. By default, derive a GHCR tag from "
+            "pyproject.toml."
+        ),
+    )
+    parser.add_argument(
+        "--owner",
+        help=(
+            "Container registry owner used by the default tag template. Defaults to "
+            "the GitHub owner in pyproject.toml's Repository URL."
+        ),
+    )
+    parser.add_argument(
+        "--tag-template",
+        default=DEFAULT_TAG_TEMPLATE,
+        help=(
+            "Default tag template when --tag is omitted. Supports {owner} and "
+            "{version}."
+        ),
     )
     parser.add_argument(
         "--comfyui-root",
@@ -45,6 +68,124 @@ def _parser() -> argparse.ArgumentParser:
     return parser
 
 
+def _project_metadata(repo_root: Path = REPO_ROOT) -> Mapping[str, Any]:
+    """Load and return the PEP 621 project metadata from pyproject.toml."""
+    pyproject_path = repo_root / "pyproject.toml"
+    try:
+        with pyproject_path.open("rb") as handle:
+            document = tomllib.load(handle)
+    except (OSError, tomllib.TOMLDecodeError) as error:
+        raise ValueError(
+            f"Unable to read project metadata from {pyproject_path}."
+        ) from error
+    project = document.get("project")
+    if not isinstance(project, Mapping):
+        raise ValueError(f"Project metadata is missing from {pyproject_path}.")
+    return project
+
+
+def _project_version(project: Mapping[str, Any]) -> str:
+    """Return the normalized project version used in the worker image tag."""
+    version = project.get("version")
+    if not isinstance(version, str) or not version.strip():
+        raise ValueError("pyproject.toml must define a non-empty project.version.")
+    normalized = version.strip()
+    if any(character in normalized for character in "/:@"):
+        raise ValueError(
+            "project.version contains characters that are invalid in a tag."
+        )
+    return normalized
+
+
+def _repository_url(project: Mapping[str, Any]) -> str:
+    """Return the repository URL declared in PEP 621 project metadata."""
+    urls = project.get("urls")
+    if not isinstance(urls, Mapping):
+        raise ValueError("pyproject.toml must define project.urls.Repository.")
+    repository = next(
+        (
+            value
+            for key, value in urls.items()
+            if str(key).casefold() == "repository" and isinstance(value, str)
+        ),
+        None,
+    )
+    if repository is None or not repository.strip():
+        raise ValueError("pyproject.toml must define project.urls.Repository.")
+    return repository.strip()
+
+
+def _github_owner(repository_url: str) -> str:
+    """Extract a normalized GitHub owner from an HTTPS, SSH, or SCP-style URL."""
+    candidate = repository_url.strip()
+    if candidate.startswith("git@github.com:"):
+        repository_path = candidate.removeprefix("git@github.com:")
+    else:
+        parsed = urlparse(candidate)
+        if parsed.hostname != "github.com":
+            raise ValueError(
+                "Repository URL must point to github.com; pass --owner or --tag."
+            )
+        repository_path = parsed.path.lstrip("/")
+    return _normalize_owner(repository_path.split("/", maxsplit=1)[0])
+
+
+def _normalize_owner(owner: str) -> str:
+    """Normalize and validate a GitHub-compatible container registry owner."""
+    normalized = owner.strip().casefold()
+    if not normalized or any(
+        not (character.isascii() and (character.isalnum() or character == "-"))
+        for character in normalized
+    ):
+        raise ValueError(
+            "Unable to infer a valid GHCR owner; pass --owner or --tag."
+        )
+    return normalized
+
+
+def _render_tag_template(tag_template: str, *, owner: str, version: str) -> str:
+    """Render a tag template containing only simple owner and version fields."""
+    try:
+        parsed_fields = tuple(Formatter().parse(tag_template))
+    except ValueError as error:
+        raise ValueError("Tag template is malformed.") from error
+    invalid_field = any(
+        field_name not in {None, "owner", "version"}
+        or bool(format_spec)
+        or conversion is not None
+        for _, field_name, format_spec, conversion in parsed_fields
+    )
+    if invalid_field:
+        raise ValueError(
+            "Tag template may contain only the {owner} and {version} fields."
+        )
+    return tag_template.format(owner=owner, version=version)
+
+
+def _resolve_image_tag(
+    explicit_tag: str | None,
+    *,
+    owner: str | None,
+    tag_template: str,
+    repo_root: Path = REPO_ROOT,
+) -> str:
+    """Return an explicit tag or render the project-derived default image tag."""
+    if explicit_tag is not None:
+        return _validate_tag(explicit_tag)
+    project = _project_metadata(repo_root)
+    resolved_owner = (
+        _github_owner(_repository_url(project))
+        if owner is None
+        else _normalize_owner(owner)
+    )
+    rendered = _render_tag_template(
+        tag_template,
+        owner=resolved_owner,
+        version=_project_version(project),
+    )
+    return _validate_tag(rendered)
+
+
 def _validate_tag(tag: str) -> str:
     """Require an explicit version tag rather than a mutable latest reference."""
     normalized = tag.strip()
@@ -154,8 +295,13 @@ def main(argv: Sequence[str] | None = None) -> int:
     """Build the requested image and print the exact runtime configuration value."""
     arguments = _parser().parse_args(argv)
     logging.basicConfig(level=logging.INFO)
-    image_reference = build_image(
+    image_tag = _resolve_image_tag(
         arguments.tag,
+        owner=arguments.owner,
+        tag_template=arguments.tag_template,
+    )
+    image_reference = build_image(
+        image_tag,
         push=arguments.push,
         comfyui_root=arguments.comfyui_root,
     )

tests/test_build_vast_worker_image.py

@@ -47,3 +47,58 @@ def test_parser_accepts_explicit_comfyui_source_root(tmp_path: Path) -> None:
     )
 
     assert arguments.comfyui_root == tmp_path
+
+
+def test_default_tag_uses_pyproject_version_and_repository_owner(
+    tmp_path: Path,
+) -> None:
+    """The normal push command should need no manually duplicated version tag."""
+    module = _module()
+    (tmp_path / "pyproject.toml").write_text(
+        """
+[project]
+name = "modal-sync"
+version = "1.2.3"
+
+[project.urls]
+Repository = "https://github.com/Example-Owner/ComfyUI-Modal"
+""".strip(),
+        encoding="utf-8",
+    )
+
+    assert module._resolve_image_tag(
+        None,
+        owner=None,
+        tag_template=module.DEFAULT_TAG_TEMPLATE,
+        repo_root=tmp_path,
+    ) == "ghcr.io/example-owner/comfy-modal-worker:v1.2.3"
+
+
+def test_default_tag_accepts_owner_and_template_overrides(tmp_path: Path) -> None:
+    """Non-default organizations and registries remain configurable without --tag."""
+    module = _module()
+    (tmp_path / "pyproject.toml").write_text(
+        """
+[project]
+name = "modal-sync"
+version = "2.0.0rc1"
+""".strip(),
+        encoding="utf-8",
+    )
+
+    assert module._resolve_image_tag(
+        None,
+        owner="Container-Org",
+        tag_template="registry.example/{owner}/vast-worker:{version}",
+        repo_root=tmp_path,
+    ) == "registry.example/container-org/vast-worker:2.0.0rc1"
+
+
+def test_parser_allows_push_without_explicit_tag() -> None:
+    """The documented one-line push command must parse without --tag."""
+    module = _module()
+
+    arguments = module._parser().parse_args(["--push"])
+
+    assert arguments.push is True
+    assert arguments.tag is None