Setting the file. One moment. 10
Setup DevOps Agent
24.7
Code Output Guide · references
RDS Oracle
(opens in a new tab)
references/model-deployment/code_templates/deploy-oss-bedrock.py
Python·130 lines·4 KB
from
sagemaker.serve.bedrock_model_builder
import
BedrockModelBuilder
16
17set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)
18
19REGION = "[REGION]"
20TRAINING_JOB_NAME = "[TRAINING_JOB_NAME]"
21ROLE_ARN = "[ROLE_ARN]"
22MODEL_NAME = "[MODEL_NAME]"
23
24sm = boto3.client("sagemaker", region_name=REGION)
25s3 = boto3.client("s3", region_name=REGION)
26
27# Cell 3: Flatten S3 Structure and Start Import
28
29# BedrockModelBuilder passes the root model artifacts path to Bedrock CMI,
30# but Bedrock expects config.json at the root of the URI. This cell copies
31# files from checkpoints/hf_merged/ to the model artifacts root (server-side).
32
33tj = sm.describe_training_job(TrainingJobName=TRAINING_JOB_NAME)
34root = tj["ModelArtifacts"]["S3ModelArtifacts"]
35parts = root.replace("s3://", "").split("/", 1)
36bucket, root_prefix = parts[0], parts[1].rstrip("/") + "/"
37hf_prefix = root_prefix + "checkpoints/hf_merged/"
38
39resp = s3.list_objects_v2(Bucket=bucket, Prefix=root_prefix + "config.json", MaxKeys=1)
40if resp.get("KeyCount", 0) > 0:
41 print("Files already at root, skipping copy")
42else:
43 paginator = s3.get_paginator("list_objects_v2")
44 copied = 0
45 for page in paginator.paginate(Bucket=bucket, Prefix=hf_prefix):
46 for obj in page.get("Contents", []):
47 filename = obj["Key"].replace(hf_prefix, "")
48 if not filename or filename.endswith("/"):
49 continue
50 s3.copy_object(
51 Bucket=bucket,
52 CopySource={"Bucket": bucket, "Key": obj["Key"]},
53 Key=root_prefix + filename,
54 )
55 copied += 1
56 print(f"Copied {copied} files to root")
57
58training_job = TrainingJob.get(training_job_name=TRAINING_JOB_NAME, region=REGION)
59builder = BedrockModelBuilder(model=training_job)
60
61result = builder.deploy(
62 job_name=MODEL_NAME,
63 imported_model_name=MODEL_NAME,
64 role_arn=ROLE_ARN,
65)
66
67job_arn = result["jobArn"]
68print(f"Import job created: {job_arn}")
69
70# Cell 4: Wait for Import to Complete
71
72bedrock = boto3.client("bedrock", region_name=REGION)
73
74while True:
75 resp = bedrock.get_model_import_job(jobIdentifier=job_arn)
76 status = resp["status"]
77 print(f"Status: {status}")
78
79 if status == "Completed":
80 model_arn = resp["importedModelArn"]
81 print(f"\nModel imported successfully!")
82 print(f"Model ARN: {model_arn}")
83 break
84 elif status in ("Failed", "Stopped"):
85 raise RuntimeError(f"Import {status}: {resp.get('failureMessage', 'Unknown error')}")
86
87 time.sleep(30)
88
89# Cell 5: Test Inference
90
91print("Testing inference (model may need a few minutes to warm up)...")
92bedrock_runtime = boto3.client("bedrock-runtime", region_name=REGION)
93
94for attempt in range(1, 25):
95 try:
96 response = bedrock_runtime.invoke_model(
97 modelId=model_arn,
98 body=json.dumps(
99 {
100 "prompt": "What is the capital of France?",
101 "max_gen_len": 50,
102 "temperature": 0.7,
103 }
104 ),
105 )
106 result = json.loads(response["body"].read())
107 print(f"Response: {json.dumps(result)[:300]}")
108 break
109 except bedrock_runtime.exceptions.ModelNotReadyException:
110 print(f" Attempt {attempt}: Model not ready, waiting 30s...")
111 time.sleep(30)
112else:
113 print("Model did not become ready after 12 minutes.")
114
115# Cell 6: Save Manifest
116# Save manifest - record output of workflow step for future reference
117from pathlib import Path
118
119manifest_dir = Path("[PROJECT_DIR]") / "manifests"
120manifest_dir.mkdir(parents=True, exist_ok=True)
121manifest_path = manifest_dir / f"deploy-{TRAINING_JOB_NAME}.json"
122manifest_path.write_text(
123 json.dumps(
124 {
125 "model_id": model_arn,
126 },
127 indent=2,
128 )
129)
130print(f"Manifest saved: {manifest_path}")