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

Summary by Bito

Implementation of a new Slurm agent for executing Flyte tasks on remote clusters, featuring job submission, monitoring, and cancellation capabilities via SSH. The enhancement adds support for output location interpolation and improved job metadata handling, with focus on dynamic script generation and better output management in the SlurmScriptAgent class.

Unit tests added: False

Estimated effort to review (1-5, lower is better): 3

Signed-off-by: jiangjiawei1103 <waynechuang97@gmail.com>

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

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

Codecov Report

Attention: Patch coverage is 50.00000% with 7 lines in your changes missing coverage. Please review.

Project coverage is 46.62%. Comparing base (9d34416) to head (d0967bd).
Report is 27 commits behind head on master.

Files with missing lines	Patch %	Lines
flytekit/extras/tasks/shell.py	36.36%	4 Missing and 3 partials ⚠️

❗ There is a different number of reports uploaded between BASE (9d34416) and HEAD (d0967bd). Click for more details.

HEAD has 24 uploads less than BASE

Flag BASE (9d34416) HEAD (d0967bd)

25 1

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #3005       +/-   ##
===========================================
- Coverage   78.20%   46.62%   -31.58%     
===========================================
  Files         292      206       -86     
  Lines       25401    21786     -3615     
  Branches     2779     2839       +60     
===========================================
- Hits        19864    10157     -9707     
- Misses       4726    11114     +6388     
+ Partials      811      515      -296

☔ 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 <waynechuang97@gmail.com>

Specifying `--raw-output-data-prefix` option handles task_module download. Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

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 eduardo@union.ai.

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

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 eduardo@union.ai.

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 <waynechuang97@gmail.com>

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

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 eduardo@union.ai.

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 <waynechuang97@gmail.com>

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 eduardo@union.ai.

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 <waynechuang97@gmail.com>

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 eduardo@union.ai.

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 <waynechuang97@gmail.com>

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 eduardo@union.ai.

Signed-off-by: pryce-turner <pryce.turner@gmail.com>

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 eduardo@union.ai.

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

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 eduardo@union.ai.

JiangJiaWei1103 · 2025-02-14T14:52:20Z

Both stdout and stderr message of the Slurm cluster can be shown on the users' terminal for SlurmTask and SlurmShellTask. Following provide local test results:

`SlurmTask`

Echo in a for loop

Before	After

Raise an exception with a python script

raise Exception("Test err msg passing!")

Before	After

`SlurmShellTask`

Echo a message

Before	After

Raise an exception with a python script

raise Exception("Test err msg passing!")

Before	After

…onTask Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

JiangJiaWei1103 · 2025-02-14T15:18:04Z

SlurmFunctionTask also supports stdout and stderr message display, as shown below:

SlurmFunctionTask with echo and python print

@task(
    task_config=SlurmFunction(
        slurm_host="aws2",
        sbatch_conf={
            ...
        },
        script="""#!/bin/bash

# == Pre-Execution ==
echo "Hello, world!"

# Setup env vars
export MY_ENV_VAR=123

# Activate virtual env
. /home/ubuntu/.cache/pypoetry/virtualenvs/demo-4A8TrTN7-py3.12/bin/activate 

# == Execute Flyte Task Function ==
{task.fn}

# == Post-Execution ==
echo "Success!!"
"""
    )
)
def plus_one(x: int) -> int: 
    print(os.getenv("MY_ENV_VAR"))
    return x + 1

Before	After

Raise an ZeroDivisionError

@task(
    task_config=SlurmFunction(
        slurm_host="aws2",
        sbatch_conf={
            ...
        },
        script="""#!/bin/bash

# == Pre-Execution ==
echo "Let's make this task fail..."

# Setup env vars
export MY_ENV_VAR=123

# Activate virtual env
. /home/ubuntu/.cache/pypoetry/virtualenvs/demo-4A8TrTN7-py3.12/bin/activate 

# == Execute Flyte Task Function ==
{task.fn}

# == Post-Execution ==
echo "Success!!"
"""
    )
)
def divide_zero(x: int) -> float: 
    print(os.getenv("MY_ENV_VAR"))
    return x / 0

As can be observed, the root cause ZeroDivisionError, can be correctly displayed on the terminal.

Before	After

flyte-bot · 2025-02-14T16:29:34Z

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 eduardo@union.ai.

