fractal-server 2.3.0a2__py3-none-any.whl → 2.3.1__py3-none-any.whl

fractal_server/__init__.py

@@ -1 +1 @@
-__VERSION__ = "2.3.0a2"
+__VERSION__ = "2.3.1"

fractal_server/app/routes/api/v1/task_collection.py

@@ -25,8 +25,8 @@ from ....schemas.v1 import TaskCollectStatusV1
 from ....security import current_active_user
 from ....security import current_active_verified_user
 from ....security import User
+from fractal_server.string_tools import slugify_task_name_for_source
 from fractal_server.tasks.utils import get_collection_log
-from fractal_server.tasks.utils import slugify_task_name
 from fractal_server.tasks.v1._TaskCollectPip import _TaskCollectPip
 from fractal_server.tasks.v1.background_operations import (
     background_collect_pip,
@@ -159,7 +159,7 @@ async def collect_tasks_pip(
 
     # Check that tasks are not already in the DB
     for new_task in task_pkg.package_manifest.task_list:
-        new_task_name_slug = slugify_task_name(new_task.name)
+        new_task_name_slug = slugify_task_name_for_source(new_task.name)
         new_task_source = f"{task_pkg.package_source}:{new_task_name_slug}"
         stm = select(Task).where(Task.source == new_task_source)
         res = await db.execute(stm)

fractal_server/app/routes/api/v2/__init__.py

@@ -12,7 +12,6 @@ from .submit import router as submit_job_router_v2
 from .task import router as task_router_v2
 from .task_collection import router as task_collection_router_v2
 from .task_collection_custom import router as task_collection_router_v2_custom
-from .task_collection_ssh import router as task_collection_router_v2_ssh
 from .task_legacy import router as task_legacy_router_v2
 from .workflow import router as workflow_router_v2
 from .workflowtask import router as workflowtask_router_v2
@@ -30,21 +29,14 @@ router_api_v2.include_router(submit_job_router_v2, tags=["V2 Job"])
 
 
 settings = Inject(get_settings)
-if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
-    router_api_v2.include_router(
-        task_collection_router_v2_ssh,
-        prefix="/task",
-        tags=["V2 Task Collection"],
-    )
-else:
-    router_api_v2.include_router(
-        task_collection_router_v2, prefix="/task", tags=["V2 Task Collection"]
-    )
-    router_api_v2.include_router(
-        task_collection_router_v2_custom,
-        prefix="/task",
-        tags=["V2 Task Collection"],
-    )
+router_api_v2.include_router(
+    task_collection_router_v2, prefix="/task", tags=["V2 Task Collection"]
+)
+router_api_v2.include_router(
+    task_collection_router_v2_custom,
+    prefix="/task",
+    tags=["V2 Task Collection"],
+)
 router_api_v2.include_router(task_router_v2, prefix="/task", tags=["V2 Task"])
 router_api_v2.include_router(
     task_legacy_router_v2, prefix="/task-legacy", tags=["V2 Task Legacy"]

fractal_server/app/routes/api/v2/submit.py

@@ -246,7 +246,7 @@ async def apply_workflow(
         worker_init=job.worker_init,
         slurm_user=user.slurm_user,
         user_cache_dir=user.cache_dir,
-        connection=request.app.state.connection,
+        fractal_ssh=request.app.state.fractal_ssh,
     )
     request.app.state.jobsV2.append(job.id)
     logger.info(

fractal_server/app/routes/api/v2/task_collection.py

@@ -7,6 +7,7 @@ from fastapi import APIRouter
 from fastapi import BackgroundTasks
 from fastapi import Depends
 from fastapi import HTTPException
+from fastapi import Request
 from fastapi import Response
 from fastapi import status
 from pydantic.error_wrappers import ValidationError
@@ -27,10 +28,10 @@ from ....schemas.v2 import TaskReadV2
 from ....security import current_active_user
 from ....security import current_active_verified_user
 from ....security import User
+from fractal_server.string_tools import slugify_task_name_for_source
 from fractal_server.tasks.utils import get_absolute_venv_path
 from fractal_server.tasks.utils import get_collection_log
 from fractal_server.tasks.utils import get_collection_path
-from fractal_server.tasks.utils import slugify_task_name
 from fractal_server.tasks.v2._TaskCollectPip import _TaskCollectPip
 from fractal_server.tasks.v2.background_operations import (
     background_collect_pip,
@@ -38,6 +39,7 @@ from fractal_server.tasks.v2.background_operations import (
 from fractal_server.tasks.v2.endpoint_operations import create_package_dir_pip
 from fractal_server.tasks.v2.endpoint_operations import download_package
 from fractal_server.tasks.v2.endpoint_operations import inspect_package
+from fractal_server.tasks.v2.utils import get_python_interpreter_v2
 
 
 router = APIRouter()
@@ -66,6 +68,7 @@ async def collect_tasks_pip(
     task_collect: TaskCollectPipV2,
     background_tasks: BackgroundTasks,
     response: Response,
+    request: Request,
     user: User = Depends(current_active_verified_user),
     db: AsyncSession = Depends(get_async_db),
 ) -> CollectionStateReadV2:
@@ -76,17 +79,26 @@ async def collect_tasks_pip(
     of a package and the collection of tasks as advertised in the manifest.
     """
 
-    logger = set_logger(logger_name="collect_tasks_pip")
+    # Get settings
+    settings = Inject(get_settings)
 
-    # Set default python version
+    # Set/check python version
     if task_collect.python_version is None:
-        settings = Inject(get_settings)
         task_collect.python_version = (
             settings.FRACTAL_TASKS_PYTHON_DEFAULT_VERSION
         )
+    try:
+        get_python_interpreter_v2(python_version=task_collect.python_version)
+    except ValueError:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                f"Python version {task_collect.python_version} is "
+                "not available for Fractal task collection."
+            ),
+        )
 
-    # Validate payload as _TaskCollectPip, which has more strict checks than
-    # TaskCollectPip
+    # Validate payload
     try:
         task_pkg = _TaskCollectPip(**task_collect.dict(exclude_unset=True))
     except ValidationError as e:
@@ -95,6 +107,37 @@ async def collect_tasks_pip(
             detail=f"Invalid task-collection object. Original error: {e}",
         )
 
+    # END of SSH/non-SSH common part
+
+    if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+
+        from fractal_server.tasks.v2.background_operations_ssh import (
+            background_collect_pip_ssh,
+        )
+
+        # Construct and return state
+        state = CollectionStateV2(
+            data=dict(
+                status=CollectionStatusV2.PENDING, package=task_collect.package
+            )
+        )
+        db.add(state)
+        await db.commit()
+
+        background_tasks.add_task(
+            background_collect_pip_ssh,
+            state.id,
+            task_pkg,
+            request.app.state.fractal_ssh,
+        )
+
+        response.status_code = status.HTTP_201_CREATED
+        return state
+
+    # Actual non-SSH endpoint
+
+    logger = set_logger(logger_name="collect_tasks_pip")
+
     with TemporaryDirectory() as tmpdir:
         try:
             # Copy or download the package wheel file to tmpdir
@@ -197,7 +240,7 @@ async def collect_tasks_pip(
 
     # Check that tasks are not already in the DB
     for new_task in task_pkg.package_manifest.task_list:
-        new_task_name_slug = slugify_task_name(new_task.name)
+        new_task_name_slug = slugify_task_name_for_source(new_task.name)
         new_task_source = f"{task_pkg.package_source}:{new_task_name_slug}"
         stm = select(TaskV2).where(TaskV2.source == new_task_source)
         res = await db.execute(stm)
@@ -253,6 +296,7 @@ async def check_collection_status(
     """
     Check status of background task collection
     """
+
     logger = set_logger(logger_name="check_collection_status")
     logger.debug(f"Querying state for state.id={state_id}")
     state = await db.get(CollectionStateV2, state_id)
@@ -263,17 +307,28 @@ async def check_collection_status(
             detail=f"No task collection info with id={state_id}",
         )
 
-    # In some cases (i.e. a successful or ongoing task collection),
-    # state.data.log is not set; if so, we collect the current logs.
-    if verbose and not state.data.get("log"):
-        if "venv_path" not in state.data.keys():
-            await db.close()
-            raise HTTPException(
-                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
-                detail=f"No 'venv_path' in CollectionStateV2[{state_id}].data",
+    settings = Inject(get_settings)
+    if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+        # FIXME SSH: add logic for when data.state["log"] is empty
+        pass
+    else:
+        # Non-SSH mode
+        # In some cases (i.e. a successful or ongoing task collection),
+        # state.data.log is not set; if so, we collect the current logs.
+        if verbose and not state.data.get("log"):
+            if "venv_path" not in state.data.keys():
+                await db.close()
+                raise HTTPException(
+                    status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                    detail=(
+                        f"No 'venv_path' in CollectionStateV2[{state_id}].data"
+                    ),
+                )
+            state.data["log"] = get_collection_log(
+                Path(state.data["venv_path"])
             )
-        state.data["log"] = get_collection_log(Path(state.data["venv_path"]))
-        state.data["venv_path"] = str(state.data["venv_path"])
+            state.data["venv_path"] = str(state.data["venv_path"])
+
     reset_logger_handlers(logger)
     await db.close()
     return state

fractal_server/app/routes/api/v2/task_collection_custom.py

@@ -42,13 +42,35 @@ async def collect_task_custom(
 
     settings = Inject(get_settings)
 
-    if task_collect.package_root is None:
-
-        if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+    if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+        if task_collect.package_root is None:
             raise HTTPException(
                 status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                 detail="Cannot infer 'package_root' with 'slurm_ssh' backend.",
             )
+    else:
+        if not Path(task_collect.python_interpreter).is_file():
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f"{task_collect.python_interpreter=} "
+                    "doesn't exist or is not a file."
+                ),
+            )
+        if (
+            task_collect.package_root is not None
+            and not Path(task_collect.package_root).is_dir()
+        ):
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f"{task_collect.package_root=} "
+                    "doesn't exist or is not a directory."
+                ),
+            )
+
+    if task_collect.package_root is None:
+
         package_name_underscore = task_collect.package_name.replace("-", "_")
         # Note that python_command is then used as part of a subprocess.run
         # statement: be careful with mixing `'` and `"`.
@@ -61,7 +83,7 @@ async def collect_task_custom(
         )
         logger.debug(
             f"Now running {python_command=} through "
-            "{task_collect.python_interpreter}."
+            f"{task_collect.python_interpreter}."
         )
         res = subprocess.run(  # nosec
             shlex.split(

fractal_server/app/runner/executors/slurm/ssh/executor.py

@@ -27,7 +27,6 @@ from typing import Sequence
 
 import cloudpickle
 from cfut import SlurmExecutor
-from fabric.connection import Connection
 from paramiko.ssh_exception import NoValidConnectionsError
 
 from ....filenames import SHUTDOWN_FILENAME
@@ -44,8 +43,7 @@ from fractal_server.app.runner.exceptions import TaskExecutionError
 from fractal_server.app.runner.executors.slurm.ssh._slurm_job import SlurmJob
 from fractal_server.config import get_settings
 from fractal_server.logger import set_logger
-from fractal_server.ssh._fabric import check_connection
-from fractal_server.ssh._fabric import run_command_over_ssh
+from fractal_server.ssh._fabric import FractalSSH
 from fractal_server.syringe import Inject
 
 logger = set_logger(__name__)
@@ -58,7 +56,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
     FIXME: docstring
 
     Attributes:
-        connection: SSH connection
+        fractal_ssh: FractalSSH connection with custom lock
         shutdown_file:
         python_remote: Equal to `settings.FRACTAL_SLURM_WORKER_PYTHON`
         wait_thread_cls: Class for waiting thread
@@ -76,7 +74,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
             Dictionary with paths of slurm-related files for active jobs
     """
 
-    connection: Connection
+    fractal_ssh: FractalSSH
 
     workflow_dir_local: Path
     workflow_dir_remote: Path
@@ -95,8 +93,8 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
     def __init__(
         self,
         *,
-        # SSH connection
-        connection: Connection,
+        # FractalSSH connection
+        fractal_ssh: FractalSSH,
         # Folders and files
         workflow_dir_local: Path,
         workflow_dir_remote: Path,
@@ -117,7 +115,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
         relevant bits of `cfut.ClusterExecutor.__init__`.
 
         Args:
-            connection:
+            fractal_ssh:
             workflow_dir_local:
             workflow_dir_remote:
             keep_pickle_files:
@@ -167,8 +165,8 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
             raise ValueError("FRACTAL_SLURM_WORKER_PYTHON is not set. Exit.")
 
         # Initialize connection and perform handshake
-        self.connection = connection
-        logger.warning(self.connection)
+        self.fractal_ssh = fractal_ssh
+        logger.warning(self.fractal_ssh)
         self.handshake()
 
         # Set/validate parameters for SLURM submission scripts
@@ -838,7 +836,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
 
         # Transfer archive
         t_0_put = time.perf_counter()
-        self.connection.put(
+        self.fractal_ssh.put(
             local=tarfile_path_local,
             remote=tarfile_path_remote,
         )
@@ -853,7 +851,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
             "fractal_server.app.runner.extract_archive "
             f"{tarfile_path_remote}"
         )
-        run_command_over_ssh(cmd=tar_command, connection=self.connection)
+        self.fractal_ssh.run_command(cmd=tar_command)
 
         # Remove local version
         t_0_rm = time.perf_counter()
@@ -875,9 +873,8 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
 
         # Submit job to SLURM, and get jobid
         sbatch_command = f"sbatch --parsable {job.slurm_script_remote}"
-        sbatch_stdout = run_command_over_ssh(
+        sbatch_stdout = self.fractal_ssh.run_command(
             cmd=sbatch_command,
-            connection=self.connection,
         )
 
         # Extract SLURM job ID from stdout
@@ -1216,9 +1213,9 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
 
         # Remove remote tarfile - FIXME SSH: is this needed?
         # rm_command = f"rm {tarfile_path_remote}"
-        # _run_command_over_ssh(cmd=rm_command, connection=self.connection)
+        # _run_command_over_ssh(cmd=rm_command, fractal_ssh=self.fractal_ssh)
         logger.warning(f"Unlink {tarfile_path_remote=} - START")
-        self.connection.sftp().unlink(tarfile_path_remote)
+        self.fractal_ssh.sftp().unlink(tarfile_path_remote)
         logger.warning(f"Unlink {tarfile_path_remote=} - STOP")
 
         # Create remote tarfile
@@ -1227,14 +1224,12 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
             "-m fractal_server.app.runner.compress_folder "
             f"{(self.workflow_dir_remote / subfolder_name).as_posix()}"
         )
-        stdout = run_command_over_ssh(
-            cmd=tar_command, connection=self.connection
-        )
+        stdout = self.fractal_ssh.run_command(cmd=tar_command)
         print(stdout)
 
         # Fetch tarfile
         t_0_get = time.perf_counter()
-        self.connection.get(
+        self.fractal_ssh.get(
             remote=tarfile_path_remote,
             local=tarfile_path_local,
         )
@@ -1331,7 +1326,6 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
         """
 
         logger.debug("Executor shutdown: start")
-        # self.connection.close()
 
         # Handle all job futures
         slurm_jobs_to_scancel = []
@@ -1354,9 +1348,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
             scancel_string = " ".join(slurm_jobs_to_scancel)
             logger.warning(f"Now scancel-ing SLURM jobs {scancel_string}")
             scancel_command = f"scancel {scancel_string}"
-            run_command_over_ssh(
-                cmd=scancel_command, connection=self.connection
-            )
+            self.fractal_ssh.run_command(cmd=scancel_command)
         logger.debug("Executor shutdown: end")
 
     def __exit__(self, *args, **kwargs):
@@ -1381,10 +1373,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
         )
         job_ids = ",".join([str(j) for j in job_ids])
         squeue_command = squeue_command.replace("__JOBS__", job_ids)
-        stdout = run_command_over_ssh(
-            cmd=squeue_command,
-            connection=self.connection,
-        )
+        stdout = self.fractal_ssh.run_command(cmd=squeue_command)
         return stdout
 
     def _jobs_finished(self, job_ids: list[str]) -> set[str]:
@@ -1458,13 +1447,13 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
         FIXME SSH: We could include further checks on version matches
         """
 
-        check_connection(self.connection)
+        self.fractal_ssh.check_connection()
 
         t_start_handshake = time.perf_counter()
 
         logger.info("[FractalSlurmSSHExecutor.ssh_handshake] START")
         cmd = f"{self.python_remote} -m fractal_server.app.runner.versions"
-        stdout = run_command_over_ssh(cmd=cmd, connection=self.connection)
+        stdout = self.fractal_ssh.run_command(cmd=cmd)
         remote_versions = json.loads(stdout.strip("\n"))
 
         # Check compatibility with local versions

fractal_server/app/runner/task_files.py

@@ -2,18 +2,7 @@ from pathlib import Path
 from typing import Optional
 from typing import Union
 
-from fractal_server.tasks.utils import slugify_task_name
-
-
-def sanitize_component(value: str) -> str:
-    """
-    Remove {" ", "/", "."} form a string, e.g. going from
-    'plate.zarr/B/03/0' to 'plate_zarr_B_03_0'.
-
-    Args:
-        value: Input strig
-    """
-    return value.replace(" ", "_").replace("/", "_").replace(".", "_")
+from fractal_server.string_tools import sanitize_string
 
 
 def task_subfolder_name(order: Union[int, str], task_name: str) -> str:
@@ -24,7 +13,7 @@ def task_subfolder_name(order: Union[int, str], task_name: str) -> str:
         order:
         task_name:
     """
-    task_name_slug = slugify_task_name(task_name)
+    task_name_slug = sanitize_string(task_name)
     return f"{order}_{task_name_slug}"
 
 
@@ -93,7 +82,7 @@ class TaskFiles:
         self.component = component
 
         if self.component is not None:
-            component_safe = sanitize_component(str(self.component))
+            component_safe = sanitize_string(str(self.component))
             component_safe = f"_par_{component_safe}"
         else:
             component_safe = ""

fractal_server/app/runner/v2/__init__.py

@@ -11,7 +11,6 @@ import traceback
 from pathlib import Path
 from typing import Optional
 
-from fabric import Connection  # FIXME SSH: try/except import
 from sqlalchemy.orm import Session as DBSyncSession
 from sqlalchemy.orm.attributes import flag_modified
 
@@ -19,6 +18,7 @@ from ....config import get_settings
 from ....logger import get_logger
 from ....logger import reset_logger_handlers
 from ....logger import set_logger
+from ....ssh._fabric import FractalSSH
 from ....syringe import Inject
 from ....utils import get_timestamp
 from ...db import DB
@@ -79,7 +79,7 @@ async def submit_workflow(
     worker_init: Optional[str] = None,
     slurm_user: Optional[str] = None,
     user_cache_dir: Optional[str] = None,
-    connection: Optional[Connection] = None,
+    fractal_ssh: Optional[FractalSSH] = None,
 ) -> None:
     """
     Prepares a workflow and applies it to a dataset
@@ -189,11 +189,8 @@ async def submit_workflow(
                     / WORKFLOW_DIR_LOCAL.name
                 )
                 # FIXME SSH: move mkdir to executor, likely within handshake
-
-                from ....ssh._fabric import _mkdir_over_ssh
-
-                _mkdir_over_ssh(
-                    folder=str(WORKFLOW_DIR_REMOTE), connection=connection
+                fractal_ssh.mkdir(
+                    folder=str(WORKFLOW_DIR_REMOTE),
                 )
                 logging.info(f"Created {str(WORKFLOW_DIR_REMOTE)} via SSH.")
             else:
@@ -299,7 +296,7 @@ async def submit_workflow(
             )
         elif FRACTAL_RUNNER_BACKEND == "slurm_ssh":
             process_workflow = slurm_ssh_process_workflow
-            backend_specific_kwargs = dict(connection=connection)
+            backend_specific_kwargs = dict(fractal_ssh=fractal_ssh)
         else:
             raise RuntimeError(
                 f"Invalid runner backend {FRACTAL_RUNNER_BACKEND=}"

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

@@ -21,8 +21,7 @@ from typing import Any
 from typing import Optional
 from typing import Union
 
-from fabric import Connection
-
+from .....ssh._fabric import FractalSSH
 from ....models.v2 import DatasetV2
 from ....models.v2 import WorkflowV2
 from ...async_wrap import async_wrap
@@ -41,7 +40,7 @@ def _process_workflow(
     workflow_dir_remote: Path,
     first_task_index: int,
     last_task_index: int,
-    connection: Connection,
+    fractal_ssh: FractalSSH,
     worker_init: Optional[Union[str, list[str]]] = None,
 ) -> dict[str, Any]:
     """
@@ -62,7 +61,7 @@ def _process_workflow(
         worker_init = worker_init.split("\n")
 
     with FractalSlurmSSHExecutor(
-        connection=connection,
+        fractal_ssh=fractal_ssh,
         workflow_dir_local=workflow_dir_local,
         workflow_dir_remote=workflow_dir_remote,
         common_script_lines=worker_init,
@@ -91,7 +90,7 @@ async def process_workflow(
     last_task_index: Optional[int] = None,
     logger_name: str,
     # Not used
-    connection: Connection,
+    fractal_ssh: FractalSSH,
     user_cache_dir: Optional[str] = None,
     slurm_user: Optional[str] = None,
     slurm_account: Optional[str] = None,
@@ -121,6 +120,6 @@ async def process_workflow(
         first_task_index=first_task_index,
         last_task_index=last_task_index,
         worker_init=worker_init,
-        connection=connection,
+        fractal_ssh=fractal_ssh,
     )
     return new_dataset_attributes

fractal_server/main.py

@@ -101,14 +101,16 @@ async def lifespan(app: FastAPI):
 
     if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
         from fractal_server.ssh._fabric import get_ssh_connection
+        from fractal_server.ssh._fabric import FractalSSH
 
-        app.state.connection = get_ssh_connection()
+        connection = get_ssh_connection()
+        app.state.fractal_ssh = FractalSSH(connection=connection)
         logger.info(
             f"Created SSH connection "
-            f"({app.state.connection.is_connected=})."
+            f"({app.state.fractal_ssh.is_connected=})."
         )
     else:
-        app.state.connection = None
+        app.state.fractal_ssh = None
 
     config_uvicorn_loggers()
     logger.info("End application startup")
@@ -120,10 +122,10 @@ async def lifespan(app: FastAPI):
     if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
         logger.info(
             f"Closing SSH connection "
-            f"(current: {app.state.connection.is_connected=})."
+            f"(current: {app.state.fractal_ssh.is_connected=})."
         )
 
-        app.state.connection.close()
+        app.state.fractal_ssh.close()
 
     logger.info(
         f"Current worker with pid {os.getpid()} is shutting down. "