fractal-server 2.2.0a0__py3-none-any.whl → 2.3.0a0__py3-none-any.whl

fractal_server/app/runner/v2/_slurm_ssh/_submit_setup.py

@@ -0,0 +1,83 @@
+# Copyright 2022 (C) Friedrich Miescher Institute for Biomedical Research and
+# University of Zurich
+#
+# Original authors:
+# Jacopo Nespolo <jacopo.nespolo@exact-lab.it>
+# Tommaso Comparin <tommaso.comparin@exact-lab.it>
+#
+# This file is part of Fractal and was originally developed by eXact lab S.r.l.
+# <exact-lab.it> under contract with Liberali Lab from the Friedrich Miescher
+# Institute for Biomedical Research and Pelkmans Lab from the University of
+# Zurich.
+"""
+Submodule to define _slurm_submit_setup, which is also the reference
+implementation of `submit_setup_call`.
+"""
+from pathlib import Path
+from typing import Literal
+
+from ...task_files import get_task_file_paths
+from .get_slurm_config import get_slurm_config
+from fractal_server.app.models.v2 import WorkflowTaskV2
+
+
+def _slurm_submit_setup(
+    *,
+    wftask: WorkflowTaskV2,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Path,
+    which_type: Literal["non_parallel", "parallel"],
+) -> dict[str, object]:
+    """
+    Collect WorfklowTask-specific configuration parameters from different
+    sources, and inject them for execution.
+
+    Here goes all the logic for reading attributes from the appropriate sources
+    and transforming them into an appropriate `SlurmConfig` object (encoding
+    SLURM configuration) and `TaskFiles` object (with details e.g. about file
+    paths or filename prefixes).
+
+    For now, this is the reference implementation for the argument
+    `submit_setup_call` of
+    [fractal_server.app.runner.v2.runner][].
+
+    Arguments:
+        wftask:
+            WorkflowTask for which the configuration is to be assembled
+        workflow_dir_local:
+            Server-owned directory to store all task-execution-related relevant
+            files (inputs, outputs, errors, and all meta files related to the
+            job execution). Note: users cannot write directly to this folder.
+        workflow_dir_remote:
+            User-side directory with the same scope as `workflow_dir_local`,
+            and where a user can write.
+
+    Returns:
+        submit_setup_dict:
+            A dictionary that will be passed on to
+            `FractalSlurmExecutor.submit` and `FractalSlurmExecutor.map`, so
+            as to set extra options.
+    """
+
+    # Get SlurmConfig object
+    slurm_config = get_slurm_config(
+        wftask=wftask,
+        workflow_dir_local=workflow_dir_local,
+        workflow_dir_remote=workflow_dir_remote,
+        which_type=which_type,
+    )
+
+    # Get TaskFiles object
+    task_files = get_task_file_paths(
+        workflow_dir_local=workflow_dir_local,
+        workflow_dir_remote=workflow_dir_remote,
+        task_order=wftask.order,
+        task_name=wftask.task.name,
+    )
+
+    # Prepare and return output dictionary
+    submit_setup_dict = dict(
+        slurm_config=slurm_config,
+        task_files=task_files,
+    )
+    return submit_setup_dict

fractal_server/app/runner/v2/_slurm_ssh/get_slurm_config.py

