from transforms.api import Input, Output, transform
from transforms.sidecar import Volume, sidecar
from myproject.datasets.sidecar_client import (
    copy_artifacts,
    copy_dataset_files,
    execute_cli,
)


@sidecar(image="simple-cli", tag="0.0.1", volumes=[Volume("shared")])
@transform(
    output=Output("/training/containers/cli-single-output"),
    source=Input("/training/containers/input"),
)
def compute(output, source, ctx):
    def run_once(_):
        def prepare_input(data_directory):
            copy_dataset_files(source, data_directory)
        with execute_cli(prepare_input) as request:
            copy_artifacts(output, request)

    ctx.spark_session.sparkContext.parallelize([1], 1).foreach(run_once)