import os

from flytekit import workflow
from flytekitplugins.slurm import Slurm, SlurmTask


echo_job = SlurmTask(
    name="echo-job-name",
    task_config=Slurm(
        slurm_host="slurm",
        # Specify a remote Slurm batch script
        batch_script_path="/home/abaowei/echo.slurm",
        batch_script_args=["--debug"],
        sbatch_conf={
            "partition": "debug",
            "job-name": "tiny-slurm",
        }
    )
    # No container_image is specified here
)


@workflow
def wf() -> None:
    echo_job()


if __name__ == "__main__":
    from flytekit.clis.sdk_in_container import pyflyte
    from click.testing import CliRunner

    runner = CliRunner()
    path = os.path.realpath(__file__)

    # Local run
    print(f">>> LOCAL EXEC <<<")
    result = runner.invoke(pyflyte.main, ["run", path, "wf"])
    print(result.output)