[Docs] Add the Slurm agent example usage

docs/integrations/index.md

@@ -116,6 +116,8 @@ Native backend plugins can be executed without any external service dependencies
   - Execute tasks on PERIAN Job Platform.
 * - {doc}`Sensor agent </auto_examples/sensor/index>`
   - Run sensor jobs in your workflows with the sensor agent.
+* - {doc}`Slurm agent </auto_examples/slurm_agent/index>`
+  - Run Slurm jobs in your workflows with the Slurm agent.
 * - {doc}`Snowflake agent </auto_examples/snowflake_agent/index>`
   - Run Snowflake jobs in your workflows with the Snowflake agent.
 ```
@@ -245,6 +247,7 @@ Memory Machine Cloud agent </auto_examples/mmcloud_agent/index>
 OpenAI batch agent </auto_examples/openai_batch_agent/index>
 PERIAN Job Platform agent </auto_examples/perian_agent/index>
 Sensor agent </auto_examples/sensor/index>
+Slurm agent </auto_examples/slurm_agent/index>
 Snowflake agent </auto_examples/snowflake_agent/index>
 ```
 

examples/memray_plugin/memray_plugin/memray_example.py

@@ -5,10 +5,11 @@
 # Memray tracks and reports memory allocations, both in python code and in compiled extension modules.
 # This Memray Profiling plugin enables memory tracking on the Flyte task level and renders a memgraph profiling graph on Flyte Deck.
 # %%
-from flytekit import workflow, task, ImageSpec
-from flytekitplugins.memray import memray_profiling
 import time
 
+from flytekit import ImageSpec, task, workflow
+from flytekitplugins.memray import memray_profiling
+
 # %% [markdown]
 # First, we use `ImageSpec` to construct a container that contains the dependencies for the
 # tasks, we want to profile:

examples/slurm_agent/Dockerfile

@@ -0,0 +1,23 @@
+# ######################
+# NOTE: For CI/CD only #
+########################
+FROM python:3.11-slim-buster
+LABEL org.opencontainers.image.source=https://github.com/flyteorg/flytesnacks
+
+WORKDIR /root
+ENV VENV /opt/venv
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8
+ENV PYTHONPATH /root
+
+# Install Python dependencies
+COPY requirements.in /root
+RUN pip install -r /root/requirements.in
+
+# Copy the actual code
+COPY . /root/
+
+# This tag is supplied by the build script and will be used to determine the version
+# when registering tasks, workflows, and launch plans
+ARG tag
+ENV FLYTE_INTERNAL_IMAGE $tag

examples/slurm_agent/README.md

@@ -0,0 +1,45 @@
+(slurm_agent)=
+
+# Slurm agent
+
+```{eval-rst}
+.. tags:: Integration, HighPerformanceComputing, Advanced
+```
+
+## Installation
+
+To install the Slurm agent, run the following command:
+
+```{eval-rst}
+.. prompt:: bash
+
+    pip install flytekitplugins-slurm
+```
+
+## Example usage
+
+For the example usage of different Slurm task types, please see {doc}`Slurm agent example usage<slurm_agent_example_usage>`.
+
+## Local testing
+
+To test the Slurm agent locally, create a class for the agent task that inherits from [AsyncAgentExecutorMixin](https://github.com/flyteorg/flytekit/blob/cd6bd01ad0ba6688afc71a33a59ece53f90e841a/flytekit/extend/backend/base_agent.py#L3). This mixin can handle asynchronous tasks and allows flytekit to mimic FlytePropeller's behavior in calling the agent. For more information, see "[Testing agents locally](https://docs.flyte.org/en/latest/flyte_agents/testing_agents_in_a_local_python_environment.html)".
+
+```{note}
+In some cases, you will need to store credentials in your local environment when testing locally.
+```
+
+## Flyte deployment configuration
+
+```{note}
+If you are using a managed deployment of Flyte, you will need to contact your deployment administrator to configure agents in your deployment.
+```
+
+To enable the Slurm agent in your Flyte deployment, see the {ref}`Slurm agent deployment guide<deployment-agent-setup-slurm>`.
+
+
+```{toctree}
+:maxdepth: -1
+:hidden:
+
+slurm_agent_example_usage
+```

examples/slurm_agent/requirements.in

@@ -0,0 +1,4 @@
+flytekitplugins-slurm
+torch
+torchvision
+tqdm