@@ -0,0 +1,182 @@
+from pathlib import Path
+from typing import Literal
+from typing import Optional
+
+from fractal_server.app.models.v2 import WorkflowTaskV2
+from fractal_server.app.runner.executors.slurm._slurm_config import (
+    _parse_mem_value,
+)
+from fractal_server.app.runner.executors.slurm._slurm_config import (
+    load_slurm_config_file,
+)
+from fractal_server.app.runner.executors.slurm._slurm_config import logger
+from fractal_server.app.runner.executors.slurm._slurm_config import SlurmConfig
+from fractal_server.app.runner.executors.slurm._slurm_config import (
+    SlurmConfigError,
+)
+
+
+def get_slurm_config(
+    wftask: WorkflowTaskV2,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Path,
+    which_type: Literal["non_parallel", "parallel"],
+    config_path: Optional[Path] = None,
+) -> SlurmConfig:
+    """
+    Prepare a `SlurmConfig` configuration object
+
+    The argument `which_type` determines whether we use `wftask.meta_parallel`
+    or `wftask.meta_non_parallel`. In the following descritpion, let us assume
+    that `which_type="parallel"`.
+
+    The sources for `SlurmConfig` attributes, in increasing priority order, are
+
+    1. The general content of the Fractal SLURM configuration file.
+    2. The GPU-specific content of the Fractal SLURM configuration file, if
+        appropriate.
+    3. Properties in `wftask.meta_parallel` (which typically include those in
+       `wftask.task.meta_parallel`). Note that `wftask.meta_parallel` may be
+       `None`.
+
+    Arguments:
+        wftask:
+            WorkflowTask for which the SLURM configuration is is to be
+            prepared.
+        workflow_dir_local:
+            Server-owned directory to store all task-execution-related relevant
+            files (inputs, outputs, errors, and all meta files related to the
+            job execution). Note: users cannot write directly to this folder.
+        workflow_dir_remote:
+            User-side directory with the same scope as `workflow_dir_local`,
+            and where a user can write.
+        config_path:
+            Path of a Fractal SLURM configuration file; if `None`, use
+            `FRACTAL_SLURM_CONFIG_FILE` variable from settings.
+        which_type:
+            Determines whether to use `meta_parallel` or `meta_non_parallel`.
+
+    Returns:
+        slurm_config:
+            The SlurmConfig object
+    """
+
+    if which_type == "non_parallel":
+        wftask_meta = wftask.meta_non_parallel
+    elif which_type == "parallel":
+        wftask_meta = wftask.meta_parallel
+    else:
+        raise ValueError(
+            f"get_slurm_config received invalid argument {which_type=}."
+        )
+
+    logger.debug(
+        f"[get_slurm_config] WorkflowTask meta attribute: {wftask_meta=}"
+    )
+
+    # Incorporate slurm_env.default_slurm_config
+    slurm_env = load_slurm_config_file(config_path=config_path)
+    slurm_dict = slurm_env.default_slurm_config.dict(
+        exclude_unset=True, exclude={"mem"}
+    )
+    if slurm_env.default_slurm_config.mem:
+        slurm_dict["mem_per_task_MB"] = slurm_env.default_slurm_config.mem
+
+    # Incorporate slurm_env.batching_config
+    for key, value in slurm_env.batching_config.dict().items():
+        slurm_dict[key] = value
+
+    # Incorporate slurm_env.user_local_exports
+    slurm_dict["user_local_exports"] = slurm_env.user_local_exports
+
+    logger.debug(
+        "[get_slurm_config] Fractal SLURM configuration file: "
+        f"{slurm_env.dict()=}"
+    )
+
+    # GPU-related options
+    # Notes about priority:
+    # 1. This block of definitions takes priority over other definitions from
+    #    slurm_env which are not under the `needs_gpu` subgroup
+    # 2. This block of definitions has lower priority than whatever comes next
+    #    (i.e. from WorkflowTask.meta).
+    if wftask_meta is not None:
+        needs_gpu = wftask_meta.get("needs_gpu", False)
+    else:
+        needs_gpu = False
+    logger.debug(f"[get_slurm_config] {needs_gpu=}")
+    if needs_gpu and slurm_env.gpu_slurm_config is not None:  # FIXME
+        for key, value in slurm_env.gpu_slurm_config.dict(
+            exclude_unset=True, exclude={"mem"}
+        ).items():
+            slurm_dict[key] = value
+        if slurm_env.gpu_slurm_config.mem:
+            slurm_dict["mem_per_task_MB"] = slurm_env.gpu_slurm_config.mem
+
+    # Number of CPUs per task, for multithreading
+    if wftask_meta is not None and "cpus_per_task" in wftask_meta:
+        cpus_per_task = int(wftask_meta["cpus_per_task"])
+        slurm_dict["cpus_per_task"] = cpus_per_task
+
+    # Required memory per task, in MB
+    if wftask_meta is not None and "mem" in wftask_meta:
+        raw_mem = wftask_meta["mem"]
+        mem_per_task_MB = _parse_mem_value(raw_mem)
+        slurm_dict["mem_per_task_MB"] = mem_per_task_MB
+
+    # Job name
+    if wftask.is_legacy_task:
+        job_name = wftask.task_legacy.name.replace(" ", "_")
+    else:
+        job_name = wftask.task.name.replace(" ", "_")
+    slurm_dict["job_name"] = job_name
+
+    # Optional SLURM arguments and extra lines
+    if wftask_meta is not None:
+        account = wftask_meta.get("account", None)
+        if account is not None:
+            error_msg = (
+                f"Invalid {account=} property in WorkflowTask `meta` "
+                "attribute.\n"
+                "SLURM account must be set in the request body of the "
+                "apply-workflow endpoint, or by modifying the user properties."
+            )
+            logger.error(error_msg)
+            raise SlurmConfigError(error_msg)
+        for key in ["time", "gres", "constraint"]:
+            value = wftask_meta.get(key, None)
+            if value:
+                slurm_dict[key] = value
+    if wftask_meta is not None:
+        extra_lines = wftask_meta.get("extra_lines", [])
+    else:
+        extra_lines = []
+    extra_lines = slurm_dict.get("extra_lines", []) + extra_lines
+    if len(set(extra_lines)) != len(extra_lines):
+        logger.debug(
+            "[get_slurm_config] Removing repeated elements "
+            f"from {extra_lines=}."
+        )
+        extra_lines = list(set(extra_lines))
+    slurm_dict["extra_lines"] = extra_lines
+
+    # Job-batching parameters (if None, they will be determined heuristically)
+    if wftask_meta is not None:
+        tasks_per_job = wftask_meta.get("tasks_per_job", None)
+        parallel_tasks_per_job = wftask_meta.get(
+            "parallel_tasks_per_job", None
+        )
+    else:
+        tasks_per_job = None
+        parallel_tasks_per_job = None
+    slurm_dict["tasks_per_job"] = tasks_per_job
+    slurm_dict["parallel_tasks_per_job"] = parallel_tasks_per_job
+
+    # Put everything together
+    logger.debug(
+        "[get_slurm_config] Now create a SlurmConfig object based "
+        f"on {slurm_dict=}"
+    )
+    slurm_config = SlurmConfig(**slurm_dict)
+
+    return slurm_config

