import os
from typing import Any, Dict

from flytekit import workflow
from flytekitplugins.slurm import SlurmTask


slurm_tiny_job = SlurmTask(
    name="demo-slurm",
    slurm_config={
        "script": "#!/bin/bash\necho Hello Slurm Agent!",
        "account": "flyte",
        "partition": "debug",
        "name": "hello-slurm-agent",
        "environment": ["PATH=/bin/:/sbin/:/usr/bin/:/usr/sbin/"],
        "current_working_directory": "/tmp"
    },
)


@workflow
def hi_slurm(dummy: str) -> Dict[str, Any]:
    """Return slurm job information."""
    res = slurm_tiny_job(dummy=dummy)

    return res


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, "hi_slurm", "--dummy", "dummy_input"])
    print(result.output)

sbatch

arun