fractal-server 2.14.5__py3-none-any.whl → 2.14.7__py3-none-any.whl

fractal_server/app/runner/executors/local/get_local_config.py

@@ -14,7 +14,6 @@ Submodule to handle the local-backend configuration for a WorkflowTask
 import json
 from pathlib import Path
 from typing import Literal
-from typing import Optional
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
@@ -45,7 +44,7 @@ class LocalBackendConfig(BaseModel):
     """
 
     model_config = ConfigDict(extra="forbid")
-    parallel_tasks_per_job: Optional[int] = None
+    parallel_tasks_per_job: int | None = None
 
     @property
     def batch_size(self) -> int:
@@ -55,7 +54,7 @@ class LocalBackendConfig(BaseModel):
 def get_local_backend_config(
     wftask: WorkflowTaskV2,
     which_type: Literal["non_parallel", "parallel"],
-    config_path: Optional[Path] = None,
+    config_path: Path | None = None,
     tot_tasks: int = 1,
 ) -> LocalBackendConfig:
     """

fractal_server/app/runner/executors/local/runner.py

@@ -44,7 +44,7 @@ def run_single_task(
     )
 
     try:
-        with open(task_files.metadiff_file_local, "r") as f:
+        with open(task_files.metadiff_file_local) as f:
             out_meta = json.load(f)
         return out_meta
     except FileNotFoundError:

fractal_server/app/runner/executors/slurm_common/_batching.py

@@ -12,7 +12,6 @@
 Submodule to determine the number of total/parallel tasks per SLURM job.
 """
 import math
-from typing import Optional
 
 from fractal_server.logger import set_logger
 
@@ -58,8 +57,8 @@ def heuristics(
     # Number of parallel components (always known)
     tot_tasks: int,
     # Optional WorkflowTask attributes:
-    tasks_per_job: Optional[int] = None,
-    parallel_tasks_per_job: Optional[int] = None,
+    tasks_per_job: int | None = None,
+    parallel_tasks_per_job: int | None = None,
     # Task requirements (multiple possible sources):
     cpus_per_task: int,
     mem_per_task: int,

fractal_server/app/runner/executors/slurm_common/_slurm_config.py

@@ -14,8 +14,6 @@ Submodule to handle the SLURM configuration for a WorkflowTask
 """
 import json
 from pathlib import Path
-from typing import Optional
-from typing import Union
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
@@ -56,16 +54,16 @@ class _SlurmConfigSet(BaseModel):
 
     model_config = ConfigDict(extra="forbid")
 
-    partition: Optional[str] = None
-    cpus_per_task: Optional[int] = None
-    mem: Optional[Union[int, str]] = None
-    constraint: Optional[str] = None
-    gres: Optional[str] = None
-    time: Optional[str] = None
-    account: Optional[str] = None
-    extra_lines: Optional[list[str]] = None
-    pre_submission_commands: Optional[list[str]] = None
-    gpus: Optional[str] = None
+    partition: str | None = None
+    cpus_per_task: int | None = None
+    mem: int | str | None = None
+    constraint: str | None = None
+    gres: str | None = None
+    time: str | None = None
+    account: str | None = None
+    extra_lines: list[str] | None = None
+    pre_submission_commands: list[str] | None = None
+    gpus: str | None = None
 
 
 class _BatchingConfigSet(BaseModel):
@@ -89,8 +87,8 @@ class _BatchingConfigSet(BaseModel):
 
     target_cpus_per_job: int
     max_cpus_per_job: int
-    target_mem_per_job: Union[int, str]
-    max_mem_per_job: Union[int, str]
+    target_mem_per_job: int | str
+    max_mem_per_job: int | str
     target_num_jobs: int
     max_num_jobs: int
 
