[DOP-22135] Add resources field to Transfer model (#214)

Author: IlyasDevelopment

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, "cpu_cores_per_task": 2, "ram_bytes_per_task": "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,9 @@ class ReadTransferSchema(BaseModel):
     transformations: list[Annotated[TransformationSchema, Field(..., discriminator="type")]] = Field(
         default_factory=list,
     )
+    resources: Resources = Field(
+        ...,
+    )
 
     class Config:
         from_attributes = True
@@ -238,6 +242,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="Transfer resources",
+    )
 
     @model_validator(mode="before")
     def validate_scheduling(cls, values):
@@ -283,6 +291,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")
+    cpu_cores_per_task: int = Field(1, ge=1, le=32, description="Cores per executor")
+    ram_bytes_per_task: 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,

syncmaster/worker/spark.py

@@ -25,7 +25,7 @@ def get_worker_spark_session(
     name = run.transfer.group.name + "_" + run.transfer.name  # noqa: WPS336
     spark_builder = SparkSession.builder.appName(f"syncmaster_{name}")
 
-    for k, v in get_spark_session_conf(source, target).items():
+    for k, v in get_spark_session_conf(source, target, run.transfer.resources).items():
         spark_builder = spark_builder.config(k, v)
 
     if source.type == "hive" or target.type == "hive":  # type: ignore
@@ -81,6 +81,7 @@ def get_excluded_packages(db_type: str) -> list[str]:
 def get_spark_session_conf(
     source: ConnectionDTO,
     target: ConnectionDTO,
+    resources: dict,
 ) -> dict:
     maven_packages: list[str] = []
     excluded_packages: list[str] = []
@@ -93,6 +94,9 @@ def get_spark_session_conf(
         "spark.jars.packages": ",".join(maven_packages),
         "spark.sql.pyspark.jvmStacktrace.enabled": "true",
         "spark.hadoop.mapreduce.fileoutputcommitter.marksuccessfuljobs": "false",
+        "spark.executor.cores": resources["cpu_cores_per_task"],
+        "spark.executor.memory": resources["ram_bytes_per_task"],
+        "spark.executor.instances": resources["max_parallel_tasks"],
     }
 
     if maven_packages: