feat(refit): add immutable revision catalog service for S3 delta weight sync

modelexpress_client/python/generate_proto.sh

@@ -13,27 +13,36 @@ SPDX_HEADER="# SPDX-FileCopyrightText: Copyright (c) 2025-${YEAR} NVIDIA CORPORA
 # SPDX-License-Identifier: Apache-2.0
 #"
 
+PROTO_NAMES=(p2p revision)
+PROTO_FILES=()
+for name in "${PROTO_NAMES[@]}"; do
+    PROTO_FILES+=("${PROTO_DIR}/${name}.proto")
+done
+
 # Generate protobuf files
-echo "Generating protobuf files from ${PROTO_DIR}/p2p.proto..."
+echo "Generating protobuf files from ${PROTO_NAMES[*]}..."
 python -m grpc_tools.protoc \
     "-I${PROTO_DIR}" \
     "--python_out=${OUT_DIR}" \
     "--grpc_python_out=${OUT_DIR}" \
-    "${PROTO_DIR}/p2p.proto"
-
-# Fix relative import in grpc file
-echo "Fixing imports in p2p_pb2_grpc.py..."
-tmp_file="$(mktemp)"
-sed 's/^import p2p_pb2 as/from . import p2p_pb2 as/' "${OUT_DIR}/p2p_pb2_grpc.py" > "${tmp_file}"
-mv "${tmp_file}" "${OUT_DIR}/p2p_pb2_grpc.py"
-
-# Add SPDX header to generated files
-for file in "${OUT_DIR}/p2p_pb2.py" "${OUT_DIR}/p2p_pb2_grpc.py"; do
-    echo "Adding SPDX header to ${file}..."
-    tmp_file=$(mktemp)
-    echo "${SPDX_HEADER}" > "${tmp_file}"
-    cat "${file}" >> "${tmp_file}"
-    mv "${tmp_file}" "${file}"
+    "${PROTO_FILES[@]}"
+
+for name in "${PROTO_NAMES[@]}"; do
+    grpc_file="${OUT_DIR}/${name}_pb2_grpc.py"
+    echo "Fixing imports in ${name}_pb2_grpc.py..."
+    tmp_file="$(mktemp)"
+    sed -E 's/^import ([a-zA-Z0-9_]+_pb2) as/from . import \1 as/' "${grpc_file}" > "${tmp_file}"
+    mv "${tmp_file}" "${grpc_file}"
+    sed -i "s/+ f' but the generated code/+ ' but the generated code/" "${grpc_file}"
+    sed -i '/^import warnings$/d' "${grpc_file}"
+
+    for file in "${OUT_DIR}/${name}_pb2.py" "${grpc_file}"; do
+        echo "Adding SPDX header to ${file}..."
+        tmp_file="$(mktemp)"
+        printf '%s\n' "${SPDX_HEADER}" > "${tmp_file}"
+        cat "${file}" >> "${tmp_file}"
+        mv "${tmp_file}" "${file}"
+    done
 done
 
 echo "Done."

modelexpress_client/python/modelexpress/p2p_pb2_grpc.py

@@ -4,7 +4,6 @@
 # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
 """Client and server classes corresponding to protobuf-defined services."""
 import grpc
-import warnings
 
 from . import p2p_pb2 as p2p__pb2
 
