# Slurm Agent Update

## task type

### Slurm function task
Goal:
1. (optional) execute shell task first
2. then run the code in the function

import os

from flytekit import task, workflow
from flytekitplugins.slurm import SlurmFunction

from flytekit.models.security import Secret 


@task(
    task_config=SlurmFunction(
        slurm_host="xxx",
        srun_conf={
            "partition": "debug",
            "job-name": "demo-slurm-function-task",
        },
           ssh_key_secret_name="SLURM_SSH_KEY",
    ssh_user=os.getenv("USER"),
        script="""
#!/bin/bash
echo "Hello, world!"
export MY_ENV_VAR=123
{flyte.fn}
exit -1
"""
    ),
    secret_requests=[Secret(name="SLURM_SSH_KEY")],
)
def plus_one(x: int) -> int: 
    print(os.getenv("MY_ENV_VAR"))
    return x + 1


@task
def greet(year: int) -> str:
    return f"Hello {year}!!!"


@workflow
def wf(x: int) -> str:
    x = plus_one(x=x)
    msg = greet(year=x)
    return msg


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", "--raw-output-data-prefix", "s3://flyte-test-slurm-0128", path, "wf", "--x", 2024])
    print(result.output)

### Slurm function task
Use format string as input now.
look at this as example.

https://github.com/flyteorg/flytekit/pull/3058

echo = SlurmShellTask(
    name="echo-slurm-shell-task",
    script="""#!/bin/bash
# We can define sbatch options here
echo "Demo Slurm agent with SlurmShellTask...\n"

# Run a demo python script on Slurm
echo "Working!"

python my_script.py --x {inputs.x} --y {inputs.y}

echo "Done!"
""",
    task_config=Slurm(
        slurm_host="ec2-18-207-193-50.compute-1.amazonaws.com",
        sbatch_conf={
            "partition": "debug",
            "job-name": "slurm-shell-task-demo",
        }
    ),
)

echo(x=1, y="2")

Note: you have to make config like the top

ssh_key_secret_name="SLURM_SSH_KEY",
ssh_user=os.getenv("USER"),