fractal-server 2.0.6__py3-none-any.whl → 2.2.0__py3-none-any.whl

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

@@ -0,0 +1,42 @@
+"""
+Submodule to define _local_submit_setup
+"""
+from pathlib import Path
+from typing import Literal
+from typing import Optional
+
+from ....models.v2 import WorkflowTaskV2
+from ._local_config import get_local_backend_config
+
+
+def _local_submit_setup(
+    *,
+    wftask: WorkflowTaskV2,
+    workflow_dir_local: Optional[Path] = None,
+    workflow_dir_remote: Optional[Path] = None,
+    which_type: Literal["non_parallel", "parallel"],
+) -> dict[str, object]:
+    """
+    Collect WorfklowTask-specific configuration parameters from different
+    sources, and inject them for execution.
+
+    Arguments:
+        wftask:
+            WorkflowTask for which the configuration is to be assembled
+        workflow_dir_local:
+            Not used in this function.
+        workflow_dir_remote:
+            Not used in this function.
+
+    Returns:
+        submit_setup_dict:
+            A dictionary that will be passed on to
+            `FractalProcessPoolExecutor.submit` and
+            `FractalProcessPoolExecutor.map`, so as to set extra options.
+    """
+
+    local_backend_config = get_local_backend_config(
+        wftask=wftask, which_type=which_type
+    )
+
+    return dict(local_backend_config=local_backend_config)

fractal_server/app/runner/v2/_local_experimental/executor.py

