fractal-server 2.16.1__py3-none-any.whl → 2.16.2a0__py3-none-any.whl

fractal_server/__init__.py

@@ -1 +1 @@
-__VERSION__ = "2.16.1"
+__VERSION__ = "2.16.2a0"

fractal_server/config.py

@@ -18,16 +18,19 @@ import sys
 from os import environ
 from os import getenv
 from pathlib import Path
+from typing import Annotated
 from typing import Literal
 from typing import TypeVar
 
 from cryptography.fernet import Fernet
 from dotenv import load_dotenv
+from pydantic import AfterValidator
 from pydantic import BaseModel
 from pydantic import EmailStr
 from pydantic import Field
 from pydantic import field_validator
 from pydantic import model_validator
+from pydantic import PositiveInt
 from pydantic import SecretStr
 from pydantic_settings import BaseSettings
 from pydantic_settings import SettingsConfigDict
@@ -36,6 +39,7 @@ from sqlalchemy.engine import URL
 import fractal_server
 from fractal_server.types import AbsolutePathStr
 from fractal_server.types import DictStrStr
+from fractal_server.types import NonEmptyStr
 
 
 class MailSettings(BaseModel):
@@ -64,6 +68,47 @@ class MailSettings(BaseModel):
     use_login: bool
 
 
+def _check_pixi_slurm_memory(mem: str) -> str:
+    if mem[-1] not in ["K", "M", "G", "T"]:
+        raise ValueError(
+            f"Invalid memory requirement {mem=} for `pixi`, "
+            "please set a K/M/G/T units suffix."
+        )
+    return mem
+
+
+class PixiSLURMConfig(BaseModel):
+    """
+    Parameters that are passed directly to a `sbatch` command.
+
+    See https://slurm.schedmd.com/sbatch.html.
+    """
+
+    partition: NonEmptyStr
+    """
+    `-p, --partition=<partition_names>`
+    """
+    cpus: PositiveInt
+    """
+    `-c, --cpus-per-task=<ncpus>
+    """
+    mem: Annotated[NonEmptyStr, AfterValidator(_check_pixi_slurm_memory)]
+    """
+    `--mem=<size>[units]` (examples: `"10M"`, `"10G"`).
+    From `sbatch` docs: Specify the real memory required per node. Default
+    units are megabytes. Different units can be specified using the suffix
+    [K|M|G|T].
+    """
+    time: NonEmptyStr
+    """
+    `-t, --time=<time>`.
+    From `sbatch` docs: "A time limit of zero requests that no time limit be
+    imposed. Acceptable time formats include "minutes", "minutes:seconds",
+    "hours:minutes:seconds", "days-hours", "days-hours:minutes" and
+    "days-hours:minutes:seconds".
+    """
+
+
 class PixiSettings(BaseModel):
     """
     Configuration for Pixi Task collection.
@@ -130,9 +175,15 @@ class PixiSettings(BaseModel):
     """
     DEFAULT_ENVIRONMENT: str = "default"
     """
+    Default pixi environment name.
     """
     DEFAULT_PLATFORM: str = "linux-64"
     """
+    Default platform for pixi.
+    """
+    SLURM_CONFIG: PixiSLURMConfig | None = None
+    """
+    Required when using pixi in a SSH/SLURM deployment.
     """
 
     @model_validator(mode="after")
@@ -784,6 +835,10 @@ class Settings(BaseSettings):
                 raise FractalConfigurationError(
                     f"Must set FRACTAL_SLURM_WORKER_PYTHON when {info}"
                 )
