Supercharge Your ML Workflows

# ZenML with SageMaker integration
from zenml import pipeline, step
from zenml.integrations.aws.flavors import SagemakerOrchestratorFlavor
from zenml.integrations.aws.sagemaker_orchestrator_settings import SagemakerOrchestratorSettings

# Set up SageMaker orchestrator
sagemaker_orchestrator = SagemakerOrchestratorFlavor(
    execution_role="arn:aws:iam::123456789012:role/SageMakerRole",
    # Other configuration options...
)

# Define custom settings for specific steps
gpu_settings = SagemakerOrchestratorSettings(
    processor_args={
        "instance_type": "ml.p3.2xlarge",
        "volume_size_in_gb": 100
    },
    input_data_s3_uri="s3://your-bucket/training-data"
)

@step
def prepare_data():
    # Data preparation logic
    return processed_data

@step(settings={"orchestrator.sagemaker": gpu_settings})
def train_model(data):
    # Training logic using SageMaker
    return model

@step
def evaluate_model(model):
    # Model evaluation logic
    return metrics

@pipeline(orchestrator=sagemaker_orchestrator)
def sagemaker_pipeline():
    data = prepare_data()
    model = train_model(data)
    metrics = evaluate_model(model)

# Run the pipeline
sagemaker_pipeline()

import boto3
import sagemaker
from sagemaker.estimator import Estimator
from sagemaker.inputs import TrainingInput

# Set up SageMaker session
sagemaker_session = sagemaker.Session()
role = sagemaker.get_execution_role()

# Define estimator
estimator = Estimator(
    image_uri="your-docker-image-uri",
    role=role,
    instance_count=1,
    instance_type="ml.m5.xlarge",
    output_path="s3://your-bucket/output"
)

# Set hyperparameters
estimator.set_hyperparameters(epochs=10, learning_rate=0.1)

# Prepare data
train_data = TrainingInput(
    s3_data="s3://your-bucket/train",
    content_type="text/csv"
)

# Train the model
estimator.fit({"train": train_data})

# Deploy the model
predictor = estimator.deploy(
    initial_instance_count=1,
    instance_type="ml.t2.medium"
)

# Make predictions
result = predictor.predict("sample input data")
print(result)