fractal_server/app/runner/v2/runner_functions_low_level.py

@@ -34,17 +34,15 @@ def _call_command_wrapper(cmd: str, log_path: Path) -> None:
         )
         raise TaskExecutionError(msg)
 
-    fp_log = open(log_path, "w")
-    try:
-        result = subprocess.run(  # nosec
-            shlex_split(cmd),
-            stderr=fp_log,
-            stdout=fp_log,
-        )
-    except Exception as e:
-        raise e
-    finally:
-        fp_log.close()
+    with open(log_path, "w") as fp_log:
+        try:
+            result = subprocess.run(  # nosec
+                shlex_split(cmd),
+                stderr=fp_log,
+                stdout=fp_log,
+            )
+        except Exception as e:
+            raise e
 
     if result.returncode > 0:
         with log_path.open("r") as fp_stderr:

fractal_server/app/runner/versions.py

@@ -0,0 +1,30 @@
+import json
+import sys
+from typing import Union
+
+import cloudpickle
+
+import fractal_server
+
+
+def get_versions() -> dict[str, Union[list[int], str]]:
+    """
+    Extract versions of Python, fractal-server and cloudpickle.
+
+    This information is useful to check compatibility of two Python
+    interpreters when running tasks: the current interpreter (which prepares
+    the input pickles and orchestrates workflow execution) and a remote
+    interpreter (e.g. the one defined in the `FRACTAL_SLURM_WORKER_PYTHON`
+    configuration variable) that executes the tasks.
+    """
+    versions = dict(
+        python=list(sys.version_info[:3]),
+        cloudpickle=cloudpickle.__version__,
+        fractal_server=fractal_server.__VERSION__,
+    )
+    return versions
+
+
+if __name__ == "__main__":
+    versions = get_versions()
+    print(json.dumps(versions))

