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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. fractal_server/__init__.py +1 -1
  2. fractal_server/app/db/__init__.py +1 -1
  3. fractal_server/app/models/v1/state.py +1 -2
  4. fractal_server/app/routes/admin/v1.py +2 -2
  5. fractal_server/app/routes/admin/v2.py +2 -2
  6. fractal_server/app/routes/api/v1/job.py +2 -2
  7. fractal_server/app/routes/api/v1/task_collection.py +4 -4
  8. fractal_server/app/routes/api/v2/__init__.py +23 -3
  9. fractal_server/app/routes/api/v2/job.py +2 -2
  10. fractal_server/app/routes/api/v2/submit.py +6 -0
  11. fractal_server/app/routes/api/v2/task_collection.py +74 -34
  12. fractal_server/app/routes/api/v2/task_collection_custom.py +170 -0
  13. fractal_server/app/routes/api/v2/task_collection_ssh.py +125 -0
  14. fractal_server/app/routes/aux/_runner.py +10 -2
  15. fractal_server/app/runner/compress_folder.py +120 -0
  16. fractal_server/app/runner/executors/slurm/__init__.py +0 -3
  17. fractal_server/app/runner/executors/slurm/_batching.py +0 -1
  18. fractal_server/app/runner/executors/slurm/_slurm_config.py +9 -9
  19. fractal_server/app/runner/executors/slurm/ssh/__init__.py +3 -0
  20. fractal_server/app/runner/executors/slurm/ssh/_executor_wait_thread.py +112 -0
  21. fractal_server/app/runner/executors/slurm/ssh/_slurm_job.py +120 -0
  22. fractal_server/app/runner/executors/slurm/ssh/executor.py +1488 -0
  23. fractal_server/app/runner/executors/slurm/sudo/__init__.py +3 -0
  24. fractal_server/app/runner/executors/slurm/{_check_jobs_status.py → sudo/_check_jobs_status.py} +1 -1
  25. fractal_server/app/runner/executors/slurm/{_executor_wait_thread.py → sudo/_executor_wait_thread.py} +1 -1
  26. fractal_server/app/runner/executors/slurm/{_subprocess_run_as_user.py → sudo/_subprocess_run_as_user.py} +1 -1
  27. fractal_server/app/runner/executors/slurm/{executor.py → sudo/executor.py} +12 -12
  28. fractal_server/app/runner/extract_archive.py +38 -0
  29. fractal_server/app/runner/v1/__init__.py +78 -40
  30. fractal_server/app/runner/v1/_slurm/__init__.py +1 -1
  31. fractal_server/app/runner/v2/__init__.py +183 -82
  32. fractal_server/app/runner/v2/_local_experimental/__init__.py +22 -12
  33. fractal_server/app/runner/v2/_local_experimental/executor.py +12 -8
  34. fractal_server/app/runner/v2/_slurm/__init__.py +1 -6
  35. fractal_server/app/runner/v2/_slurm_ssh/__init__.py +125 -0
  36. fractal_server/app/runner/v2/_slurm_ssh/_submit_setup.py +83 -0
  37. fractal_server/app/runner/v2/_slurm_ssh/get_slurm_config.py +182 -0
  38. fractal_server/app/runner/v2/runner_functions_low_level.py +9 -11
  39. fractal_server/app/runner/versions.py +30 -0
  40. fractal_server/app/schemas/v1/__init__.py +1 -0
  41. fractal_server/app/schemas/{state.py → v1/state.py} +4 -21
  42. fractal_server/app/schemas/v2/__init__.py +4 -1
  43. fractal_server/app/schemas/v2/task_collection.py +101 -30
  44. fractal_server/config.py +222 -21
  45. fractal_server/main.py +27 -1
  46. fractal_server/migrations/env.py +1 -1
  47. fractal_server/ssh/__init__.py +4 -0
  48. fractal_server/ssh/_fabric.py +245 -0
  49. fractal_server/tasks/utils.py +12 -64
  50. fractal_server/tasks/v1/background_operations.py +2 -2
  51. fractal_server/tasks/{endpoint_operations.py → v1/endpoint_operations.py} +7 -12
  52. fractal_server/tasks/v1/utils.py +67 -0
  53. fractal_server/tasks/v2/_TaskCollectPip.py +61 -32
  54. fractal_server/tasks/v2/_venv_pip.py +195 -0
  55. fractal_server/tasks/v2/background_operations.py +257 -295
  56. fractal_server/tasks/v2/background_operations_ssh.py +317 -0
  57. fractal_server/tasks/v2/endpoint_operations.py +136 -0
  58. fractal_server/tasks/v2/templates/_1_create_venv.sh +46 -0
  59. fractal_server/tasks/v2/templates/_2_upgrade_pip.sh +30 -0
  60. fractal_server/tasks/v2/templates/_3_pip_install.sh +32 -0
  61. fractal_server/tasks/v2/templates/_4_pip_freeze.sh +21 -0
  62. fractal_server/tasks/v2/templates/_5_pip_show.sh +59 -0
  63. fractal_server/tasks/v2/utils.py +54 -0
  64. {fractal_server-2.2.0a0.dist-info → fractal_server-2.3.0.dist-info}/METADATA +6 -2
  65. {fractal_server-2.2.0a0.dist-info → fractal_server-2.3.0.dist-info}/RECORD +68 -44
  66. fractal_server/tasks/v2/get_collection_data.py +0 -14
  67. {fractal_server-2.2.0a0.dist-info → fractal_server-2.3.0.dist-info}/LICENSE +0 -0
  68. {fractal_server-2.2.0a0.dist-info → fractal_server-2.3.0.dist-info}/WHEEL +0 -0
  69. {fractal_server-2.2.0a0.dist-info → fractal_server-2.3.0.dist-info}/entry_points.txt +0 -0