@@ -0,0 +1,156 @@
+"""
+Custom version of Python
+[ProcessPoolExecutor](https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ProcessPoolExecutor)).
+"""
+import threading
+import time
+from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures.process import BrokenProcessPool
+from pathlib import Path
+from typing import Callable
+from typing import Iterable
+from typing import Optional
+from typing import Sequence
+
+import psutil
+
+from ._local_config import get_default_local_backend_config
+from ._local_config import LocalBackendConfig
+from fractal_server.app.runner.exceptions import JobExecutionError
+from fractal_server.logger import set_logger
+
+logger = set_logger("FractalProcessPoolExecutor")
+
+
+class FractalProcessPoolExecutor(ProcessPoolExecutor):
+
+    shutdown_file: Path
+    interval: float
+    _shutdown: bool
+    _shutdown_file_thread: threading.Thread
+
+    def __init__(
+        self, shutdown_file: Path, interval: float = 1.0, *args, **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.shutdown_file = Path(shutdown_file)
+        self.interval = float(interval)
+        logger.debug(
+            f"Start monitoring {shutdown_file} every {interval} seconds"
+        )
+        self._shutdown = False
+        self._shutdown_file_thread = threading.Thread(
+            target=self._run, daemon=True
+        )
+        self._shutdown_file_thread.start()
+
+    def _run(self):
+        """
+        Running on '_shutdown_file_thread'.
+        """
+        while True:
+            if self.shutdown_file.exists() or self._shutdown:
+                try:
+                    self._terminate_processes()
+                except Exception as e:
+                    logger.error(
+                        "Terminate processes failed. "
+                        f"Original error: {str(e)}."
+                    )
+                finally:
+                    return
+            time.sleep(self.interval)
+
+    def _terminate_processes(self):
+        """
+        Running on '_shutdown_file_thread'.
+        """
+
+        logger.info("Start terminating FractalProcessPoolExecutor processes.")
+        # We use 'psutil' in order to easily access the PIDs of the children.
+        if self._processes is not None:
+            for pid in self._processes.keys():
+                parent = psutil.Process(pid)
+                children = parent.children(recursive=True)
+                for child in children:
+                    child.kill()
+                parent.kill()
+                logger.info(f"Process {pid} and its children terminated.")
+        logger.info("FractalProcessPoolExecutor processes terminated.")
+
+    def shutdown(self, *args, **kwargs) -> None:
+        self._shutdown = True
+        self._shutdown_file_thread.join()
+        return super().shutdown(*args, **kwargs)
+
+    def submit(
+        self,
+        *args,
+        local_backend_config: Optional[LocalBackendConfig] = None,
+        **kwargs,
+    ):
+        """
+        Compared to the `ProcessPoolExecutor` method, here we accept an
+        additional keyword argument (`local_backend_config`), which is then
+        simply ignored.
+        """
+        return super().submit(*args, **kwargs)
+
+    def map(
+        self,
+        fn: Callable,
+        *iterables: Sequence[Iterable],
+        local_backend_config: Optional[LocalBackendConfig] = None,
+    ):
+        """
+        Custom version of the `Executor.map` method
+
+        The main change with the respect to the original `map` method is that
+        the list of tasks to be executed is split into chunks, and then
+        `super().map` is called (sequentially) on each chunk. The goal of this
+        change is to limit parallelism, e.g. due to limited computational
+        resources.
+
+        Other changes from the `concurrent.futures` `map` method:
+
+        1. Removed `timeout` argument;
+        2. Removed `chunksize`;
+        3. All iterators (both inputs and output ones) are transformed into
+           lists.
+
+        Args:
+            fn: A callable function.
+            iterables: The argument iterables (one iterable per argument of
+                       `fn`).
+           local_backend_config: The backend configuration, needed to extract
+                                 `parallel_tasks_per_job`.
+        """
+        # Preliminary check
+        iterable_lengths = [len(it) for it in iterables]
+        if not len(set(iterable_lengths)) == 1:
+            raise ValueError("Iterables have different lengths.")
+        # Set total number of arguments
+        n_elements = len(iterables[0])
+
+        # Set parallel_tasks_per_job
+        if local_backend_config is None:
+            local_backend_config = get_default_local_backend_config()
+        parallel_tasks_per_job = local_backend_config.parallel_tasks_per_job
+        if parallel_tasks_per_job is None:
+            parallel_tasks_per_job = n_elements
+
+        # Execute tasks, in chunks of size parallel_tasks_per_job
+        results = []
+        for ind_chunk in range(0, n_elements, parallel_tasks_per_job):
+            chunk_iterables = [
+                it[ind_chunk : ind_chunk + parallel_tasks_per_job]  # noqa
+                for it in iterables
+            ]
+            map_iter = super().map(fn, *chunk_iterables)
+
+            try:
+                results.extend(list(map_iter))
+            except BrokenProcessPool as e:
+                raise JobExecutionError(info=e.args[0])
+
+        return iter(results)

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

@@ -40,8 +40,8 @@ def _process_workflow(
     workflow: WorkflowV2,
     dataset: DatasetV2,
     logger_name: str,
-    workflow_dir: Path,
-    workflow_dir_user: Path,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Path,
     first_task_index: int,
     last_task_index: int,
     slurm_user: Optional[str] = None,
@@ -76,8 +76,8 @@ def _process_workflow(
         keep_logs=True,
         slurm_user=slurm_user,
         user_cache_dir=user_cache_dir,
-        working_dir=workflow_dir,
-        working_dir_user=workflow_dir_user,
+        workflow_dir_local=workflow_dir_local,
+        workflow_dir_remote=workflow_dir_remote,
         common_script_lines=worker_init,
         slurm_account=slurm_account,
     ) as executor:
@@ -87,8 +87,8 @@ def _process_workflow(
             ],  # noqa
             dataset=dataset,
             executor=executor,
-            workflow_dir=workflow_dir,
-            workflow_dir_user=workflow_dir_user,
+            workflow_dir_local=workflow_dir_local,
+            workflow_dir_remote=workflow_dir_remote,
             logger_name=logger_name,
             submit_setup_call=_slurm_submit_setup,
         )
@@ -99,8 +99,8 @@ async def process_workflow(
     *,
     workflow: WorkflowV2,
     dataset: DatasetV2,
-    workflow_dir: Path,
-    workflow_dir_user: Optional[Path] = None,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Optional[Path] = None,
     first_task_index: Optional[int] = None,
     last_task_index: Optional[int] = None,
     logger_name: str,
@@ -129,8 +129,8 @@ async def process_workflow(
         workflow=workflow,
         dataset=dataset,
         logger_name=logger_name,
-        workflow_dir=workflow_dir,
-        workflow_dir_user=workflow_dir_user,
+        workflow_dir_local=workflow_dir_local,
+        workflow_dir_remote=workflow_dir_remote,
         first_task_index=first_task_index,
         last_task_index=last_task_index,
         user_cache_dir=user_cache_dir,

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

@@ -24,8 +24,8 @@ from fractal_server.app.models.v2 import WorkflowTaskV2
 def _slurm_submit_setup(
     *,
     wftask: WorkflowTaskV2,
-    workflow_dir: Path,
-    workflow_dir_user: Path,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Path,
     which_type: Literal["non_parallel", "parallel"],
 ) -> dict[str, object]:
     """
@@ -44,13 +44,13 @@ def _slurm_submit_setup(
     Arguments:
         wftask:
             WorkflowTask for which the configuration is to be assembled
-        workflow_dir:
+        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_user:
-            User-side directory with the same scope as `workflow_dir`, and
-            where a user can write.
+        workflow_dir_remote:
+            User-side directory with the same scope as `workflow_dir_local`,
+            and where a user can write.
 
     Returns:
         submit_setup_dict:
@@ -62,16 +62,17 @@ def _slurm_submit_setup(
     # Get SlurmConfig object
     slurm_config = get_slurm_config(
         wftask=wftask,
-        workflow_dir=workflow_dir,
-        workflow_dir_user=workflow_dir_user,
+        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=workflow_dir,
-        workflow_dir_user=workflow_dir_user,
+        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

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

@@ -18,8 +18,8 @@ from fractal_server.app.runner.executors.slurm._slurm_config import (
 
 def get_slurm_config(
     wftask: WorkflowTaskV2,
-    workflow_dir: Path,
-    workflow_dir_user: Path,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Path,
     which_type: Literal["non_parallel", "parallel"],
     config_path: Optional[Path] = None,
 ) -> SlurmConfig:
@@ -43,13 +43,13 @@ def get_slurm_config(
         wftask:
             WorkflowTask for which the SLURM configuration is is to be
             prepared.
-        workflow_dir:
+        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_user:
-            User-side directory with the same scope as `workflow_dir`, and
-            where a user can write.
+        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.

fractal_server/app/runner/v2/runner.py

@@ -32,16 +32,18 @@ def execute_tasks_v2(
     wf_task_list: list[WorkflowTaskV2],
     dataset: DatasetV2,
     executor: ThreadPoolExecutor,
-    workflow_dir: Path,
-    workflow_dir_user: Optional[Path] = None,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Optional[Path] = None,
     logger_name: Optional[str] = None,
     submit_setup_call: Callable = no_op_submit_setup_call,
 ) -> DatasetV2:
 
     logger = logging.getLogger(logger_name)
 
-    if not workflow_dir.exists():  # FIXME: this should have already happened
-        workflow_dir.mkdir()
+    if (
+        not workflow_dir_local.exists()
+    ):  # FIXME: this should have already happened
+        workflow_dir_local.mkdir()
 
     # Initialize local dataset attributes
     zarr_dir = dataset.zarr_dir
@@ -94,8 +96,8 @@ def execute_tasks_v2(
                     zarr_dir=zarr_dir,
                     wftask=wftask,
                     task=task,
-                    workflow_dir=workflow_dir,
-                    workflow_dir_user=workflow_dir_user,
+                    workflow_dir_local=workflow_dir_local,
+                    workflow_dir_remote=workflow_dir_remote,
                     executor=executor,
                     logger_name=logger_name,
                     submit_setup_call=submit_setup_call,
@@ -105,8 +107,8 @@ def execute_tasks_v2(
                     images=filtered_images,
                     wftask=wftask,
                     task=task,
-                    workflow_dir=workflow_dir,
-                    workflow_dir_user=workflow_dir_user,
+                    workflow_dir_local=workflow_dir_local,
+                    workflow_dir_remote=workflow_dir_remote,
                     executor=executor,
                     logger_name=logger_name,
                     submit_setup_call=submit_setup_call,
@@ -117,8 +119,8 @@ def execute_tasks_v2(
                     zarr_dir=zarr_dir,
                     wftask=wftask,
                     task=task,
-                    workflow_dir=workflow_dir,
-                    workflow_dir_user=workflow_dir_user,
+                    workflow_dir_local=workflow_dir_local,
+                    workflow_dir_remote=workflow_dir_remote,
                     executor=executor,
                     logger_name=logger_name,
                     submit_setup_call=submit_setup_call,
@@ -133,8 +135,8 @@ def execute_tasks_v2(
                 task_legacy=task_legacy,
                 executor=executor,
                 logger_name=logger_name,
-                workflow_dir=workflow_dir,
-                workflow_dir_user=workflow_dir_user,
+                workflow_dir_local=workflow_dir_local,
+                workflow_dir_remote=workflow_dir_remote,
                 submit_setup_call=submit_setup_call,
             )
 
@@ -323,11 +325,11 @@ def execute_tasks_v2(
         # temporary files which can be used (1) to retrieve the latest state
         # when the job fails, (2) from within endpoints that need up-to-date
         # information
-        with open(workflow_dir / HISTORY_FILENAME, "w") as f:
+        with open(workflow_dir_local / HISTORY_FILENAME, "w") as f:
             json.dump(tmp_history, f, indent=2)
-        with open(workflow_dir / FILTERS_FILENAME, "w") as f:
+        with open(workflow_dir_local / FILTERS_FILENAME, "w") as f:
             json.dump(tmp_filters, f, indent=2)
-        with open(workflow_dir / IMAGES_FILENAME, "w") as f:
+        with open(workflow_dir_local / IMAGES_FILENAME, "w") as f:
             json.dump(tmp_images, f, indent=2)
 
         logger.debug(f'END    {wftask.order}-th task (name="{task_name}")')

fractal_server/app/runner/v2/runner_functions.py

@@ -65,8 +65,8 @@ def _cast_and_validate_InitTaskOutput(
 def no_op_submit_setup_call(
     *,
     wftask: WorkflowTaskV2,
-    workflow_dir: Path,
-    workflow_dir_user: Path,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Path,
     which_type: Literal["non_parallel", "parallel"],
 ) -> dict:
     """
@@ -79,16 +79,16 @@ def no_op_submit_setup_call(
 def _get_executor_options(
     *,
     wftask: WorkflowTaskV2,
-    workflow_dir: Path,
-    workflow_dir_user: Path,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Path,
     submit_setup_call: Callable,
     which_type: Literal["non_parallel", "parallel"],
 ) -> dict:
     try:
         options = submit_setup_call(
             wftask=wftask,
-            workflow_dir=workflow_dir,
-            workflow_dir_user=workflow_dir_user,
+            workflow_dir_local=workflow_dir_local,
+            workflow_dir_remote=workflow_dir_remote,
             which_type=which_type,
         )
     except Exception as e:
@@ -115,8 +115,8 @@ def run_v2_task_non_parallel(
     zarr_dir: str,
     task: TaskV2,
     wftask: WorkflowTaskV2,
-    workflow_dir: Path,
-    workflow_dir_user: Optional[Path] = None,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Optional[Path] = None,
     executor: Executor,
     logger_name: Optional[str] = None,
     submit_setup_call: Callable = no_op_submit_setup_call,
@@ -125,17 +125,17 @@ def run_v2_task_non_parallel(
     This runs server-side (see `executor` argument)
     """
 
-    if workflow_dir_user is None:
-        workflow_dir_user = workflow_dir
+    if workflow_dir_remote is None:
+        workflow_dir_remote = workflow_dir_local
         logging.warning(
-            "In `run_single_task`, workflow_dir_user=None. Is this right?"
+            "In `run_single_task`, workflow_dir_remote=None. Is this right?"
         )
-        workflow_dir_user = workflow_dir
+        workflow_dir_remote = workflow_dir_local
 
     executor_options = _get_executor_options(
         wftask=wftask,
-        workflow_dir=workflow_dir,
-        workflow_dir_user=workflow_dir_user,
+        workflow_dir_local=workflow_dir_local,
+        workflow_dir_remote=workflow_dir_remote,
         submit_setup_call=submit_setup_call,
         which_type="non_parallel",
     )
@@ -150,8 +150,8 @@ def run_v2_task_non_parallel(
             run_single_task,
             wftask=wftask,
             command=task.command_non_parallel,
-            workflow_dir=workflow_dir,
-            workflow_dir_user=workflow_dir_user,
+            workflow_dir_local=workflow_dir_local,
+            workflow_dir_remote=workflow_dir_remote,
         ),
         function_kwargs,
         **executor_options,
@@ -169,8 +169,8 @@ def run_v2_task_parallel(
     task: TaskV2,
     wftask: WorkflowTaskV2,
     executor: Executor,
-    workflow_dir: Path,
-    workflow_dir_user: Optional[Path] = None,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Optional[Path] = None,
     logger_name: Optional[str] = None,
     submit_setup_call: Callable = no_op_submit_setup_call,
 ) -> TaskOutput:
@@ -182,8 +182,8 @@ def run_v2_task_parallel(
 
     executor_options = _get_executor_options(
         wftask=wftask,
-        workflow_dir=workflow_dir,
-        workflow_dir_user=workflow_dir_user,
+        workflow_dir_local=workflow_dir_local,
+        workflow_dir_remote=workflow_dir_remote,
         submit_setup_call=submit_setup_call,
         which_type="parallel",
     )
@@ -203,8 +203,8 @@ def run_v2_task_parallel(
             run_single_task,
             wftask=wftask,
             command=task.command_parallel,
-            workflow_dir=workflow_dir,
-            workflow_dir_user=workflow_dir_user,
+            workflow_dir_local=workflow_dir_local,
+            workflow_dir_remote=workflow_dir_remote,
         ),
         list_function_kwargs,
         **executor_options,
@@ -230,23 +230,23 @@ def run_v2_task_compound(
     task: TaskV2,
     wftask: WorkflowTaskV2,
     executor: Executor,
-    workflow_dir: Path,
-    workflow_dir_user: Optional[Path] = None,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Optional[Path] = None,
     logger_name: Optional[str] = None,
     submit_setup_call: Callable = no_op_submit_setup_call,
 ) -> TaskOutput:
 
     executor_options_init = _get_executor_options(
         wftask=wftask,
-        workflow_dir=workflow_dir,
-        workflow_dir_user=workflow_dir_user,
+        workflow_dir_local=workflow_dir_local,
+        workflow_dir_remote=workflow_dir_remote,
         submit_setup_call=submit_setup_call,
         which_type="non_parallel",
     )
     executor_options_compute = _get_executor_options(
         wftask=wftask,
-        workflow_dir=workflow_dir,
-        workflow_dir_user=workflow_dir_user,
+        workflow_dir_local=workflow_dir_local,
+        workflow_dir_remote=workflow_dir_remote,
         submit_setup_call=submit_setup_call,
         which_type="parallel",
     )
@@ -262,8 +262,8 @@ def run_v2_task_compound(
             run_single_task,
             wftask=wftask,
             command=task.command_non_parallel,
-            workflow_dir=workflow_dir,
-            workflow_dir_user=workflow_dir_user,
+            workflow_dir_local=workflow_dir_local,
+            workflow_dir_remote=workflow_dir_remote,
         ),
         function_kwargs,
         **executor_options_init,
@@ -298,8 +298,8 @@ def run_v2_task_compound(
             run_single_task,
             wftask=wftask,
             command=task.command_parallel,
-            workflow_dir=workflow_dir,
-            workflow_dir_user=workflow_dir_user,
+            workflow_dir_local=workflow_dir_local,
+            workflow_dir_remote=workflow_dir_remote,
         ),
         list_function_kwargs,
         **executor_options_compute,
@@ -325,8 +325,8 @@ def run_v1_task_parallel(
     task_legacy: TaskV1,
     wftask: WorkflowTaskV2,
     executor: Executor,
-    workflow_dir: Path,
-    workflow_dir_user: Optional[Path] = None,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Optional[Path] = None,
     logger_name: Optional[str] = None,
     submit_setup_call: Callable = no_op_submit_setup_call,
 ) -> TaskOutput:
@@ -335,8 +335,8 @@ def run_v1_task_parallel(
 
     executor_options = _get_executor_options(
         wftask=wftask,
-        workflow_dir=workflow_dir,
-        workflow_dir_user=workflow_dir_user,
+        workflow_dir_local=workflow_dir_local,
+        workflow_dir_remote=workflow_dir_remote,
         submit_setup_call=submit_setup_call,
         which_type="parallel",
     )
@@ -359,8 +359,8 @@ def run_v1_task_parallel(
             run_single_task,
             wftask=wftask,
             command=task_legacy.command,
-            workflow_dir=workflow_dir,
-            workflow_dir_user=workflow_dir_user,
+            workflow_dir_local=workflow_dir_local,
+            workflow_dir_remote=workflow_dir_remote,
             is_task_v1=True,
         ),
         list_function_kwargs,

fractal_server/app/runner/v2/runner_functions_low_level.py

@@ -60,8 +60,8 @@ def run_single_task(
     args: dict[str, Any],
     command: str,
     wftask: WorkflowTaskV2,
-    workflow_dir: Path,
-    workflow_dir_user: Optional[Path] = None,
+    workflow_dir_local: Path,
+    workflow_dir_remote: Optional[Path] = None,
     logger_name: Optional[str] = None,
     is_task_v1: bool = False,
 ) -> dict[str, Any]:
@@ -72,14 +72,20 @@ def run_single_task(
     logger = logging.getLogger(logger_name)
     logger.debug(f"Now start running {command=}")
 
-    if not workflow_dir_user:
-        workflow_dir_user = workflow_dir
+    if not workflow_dir_remote:
+        workflow_dir_remote = workflow_dir_local
+
+    if is_task_v1:
+        task_name = wftask.task_legacy.name
+    else:
+        task_name = wftask.task.name
 
     component = args.pop(_COMPONENT_KEY_, None)
     task_files = get_task_file_paths(
-        workflow_dir=workflow_dir,
-        workflow_dir_user=workflow_dir_user,
+        workflow_dir_local=workflow_dir_local,
+        workflow_dir_remote=workflow_dir_remote,
         task_order=wftask.order,
+        task_name=task_name,
         component=component,
     )
 

fractal_server/app/security/__init__.py

@@ -57,9 +57,8 @@ from sqlmodel import select
 from ...config import get_settings
 from ...syringe import Inject
 from ..db import get_async_db
-from ..models.security import OAuthAccount
-from ..models.security import UserOAuth as User
-from fractal_server.app.models.security import UserOAuth
+from fractal_server.app.models.security import OAuthAccount
+from fractal_server.app.models.security import UserOAuth as User
 from fractal_server.app.schemas.user import UserCreate
 from fractal_server.logger import get_logger
 
@@ -287,8 +286,8 @@ async def _create_first_user(
 
             if is_superuser is True:
                 # If a superuser already exists, exit
-                stm = select(UserOAuth).where(
-                    UserOAuth.is_superuser == True  # noqa: E712
+                stm = select(User).where(
+                    User.is_superuser == True  # noqa E712
                 )
                 res = await session.execute(stm)
                 existing_superuser = res.scalars().first()