{
  "$schema": "https://json-schema.org/draft/2020-12/schema",
  "$id": "https://raw.githubusercontent.com/api-evangelist/amazon-neptune/refs/heads/main/json-schema/data-start-ml-model-training-job-input-schema.json",
  "title": "StartMLModelTrainingJobInput",
  "description": "StartMLModelTrainingJobInput schema from Neptune",
  "type": "object",
  "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"
    },
    "sagemakerIamRoleArn": {
      "type": "string"
    },
    "neptuneIamRoleArn": {
      "type": "string"
    },
    "modelName": {
      "type": "string",
      "description": "The model type to use.",
      "enum": [
        "rgcn",
        "transe",
        "distmult",
        "rotate",
        "custom"
      ]
    },
    "baseProcessingInstanceType": {
      "type": "string"
    },
    "trainingInstanceType": {
      "type": "string",
      "default": "ml.p3.2xlarge"
    },
    "trainingInstanceVolumeSizeInGB": {
      "type": "integer"
    },
    "trainingTimeOutInSeconds": {
      "type": "integer",
      "default": 86400
    },
    "maxHPONumberOfTrainingJobs": {
      "type": "integer",
      "default": 2
    },
    "maxHPOParallelTrainingJobs": {
      "type": "integer",
      "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"
        }
      }
    }
  },
  "required": [
    "dataProcessingJobId",
    "trainModelS3Location"
  ]
}