Skip to content

Commit

Permalink
Merge branch 'main' into cherry-pick-ce-commit-38bcb06616369ecc46eb5f…
Browse files Browse the repository at this point in the history
…231ad49f48ee3c4456
  • Loading branch information
macintushar committed Sep 6, 2024
2 parents 68cfa76 + 3fba113 commit 20b0b50
Show file tree
Hide file tree
Showing 100 changed files with 2,667 additions and 914 deletions.
4 changes: 4 additions & 0 deletions integrations/Gemfile
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,10 @@ gem "aws-sdk-sts"

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

gem "aws-sdk-sagemaker"

gem "aws-sdk-sagemakerruntime"

group :development, :test do
gem "simplecov", require: false
gem "simplecov_json_formatter", require: false
Expand Down
15 changes: 14 additions & 1 deletion integrations/Gemfile.lock
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,14 @@ GIT
PATH
remote: .
specs:
multiwoven-integrations (0.9.1)
multiwoven-integrations (0.10.0)
activesupport
async-websocket
aws-sdk-athena
aws-sdk-cloudwatchlogs
aws-sdk-s3
aws-sdk-sts
aws-sigv4
csv
dry-schema
dry-struct
Expand Down Expand Up @@ -67,6 +69,9 @@ GEM
aws-sdk-athena (1.83.0)
aws-sdk-core (~> 3, >= 3.193.0)
aws-sigv4 (~> 1.1)
aws-sdk-cloudwatchlogs (1.82.0)
aws-sdk-core (~> 3, >= 3.193.0)
aws-sigv4 (~> 1.1)
aws-sdk-core (3.196.1)
aws-eventstream (~> 1, >= 1.3.0)
aws-partitions (~> 1, >= 1.651.0)
Expand All @@ -79,6 +84,12 @@ GEM
aws-sdk-core (~> 3, >= 3.194.0)
aws-sdk-kms (~> 1)
aws-sigv4 (~> 1.8)
aws-sdk-sagemaker (1.229.0)
aws-sdk-core (~> 3, >= 3.188.0)
aws-sigv4 (~> 1.1)
aws-sdk-sagemakerruntime (1.63.0)
aws-sdk-core (~> 3, >= 3.193.0)
aws-sigv4 (~> 1.1)
aws-sdk-sts (1.11.0)
aws-sdk-core (~> 3, >= 3.110.0)
aws-sigv4 (~> 1.1)
Expand Down Expand Up @@ -338,6 +349,8 @@ DEPENDENCIES
async-websocket (~> 0.8.0)
aws-sdk-athena
aws-sdk-s3
aws-sdk-sagemaker
aws-sdk-sagemakerruntime
aws-sdk-sts
byebug
csv
Expand Down
3 changes: 3 additions & 0 deletions integrations/lib/multiwoven/integrations.rb
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
require "iterable-api-client"
require "aws-sdk-sts"
require "ruby-oci8"
require "aws-sdk-sagemaker"
require "aws-sdk-sagemakerruntime"

# Service
require_relative "integrations/config"
Expand Down Expand Up @@ -63,6 +65,7 @@
require_relative "integrations/source/maria_db/client"
require_relative "integrations/source/oracle_db/client"
require_relative "integrations/source/databrics_model/client"
require_relative "integrations/source/aws_sagemaker_model/client"

# Destination
require_relative "integrations/destination/klaviyo/client"
Expand Down
3 changes: 0 additions & 3 deletions integrations/lib/multiwoven/integrations/core/constants.rb
Original file line number Diff line number Diff line change
Expand Up @@ -50,9 +50,6 @@ module Constants
DATABRICKS_HEALTH_URL = "https://%<databricks_host>s/api/2.0/serving-endpoints/%<endpoint_name>s"
DATABRICKS_SERVING_URL = "https://%<databricks_host>s/serving-endpoints/%<endpoint_name>s/invocations"

AWS_ACCESS_KEY_ID = ENV["AWS_ACCESS_KEY_ID"]
AWS_SECRET_ACCESS_KEY = ENV["AWS_SECRET_ACCESS_KEY"]

# HTTP
HTTP_GET = "GET"
HTTP_POST = "POST"
Expand Down
3 changes: 2 additions & 1 deletion integrations/lib/multiwoven/integrations/rollout.rb
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

module Multiwoven
module Integrations
VERSION = "0.9.1"
VERSION = "0.10.0"

ENABLED_SOURCES = %w[
Snowflake
Expand All @@ -17,6 +17,7 @@ module Integrations
MariaDB
Oracle
DatabricksModel
AwsSagemakerModel
].freeze

