feat: bedrock model reuse for OSS models

sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py

@@ -26,6 +26,8 @@
     build_source_tag,
     find_active_bedrock_deployment_for_model,
     find_existing_bedrock_model,
+    find_existing_imported_model,
+    find_existing_model_import_job,
 )
 from sagemaker.core.training.utils import (
     build_nova_manifest_s3_uri,
@@ -167,7 +169,7 @@ def _get_sagemaker_client(self):
             self._sagemaker_client = self.boto_session.client("sagemaker")
         return self._sagemaker_client
 
-    def _resolve_model_source_id(self) -> Optional[str]:
+    def _resolve_nova_model_source_id(self) -> Optional[str]:
         """Determine the model source identifier for reuse lookups.
 
         Resolution order:
@@ -349,7 +351,7 @@ def deploy(
                 sagemaker_session=self.sagemaker_session,
             )
 
-            source_id = self._resolve_model_source_id()
+            source_id = self._resolve_nova_model_source_id()
 
             if source_id and reuse_resources:
                 existing_arn = find_existing_bedrock_model(
@@ -435,23 +437,104 @@ def deploy(
                 role_type="bedrock",
                 sagemaker_session=self.sagemaker_session,
             )
-            model_data_source = {"s3DataSource": {"s3Uri": self.s3_model_artifacts}}
+
+            # Resolve model source identifier for reuse tagging.
+            # Priority: model package ARN > S3 artifact URI > None (with warning).
+            oss_source_id = None
+            if self.model_package:
+                mp_arn = getattr(self.model_package, "model_package_arn", None)
+                if mp_arn and isinstance(mp_arn, str):
+                    oss_source_id = mp_arn
+            if not oss_source_id and self.s3_model_artifacts and isinstance(self.s3_model_artifacts, str):
+                oss_source_id = self.s3_model_artifacts
+            if not oss_source_id:
+                logger.warning(
+                    "Cannot determine model source identifier for OSS model resource reuse. "
+                    "Neither Model package ARN nor model artifacts S3 URI is available. "
+                )
+
+            # Reuse: first look for an already-completed imported model, then
+            # fall back to an in-progress import job for the same source.
+            if oss_source_id and reuse_resources:
+                # 1. A completed imported model can be reused directly; there is
+                #    no import job to wait on.
+                model_arn = find_existing_imported_model(
+                    self._get_bedrock_client(),
+                    oss_source_id,
+                )
+                if model_arn:
+                    logger.info(
+                        "Reusing existing imported model %s (matched model-source tag). "
+                        "No new import job was created. Pass reuse_resources=False to "
+                        "force a new import.",
+                        model_arn,
+                    )
+                    model_details = self._get_bedrock_client().get_imported_model(
+                        modelIdentifier=model_arn
+                    )
+                    self._imported_model_id = model_details.get("modelName")
+                    return model_details
+
+                # 2. Otherwise, an import job may already be running for this
+                #    source; wait for it to complete instead of starting a new one.
+                job_arn = find_existing_model_import_job(
+                    self._get_bedrock_client(),
+                    oss_source_id,
+                )
+                if job_arn:
+                    logger.info(
+                        "Reusing in-progress import job %s (matched model-source tag). "
+                        "No new import job was created. Pass reuse_resources=False to "
+                        "force a new import.",
+                        job_arn,
+                    )
+                    self._wait_for_import_job_complete(job_arn)
+                    job_details = self._get_bedrock_client().get_model_import_job(
+                        jobIdentifier=job_arn
+                    )
+                    self._imported_model_id = job_details.get("importedModelName")
+                    return job_details
+
+
             # If artifacts are a tar.gz, extract to S3 first (Bedrock requires uncompressed format)
             if self.s3_model_artifacts.endswith(".tar.gz") or self.s3_model_artifacts.endswith(".tar.gz/"):
                 extracted_uri = self._extract_tar_gz_to_s3(self.s3_model_artifacts.rstrip("/"))
                 resolved_uri = self._resolve_hf_model_path(extracted_uri)
                 model_data_source = {"s3DataSource": {"s3Uri": resolved_uri}}
+            else:
+                resolved_uri = self._resolve_hf_model_path(self.s3_model_artifacts)
+                model_data_source = {"s3DataSource": {"s3Uri": resolved_uri}}
+
             # Auto-generate job_name if not provided
             if not job_name:
                 import time
                 job_name = f"{imported_model_name or 'import'}-{int(time.time())}"
+
+            # Inject the source tag into both the imported model tags and the
+            # import job tags. The model tags let a completed model be reused;
+            # the job tags let an in-progress import job be discovered and reused
+            # (reuse discovery matches the tag on the job ARN while the model
+            # does not yet exist).
+            merged_imported_tags = list(imported_model_tags) if imported_model_tags else []
+            merged_job_tags = list(job_tags) if job_tags else []
+            if oss_source_id:
+                source_tag = build_source_tag(oss_source_id)
+                merged_imported_tags = [
+                    t for t in merged_imported_tags if t.get("key") != source_tag["key"]
+                ]
+                merged_imported_tags.append(source_tag)
+                merged_job_tags = [
+                    t for t in merged_job_tags if t.get("key") != source_tag["key"]
+                ]
+                merged_job_tags.append(source_tag)
+
             params = {
                 "jobName": job_name,
                 "importedModelName": imported_model_name,
                 "roleArn": role_arn,
                 "modelDataSource": model_data_source,
-                "jobTags": job_tags,
-                "importedModelTags": imported_model_tags,
+                "jobTags": merged_job_tags if merged_job_tags else None,
+                "importedModelTags": merged_imported_tags if merged_imported_tags else None,
                 "clientRequestToken": client_request_token,
                 "importedModelKmsKeyId": imported_model_kms_key_id,
             }
@@ -902,6 +985,18 @@ def _resolve_hf_model_path(self, s3_uri: str) -> str:
 
         print(f"[BedrockModelBuilder] Base s3_uri from model package: {s3_uri}")
 
+        # Idempotency guard: if the given URI already points directly at a
+        # resolved model directory (contains config.json), it is already
+        # correct. Return it as-is instead of appending another checkpoints/
+        # prefix, so repeated calls are a no-op.
+        base_config_key = parsed_base.path.lstrip("/") + "config.json"
+        try:
+            s3_client.head_object(Bucket=bucket, Key=base_config_key)
+            logger.info("s3_uri already resolved (config.json present) at %s", s3_uri)
+            return s3_uri.rstrip("/")
+        except Exception as e:
+            logger.debug(f"[BedrockModelBuilder]{s3_uri} Not a resolved dir, continuing: {e}")
+
         hf_merged_uri = s3_uri + "checkpoints/hf_merged/"
         merged_config_key = urlparse(hf_merged_uri).path.lstrip("/") + "config.json"
         print(f"[BedrockModelBuilder] Probing for hf_merged: s3://{bucket}/{merged_config_key}")

sagemaker-serve/src/sagemaker/serve/model_reuse.py

@@ -153,6 +153,75 @@ def find_active_bedrock_deployment_for_model(bedrock_client, model_arn: str) ->
         return None
 
 
+def find_existing_imported_model(
+    bedrock_client,
+    source_id: str,
+) -> Optional[str]:
+    """Find an existing completed Bedrock imported model matching a source id.
+
+    Enumerates imported models (via ``list_imported_models``) and matches on the
+    ``sagemaker.amazonaws.com/model-source`` tag.
+
+    Args:
+        bedrock_client: A boto3 Bedrock client.
+        source_id: Raw source identifier (will be normalized).
+
+    Returns:
+        The imported-model ARN (``.../imported-model/...``) if a match is found,
+        None otherwise.
+    """
+    tag_value = normalize_tag_value(source_id)
+
+    try:
+        resource_arn = _find_imported_model_arn_by_tag(bedrock_client, tag_value)
+    except ClientError as e:
+        _reraise_if_access_denied(e, "bedrock:ListTagsForResource")
+        logger.warning("Could not list Bedrock imported models: %s. Proceeding without.", e)
+        return None
+    except Exception as e:
+        logger.warning("Could not list Bedrock imported models: %s. Proceeding without.", e)
+        return None
+
+    return resource_arn
+
+
+def find_existing_model_import_job(
+    bedrock_client,
+    source_id: str,
+) -> Optional[str]:
+    """Find an in-progress Bedrock model import job matching a source id.
+
+    Enumerates in-progress import jobs (via ``list_model_import_jobs``) and
+    matches on the ``sagemaker.amazonaws.com/model-source`` tag. Use this when
+    ``find_existing_imported_model`` returns None to detect an import that is
+    already running for the same source.
+
+    Args:
+        bedrock_client: A boto3 Bedrock client.
+        source_id: Raw source identifier (will be normalized).
+
+    Returns:
+        The import-job ARN (``.../model-import-job/...``) if a matching
+        in-progress job is found, None otherwise.
+    """
+    tag_value = normalize_tag_value(source_id)
+
+    try:
+        job_arn = _find_in_progress_import_job_by_tag(bedrock_client, tag_value)
+    except ClientError as e:
+        _reraise_if_access_denied(e, "bedrock:ListTagsForResource")
+        logger.warning("Could not list Bedrock import jobs: %s. Proceeding without.", e)
+        return None
+    except Exception as e:
+        logger.warning("Could not list Bedrock import jobs: %s. Proceeding without.", e)
+        return None
+
+    if job_arn:
+        logger.info("Found in-progress import job %s with matching model-source tag.", job_arn)
+
+    return job_arn
+
+
 def find_existing_sagemaker_endpoint(
     sagemaker_client,
     source_id: str,
@@ -211,6 +280,47 @@ def _find_bedrock_model_arn_by_tag(bedrock_client, tag_value: str) -> Optional[s
             return None
 
 
+def _find_imported_model_arn_by_tag(bedrock_client, tag_value: str) -> Optional[str]:
+    """Return the ARN of the first Bedrock imported model carrying the source tag."""
+    next_token = None
+    while True:
+        kwargs = {"nextToken": next_token} if next_token else {}
+        response = bedrock_client.list_imported_models(**kwargs)
+        for summary in response.get("modelSummaries", []):
+            arn = summary.get("modelArn")
+            if arn and _bedrock_resource_has_tag(bedrock_client, arn, tag_value):
+                return arn
+        next_token = response.get("nextToken")
+        if not next_token:
+            return None
+
+
+# The Bedrock ListModelImportJobs API only accepts the enum values
+# {Completed, InProgress, Failed} for statusEquals.
+_IMPORT_JOB_IN_PROGRESS_STATUSES = {"InProgress"}
+
+def _find_in_progress_import_job_by_tag(bedrock_client, tag_value: str) -> Optional[str]:
+    """Return the job ARN of an in-progress import job carrying the source tag.
+
+    Searches jobs in the InProgress state.
+    """
+    for status_filter in _IMPORT_JOB_IN_PROGRESS_STATUSES:
+        next_token = None
+        while True:
+            kwargs = {"statusEquals": status_filter}
+            if next_token:
+                kwargs["nextToken"] = next_token
+            response = bedrock_client.list_model_import_jobs(**kwargs)
+            for summary in response.get("modelImportJobSummaries", []):
+                job_arn = summary.get("jobArn")
+                if job_arn and _bedrock_resource_has_tag(bedrock_client, job_arn, tag_value):
+                    return job_arn
+            next_token = response.get("nextToken")
+            if not next_token:
+                break
+    return None
+
+
 def _bedrock_resource_has_tag(bedrock_client, resource_arn: str, tag_value: str) -> bool:
     """Return True if the Bedrock resource carries the source tag with tag_value."""
     tags = bedrock_client.list_tags_for_resource(resourceARN=resource_arn).get("tags", [])