https://flyte.org logo
Join Slack
Powered by
Hi Team, Please take a look at this comment. This...
# slurm-flyte-wg
d
Hi Team, Please take a look at this comment. This video demo showcases two task types on Slurm, running on an AWS EC2 instance: 1. Slurm Python Task – Executes a shell script on the Slurm host. 2. Slurm Shell Task – Executes a shell script provided by my computer. Let me know how @creamy-shampoo-53278 and I can help! PR Link cc @eager-processor-63090
c
Thanks Han-Ru, will also add back 
PythonFunctionTask
tomorrow!
e
This is awesome. I think once you re-add PythonFunctionTask we should have more than enough task types for this agent. We should then focus on QoL improvements like log streaming, stability, and informative error messages. I'm working on building out a representative workflow so I'll be able to surface and fix these as they come up. Excellent work so far!!
c
(Update, the latest commit) 
PythonFunctionTask
is now added! We'll briefly describe how tasks/agents are categorized. We first split them into 
script/
and 
function/
under slurm plugin root. For the former, we have: 1. 
SlurmTask
: Executes a batch/shell script on the Slurm host 2. `SlurmShellTask`: Execute an user-defined batch/shell script from localhost (e.g., our local dev pc) As for the latter, we write an additional 
SlurmFunctionTask
which is inherited from 
PythonFunctionTask
. This supports running an user-defined function in the 
task
body by executing flytekit entrypoints on the Slurm host. As the structure is slightly changed and tasks are renamed, we need to modify the interface a little bit!
Hi Pryce, Sounds amazing. Thanks for encouragement! I'll provide python scripts we use to test these 3 types of tasks here to keep the PR clean, just as you mentioned in another thread. Then, we can move on to QoL improvements!
SlurmTask
Copy code
import os

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



echo_job = SlurmTask(
    name="<task-name>",
    task_config=SlurmRemoteScript(
        slurm_host="<slurm-host>",
        # Specify a remote Slurm batch script
        batch_script_path="<remote-batch-script-path>",
        sbatch_conf={
            "partition": "debug",
            "job-name": "tiny-slurm",
        }
    )
)


@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)
SlurmShellTask
Copy code
import os

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


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!"
""",
    task_config=Slurm(
        slurm_host="<slurm-host>",
        sbatch_conf={
            "partition": "debug",
            "job-name": "tiny-slurm",
        }
    ),
)


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



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)
SlurmFunctionTask
Copy code
import os
from typing import Any, Dict

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


@task(
    task_config=SlurmFunction(
        slurm_host="<slurm-host>",
        srun_conf={
            "partition": "debug",
            "job-name": "demo-slurm-function-task",
            "chdir": "<working-dir>"
        }
    )
)
def plus_one(x: int) -> int: 
    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-url>", path, "wf", "--x", 2024])
    print(result.output)
Open in Slack
PreviousNext
Powered by