ENABLED_DESTINATIONS = %w[
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
# frozen_string_literal: true

module Multiwoven::Integrations::Source
module AwsSagemakerModel
include Multiwoven::Integrations::Core
class Client < SourceConnector
def check_connection(connection_config)
connection_config = connection_config.with_indifferent_access
create_connection(connection_config)
response = @client.describe_endpoint(endpoint_name: connection_config[:endpoint_name])
if response.endpoint_status == "InService"
success_status
else
failure_status
end
rescue StandardError => e
ConnectionStatus.new(status: ConnectionStatusType["failed"], message: e.message).to_multiwoven_message
end

def discover(_connection_config)
catalog_json = read_json(CATALOG_SPEC_PATH)
catalog = build_catalog(catalog_json)
catalog.to_multiwoven_message
rescue StandardError => e
handle_exception(e, {
context: "AWS:SAGEMAKER MODEL:DISCOVER:EXCEPTION",
type: "error"
})
end

def read(sync_config)
connection_config = sync_config.source.connection_specification
connection_config = connection_config.with_indifferent_access
payload = sync_config.model.query
create_connection(connection_config)
run_model(connection_config, payload)
rescue StandardError => e
handle_exception(e, {
context: "AWS:SAGEMAKER MODEL:READ:EXCEPTION",
type: "error",
sync_id: sync_config.sync_id,
sync_run_id: sync_config.sync_run_id
})
end

private

def create_connection(connection_config)
@client = Aws::SageMaker::Client.new(
region: connection_config[:region],
access_key_id: connection_config[:access_key],
secret_access_key: connection_config[:secret_access_key]
)

@client_runtime = Aws::SageMakerRuntime::Client.new(
region: connection_config[:region],
access_key_id: connection_config[:access_key],
secret_access_key: connection_config[:secret_access_key]
)
end

def run_model(connection_config, payload)
response = @client_runtime.invoke_endpoint(
endpoint_name: connection_config[:endpoint_name],
content_type: "application/json",
body: payload
)
process_response(response)
rescue StandardError => e
handle_exception(e, context: "AWS:SAGEMAKER MODEL:RUN_MODEL:EXCEPTION", type: "error")
end

def process_response(response)
data = JSON.parse(response.body.read)
[RecordMessage.new(data: { response: data }, emitted_at: Time.now.to_i).to_multiwoven_message]
end
end
end
end
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
{
"request_rate_limit": 600,
"request_rate_limit_unit": "minute",
"request_rate_concurrency": 10,
"streams": []
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
{
"data": {
"name": "AwsSagemakerModel",
"title": "AWS Sagemaker Model",
"connector_type": "source",
"category": "AI Model",
"documentation_url": "https://docs.mutliwoven.com",
"github_issue_label": "source-aws-sagemaker-model",
"icon": "icon.svg",
"license": "MIT",
"release_stage": "alpha",
"support_level": "community",
"tags": ["language:ruby", "multiwoven"]
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
{
"documentation_url": "https://docs.multiwoven.com/integrations/sources/aws_sagemaker-model",
"stream_type": "user_defined",
"connector_query_type": "ai_ml",
"connection_specification": {
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "AWS Sagemaker Model",
"type": "object",
"required": ["access_key", "secret_access_key", "region", "endpoint_name", "request_format", "response_format"],
"properties": {
"access_key": {
"description": "The AWS Access Key ID to use for authentication",
"type": "string",
"title": "Personal Access Key",
"order": 0
},
"secret_access_key": {
"description": "The AWS Secret Access Key to use for authentication",
"type": "string",
"multiwoven_secret": true,
"title": "Secret Access Key",
"order": 1
},
"region": {
"description": "AWS region",
"type": "string",
"title": "Region",
"order": 2
},
"endpoint_name": {
"description": "Endpoint name for AWS Sagemaker",
"type": "string",
"title": "Endpoint name",
"order": 3
},
"request_format": {
"description": "Sample Request Format",
"type": "string",
"title": "Request Format",
"order": 4
},
"response_format": {
"description": "Sample Response Format",
"type": "string",
"title": "Response Format",
"order": 5
}
}
}
}
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "Databricks Model",
"type": "object",
"required": ["databricks_host", "token", "endpoint"],
"required": ["databricks_host", "token", "endpoint", "request_format", "response_format"],
"properties": {
"databricks_host": {
"title": "databricks_host",
Expand Down
2 changes: 2 additions & 0 deletions integrations/multiwoven-integrations.gemspec
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,10 @@ Gem::Specification.new do |spec|
spec.add_runtime_dependency "activesupport"
spec.add_runtime_dependency "async-websocket"
spec.add_runtime_dependency "aws-sdk-athena"
spec.add_runtime_dependency "aws-sdk-cloudwatchlogs"
spec.add_runtime_dependency "aws-sdk-s3"
spec.add_runtime_dependency "aws-sdk-sts"
spec.add_runtime_dependency "aws-sigv4"
spec.add_runtime_dependency "csv"
spec.add_runtime_dependency "dry-schema"
spec.add_runtime_dependency "dry-struct"
Expand Down
Loading

0 comments on commit 20b0b50

Please sign in to comment.