+            if self.pixi and self.pixi.SLURM_CONFIG is None:
+                raise FractalConfigurationError(
+                    "Pixi config must include SLURM_CONFIG."
+                )
 
             from fractal_server.app.runner.executors.slurm_common._slurm_config import (  # noqa: E501
                 load_slurm_config_file,

fractal_server/tasks/v2/ssh/_pixi_slurm_ssh.py

@@ -0,0 +1,115 @@
+import os
+import time
+from pathlib import Path
+
+from fractal_server.config import get_settings
+from fractal_server.config import PixiSLURMConfig
+from fractal_server.logger import get_logger
+from fractal_server.ssh._fabric import FractalSSH
+from fractal_server.syringe import Inject
+
+
+# https://slurm.schedmd.com/squeue.html#lbAG
+STATES_FINISHED = {
+    "BOOT_FAIL",
+    "CANCELLED",
+    "COMPLETED",
+    "DEADLINE",
+    "FAILED",
+    "NODE_FAIL",
+    "OUT_OF_MEMORY",
+    "PREEMPTED",
+    "SPECIAL_EXIT",
+    "TIMEOUT",
+}
+
+
+def run_script_on_remote_slurm(
+    *,
+    script_path: str,
+    slurm_config: PixiSLURMConfig,
+    fractal_ssh: FractalSSH,
+    logger_name: str,
+    prefix: str,
+):
+    """
+    FIXME
+
+
+    NOTE: This is called from within a try/except, thus we can use exceptions
+    as a mechanism to propagate failure/errors.
+    """
+
+    logger = get_logger(logger_name=logger_name)
+    settings = Inject(get_settings)
+
+    # (1) Prepare remote submission script
+    workdir_remote = Path(script_path).parent.as_posix()
+    submission_script_remote = os.path.join(
+        workdir_remote, f"{prefix}-submit.sh"
+    )
+    stderr_remote = os.path.join(workdir_remote, f"{prefix}-err.txt")
+    stdout_remote = os.path.join(workdir_remote, f"{prefix}-out.txt")
+    success_file_remote = os.path.join(workdir_remote, f"{prefix}-success.txt")
+    script_lines = [
+        "#!/bin/bash",
+        f"#SBATCH --partition={slurm_config.partition}",
+        f"#SBATCH --cpus-per-task={slurm_config.cpus}",
+        f"#SBATCH --mem={slurm_config.mem}",
+        f"#SBATCH --time={slurm_config.time}",
+        f"#SBATCH --err={stderr_remote}",
+        f"#SBATCH --out={stdout_remote}",
+        f"#SBATCH -D {workdir_remote}",
+        "",
+        f"bash {script_path}",
+        f"touch {success_file_remote}",
+        "",
+    ]
+    script_contents = "\n".join(script_lines)
+    fractal_ssh.write_remote_file(
+        path=submission_script_remote,
+        content=script_contents,
+    )
+
+    # (2) Submit SLURM job
+    sbatch_cmd = f"sbatch --parsable {submission_script_remote} "
+    try:
+        stdout = fractal_ssh.run_command(cmd=sbatch_cmd)
+    except Exception as e:
+        logger.error(
+            f"Submission of {submission_script_remote} failed. "
+            f"Original error: {str(e)}"
+        )
+        raise e
+    logger.debug(f"Now submit job {submission_script_remote} to SLURM.")
+    job_id = int(stdout)
+    logger.debug(f"SLURM-job submission successful ({job_id=}).")
+
+    # (3) Monitor job
+    squeue_cmd = (
+        f"squeue --noheader --format='%i %T' --states=all --jobs={job_id}"
+    )
+    while True:
+        try:
+            stdout = fractal_ssh.run_command(cmd=squeue_cmd)
+        except Exception as e:
+            # FIXME: review this logic
+            logger.info(
+                f"`squeue` command failed (original error: {e}), "
+                "consider the job as complete."
+            )
+            break
+        state = stdout.strip().split()[1]
+        logger.debug(f"Status of SLURM job {job_id}: {state}")
+        if state in STATES_FINISHED:
+            logger.debug(f"Exit retrieval loop ({state=}).")
+            break
+        time.sleep(settings.FRACTAL_SLURM_POLL_INTERVAL)
+
+    if fractal_ssh.remote_exists(path=success_file_remote):
+        logger.info(f"{success_file_remote=} exists.")
+    else:
+        raise RuntimeError(
+            "SLURM job did not complete correctly "
+            f"({success_file_remote=} missing)."
+        )

fractal_server/tasks/v2/ssh/_utils.py

@@ -17,29 +17,26 @@ from fractal_server.tasks.v2.utils_templates import customize_template
 logger = set_logger(__name__)
 
 
-def _customize_and_run_template(
+def _customize_and_send_template(
     *,
     template_filename: str,
     replacements: set[tuple[str, str]],
     script_dir_local: str,
+    script_dir_remote: str,
     prefix: str,
     fractal_ssh: FractalSSH,
-    script_dir_remote: str,
     logger_name: str,
-    login_shell: bool = False,
 ) -> str:
     """
-    Customize one of the template bash scripts, transfer it to the remote host
-    via SFTP and then run it via SSH.
+    Customize a template bash scripts and transfer it to the remote host.
 
     Args:
         template_filename: Filename of the template file (ends with ".sh").
         replacements: Dictionary of replacements.
         script_dir_local: Local folder where the script will be placed.
+        script_dir_remote: Remote scripts directory
         prefix: Prefix for the script filename.
         fractal_ssh: FractalSSH object
-        script_dir_remote: Remote scripts directory
-        login_shell: If `True`, use `bash --login` for remote script execution.
     """
     logger = get_logger(logger_name=logger_name)
     logger.debug(f"_customize_and_run_template {template_filename} - START")
@@ -67,10 +64,46 @@ def _customize_and_run_template(
         local=script_path_local,
         remote=script_path_remote,
     )
+    return script_path_remote
+
+
+def _customize_and_run_template(
+    *,
+    template_filename: str,
+    replacements: set[tuple[str, str]],
+    script_dir_local: str,
+    script_dir_remote: str,
+    prefix: str,
+    fractal_ssh: FractalSSH,
+    logger_name: str,
+) -> str:
+    """
+    Customize one of the template bash scripts, transfer it to the remote host
+    via SFTP and then run it via SSH.
+
+    Args:
+        template_filename: Filename of the template file (ends with ".sh").
+        replacements: Dictionary of replacements.
+        script_dir_remote: Remote scripts directory
+        script_dir_local: Local folder where the script will be placed.
+        prefix: Prefix for the script filename.
+        fractal_ssh: FractalSSH object
+    """
+    logger = get_logger(logger_name=logger_name)
+    logger.debug(f"_customize_and_run_template {template_filename} - START")
+
+    script_path_remote = _customize_and_send_template(
+        template_filename=template_filename,
+        replacements=replacements,
+        script_dir_local=script_dir_local,
+        script_dir_remote=script_dir_remote,
+        prefix=prefix,
+        fractal_ssh=fractal_ssh,
+        logger_name=logger_name,
+    )
 
     # Execute script remotely
-    bash = "bash --login" if login_shell else "bash"
-    cmd = f"{bash} {script_path_remote}"
+    cmd = f"bash {script_path_remote}"
     logger.debug(f"Now run '{cmd}' over SSH.")
     stdout = fractal_ssh.run_command(cmd=cmd)
 

fractal_server/tasks/v2/ssh/collect_pixi.py

@@ -8,6 +8,7 @@ from ..utils_background import prepare_tasks_metadata
 from ..utils_database import create_db_tasks_and_update_task_group_sync
 from ..utils_pixi import parse_collect_stdout
 from ..utils_pixi import SOURCE_DIR_NAME
+from ._pixi_slurm_ssh import run_script_on_remote_slurm
 from ._utils import check_ssh_or_fail_and_cleanup
 from ._utils import edit_pyproject_toml_in_place_ssh
 from fractal_server.app.db import get_sync_db
@@ -22,6 +23,7 @@ from fractal_server.ssh._fabric import SingleUseFractalSSH
 from fractal_server.ssh._fabric import SSHConfig
 from fractal_server.syringe import Inject
 from fractal_server.tasks.v2.ssh._utils import _customize_and_run_template
+from fractal_server.tasks.v2.ssh._utils import _customize_and_send_template
 from fractal_server.tasks.v2.utils_background import add_commit_refresh
 from fractal_server.tasks.v2.utils_background import get_current_log
 from fractal_server.tasks.v2.utils_templates import SCRIPTS_SUBFOLDER
@@ -201,16 +203,30 @@ def collect_ssh_pixi(
                         pyproject_toml_path=pyproject_toml_path,
                     )
 
-                    stdout = _customize_and_run_template(
+                    # Run script 2 - run pixi-install command
+                    remote_script2_path = _customize_and_send_template(
                         template_filename="pixi_2_install.sh",
                         replacements=replacements,
-                        login_shell=True,
                         **common_args,
                     )
-                    logger.debug(f"STDOUT: {stdout}")
+                    logger.debug(
+                        "Installation script written to "
+                        f"{remote_script2_path=}."
+                    )
+                    activity.log = get_current_log(log_file_path)
+                    activity = add_commit_refresh(obj=activity, db=db)
+
+                    run_script_on_remote_slurm(
+                        script_path=remote_script2_path,
+                        slurm_config=settings.pixi.SLURM_CONFIG,
+                        fractal_ssh=fractal_ssh,
+                        logger_name=LOGGER_NAME,
+                        prefix=common_args["prefix"],
+                    )
                     activity.log = get_current_log(log_file_path)
                     activity = add_commit_refresh(obj=activity, db=db)
 
+                    # Run script 3 - post-install
                     stdout = _customize_and_run_template(
                         template_filename="pixi_3_post_install.sh",
                         replacements=replacements,

fractal_server/tasks/v2/ssh/reactivate_pixi.py

@@ -5,6 +5,7 @@ from tempfile import TemporaryDirectory
 from ..utils_background import fail_and_cleanup
 from ..utils_background import get_activity_and_task_group
 from ..utils_pixi import SOURCE_DIR_NAME
+from ._pixi_slurm_ssh import run_script_on_remote_slurm
 from ._utils import check_ssh_or_fail_and_cleanup
 from ._utils import edit_pyproject_toml_in_place_ssh
 from fractal_server.app.db import get_sync_db
@@ -18,6 +19,7 @@ from fractal_server.ssh._fabric import SSHConfig
 from fractal_server.syringe import Inject
 from fractal_server.tasks.utils import get_log_path
 from fractal_server.tasks.v2.ssh._utils import _customize_and_run_template
+from fractal_server.tasks.v2.ssh._utils import _customize_and_send_template
 from fractal_server.tasks.v2.utils_background import add_commit_refresh
 from fractal_server.tasks.v2.utils_background import get_current_log
 from fractal_server.tasks.v2.utils_templates import SCRIPTS_SUBFOLDER
@@ -187,13 +189,25 @@ def reactivate_ssh_pixi(
                     )
 
                     # Run script 2 - run pixi-install command
-                    stdout = _customize_and_run_template(
+                    remote_script2_path = _customize_and_send_template(
                         template_filename="pixi_2_install.sh",
                         replacements=replacements,
-                        login_shell=True,
                         **common_args,
                     )
-                    logger.debug(f"STDOUT: {stdout}")
+                    logger.debug(
+                        "Installation script written to "
+                        f"{remote_script2_path=}."
+                    )
+                    activity.log = get_current_log(log_file_path)
+                    activity = add_commit_refresh(obj=activity, db=db)
+
+                    run_script_on_remote_slurm(
+                        script_path=remote_script2_path,
+                        slurm_config=settings.pixi.SLURM_CONFIG,
+                        fractal_ssh=fractal_ssh,
+                        logger_name=LOGGER_NAME,
+                        prefix=common_args["prefix"],
+                    )
                     activity.log = get_current_log(log_file_path)
                     activity = add_commit_refresh(obj=activity, db=db)
 

{fractal_server-2.16.1.dist-info → fractal_server-2.16.2a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: fractal-server
-Version: 2.16.1
+Version: 2.16.2a0
 Summary: Backend component of the Fractal analytics platform
 License: BSD-3-Clause
 Author: Tommaso Comparin

{fractal_server-2.16.1.dist-info → fractal_server-2.16.2a0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-fractal_server/__init__.py,sha256=Slaz4BSSaAhGoGVOZBGgBrBLlFksmPxY8vXizrOcKgE,23
+fractal_server/__init__.py,sha256=ywvZY_b51fGszgBtc9wZwOVbz6O5w5tzfaYERx4Ui_s,25
 fractal_server/__main__.py,sha256=rkM8xjY1KeS3l63irB8yCrlVobR-73uDapC4wvrIlxI,6957
 fractal_server/alembic.ini,sha256=MWwi7GzjzawI9cCAK1LW7NxIBQDUqD12-ptJoq5JpP0,3153
 fractal_server/app/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -129,7 +129,7 @@ fractal_server/app/schemas/v2/workflowtask.py,sha256=6eweAMyziwaoMT-7R1fVJYunIeZ
 fractal_server/app/security/__init__.py,sha256=oJ8RVglpOvWPQY4RokiE2YA72Nqo42dZEjywWTt8xr8,14032
 fractal_server/app/security/signup_email.py,sha256=Xd6QYxcdmg0PHpDwmUE8XQmPcOj3Xjy5oROcIMhmltM,1472
 fractal_server/app/user_settings.py,sha256=OP1yiYKtPadxwM51_Q0hdPk3z90TCN4z1BLpQsXyWiU,1316
-fractal_server/config.py,sha256=EdgVC0uAEUKiLft7k-IQd9W_TxVGrBNi9TflKEcNKu8,29177
+fractal_server/config.py,sha256=q0f1DsSgXKE9gyzty7y4NGiQyttV6H13LUMPZsOeGH8,30896
 fractal_server/data_migrations/2_14_10.py,sha256=jzMg2c1zNO8C_Nho_9_EZJD6kR1-gkFNpNrMR5Hr8hM,1598
 fractal_server/data_migrations/README.md,sha256=_3AEFvDg9YkybDqCLlFPdDmGJvr6Tw7HRI14aZ3LOIw,398
 fractal_server/data_migrations/tools.py,sha256=LeMeASwYGtEqd-3wOLle6WARdTGAimoyMmRbbJl-hAM,572
@@ -204,14 +204,15 @@ fractal_server/tasks/v2/local/delete.py,sha256=TbNcg0IhbJNXzXsO2Vgv6Dq4jWmFUVtSu
 fractal_server/tasks/v2/local/reactivate.py,sha256=Q43DOadNeFyyfgNP67lUqaXmZsS6onv67XwxH_-5ANA,5756
 fractal_server/tasks/v2/local/reactivate_pixi.py,sha256=IuxDRaj8i6Rc582TIbc9HVKQ9pOvR4IepyXSaJx2PfQ,7565
 fractal_server/tasks/v2/ssh/__init__.py,sha256=dPK6BtEZVh1GiFP05j1RKTEnZvjJez8o2KkMC2hWXaw,339
-fractal_server/tasks/v2/ssh/_utils.py,sha256=V8ZcjEMCNSHivUyu1rnGG_V96Ylq1RRvocQDN1n6BAw,5120
+fractal_server/tasks/v2/ssh/_pixi_slurm_ssh.py,sha256=P-QyoKlQ9rbq3DPuG0DjgnbJ9KOCBXqbFcJZ2qB4VzA,3484
+fractal_server/tasks/v2/ssh/_utils.py,sha256=AjbH4o6-ENe-ZiHvE_pGe_yw8e2lapGMYtLX7VuBkA8,6111
 fractal_server/tasks/v2/ssh/collect.py,sha256=WesUBtNaax9ST7CtgqAD2qd-ZII4t1JFPzRzltEOuM8,14333
-fractal_server/tasks/v2/ssh/collect_pixi.py,sha256=unZXs3qjBoWYBIa2xWiIKz2CVMVBAhcloQ_9YJ97n5k,14211
+fractal_server/tasks/v2/ssh/collect_pixi.py,sha256=5CRr_XiuRaFjydZqxgfoMMj3p1TfT7aqNOxHcVxrOQI,15020
 fractal_server/tasks/v2/ssh/deactivate.py,sha256=7gJ2mOah0wKwUnK0S9QpaXg08_WE95P0rC-oExAGhLE,12438
 fractal_server/tasks/v2/ssh/deactivate_pixi.py,sha256=K0yK_NPUqhFMj6cp6G_0Kfn0Yo7oQux4kT5dFPulnos,4748
 fractal_server/tasks/v2/ssh/delete.py,sha256=9DCe5QXZ14pY8NRb_FCqpIr31r5Gboz6bnFrS8Oa044,4277
 fractal_server/tasks/v2/ssh/reactivate.py,sha256=NJIgMNFKaXMhbvK0iZOsMwMtsms6Boj9f8N4L01X9Bo,8271
-fractal_server/tasks/v2/ssh/reactivate_pixi.py,sha256=UOlG01wOv-eQCtUxSwphEY-Ey7OoUm5PMCSP9J-i1rY,10323
+fractal_server/tasks/v2/ssh/reactivate_pixi.py,sha256=LZk7WaC2uCo3cHTwrpRcc4F_I7HUDA-8rE-RJ2agLZA,11020
 fractal_server/tasks/v2/templates/1_create_venv.sh,sha256=PK0jdHKtQpda1zULebBaVPORt4t6V17wa4N1ohcj5ac,548
 fractal_server/tasks/v2/templates/2_pip_install.sh,sha256=vN9DX-eucJjB-XCuQuZmkuATGBzL4FlDQJTQdVewJG8,2155
 fractal_server/tasks/v2/templates/3_pip_freeze.sh,sha256=JldREScEBI4cD_qjfX4UK7V4aI-FnX9ZvVNxgpSOBFc,168
@@ -235,8 +236,8 @@ fractal_server/types/validators/_workflow_task_arguments_validators.py,sha256=HL
 fractal_server/urls.py,sha256=QjIKAC1a46bCdiPMu3AlpgFbcv6a4l3ABcd5xz190Og,471
 fractal_server/utils.py,sha256=Vn35lApt1T1J8nc09sAVqd10Cy0sa3dLipcljI-hkuk,2185
 fractal_server/zip_tools.py,sha256=H0w7wS5yE4ebj7hw1_77YQ959dl2c-L0WX6J_ro1TY4,4884
-fractal_server-2.16.1.dist-info/LICENSE,sha256=QKAharUuhxL58kSoLizKJeZE3mTCBnX6ucmz8W0lxlk,1576
-fractal_server-2.16.1.dist-info/METADATA,sha256=eaTdwJo3-jLcC4s9w7iJgXLVjn6axnfpWBRgxhq1jRA,4334
-fractal_server-2.16.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-fractal_server-2.16.1.dist-info/entry_points.txt,sha256=8tV2kynvFkjnhbtDnxAqImL6HMVKsopgGfew0DOp5UY,58
-fractal_server-2.16.1.dist-info/RECORD,,
+fractal_server-2.16.2a0.dist-info/LICENSE,sha256=QKAharUuhxL58kSoLizKJeZE3mTCBnX6ucmz8W0lxlk,1576
+fractal_server-2.16.2a0.dist-info/METADATA,sha256=6x9kEFUU9jQkukimP25U3iu8UKP1JY2bCA67qTSxx2I,4336
+fractal_server-2.16.2a0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+fractal_server-2.16.2a0.dist-info/entry_points.txt,sha256=8tV2kynvFkjnhbtDnxAqImL6HMVKsopgGfew0DOp5UY,58
+fractal_server-2.16.2a0.dist-info/RECORD,,