@juspay/neurolink
Version:
Universal AI Development Platform with working MCP integration, multi-provider support, and professional CLI. Built-in tools operational, 58+ external MCP servers discoverable. Connect to filesystem, GitHub, database operations, and more. Build, test, and
463 lines (462 loc) • 18.3 kB
JavaScript
/**
* AWS SageMaker Runtime Client Wrapper
*
* This module provides a wrapper around the AWS SDK SageMaker Runtime client
* with enhanced error handling, retry logic, and NeuroLink-specific features.
*/
import { SageMakerRuntimeClient as AWSClient, InvokeEndpointCommand, InvokeEndpointWithResponseStreamCommand, } from "@aws-sdk/client-sagemaker-runtime";
import { handleSageMakerError, SageMakerError, isRetryableError, getRetryDelay, } from "./errors.js";
import { logger } from "../../utils/logger.js";
/**
* Enhanced SageMaker Runtime client with retry logic and error handling
*/
export class SageMakerRuntimeClient {
client;
config;
isDisposed = false;
constructor(config) {
this.config = config;
// Initialize AWS SDK client with configuration
this.client = new AWSClient({
region: config.region,
credentials: {
accessKeyId: config.accessKeyId,
secretAccessKey: config.secretAccessKey,
sessionToken: config.sessionToken,
},
maxAttempts: config.maxRetries || 3,
requestHandler: {
requestTimeout: config.timeout || 30000,
httpsAgent: {
// Keep connections alive for better performance
keepAlive: true,
maxSockets: 50,
},
},
...(config.endpoint && { endpoint: config.endpoint }),
});
logger.debug("SageMaker Runtime client initialized", {
region: config.region,
timeout: config.timeout,
maxRetries: config.maxRetries,
hasSessionToken: !!config.sessionToken,
customEndpoint: !!config.endpoint,
});
}
/**
* Invoke a SageMaker endpoint for synchronous inference
*
* @param params - Endpoint invocation parameters
* @returns Promise resolving to the inference response
* @throws {SageMakerError} When the request fails
*/
async invokeEndpoint(params) {
this.ensureNotDisposed();
const startTime = Date.now();
try {
logger.debug("Invoking SageMaker endpoint", {
endpointName: params.EndpointName,
contentType: params.ContentType,
bodySize: typeof params.Body === "string"
? params.Body.length
: params.Body?.length || 0,
});
// Prepare the command input
const input = {
EndpointName: params.EndpointName,
Body: params.Body,
ContentType: params.ContentType || "application/json",
Accept: params.Accept || "application/json",
CustomAttributes: params.CustomAttributes,
TargetModel: params.TargetModel,
TargetVariant: params.TargetVariant,
InferenceId: params.InferenceId,
};
const command = new InvokeEndpointCommand(input);
const client = this.client;
if (!client) {
throw new Error("SageMaker client has been disposed");
}
const response = (await this.executeWithRetry(() => client.send(command), params.EndpointName));
const duration = Date.now() - startTime;
logger.debug("SageMaker endpoint invocation successful", {
endpointName: params.EndpointName,
duration,
responseSize: response.Body?.length || 0,
invokedVariant: response.InvokedProductionVariant,
});
return {
Body: response.Body,
ContentType: response.ContentType,
InvokedProductionVariant: response.InvokedProductionVariant,
CustomAttributes: response.CustomAttributes,
};
}
catch (error) {
const duration = Date.now() - startTime;
logger.error("SageMaker endpoint invocation failed", {
endpointName: params.EndpointName,
duration,
error: error instanceof Error ? error.message : String(error),
});
throw handleSageMakerError(error, params.EndpointName);
}
}
/**
* Invoke a SageMaker endpoint with streaming response
*
* @param params - Endpoint invocation parameters for streaming
* @returns Promise resolving to async iterable of response chunks
* @throws {SageMakerError} When the request fails
*/
async invokeEndpointWithStreaming(params) {
this.ensureNotDisposed();
const startTime = Date.now();
try {
logger.debug("Starting SageMaker streaming invocation", {
endpointName: params.EndpointName,
contentType: params.ContentType,
bodySize: typeof params.Body === "string"
? params.Body.length
: params.Body?.length || 0,
});
// Prepare the command input for streaming
const input = {
EndpointName: params.EndpointName,
Body: params.Body,
ContentType: params.ContentType || "application/json",
Accept: params.Accept || "application/json",
CustomAttributes: params.CustomAttributes,
// Note: TargetModel, TargetVariant, InferenceId not available in streaming interface
};
const command = new InvokeEndpointWithResponseStreamCommand(input);
const client = this.client;
if (!client) {
throw new Error("SageMaker client has been disposed");
}
const response = (await this.executeWithRetry(() => client.send(command), params.EndpointName));
logger.debug("SageMaker streaming invocation started", {
endpointName: params.EndpointName,
setupDuration: Date.now() - startTime,
invokedVariant: response.InvokedProductionVariant,
});
// Return the response with streaming body
if (!response.Body) {
throw new SageMakerError("No response body received from streaming endpoint", "MODEL_ERROR", 500, undefined, params.EndpointName);
}
// Convert AWS response stream to async iterable of Uint8Array
const streamIterable = this.convertAWSStreamToIterable(response.Body);
return {
Body: streamIterable,
ContentType: response.ContentType,
InvokedProductionVariant: response.InvokedProductionVariant,
};
}
catch (error) {
const duration = Date.now() - startTime;
logger.error("SageMaker streaming invocation failed", {
endpointName: params.EndpointName,
duration,
error: error instanceof Error ? error.message : String(error),
});
throw handleSageMakerError(error, params.EndpointName);
}
}
/**
* Execute a request with automatic retry logic
*
* @param operation - Function that executes the AWS SDK command
* @param endpointName - Endpoint name for error context
* @param attempt - Current attempt number (for recursive retries)
* @returns Promise resolving to the operation result
*/
async executeWithRetry(operation, endpointName, attempt = 1) {
try {
return await operation();
}
catch (error) {
const maxRetries = this.config.maxRetries || 3;
// Check if we should retry
if (attempt < maxRetries && isRetryableError(error)) {
const delay = getRetryDelay(error, attempt);
logger.warn(`SageMaker request failed, retrying in ${delay}ms`, {
endpointName,
attempt,
maxRetries,
error: error instanceof Error ? error.message : String(error),
});
// Wait before retrying
await new Promise((resolve) => setTimeout(resolve, delay));
// Recursive retry
return this.executeWithRetry(operation, endpointName, attempt + 1);
}
// No more retries or not retryable
throw error;
}
}
/**
* Validate endpoint connectivity and permissions
*
* @param endpointName - Name of the endpoint to validate
* @returns Promise resolving to validation result
*/
async validateEndpoint(endpointName) {
this.ensureNotDisposed();
try {
// Try a minimal test request to validate endpoint
const testPayload = JSON.stringify({ test: true });
await this.invokeEndpoint({
EndpointName: endpointName,
Body: testPayload,
ContentType: "application/json",
Accept: "application/json",
});
return { isValid: true };
}
catch (error) {
if (error instanceof SageMakerError) {
return {
isValid: false,
error: error.message,
};
}
return {
isValid: false,
error: error instanceof Error ? error.message : "Unknown validation error",
};
}
}
/**
* Get client configuration summary for debugging
*
* @returns Configuration summary (with sensitive data masked)
*/
getConfigSummary() {
this.ensureNotDisposed();
return {
region: this.config.region,
timeout: this.config.timeout,
maxRetries: this.config.maxRetries,
hasCustomEndpoint: !!this.config.endpoint,
credentialsConfigured: !!(this.config.accessKeyId && this.config.secretAccessKey),
hasSessionToken: !!this.config.sessionToken,
};
}
/**
* Check if the client is properly configured
*
* @returns True if client appears to be properly configured
*/
isConfigured() {
this.ensureNotDisposed();
return !!(this.config.region &&
this.config.accessKeyId &&
this.config.secretAccessKey);
}
/**
* Convert AWS SDK async iterable stream with payload structure
*/
async *convertAsyncIterableStream(awsStream) {
for await (const chunk of awsStream) {
// Handle AWS SDK payload structure
if (chunk && typeof chunk === "object" && "PayloadPart" in chunk) {
const payloadChunk = chunk;
if (payloadChunk.PayloadPart?.Data) {
yield payloadChunk.PayloadPart.Data;
}
}
else if (chunk instanceof Uint8Array) {
yield chunk;
}
}
}
/**
* Convert Node.js readable stream with reader interface
*/
async *convertReadableStream(streamObj) {
const reader = typeof streamObj.getReader === "function"
? streamObj.getReader()
: undefined;
if (!reader) {
return; // No valid reader available
}
try {
while (true) {
const { done, value } = await reader.read();
if (done) {
break;
}
if (value instanceof Uint8Array) {
yield value;
}
else if (typeof value === "string") {
yield new TextEncoder().encode(value);
}
}
}
finally {
reader.releaseLock?.();
}
}
/**
* Convert non-stream data to single Uint8Array chunk as fallback
*/
*convertFallbackData(awsStream) {
logger.warn("Unsupported stream type, treating as single chunk", {
type: typeof awsStream,
isUint8Array: awsStream instanceof Uint8Array,
});
if (awsStream instanceof Uint8Array) {
yield awsStream;
}
else if (typeof awsStream === "string") {
yield new TextEncoder().encode(awsStream);
}
else {
yield new TextEncoder().encode(JSON.stringify(awsStream));
}
}
/**
* Convert AWS response stream to async iterable of Uint8Array chunks
* Refactored into smaller focused methods for different stream types
*/
async *convertAWSStreamToIterable(awsStream) {
try {
// AWS SDK streaming response handling
if (awsStream &&
typeof awsStream === "object" &&
Symbol.asyncIterator in awsStream) {
// Direct async iterable (AWS SDK event stream)
yield* this.convertAsyncIterableStream(awsStream);
}
else if (awsStream &&
typeof awsStream === "object" &&
"pipe" in awsStream) {
// Node.js stream conversion (readable stream with reader)
yield* this.convertReadableStream(awsStream);
}
else {
// Fallback: treat as single response (non-streaming data)
yield* this.convertFallbackData(awsStream);
}
}
catch (error) {
logger.error("Error converting AWS stream", {
error: error instanceof Error ? error.message : String(error),
streamType: typeof awsStream,
});
throw new SageMakerError(`Stream conversion failed: ${error instanceof Error ? error.message : String(error)}`, "NETWORK_ERROR", 500);
}
}
/**
* Check if the client has been disposed
*/
get disposed() {
return this.isDisposed;
}
/**
* Dispose of the client and clean up resources using explicit disposed state pattern
*
* AWS SDK v3 Automatic Resource Management:
* ========================================
*
* The AWS SDK v3 uses automatic resource cleanup and doesn't require explicit disposal
* of client instances in most cases. Here's how it works:
*
* 1. **HTTP Connection Pools**: AWS SDK v3 uses Node.js's built-in HTTP agent with
* connection pooling. These connections are automatically managed and will be
* closed when the Node.js process exits or becomes idle.
*
* 2. **Memory Management**: SDK clients don't hold significant resources that require
* manual cleanup. The JavaScript garbage collector handles memory deallocation
* when client references are removed.
*
* 3. **Background Timers**: Any internal timers (for retries, timeouts) are automatically
* cleared when operations complete or the client goes out of scope.
*
* 4. **Keep-Alive Connections**: HTTP keep-alive connections are managed by the
* underlying HTTP agent and will timeout automatically based on the configured
* keep-alive timeout (typically 15 seconds).
*
* Why We Still Implement dispose():
* =================================
*
* 1. **Explicit State Management**: Provides clear lifecycle control and prevents
* accidental usage of disposed clients.
*
* 2. **Resource Tracking**: Allows our application to track when clients are no
* longer needed, which is useful for debugging and monitoring.
*
* 3. **Defensive Programming**: Ensures we don't rely on automatic cleanup in
* environments where it might not work as expected.
*
* 4. **Future Compatibility**: If future SDK versions require explicit cleanup,
* we already have the infrastructure in place.
*
* For more information, see:
* - https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/
* - https://aws.amazon.com/blogs/developer/node-js-configuring-maxsockets-in-sdk-for-javascript/
*/
dispose() {
// Check for race condition - already disposed using explicit state
if (this.isDisposed) {
logger.debug("SageMaker Runtime client already disposed");
return;
}
// Mark as disposed first to prevent race conditions
this.isDisposed = true;
// Clear our client reference to enable garbage collection
// Note: AWS SDK v3 handles all internal resource cleanup automatically
this.client = null;
logger.debug("SageMaker Runtime client disposed", {
note: "AWS SDK v3 handles internal resource cleanup automatically",
});
}
/**
* Ensure client is not disposed before operations
*/
ensureNotDisposed() {
if (this.isDisposed) {
throw new SageMakerError("Cannot perform operation on disposed SageMaker client", "VALIDATION_ERROR", 400);
}
}
}
/**
* Factory function to create a SageMaker Runtime client
*
* @param config - SageMaker configuration
* @returns Configured SageMakerRuntimeClient instance
*/
export function createSageMakerRuntimeClient(config) {
return new SageMakerRuntimeClient(config);
}
/**
* Utility function to test SageMaker connectivity
*
* @param config - SageMaker configuration
* @param endpointName - Endpoint to test
* @returns Promise resolving to connectivity test result
*/
export async function testSageMakerConnectivity(config, endpointName) {
const client = new SageMakerRuntimeClient(config);
const startTime = Date.now();
try {
const result = await client.validateEndpoint(endpointName);
const latency = Date.now() - startTime;
if (result.isValid) {
return { connected: true, latency };
}
else {
return { connected: false, error: result.error };
}
}
catch (error) {
return {
connected: false,
error: error instanceof Error ? error.message : "Unknown connectivity error",
};
}
finally {
client.dispose();
}
}