#!/usr/bin/env Rscript

library(this.path)
library(renv)
load(here(".."))

# Import training logic
source(here("training.R"))

# Setup parameters
# Container directories
prefix <- "/opt/ml"
input_data_path <- paste(prefix, "input/data/training", sep = "/")
model_save_path <- paste(prefix, "model", sep = "/")
failure_output <- paste(prefix, "failure", sep = "/")

tryCatch(
    {
        train_function(input_data_path = input_data_path, model_save_path = model_save_path)

        # A zero exit code causes the job to be marked a Succeeded.
        q(status = 0)
    },
    error = function(e) {
        # Handle error here
        write(e$message, file = paste(failure_output, "failure", sep = "/"))
        cat(e$message, file = stderr())

        # A non-zero exit code causes the training job to be marked as Failed.
        q(status = 255)
    }
)
