deploy_sagemaker.py

import os
import boto3
import sagemaker
from sagemaker.pytorch import PyTorchModel
from sagemaker.predictor import Predictor
from sagemaker.serializers import JSONSerializer
from sagemaker.deserializers import JSONDeserializer
import json
import time
from datetime import datetime

class TransPolymerSageMakerDeployer:
    """
    Deploy TransPolymer model to Amazon SageMaker
    """
    
    def __init__(self, region_name='us-east-1', role_arn=None):
        """
        Initialize the deployer
        
        Args:
            region_name: AWS region name
            role_arn: IAM role ARN for SageMaker (if None, will try to get default)
        """
        self.region_name = region_name
        self.session = sagemaker.Session(boto_session=boto3.Session(region_name=region_name))
        
        # Get SageMaker execution role
        if role_arn:
            self.role = role_arn
        else:
            self.role = sagemaker.get_execution_role()
        
        self.bucket = self.session.default_bucket()
        print(f"Using S3 bucket: {self.bucket}")
        print(f"Using IAM role: {self.role}")
        
    def prepare_model_artifacts(self, model_dir="model_artifacts"):
        """
        Prepare model artifacts for deployment
        
        Args:
            model_dir: Directory to store model artifacts
        """
        print("Preparing model artifacts...")
        
        # Create model directory
        os.makedirs(model_dir, exist_ok=True)
        
        # Copy necessary files to model directory
        import shutil
        
        # Copy inference script
        shutil.copy("inference.py", os.path.join(model_dir, "code", "inference.py"))
        
        # Copy custom tokenizer
        shutil.copy("PolymerSmilesTokenization.py", os.path.join(model_dir, "code", "PolymerSmilesTokenization.py"))
        
        # Copy dataset module (needed for imports)
        if os.path.exists("dataset.py"):
            shutil.copy("dataset.py", os.path.join(model_dir, "code", "dataset.py"))
        
        # Copy model files if they exist
        if os.path.exists("ckpt/pretrain.pt/config.json"):
            shutil.copy("ckpt/pretrain.pt/config.json", os.path.join(model_dir, "config.json"))
        
        if os.path.exists("ckpt/pretrain.pt/pytorch_model.bin"):
            shutil.copy("ckpt/pretrain.pt/pytorch_model.bin", os.path.join(model_dir, "pytorch_model.bin"))
        
        # Copy requirements
        shutil.copy("requirements.txt", os.path.join(model_dir, "code", "requirements.txt"))
        
        # Create model.tar.gz
        import tarfile
        model_path = f"{model_dir}.tar.gz"
        
        with tarfile.open(model_path, "w:gz") as tar:
            tar.add(model_dir, arcname=".")
        
        print(f"Model artifacts prepared: {model_path}")
        return model_path
    
    def upload_model_to_s3(self, model_path, s3_key_prefix="transpolymer-model"):
        """
        Upload model artifacts to S3
        
        Args:
            model_path: Path to model.tar.gz file
            s3_key_prefix: S3 key prefix
        """
        print("Uploading model to S3...")
        
        timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        s3_key = f"{s3_key_prefix}/{timestamp}/model.tar.gz"
        
        # Upload to S3
        s3_uri = self.session.upload_data(
            path=model_path,
            bucket=self.bucket,
            key_prefix=f"{s3_key_prefix}/{timestamp}"
        )
        
        print(f"Model uploaded to: {s3_uri}")
        return s3_uri
    
    def create_model(self, model_s3_uri, model_name=None):
        """
        Create SageMaker model
        
        Args:
            model_s3_uri: S3 URI of model artifacts
            model_name: Name for the model (if None, will generate)
        """
        if model_name is None:
            timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
            model_name = f"transpolymer-model-{timestamp}"
        
        print(f"Creating SageMaker model: {model_name}")
        
        # Create PyTorch model
        model = PyTorchModel(
            model_data=model_s3_uri,
            role=self.role,
            entry_point="inference.py",
            framework_version="1.12.0",
            py_version="py38",
            name=model_name,
            sagemaker_session=self.session
        )
        
        self.model = model
        self.model_name = model_name
        print(f"Model created successfully: {model_name}")
        return model
    
    def deploy_endpoint(self, instance_type="ml.m5.large", initial_instance_count=1, 
                       endpoint_name=None, auto_scaling=True):
        """
        Deploy model to SageMaker endpoint
        
        Args:
            instance_type: EC2 instance type for endpoint
            initial_instance_count: Initial number of instances
            endpoint_name: Name for the endpoint (if None, will generate)
            auto_scaling: Whether to enable auto-scaling
        """
        if endpoint_name is None:
            timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
            endpoint_name = f"transpolymer-endpoint-{timestamp}"
        
        print(f"Deploying endpoint: {endpoint_name}")
        print(f"Instance type: {instance_type}")
        print(f"Initial instance count: {initial_instance_count}")
        
        # Deploy the model
        predictor = self.model.deploy(
            initial_instance_count=initial_instance_count,
            instance_type=instance_type,
            endpoint_name=endpoint_name,
            serializer=JSONSerializer(),
            deserializer=JSONDeserializer()
        )
        
        self.predictor = predictor
        self.endpoint_name = endpoint_name
        
        # Configure auto-scaling if requested
        if auto_scaling:
            self._configure_auto_scaling()
        
        print(f"Endpoint deployed successfully: {endpoint_name}")
        return predictor
    
    def _configure_auto_scaling(self, min_capacity=1, max_capacity=5, target_value=70.0):
        """
        Configure auto-scaling for the endpoint
        """
        print("Configuring auto-scaling...")
        
        autoscaling_client = boto3.client('application-autoscaling', region_name=self.region_name)
        
        # Register scalable target
        resource_id = f"endpoint/{self.endpoint_name}/variant/AllTraffic"
        
        try:
            autoscaling_client.register_scalable_target(
                ServiceNamespace='sagemaker',
                ResourceId=resource_id,
                ScalableDimension='sagemaker:variant:DesiredInstanceCount',
                MinCapacity=min_capacity,
                MaxCapacity=max_capacity
            )
            
            # Create scaling policy
            autoscaling_client.put_scaling_policy(
                PolicyName=f'{self.endpoint_name}-scaling-policy',
                ServiceNamespace='sagemaker',
                ResourceId=resource_id,
                ScalableDimension='sagemaker:variant:DesiredInstanceCount',
                PolicyType='TargetTrackingScaling',
                TargetTrackingScalingPolicyConfiguration={
                    'TargetValue': target_value,
                    'PredefinedMetricSpecification': {
                        'PredefinedMetricType': 'SageMakerVariantInvocationsPerInstance'
                    },
                    'ScaleOutCooldown': 300,
                    'ScaleInCooldown': 300
                }
            )
            
            print("Auto-scaling configured successfully")
            
        except Exception as e:
            print(f"Warning: Could not configure auto-scaling: {e}")
    
    def test_endpoint(self, test_data=None):
        """
        Test the deployed endpoint
        
        Args:
            test_data: Test data for prediction (if None, will use default)
        """
        if test_data is None:
            # Default test data - polymer SMILES examples
            test_data = [
                {
                    "smiles": "CC(C)(C)OC(=O)NC1=CC=CC=C1",
                    "property": "conductivity",
                    "model_type": "PE_I"
                },
                {
                    "smiles": "c1ccc2c(c1)oc1ccccc12",
                    "property": "band_gap",
                    "model_type": "Egc"
                }
            ]
        
        print("Testing endpoint...")
        print(f"Test data: {test_data}")
        
        try:
            # Make prediction
            result = self.predictor.predict(test_data)
            print("Prediction successful!")
            print(f"Results: {json.dumps(result, indent=2)}")
            return result
            
        except Exception as e:
            print(f"Error during prediction: {e}")
            raise
    
    def cleanup_endpoint(self):
        """
        Delete the endpoint to avoid charges
        """
        if hasattr(self, 'predictor'):
            print(f"Deleting endpoint: {self.endpoint_name}")
            self.predictor.delete_endpoint()
            print("Endpoint deleted successfully")
    
    def get_endpoint_info(self):
        """
        Get information about the deployed endpoint
        """
        if hasattr(self, 'endpoint_name'):
            return {
                "endpoint_name": self.endpoint_name,
                "model_name": self.model_name,
                "region": self.region_name,
                "endpoint_url": f"https://runtime.sagemaker.{self.region_name}.amazonaws.com/endpoints/{self.endpoint_name}/invocations"
            }
        else:
            return None