fractal_server/app/schemas/v1/__init__.py

@@ -17,6 +17,7 @@ from .manifest import TaskManifestV1  # noqa: F401
 from .project import ProjectCreateV1  # noqa: F401
 from .project import ProjectReadV1  # noqa: F401
 from .project import ProjectUpdateV1  # noqa: F401
+from .state import StateRead  # noqa: F401
 from .task import TaskCreateV1  # noqa: F401
 from .task import TaskImportV1  # noqa: F401
 from .task import TaskReadV1  # noqa: F401

fractal_server/app/schemas/{state.py → v1/state.py}

@@ -5,29 +5,10 @@ from typing import Optional
 from pydantic import BaseModel
 from pydantic import validator
 
-from ._validators import valutc
+from fractal_server.app.schemas._validators import valutc
 
-__all__ = (
-    "_StateBase",
-    "StateRead",
-)
 
-
-class _StateBase(BaseModel):
-    """
-    Base class for `State`.
-
-    Attributes:
-        id: Primary key
-        data: Content of the state
-        timestamp: Time stamp of the state
-    """
-
-    data: dict[str, Any]
-    timestamp: datetime
-
-
-class StateRead(_StateBase):
+class StateRead(BaseModel):
     """
     Class for `State` read from database.
 
@@ -36,5 +17,7 @@ class StateRead(_StateBase):
     """
 
     id: Optional[int]
+    data: dict[str, Any]
+    timestamp: datetime
 
     _timestamp = validator("timestamp", allow_reuse=True)(valutc("timestamp"))

fractal_server/app/schemas/v2/__init__.py

@@ -13,6 +13,7 @@ from .job import JobReadV2  # noqa F401
 from .job import JobStatusTypeV2  # noqa F401
 from .job import JobUpdateV2  # noqa F401
 from .manifest import ManifestV2  # noqa F401
+from .manifest import TaskManifestV2  # noqa F401
 from .project import ProjectCreateV2  # noqa F401
 from .project import ProjectReadV2  # noqa F401
 from .project import ProjectUpdateV2  # noqa F401
@@ -22,8 +23,10 @@ from .task import TaskImportV2  # noqa F401
 from .task import TaskLegacyReadV2  # noqa F401
 from .task import TaskReadV2  # noqa F401
 from .task import TaskUpdateV2  # noqa F401
+from .task_collection import CollectionStateReadV2  # noqa F401
+from .task_collection import CollectionStatusV2  # noqa F401
+from .task_collection import TaskCollectCustomV2  # noqa F401
 from .task_collection import TaskCollectPipV2  # noqa F401
-from .task_collection import TaskCollectStatusV2  # noqa F401
 from .workflow import WorkflowCreateV2  # noqa F401
 from .workflow import WorkflowExportV2  # noqa F401
 from .workflow import WorkflowImportV2  # noqa F401

fractal_server/app/schemas/v2/task_collection.py

@@ -1,14 +1,26 @@
+from datetime import datetime
+from enum import Enum
 from pathlib import Path
+from typing import Any
 from typing import Literal
 from typing import Optional
 
 from pydantic import BaseModel
-from pydantic import Field
+from pydantic import root_validator
 from pydantic import validator
 
 from .._validators import valdictkeys
 from .._validators import valstr
-from .task import TaskReadV2
+from fractal_server.app.schemas._validators import valutc
+from fractal_server.app.schemas.v2 import ManifestV2
+
+
+class CollectionStatusV2(str, Enum):
+    PENDING = "pending"
+    INSTALLING = "installing"
+    COLLECTING = "collecting"
+    FAIL = "fail"
+    OK = "OK"
 
 
 class TaskCollectPipV2(BaseModel):
@@ -41,7 +53,7 @@ class TaskCollectPipV2(BaseModel):
     package: str
     package_version: Optional[str] = None
     package_extras: Optional[str] = None