@@ -21,7 +20,7 @@
 if _version_not_supported:
     raise RuntimeError(
         f'The grpc package installed is at version {GRPC_VERSION},'
-        + f' but the generated code in p2p_pb2_grpc.py depends on'
+        + ' but the generated code in p2p_pb2_grpc.py depends on'
         + f' grpcio>={GRPC_GENERATED_VERSION}.'
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'

modelexpress_client/python/modelexpress/refit/catalog.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Typed boundary over the three minimal revision-catalog RPCs."""
+
+from __future__ import annotations
+
+import math
+from typing import Protocol, runtime_checkable
+
+import grpc
+
+from modelexpress import revision_pb2, revision_pb2_grpc
+
+from .manifest import RevisionManifest, RevisionRecord
+
+
+@runtime_checkable
+class RevisionCatalog(Protocol):
+    """Exact metadata operations available to the publisher and orchestrator."""
+
+    def publish_revision(self, manifest: RevisionManifest) -> RevisionRecord: ...
+
+    def get_revision(
+        self, model_id: str, target_version: str
+    ) -> RevisionRecord: ...
+
+    def commit_revision(
+        self, model_id: str, target_version: str
+    ) -> RevisionRecord: ...
+
+
+class GrpcRevisionCatalog:
+    """Concrete :class:`RevisionCatalog` over the generated gRPC service."""
+
+    def __init__(
+        self,
+        endpoint: str | None = None,
+        stub=None,
+        timeout: float = 10.0,
+    ) -> None:
+        if not math.isfinite(timeout) or timeout <= 0:
+            raise ValueError("catalog RPC timeout must be finite and positive")
+        if (endpoint is None) == (stub is None):
+            raise ValueError(
+                "GrpcRevisionCatalog needs exactly one of endpoint or stub"
+            )
+        self._channel = None
+        if stub is None:
+            assert endpoint is not None
+            if endpoint.startswith("https://"):
+                target = endpoint.removeprefix("https://")
+                self._channel = grpc.secure_channel(
+                    target, grpc.ssl_channel_credentials()
+                )
+            else:
+                target = endpoint.removeprefix("http://")
+                self._channel = grpc.insecure_channel(target)
+            stub = revision_pb2_grpc.RevisionCatalogServiceStub(self._channel)
+        self._stub = stub
+        self._timeout = timeout
+
+    def __enter__(self) -> GrpcRevisionCatalog:
+        return self
+
+    def __exit__(self, *_exc_info) -> None:
+        self.close()
+
+    def close(self) -> None:
+        if self._channel is not None:
+            self._channel.close()
+            self._channel = None
+
+    def publish_revision(self, manifest: RevisionManifest) -> RevisionRecord:
+        response = self._stub.PublishRevision(
+            revision_pb2.PublishRevisionRequest(manifest=manifest.to_proto()),
+            timeout=self._timeout,
+        )
+        return RevisionRecord.from_proto(response)
+
+    def get_revision(self, model_id: str, target_version: str) -> RevisionRecord:
+        response = self._stub.GetRevision(
+            revision_pb2.GetRevisionRequest(
+                model_id=model_id,
+                target_version=target_version,
+            ),
+            timeout=self._timeout,
+        )
+        return RevisionRecord.from_proto(response)
+
+    def commit_revision(self, model_id: str, target_version: str) -> RevisionRecord:
+        response = self._stub.CommitRevision(
+            revision_pb2.CommitRevisionRequest(
+                model_id=model_id,
+                target_version=target_version,
+            ),
+            timeout=self._timeout,
+        )
+        return RevisionRecord.from_proto(response)

modelexpress_client/python/modelexpress/refit/manifest.py

@@ -0,0 +1,105 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Minimal revision-catalog DTOs and their exact protobuf mapping."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import IntEnum
+
+from modelexpress import revision_pb2
+
+
+class RevisionState(IntEnum):
+    UNSPECIFIED = revision_pb2.REVISION_STATE_UNSPECIFIED
+    READY = revision_pb2.REVISION_STATE_READY
+    COMMITTED = revision_pb2.REVISION_STATE_COMMITTED
+
+
+@dataclass(frozen=True)
+class S3Object:
+    bucket: str
+    key: str
+    checksum: str
+    object_version: str | None = None
+
+    def to_proto(self) -> revision_pb2.S3Object:
+        fields = {
+            "bucket": self.bucket,
+            "key": self.key,
+            "checksum": self.checksum,
+        }
+        if self.object_version is not None:
+            fields["object_version"] = self.object_version
+        return revision_pb2.S3Object(**fields)
+
+    @classmethod
+    def from_proto(cls, proto: revision_pb2.S3Object) -> S3Object:
+        return cls(
+            bucket=proto.bucket,
+            key=proto.key,
+            checksum=proto.checksum,
+            object_version=(
+                proto.object_version if proto.HasField("object_version") else None
+            ),
+        )
+
+
+@dataclass(frozen=True)
+class RevisionManifest:
+    model_id: str
+    target_version: str
+    target_digest: str
+    format_digest: str
+    base_version: str | None = None
+    base_digest: str | None = None
+    payload: S3Object | None = None
+
+    def to_proto(self) -> revision_pb2.RevisionManifest:
+        fields = {
+            "model_id": self.model_id,
+            "target_version": self.target_version,
+            "target_digest": self.target_digest,
+            "format_digest": self.format_digest,
+        }
+        if self.base_version is not None:
+            fields["base_version"] = self.base_version
+        if self.base_digest is not None:
+            fields["base_digest"] = self.base_digest
+        if self.payload is not None:
+            fields["payload"] = self.payload.to_proto()
+        return revision_pb2.RevisionManifest(**fields)
+
+    @classmethod
+    def from_proto(cls, proto: revision_pb2.RevisionManifest) -> RevisionManifest:
+        return cls(
+            model_id=proto.model_id,
+            target_version=proto.target_version,
+            target_digest=proto.target_digest,
+            format_digest=proto.format_digest,
+            base_version=(proto.base_version if proto.HasField("base_version") else None),
+            base_digest=(proto.base_digest if proto.HasField("base_digest") else None),
+            payload=(S3Object.from_proto(proto.payload) if proto.HasField("payload") else None),
+        )
+
+
+@dataclass(frozen=True)
+class RevisionRecord:
+    manifest: RevisionManifest
+    state: RevisionState
+
+    def to_proto(self) -> revision_pb2.RevisionRecord:
+        return revision_pb2.RevisionRecord(
+            manifest=self.manifest.to_proto(),
+            state=int(self.state),
+        )
+
+    @classmethod
+    def from_proto(cls, proto: revision_pb2.RevisionRecord) -> RevisionRecord:
+        if not proto.HasField("manifest"):
+            raise ValueError("revision record is missing manifest")
+        return cls(
+            manifest=RevisionManifest.from_proto(proto.manifest),
+            state=RevisionState(proto.state),
+        )

modelexpress_client/python/modelexpress/revision_pb2.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: revision.proto
+# Protobuf Python Version: 5.27.2
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    5,
+    27,
+    2,
+    '',
+    'revision.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0erevision.proto\x12\x16model_express.revision\"i\n\x08S3Object\x12\x0e\n\x06\x62ucket\x18\x01 \x01(\t\x12\x0b\n\x03key\x18\x02 \x01(\t\x12\x1b\n\x0eobject_version\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x10\n\x08\x63hecksum\x18\x04 \x01(\tB\x11\n\x0f_object_version\"\xf3\x01\n\x10RevisionManifest\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x16\n\x0etarget_version\x18\x02 \x01(\t\x12\x19\n\x0c\x62\x61se_version\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x18\n\x0b\x62\x61se_digest\x18\x04 \x01(\tH\x01\x88\x01\x01\x12\x15\n\rtarget_digest\x18\x05 \x01(\t\x12\x15\n\rformat_digest\x18\x06 \x01(\t\x12\x31\n\x07payload\x18\x07 \x01(\x0b\x32 .model_express.revision.S3ObjectB\x0f\n\r_base_versionB\x0e\n\x0c_base_digest\"\x82\x01\n\x0eRevisionRecord\x12:\n\x08manifest\x18\x01 \x01(\x0b\x32(.model_express.revision.RevisionManifest\x12\x34\n\x05state\x18\x02 \x01(\x0e\x32%.model_express.revision.RevisionState\"T\n\x16PublishRevisionRequest\x12:\n\x08manifest\x18\x01 \x01(\x0b\x32(.model_express.revision.RevisionManifest\">\n\x12GetRevisionRequest\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x16\n\x0etarget_version\x18\x02 \x01(\t\"A\n\x15\x43ommitRevisionRequest\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x16\n\x0etarget_version\x18\x02 \x01(\t*g\n\rRevisionState\x12\x1e\n\x1aREVISION_STATE_UNSPECIFIED\x10\x00\x12\x18\n\x14REVISION_STATE_READY\x10\x01\x12\x1c\n\x18REVISION_STATE_COMMITTED\x10\x02\x32\xcf\x02\n\x16RevisionCatalogService\x12i\n\x0fPublishRevision\x12..model_express.revision.PublishRevisionRequest\x1a&.model_express.revision.RevisionRecord\x12\x61\n\x0bGetRevision\x12*.model_express.revision.GetRevisionRequest\x1a&.model_express.revision.RevisionRecord\x12g\n\x0e\x43ommitRevision\x12-.model_express.revision.CommitRevisionRequest\x1a&.model_express.revision.RevisionRecordb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'revision_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_REVISIONSTATE']._serialized_start=745
+  _globals['_REVISIONSTATE']._serialized_end=848
+  _globals['_S3OBJECT']._serialized_start=42
+  _globals['_S3OBJECT']._serialized_end=147
+  _globals['_REVISIONMANIFEST']._serialized_start=150
+  _globals['_REVISIONMANIFEST']._serialized_end=393
+  _globals['_REVISIONRECORD']._serialized_start=396
+  _globals['_REVISIONRECORD']._serialized_end=526
+  _globals['_PUBLISHREVISIONREQUEST']._serialized_start=528
+  _globals['_PUBLISHREVISIONREQUEST']._serialized_end=612
+  _globals['_GETREVISIONREQUEST']._serialized_start=614
+  _globals['_GETREVISIONREQUEST']._serialized_end=676
+  _globals['_COMMITREVISIONREQUEST']._serialized_start=678
+  _globals['_COMMITREVISIONREQUEST']._serialized_end=743
+  _globals['_REVISIONCATALOGSERVICE']._serialized_start=851
+  _globals['_REVISIONCATALOGSERVICE']._serialized_end=1186
+# @@protoc_insertion_point(module_scope)