Amazon Neptune Ml Job Structure

{
  "$schema": "https://json-structure.org/meta/core/v0/#",
  "$id": "https://github.com/api-search/amazon-neptune/json-structure/amazon-neptune-ml-job-structure.json",
  "name": "Amazon Neptune ML Job",
  "description": "Represents Neptune ML job configurations and status for data processing, model training, model transform, and inference endpoint operations powered by Amazon SageMaker.",
  "oneOf": [
    {
      "$ref": "#/$defs/DataProcessingJob"
    },
    {
      "$ref": "#/$defs/ModelTrainingJob"
    },
    {
      "$ref": "#/$defs/ModelTransformJob"
    },
    {
      "$ref": "#/$defs/InferenceEndpoint"
    }
  ],
  "definitions": {
    "DataProcessingJob": {
      "type": "object",
      "name": "DataProcessingJob",
      "description": "A Neptune ML data processing job that exports and prepares graph data for model training.",
      "required": [
        "inputDataS3Location",
        "processedDataS3Location"
      ],
      "properties": {
        "id": {
          "type": "string",
          "description": "Unique identifier for the job."
        },
        "inputDataS3Location": {
          "type": "string",
          "description": "S3 URI for the input graph data."
        },
        "processedDataS3Location": {
          "type": "string",
          "description": "S3 URI where processed output is written."
        },
        "previousDataProcessingJobId": {
          "type": "string",
          "description": "Job ID of a previous job for incremental processing."
        },
        "sagemakerIamRoleArn": {
          "type": "string",
          "description": "IAM role ARN for SageMaker execution."
        },
        "neptuneIamRoleArn": {
          "type": "string",
          "description": "IAM role ARN for Neptune access."
        },
        "processingInstanceType": {
          "type": "string",
          "description": "The SageMaker ML instance type for processing.",
          "examples": [
            "ml.r5.xlarge",
            "ml.r5.2xlarge",
            "ml.r5.4xlarge"
          ]
        },
        "processingInstanceVolumeSizeInGB": {
          "type": "int32",
          "description": "Disk volume size in GB.",
          "default": 0
        },
        "processingTimeOutInSeconds": {
          "type": "int32",
          "description": "Processing timeout in seconds.",
          "default": 86400
        },
        "modelType": {
          "type": "string",
          "description": "The model type to prepare data for.",
          "enum": [
            "heterogeneous",
            "kge"
          ]
        },
        "configFileName": {
          "type": "string",
          "description": "The data specification configuration file name.",
          "default": "training-data-configuration.json"
        },
        "subnets": {
          "type": "array",
          "items": {
            "type": "string"
          }
        },
        "securityGroupIds": {
          "type": "array",
          "items": {
            "type": "string"
          }
        },
        "volumeEncryptionKMSKey": {
          "type": "string"
        },
        "s3OutputEncryptionKMSKey": {
          "type": "string"
        },
        "enableInterContainerTrafficEncryption": {
          "type": "boolean",
          "default": true
        }
      }
    },
    "ModelTrainingJob": {
      "type": "object",
      "name": "ModelTrainingJob",
      "description": "A Neptune ML model training job that trains a graph neural network model using Amazon SageMaker.",
      "required": [
        "dataProcessingJobId",
        "trainModelS3Location"
      ],
      "properties": {
        "id": {
          "type": "string"
        },
        "dataProcessingJobId": {
          "type": "string",
          "description": "Job ID of the completed data processing job."
        },
        "trainModelS3Location": {
          "type": "string",
          "description": "S3 location for model artifacts."
        },
        "previousModelTrainingJobId": {
          "type": "string",
          "description": "Job ID for incremental training."
        },
        "sagemakerIamRoleArn": {
          "type": "string"
        },
        "neptuneIamRoleArn": {
          "type": "string"
        },
        "modelName": {
          "type": "string",
          "description": "The model architecture to train.",
          "enum": [
            "rgcn",
            "transe",
            "distmult",
            "rotate",
            "custom"
          ]
        },
        "baseProcessingInstanceType": {
          "type": "string"
        },
        "trainingInstanceType": {
          "type": "string",
          "default": "ml.p3.2xlarge"
        },
        "trainingInstanceVolumeSizeInGB": {
          "type": "int32"
        },
        "trainingTimeOutInSeconds": {
          "type": "int32",
          "default": 86400
        },
        "maxHPONumberOfTrainingJobs": {
          "type": "int32",
          "description": "Max total training jobs for hyperparameter tuning.",
          "default": 2
        },
        "maxHPOParallelTrainingJobs": {
          "type": "int32",
          "default": 2
        },
        "subnets": {
          "type": "array",
          "items": {
            "type": "string"
          }
        },
        "securityGroupIds": {
          "type": "array",
          "items": {
            "type": "string"
          }
        },
        "volumeEncryptionKMSKey": {
          "type": "string"
        },
        "s3OutputEncryptionKMSKey": {
          "type": "string"
        },
        "enableInterContainerTrafficEncryption": {
          "type": "boolean",
          "default": true
        },
        "enableManagedSpotTraining": {
          "type": "boolean",
          "default": false
        },
        "customModelTrainingParameters": {
          "type": "object",
          "properties": {
            "sourceS3DirectoryPath": {
              "type": "string"
            },
            "trainingEntryPointScript": {
              "type": "string"
            },
            "transformEntryPointScript": {
              "type": "string"
            }
          }
        }
      }
    },
    "ModelTransformJob": {
      "type": "object",
      "name": "ModelTransformJob",
      "description": "A Neptune ML model transform job that generates model artifacts for inference.",
      "required": [
        "modelTransformOutputS3Location"
      ],
      "properties": {
        "id": {
          "type": "string"
        },
        "dataProcessingJobId": {
          "type": "string"
        },
        "mlModelTrainingJobId": {
          "type": "string"
        },
        "trainingJobName": {
          "type": "string"
        },
        "modelTransformOutputS3Location": {
          "type": "string"
        },
        "sagemakerIamRoleArn": {
          "type": "string"
        },
        "neptuneIamRoleArn": {
          "type": "string"
        },
        "baseProcessingInstanceType": {
          "type": "string"
        },
        "baseProcessingInstanceVolumeSizeInGB": {
          "type": "int32"
        },
        "subnets": {
          "type": "array",
          "items": {
            "type": "string"
          }
        },
        "securityGroupIds": {
          "type": "array",
          "items": {
            "type": "string"
          }
        },
        "volumeEncryptionKMSKey": {
          "type": "string"
        },
        "s3OutputEncryptionKMSKey": {
          "type": "string"
        },
        "enableInterContainerTrafficEncryption": {
          "type": "boolean",
          "default": true
        }
      }
    },
    "InferenceEndpoint": {
      "type": "object",
      "name": "InferenceEndpoint",
      "description": "A Neptune ML inference endpoint backed by Amazon SageMaker for real-time graph predictions.",
      "properties": {
        "id": {
          "type": "string"
        },
        "mlModelTrainingJobId": {
          "type": "string"
        },
        "mlModelTransformJobId": {
          "type": "string"
        },
        "update": {
          "type": "boolean",
          "default": false
        },
        "neptuneIamRoleArn": {
          "type": "string"
        },
        "modelName": {
          "type": "string",
          "enum": [
            "rgcn",
            "kge",
            "transe",
            "distmult",
            "rotate"
          ]
        },
        "instanceType": {
          "type": "string",
          "default": "ml.m5.xlarge"
        },
        "instanceCount": {
          "type": "int32",
          "default": 1,
          "minimum": 1
        },
        "volumeEncryptionKMSKey": {
          "type": "string"
        }
      }
    }
  }
}