1. Make SSH `host` and `username` required fields 2. Support SSH connection based on the default OpenSSH client config file `~/.ssh/config` 3. Support SSH connection via public key auth either by user-specified `client_keys` or the secret for key `FLYTE_SLURM_PRIVATE_KEY` Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

flyte-bot · 2025-02-16T16:22:29Z

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 eduardo@union.ai.

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

flyte-bot · 2025-02-17T15:32:43Z

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 eduardo@union.ai.

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

flyte-bot · 2025-02-18T15:00:57Z

Code Review Agent Run #e8c545

Actionable Suggestions - 8

plugins/flytekit-slurm/flytekitplugins/slurm/script/task.py - 1
- Consider proper interface initialization · Line 61-61
flytekit/extras/tasks/shell.py - 1
- Consider using fully qualified class name · Line 265-265
plugins/flytekit-slurm/flytekitplugins/slurm/function/agent.py - 1
- Insecure temporary file usage · Line 39-39
plugins/flytekit-slurm/flytekitplugins/slurm/script/agent.py - 3
- Use secure temporary file creation · Line 40-40
- Initialize msg variable with default value · Line 96-100
- Consider using shlex.join for commands · Line 164-164
plugins/flytekit-slurm/flytekitplugins/slurm/ssh_utils.py - 2
- Consider secure private key handling · Line 92-96
- Replace blocking IO with async alternative · Line 94-95

Additional Suggestions - 2

flytekit/extras/tasks/shell.py - 2
- Consider using guard clause pattern · Line 381-384
- Simplify task_type conditional expression · Line 284-286

Review Details

Files reviewed - 10 · Commit Range: 421d1b8..f80ccd4
- flytekit/extend/backend/base_agent.py
- flytekit/extend/backend/utils.py
- flytekit/extras/tasks/shell.py
- plugins/flytekit-slurm/flytekitplugins/slurm/__init__.py
- plugins/flytekit-slurm/flytekitplugins/slurm/function/agent.py
- plugins/flytekit-slurm/flytekitplugins/slurm/function/task.py
- plugins/flytekit-slurm/flytekitplugins/slurm/script/agent.py
- plugins/flytekit-slurm/flytekitplugins/slurm/script/task.py
- plugins/flytekit-slurm/flytekitplugins/slurm/ssh_utils.py
- plugins/flytekit-slurm/setup.py
Files skipped - 2
- plugins/flytekit-slurm/README.md - Reason: Filter setting
- plugins/flytekit-slurm/demo.md - Reason: Filter setting
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-02-18T15:13:32Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
New Feature - Slurm Agent Implementation	- `agent.py` - Implements core Slurm function agent for job submission and monitoring - `task.py` - Defines Slurm function task types and configurations - `agent.py` - Implements Slurm script agent for executing batch scripts - `task.py` - Defines Slurm script task types and shell task integration - `ssh_utils.py` - Provides SSH utilities for Slurm cluster communication - `setup.py` - Sets up Slurm plugin package configuration
Feature Improvement - Task State Handling Enhancement	- `base_agent.py` - Updates image configuration handling - `utils.py` - Enhances task state conversion with additional states - `shell.py` - Improves shell task integration with Slurm support

flyte-bot · 2025-02-18T15:13:34Z

plugins/flytekit-slurm/flytekitplugins/slurm/script/task.py

+            name=name,
+            task_config=task_config,
+            # Dummy interface, will support this after discussion
+            interface=Interface(None, None),


Consider proper interface initialization

The interface initialization with None values in SlurmTask.__init__() may cause issues. Consider defining a proper interface with input/output types or providing a clearer explanation for using None values.

Code suggestion

Check the AI-generated fix before applying

Suggested change

interface=Interface(None, None),

interface=Interface(inputs={}, outputs={}),

Code Review Run #e8c545

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

Yes, avoid them

flyte-bot · 2025-02-18T15:13:35Z

flytekit/extras/tasks/shell.py

-        # tasks should not be serialized at all, but we don't currently have a mechanism for skipping Flyte entities
-        # at serialization time.
-        self._config_task_instance._name = f"_bash.{name}"
+        if plugin_class.__name__ in ["SlurmShellTask"]:


Consider using fully qualified class name

Consider using a more explicit condition by checking the fully qualified class name instead of just the class name. The current check plugin_class.__name__ in ["SlurmShellTask"] may be fragile if class names change or if there are name collisions.

