Amazon Neptuneに接続するAmazon SagemakerノートブックインスタンスをCloudFormationで作成する

AWSTemplateFormatVersion: '2010-09-09'
Description: SageMaker Notebook Instance for Amazon Neptune connection

Parameters:
  VpcId:
    Type: AWS::EC2::VPC::Id
    Description: VPC ID where Neptune cluster is deployed

  SubnetId:
    Type: AWS::EC2::Subnet::Id
    Description: Subnet ID where the SageMaker Notebook will be deployed

  NeptuneClusterEndpoint:
    Type: String
    Description: Neptune cluster endpoint (e.g., your-cluster.cluster-xxxxxxxxx.region.neptune.amazonaws.com)

  NeptuneClusterResourceId:
    Type: String
    Description: Neptune cluster resource ID (e.g., cluster-xxxxxxxxx)

  InstanceType:
    Type: String
    Default: ml.t3.medium
    Description: SageMaker Notebook instance type

  NotebookInstanceName:
    Type: String
    Default: neptune-notebook
    Description: Name for the SageMaker Notebook instance

Resources:
  # SageMaker Notebook Instance用のIAMロール
  SageMakerExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: sagemaker.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: NeptuneNotebookPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: S3ListBucketAndGetObject
                Effect: Allow
                Action:
                  - s3:GetObject
                  - s3:ListBucket
                Resource:
                  - !Sub arn:aws:s3:::aws-neptune-notebook-${AWS::Region}
                  - !Sub arn:aws:s3:::aws-neptune-notebook-${AWS::Region}/*
              - Sid: DBAccess
                Effect: Allow
                Action: neptune-db:*
                Resource:
                  - !Sub arn:aws:neptune-db:${AWS::Region}:${AWS::AccountId}:${NeptuneClusterResourceId}/*
              - Sid: LogGroupAccess
                Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource:
                  - !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/sagemaker/*
              - Sid: SageMakerNotebookAccess
                Effect: Allow
                Action: sagemaker:DescribeNotebookInstance
                Resource:
                  - !Sub arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:notebook-instance/*

  # SageMaker Notebook Instance用のセキュリティグループ
  SageMakerSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Security group for SageMaker Notebook
      VpcId: !Ref VpcId
      SecurityGroupEgress:
        - IpProtocol: -1
          CidrIp: 0.0.0.0/0
      Tags:
        - Key: Name
          Value: !Sub ${NotebookInstanceName}-sg

  # ライフサイクル設定（Neptune接続用のセットアップ）
  NotebookLifecycleConfig:
    Type: AWS::SageMaker::NotebookInstanceLifecycleConfig
    Properties:
      NotebookInstanceLifecycleConfigName: !Sub ${NotebookInstanceName}-lifecycle
      OnStart:
        - Content:
            Fn::Base64: !Sub |
              #!/bin/bash

              sudo -u ec2-user -i <<'EOF'

              echo "export GRAPH_NOTEBOOK_AUTH_MODE=DEFAULT" >> ~/.bashrc
              echo "export GRAPH_NOTEBOOK_SSL=True" >> ~/.bashrc
              echo "export GRAPH_NOTEBOOK_SERVICE=neptune-db" >> ~/.bashrc
              echo "export GRAPH_NOTEBOOK_HOST=${NeptuneClusterEndpoint}" >> ~/.bashrc
              echo "export GRAPH_NOTEBOOK_PORT=8182" >> ~/.bashrc
              echo "export NEPTUNE_LOAD_FROM_S3_ROLE_ARN=" >> ~/.bashrc
              echo "export AWS_REGION=${AWS::Region}" >> ~/.bashrc

              aws s3 cp s3://aws-neptune-notebook-${AWS::Region}/graph_notebook.tar.gz /tmp/graph_notebook.tar.gz
              rm -rf /tmp/graph_notebook
              tar -zxvf /tmp/graph_notebook.tar.gz -C /tmp
              chmod +x /tmp/graph_notebook/install_jl4x.sh
              /tmp/graph_notebook/install_jl4x.sh

              EOF

  # SageMaker Notebook Instance
  SageMakerNotebookInstance:
    Type: AWS::SageMaker::NotebookInstance
    Properties:
      NotebookInstanceName: !Ref NotebookInstanceName
      InstanceType: !Ref InstanceType
      RoleArn: !GetAtt SageMakerExecutionRole.Arn
      SubnetId: !Ref SubnetId
      SecurityGroupIds:
        - !Ref SageMakerSecurityGroup
      DirectInternetAccess: Enabled
      LifecycleConfigName: !GetAtt NotebookLifecycleConfig.NotebookInstanceLifecycleConfigName
      Tags:
        - Key: Name
          Value: !Ref NotebookInstanceName

Outputs:
  NotebookInstanceArn:
    Description: ARN of the SageMaker Notebook instance
    Value: !Ref SageMakerNotebookInstance

  NotebookInstanceName:
    Description: Name of the SageMaker Notebook instance
    Value: !Ref NotebookInstanceName

  SageMakerSecurityGroupId:
    Description: Security Group ID of the SageMaker Notebook
    Value: !Ref SageMakerSecurityGroup

Fn::Base64: !Sub |
  #!/bin/bash

  sudo -u ec2-user -i <<'EOF'

  echo "export GRAPH_NOTEBOOK_AUTH_MODE=DEFAULT" >> ~/.bashrc
  echo "export GRAPH_NOTEBOOK_SSL=True" >> ~/.bashrc
  echo "export GRAPH_NOTEBOOK_SERVICE=neptune-db" >> ~/.bashrc
  echo "export GRAPH_NOTEBOOK_HOST=${NeptuneClusterEndpoint}" >> ~/.bashrc
  echo "export GRAPH_NOTEBOOK_PORT=8182" >> ~/.bashrc
  echo "export NEPTUNE_LOAD_FROM_S3_ROLE_ARN=" >> ~/.bashrc
  echo "export AWS_REGION=${AWS::Region}" >> ~/.bashrc

  aws s3 cp s3://aws-neptune-notebook-${AWS::Region}/graph_notebook.tar.gz /tmp/graph_notebook.tar.gz
  rm -rf /tmp/graph_notebook
  tar -zxvf /tmp/graph_notebook.tar.gz -C /tmp
  chmod +x /tmp/graph_notebook/install_jl4x.sh
  /tmp/graph_notebook/install_jl4x.sh

  EOF