diff --git a/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/queries/common.ts b/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/queries/common.ts index 4d553521aea43..2b90edb3b00fc 100644 --- a/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/queries/common.ts +++ b/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/queries/common.ts @@ -50,3 +50,4 @@ export type UiServiceUpdateWorkerMaintenanceMutationResult = Awaited>; export type UiServiceDeleteWorkerMutationResult = Awaited>; export type UiServiceRemoveWorkerQueueMutationResult = Awaited>; +export type UiServiceSetWorkerConcurrencyLimitMutationResult = Awaited>; diff --git a/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/queries/queries.ts b/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/queries/queries.ts index 782a26bbf9114..83cae99c134b1 100644 --- a/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/queries/queries.ts +++ b/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/queries/queries.ts @@ -2,7 +2,7 @@ import { UseMutationOptions, UseQueryOptions, useMutation, useQuery } from "@tanstack/react-query"; import { JobsService, LogsService, MonitorService, UiService, WorkerService } from "../requests/services.gen"; -import { EdgeWorkerState, MaintenanceRequest, PushLogsBody, TaskInstanceState, WorkerQueueUpdateBody, WorkerQueuesBody, WorkerStateBody } from "../requests/types.gen"; +import { ConcurrencyRequest, EdgeWorkerState, MaintenanceRequest, PushLogsBody, TaskInstanceState, WorkerQueueUpdateBody, WorkerQueuesBody, WorkerStateBody } from "../requests/types.gen"; import * as Common from "./common"; export const useLogsServiceLogfilePath = = unknown[]>({ authorization, dagId, mapIndex, runId, taskId, tryNumber }: { authorization: string; @@ -139,3 +139,10 @@ export const useUiServiceRemoveWorkerQueue = ({ mutationFn: ({ queueName, workerName }) => UiService.removeWorkerQueue({ queueName, workerName }) as unknown as Promise, ...options }); +export const useUiServiceSetWorkerConcurrencyLimit = (options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ requestBody, workerName }) => UiService.setWorkerConcurrencyLimit({ requestBody, workerName }) as unknown as Promise, ...options }); diff --git a/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/schemas.gen.ts b/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/schemas.gen.ts index 11e48528f69e6..51756edfb4452 100644 --- a/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/schemas.gen.ts +++ b/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/schemas.gen.ts @@ -257,6 +257,21 @@ export const $JobCollectionResponse = { description: 'Job Collection serializer.' } as const; +export const $ConcurrencyRequest = { + properties: { + concurrency: { + type: 'integer', + exclusiveMinimum: 0, + title: 'Concurrency', + description: 'New concurrency limit for the worker.' + } + }, + type: 'object', + required: ['concurrency'], + title: 'ConcurrencyRequest', + description: 'Request body for worker concurrency update.' +} as const; + export const $MaintenanceRequest = { properties: { maintenance_comment: { diff --git a/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/services.gen.ts b/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/services.gen.ts index da00585e04498..2dd062ec1ff56 100644 --- a/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/services.gen.ts +++ b/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/services.gen.ts @@ -3,7 +3,7 @@ import type { CancelablePromise } from './core/CancelablePromise'; import { OpenAPI } from './core/OpenAPI'; import { request as __request } from './core/request'; -import type { FetchData, FetchResponse, StateData, StateResponse, LogfilePathData, LogfilePathResponse, PushLogsData, PushLogsResponse, RegisterData, RegisterResponse, SetStateData, SetStateResponse, UpdateQueuesData, UpdateQueuesResponse, HealthResponse, WorkerData, WorkerResponse, JobsResponse, RequestWorkerMaintenanceData, RequestWorkerMaintenanceResponse, UpdateWorkerMaintenanceData, UpdateWorkerMaintenanceResponse, ExitWorkerMaintenanceData, ExitWorkerMaintenanceResponse, RequestWorkerShutdownData, RequestWorkerShutdownResponse, DeleteWorkerData, DeleteWorkerResponse, AddWorkerQueueData, AddWorkerQueueResponse, RemoveWorkerQueueData, RemoveWorkerQueueResponse } from './types.gen'; +import type { FetchData, FetchResponse, StateData, StateResponse, LogfilePathData, LogfilePathResponse, PushLogsData, PushLogsResponse, RegisterData, RegisterResponse, SetStateData, SetStateResponse, UpdateQueuesData, UpdateQueuesResponse, HealthResponse, WorkerData, WorkerResponse, JobsResponse, RequestWorkerMaintenanceData, RequestWorkerMaintenanceResponse, UpdateWorkerMaintenanceData, UpdateWorkerMaintenanceResponse, ExitWorkerMaintenanceData, ExitWorkerMaintenanceResponse, RequestWorkerShutdownData, RequestWorkerShutdownResponse, DeleteWorkerData, DeleteWorkerResponse, AddWorkerQueueData, AddWorkerQueueResponse, RemoveWorkerQueueData, RemoveWorkerQueueResponse, SetWorkerConcurrencyLimitData, SetWorkerConcurrencyLimitResponse } from './types.gen'; export class JobsService { /** @@ -472,5 +472,29 @@ export class UiService { } }); } - + + /** + * Set Worker Concurrency Limit + * Set the concurrency limit for an edge worker. + * @param data The data for the request. + * @param data.workerName + * @param data.requestBody + * @returns unknown Successful Response + * @throws ApiError + */ + public static setWorkerConcurrencyLimit(data: SetWorkerConcurrencyLimitData): CancelablePromise { + return __request(OpenAPI, { + method: 'PATCH', + url: '/edge_worker/ui/worker/{worker_name}/concurrency', + path: { + worker_name: data.workerName + }, + body: data.requestBody, + mediaType: 'application/json', + errors: { + 422: 'Validation Error' + } + }); + } + } \ No newline at end of file diff --git a/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/types.gen.ts b/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/types.gen.ts index 12e2a71fd645b..4ac26c50b1a40 100644 --- a/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/types.gen.ts +++ b/providers/edge3/src/airflow/providers/edge3/plugins/www/openapi-gen/requests/types.gen.ts @@ -523,6 +523,20 @@ export type RemoveWorkerQueueData = { export type RemoveWorkerQueueResponse = unknown; +export type ConcurrencyRequest = { + /** + * New concurrency limit for the worker. + */ + concurrency: number; +}; + +export type SetWorkerConcurrencyLimitData = { + workerName: string; + requestBody: ConcurrencyRequest; +}; + +export type SetWorkerConcurrencyLimitResponse = unknown; + export type $OpenApiTs = { '/edge_worker/v1/jobs/fetch/{worker_name}': { post: { @@ -831,4 +845,19 @@ export type $OpenApiTs = { }; }; }; + '/edge_worker/ui/worker/{worker_name}/concurrency': { + patch: { + req: SetWorkerConcurrencyLimitData; + res: { + /** + * Successful Response + */ + 200: unknown; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; }; \ No newline at end of file diff --git a/providers/edge3/src/airflow/providers/edge3/plugins/www/src/components/WorkerConcurrencyButton.tsx b/providers/edge3/src/airflow/providers/edge3/plugins/www/src/components/WorkerConcurrencyButton.tsx new file mode 100644 index 0000000000000..90b8404cb8e50 --- /dev/null +++ b/providers/edge3/src/airflow/providers/edge3/plugins/www/src/components/WorkerConcurrencyButton.tsx @@ -0,0 +1,157 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { + Button, + CloseButton, + Dialog, + IconButton, + Input, + Portal, + Text, + VStack, + useDisclosure, +} from "@chakra-ui/react"; +import { useUiServiceSetWorkerConcurrencyLimit } from "openapi/queries"; +import type { Worker } from "openapi/requests/types.gen"; +import { useState } from "react"; +import { LuSlidersHorizontal } from "react-icons/lu"; + +interface WorkerConcurrencyButtonProps { + onConcurrencyUpdate: (toast: Record) => void; + worker: Worker; +} + +export const WorkerConcurrencyButton = ({ + onConcurrencyUpdate, + worker, +}: WorkerConcurrencyButtonProps) => { + const { onClose, onOpen, open } = useDisclosure(); + const workerName = worker.worker_name; + const currentConcurrency = worker.sysinfo?.concurrency; + const [concurrency, setConcurrency] = useState( + currentConcurrency !== undefined ? String(currentConcurrency) : "", + ); + + const setConcurrencyMutation = useUiServiceSetWorkerConcurrencyLimit({ + onError: (error: unknown) => { + onConcurrencyUpdate({ + description: `Unable to set concurrency for worker ${workerName}: ${error}`, + title: "Set Concurrency Failed", + type: "error", + }); + }, + onSuccess: () => { + onConcurrencyUpdate({ + description: `Concurrency for worker ${workerName} set to ${concurrency}.`, + title: "Concurrency Updated", + type: "success", + }); + onClose(); + }, + }); + + const handleSetConcurrency = () => { + const value = parseInt(concurrency, 10); + + if (!concurrency.trim() || isNaN(value) || value <= 0) { + onConcurrencyUpdate({ + description: "Please enter a valid concurrency value greater than 0.", + title: "Invalid Input", + type: "error", + }); + return; + } + + setConcurrencyMutation.mutate({ + requestBody: { concurrency: value }, + workerName, + }); + }; + + const handleOpen = () => { + setConcurrency(currentConcurrency !== undefined ? String(currentConcurrency) : ""); + onOpen(); + }; + + const concurrencyValue = parseInt(concurrency, 10); + const isValid = concurrency.trim() !== "" && !isNaN(concurrencyValue) && concurrencyValue > 0; + + return ( + <> + + + + + + + + + + + Set Concurrency for {workerName} + + + + Enter the new concurrency limit for this worker: + setConcurrency(e.target.value)} + onKeyDown={(e) => { + if (e.key === "Enter") { + handleSetConcurrency(); + } + }} + placeholder="Concurrency limit" + type="number" + value={concurrency} + /> + + + + + + + + + + + + + + + + + ); +}; diff --git a/providers/edge3/src/airflow/providers/edge3/plugins/www/src/components/WorkerOperations.tsx b/providers/edge3/src/airflow/providers/edge3/plugins/www/src/components/WorkerOperations.tsx index 0d0aa727eea73..1b5322ce08bf5 100644 --- a/providers/edge3/src/airflow/providers/edge3/plugins/www/src/components/WorkerOperations.tsx +++ b/providers/edge3/src/airflow/providers/edge3/plugins/www/src/components/WorkerOperations.tsx @@ -26,6 +26,7 @@ import { MaintenanceEditCommentButton } from "./MaintenanceEditCommentButton"; import { MaintenanceEnterButton } from "./MaintenanceEnterButton"; import { MaintenanceExitButton } from "./MaintenanceExitButton"; import { RemoveQueueButton } from "./RemoveQueueButton"; +import { WorkerConcurrencyButton } from "./WorkerConcurrencyButton"; import { WorkerDeleteButton } from "./WorkerDeleteButton"; import { WorkerShutdownButton } from "./WorkerShutdownButton"; @@ -48,6 +49,7 @@ export const WorkerOperations = ({ onOperations, worker }: WorkerOperationsProps + diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/datamodels_ui.py b/providers/edge3/src/airflow/providers/edge3/worker_api/datamodels_ui.py index 3aa2e12e9dfed..e38671928a60a 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/datamodels_ui.py +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/datamodels_ui.py @@ -77,3 +77,9 @@ class QueueUpdateRequest(BaseModel): """Request body for queue operations.""" queue_name: Annotated[str, Field(description="Name of the queue to add or remove.")] + + +class ConcurrencyRequest(BaseModel): + """Request body for worker concurrency update.""" + + concurrency: Annotated[int, Field(description="New concurrency limit for the worker.", gt=0)] diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/ui.py b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/ui.py index f24b1e19729fa..996a3a261283a 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/ui.py +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/ui.py @@ -38,8 +38,10 @@ remove_worker_queues, request_maintenance, request_shutdown, + set_worker_concurrency, ) from airflow.providers.edge3.worker_api.datamodels_ui import ( + ConcurrencyRequest, Job, JobCollectionResponse, MaintenanceRequest, @@ -325,3 +327,26 @@ def remove_worker_queue( remove_worker_queues(worker_name, [queue_name], session=session) except Exception as e: raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=str(e)) + + +@ui_router.patch( + "/worker/{worker_name}/concurrency", + dependencies=[ + Depends(requires_access_view(access_view=AccessView.JOBS)), + ], +) +def set_worker_concurrency_limit( + worker_name: str, + concurrency_request: ConcurrencyRequest, + session: SessionDep, +) -> None: + """Set the concurrency limit for an edge worker.""" + worker_query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name) + worker = session.scalar(worker_query) + if not worker: + raise HTTPException(status.HTTP_404_NOT_FOUND, detail=f"Worker {worker_name} not found") + + try: + set_worker_concurrency(worker_name, concurrency_request.concurrency, session=session) + except Exception as e: + raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=str(e)) diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml index a5daeacb86fbd..2904a2e0d3d5d 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml @@ -912,6 +912,41 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + /edge_worker/ui/worker/{worker_name}/concurrency: + patch: + tags: + - UI + summary: Set Worker Concurrency Limit + description: Set the concurrency limit for an edge worker. + operationId: set_worker_concurrency_limit + security: + - OAuth2PasswordBearer: [] + - HTTPBearer: [] + parameters: + - name: worker_name + in: path + required: true + schema: + type: string + title: Worker Name + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/ConcurrencyRequest' + responses: + '200': + description: Successful Response + content: + application/json: + schema: {} + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' components: schemas: BundleInfo: @@ -929,6 +964,18 @@ components: - name title: BundleInfo description: Schema for telling task which bundle to run with. + ConcurrencyRequest: + properties: + concurrency: + type: integer + exclusiveMinimum: 0.0 + title: Concurrency + description: New concurrency limit for the worker. + type: object + required: + - concurrency + title: ConcurrencyRequest + description: Request body for worker concurrency update. EdgeJobFetched: properties: dag_id: diff --git a/providers/edge3/tests/unit/edge3/worker_api/routes/test_ui.py b/providers/edge3/tests/unit/edge3/worker_api/routes/test_ui.py index f1934b5d0f6dd..256cdd977d396 100644 --- a/providers/edge3/tests/unit/edge3/worker_api/routes/test_ui.py +++ b/providers/edge3/tests/unit/edge3/worker_api/routes/test_ui.py @@ -47,3 +47,30 @@ def test_worker(self, session: Session): assert worker_response.total_entries == 1 assert len(worker_response.workers) == 1 assert worker_response.workers[0].worker_name == "worker1" + + def test_set_worker_concurrency_limit(self, session: Session): + from airflow.providers.edge3.worker_api.datamodels_ui import ConcurrencyRequest + from airflow.providers.edge3.worker_api.routes.ui import set_worker_concurrency_limit + + set_worker_concurrency_limit( + worker_name="worker1", + concurrency_request=ConcurrencyRequest(concurrency=4), + session=session, + ) + worker_model = session.get(EdgeWorkerModel, "worker1") + assert worker_model is not None + assert worker_model.concurrency == 4 + + def test_set_worker_concurrency_limit_not_found(self, session: Session): + from fastapi import HTTPException + + from airflow.providers.edge3.worker_api.datamodels_ui import ConcurrencyRequest + from airflow.providers.edge3.worker_api.routes.ui import set_worker_concurrency_limit + + with pytest.raises(HTTPException) as exc_info: + set_worker_concurrency_limit( + worker_name="nonexistent_worker", + concurrency_request=ConcurrencyRequest(concurrency=4), + session=session, + ) + assert exc_info.value.status_code == 404