[WIP] Slurm agent #3005

JiangJiaWei1103 · 2024-12-16T15:23:33Z

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

Implement the Slurm agent, which submits the user-defined flytekit task to a remote Slurm cluster to run. Following describe three core methods:

create: Submit a Slurm job with sbatch to run a batch script on Slurm cluster
get: Check the Slurm job state
delete (haven't been tested): Cancel the Slurm job

How was this patch tested?

We test create and get in the development environment described as follows:

Local: MacBook with flytekit installed
- Test slurm agent locally following this guide
Remote: Single Ubuntu server with slurmctld and slurmd running
- We plan to write a single-host setup tutorial and organize useful resources here
Communication is done with asyncssh

Suppose we have a batch script to run on Slurm cluster:

#!/bin/bash

echo "Working!" >> ./remote_touch.txt

We use the following python script to test Slurm agent on the client side:

import os

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


echo_job = SlurmTask(
    name="echo-job-name",
    task_config=Slurm(
        slurm_host="<host_alias>",
        batch_script_path="<path_to_batch_script_within_slurm_cluster>",
        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)

The test result is shown as follows:

Setup process

As stated above

Check all the applicable boxes

I updated the documentation accordingly.
All new and existing tests passed.
All commits are signed-off.

Related PRs

Docs link

Signed-off-by: jiangjiawei1103 <[email protected]>

Signed-off-by: JiaWei Jiang <[email protected]>

codecov · 2024-12-19T15:01:17Z

Codecov Report

Attention: Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.

Project coverage is 76.63%. Comparing base (f99d50e) to head (9e6d8a6).
Report is 18 commits behind head on master.

Files with missing lines	Patch %	Lines
flytekit/extend/backend/utils.py	0.00%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #3005       +/-   ##
===========================================
+ Coverage   51.08%   76.63%   +25.55%     
===========================================
  Files         201      201               
  Lines       21231    21274       +43     
  Branches     2731     2733        +2     
===========================================
+ Hits        10846    16304     +5458     
+ Misses       9787     4219     -5568     
- Partials      598      751      +153

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Successfully submit and run the user-defined task as a normal python function on a remote Slurm cluster. 1. Inherit from PythonFunctionTask instead of PythonTask 2. Transfer the task module through sftp 3. Interact with amazon s3 bucket on both localhost and Slurm cluster Signed-off-by: JiaWei Jiang <[email protected]>

Specifying `--raw-output-data-prefix` option handles task_module download. Signed-off-by: JiaWei Jiang <[email protected]>

Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2024-12-28T06:14:14Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2024-12-30T15:08:45Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

Add `ssh_conf` filed to let users specify connection secret Note that reconnection is done in both `get` and `delete`. This is just a temporary workaround. Signed-off-by: JiaWei Jiang <[email protected]>

Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2024-12-31T15:08:45Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

For data scientists and MLEs developing flyte wf with Slurm agent, they don't actually need to know ssh connection details. We assume they only need to specify which Slurm cluster to use by hostname. Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2025-01-07T14:27:19Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

1. Write user-defined batch script to a tmp file 2. Transfer the batch script through sftp 3. Construct sbatch command to run on Slurm cluster Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2025-01-07T17:07:18Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

1. Remove SFTP for batch script transfer * Assume Slurm batch script is present on Slurm cluster 2. Support directly specifying a remote batch script path Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2025-01-09T16:27:41Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

Signed-off-by: pryce-turner <[email protected]>

flyte-bot · 2025-01-10T01:17:53Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2025-01-11T14:31:05Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2025-01-11T16:31:05Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

Future-Outlier · 2025-01-14T12:31:41Z

flytekit/extend/backend/base_agent.py

@@ -326,7 +326,7 @@ class AsyncAgentExecutorMixin:

    def execute(self: PythonTask, **kwargs) -> LiteralMap:
        ctx = FlyteContext.current_context()
-        ss = ctx.serialization_settings or SerializationSettings(ImageConfig())
+        ss = ctx.serialization_settings or SerializationSettings(ImageConfig.auto_default_image())


why do we need this?
is this for shell task?

If we define a SlurmTask without specifying container_image (as the example python script provided above), ctx.serialization_settings will be None. Then, an error is raised which describes that PythonAutoContainerTask needs an image.

I think this is just a temporary workaround for local test and I'm still pondering how to better handle this issue.

Future-Outlier · 2025-01-14T12:40:07Z

plugins/flytekit-slurm/assets/basic_arch.png

Amazing Graph.

Thanks, bro.

`SlurmTask` and `SlurmShellTask` now share the same agent. Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2025-01-14T17:57:38Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

Future-Outlier

1. Inherited from `PythonTask` for cases in which the batch script is already on the Slurm cluster 2. Use a dummy `Interface` as a tmp workaround Signed-off-by: JiaWei Jiang <[email protected]>

Future-Outlier

This video demo 2 things.

Slurm Python task (executes a shell script on the Slurm host).
Slurm shell task (executes a shell script provided on my computer).

output.mp4

…ript Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2025-01-15T15:43:37Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

JiangJiaWei1103 · 2025-01-15T16:02:32Z

This video demo 2 things.

* Slurm Python task (executes a shell script on the Slurm host).

* Slurm shell task (executes a shell script provided on my computer).

output.mp4

Nice bro!

I'll push PythonFunctionTask tomorrow (got some errors to solve).

pryce-turner · 2025-01-15T19:45:45Z

Now that we can reason more sensibly about the inner-workings of this agent I want to get a conversation going early about the object store and where it fits into all of this. IMO having a consistent persistence layer is the killer feature here and required for composing workflows with Slurm tasks alongside all the other flyte task types.

For this initial implementation we should assume workflows composed entirely of slurm tasks on the local filesystem. We can then orchestrate tasks by passing filepaths around as input and output. This is a fairly naive implementation but will still have the benefit of the console, remote execution, versioning, logging, etc. We should nevertheless build with a consistent object store in mind down the road.

Some high perf filesystems will have the option to expose an S3 interface but I don't think we should assume that's always the case. We may need to fall back to some awscli dependency for getting/putting inputs and outputs. This is all down the road but let's keep it in mind! Usecases where people have an on-prem slurm cluster but no GPUs for example can run CPU bound tasks there and then seamlessly offload accelerated tasks, all in the same workflow, is a killer feature.

Let me know your thoughts.

1. Add back `PythonFunctionTask` to support running user-defined functions on Slurm 2. Categorize task types into `script/` and `function/` Signed-off-by: JiaWei Jiang <[email protected]>

flyte-bot · 2025-01-16T13:57:35Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

JiangJiaWei1103 added 5 commits December 14, 2024 19:18

Add slurm plugin blank components

421d1b8

Signed-off-by: jiangjiawei1103 <[email protected]>

feat: Add naive slurm agent create and get with rest api

1d1f806

Signed-off-by: jiangjiawei1103 <[email protected]>

Use asyncssh instead of REST API

5d97126

Signed-off-by: jiangjiawei1103 <[email protected]>

Test ssh communication and run sbatch

2e7f0f2

Signed-off-by: JiaWei Jiang <[email protected]>

Add delete method and support slurm job state

9644b99

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 added 3 commits December 27, 2024 15:29

refactor: Remove redundant task_module transfer

6db24dc

Specifying `--raw-output-data-prefix` option handles task_module download. Signed-off-by: JiaWei Jiang <[email protected]>

refactor: Remove redundant env var

122c7f1

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 added 5 commits December 30, 2024 22:07

docs: Add env setup guide for local test

e9760a7

Signed-off-by: JiaWei Jiang <[email protected]>

docs: Add links and figures

e68fda9

Signed-off-by: JiaWei Jiang <[email protected]>

docs: Fix commit sha

470637c

Signed-off-by: JiaWei Jiang <[email protected]>

docs: Fix commit sha for demo guide

1579ab4

Signed-off-by: JiaWei Jiang <[email protected]>

docs: Fix links

0e538f0

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 added 2 commits December 31, 2024 21:56

feat: Support SSH config in task config

8229418

Add `ssh_conf` filed to let users specify connection secret Note that reconnection is done in both `get` and `delete`. This is just a temporary workaround. Signed-off-by: JiaWei Jiang <[email protected]>

docs: Include ssh config in demo example

9e6d8a6

Signed-off-by: JiaWei Jiang <[email protected]>

feat: Support Slurm agent with ShellTask

3a7eb6d

1. Write user-defined batch script to a tmp file 2. Transfer the batch script through sftp 3. Construct sbatch command to run on Slurm cluster Signed-off-by: JiaWei Jiang <[email protected]>

feat: Simplify Slurm job submission logic

a815fd9

1. Remove SFTP for batch script transfer * Assume Slurm batch script is present on Slurm cluster 2. Support directly specifying a remote batch script path Signed-off-by: JiaWei Jiang <[email protected]>

Added script args to agent and task

a3ea014

Signed-off-by: pryce-turner <[email protected]>

JiangJiaWei1103 mentioned this pull request Jan 10, 2025

[BUG] No default image is assigned to an agent running a containerized task flyteorg/flyte#6157

Open

2 tasks

Add asyncssh to dependencies

a109bd8

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 added 4 commits January 11, 2025 23:22

docs: Update setup and demo for a basic use case

e5da665

Signed-off-by: JiaWei Jiang <[email protected]>

docs: Update basic arch figure path

0a3d9f1

Signed-off-by: JiaWei Jiang <[email protected]>

docs: Fix typo and hyperlink

1b0f6df

Signed-off-by: JiaWei Jiang <[email protected]>

fix: A tmp workaround to test agent locally without container_image

26cc201

Signed-off-by: JiaWei Jiang <[email protected]>

Future-Outlier self-assigned this Jan 13, 2025

Future-Outlier reviewed Jan 14, 2025

View reviewed changes

feat: Support user-defined batch script content with SlurmShellTask

16d953e

`SlurmTask` and `SlurmShellTask` now share the same agent. Signed-off-by: JiaWei Jiang <[email protected]>

Future-Outlier reviewed Jan 15, 2025

View reviewed changes

feat: Fall back to PythonTask for naive use cases

c743917

1. Inherited from `PythonTask` for cases in which the batch script is already on the Slurm cluster 2. Use a dummy `Interface` as a tmp workaround Signed-off-by: JiaWei Jiang <[email protected]>

Future-Outlier reviewed Jan 15, 2025

View reviewed changes

refactor: Define Slurm as a base task config and extend for remote sc…

e365dee

…ript Signed-off-by: JiaWei Jiang <[email protected]>

feat: Support PythonFunctionTask and reorganize agent structure

c1064d4

1. Add back `PythonFunctionTask` to support running user-defined functions on Slurm 2. Categorize task types into `script/` and `function/` Signed-off-by: JiaWei Jiang <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WIP] Slurm agent #3005

[WIP] Slurm agent #3005

JiangJiaWei1103 commented Dec 16, 2024 •

edited

Loading

codecov bot commented Dec 19, 2024 •

edited

Loading

flyte-bot commented Dec 28, 2024

flyte-bot commented Dec 30, 2024

flyte-bot commented Dec 31, 2024

flyte-bot commented Jan 7, 2025

flyte-bot commented Jan 7, 2025

flyte-bot commented Jan 9, 2025

flyte-bot commented Jan 10, 2025

flyte-bot commented Jan 11, 2025

flyte-bot commented Jan 11, 2025

Future-Outlier Jan 14, 2025

JiangJiaWei1103 Jan 14, 2025 •

edited

Loading

Future-Outlier Jan 14, 2025

JiangJiaWei1103 Jan 15, 2025

flyte-bot commented Jan 14, 2025

Future-Outlier left a comment

Future-Outlier left a comment •

edited

Loading

flyte-bot commented Jan 15, 2025

JiangJiaWei1103 commented Jan 15, 2025 •

edited

Loading

pryce-turner commented Jan 15, 2025

flyte-bot commented Jan 16, 2025

[WIP] Slurm agent #3005

Are you sure you want to change the base?

[WIP] Slurm agent #3005

Conversation

JiangJiaWei1103 commented Dec 16, 2024 • edited Loading

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Check all the applicable boxes

Related PRs

Docs link

codecov bot commented Dec 19, 2024 • edited Loading

Codecov Report

flyte-bot commented Dec 28, 2024

Code Review Agent Run Status

flyte-bot commented Dec 30, 2024

Code Review Agent Run Status

flyte-bot commented Dec 31, 2024

Code Review Agent Run Status

flyte-bot commented Jan 7, 2025

Code Review Agent Run Status

flyte-bot commented Jan 7, 2025

Code Review Agent Run Status

flyte-bot commented Jan 9, 2025

Code Review Agent Run Status

flyte-bot commented Jan 10, 2025

Code Review Agent Run Status

flyte-bot commented Jan 11, 2025

Code Review Agent Run Status

flyte-bot commented Jan 11, 2025

Code Review Agent Run Status

Future-Outlier Jan 14, 2025

Choose a reason for hiding this comment

JiangJiaWei1103 Jan 14, 2025 • edited Loading

Choose a reason for hiding this comment

Future-Outlier Jan 14, 2025

Choose a reason for hiding this comment

JiangJiaWei1103 Jan 15, 2025

Choose a reason for hiding this comment

flyte-bot commented Jan 14, 2025

Code Review Agent Run Status

Future-Outlier left a comment

Choose a reason for hiding this comment

Future-Outlier left a comment • edited Loading

Choose a reason for hiding this comment

flyte-bot commented Jan 15, 2025

Code Review Agent Run Status

JiangJiaWei1103 commented Jan 15, 2025 • edited Loading

pryce-turner commented Jan 15, 2025

flyte-bot commented Jan 16, 2025

Code Review Agent Run Status

JiangJiaWei1103 commented Dec 16, 2024 •

edited

Loading

codecov bot commented Dec 19, 2024 •

edited

Loading

JiangJiaWei1103 Jan 14, 2025 •

edited

Loading

Future-Outlier left a comment •

edited

Loading

JiangJiaWei1103 commented Jan 15, 2025 •

edited

Loading