@@ -143,13 +141,13 @@ class SlurmConfigFile(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
     default_slurm_config: _SlurmConfigSet
-    gpu_slurm_config: Optional[_SlurmConfigSet] = None
+    gpu_slurm_config: _SlurmConfigSet | None = None
     batching_config: _BatchingConfigSet
-    user_local_exports: Optional[dict[str, str]] = None
+    user_local_exports: dict[str, str] | None = None
 
 
 def load_slurm_config_file(
-    config_path: Optional[Path] = None,
+    config_path: Path | None = None,
 ) -> SlurmConfigFile:
     """
     Load a SLURM configuration file and validate its content with
@@ -264,23 +262,23 @@ class SlurmConfig(BaseModel):
     shebang_line: str = "#!/bin/sh"
 
     # Optional SLURM parameters
-    job_name: Optional[str] = None
-    constraint: Optional[str] = None
-    gres: Optional[str] = None
-    gpus: Optional[str] = None
-    time: Optional[str] = None
-    account: Optional[str] = None
+    job_name: str | None = None
+    constraint: str | None = None
+    gres: str | None = None
+    gpus: str | None = None
+    time: str | None = None
+    account: str | None = None
 
     # Free-field attribute for extra lines to be added to the SLURM job
     # preamble
-    extra_lines: Optional[list[str]] = Field(default_factory=list)
+    extra_lines: list[str] | None = Field(default_factory=list)
 
     # Variables that will be `export`ed in the SLURM submission script
-    user_local_exports: Optional[dict[str, str]] = None
+    user_local_exports: dict[str, str] | None = None
 
     # Metaparameters needed to combine multiple tasks in each SLURM job
-    tasks_per_job: Optional[int] = None
-    parallel_tasks_per_job: Optional[int] = None
+    tasks_per_job: int | None = None
+    parallel_tasks_per_job: int | None = None
     target_cpus_per_job: int
     max_cpus_per_job: int
     target_mem_per_job: int
@@ -328,7 +326,7 @@ class SlurmConfig(BaseModel):
 
     def to_sbatch_preamble(
         self,
-        remote_export_dir: Optional[str] = None,
+        remote_export_dir: str | None = None,
     ) -> list[str]:
         """
         Compile `SlurmConfig` object into the preamble of a SLURM submission
@@ -410,7 +408,7 @@ class SlurmConfig(BaseModel):
         return self.tasks_per_job
 
 
-def _parse_mem_value(raw_mem: Union[str, int]) -> int:
+def _parse_mem_value(raw_mem: str | int) -> int:
     """
     Convert a memory-specification string into an integer (in MB units), or
     simply return the input if it is already an integer.

fractal_server/app/runner/executors/slurm_common/base_slurm_runner.py

@@ -5,7 +5,6 @@ import time
 from pathlib import Path
 from typing import Any
 from typing import Literal
-from typing import Optional
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
@@ -74,6 +73,7 @@ class BaseSlurmRunner(BaseRunner):
     jobs: dict[str, SlurmJob]
     python_worker_interpreter: str
     slurm_runner_type: Literal["ssh", "sudo"]
+    slurm_account: str | None = None
 
     def __init__(
         self,
@@ -81,9 +81,10 @@ class BaseSlurmRunner(BaseRunner):
         root_dir_remote: Path,
         slurm_runner_type: Literal["ssh", "sudo"],
         python_worker_interpreter: str,
-        common_script_lines: Optional[list[str]] = None,
-        user_cache_dir: Optional[str] = None,
-        poll_interval: Optional[int] = None,
+        common_script_lines: list[str] | None = None,
+        user_cache_dir: str | None = None,
+        poll_interval: int | None = None,
+        slurm_account: str | None = None,
     ):
         self.slurm_runner_type = slurm_runner_type
         self.root_dir_local = root_dir_local
@@ -92,6 +93,7 @@ class BaseSlurmRunner(BaseRunner):
         self._check_slurm_account()
         self.user_cache_dir = user_cache_dir
         self.python_worker_interpreter = python_worker_interpreter
+        self.slurm_account = slurm_account
 
         settings = Inject(get_settings)
 
@@ -186,6 +188,24 @@ class BaseSlurmRunner(BaseRunner):
     ) -> str:
         logger.debug("[_submit_single_sbatch] START")
 
+        # Include SLURM account in `slurm_config`. Note: we make this change
+        # here, rather than exposing a new argument of `get_slurm_config`,
+        # because it's a backend-specific argument while `get_slurm_config` has
+        # a generic interface.
+        if self.slurm_account is not None:
+            slurm_config.account = self.slurm_account
+
+        # Include common_script_lines in extra_lines
+        if len(self.common_script_lines) > 0:
+            logger.debug(
+                f"Add {self.common_script_lines} to "
+                f"{slurm_config.extra_lines=}."
+            )
+            current_extra_lines = slurm_config.extra_lines or []
+            slurm_config.extra_lines = (
+                current_extra_lines + self.common_script_lines
+            )
+
         for task in slurm_job.tasks:
             # Write input file
             if self.slurm_runner_type == "ssh":
@@ -241,13 +261,11 @@ class BaseSlurmRunner(BaseRunner):
                 input_file = task.input_file_local
             output_file = task.output_file_remote
             cmdlines.append(
-                (
-                    f"{self.python_worker_interpreter}"
-                    " -m fractal_server.app.runner."
-                    "executors.slurm_common.remote "
-                    f"--input-file {input_file} "
-                    f"--output-file {output_file}"
-                )
+                f"{self.python_worker_interpreter}"
+                " -m fractal_server.app.runner."
+                "executors.slurm_common.remote "
+                f"--input-file {input_file} "
+                f"--output-file {output_file}"
             )
 
         # Set ntasks
@@ -390,7 +408,7 @@ class BaseSlurmRunner(BaseRunner):
         was_job_scancelled: bool = False,
     ) -> tuple[Any, Exception]:
         try:
-            with open(task.output_file_local, "r") as f:
+            with open(task.output_file_local) as f:
                 output = json.load(f)
             success = output[0]
             if success:
@@ -885,8 +903,8 @@ class BaseSlurmRunner(BaseRunner):
         """
         Check that a list of `SlurmJob`s have homogeneous working folders.
         """
-        set_workdir_local = set(_job.workdir_local for _job in slurm_jobs)
-        set_workdir_remote = set(_job.workdir_remote for _job in slurm_jobs)
+        set_workdir_local = {_job.workdir_local for _job in slurm_jobs}
+        set_workdir_remote = {_job.workdir_remote for _job in slurm_jobs}
         if len(set_workdir_local) > 1:
             raise ValueError(f"Non-unique values in {set_workdir_local=}.")
         if len(set_workdir_remote) > 1:

fractal_server/app/runner/executors/slurm_common/get_slurm_config.py

@@ -1,6 +1,5 @@
 from pathlib import Path
 from typing import Literal
-from typing import Optional
 
 from ._batching import heuristics
 from ._slurm_config import _parse_mem_value
@@ -14,7 +13,7 @@ from fractal_server.app.models.v2 import WorkflowTaskV2
 def get_slurm_config_internal(
     wftask: WorkflowTaskV2,
     which_type: Literal["non_parallel", "parallel"],
-    config_path: Optional[Path] = None,
+    config_path: Path | None = None,
 ) -> SlurmConfig:
     """
     Prepare a `SlurmConfig` configuration object
@@ -168,7 +167,7 @@ def get_slurm_config_internal(
 def get_slurm_config(
     wftask: WorkflowTaskV2,
     which_type: Literal["non_parallel", "parallel"],
-    config_path: Optional[Path] = None,
+    config_path: Path | None = None,
     tot_tasks: int = 1,
 ) -> SlurmConfig:
     config = get_slurm_config_internal(

fractal_server/app/runner/executors/slurm_common/remote.py

@@ -37,7 +37,7 @@ def worker(
 
     # Execute the job and capture exceptions
     try:
-        with open(in_fname, "r") as f:
+        with open(in_fname) as f:
             input_data = json.load(f)
 
         server_python_version = input_data["python_version"]
@@ -68,7 +68,7 @@ def worker(
         call_command_wrapper(cmd=full_command, log_path=log_path)
 
         try:
-            with open(metadiff_file_remote, "r") as f:
+            with open(metadiff_file_remote) as f:
                 out_meta = json.load(f)
             result = (True, out_meta)
         except FileNotFoundError:

fractal_server/app/runner/executors/slurm_common/slurm_job_task_models.py

@@ -1,6 +1,5 @@
 from pathlib import Path
 from typing import Any
-from typing import Optional
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
@@ -15,7 +14,7 @@ class SlurmTask(BaseModel):
     workdir_local: Path
     workdir_remote: Path
     parameters: dict[str, Any]
-    zarr_url: Optional[str] = None
+    zarr_url: str | None = None
     task_files: TaskFiles
     index: int
 
@@ -65,7 +64,7 @@ class SlurmTask(BaseModel):
 
 
 class SlurmJob(BaseModel):
-    slurm_job_id: Optional[str] = None
+    slurm_job_id: str | None = None
     prefix: str
     workdir_local: Path
     workdir_remote: Path

fractal_server/app/runner/executors/slurm_ssh/run_subprocess.py

@@ -1,6 +1,5 @@
 import shlex
 import subprocess  # nosec
-from typing import Optional
 
 from fractal_server.logger import get_logger
 from fractal_server.string_tools import validate_cmd
@@ -8,8 +7,8 @@ from fractal_server.string_tools import validate_cmd
 
 def run_subprocess(
     cmd: str,
-    allow_char: Optional[str] = None,
-    logger_name: Optional[str] = None,
+    allow_char: str | None = None,
+    logger_name: str | None = None,
 ) -> subprocess.CompletedProcess:
     validate_cmd(cmd, allow_char=allow_char)
     logger = get_logger(logger_name)

fractal_server/app/runner/executors/slurm_ssh/runner.py

@@ -1,6 +1,5 @@
 import time
 from pathlib import Path
-from typing import Optional
 
 from ..slurm_common.base_slurm_runner import BaseSlurmRunner
 from ..slurm_common.slurm_job_task_models import SlurmJob
@@ -27,10 +26,11 @@ class SlurmSSHRunner(BaseSlurmRunner):
         # Common
         root_dir_local: Path,
         root_dir_remote: Path,
-        common_script_lines: Optional[list[str]] = None,
-        user_cache_dir: Optional[str] = None,
-        poll_interval: Optional[int] = None,
+        common_script_lines: list[str] | None = None,
+        user_cache_dir: str | None = None,
+        poll_interval: int | None = None,
         # Specific
+        slurm_account: str | None = None,
         fractal_ssh: FractalSSH,
     ) -> None:
         """
@@ -50,6 +50,7 @@ class SlurmSSHRunner(BaseSlurmRunner):
             user_cache_dir=user_cache_dir,
             poll_interval=poll_interval,
             python_worker_interpreter=settings.FRACTAL_SLURM_WORKER_PYTHON,
+            slurm_account=slurm_account,
         )
 
     def _mkdir_local_folder(self, folder: str) -> None:

fractal_server/app/runner/executors/slurm_sudo/_subprocess_run_as_user.py

@@ -17,7 +17,6 @@ another user. Note that this requires appropriate sudo permissions.
 """
 import shlex
 import subprocess  # nosec
-from typing import Optional
 
 from fractal_server.logger import set_logger
 from fractal_server.string_tools import validate_cmd
@@ -28,7 +27,7 @@ logger = set_logger(__name__)
 def _run_command_as_user(
     *,
     cmd: str,
-    user: Optional[str] = None,
+    user: str | None = None,
     check: bool = False,
 ) -> subprocess.CompletedProcess:
     """

fractal_server/app/runner/executors/slurm_sudo/runner.py

@@ -5,7 +5,6 @@ import subprocess  # nosec
 import sys
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
-from typing import Optional
 
 from ..slurm_common.base_slurm_runner import BaseSlurmRunner
 from ..slurm_common.slurm_job_task_models import SlurmJob
@@ -21,7 +20,7 @@ logger = set_logger(__name__)
 
 def _subprocess_run_or_raise(
     full_command: str,
-) -> Optional[subprocess.CompletedProcess]:
+) -> subprocess.CompletedProcess | None:
     try:
         output = subprocess.run(  # nosec
             shlex.split(full_command),
@@ -43,7 +42,7 @@ def _subprocess_run_or_raise(
 
 class SudoSlurmRunner(BaseSlurmRunner):
     slurm_user: str
-    slurm_account: Optional[str] = None
+    slurm_account: str | None = None
 
     def __init__(
         self,
@@ -51,11 +50,11 @@ class SudoSlurmRunner(BaseSlurmRunner):
         # Common
         root_dir_local: Path,
         root_dir_remote: Path,
-        common_script_lines: Optional[list[str]] = None,
-        user_cache_dir: Optional[str] = None,
-        poll_interval: Optional[int] = None,
+        common_script_lines: list[str] | None = None,
+        user_cache_dir: str | None = None,
+        poll_interval: int | None = None,
         # Specific
-        slurm_account: Optional[str] = None,
+        slurm_account: str | None = None,
         slurm_user: str,
     ) -> None:
         """
@@ -64,7 +63,6 @@ class SudoSlurmRunner(BaseSlurmRunner):
         """
 
         self.slurm_user = slurm_user
-        self.slurm_account = slurm_account
         settings = Inject(get_settings)
 
         super().__init__(
@@ -77,6 +75,7 @@ class SudoSlurmRunner(BaseSlurmRunner):
             python_worker_interpreter=(
                 settings.FRACTAL_SLURM_WORKER_PYTHON or sys.executable
             ),
+            slurm_account=slurm_account,
         )
 
     def _mkdir_local_folder(self, folder: str) -> None:

fractal_server/app/runner/set_start_and_last_task_index.py

@@ -1,10 +1,7 @@
-from typing import Optional
-
-
 def set_start_and_last_task_index(
     num_tasks: int,
-    first_task_index: Optional[int] = None,
-    last_task_index: Optional[int] = None,
+    first_task_index: int | None = None,
+    last_task_index: int | None = None,
 ) -> tuple[int, int]:
     """
     Handle `first_task_index` and `last_task_index`, by setting defaults and

fractal_server/app/runner/shutdown.py

@@ -38,21 +38,15 @@ async def cleanup_after_shutdown(*, jobsV2: list[int], logger_name: str):
 
             if len(jobsV2_db) == 0:
                 logger.info(
-                    (
-                        "All jobs associated to this app are "
-                        "either done or failed. Exit."
-                    )
+                    "All jobs associated to this app are "
+                    "either done or failed. Exit."
                 )
                 return
             else:
-                logger.info(
-                    (f"Some jobs are still 'submitted' " f"{jobsV2_db=}")
-                )
+                logger.info(f"Some jobs are still 'submitted' {jobsV2_db=}")
         logger.info(
-            (
-                "Graceful shutdown reached its maximum time, "
-                "but some jobs are still submitted"
-            )
+            "Graceful shutdown reached its maximum time, "
+            "but some jobs are still submitted"
         )
 
         for job in jobsV2_db:

fractal_server/app/runner/task_files.py

@@ -1,6 +1,4 @@
 from pathlib import Path
-from typing import Optional
-from typing import Union
 
 from pydantic import BaseModel
 
@@ -12,7 +10,7 @@ MULTISUBMIT_PREFIX = "par"
 
 
 def task_subfolder_name(
-    order: Union[int, str],
+    order: int | str,
     task_name: str,
 ) -> str:
     """
@@ -48,8 +46,8 @@ class TaskFiles(BaseModel):
     task_order: int
 
     # Per-single-component
-    component: Optional[str] = None
-    prefix: Optional[str] = None
+    component: str | None = None
+    prefix: str | None = None
 
     def _check_component(self):
         if self.component is None:

fractal_server/app/runner/v2/_local.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-from typing import Optional
 
 from ...models.v2 import DatasetV2
 from ...models.v2 import WorkflowV2
@@ -16,9 +15,9 @@ def process_workflow(
     dataset: DatasetV2,
     workflow_dir_local: Path,
     job_id: int,
-    workflow_dir_remote: Optional[Path] = None,
-    first_task_index: Optional[int] = None,
-    last_task_index: Optional[int] = None,
+    workflow_dir_remote: Path | None = None,
+    first_task_index: int | None = None,
+    last_task_index: int | None = None,
     logger_name: str,
     job_attribute_filters: AttributeFilters,
     job_type_filters: dict[str, bool],

fractal_server/app/runner/v2/_slurm_ssh.py

@@ -16,7 +16,6 @@ Slurm Backend
 This backend runs fractal workflows in a SLURM cluster.
 """
 from pathlib import Path
-from typing import Optional
 
 from ....ssh._fabric import FractalSSH
 from ...models.v2 import DatasetV2
@@ -38,16 +37,17 @@ def process_workflow(
     dataset: DatasetV2,
     workflow_dir_local: Path,
     job_id: int,
-    workflow_dir_remote: Optional[Path] = None,
-    first_task_index: Optional[int] = None,
-    last_task_index: Optional[int] = None,
+    workflow_dir_remote: Path | None = None,
+    first_task_index: int | None = None,
+    last_task_index: int | None = None,
     logger_name: str,
     job_attribute_filters: AttributeFilters,
     job_type_filters: dict[str, bool],
-    fractal_ssh: FractalSSH,
-    worker_init: Optional[str] = None,
     user_id: int,
-    **kwargs,  # not used
+    # SLURM-ssh-specific
+    fractal_ssh: FractalSSH,
+    slurm_account: str | None = None,
+    worker_init: str | None = None,
 ) -> None:
     """
     Process workflow (SLURM backend public interface)
@@ -80,6 +80,7 @@ def process_workflow(
         fractal_ssh=fractal_ssh,
         root_dir_local=workflow_dir_local,
         root_dir_remote=workflow_dir_remote,
+        slurm_account=slurm_account,
         common_script_lines=worker_init,
     ) as runner:
         execute_tasks_v2(

fractal_server/app/runner/v2/_slurm_sudo.py

@@ -16,7 +16,6 @@ Slurm Backend
 This backend runs fractal workflows in a SLURM cluster.
 """
 from pathlib import Path
-from typing import Optional
 
 from ...models.v2 import DatasetV2
 from ...models.v2 import WorkflowV2
@@ -33,18 +32,18 @@ def process_workflow(
     dataset: DatasetV2,
     workflow_dir_local: Path,
     job_id: int,
-    workflow_dir_remote: Optional[Path] = None,
-    first_task_index: Optional[int] = None,
-    last_task_index: Optional[int] = None,
+    workflow_dir_remote: Path | None = None,
+    first_task_index: int | None = None,
+    last_task_index: int | None = None,
     logger_name: str,
     job_attribute_filters: AttributeFilters,
     job_type_filters: dict[str, bool],
     user_id: int,
-    # Slurm-specific
-    user_cache_dir: Optional[str] = None,
-    slurm_user: Optional[str] = None,
-    slurm_account: Optional[str] = None,
-    worker_init: Optional[str] = None,
+    # SLURM-sudo-specific
+    user_cache_dir: str | None = None,
+    slurm_user: str | None = None,
+    slurm_account: str | None = None,
+    worker_init: str | None = None,
 ) -> None:
     """
     Process workflow (SLURM backend public interface).

fractal_server/app/runner/v2/runner.py

@@ -1,11 +1,10 @@
 import logging
+from collections.abc import Callable
 from copy import copy
 from copy import deepcopy
 from pathlib import Path
 from typing import Any
-from typing import Callable
 from typing import Literal
-from typing import Optional
 
 from sqlalchemy.orm.attributes import flag_modified
 from sqlmodel import delete
@@ -76,13 +75,13 @@ def execute_tasks_v2(
     user_id: int,
     workflow_dir_local: Path,
     job_id: int,
-    workflow_dir_remote: Optional[Path] = None,
-    logger_name: Optional[str] = None,
+    workflow_dir_remote: Path | None = None,
+    logger_name: str | None = None,
     get_runner_config: Callable[
         [
             WorkflowTaskV2,
             Literal["non_parallel", "parallel"],
-            Optional[Path],
+            Path | None,
         ],
         Any,
     ],