Hi, @freezing-airport-6809 You mentioned that you want the slurm python function task to be like this.

@task(
    task_config=SlurmFunction(
        slurm_host="aws",
        srun_conf={
            "partition": "debug",
            "job-name": "tiny-slurm",
        },
        script="""
#!/bin/bash
# Pre-execute
echo "Hello, world!"
export MY_ENV_VAR=123

# Run the python function here
{flyte.fn}

# Post-execute
exit -1
"""
    )
)
def plus_one(x: int) -> int: 
    print(os.getenv("MY_ENV_VAR"))
    return x + 1

so the execution order will be (1) script (pre-execute) (2) function I am wondering about in agent's operation framework (create, get, and delete) the pre-execute part should be put in

create

operation, right? or every operation should be put in

get

operation? in agent's each operation we have a timeout mechanism, so we need to figure out a way to implement this. TLDR: the agent framework is for 1 execution, but in your propose, this will be more than 1 execution, which is still possible to implement it but ugly

@damp-lion-88352 yes, but it should just be one script that is run. it should not be 3 different calls, only one call

We've supported

SlurmFunctionTask

with an optional user-defined

script

, following shows the corresponding commit: https://github.com/flyteorg/flytekit/pull/3005/commits/9985305c1f5fe05e9dcab33f4116e0ef5bca21a1 This is an example we used to test it:

import os

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


@task(
    task_config=SlurmFunction(
        slurm_host="aws",
        srun_conf={
            "partition": "debug",
            "job-name": "fn-task",
            "output": "/home/ubuntu/fn_task.log"
        },
        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-poetry-RLi6T71_-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


@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://my-flyte-slurm-agent>", path, "wf", "--x", 2024])
    print(result.output)

The example can be accessed here.

hi

No problem!