Code suggestion

Check the AI-generated fix before applying

Suggested change

if plugin_class.__name__ in ["SlurmShellTask"]:

plugin_class_fqn = f"{plugin_class.__module__}.{plugin_class.__name__}"

if plugin_class_fqn in ["flytekitplugins.slurm.script.task.SlurmShellTask"]:

Code Review Run #e8c545

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

Yes, avoid them

flyte-bot · 2025-02-18T15:13:37Z

plugins/flytekit-slurm/flytekitplugins/slurm/function/agent.py

+    _conn: Optional[SSHClientConnection] = None
+
+    # Tmp remote path of the batch script
+    REMOTE_PATH = "/tmp/task.slurm"


Insecure temporary file usage

Using hardcoded temporary file path /tmp/task.slurm could be insecure. Consider using tempfile.mkstemp() instead.

Code suggestion

Check the AI-generated fix before applying

Suggested change

REMOTE_PATH = "/tmp/task.slurm"

_, REMOTE_PATH = tempfile.mkstemp(suffix='.slurm')

Code Review Run #e8c545

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

Yes, avoid them

flyte-bot · 2025-02-18T15:13:37Z

plugins/flytekit-slurm/flytekitplugins/slurm/script/agent.py

+    _conn: Optional[SSHClientConnection] = None
+
+    # Tmp remote path of the batch script
+    REMOTE_PATH = "/tmp/echo_shell.slurm"


Use secure temporary file creation

The hardcoded temporary file path /tmp/echo_shell.slurm could pose a security risk. Consider using tempfile.mkstemp() to generate secure temporary files.

Code suggestion

Check the AI-generated fix before applying

Suggested change

REMOTE_PATH = "/tmp/echo_shell.slurm"

_tmp_fd, REMOTE_PATH = tempfile.mkstemp(suffix='.slurm')

os.close(_tmp_fd)

Code Review Run #e8c545

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

Yes, avoid them

flyte-bot · 2025-02-18T15:13:39Z

plugins/flytekit-slurm/flytekitplugins/slurm/script/agent.py

+        # Determine the current flyte phase from Slurm job state
+        job_state = "running"
+        for o in job_res.stdout.split(" "):
+            if "JobState" in o:
+                job_state = o.split("=")[1].strip().lower()


Initialize msg variable with default value

Consider handling the case where StdOut path is not found in the job output. Currently, msg will be undefined if the StdOut path is not present in the output, which could lead to a runtime error. Consider initializing msg with a default value.

Code suggestion

Check the AI-generated fix before applying

Suggested change

# Determine the current flyte phase from Slurm job state

job_state = "running"

for o in job_res.stdout.split(" "):

if "JobState" in o:

job_state = o.split("=")[1].strip().lower()

# Determine the current flyte phase from Slurm job state

job_state = "running"

msg = "No output available"

for o in job_res.stdout.split(" "):

if "JobState" in o:

job_state = o.split("=")[1].strip().lower()

Code Review Run #e8c545

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

Yes, avoid them

flyte-bot · 2025-02-18T15:13:40Z

plugins/flytekit-slurm/flytekitplugins/slurm/script/agent.py

+        for arg in batch_script_args:
+            cmd.append(arg)
+
+    cmd = " ".join(cmd)


Consider using shlex.join for commands

Consider using shlex.join() instead of " ".join() for shell command construction to properly handle arguments containing spaces or special characters.

Code suggestion

Check the AI-generated fix before applying

Suggested change

cmd = " ".join(cmd)

import shlex

cmd = shlex.join(cmd)

Code Review Run #e8c545

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

Yes, avoid them

flyte-bot · 2025-02-18T15:13:41Z

plugins/flytekit-slurm/flytekitplugins/slurm/ssh_utils.py

+        # Write the private key to a local path
+        # This may not be a good practice...
+        with open("./slurm_private_key", "w") as f:
+            f.write(default_client_key)
+        client_keys.append("./slurm_private_key")


Consider secure private key handling

Consider using a more secure approach for storing the private key. Writing sensitive data to a local file at ./slurm_private_key could pose security risks. Consider using environment variables or a secure key management system.

Code suggestion

Check the AI-generated fix before applying

Suggested change

# Write the private key to a local path

# This may not be a good practice...

with open("./slurm_private_key", "w") as f:

f.write(default_client_key)

client_keys.append("./slurm_private_key")