def main():
    """
    Main deployment function
    """
    print("🚀 Starting TransPolymer SageMaker Deployment")
    
    # Initialize deployer
    deployer = TransPolymerSageMakerDeployer()
    
    try:
        # Step 1: Prepare model artifacts
        model_path = deployer.prepare_model_artifacts()
        
        # Step 2: Upload to S3
        s3_uri = deployer.upload_model_to_s3(model_path)
        
        # Step 3: Create SageMaker model
        model = deployer.create_model(s3_uri)
        
        # Step 4: Deploy endpoint
        predictor = deployer.deploy_endpoint(
            instance_type="ml.m5.large",  # Change to ml.g4dn.xlarge for GPU
            initial_instance_count=1
        )
        
        # Step 5: Test endpoint
        deployer.test_endpoint()
        
        # Step 6: Print endpoint information
        endpoint_info = deployer.get_endpoint_info()
        print("\n✅ Deployment completed successfully!")
        print(f"Endpoint info: {json.dumps(endpoint_info, indent=2)}")
        
        print("\n📋 Usage example:")
        print(f"""
import boto3
import json

runtime = boto3.client('sagemaker-runtime', region_name='{deployer.region_name}')

payload = {{
    "smiles": "CC(C)(C)OC(=O)NC1=CC=CC=C1",
    "property": "conductivity",
    "model_type": "PE_I"
}}

response = runtime.invoke_endpoint(
    EndpointName='{deployer.endpoint_name}',
    ContentType='application/json',
    Body=json.dumps(payload)
)

result = json.loads(response['Body'].read().decode())
print(result)
        """)
        
        # Ask user if they want to keep the endpoint running
        keep_running = input("\nDo you want to keep the endpoint running? (y/n): ")
        if keep_running.lower() != 'y':
            deployer.cleanup_endpoint()
        
    except Exception as e:
        print(f"❌ Deployment failed: {e}")
        # Cleanup on failure
        if hasattr(deployer, 'predictor'):
            deployer.cleanup_endpoint()
        raise

if __name__ == "__main__":
    main()