-    python_version: Optional[str] = None
+    python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12"]] = None
     pinned_package_versions: Optional[dict[str, str]] = None
 
     _pinned_package_versions = validator(
@@ -70,40 +82,98 @@ class TaskCollectPipV2(BaseModel):
 
     @validator("package_version")
     def package_version_validator(cls, v, values):
-
-        valstr("package_version")(v)
-
+        v = valstr("package_version")(v)
         if values["package"].endswith(".whl"):
             raise ValueError(
-                "Cannot provide version when package is a Wheel file."
+                "Cannot provide package version when package is a wheel file."
             )
         return v
 
 
-class TaskCollectStatusV2(BaseModel):
+class TaskCollectCustomV2(BaseModel):
     """
-    TaskCollectStatus class
-
     Attributes:
-        status:
-        package:
-        venv_path:
-        task_list:
-        log:
-        info:
+        manifest: Manifest of a Fractal task package (this is typically the
+            content of `__FRACTAL_MANIFEST__.json`).
+        python_interpreter: Absolute path to the Python interpreter to be used
+            for running tasks.
+        source: A common label identifying this package.
+        package_root: The folder where the package is installed.
+            If not provided, it will be extracted via `pip show`
+            (requires `package_name` to be set).
+        package_name: Name of the package, as used for `import <package_name>`;
+            this is then used to extract the package directory (`package_root`)
+            via `pip show <package_name>`.
+        version: Optional version of tasks to be collected.
     """
 
-    status: Literal["pending", "installing", "collecting", "fail", "OK"]
-    package: str
-    venv_path: Path
-    task_list: Optional[list[TaskReadV2]] = Field(default=[])
-    log: Optional[str]
-    info: Optional[str]
+    manifest: ManifestV2
+    python_interpreter: str
+    source: str
+    package_root: Optional[str]
+    package_name: Optional[str]
+    version: Optional[str]
+
+    @root_validator(pre=True)
+    def one_of_package_root_or_name(cls, values):
+        package_root = values["package_root"]
+        package_name = values["package_name"]
+        if (package_root is None and package_name is None) or (
+            package_root is not None and package_name is not None
+        ):
+            raise ValueError(
+                "One and only one must be set between "
+                "'package_root' and 'package_name'"
+            )
+        return values
 
-    def sanitised_dict(self):
+    @validator("package_name")
+    def package_name_prevent_injection(cls, value: str):
         """
-        Return `self.dict()` after casting `self.venv_path` to a string
+        Remove all whitespace characters, and reject values containing `;`.
         """
-        d = self.dict()
-        d["venv_path"] = str(self.venv_path)
-        return d
+        if value is not None:
+            if ";" in value:
+                raise ValueError(f"Invalid package_name: {value}")
+            value = value.replace(" ", "")
+        return value
+
+    @validator("package_root")
+    def package_root_validator(cls, value):
+        if (value is not None) and (not Path(value).is_absolute()):
+            raise ValueError(
+                f"'package_root' must be an absolute path: (given {value})."
+            )
+        return value
+
+    @validator("python_interpreter")
+    def python_interpreter_validator(cls, value):
+        if not Path(value).is_absolute():
+            raise ValueError(
+                f"Python interpreter path must be absolute: (given {value})."
+            )
+        return value
+
+    # Valstr
+    _python_interpreter = validator("python_interpreter", allow_reuse=True)(
+        valstr("python_interpreter")
+    )
+    _source = validator("source", allow_reuse=True)(valstr("source"))
+    _package_root = validator("package_root", allow_reuse=True)(
+        valstr("package_root", accept_none=True)
+    )
+    _package_name = validator("package_name", allow_reuse=True)(
+        valstr("package_name", accept_none=True)
+    )
+    _version = validator("version", allow_reuse=True)(
+        valstr("version", accept_none=True)
+    )
+
+
+class CollectionStateReadV2(BaseModel):
+
+    id: Optional[int]
+    data: dict[str, Any]
+    timestamp: datetime
+
+    _timestamp = validator("timestamp", allow_reuse=True)(valutc("timestamp"))