Skip to content

Commit b53ec87

Browse files
github-actions[bot]TivonB-AI2
github-actions[bot]
and
TivonB-AI2
authored
feat(CE): add aws sagemaker model source connector (#352)
Co-authored-by: TivonB-AI2 <[email protected]>
1 parent 4d5d740 commit b53ec87

File tree

12 files changed

+339
-5
lines changed

12 files changed

+339
-5
lines changed

‎integrations/Gemfile

+4
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,10 @@ gem "aws-sdk-sts"
6969

7070
gem "ruby-oci8", "~> 2.2.12"
7171

72+
gem "aws-sdk-sagemaker"
73+
74+
gem "aws-sdk-sagemakerruntime"
75+
7276
group :development, :test do
7377
gem "simplecov", require: false
7478
gem "simplecov_json_formatter", require: false

‎integrations/Gemfile.lock

+14-1
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,14 @@ GIT
77
PATH
88
remote: .
99
specs:
10-
multiwoven-integrations (0.9.2)
10+
multiwoven-integrations (0.10.0)
1111
activesupport
1212
async-websocket
1313
aws-sdk-athena
14+
aws-sdk-cloudwatchlogs
1415
aws-sdk-s3
1516
aws-sdk-sts
17+
aws-sigv4
1618
csv
1719
dry-schema
1820
dry-struct
@@ -67,6 +69,9 @@ GEM
6769
aws-sdk-athena (1.83.0)
6870
aws-sdk-core (~> 3, >= 3.193.0)
6971
aws-sigv4 (~> 1.1)
72+
aws-sdk-cloudwatchlogs (1.82.0)
73+
aws-sdk-core (~> 3, >= 3.193.0)
74+
aws-sigv4 (~> 1.1)
7075
aws-sdk-core (3.196.1)
7176
aws-eventstream (~> 1, >= 1.3.0)
7277
aws-partitions (~> 1, >= 1.651.0)
@@ -79,6 +84,12 @@ GEM
7984
aws-sdk-core (~> 3, >= 3.194.0)
8085
aws-sdk-kms (~> 1)
8186
aws-sigv4 (~> 1.8)
87+
aws-sdk-sagemaker (1.229.0)
88+
aws-sdk-core (~> 3, >= 3.188.0)
89+
aws-sigv4 (~> 1.1)
90+
aws-sdk-sagemakerruntime (1.63.0)
91+
aws-sdk-core (~> 3, >= 3.193.0)
92+
aws-sigv4 (~> 1.1)
8293
aws-sdk-sts (1.11.0)
8394
aws-sdk-core (~> 3, >= 3.110.0)
8495
aws-sigv4 (~> 1.1)
@@ -338,6 +349,8 @@ DEPENDENCIES
338349
async-websocket (~> 0.8.0)
339350
aws-sdk-athena
340351
aws-sdk-s3
352+
aws-sdk-sagemaker
353+
aws-sdk-sagemakerruntime
341354
aws-sdk-sts
342355
byebug
343356
csv

‎integrations/lib/multiwoven/integrations.rb

+3
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
3232
require "iterable-api-client"
3333
require "aws-sdk-sts"
3434
require "ruby-oci8"
35+
require "aws-sdk-sagemaker"
36+
require "aws-sdk-sagemakerruntime"
3537

3638
# Service
3739
require_relative "integrations/config"
@@ -63,6 +65,7 @@
6365
require_relative "integrations/source/maria_db/client"
6466
require_relative "integrations/source/oracle_db/client"
6567
require_relative "integrations/source/databrics_model/client"
68+
require_relative "integrations/source/aws_sagemaker_model/client"
6669

6770
# Destination
6871
require_relative "integrations/destination/klaviyo/client"

‎integrations/lib/multiwoven/integrations/core/constants.rb

-3
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,6 @@ module Constants
5050
DATABRICKS_HEALTH_URL = "https://%<databricks_host>s/api/2.0/serving-endpoints/%<endpoint_name>s"
5151
DATABRICKS_SERVING_URL = "https://%<databricks_host>s/serving-endpoints/%<endpoint_name>s/invocations"
5252

53-
AWS_ACCESS_KEY_ID = ENV["AWS_ACCESS_KEY_ID"]
54-
AWS_SECRET_ACCESS_KEY = ENV["AWS_SECRET_ACCESS_KEY"]
55-
5653
# HTTP
5754
HTTP_GET = "GET"
5855
HTTP_POST = "POST"

‎integrations/lib/multiwoven/integrations/rollout.rb

+2-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
module Multiwoven
44
module Integrations
5-
VERSION = "0.9.2"
5+
VERSION = "0.10.0"
66

77
ENABLED_SOURCES = %w[
88
Snowflake
@@ -17,6 +17,7 @@ module Integrations
1717
MariaDB
1818
Oracle
1919
DatabricksModel
20+
AwsSagemakerModel
2021
].freeze
2122

2223
ENABLED_DESTINATIONS = %w[

‎integrations/lib/multiwoven/integrations/source/aws_sagemaker_model/client.rb

+79
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
# frozen_string_literal: true
2+
3+
module Multiwoven::Integrations::Source
4+
module AwsSagemakerModel
5+
include Multiwoven::Integrations::Core
6+
class Client < SourceConnector
7+
def check_connection(connection_config)
8+
connection_config = connection_config.with_indifferent_access
9+
create_connection(connection_config)
10+
response = @client.describe_endpoint(endpoint_name: connection_config[:endpoint_name])
11+
if response.endpoint_status == "InService"
12+
success_status
13+
else
14+
failure_status
15+
end
16+
rescue StandardError => e
17+
ConnectionStatus.new(status: ConnectionStatusType["failed"], message: e.message).to_multiwoven_message
18+
end
19+
20+
def discover(_connection_config)
21+
catalog_json = read_json(CATALOG_SPEC_PATH)
22+
catalog = build_catalog(catalog_json)
23+
catalog.to_multiwoven_message
24+
rescue StandardError => e
25+
handle_exception(e, {
26+
context: "AWS:SAGEMAKER MODEL:DISCOVER:EXCEPTION",
27+
type: "error"
28+
})
29+
end
30+
31+
def read(sync_config)
32+
connection_config = sync_config.source.connection_specification
33+
connection_config = connection_config.with_indifferent_access
34+
payload = sync_config.model.query
35+
create_connection(connection_config)
36+
run_model(connection_config, payload)
37+
rescue StandardError => e
38+
handle_exception(e, {
39+
context: "AWS:SAGEMAKER MODEL:READ:EXCEPTION",
40+
type: "error",
41+
sync_id: sync_config.sync_id,
42+
sync_run_id: sync_config.sync_run_id
43+
})
44+
end
45+
46+
private
47+
48+
def create_connection(connection_config)
49+
@client = Aws::SageMaker::Client.new(
50+
region: connection_config[:region],
51+
access_key_id: connection_config[:access_key],
52+
secret_access_key: connection_config[:secret_access_key]
53+
)
54+
55+
@client_runtime = Aws::SageMakerRuntime::Client.new(
56+
region: connection_config[:region],
57+
access_key_id: connection_config[:access_key],
58+
secret_access_key: connection_config[:secret_access_key]
59+
)
60+
end
61+
62+
def run_model(connection_config, payload)
63+
response = @client_runtime.invoke_endpoint(
64+
endpoint_name: connection_config[:endpoint_name],
65+
content_type: "application/json",
66+
body: payload
67+
)
68+
process_response(response)
69+
rescue StandardError => e
70+
handle_exception(e, context: "AWS:SAGEMAKER MODEL:RUN_MODEL:EXCEPTION", type: "error")
71+
end
72+
73+
def process_response(response)
74+
data = JSON.parse(response.body.read)
75+
[RecordMessage.new(data: { response: data }, emitted_at: Time.now.to_i).to_multiwoven_message]
76+
end
77+
end
78+
end
79+
end

‎integrations/lib/multiwoven/integrations/source/aws_sagemaker_model/config/catalog.json

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"request_rate_limit": 600,
3+
"request_rate_limit_unit": "minute",
4+
"request_rate_concurrency": 10,
5+
"streams": []
6+
}

‎integrations/lib/multiwoven/integrations/source/aws_sagemaker_model/config/meta.json

+15
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
{
2+
"data": {
3+
"name": "AwsSagemakerModel",
4+
"title": "AWS Sagemaker Model",
5+
"connector_type": "source",
6+
"category": "AI Model",
7+
"documentation_url": "https://docs.mutliwoven.com",
8+
"github_issue_label": "source-aws-sagemaker-model",
9+
"icon": "icon.svg",
10+
"license": "MIT",
11+
"release_stage": "alpha",
12+
"support_level": "community",
13+
"tags": ["language:ruby", "multiwoven"]
14+
}
15+
}

‎integrations/lib/multiwoven/integrations/source/aws_sagemaker_model/config/spec.json

+50
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
{
2+
"documentation_url": "https://docs.multiwoven.com/integrations/sources/aws_sagemaker-model",
3+
"stream_type": "user_defined",
4+
"connector_query_type": "ai_ml",
5+
"connection_specification": {
6+
"$schema": "http://json-schema.org/draft-07/schema#",
7+
"title": "AWS Sagemaker Model",
8+
"type": "object",
9+
"required": ["access_key", "secret_access_key", "region", "endpoint_name", "request_format", "response_format"],
10+
"properties": {
11+
"access_key": {
12+
"description": "The AWS Access Key ID to use for authentication",
13+
"type": "string",
14+
"title": "Personal Access Key",
15+
"order": 0
16+
},
17+
"secret_access_key": {
18+
"description": "The AWS Secret Access Key to use for authentication",
19+
"type": "string",
20+
"multiwoven_secret": true,
21+
"title": "Secret Access Key",
22+
"order": 1
23+
},
24+
"region": {
25+
"description": "AWS region",
26+
"type": "string",
27+
"title": "Region",
28+
"order": 2
29+
},
30+
"endpoint_name": {
31+
"description": "Endpoint name for AWS Sagemaker",
32+
"type": "string",
33+
"title": "Endpoint name",
34+
"order": 3
35+
},
36+
"request_format": {
37+
"description": "Sample Request Format",
38+
"type": "string",
39+
"title": "Request Format",
40+
"order": 4
41+
},
42+
"response_format": {
43+
"description": "Sample Response Format",
44+
"type": "string",
45+
"title": "Response Format",
46+
"order": 5
47+
}
48+
}
49+
}
50+
}

‎integrations/lib/multiwoven/integrations/source/aws_sagemaker_model/icon.svg

+18
Loading

‎integrations/multiwoven-integrations.gemspec

+2
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,10 @@ Gem::Specification.new do |spec|
3636
spec.add_runtime_dependency "activesupport"
3737
spec.add_runtime_dependency "async-websocket"
3838
spec.add_runtime_dependency "aws-sdk-athena"
39+
spec.add_runtime_dependency "aws-sdk-cloudwatchlogs"
3940
spec.add_runtime_dependency "aws-sdk-s3"
4041
spec.add_runtime_dependency "aws-sdk-sts"
42+
spec.add_runtime_dependency "aws-sigv4"
4143
spec.add_runtime_dependency "csv"
4244
spec.add_runtime_dependency "dry-schema"
4345
spec.add_runtime_dependency "dry-struct"

0 commit comments

Comments
 (0)