# Store private key in protected directory with secure permissions

key_dir = os.path.expanduser("~/.ssh")

os.makedirs(key_dir, mode=0o700, exist_ok=True)

key_path = os.path.join(key_dir, "slurm_private_key")

with open(key_path, "w") as f:

os.chmod(key_path, 0o600)

f.write(default_client_key)

client_keys.append(key_path)

Code Review Run #e8c545

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

Yes, avoid them

flyte-bot · 2025-02-18T15:13:42Z

plugins/flytekit-slurm/flytekitplugins/slurm/ssh_utils.py

+        with open("./slurm_private_key", "w") as f:
+            f.write(default_client_key)


Replace blocking IO with async alternative

Using blocking open() in an async function can impact performance. Consider using aiofiles for async file operations.

Code suggestion

Check the AI-generated fix before applying

Suggested change

with open("./slurm_private_key", "w") as f:

f.write(default_client_key)

async with aiofiles.open("./slurm_private_key", "w") as f:

await f.write(default_client_key)

Code Review Run #e8c545

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

Yes, avoid them

Support passing files across multiple `SlurmShellTask` Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

flyte-bot · 2025-02-18T17:31:10Z

Code Review Agent Run #0e549b

Actionable Suggestions - 1

plugins/flytekit-slurm/flytekitplugins/slurm/script/agent.py - 1
- Consider validating outputs before usage · Line 112-112

Review Details

Files reviewed - 2 · Commit Range: f80ccd4..ac25446
- plugins/flytekit-slurm/flytekitplugins/slurm/script/agent.py
- plugins/flytekit-slurm/flytekitplugins/slurm/script/task.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-02-18T17:35:14Z

plugins/flytekit-slurm/flytekitplugins/slurm/script/agent.py

+                msg = msg_res.stdout
+        cur_phase = convert_to_flyte_phase(job_state)
+
+        return Resource(phase=cur_phase, message=msg, outputs=resource_meta.outputs)


Consider validating outputs before usage

Consider validating resource_meta.outputs before using it in the Resource constructor. The outputs field could potentially be None which might cause issues downstream.

Code suggestion

Check the AI-generated fix before applying

Suggested change

return Resource(phase=cur_phase, message=msg, outputs=resource_meta.outputs)

outputs = resource_meta.outputs if resource_meta.outputs is not None else {}

return Resource(phase=cur_phase, message=msg, outputs=outputs)

Code Review Run #0e549b

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

Yes, avoid them

JiangJiaWei1103 added 5 commits December 14, 2024 19:18

Add slurm plugin blank components

421d1b8

Signed-off-by: jiangjiawei1103 <waynechuang97@gmail.com>

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

1d1f806

Signed-off-by: jiangjiawei1103 <waynechuang97@gmail.com>

Use asyncssh instead of REST API

5d97126

Signed-off-by: jiangjiawei1103 <waynechuang97@gmail.com>

Test ssh communication and run sbatch

2e7f0f2

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

Add delete method and support slurm job state

9644b99

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

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 <waynechuang97@gmail.com>

refactor: Remove redundant env var

122c7f1

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

JiangJiaWei1103 added 5 commits December 30, 2024 22:07

docs: Add env setup guide for local test

e9760a7

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

docs: Add links and figures

e68fda9

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

docs: Fix commit sha

470637c

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

docs: Fix commit sha for demo guide

1579ab4

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

docs: Fix links

0e538f0

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

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 <waynechuang97@gmail.com>

docs: Include ssh config in demo example

9e6d8a6

Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>

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 <waynechuang97@gmail.com>

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 <waynechuang97@gmail.com>

Added script args to agent and task

a3ea014

Signed-off-by: pryce-turner <pryce.turner@gmail.com>

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 <waynechuang97@gmail.com>

feat: Show stdout and stderr msg of the Slurm cluster for SlurmFuncti…

a465607

…onTask Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

Clarify SSH connection logic

d0967bd

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

feat: Interpolate the script with dynamic input values

f80ccd4

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

Future-Outlier changed the title ~~[WIP] Slurm agent~~ Slurm agent Feb 18, 2025

Future-Outlier marked this pull request as ready for review February 18, 2025 14:50

Future-Outlier requested review from wild-endeavor, eapolinario, pingsutw, cosmicBboy, samhita-alla and thomasjpfan as code owners February 18, 2025 14:50