@@ -1 +1 @@
1
- __VERSION__ = "2.2.0a0"
1
+ __VERSION__ = "2.3.0"
@@ -63,7 +63,7 @@ class DB:
63
63
  }
64
64
 
65
65
  cls._engine_async = create_async_engine(
66
- settings.DATABASE_URL,
66
+ settings.DATABASE_ASYNC_URL,
67
67
  echo=settings.DB_ECHO,
68
68
  future=True,
69
69
  **engine_kwargs_async,
@@ -9,10 +9,9 @@ from sqlmodel import Field
9
9
  from sqlmodel import SQLModel
10
10
 
11
11
  from ....utils import get_timestamp
12
- from ...schemas.state import _StateBase
13
12
 
14
13
 
15
- class State(_StateBase, SQLModel, table=True):
14
+ class State(SQLModel, table=True):
16
15
  """
17
16
  Store arbitrary data in the database
18
17
 
@@ -35,7 +35,7 @@ from ...schemas.v1 import WorkflowReadV1
35
35
  from ...security import current_active_superuser
36
36
  from ..aux._job import _write_shutdown_file
37
37
  from ..aux._job import _zip_folder_to_byte_stream
38
- from ..aux._runner import _is_shutdown_available
38
+ from ..aux._runner import _check_shutdown_is_supported
39
39
 
40
40
  router_admin_v1 = APIRouter()
41
41
 
@@ -351,7 +351,7 @@ async def stop_job(
351
351
  Stop execution of a workflow job.
352
352
  """
353
353
 
354
- _is_shutdown_available()
354
+ _check_shutdown_is_supported()
355
355
 
356
356
  job = await db.get(ApplyWorkflow, job_id)
357
357
  if job is None:
@@ -38,7 +38,7 @@ from ...schemas.v2 import ProjectReadV2
38
38
  from ...security import current_active_superuser
39
39
  from ..aux._job import _write_shutdown_file
40
40
  from ..aux._job import _zip_folder_to_byte_stream
41
- from ..aux._runner import _is_shutdown_available
41
+ from ..aux._runner import _check_shutdown_is_supported
42
42
 
43
43
  router_admin_v2 = APIRouter()
44
44
 
@@ -238,7 +238,7 @@ async def stop_job(
238
238
  Stop execution of a workflow job.
239
239
  """
240
240
 
241
- _is_shutdown_available()
241
+ _check_shutdown_is_supported()
242
242
 
243
243
  job = await db.get(JobV2, job_id)
244
244
  if job is None:
@@ -19,7 +19,7 @@ from ....security import current_active_user
19
19
  from ....security import User
20
20
  from ...aux._job import _write_shutdown_file
21
21
  from ...aux._job import _zip_folder_to_byte_stream
22
- from ...aux._runner import _is_shutdown_available
22
+ from ...aux._runner import _check_shutdown_is_supported
23
23
  from ._aux_functions import _get_job_check_owner
24
24
  from ._aux_functions import _get_project_check_owner
25
25
  from ._aux_functions import _get_workflow_check_owner
@@ -180,7 +180,7 @@ async def stop_job(
180
180
  Stop execution of a workflow job.
181
181
  """
182
182
 
183
- _is_shutdown_available()
183
+ _check_shutdown_is_supported()
184
184
 
185
185
  # Get job from DB
186
186
  output = await _get_job_check_owner(
@@ -19,21 +19,21 @@ from ....db import AsyncSession
19
19
  from ....db import get_async_db
20
20
  from ....models.v1 import State
21
21
  from ....models.v1 import Task
22
- from ....schemas.state import StateRead
22
+ from ....schemas.v1 import StateRead
23
23
  from ....schemas.v1 import TaskCollectPipV1
24
24
  from ....schemas.v1 import TaskCollectStatusV1
25
25
  from ....security import current_active_user
26
26
  from ....security import current_active_verified_user
27
27
  from ....security import User
28
- from fractal_server.tasks.endpoint_operations import create_package_dir_pip
29
- from fractal_server.tasks.endpoint_operations import download_package
30
- from fractal_server.tasks.endpoint_operations import inspect_package
31
28
  from fractal_server.tasks.utils import get_collection_log
32
29
  from fractal_server.tasks.utils import slugify_task_name
33
30
  from fractal_server.tasks.v1._TaskCollectPip import _TaskCollectPip
34
31
  from fractal_server.tasks.v1.background_operations import (
35
32
  background_collect_pip,
36
33
  )
34
+ from fractal_server.tasks.v1.endpoint_operations import create_package_dir_pip
35
+ from fractal_server.tasks.v1.endpoint_operations import download_package
36
+ from fractal_server.tasks.v1.endpoint_operations import inspect_package
37
37
  from fractal_server.tasks.v1.get_collection_data import get_collection_data
38
38
 
39
39
  router = APIRouter()
@@ -11,9 +11,14 @@ from .status import router as status_router_v2
11
11
  from .submit import router as submit_job_router_v2
12
12
  from .task import router as task_router_v2
13
13
  from .task_collection import router as task_collection_router_v2
14
+ from .task_collection_custom import router as task_collection_router_v2_custom
15
+ from .task_collection_ssh import router as task_collection_router_v2_ssh
14
16
  from .task_legacy import router as task_legacy_router_v2
15
17
  from .workflow import router as workflow_router_v2
16
18
  from .workflowtask import router as workflowtask_router_v2
19
+ from fractal_server.config import get_settings
20
+ from fractal_server.syringe import Inject
21
+
17
22
 
18
23
  router_api_v2 = APIRouter()
19
24
 
@@ -22,9 +27,24 @@ router_api_v2.include_router(job_router_v2, tags=["V2 Job"])
22
27
  router_api_v2.include_router(images_routes_v2, tags=["V2 Images"])
23
28
  router_api_v2.include_router(project_router_v2, tags=["V2 Project"])
24
29
  router_api_v2.include_router(submit_job_router_v2, tags=["V2 Job"])
25
- router_api_v2.include_router(
26
- task_collection_router_v2, prefix="/task", tags=["V2 Task Collection"]
27
- )
30
+
31
+
32
+ settings = Inject(get_settings)
33
+ if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
34
+ router_api_v2.include_router(
35
+ task_collection_router_v2_ssh,
36
+ prefix="/task",
37
+ tags=["V2 Task Collection"],
38
+ )
39
+ else:
40
+ router_api_v2.include_router(
41
+ task_collection_router_v2, prefix="/task", tags=["V2 Task Collection"]
42
+ )
43
+ router_api_v2.include_router(
44
+ task_collection_router_v2_custom,
45
+ prefix="/task",
46
+ tags=["V2 Task Collection"],
47
+ )
28
48
  router_api_v2.include_router(task_router_v2, prefix="/task", tags=["V2 Task"])
29
49
  router_api_v2.include_router(
30
50
  task_legacy_router_v2, prefix="/task-legacy", tags=["V2 Task Legacy"]
@@ -19,7 +19,7 @@ from ....security import current_active_user
19
19
  from ....security import User
20
20
  from ...aux._job import _write_shutdown_file
21
21
  from ...aux._job import _zip_folder_to_byte_stream
22
- from ...aux._runner import _is_shutdown_available
22
+ from ...aux._runner import _check_shutdown_is_supported
23
23
  from ._aux_functions import _get_job_check_owner
24
24
  from ._aux_functions import _get_project_check_owner
25
25
  from ._aux_functions import _get_workflow_check_owner
@@ -183,7 +183,7 @@ async def stop_job(
183
183
  Stop execution of a workflow job.
184
184
  """
185
185
 
186
- _is_shutdown_available()
186
+ _check_shutdown_is_supported()
187
187
 
188
188
  # Get job from DB
189
189
  output = await _get_job_check_owner(
@@ -226,6 +226,11 @@ async def apply_workflow(
226
226
  WORKFLOW_DIR_REMOTE = (
227
227
  Path(user.cache_dir) / f"{WORKFLOW_DIR_LOCAL.name}"
228
228
  )
229
+ elif FRACTAL_RUNNER_BACKEND == "slurm_ssh":
230
+ WORKFLOW_DIR_REMOTE = (
231
+ Path(settings.FRACTAL_SLURM_SSH_WORKING_BASE_DIR)
232
+ / f"{WORKFLOW_DIR_LOCAL.name}"
233
+ )
229
234
 
230
235
  # Update job folders in the db
231
236
  job.working_dir = WORKFLOW_DIR_LOCAL.as_posix()
@@ -241,6 +246,7 @@ async def apply_workflow(
241
246
  worker_init=job.worker_init,
242
247
  slurm_user=user.slurm_user,
243
248
  user_cache_dir=user.cache_dir,
249
+ fractal_ssh=request.app.state.fractal_ssh,
244
250
  )
245
251
  request.app.state.jobsV2.append(job.id)
246
252
  logger.info(
@@ -1,3 +1,4 @@
1
+ import json
1
2
  from pathlib import Path
2
3
  from shutil import copy as shell_copy
3
4
  from tempfile import TemporaryDirectory
@@ -19,22 +20,25 @@ from ....db import AsyncSession
19
20
  from ....db import get_async_db
20
21
  from ....models.v2 import CollectionStateV2
21
22
  from ....models.v2 import TaskV2
22
- from ....schemas.state import StateRead
23
+ from ....schemas.v2 import CollectionStateReadV2
24
+ from ....schemas.v2 import CollectionStatusV2
23
25
  from ....schemas.v2 import TaskCollectPipV2
24
- from ....schemas.v2 import TaskCollectStatusV2
26
+ from ....schemas.v2 import TaskReadV2
25
27
  from ....security import current_active_user
26
28
  from ....security import current_active_verified_user
27
29
  from ....security import User
28
- from fractal_server.tasks.endpoint_operations import create_package_dir_pip
29
- from fractal_server.tasks.endpoint_operations import download_package
30
- from fractal_server.tasks.endpoint_operations import inspect_package
30
+ from fractal_server.tasks.utils import get_absolute_venv_path
31
31
  from fractal_server.tasks.utils import get_collection_log
32
+ from fractal_server.tasks.utils import get_collection_path
32
33
  from fractal_server.tasks.utils import slugify_task_name
33
34
  from fractal_server.tasks.v2._TaskCollectPip import _TaskCollectPip
34
35
  from fractal_server.tasks.v2.background_operations import (
35
36
  background_collect_pip,
36
37
  )
37
- from fractal_server.tasks.v2.get_collection_data import get_collection_data
38
+ from fractal_server.tasks.v2.endpoint_operations import create_package_dir_pip
39
+ from fractal_server.tasks.v2.endpoint_operations import download_package
40
+ from fractal_server.tasks.v2.endpoint_operations import inspect_package
41
+
38
42
 
39
43
  router = APIRouter()
40
44
 
@@ -43,7 +47,7 @@ logger = set_logger(__name__)
43
47
 
44
48
  @router.post(
45
49
  "/collect/pip/",
46
- response_model=StateRead,
50
+ response_model=CollectionStateReadV2,
47
51
  responses={
48
52
  201: dict(
49
53
  description=(
@@ -64,7 +68,7 @@ async def collect_tasks_pip(
64
68
  response: Response,
65
69
  user: User = Depends(current_active_verified_user),
66
70
  db: AsyncSession = Depends(get_async_db),
67
- ) -> StateRead: # State[TaskCollectStatus]
71
+ ) -> CollectionStateReadV2:
68
72
  """
69
73
  Task collection endpoint
70
74
 
@@ -74,6 +78,13 @@ async def collect_tasks_pip(
74
78
 
75
79
  logger = set_logger(logger_name="collect_tasks_pip")
76
80
 
81
+ # Set default python version
82
+ if task_collect.python_version is None:
83
+ settings = Inject(get_settings)
84
+ task_collect.python_version = (
85
+ settings.FRACTAL_TASKS_PYTHON_DEFAULT_VERSION
86
+ )
87
+
77
88
  # Validate payload as _TaskCollectPip, which has more strict checks than
78
89
  # TaskCollectPip
79
90
  try:
@@ -89,18 +100,18 @@ async def collect_tasks_pip(
89
100
  # Copy or download the package wheel file to tmpdir
90
101
  if task_pkg.is_local_package:
91
102
  shell_copy(task_pkg.package_path.as_posix(), tmpdir)
92
- pkg_path = Path(tmpdir) / task_pkg.package_path.name
103
+ wheel_path = Path(tmpdir) / task_pkg.package_path.name
93
104
  else:
94
- pkg_path = await download_package(
105
+ logger.info(f"Now download {task_pkg}")
106
+ wheel_path = await download_package(
95
107
  task_pkg=task_pkg, dest=tmpdir
96
108
  )
97
109
  # Read package info from wheel file, and override the ones coming
98
- # from the request body
99
- pkg_info = inspect_package(pkg_path)
100
- task_pkg.package_name = pkg_info["pkg_name"]
110
+ # from the request body. Note that `package_name` was already set
111
+ # (and normalized) as part of `_TaskCollectPip` initialization.
112
+ pkg_info = inspect_package(wheel_path)
101
113
  task_pkg.package_version = pkg_info["pkg_version"]
102
114
  task_pkg.package_manifest = pkg_info["pkg_manifest"]
103
- task_pkg.check()
104
115
  except Exception as e:
105
116
  raise HTTPException(
106
117
  status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
@@ -112,8 +123,35 @@ async def collect_tasks_pip(
112
123
  except FileExistsError:
113
124
  venv_path = create_package_dir_pip(task_pkg=task_pkg, create=False)
114
125
  try:
115
- task_collect_status = get_collection_data(venv_path)
116
- for task in task_collect_status.task_list:
126
+ package_path = get_absolute_venv_path(venv_path)
127
+ collection_path = get_collection_path(package_path)
128
+ with collection_path.open("r") as f:
129
+ task_collect_data = json.load(f)
130
+
131
+ err_msg = (
132
+ "Cannot collect package, possible reason: an old version of "
133
+ "the same package has already been collected.\n"
134
+ f"{str(collection_path)} has invalid content: "
135
+ )
136
+ if not isinstance(task_collect_data, dict):
137
+ raise HTTPException(
138
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
139
+ detail=f"{err_msg} it's not a Python dictionary.",
140
+ )
141
+ if "task_list" not in task_collect_data.keys():
142
+ raise HTTPException(
143
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
144
+ detail=f"{err_msg} it has no key 'task_list'.",
145
+ )
146
+ if not isinstance(task_collect_data["task_list"], list):
147
+ raise HTTPException(
148
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
149
+ detail=f"{err_msg} 'task_list' is not a Python list.",
150
+ )
151
+
152
+ for task_dict in task_collect_data["task_list"]:
153
+
154
+ task = TaskReadV2(**task_dict)
117
155
  db_task = await db.get(TaskV2, task.id)
118
156
  if (
119
157
  (not db_task)
@@ -150,8 +188,8 @@ async def collect_tasks_pip(
150
188
  f"Original ValidationError: {e}"
151
189
  ),
152
190
  )
153
- task_collect_status.info = "Already installed"
154
- state = CollectionStateV2(data=task_collect_status.sanitised_dict())
191
+ task_collect_data["info"] = "Already installed"
192
+ state = CollectionStateV2(data=task_collect_data)
155
193
  response.status_code == status.HTTP_200_OK
156
194
  await db.close()
157
195
  return state
@@ -173,15 +211,12 @@ async def collect_tasks_pip(
173
211
  )
174
212
 
175
213
  # All checks are OK, proceed with task collection
176
- full_venv_path = venv_path.relative_to(settings.FRACTAL_TASKS_DIR)
177
- collection_status = TaskCollectStatusV2(
178
- status="pending", venv_path=full_venv_path, package=task_pkg.package
214
+ collection_status = dict(
215
+ status=CollectionStatusV2.PENDING,
216
+ venv_path=venv_path.relative_to(settings.FRACTAL_TASKS_DIR).as_posix(),
217
+ package=task_pkg.package,
179
218
  )
180
-
181
- # Create State object (after casting venv_path to string)
182
- collection_status_dict = collection_status.dict()
183
- collection_status_dict["venv_path"] = str(collection_status.venv_path)
184
- state = CollectionStateV2(data=collection_status_dict)
219
+ state = CollectionStateV2(data=collection_status)
185
220
  db.add(state)
186
221
  await db.commit()
187
222
  await db.refresh(state)
@@ -208,13 +243,13 @@ async def collect_tasks_pip(
208
243
  return state
209
244
 
210
245
 
211
- @router.get("/collect/{state_id}/", response_model=StateRead)
246
+ @router.get("/collect/{state_id}/", response_model=CollectionStateReadV2)
212
247
  async def check_collection_status(
213
248
  state_id: int,
214
249
  user: User = Depends(current_active_user),
215
250
  verbose: bool = False,
216
251
  db: AsyncSession = Depends(get_async_db),
217
- ) -> StateRead: # State[TaskCollectStatus]
252
+ ) -> CollectionStateReadV2: # State[TaskCollectStatus]
218
253
  """
219
254
  Check status of background task collection
220
255
  """
@@ -227,13 +262,18 @@ async def check_collection_status(
227
262
  status_code=status.HTTP_404_NOT_FOUND,
228
263
  detail=f"No task collection info with id={state_id}",
229
264
  )
230
- data = TaskCollectStatusV2(**state.data)
231
265
 
232
- # In some cases (i.e. a successful or ongoing task collection), data.log is
233
- # not set; if so, we collect the current logs
234
- if verbose and not data.log:
235
- data.log = get_collection_log(data.venv_path)
236
- state.data = data.sanitised_dict()
266
+ # In some cases (i.e. a successful or ongoing task collection),
267
+ # state.data.log is not set; if so, we collect the current logs.
268
+ if verbose and not state.data.get("log"):
269
+ if "venv_path" not in state.data.keys():
270
+ await db.close()
271
+ raise HTTPException(
272
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
273
+ detail=f"No 'venv_path' in CollectionStateV2[{state_id}].data",
274
+ )
275
+ state.data["log"] = get_collection_log(Path(state.data["venv_path"]))
276
+ state.data["venv_path"] = str(state.data["venv_path"])
237
277
  reset_logger_handlers(logger)
238
278
  await db.close()
239
279
  return state
@@ -0,0 +1,170 @@
1
+ import shlex
2
+ import subprocess # nosec
3
+ from pathlib import Path
4
+
5
+ from fastapi import APIRouter
6
+ from fastapi import Depends
7
+ from fastapi import HTTPException
8
+ from fastapi import status
9
+ from sqlmodel import select
10
+
11
+ from .....config import get_settings
12
+ from .....logger import set_logger
13
+ from .....syringe import Inject
14
+ from ....db import DBSyncSession
15
+ from ....db import get_sync_db
16
+ from ....models.v1 import Task as TaskV1
17
+ from ....models.v2 import TaskV2
18
+ from ....schemas.v2 import TaskCollectCustomV2
19
+ from ....schemas.v2 import TaskCreateV2
20
+ from ....schemas.v2 import TaskReadV2
21
+ from ....security import current_active_verified_user
22
+ from ....security import User
23
+ from fractal_server.tasks.v2.background_operations import _insert_tasks
24
+ from fractal_server.tasks.v2.background_operations import (
25
+ _prepare_tasks_metadata,
26
+ )
27
+
28
+
29
+ router = APIRouter()
30
+
31
+ logger = set_logger(__name__)
32
+
33
+
34
+ @router.post(
35
+ "/collect/custom/", status_code=201, response_model=list[TaskReadV2]
36
+ )
37
+ async def collect_task_custom(
38
+ task_collect: TaskCollectCustomV2,
39
+ user: User = Depends(current_active_verified_user),
40
+ db: DBSyncSession = Depends(get_sync_db),
41
+ ) -> list[TaskReadV2]:
42
+
43
+ settings = Inject(get_settings)
44
+
45
+ if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
46
+ if task_collect.package_root is None:
47
+ raise HTTPException(
48
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
49
+ detail="Cannot infer 'package_root' with 'slurm_ssh' backend.",
50
+ )
51
+ else:
52
+ if not Path(task_collect.python_interpreter).is_file():
53
+ raise HTTPException(
54
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
55
+ detail=(
56
+ f"{task_collect.python_interpreter=} "
57
+ "doesn't exist or is not a file."
58
+ ),
59
+ )
60
+ if (
61
+ task_collect.package_root is not None
62
+ and not Path(task_collect.package_root).is_dir()
63
+ ):
64
+ raise HTTPException(
65
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
66
+ detail=(
67
+ f"{task_collect.package_root=} "
68
+ "doesn't exist or is not a directory."
69
+ ),
70
+ )
71
+
72
+ if task_collect.package_root is None:
73
+
74
+ package_name_underscore = task_collect.package_name.replace("-", "_")
75
+ # Note that python_command is then used as part of a subprocess.run
76
+ # statement: be careful with mixing `'` and `"`.
77
+ python_command = (
78
+ "import importlib.util; "
79
+ "from pathlib import Path; "
80
+ "init_path=importlib.util.find_spec"
81
+ f'("{package_name_underscore}").origin; '
82
+ "print(Path(init_path).parent.as_posix())"
83
+ )
84
+ logger.debug(
85
+ f"Now running {python_command=} through "
86
+ f"{task_collect.python_interpreter}."
87
+ )
88
+ res = subprocess.run( # nosec
89
+ shlex.split(
90
+ f"{task_collect.python_interpreter} -c '{python_command}'"
91
+ ),
92
+ capture_output=True,
93
+ encoding="utf8",
94
+ )
95
+
96
+ if (
97
+ res.returncode != 0
98
+ or res.stdout is None
99
+ or ("\n" in res.stdout.strip("\n"))
100
+ ):
101
+ raise HTTPException(
102
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
103
+ detail=(
104
+ "Cannot determine 'package_root'.\n"
105
+ f"Original output: {res.stdout}\n"
106
+ f"Original error: {res.stderr}"
107
+ ),
108
+ )
109
+ package_root = Path(res.stdout.strip("\n"))
110
+ else:
111
+ package_root = Path(task_collect.package_root)
112
+
113
+ # Set task.owner attribute
114
+ owner = user.username or user.slurm_user
115
+ if owner is None:
116
+ raise HTTPException(
117
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
118
+ detail=(
119
+ "Cannot add a new task because current user does not "
120
+ "have `username` or `slurm_user` attributes."
121
+ ),
122
+ )
123
+ source = f"{owner}:{task_collect.source}"
124
+
125
+ task_list: list[TaskCreateV2] = _prepare_tasks_metadata(
126
+ package_manifest=task_collect.manifest,
127
+ package_source=source,
128
+ python_bin=Path(task_collect.python_interpreter),
129
+ package_root=package_root,
130
+ package_version=task_collect.version,
131
+ )
132
+ # Verify that source is not already in use (note: this check is only useful
133
+ # to provide a user-friendly error message, but `task.source` uniqueness is
134
+ # already guaranteed by a constraint in the table definition).
135
+ sources = [task.source for task in task_list]
136
+ stm = select(TaskV2).where(TaskV2.source.in_(sources))
137
+ res = db.execute(stm)
138
+ overlapping_sources_v2 = res.scalars().all()
139
+ if overlapping_sources_v2:
140
+ overlapping_tasks_v2_source_and_id = [
141
+ f"TaskV2 with ID {task.id} already has source='{task.source}'"
142
+ for task in overlapping_sources_v2
143
+ ]
144
+ raise HTTPException(
145
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
146
+ detail="\n".join(overlapping_tasks_v2_source_and_id),
147
+ )
148
+ stm = select(TaskV1).where(TaskV1.source.in_(sources))
149
+ res = db.execute(stm)
150
+ overlapping_sources_v1 = res.scalars().all()
151
+ if overlapping_sources_v1:
152
+ overlapping_tasks_v1_source_and_id = [
153
+ f"TaskV1 with ID {task.id} already has source='{task.source}'\n"
154
+ for task in overlapping_sources_v1
155
+ ]
156
+ raise HTTPException(
157
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
158
+ detail="\n".join(overlapping_tasks_v1_source_and_id),
159
+ )
160
+
161
+ task_list_db: list[TaskV2] = _insert_tasks(
162
+ task_list=task_list, owner=owner, db=db
163
+ )
164
+
165
+ logger.debug(
166
+ f"Custom-environment task collection by user {user.email} completed, "
167
+ f"for package with {source=}"
168
+ )
169
+
170
+ return task_list_db