Add created IDs to async-upload job termination message (#2173)

* add entityIds to termination msg

Signed-off-by: Adysen Rothman <85646824+adysenrothman@users.noreply.github.com>

* add intent type to message

Signed-off-by: Adysen Rothman <85646824+adysenrothman@users.noreply.github.com>

* optional update_artifact intent ids

Signed-off-by: Adysen Rothman <85646824+adysenrothman@users.noreply.github.com>

* tests

Signed-off-by: Adysen Rothman <85646824+adysenrothman@users.noreply.github.com>

* update test

Signed-off-by: Adysen Rothman <85646824+adysenrothman@users.noreply.github.com>

---------

Signed-off-by: Adysen Rothman <85646824+adysenrothman@users.noreply.github.com>

Author: adysenrothman

jobs/async-upload/job/entrypoint.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 import logging
 import os
 
@@ -12,6 +13,7 @@
     create_version_and_artifact,
     validate_create_model_intent,
     validate_create_version_intent,
+    CreatedEntityIds,
 )
 from .models import CreateModelIntent, CreateVersionIntent, UpdateArtifactIntent
 from .download import perform_download
@@ -44,6 +46,29 @@ def record_error(exc):
     logger.error(message)
 
 
+def write_success_result(entity_ids: CreatedEntityIds, intent_type: str) -> None:
+    """Write success result with entity IDs to termination message path.
+    
+    Args:
+        entity_ids: CreatedEntityIds object containing the IDs of created/updated entities.
+        intent_type: The intent type string (e.g., "update_artifact", "create_model", "create_version").
+    """
+    result_dict = {"intent": intent_type}
+    
+    if entity_ids.registered_model_id:
+        result_dict["RegisteredModel"] = {"id": entity_ids.registered_model_id}
+    
+    if entity_ids.model_version_id:
+        result_dict["ModelVersion"] = {"id": entity_ids.model_version_id}
+    
+    if entity_ids.model_artifact_id:
+        result_dict["ModelArtifact"] = {"id": entity_ids.model_artifact_id}
+    
+    result_json = json.dumps(result_dict, indent=2)
+    write_to_termination_message_path(result_json)
+    logger.info(f"✅ Success result written to termination message path: {result_json}")
+
+
 async def main() -> None:
     """
     Main entrypoint for the async upload job.
@@ -59,12 +84,21 @@ async def main() -> None:
         client = validate_and_get_model_registry_client(config.registry)
 
         intent = config.model.intent
+        entity_ids: CreatedEntityIds | None = None
+        
         if isinstance(intent, UpdateArtifactIntent):
             logger.info("📋 Processing update_artifact intent")
             await set_artifact_pending(client, intent.artifact_id)
             perform_download(config)
             uri = perform_upload(config)
-            await update_model_artifact_uri(client, intent.artifact_id, uri)
+            # Pass through optional model_id and version_id if provided
+            entity_ids = await update_model_artifact_uri(
+                client,
+                intent.artifact_id,
+                uri,
+                registered_model_id=intent.model_id,
+                model_version_id=intent.version_id,
+            )
         elif isinstance(intent, CreateModelIntent):
             logger.info("📋 Processing create_model intent")
             if not config.metadata:
@@ -74,7 +108,7 @@ async def main() -> None:
             await validate_create_model_intent(client, config.metadata)
             perform_download(config)
             uri = perform_upload(config)
-            await create_model_and_artifact(client, config.metadata, uri)
+            entity_ids = await create_model_and_artifact(client, config.metadata, uri)
         elif isinstance(intent, CreateVersionIntent):
             logger.info("📋 Processing create_version intent")
             if not config.metadata:
@@ -84,9 +118,13 @@ async def main() -> None:
             await validate_create_version_intent(client, intent.model_id, config.metadata)
             perform_download(config)
             uri = perform_upload(config)
-            await create_version_and_artifact(client, intent.model_id, config.metadata, uri)
+            entity_ids = await create_version_and_artifact(client, intent.model_id, config.metadata, uri)
         else:
             raise ValueError(f"Unknown intent type: {type(intent)}")
+        
+        # Write success result to termination message path
+        if entity_ids:
+            write_success_result(entity_ids, intent.intent_type)
     except BaseException as e:
         record_error(e)
         raise

jobs/async-upload/job/models.py

@@ -123,6 +123,9 @@ class CreateVersionIntent(BaseModel):
 class UpdateArtifactIntent(BaseModel):
     intent_type: Literal[UploadIntent.update_artifact] = UploadIntent.update_artifact
     artifact_id: str = Field(..., description="Model artifact ID to update")
+    # Optional IDs to pass through to termination message output
+    model_id: str | None = Field(default=None, description="Optional registered model ID to include in output")
+    version_id: str | None = Field(default=None, description="Optional model version ID to include in output")
 
 
 IntentConfig = Union[CreateModelIntent, CreateVersionIntent, UpdateArtifactIntent]

jobs/async-upload/job/mr_client.py

@@ -1,4 +1,5 @@
 import logging
+from dataclasses import dataclass
 
 from model_registry import ModelRegistry
 from model_registry.types import ArtifactState
@@ -14,6 +15,14 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class CreatedEntityIds:
+    """IDs of created/updated entities in the model registry."""
+    registered_model_id: str | None = None
+    model_version_id: str | None = None
+    model_artifact_id: str | None = None
+
+
 def validate_and_get_model_registry_client(config: RegistryConfig) -> ModelRegistry:
     """
     Validates the model registry client configuration and returns a ModelRegistry client.
