[DOP-22135] Add resources field to Transfer model

Author: Ilyas Gasanov

.readthedocs.yaml

@@ -4,6 +4,7 @@ build:
   os: ubuntu-22.04
   apt_packages:
     - make
+    - libkrb5-dev
   tools:
     python: '3.12'
   jobs:

docs/changelog/next_release/214.feature.rst

@@ -0,0 +1 @@
+Add `resources` field to `Transfer` model

docs/design/entities.rst

@@ -49,7 +49,7 @@ Example:
 .. code-block:: json
 
     {
-        "group_id": "1",
+        "group_id": 1,
         "name": "Beautiful name",
         "description": "What a great connection !",
         "type": "postgres",
@@ -93,17 +93,18 @@ Example:
 .. code-block:: json
 
     {
-        "group_id": "1",
-        "queue_id": "1",
+        "group_id": 1,
+        "queue_id": 1,
         "name": "My beautiful transfer.",
         "description": "What a great transfer !",
-        "is_scheduled": false,
-        "schedule": "",
-        "source_connection_id": "1",
-        "target_connection_id": "2",
-        "source_params": "{'type': 'postgres', 'table_name': 'source_table'}",
-        "target_params": "{'type': 'postgres', 'table_name': 'target_table'}",
-        "strategy_params": "{'type': 'full'}",
+        "is_scheduled": true,
+        "schedule": "0 0 * * *",
+        "source_connection_id": 1,
+        "target_connection_id": 2,
+        "source_params": {"type": "postgres", "table_name": "source_table"},
+        "target_params": {"type": "mysql", "table_name": "target_table"},
+        "strategy_params": {"type": "full"},
+        "resources": {"max_parallel_tasks": 2, "per_task_cores": 2, "per_task_ram_bytes": "1 GiB"}
     }
 
 Run

syncmaster/db/migrations/versions/2023-11-23_0007_create_transfer_table.py

@@ -45,6 +45,7 @@ def upgrade():
         sa.Column("source_params", sa.JSON(), nullable=False),
         sa.Column("target_params", sa.JSON(), nullable=False),
         sa.Column("transformations", sa.JSON(), nullable=False),
+        sa.Column("resources", sa.JSON(), nullable=False),
         sa.Column("is_scheduled", sa.Boolean(), nullable=False),
         sa.Column("schedule", sa.String(length=32), nullable=False),
         sa.Column("queue_id", sa.BigInteger(), nullable=False),

syncmaster/db/models/transfer.py

@@ -47,6 +47,7 @@ class Transfer(
     source_params: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default={})
     target_params: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default={})
     transformations: Mapped[list[dict[str, Any]]] = mapped_column(JSON, nullable=False, default=list)
+    resources: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default={})
     is_scheduled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
     schedule: Mapped[str] = mapped_column(String(32), nullable=False, default="")
     queue_id: Mapped[int] = mapped_column(

syncmaster/db/repositories/transfer.py

@@ -116,6 +116,7 @@ async def create(
         target_params: dict[str, Any],
         strategy_params: dict[str, Any],
         transformations: list[dict[str, Any]],
+        resources: dict[str, Any],
         queue_id: int,
         is_scheduled: bool,
         schedule: str | None,
@@ -132,6 +133,7 @@ async def create(
                 target_params=target_params,
                 strategy_params=strategy_params,
                 transformations=transformations,
+                resources=resources,
                 queue_id=queue_id,
                 is_scheduled=is_scheduled,
                 schedule=schedule or "",
@@ -157,6 +159,7 @@ async def update(
         target_params: dict[str, Any],
         strategy_params: dict[str, Any],
         transformations: list[dict[str, Any]],
+        resources: dict[str, Any],
         is_scheduled: bool | None,
         schedule: str | None,
         new_queue_id: int | None,
@@ -183,6 +186,7 @@ async def update(
                 source_params=source_params,
                 target_params=target_params,
                 transformations=transformations or transfer.transformations,
+                resources=resources or transfer.resources,
                 queue_id=new_queue_id or transfer.queue_id,
             )
         except IntegrityError as e:

syncmaster/schemas/v1/transfers/__init__.py

@@ -59,6 +59,7 @@
     WebDAVReadTransferSource,
     WebDAVReadTransferTarget,
 )
+from syncmaster.schemas.v1.transfers.resources import Resources
 from syncmaster.schemas.v1.transfers.strategy import FullStrategy, IncrementalStrategy
 from syncmaster.schemas.v1.transfers.transformations.dataframe_columns_filter import (
     DataframeColumnsFilter,
@@ -206,6 +207,10 @@ class ReadTransferSchema(BaseModel):
     transformations: list[Annotated[TransformationSchema, Field(..., discriminator="type")]] = Field(
         default_factory=list,
     )
+    resources: Resources = Field(
+        ...,
+        description="Spark resources",
+    )
 
     class Config:
         from_attributes = True
@@ -238,6 +243,10 @@ class CreateTransferSchema(BaseModel):
     transformations: list[
         Annotated[TransformationSchema, Field(None, discriminator="type", description="List of transformations")]
     ] = Field(default_factory=list)
+    resources: Resources = Field(
+        default_factory=Resources,
+        description="Spark resources",
+    )
 
     @model_validator(mode="before")
     def validate_scheduling(cls, values):
@@ -283,6 +292,7 @@ class UpdateTransferSchema(BaseModel):
     target_params: UpdateTransferSchemaTarget = Field(discriminator="type", default=None)
     strategy_params: FullStrategy | IncrementalStrategy | None = Field(discriminator="type", default=None)
     transformations: list[Annotated[TransformationSchema, Field(discriminator="type", default=None)]] = None
+    resources: Resources | None = Field(default=None)
 
 
 class ReadFullTransferSchema(ReadTransferSchema):

syncmaster/schemas/v1/transfers/resources.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from pydantic import BaseModel, ByteSize, Field
+
+
+class Resources(BaseModel):
+    max_parallel_tasks: int = Field(1, ge=1, le=100, description="Parallel executors")
+    per_task_cores: int = Field(1, ge=1, le=32, description="Cores per executor")
+    per_task_ram_bytes: ByteSize = Field(
+        1024**3,
+        ge=512 * 1024**2,
+        le=64 * 1024**3,
+        description="RAM per executor",
+    )

syncmaster/schemas/v1/transfers/strategy.py

@@ -1,7 +1,5 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
 from pydantic import BaseModel
 
 from syncmaster.schemas.v1.transfer_types import FULL_TYPE, INCREMENTAL_TYPE

syncmaster/server/api/v1/transfers.py

@@ -131,6 +131,7 @@ async def create_transfer(
             target_params=transfer_data.target_params.dict(),
             strategy_params=transfer_data.strategy_params.dict(),
             transformations=[tr.dict() for tr in transfer_data.transformations],
+            resources=transfer_data.resources.dict(),
             queue_id=transfer_data.queue_id,
             is_scheduled=transfer_data.is_scheduled,
             schedule=transfer_data.schedule,
@@ -330,6 +331,7 @@ async def update_transfer(
             transformations=(
                 [tr.dict() for tr in transfer_data.transformations] if transfer_data.transformations else []
             ),
+            resources=transfer_data.resources.dict() if transfer_data.resources else {},
             is_scheduled=transfer_data.is_scheduled,
             schedule=transfer_data.schedule,
             new_queue_id=transfer_data.new_queue_id,