flytekitplugins-slurm 1.15.1__tar.gz

flytekitplugins_slurm-1.15.1/PKG-INFO

@@ -0,0 +1,30 @@
+Metadata-Version: 2.2
+Name: flytekitplugins-slurm
+Version: 1.15.1
+Summary: This package holds the Slurm plugins for flytekit
+Author: flyteorg
+Author-email: admin@flyte.org
+License: apache2
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Requires-Dist: flytekit>=1.15.0
+Requires-Dist: flyteidl>=1.15.0
+Requires-Dist: asyncssh
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

flytekitplugins_slurm-1.15.1/README.md

@@ -0,0 +1,5 @@
+# Flytekit Slurm Plugin
+
+The Slurm agent is designed to integrate Flyte workflows with Slurm-managed high-performance computing (HPC) clusters, enabling users to leverage Slurm's capability of compute resource allocation, scheduling, and monitoring.
+
+This [guide](https://github.com/JiangJiaWei1103/flytekit/blob/slurm-agent-dev/plugins/flytekit-slurm/demo.md) provides a concise overview of the design philosophy behind the Slurm agent and explains how to set up a local environment for testing the agent.

flytekitplugins_slurm-1.15.1/flytekitplugins/slurm/__init__.py

@@ -0,0 +1,4 @@
+from .function.agent import SlurmFunctionAgent
+from .function.task import SlurmFunction, SlurmFunctionTask
+from .script.agent import SlurmScriptAgent
+from .script.task import Slurm, SlurmRemoteScript, SlurmShellTask, SlurmTask

flytekitplugins_slurm-1.15.1/flytekitplugins/slurm/function/agent.py

@@ -0,0 +1,190 @@
+import tempfile
+import uuid
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union
+
+from asyncssh import SSHClientConnection
+
+from flytekit import logger
+from flytekit.extend.backend.base_agent import AgentRegistry, AsyncAgentBase, Resource, ResourceMeta
+from flytekit.extend.backend.utils import convert_to_flyte_phase
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+
+from ..ssh_utils import ssh_connect
+
+
+@dataclass
+class SlurmJobMetadata(ResourceMeta):
+    """Slurm job metadata.
+
+    Args:
+        job_id: Slurm job id.
+        ssh_config: Options of SSH client connection. For available options, please refer to
+            the ssh_utils module.
+
+    Attributes:
+        job_id (str): Slurm job id.
+        ssh_config (Dict[str, Union[str, List[str], Tuple[str, ...]]]): SSH configuration options
+            for establishing client connections.
+    """
+
+    job_id: str
+    ssh_config: Dict[str, Union[str, List[str], Tuple[str, ...]]]
+
+
+@dataclass
+class SlurmCluster:
+    host: str
+    username: Optional[str] = None
+
+    def __hash__(self):
+        return hash((self.host, self.username))
+
+
+class SlurmFunctionAgent(AsyncAgentBase):
+    name = "Slurm Function Agent"
+
+    # SSH connection pool for multi-host environment
+    ssh_config_to_ssh_conn: Dict[SlurmCluster, SSHClientConnection] = {}
+
+    def __init__(self) -> None:
+        super(SlurmFunctionAgent, self).__init__(task_type_name="slurm_fn", metadata_type=SlurmJobMetadata)
+
+    async def create(
+        self,
+        task_template: TaskTemplate,
+        inputs: Optional[LiteralMap] = None,
+        **kwargs,
+    ) -> SlurmJobMetadata:
+        unique_script_name = f"/tmp/task_{uuid.uuid4().hex}.slurm"
+
+        # Retrieve task config
+        ssh_config = task_template.custom["ssh_config"]
+        sbatch_conf = task_template.custom["sbatch_conf"]
+        script = task_template.custom["script"]
+
+        # Construct command for Slurm cluster
+        cmd, script = _get_sbatch_cmd_and_script(
+            sbatch_conf=sbatch_conf,
+            entrypoint=" ".join(task_template.container.args),
+            script=script,
+            batch_script_path=unique_script_name,
+        )
+
+        # Run Slurm job
+        conn = await self._get_or_create_ssh_connection(ssh_config)
+        with tempfile.NamedTemporaryFile("w") as f:
+            f.write(script)
+            f.flush()
+            async with conn.start_sftp_client() as sftp:
+                await sftp.put(f.name, unique_script_name)
+        res = await conn.run(cmd, check=True)
+
+        # Retrieve Slurm job id
+        job_id = res.stdout.split()[-1]
+
+        return SlurmJobMetadata(job_id=job_id, ssh_config=ssh_config)
+
+    async def get(self, resource_meta: SlurmJobMetadata, **kwargs) -> Resource:
+        ssh_config = resource_meta.ssh_config
+        conn = await self._get_or_create_ssh_connection(ssh_config)
+        job_res = await conn.run(f"scontrol show job {resource_meta.job_id}", check=True)
+
+        # Determine the current flyte phase from Slurm job state
+        job_state = "running"
+        msg = "No stdout available"
+        for o in job_res.stdout.split(" "):
+            if "JobState" in o:
+                job_state = o.split("=")[1].strip().lower()
+            elif "StdOut" in o:
+                stdout_path = o.split("=")[1].strip()
+                msg_res = await conn.run(f"cat {stdout_path}", check=True)
+                msg = msg_res.stdout
+
+        cur_phase = convert_to_flyte_phase(job_state)
+
+        return Resource(phase=cur_phase, message=msg)
+
+    async def delete(self, resource_meta: SlurmJobMetadata, **kwargs) -> None:
+        conn = await self._get_or_create_ssh_connection(resource_meta.ssh_config)
+        _ = await conn.run(f"scancel {resource_meta.job_id}", check=True)
+
+    async def _get_or_create_ssh_connection(
+        self, ssh_config: Dict[str, Union[str, List[str], Tuple[str, ...]]]
+    ) -> SSHClientConnection:
+        """Get an existing SSH connection or create a new one if needed.
+
+        Args:
+            ssh_config (Dict[str, Union[str, List[str], Tuple[str, ...]]]): SSH configuration dictionary.
+
+        Returns:
+            SSHClientConnection: An active SSH connection, either pre-existing or newly established.
+        """
+        host = ssh_config.get("host")
+        username = ssh_config.get("username")
+
+        ssh_cluster_config = SlurmCluster(host=host, username=username)
+        if self.ssh_config_to_ssh_conn.get(ssh_cluster_config) is None:
+            logger.info("ssh connection key not found, creating new connection")
+            conn = await ssh_connect(ssh_config=ssh_config)
+            self.ssh_config_to_ssh_conn[ssh_cluster_config] = conn
+        else:
+            conn = self.ssh_config_to_ssh_conn[ssh_cluster_config]
+            try:
+                await conn.run("echo [TEST] SSH connection", check=True)
+                logger.info("re-using new connection")
+            except Exception as e:
+                logger.info(f"Re-establishing SSH connection due to error: {e}")
+                conn = await ssh_connect(ssh_config=ssh_config)
+                self.ssh_config_to_ssh_conn[ssh_cluster_config] = conn
+
+        return conn
+
+
+def _get_sbatch_cmd_and_script(
+    sbatch_conf: Dict[str, str],
+    entrypoint: str,
+    script: Optional[str] = None,
+    batch_script_path: str = "/tmp/task.slurm",
+) -> Tuple[str, str]:
+    """Construct the Slurm sbatch command and the batch script content.
+
+    Flyte entrypoint, pyflyte-execute, is run within a bash shell process.
+
+    Args:
+        sbatch_conf (Dict[str, str]): Options of sbatch command.
+        entrypoint (str): Flyte entrypoint.
+        script (Optional[str], optional): User-defined script where "{task.fn}" serves as a placeholder for the
+            task function execution. Users should insert "{task.fn}" at the desired
+            execution point within the script. If the script is not provided, the task
+            function will be executed directly. Defaults to None.
+        batch_script_path (str, optional): Absolute path of the batch script on Slurm cluster.
+            Defaults to "/tmp/task.slurm".
+
+    Returns:
+        Tuple[str, str]: A tuple containing:
+            - cmd: Slurm sbatch command
+            - script: The batch script content
+    """
+    # Setup sbatch options
+    cmd = ["sbatch"]
+    for opt, val in sbatch_conf.items():
+        cmd.extend([f"--{opt}", str(val)])
+
+    # Assign the batch script to run
+    cmd.append(batch_script_path)
+
+    if script is None:
+        script = f"""#!/bin/bash -i
+        {entrypoint}
+        """
+    else:
+        script = script.replace("{task.fn}", entrypoint)
+
+    cmd = " ".join(cmd)
+
+    return cmd, script
+
+
+AgentRegistry.register(SlurmFunctionAgent())

flytekitplugins_slurm-1.15.1/flytekitplugins/slurm/function/task.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+from flytekit import FlyteContextManager, PythonFunctionTask
+from flytekit.configuration import SerializationSettings
+from flytekit.extend import TaskPlugins
+from flytekit.extend.backend.base_agent import AsyncAgentExecutorMixin
+from flytekit.image_spec import ImageSpec
+
+
+@dataclass
+class SlurmFunction(object):
+    """Configure Slurm settings. Note that we focus on sbatch command now.
+
+    Args:
+        ssh_config: Options of SSH client connection. For available options, please refer to
+            <newly-added-ssh-utils-file>
+        sbatch_conf: Options of sbatch command. If not provided, defaults to an empty dict.
+        script: User-defined script where "{task.fn}" serves as a placeholder for the
+            task function execution. Users should insert "{task.fn}" at the desired
+            execution point within the script. If the script is not provided, the task
+            function will be executed directly.
+
+    Attributes:
+        ssh_config (Dict[str, Union[str, List[str], Tuple[str, ...]]]): SSH client configuration options.
+        sbatch_conf (Optional[Dict[str, str]]): Slurm sbatch command options.
+        script (Optional[str]): Custom script template for task execution.
+    """
+
+    ssh_config: Dict[str, Union[str, List[str], Tuple[str, ...]]]
+    sbatch_conf: Optional[Dict[str, str]] = None
+    script: Optional[str] = None
+
+    def __post_init__(self):
+        assert self.ssh_config["host"] is not None, "'host' must be specified in ssh_config."
+        if self.sbatch_conf is None:
+            self.sbatch_conf = {}
+
+
+class SlurmFunctionTask(AsyncAgentExecutorMixin, PythonFunctionTask[SlurmFunction]):
+    """
+    Actual Plugin that transforms the local python code for execution within a slurm context...
+    """
+
+    _TASK_TYPE = "slurm_fn"
+
+    def __init__(
+        self,
+        task_config: SlurmFunction,
+        task_function: Callable,
+        container_image: Optional[Union[str, ImageSpec]] = None,
+        **kwargs,
+    ):
+        super(SlurmFunctionTask, self).__init__(
+            task_config=task_config,
+            task_type=self._TASK_TYPE,
+            task_function=task_function,
+            container_image=container_image,
+            **kwargs,
+        )
+
+    def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
+        return {
+            "ssh_config": self.task_config.ssh_config,
+            "sbatch_conf": self.task_config.sbatch_conf,
+            "script": self.task_config.script,
+        }
+
+    def execute(self, **kwargs) -> Any:
+        ctx = FlyteContextManager.current_context()
+        if ctx.execution_state and ctx.execution_state.is_local_execution():
+            # Mimic the propeller's behavior in local agent test
+            return AsyncAgentExecutorMixin.execute(self, **kwargs)
+        else:
+            # Execute the task with a direct python function call
+            return PythonFunctionTask.execute(self, **kwargs)
+
+
+TaskPlugins.register_pythontask_plugin(SlurmFunction, SlurmFunctionTask)

flytekitplugins_slurm-1.15.1/flytekitplugins/slurm/script/agent.py

@@ -0,0 +1,191 @@
+import tempfile
+import uuid
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Type
+
+from asyncssh import SSHClientConnection
+
+import flytekit
+from flytekit.core.type_engine import TypeEngine
+from flytekit.extend.backend.base_agent import AgentRegistry, AsyncAgentBase, Resource, ResourceMeta
+from flytekit.extend.backend.utils import convert_to_flyte_phase
+from flytekit.extras.tasks.shell import OutputLocation, _PythonFStringInterpolizer
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+
+from ..ssh_utils import SlurmCluster, get_ssh_conn
+
+
+@dataclass
+class SlurmJobMetadata(ResourceMeta):
+    """
+    Slurm job metadata.
+
+    Attributes:
+        job_id (str): Slurm job id.
+        ssh_config (Dict[str, Any]): SSH connection configuration options.
+        outputs (Dict[str, str]): Mapping from the output variable name to the output location.
+    """
+
+    job_id: str
+    ssh_config: Dict[str, Any]
+    outputs: Dict[str, str]
+
+
+class SlurmScriptAgent(AsyncAgentBase):
+    name = "Slurm Script Agent"
+
+    # SSH connection pool for multi-host environment
+    slurm_cluster_to_ssh_conn: Dict[SlurmCluster, SSHClientConnection] = {}
+
+    # Dummy script content
+    DUMMY_SCRIPT = "#!/bin/bash"
+
+    def __init__(self) -> None:
+        super(SlurmScriptAgent, self).__init__(task_type_name="slurm", metadata_type=SlurmJobMetadata)
+
+    async def create(
+        self,
+        task_template: TaskTemplate,
+        inputs: Optional[LiteralMap] = None,
+        **kwargs,
+    ) -> SlurmJobMetadata:
+        uniq_script_path = f"/tmp/task_{uuid.uuid4().hex}.slurm"
+        outputs = {}
+
+        # Retrieve task config
+        ssh_config = task_template.custom["ssh_config"]
+        batch_script_args = task_template.custom["batch_script_args"]
+        sbatch_conf = task_template.custom["sbatch_conf"]
+
+        # Construct sbatch command for Slurm cluster
+        upload_script = False
+        if "script" in task_template.custom:
+            script = task_template.custom["script"]
+            assert script != self.DUMMY_SCRIPT, "Please write the user-defined batch script content."
+            script, outputs = self._interpolate_script(
+                script,
+                input_literal_map=inputs,
+                python_input_types=task_template.custom["python_input_types"],
+                output_locs=task_template.custom["output_locs"],
+            )
+
+            batch_script_path = uniq_script_path
+            upload_script = True
+        else:
+            # Assume the batch script is already on Slurm
+            batch_script_path = task_template.custom["batch_script_path"]
+        cmd = _get_sbatch_cmd(
+            sbatch_conf=sbatch_conf, batch_script_path=batch_script_path, batch_script_args=batch_script_args
+        )
+
+        # Run Slurm job
+        conn = await get_ssh_conn(ssh_config=ssh_config, slurm_cluster_to_ssh_conn=self.slurm_cluster_to_ssh_conn)
+        if upload_script:
+            with tempfile.NamedTemporaryFile("w") as f:
+                f.write(script)
+                f.flush()
+                async with conn.start_sftp_client() as sftp:
+                    await sftp.put(f.name, batch_script_path)
+        res = await conn.run(cmd, check=True)
+
+        # Retrieve Slurm job id
+        job_id = res.stdout.split()[-1]
+
+        return SlurmJobMetadata(job_id=job_id, ssh_config=ssh_config, outputs=outputs)
+
+    async def get(self, resource_meta: SlurmJobMetadata, **kwargs) -> Resource:
+        conn = await get_ssh_conn(
+            ssh_config=resource_meta.ssh_config, slurm_cluster_to_ssh_conn=self.slurm_cluster_to_ssh_conn
+        )
+        job_res = await conn.run(f"scontrol show job {resource_meta.job_id}", check=True)
+
+        # Determine the current flyte phase from Slurm job state
+        msg = ""
+        job_state = "running"
+        for o in job_res.stdout.split(" "):
+            if "JobState" in o:
+                job_state = o.split("=")[1].strip().lower()
+            elif "StdOut" in o:
+                stdout_path = o.split("=")[1].strip()
+                msg_res = await conn.run(f"cat {stdout_path}", check=True)
+                msg = msg_res.stdout
+        cur_phase = convert_to_flyte_phase(job_state)
+
+        return Resource(phase=cur_phase, message=msg, outputs=resource_meta.outputs)
+
+    async def delete(self, resource_meta: SlurmJobMetadata, **kwargs) -> None:
+        conn = await get_ssh_conn(
+            ssh_config=resource_meta.ssh_config, slurm_cluster_to_ssh_conn=self.slurm_cluster_to_ssh_conn
+        )
+        _ = await conn.run(f"scancel {resource_meta.job_id}", check=True)
+
+    def _interpolate_script(
+        self,
+        script: str,
+        input_literal_map: Optional[LiteralMap] = None,
+        python_input_types: Optional[Dict[str, Type]] = None,
+        output_locs: Optional[List[OutputLocation]] = None,
+    ) -> Tuple[str, Dict[str, str]]:
+        """
+        Interpolate the user-defined script with the specified input and output arguments.
+
+        Args:
+            script (str): The user-defined script with placeholders for dynamic input/output.
+            input_literal_map (Optional[LiteralMap]): The Flyte LiteralMap of inputs.
+            python_input_types (Optional[Dict[str, Type]]): Mapping of input names to their Python/typing types.
+            output_locs (Optional[List[OutputLocation]]): List of output locations to be interpolated.
+
+        Returns:
+            Tuple[str, Dict[str, str]]:
+                - A two-element tuple in which the first element is the interpolated script (str),
+                and the second is a dictionary mapping each output variable name to its final location (str).
+        """
+        input_kwargs = TypeEngine.literal_map_to_kwargs(
+            flytekit.current_context(),
+            lm=input_literal_map,
+            python_types={} if python_input_types is None else python_input_types,
+        )
+        interpolizer = _PythonFStringInterpolizer()
+
+        # Interpolate output locations with input values
+        outputs = {}
+        if output_locs is not None:
+            for oloc in output_locs:
+                outputs[oloc.var] = interpolizer.interpolate(oloc.location, inputs=input_kwargs)
+
+        # Interpolate the script
+        script = interpolizer.interpolate(script, inputs=input_kwargs, outputs=outputs)
+
+        return script, outputs
+
+
+def _get_sbatch_cmd(sbatch_conf: Dict[str, str], batch_script_path: str, batch_script_args: List[str] = None) -> str:
+    """
+    Construct the Slurm sbatch command.
+
+    Args:
+        sbatch_conf (Dict[str, str]): Slurm sbatch configuration options (e.g., partition, job-name, etc.).
+        batch_script_path (str): Absolute path on the Slurm cluster of the script to run.
+        batch_script_args (List[str], optional): Additional arguments to pass to the batch script.
+
+    Returns:
+        str: The sbatch command string that can be executed on the Slurm cluster.
+    """
+    cmd = ["sbatch"]
+    for opt, val in sbatch_conf.items():
+        cmd.extend([f"--{opt}", str(val)])
+
+    # Assign the batch script to run
+    cmd.append(batch_script_path)
+
+    # Add args if present
+    if batch_script_args:
+        for arg in batch_script_args:
+            cmd.append(arg)
+
+    cmd = " ".join(cmd)
+    return cmd
+
+
+AgentRegistry.register(SlurmScriptAgent())

flytekitplugins_slurm-1.15.1/flytekitplugins/slurm/script/task.py

@@ -0,0 +1,136 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Type
+
+from flytekit.configuration import SerializationSettings
+from flytekit.core.base_task import PythonTask
+from flytekit.core.interface import Interface
+from flytekit.extend import TaskPlugins
+from flytekit.extend.backend.base_agent import AsyncAgentExecutorMixin
+from flytekit.extras.tasks.shell import OutputLocation
+from flytekit.types.directory import FlyteDirectory
+from flytekit.types.file import FlyteFile
+
+
+@dataclass
+class Slurm(object):
+    """
+    Configure Slurm settings. Note that we focus on sbatch command now.
+
+    Compared with spark, please refer to:
+    https://api-docs.databricks.com/python/pyspark/latest/api/pyspark.SparkContext.html.
+
+    Attributes:
+        ssh_config (Dict[str, Any]): Options of SSH client connection. For available options, please refer to
+            <newly-added-ssh-utils-file>.
+        sbatch_conf (Optional[Dict[str, str]]): Options of sbatch command. For available options, please refer to
+            https://slurm.schedmd.com/sbatch.html.
+        batch_script_args (Optional[List[str]]): Additional args for the batch script on Slurm cluster.
+    """
+
+    ssh_config: Dict[str, Any]
+    sbatch_conf: Optional[Dict[str, str]] = None
+    batch_script_args: Optional[List[str]] = None
+
+    def __post_init__(self):
+        if self.sbatch_conf is None:
+            self.sbatch_conf = {}
+
+
+@dataclass
+class SlurmRemoteScript(Slurm):
+    """Encounter collision if Slurm is shared btw SlurmTask and SlurmShellTask."""
+
+    batch_script_path: str = field(default=None)
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.batch_script_path is None:
+            raise ValueError("batch_script_path must be provided")
+
+
+class SlurmTask(AsyncAgentExecutorMixin, PythonTask[SlurmRemoteScript]):
+    _TASK_TYPE = "slurm"
+
+    def __init__(
+        self,
+        name: str,
+        task_config: SlurmRemoteScript,
+        **kwargs,
+    ):
+        super(SlurmTask, self).__init__(
+            task_type=self._TASK_TYPE,
+            name=name,
+            task_config=task_config,
+            # Dummy interface, will support this after discussion
+            interface=Interface(None, None),
+            **kwargs,
+        )
+
+    def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
+        return {
+            "ssh_config": self.task_config.ssh_config,
+            "batch_script_path": self.task_config.batch_script_path,
+            "batch_script_args": self.task_config.batch_script_args,
+            "sbatch_conf": self.task_config.sbatch_conf,
+        }
+
+
+class SlurmShellTask(AsyncAgentExecutorMixin, PythonTask[Slurm]):
+    _TASK_TYPE = "slurm"
+
+    def __init__(
+        self,
+        name: str,
+        task_config: Slurm,
+        script: str,
+        inputs: Optional[Dict[str, Type]] = None,
+        output_locs: Optional[List[OutputLocation]] = None,
+        **kwargs,
+    ):
+        self._inputs = inputs
+        self._output_locs = output_locs if output_locs is not None else []
+        self._script = script
+
+        outputs = self._validate_output_locs()
+
+        super().__init__(
+            name=name,
+            task_type=self._TASK_TYPE,
+            task_config=task_config,
+            interface=Interface(inputs=inputs, outputs=outputs),
+            **kwargs,
+        )
+
+    def _validate_output_locs(self) -> Dict[str, Type]:
+        outputs = {}
+        for v in self._output_locs:
+            if v is None:
+                raise ValueError("OutputLocation cannot be none")
+            if not isinstance(v, OutputLocation):
+                raise ValueError("Every output type should be an output location on the file-system")
+            if v.location is None:
+                raise ValueError(f"Output Location not provided for output var {v.var}")
+            if not issubclass(v.var_type, FlyteFile) and not issubclass(v.var_type, FlyteDirectory):
+                raise ValueError(
+                    "Currently only outputs of type FlyteFile/FlyteDirectory and their derived types are supported"
+                )
+            outputs[v.var] = v.var_type
+        return outputs
+
+    def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
+        return {
+            "ssh_config": self.task_config.ssh_config,
+            "batch_script_args": self.task_config.batch_script_args,
+            "sbatch_conf": self.task_config.sbatch_conf,
+            "script": self._script,
+            "python_input_types": self._inputs,
+            "output_locs": self._output_locs,
+        }
+
+    @property
+    def script(self) -> str:
+        return self._script
+
+
+TaskPlugins.register_pythontask_plugin(SlurmRemoteScript, SlurmTask)
+TaskPlugins.register_pythontask_plugin(Slurm, SlurmShellTask)

flytekitplugins_slurm-1.15.1/flytekitplugins/slurm/ssh_utils.py

@@ -0,0 +1,190 @@
+"""
+Utilities of asyncssh connections.
+"""
+
+import os
+import sys
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
+
+import asyncssh
+from asyncssh import SSHClientConnection
+
+from flytekit import logger
+from flytekit.extend.backend.utils import get_agent_secret
+
+T = TypeVar("T", bound="SSHConfig")
+SLURM_PRIVATE_KEY = "FLYTE_SLURM_PRIVATE_KEY"
+
+
+@dataclass
+class SlurmCluster:
+    """A Slurm cluster instance is defined by a pair of (Slurm host, username).
+
+    Attributes:
+        host (str): The hostname or address to connect to.
+        username (Optional[str]): The username to authenticate as on the server.
+    """
+
+    host: str
+    username: Optional[str] = None
+
+    def __hash__(self):
+        return hash((self.host, self.username))
+
+
+@dataclass(frozen=True)
+class SSHConfig:
+    """A customized version of SSHClientConnectionOptions, tailored to specific needs.
+
+    This config is based on the official SSHClientConnectionOptions but includes
+    only a subset of options, with some fields adjusted to be optional or required.
+    For the official options, please refer to:
+    https://asyncssh.readthedocs.io/en/latest/api.html#asyncssh.SSHClientConnectionOptions
+
+    Attributes:
+        host (str): The hostname or address to connect to.
+        username (Optional[str]): The username to authenticate as on the server.
+        client_keys (Union[str, List[str], Tuple[str, ...]]): File paths to private keys which will be used to authenticate the
+            client via public key authentication. The default value is an empty tuple since
+            client public key authentication is mandatory.
+    """
+
+    host: str
+    username: Optional[str] = None
+    client_keys: Union[str, List[str], Tuple[str, ...]] = ()
+
+    @classmethod
+    def from_dict(cls: Type[T], ssh_config: Dict[str, Any]) -> T:
+        return cls(**ssh_config)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+    def __eq__(self, other):
+        if not isinstance(other, SSHConfig):
+            return False
+        return self.host == other.host and self.username == other.username and self.client_keys == other.client_keys
+
+
+async def ssh_connect(ssh_config: Dict[str, Any]) -> SSHClientConnection:
+    """Make an SSH client connection.
+
+    Args:
+        ssh_config (Dict[str, Any]): Options of SSH client connection defined in SSHConfig.
+
+    Returns:
+        SSHClientConnection: An SSH client connection object.
+
+    Raises:
+        ValueError: If both FLYTE_SLURM_PRIVATE_KEY secret and ssh_config['private_key'] are missing.
+    """
+    # Validate ssh_config
+    ssh_config = SSHConfig.from_dict(ssh_config).to_dict()
+    # This is required to avoid the error "asyncssh.misc.HostKeyNotVerifiable" when connecting to a new host.
+    ssh_config["known_hosts"] = None
+
+    # Make the first SSH connection using either OpenSSH client config files or
+    # a user-defined private key. If using OpenSSH config, it will attempt to
+    # load settings from ~/.ssh/config.
+    try:
+        conn = await asyncssh.connect(**ssh_config)
+        return conn
+    except Exception as e:
+        logger.info(
+            "Failed to make an SSH connection using the default OpenSSH client config (~/.ssh/config) or "
+            f"the provided private keys. Error details:\n{e}"
+        )
+
+    try:
+        default_client_key = get_agent_secret(secret_key=SLURM_PRIVATE_KEY)
+    except ValueError:
+        logger.info("The secret for key FLYTE_SLURM_PRIVATE_KEY is not set.")
+        default_client_key = None
+
+    if default_client_key is None and ssh_config.get("client_keys") == ():
+        raise ValueError(
+            "Both the secret for key FLYTE_SLURM_PRIVATE_KEY and ssh_config['private_key'] are missing. "
+            "At least one must be set."
+        )
+
+    client_keys = []
+    if default_client_key is not None:
+        # Write the private key to a local path
+        # This may not be a good practice...
+        private_key_path = os.path.abspath("./slurm_private_key")
+        with open(private_key_path, "w") as f:
+            f.write(default_client_key)
+        client_keys.append(private_key_path)
+
+    user_client_keys = ssh_config.get("client_keys")
+    if user_client_keys is not None:
+        client_keys.extend([user_client_keys] if isinstance(user_client_keys, str) else user_client_keys)
+
+    ssh_config["client_keys"] = client_keys
+    logger.info(f"Updated SSH config: {ssh_config}")
+    try:
+        conn = await asyncssh.connect(**ssh_config)
+        return conn
+    except Exception as e:
+        logger.info(
+            "Failed to make an SSH connection using the provided private keys. Please verify your setup."
+            f"Error details:\n{e}"
+        )
+        sys.exit(1)
+
+
+async def get_ssh_conn(
+    ssh_config: Dict[str, Union[str, List[str], Tuple[str, ...]]],
+    slurm_cluster_to_ssh_conn: Dict[SlurmCluster, SSHClientConnection],
+) -> Tuple[SlurmCluster, SSHClientConnection]:
+    """
+    Get an existing SSH connection or create a new one if needed.
+
+    Args:
+        ssh_config (Dict[str, Union[str, List[str], Tuple[str, ...]]]):
+            SSH configuration dictionary, including host and username.
+        slurm_cluster_to_ssh_conn (Dict[SlurmCluster, SSHClientConnection]):
+            A mapping of SlurmCluster to existing SSHClientConnection objects.
+
+    Returns:
+        Tuple[SlurmCluster, SSHClientConnection]:
+            A tuple containing (SlurmCluster, SSHClientConnection). If no connection
+            for the given SlurmCluster exists, a new one is created and cached.
+    """
+
+    # (Optional) normal code comment instead of docstring line:
+    # Is it necessary to ensure immutability in this function?
+
+    host = ssh_config.get("host")
+    username = ssh_config.get("username")
+    slurm_cluster = SlurmCluster(host=host, username=username)
+
+    if slurm_cluster_to_ssh_conn.get(slurm_cluster) is None:
+        logger.info("SSH connection key not found, creating new connection")
+        conn = await ssh_connect(ssh_config=ssh_config)
+        slurm_cluster_to_ssh_conn[slurm_cluster] = conn
+    else:
+        conn = slurm_cluster_to_ssh_conn[slurm_cluster]
+        try:
+            await conn.run("echo [TEST] SSH connection", check=True)
+            logger.info("Re-using new connection")
+        except Exception as e:
+            logger.info(f"Re-establishing SSH connection due to error: {e}")
+            conn = await ssh_connect(ssh_config=ssh_config)
+            slurm_cluster_to_ssh_conn[slurm_cluster] = conn
+
+    return conn
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    async def test_connect():
+        conn = await ssh_connect({"host": "aws2", "username": "ubuntu"})
+        res = await conn.run("echo [TEST] SSH connection", check=True)
+        out = res.stdout
+
+        return out
+
+    logger.info(asyncio.run(test_connect()))

flytekitplugins_slurm-1.15.1/flytekitplugins_slurm.egg-info/PKG-INFO

@@ -0,0 +1,30 @@
+Metadata-Version: 2.2
+Name: flytekitplugins-slurm
+Version: 1.15.1
+Summary: This package holds the Slurm plugins for flytekit
+Author: flyteorg
+Author-email: admin@flyte.org
+License: apache2
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Requires-Dist: flytekit>=1.15.0
+Requires-Dist: flyteidl>=1.15.0
+Requires-Dist: asyncssh
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

flytekitplugins_slurm-1.15.1/flytekitplugins_slurm.egg-info/SOURCES.txt

@@ -0,0 +1,17 @@
+README.md
+setup.py
+flytekitplugins/slurm/__init__.py
+flytekitplugins/slurm/ssh_utils.py
+flytekitplugins/slurm/function/agent.py
+flytekitplugins/slurm/function/task.py
+flytekitplugins/slurm/script/agent.py
+flytekitplugins/slurm/script/task.py
+flytekitplugins_slurm.egg-info/PKG-INFO
+flytekitplugins_slurm.egg-info/SOURCES.txt
+flytekitplugins_slurm.egg-info/dependency_links.txt
+flytekitplugins_slurm.egg-info/entry_points.txt
+flytekitplugins_slurm.egg-info/namespace_packages.txt
+flytekitplugins_slurm.egg-info/requires.txt
+flytekitplugins_slurm.egg-info/top_level.txt
+tests/test_slurm_fn_task.py
+tests/test_slurm_shell_task.py

flytekitplugins_slurm-1.15.1/flytekitplugins_slurm.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

flytekitplugins_slurm-1.15.1/flytekitplugins_slurm.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[flytekit.plugins]
+slurm = flytekitplugins.slurm

flytekitplugins_slurm-1.15.1/flytekitplugins_slurm.egg-info/namespace_packages.txt

@@ -0,0 +1 @@
+flytekitplugins

flytekitplugins_slurm-1.15.1/flytekitplugins_slurm.egg-info/requires.txt

@@ -0,0 +1,3 @@
+flytekit>=1.15.0
+flyteidl>=1.15.0
+asyncssh

flytekitplugins_slurm-1.15.1/flytekitplugins_slurm.egg-info/top_level.txt

@@ -0,0 +1 @@
+flytekitplugins

flytekitplugins_slurm-1.15.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

flytekitplugins_slurm-1.15.1/setup.py

@@ -0,0 +1,41 @@
+from setuptools import setup
+
+PLUGIN_NAME = "slurm"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit>=1.15.0", "flyteidl>=1.15.0", "asyncssh"]
+
+__version__ = "1.15.1"
+
+setup(
+    name=microlib_name,
+    version=__version__,
+    author="flyteorg",
+    author_email="admin@flyte.org",
+    description="This package holds the Slurm plugins for flytekit",
+    namespace_packages=["flytekitplugins"],
+    packages=[
+        f"flytekitplugins.{PLUGIN_NAME}",
+        f"flytekitplugins.{PLUGIN_NAME}.script",
+        f"flytekitplugins.{PLUGIN_NAME}.function",
+    ],
+    install_requires=plugin_requires,
+    license="apache2",
+    python_requires=">=3.9",
+    classifiers=[
+        "Intended Audience :: Science/Research",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Software Development",
+        "Topic :: Software Development :: Libraries",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+    entry_points={"flytekit.plugins": [f"{PLUGIN_NAME}=flytekitplugins.{PLUGIN_NAME}"]},
+)

flytekitplugins_slurm-1.15.1/tests/test_slurm_fn_task.py

@@ -0,0 +1,55 @@
+import os.path
+from unittest import mock
+from flytekit.core import context_manager
+import flytekit
+from flytekit import StructuredDataset, StructuredDatasetTransformerEngine, task, ImageSpec
+from flytekit.configuration import Image, ImageConfig, SerializationSettings, FastSerializationSettings, DefaultImages
+from flytekit.core.context_manager import ExecutionParameters, FlyteContextManager, ExecutionState
+from flytekitplugins.slurm import SlurmFunction
+
+
+def test_slurm_task():
+    script_file = """#!/bin/bash -i
+
+    echo Run function with sbatch...
+
+    # Run the user-defined task function
+    {task.fn}
+    """
+
+    @task(
+    # container_image=image,
+    task_config=SlurmFunction(
+        ssh_config={
+            "host": "your-slurm-host",
+            "username": "ubuntu",
+        },
+        sbatch_conf={
+            "partition": "debug",
+            "job-name": "tiny-slurm",
+            "output": "/home/ubuntu/fn_task.log"
+        },
+        script=script_file
+        )
+    )
+    def plus_one(x: int) -> int:
+        return x + 1
+
+    assert plus_one.task_config is not None
+    assert plus_one.task_config.ssh_config == {"host": "your-slurm-host", "username": "ubuntu"}
+    assert plus_one.task_config.sbatch_conf == {"partition": "debug", "job-name": "tiny-slurm", "output": "/home/ubuntu/fn_task.log"}
+    assert plus_one.task_config.script == script_file
+
+    default_img = Image(name="default", fqn="test", tag="tag")
+    settings = SerializationSettings(
+        project="project",
+        domain="domain",
+        version="version",
+        env={"FOO": "baz"},
+        image_config=ImageConfig(default_image=default_img, images=[default_img]),
+    )
+
+    retrieved_settings = plus_one.get_custom(settings)
+    assert retrieved_settings["ssh_config"] == {"host": "your-slurm-host", "username": "ubuntu"}
+    assert retrieved_settings["sbatch_conf"] == {"partition": "debug", "job-name": "tiny-slurm", "output": "/home/ubuntu/fn_task.log"}
+    assert retrieved_settings["script"] == script_file

flytekitplugins_slurm-1.15.1/tests/test_slurm_shell_task.py

@@ -0,0 +1,77 @@
+from flytekit.configuration import Image, ImageConfig, SerializationSettings
+from flytekitplugins.slurm import Slurm, SlurmRemoteScript, SlurmTask, SlurmShellTask
+
+
+def test_slurm_shell_task():
+    """Test SlurmTask and SlurmShellTask."""
+    batch_script_path = "/script/path/on/the/slurm/cluster"
+    script = """#!/bin/bash -i
+
+    echo "Run a Flyte SlurmShellTask...\n"
+    """
+
+    # SlurmTask
+    slurm_task = SlurmTask(
+        name="test-slurm-task",
+        task_config=SlurmRemoteScript(
+            ssh_config={
+                "host": "<your-slurm-host>",
+                "username": "ubuntu",
+            },
+            sbatch_conf={
+                "partition": "debug",
+                "job-name": "tiny-slurm",
+                "output": "/home/ubuntu/slurm_task.log"
+            },
+            batch_script_path=batch_script_path
+        )
+    )
+
+    assert slurm_task.task_config is not None
+    assert slurm_task.task_config.ssh_config == {"host": "<your-slurm-host>", "username": "ubuntu"}
+    assert slurm_task.task_config.sbatch_conf == {"partition": "debug", "job-name": "tiny-slurm", "output": "/home/ubuntu/slurm_task.log"}
+    assert slurm_task.task_config.batch_script_path == batch_script_path
+
+    # SlurmShellTask
+    slurm_shell_task = SlurmShellTask(
+        name="test-slurm-shell-task",
+        script=script,
+        task_config=Slurm(
+            ssh_config={
+                "host": "<your-slurm-host>",
+                "username": "ubuntu",
+            },
+            sbatch_conf={
+                "partition": "debug",
+                "job-name": "tiny-slurm",
+                "output": "/home/ubuntu/slurm_shell_task.log"
+            }
+        )
+    )
+
+    assert slurm_shell_task.task_config is not None
+    assert slurm_shell_task.task_config.ssh_config == {"host": "<your-slurm-host>", "username": "ubuntu"}
+    assert slurm_shell_task.task_config.sbatch_conf == {"partition": "debug", "job-name": "tiny-slurm", "output": "/home/ubuntu/slurm_shell_task.log"}
+    assert slurm_shell_task.script == script
+
+    # Define dummy SerializationSettings
+    default_img = Image(name="default", fqn="test", tag="tag")
+    settings = SerializationSettings(
+        project="project",
+        domain="domain",
+        version="version",
+        env={"FOO": "baz"},
+        image_config=ImageConfig(default_image=default_img, images=[default_img])
+    )
+
+    custom = slurm_task.get_custom(settings)
+    assert custom["ssh_config"] == {"host": "<your-slurm-host>", "username": "ubuntu"}
+    assert custom["sbatch_conf"] == {"partition": "debug", "job-name": "tiny-slurm", "output": "/home/ubuntu/slurm_task.log"}
+    assert custom["batch_script_path"] == batch_script_path
+    assert custom["batch_script_args"] is None
+
+    shell_custom = slurm_shell_task.get_custom(settings)
+    assert shell_custom["ssh_config"] == {"host": "<your-slurm-host>", "username": "ubuntu"}
+    assert shell_custom["sbatch_conf"] == {"partition": "debug", "job-name": "tiny-slurm", "output": "/home/ubuntu/slurm_shell_task.log"}
+    assert shell_custom["script"] == script
+    assert shell_custom["batch_script_args"] is None