flyte-bot reviewed Feb 18, 2025

View reviewed changes

feat: Interpolate the script with dynamic output values

ac25446

Support passing files across multiple `SlurmShellTask` Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

flyte-bot reviewed Feb 18, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Slurm agent #3005

Slurm agent #3005

JiangJiaWei1103 commented Dec 16, 2024 •

edited by flyte-bot

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

JiangJiaWei1103 commented Feb 14, 2025 •

edited

Loading

JiangJiaWei1103 commented Feb 14, 2025

flyte-bot commented Feb 14, 2025

flyte-bot commented Feb 16, 2025

flyte-bot commented Feb 17, 2025

flyte-bot commented Feb 18, 2025 •

edited

Loading

Code Review Agent Run #e8c545

flyte-bot commented Feb 18, 2025 •

edited

Loading

Changelist by Bito

flyte-bot Feb 18, 2025

flyte-bot Feb 18, 2025

flyte-bot Feb 18, 2025

flyte-bot Feb 18, 2025

flyte-bot Feb 18, 2025

flyte-bot Feb 18, 2025

flyte-bot Feb 18, 2025

flyte-bot Feb 18, 2025

flyte-bot commented Feb 18, 2025 •

edited

Loading

Code Review Agent Run #0e549b

flyte-bot Feb 18, 2025

	interface=Interface(None, None),
	interface=Interface(inputs={}, outputs={}),

	if plugin_class.__name__ in ["SlurmShellTask"]:
	plugin_class_fqn = f"{plugin_class.__module__}.{plugin_class.__name__}"
	if plugin_class_fqn in ["flytekitplugins.slurm.script.task.SlurmShellTask"]:

	REMOTE_PATH = "/tmp/task.slurm"
	_, REMOTE_PATH = tempfile.mkstemp(suffix='.slurm')

	REMOTE_PATH = "/tmp/echo_shell.slurm"
	_tmp_fd, REMOTE_PATH = tempfile.mkstemp(suffix='.slurm')
	os.close(_tmp_fd)

-        # Write the private key to a local path
-        # This may not be a good practice...
-        with open("./slurm_private_key", "w") as f:
-            f.write(default_client_key)
-        client_keys.append("./slurm_private_key")
+        # Store private key in protected directory with secure permissions
+        key_dir = os.path.expanduser("~/.ssh")
+        os.makedirs(key_dir, mode=0o700, exist_ok=True)
+        key_path = os.path.join(key_dir, "slurm_private_key")
+        with open(key_path, "w") as f:
+            os.chmod(key_path, 0o600)
+            f.write(default_client_key)
+        client_keys.append(key_path)

		with open("./slurm_private_key", "w") as f:
		f.write(default_client_key)

	return Resource(phase=cur_phase, message=msg, outputs=resource_meta.outputs)
	outputs = resource_meta.outputs if resource_meta.outputs is not None else {}
	return Resource(phase=cur_phase, message=msg, outputs=outputs)

Slurm agent #3005

Are you sure you want to change the base?

Slurm agent #3005

Conversation

JiangJiaWei1103 commented Dec 16, 2024 • edited by flyte-bot 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

Summary by Bito

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

JiangJiaWei1103 commented Feb 14, 2025 • edited Loading

SlurmTask

SlurmShellTask

JiangJiaWei1103 commented Feb 14, 2025

flyte-bot commented Feb 14, 2025

Code Review Agent Run Status

flyte-bot commented Feb 16, 2025

Code Review Agent Run Status

flyte-bot commented Feb 17, 2025

Code Review Agent Run Status

flyte-bot commented Feb 18, 2025 • edited Loading

Code Review Agent Run #e8c545

flyte-bot commented Feb 18, 2025 • edited Loading

Changelist by Bito

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

flyte-bot commented Feb 18, 2025 • edited Loading

Code Review Agent Run #0e549b

Choose a reason for hiding this comment

JiangJiaWei1103 commented Dec 16, 2024 •

edited by flyte-bot

Loading

codecov bot commented Dec 19, 2024 •

edited

Loading

JiangJiaWei1103 commented Feb 14, 2025 •

edited

Loading

`SlurmTask`

`SlurmShellTask`

flyte-bot commented Feb 18, 2025 •

edited

Loading

flyte-bot commented Feb 18, 2025 •

edited

Loading

flyte-bot commented Feb 18, 2025 •

edited

Loading