{
  "$schema": "https://json-schema.org/draft/2020-12/schema",
  "$id": "https://raw.githubusercontent.com/api-evangelist/amazon-neptune/refs/heads/main/json-schema/ml-create-model-training-job-request-schema.json",
  "title": "CreateModelTrainingJobRequest",
  "description": "CreateModelTrainingJobRequest schema from Neptune",
  "type": "object",
  "properties": {
    "id": {
      "type": "string",
      "description": "Unique identifier for the job."
    },
    "dataProcessingJobId": {
      "type": "string",
      "description": "Job ID of the completed data processing job."
    },
    "trainModelS3Location": {
      "type": "string",
      "description": "S3 location for model artifacts output."
    },
    "previousModelTrainingJobId": {
      "type": "string",
      "description": "Job ID of a previous training job for incremental training."
    },
    "sagemakerIamRoleArn": {
      "type": "string"
    },
    "neptuneIamRoleArn": {
      "type": "string"
    },
    "modelName": {
      "type": "string",
      "description": "The model type to train: rgcn (relational graph convolutional network), transe, distmult, rotate, or custom.",
      "enum": [
        "rgcn",
        "transe",
        "distmult",
        "rotate",
        "custom"
      ]
    },
    "baseProcessingInstanceType": {
      "type": "string",
      "description": "ML instance type for data preparation step."
    },
    "trainingInstanceType": {
      "type": "string",
      "description": "ML instance type for the training step.",
      "default": "ml.p3.2xlarge"
    },
    "trainingInstanceVolumeSizeInGB": {
      "type": "integer",
      "description": "Disk volume size for training instance in GB."
    },
    "trainingTimeOutInSeconds": {
      "type": "integer",
      "description": "Training job timeout in seconds.",
      "default": 86400
    },
    "maxHPONumberOfTrainingJobs": {
      "type": "integer",
      "description": "Maximum total training jobs for hyperparameter tuning. Minimum 10 recommended for meaningful results.",
      "default": 2
    },
    "maxHPOParallelTrainingJobs": {
      "type": "integer",
      "description": "Maximum parallel training jobs.",
      "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",
      "description": "Whether to use EC2 spot instances for training.",
      "default": false
    },
    "customModelTrainingParameters": {
      "type": "object",
      "description": "Custom model training configuration.",
      "properties": {
        "sourceS3DirectoryPath": {
          "type": "string",
          "description": "S3 path to the custom training script directory."
        },
        "trainingEntryPointScript": {
          "type": "string",
          "description": "Name of the training entry point script."
        },
        "transformEntryPointScript": {
          "type": "string",
          "description": "Name of the transform entry point script."
        }
      }
    }
  },
  "required": [
    "dataProcessingJobId",
    "trainModelS3Location"
  ]
}