@@ -94,27 +103,64 @@ async def validate_create_version_intent(client: ModelRegistry, model_id: str, m
     logger.debug("✅ create_version intent validation passed")
 
 
-async def create_model_and_artifact(client: ModelRegistry, metadata: ConfigMapMetadata, uri: str) -> None:
-    """Creates a new registered model, model version, and model artifact."""
+async def create_model_and_artifact(client: ModelRegistry, metadata: ConfigMapMetadata, uri: str) -> CreatedEntityIds:
+    """Creates a new registered model, model version, and model artifact.
+    
+    Returns:
+        CreatedEntityIds: IDs of the created registered model, model version, and model artifact.
+    """
     logger.debug("🔍 Creating new registered model, version, and artifact")
     rm = await _create_registered_model(client, metadata.registered_model)
-    await _create_version_and_artifact_for_model(client, rm, uri, metadata)
+    mv, artifact = await _create_version_and_artifact_for_model(client, rm, uri, metadata)
+    return CreatedEntityIds(
+        registered_model_id=str(rm.id) if rm.id else None,
+        model_version_id=str(mv.id) if mv.id else None,
+        model_artifact_id=str(artifact.id) if artifact.id else None,
+    )
 
 
 async def create_version_and_artifact(
     client: ModelRegistry, model_id: str, metadata: ConfigMapMetadata, uri: str
-) -> None:
-    """Creates a new model version and model artifact under an existing registered model."""
+) -> CreatedEntityIds:
+    """Creates a new model version and model artifact under an existing registered model.
+    
+    Returns:
+        CreatedEntityIds: IDs of the existing registered model, and created model version and model artifact.
+    """
     logger.debug("🔍 Creating new version and artifact for model ID: %s", model_id)
 
     rm = await client._api.get_registered_model_by_id(model_id)
     if not rm:
         raise ValueError(f"RegisteredModel with ID '{model_id}' not found")
 
-    await _create_version_and_artifact_for_model(client, rm, uri, metadata)
+    mv, artifact = await _create_version_and_artifact_for_model(client, rm, uri, metadata)
+    return CreatedEntityIds(
+        registered_model_id=str(rm.id) if rm.id else None,
+        model_version_id=str(mv.id) if mv.id else None,
+        model_artifact_id=str(artifact.id) if artifact.id else None,
+    )
 
 
-async def update_model_artifact_uri(client: ModelRegistry, artifact_id: str, uri: str) -> None:
+async def update_model_artifact_uri(
+    client: ModelRegistry,
+    artifact_id: str,
+    uri: str,
+    registered_model_id: str | None = None,
+    model_version_id: str | None = None,
+) -> CreatedEntityIds:
+    """Updates the model artifact URI and sets state to LIVE.
+    
+    Args:
+        client: Model registry client.
+        artifact_id: ID of the artifact to update.
+        uri: New URI for the artifact.
+        registered_model_id: Optional registered model ID to pass through to output.
+        model_version_id: Optional model version ID to pass through to output.
+    
+    Returns:
+        CreatedEntityIds: IDs passed through to output. For update_artifact intent,
+        model and version IDs will be in the output if and only if they are passed in.
+    """
     logger.debug("🔍 Updating model artifact URI: %s", uri)
     artifact = await client._api.get_model_artifact_by_id(artifact_id)
 
@@ -126,6 +172,13 @@ async def update_model_artifact_uri(client: ModelRegistry, artifact_id: str, uri
     artifact.uri = uri
     await client._api.upsert_model_artifact(artifact)
     logger.debug("✅ Model artifact URI updated: %s", uri)
+    
+    # Pass through the IDs that were provided
+    return CreatedEntityIds(
+        registered_model_id=registered_model_id,
+        model_version_id=model_version_id,
+        model_artifact_id=artifact_id,
+    )
 
 
 async def _create_registered_model(client: ModelRegistry, rm_metadata):
@@ -142,8 +195,12 @@ async def _create_registered_model(client: ModelRegistry, rm_metadata):
 
 async def _create_version_and_artifact_for_model(
     client: ModelRegistry, rm, uri: str, metadata: ConfigMapMetadata
-) -> None:
-    """Creates a model version and artifact under the given registered model."""
+) -> tuple:
+    """Creates a model version and artifact under the given registered model.
+    
+    Returns:
+        tuple: (ModelVersion, ModelArtifact) - The created model version and artifact objects.
+    """
     mv_metadata = metadata.model_version
     version_name = mv_metadata.name or "1.0.0"
 
@@ -181,3 +238,5 @@ async def _create_version_and_artifact_for_model(
     artifact.state = ArtifactState.LIVE
     await client._api.upsert_model_artifact(artifact)
     logger.debug("✅ Updated ModelArtifact state to LIVE: %s", artifact.id)
+    
+    return (mv, artifact)