apache-airflow-providers-edge3 1.2.0__tar.gz → 1.3.0rc1__tar.gz

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-edge3
-Version: 1.2.0
+Version: 1.3.0rc1
 Summary: Provider package apache-airflow-providers-edge3 for Apache Airflow
 Keywords: airflow-provider,edge3,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -20,12 +20,12 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: System :: Monitoring
-Requires-Dist: apache-airflow>=2.10.0
+Requires-Dist: apache-airflow>=2.10.0rc1
 Requires-Dist: pydantic>=2.11.0
 Requires-Dist: retryhttp>=1.2.0,!=1.3.0
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.2.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.2.0
+Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-edge3/1.3.0/changelog.html
+Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-edge3/1.3.0
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
@@ -56,7 +56,7 @@ Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 
 Package ``apache-airflow-providers-edge3``
 
-Release: ``1.2.0``
+Release: ``1.3.0``
 
 
 Handle edge workers on remote sites via HTTP(s) connection and orchestrates work over distributed sites.
@@ -81,7 +81,7 @@ This is a provider package for ``edge3`` provider. All classes for this provider
 are in ``airflow.providers.edge3`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.2.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.3.0/>`_.
 
 Installation
 ------------
@@ -104,5 +104,5 @@ PIP package         Version required
 ==================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.2.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.3.0/changelog.html>`_.
 

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/README.rst

@@ -23,7 +23,7 @@
 
 Package ``apache-airflow-providers-edge3``
 
-Release: ``1.2.0``
+Release: ``1.3.0``
 
 
 Handle edge workers on remote sites via HTTP(s) connection and orchestrates work over distributed sites.
@@ -48,7 +48,7 @@ This is a provider package for ``edge3`` provider. All classes for this provider
 are in ``airflow.providers.edge3`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.2.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.3.0/>`_.
 
 Installation
 ------------
@@ -71,4 +71,4 @@ PIP package         Version required
 ==================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.2.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.3.0/changelog.html>`_.

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/docs/architecture.rst

@@ -154,7 +154,6 @@ The following features are known missing and will be implemented in increments:
 - API token per worker: Today there is a global API token available only
 - Edge Worker Plugin
 
-  - Make plugin working on Airflow 3.0, depending on AIP-68
   - Overview about queues / jobs per queue
   - Allow starting Edge Worker REST API separate to api-server
   - Add some hints how to setup an additional worker

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/docs/changelog.rst

@@ -27,6 +27,47 @@
 Changelog
 ---------
 
+1.3.0
+.....
+
+
+Release Date: ``|PypiReleaseDate|``
+
+Features
+~~~~~~~~
+
+* ``Add Color Scheme to all Operations Buttons in Edge React UI (#55637)``
+* ``Add shutdown-all-workers command to Edge CLI (#55626)``
+* ``add yellow hover background to enter maintenance icon (#55631)``
+* ``Add queue management UI buttons for Edge workers (#55625)``
+* ``Add Worker Maintenance Comment Change to React UI (#55547)``
+* ``feat: Add delete button for offline edge workers (#55529)``
+* ``feat: Add shutdown button for edge workers with confirmation dialog (#55513)``
+* ``Add confirmation dialog for exit maintenance action in Edge Worker Page (#55400)``
+* ``Add Links to Edge React UI (#55356)``
+* ``Feature/edge maintenance plugin beautification (#55348)``
+* ``Add worker maintenance mode functionality to Edge3 provider UI (#55301)``
+
+Bug Fixes
+~~~~~~~~~
+
+* ``Fix EdgeWorker multiprocessing pickle error on Windows (#55284)``
+
+Misc
+~~~~
+
+* ``Pick/vite vitest full (#55623)``
+* ``Bump React and React-DOM (#55598)``
+* ``Bump axios from 1.11.0 to 1.12.0 in /providers/edge3/src/airflow/providers/edge3/plugins/www (#55550)``
+* ``Some small UI polishing for Edge React UI (#55545)``
+* ``Adjust Edge color scheme after merge of PR 53981 (#55485)``
+* ``Bump vite from 5.4.19 to 5.4.20 in /providers/edge3/src/airflow/providers/edge3/plugins/www (#55449)``
+
+Doc-only
+~~~~~~~~
+
+* ``Update docs for new Airflow 3.1 UI Plugins (#55654)``
+
 1.2.0
 .....
 

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/docs/deployment.rst

@@ -130,14 +130,14 @@ Also please note if the worker detects that the Airflow or Edge provider package
 is not the same as the one running on the API server, it will stop accepting new tasks and shut down gracefully.
 This is to prevent running tasks with different versions of the code.
 
-Worker status can be checked via the web UI in the "Admin" - "Edge Worker Hosts" page.
+Worker status can be checked via the web UI in the "Admin" - "Edge Worker" page. Maintenance can be set on a per worker basis.
 
 .. image:: img/worker_hosts.png
 
 .. note::
 
-    As of time of writing the web UI to see edge jobs and manage workers is not ported to Airflow 3.0.
-    Until this is available you can use the CLI commands as described in :ref:`deployment:maintenance-mgmt-cli`.
+    Airflow 3.0 does not support UI plugins. The UI plugin is only available in Airflow 2.10 and in 3.1 and newer.
+    Alternatively you can use the CLI commands as described in :ref:`deployment:maintenance-mgmt-cli`.
 
 
 Worker maintenance can also be triggered via the CLI command on the machine that runs the worker.
@@ -160,9 +160,7 @@ use the command
 This will show the status of the local worker instance as JSON and the tasks running on it.
 
 The status and maintenance comments will also be shown in the web UI
-in the "Admin" - "Edge Worker Hosts" page.
-
-.. image:: img/worker_maintenance.png
+in the "Admin" - "Edge Worker" page.
 
 The local worker instance can be started to fetch new tasks via the command
 
@@ -189,6 +187,7 @@ instance. The commands are:
 - ``airflow edge remote-edge-worker-update-maintenance-comment``: Updates the maintenance comment for a remote edge worker
 - ``airflow edge remote-edge-worker-exit-maintenance``: Request a remote edge worker to exit maintenance mode
 - ``airflow edge shutdown-remote-edge-worker``: Shuts down a remote edge worker gracefully
+- ``airflow edge shutdown-all-workers``: Request graceful shutdown of all registered edge workers
 - ``airflow edge remove-remote-edge-worker``: Remove a worker instance from the cluster
 - ``airflow edge add-worker-queues``: Add queues to an edge worker
 - ``airflow edge remove-worker-queues``: Remove queues from an edge worker

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/docs/edge_executor.rst

@@ -109,4 +109,4 @@ Current Limitations Edge Executor
     the logs will be scattered across the webserver instances.
   - Performance: No extensive performance assessment and scaling tests have been made. The edge executor package is
     optimized for stability. This will be incrementally improved in future releases. Setups have reported stable
-    operation with ~50 workers until now. Note that executed tasks require more webserver API capacity.
+    operation with ~80 workers until now. Note that executed tasks require more webserver API capacity.

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/docs/index.rst

@@ -90,7 +90,7 @@ Additional REST API endpoints are provided to distribute tasks and manage the ed
 are provided by the API server.
 
 
-Release: 1.2.0
+Release: 1.3.0
 
 Provider package
 ----------------

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/docs/install_on_windows.rst

@@ -25,6 +25,8 @@ Install Edge Worker on Windows
     Windows-based setup is intended solely for testing at your own risk. It is technically limited
     due to Python OS restrictions and if currently of Proof-of-Concept quality.
 
+    It is known that due to changes in Task SDK running in Windows is not possible.
+
 
 The setup was tested on Windows 10 with Python 3.12.8, 64-bit. Backend for tests was Airflow 2.10.5.
 To setup a instance of Edge Worker on Windows, you need to follow the steps below:

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/docs/ui_plugin.rst

@@ -22,8 +22,7 @@ The Edge provider uses a Plugin to
 
 - Extend the REST API endpoints for connecting workers to the Airflow cluster
 - Provide a web UI for managing the workers and monitoring their status and tasks
-  (Note: The UI is currently only available in Airflow 2.10+, implementation for
-  Airflow 3.0 depends on completion of AIP-68)
+  (Note: Airflow 3.0 does not have support for UI plugins. The UI plugin is only available in Airflow 2.10 and in 3.1 and newer.)
 
 REST API endpoints
 ------------------
@@ -38,29 +37,27 @@ The Edge provider adds the following REST API endpoints to the Airflow API:
 To see full documentation of the API endpoints open the Airflow web UI and navigate to
 the sub-path ``/edge_worker/docs`` (Airflow 3.0) or ``/edge_worker/v1/ui`` (Airflow 2.10).
 
-Web UI Plugin (Airflow 2.10 only)
----------------------------------
+Web UI Plugin
+-------------
 
 .. note::
 
-    As of time of writing the web UI to see edge jobs and manage workers is not ported to Airflow 3.0.
-    Until this is available you can use the CLI commands as described in :ref:`deployment:maintenance-mgmt-cli`.
+    Airflow 3.0 does not support UI plugins. The UI plugin is only available in Airflow 2.10 and in 3.1 and newer.
+    Alternatively you can use the CLI commands as described in :ref:`deployment:maintenance-mgmt-cli`.
 
 The Edge provider adds a web UI plugin to the Airflow web UI. The plugin is
 made to be able to see job queue and Edge Worker status.
 
 Pending and processes tasks can be checked in "Admin" - "Edge Worker Jobs" page.
 
-Worker status can be checked via the web UI in the "Admin" - "Edge Worker Hosts" page.
+Worker status can be checked via the web UI in the "Admin" - "Edge Worker" page.
 
 .. image:: img/worker_hosts.png
 
 Via the UI you can also set the status of the worker to "Maintenance" or "Active".
 
 The status and maintenance comments will also be shown in the web UI
-in the "Admin" - "Edge Worker Hosts" page.
-
-.. image:: img/worker_maintenance.png
+in the "Admin" - "Edge Worker" page.
 
 Note that maintenance mode can also be adjusted via CLI.
 See :ref:`deployment:maintenance` for more details.

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/provider.yaml

@@ -33,13 +33,14 @@ description: |
   are provided by the API server.
 
 state: ready
-source-date-epoch: 1756876811
+source-date-epoch: 1757950232
 
 # Note that those versions are maintained by release manager - do not update them manually
 # with the exception of case where other provider in sources has >= new provider version.
 # In such case adding >= NEW_VERSION and bumping to NEW_VERSION in a provider have
 # to be done in the same PR
 versions:
+  - 1.3.0
   - 1.2.0
   - 1.1.3
   - 1.1.2

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-edge3"
-version = "1.2.0"
+version = "1.3.0rc1"
 description = "Provider package apache-airflow-providers-edge3 for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -57,7 +57,7 @@ requires-python = ">=3.10"
 # Make sure to run ``prek update-providers-dependencies --all-files``
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
-    "apache-airflow>=2.10.0",
+    "apache-airflow>=2.10.0rc1",
     "pydantic>=2.11.0",
     "retryhttp>=1.2.0,!=1.3.0",
 ]
@@ -96,8 +96,8 @@ apache-airflow-providers-common-sql = {workspace = true}
 apache-airflow-providers-standard = {workspace = true}
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.2.0"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.2.0/changelog.html"
+"Documentation" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-edge3/1.3.0"
+"Changelog" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-edge3/1.3.0/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/src/airflow/providers/edge3/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "1.2.0"
+__version__ = "1.3.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/src/airflow/providers/edge3/cli/edge_command.py

@@ -351,6 +351,37 @@ def remote_worker_request_shutdown(args) -> None:
     logger.info("Requested shutdown of Edge Worker host %s by %s.", args.edge_hostname, getuser())
 
 
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def shutdown_all_workers(args) -> None:
+    """Request graceful shutdown of all edge workers."""
+    _check_valid_db_connection()
+    if not (
+        args.yes
+        or input("This will shutdown all active edge workers, this cannot be undone! Proceed? (y/n)").upper()
+        == "Y"
+    ):
+        raise SystemExit("Cancelled")
+
+    from airflow.providers.edge3.models.edge_worker import get_registered_edge_hosts, request_shutdown
+
+    all_hosts = list(get_registered_edge_hosts())
+    if not all_hosts:
+        logger.info("No edge workers found to shutdown.")
+        return
+
+    shutdown_count = 0
+    for host in all_hosts:
+        try:
+            request_shutdown(host.worker_name)
+            logger.info("Requested shutdown of Edge Worker host %s", host.worker_name)
+            shutdown_count += 1
+        except Exception as e:
+            logger.error("Failed to shutdown Edge Worker host %s: %s", host.worker_name, e)
+
+    logger.info("Requested shutdown of %d edge workers by %s.", shutdown_count, getuser())
+
+
 @cli_utils.action_cli(check_db=False)
 @providers_configuration_loaded
 def add_worker_queues(args) -> None:
@@ -468,6 +499,12 @@ ARG_UMASK = Arg(
 ARG_STDERR = Arg(("--stderr",), help="Redirect stderr to this file if run in daemon mode")
 ARG_STDOUT = Arg(("--stdout",), help="Redirect stdout to this file if run in daemon mode")
 ARG_LOG_FILE = Arg(("-l", "--log-file"), help="Location of the log file if run in daemon mode")
+ARG_YES = Arg(
+    ("-y", "--yes"),
+    help="Skip confirmation prompt and proceed with shutdown",
+    action="store_true",
+    default=False,
+)
 
 EDGE_COMMANDS: list[ActionCommand] = [
     ActionCommand(
@@ -581,4 +618,10 @@ EDGE_COMMANDS: list[ActionCommand] = [
             ARG_QUEUES_MANAGE,
         ),
     ),
+    ActionCommand(
+        name="shutdown-all-workers",
+        help=shutdown_all_workers.__doc__,
+        func=shutdown_all_workers,
+        args=(ARG_YES,),
+    ),
 ]

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/src/airflow/providers/edge3/cli/worker.py

@@ -175,48 +175,46 @@ class EdgeWorker:
             return EdgeWorkerState.MAINTENANCE_MODE
         return EdgeWorkerState.IDLE
 
-    def _launch_job_af3(self, edge_job: EdgeJobFetched) -> tuple[Process, Path]:
-        if TYPE_CHECKING:
-            from airflow.executors.workloads import ExecuteTask
+    @staticmethod
+    def _run_job_via_supervisor(workload) -> int:
+        from airflow.sdk.execution_time.supervisor import supervise
 
-        def _run_job_via_supervisor(
-            workload: ExecuteTask,
-        ) -> int:
-            from airflow.sdk.execution_time.supervisor import supervise
+        # Ignore ctrl-c in this process -- we don't want to kill _this_ one. we let tasks run to completion
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
 
-            # Ignore ctrl-c in this process -- we don't want to kill _this_ one. we let tasks run to completion
-            signal.signal(signal.SIGINT, signal.SIG_IGN)
+        logger.info("Worker starting up pid=%d", os.getpid())
+        setproctitle(f"airflow edge worker: {workload.ti.key}")
 
-            logger.info("Worker starting up pid=%d", os.getpid())
-            setproctitle(f"airflow edge worker: {workload.ti.key}")
+        try:
+            api_url = conf.get("edge", "api_url")
+            execution_api_server_url = conf.get("core", "execution_api_server_url", fallback="")
+            if not execution_api_server_url:
+                parsed = urlparse(api_url)
+                execution_api_server_url = f"{parsed.scheme}://{parsed.netloc}/execution/"
+
+            supervise(
+                # This is the "wrong" ti type, but it duck types the same. TODO: Create a protocol for this.
+                # Same like in airflow/executors/local_executor.py:_execute_work()
+                ti=workload.ti,  # type: ignore[arg-type]
+                dag_rel_path=workload.dag_rel_path,
+                bundle_info=workload.bundle_info,
+                token=workload.token,
+                server=execution_api_server_url,
+                log_path=workload.log_path,
+            )
+            return 0
+        except Exception as e:
+            logger.exception("Task execution failed: %s", e)
+            return 1
 
-            try:
-                api_url = conf.get("edge", "api_url")
-                execution_api_server_url = conf.get("core", "execution_api_server_url", fallback="")
-                if not execution_api_server_url:
-                    parsed = urlparse(api_url)
-                    execution_api_server_url = f"{parsed.scheme}://{parsed.netloc}/execution/"
-
-                logger.info("Worker starting up server=execution_api_server_url=%s", execution_api_server_url)
-
-                supervise(
-                    # This is the "wrong" ti type, but it duck types the same. TODO: Create a protocol for this.
-                    # Same like in airflow/executors/local_executor.py:_execute_work()
-                    ti=workload.ti,  # type: ignore[arg-type]
-                    dag_rel_path=workload.dag_rel_path,
-                    bundle_info=workload.bundle_info,
-                    token=workload.token,
-                    server=execution_api_server_url,
-                    log_path=workload.log_path,
-                )
-                return 0
-            except Exception as e:
-                logger.exception("Task execution failed: %s", e)
-                return 1
+    @staticmethod
+    def _launch_job_af3(edge_job: EdgeJobFetched) -> tuple[Process, Path]:
+        if TYPE_CHECKING:
+            from airflow.executors.workloads import ExecuteTask
 
         workload: ExecuteTask = edge_job.command
         process = Process(
-            target=_run_job_via_supervisor,
+            target=EdgeWorker._run_job_via_supervisor,
             kwargs={"workload": workload},
         )
         process.start()
@@ -226,7 +224,8 @@ class EdgeWorker:
         logfile = Path(base_log_folder, workload.log_path)
         return process, logfile
 
-    def _launch_job_af2_10(self, edge_job: EdgeJobFetched) -> tuple[Popen, Path]:
+    @staticmethod
+    def _launch_job_af2_10(edge_job: EdgeJobFetched) -> tuple[Popen, Path]:
         """Compatibility for Airflow 2.10 Launch."""
         env = os.environ.copy()
         env["AIRFLOW__CORE__DATABASE_ACCESS_ISOLATION"] = "True"
@@ -237,14 +236,15 @@ class EdgeWorker:
         logfile = logs_logfile_path(edge_job.key)
         return process, logfile
 
-    def _launch_job(self, edge_job: EdgeJobFetched):
+    @staticmethod
+    def _launch_job(edge_job: EdgeJobFetched):
         """Get the received job executed."""
         process: Popen | Process
         if AIRFLOW_V_3_0_PLUS:
-            process, logfile = self._launch_job_af3(edge_job)
+            process, logfile = EdgeWorker._launch_job_af3(edge_job)
         else:
             # Airflow 2.10
-            process, logfile = self._launch_job_af2_10(edge_job)
+            process, logfile = EdgeWorker._launch_job_af2_10(edge_job)
         EdgeWorker.jobs.append(Job(edge_job, process, logfile, 0))
 
     def start(self):
@@ -316,7 +316,7 @@ class EdgeWorker:
         edge_job = jobs_fetch(self.hostname, self.queues, self.free_concurrency)
         if edge_job:
             logger.info("Received job: %s", edge_job)
-            self._launch_job(edge_job)
+            EdgeWorker._launch_job(edge_job)
             jobs_set_state(edge_job.key, TaskInstanceState.RUNNING)
             return True
 

{apache_airflow_providers_edge3-1.2.0 → apache_airflow_providers_edge3-1.3.0rc1}/src/airflow/providers/edge3/models/edge_worker.py

@@ -56,21 +56,21 @@ class EdgeWorkerState(str, Enum):
     IDLE = "idle"
     """Edge Worker is active and waiting for a task."""
     SHUTDOWN_REQUEST = "shutdown request"
-    """Request to shutdown Edge Worker."""
+    """Request to shutdown Edge Worker is issued. It will be picked-up on the next heartbeat, tasks will drain and then worker will terminate."""
     TERMINATING = "terminating"
-    """Edge Worker is completing work and stopping."""
+    """Edge Worker is completing work (draining running tasks) and stopping."""
     OFFLINE = "offline"
     """Edge Worker was shut down."""
     UNKNOWN = "unknown"
-    """No heartbeat signal from worker for some time, Edge Worker probably down."""
+    """No heartbeat signal from worker for some time, Edge Worker probably down or got disconnected."""
     MAINTENANCE_REQUEST = "maintenance request"
-    """Worker was requested to enter maintenance mode. Once worker receives this it will pause fetching jobs."""
+    """Worker was requested to enter maintenance mode. Once worker receives this message it will pause fetching tasks and drain tasks."""
     MAINTENANCE_PENDING = "maintenance pending"
-    """Edge worker received the request for maintenance, waiting for jobs to finish. Once jobs are finished will move to 'maintenance mode'."""
+    """Edge Worker received the request for maintenance, waiting for tasks to finish. Once tasks are finished will move to 'maintenance mode'."""
     MAINTENANCE_MODE = "maintenance mode"
-    """Edge worker is in maintenance mode. It is online but pauses fetching jobs."""
+    """Edge Worker is in maintenance mode. It is online but pauses fetching tasks."""
     MAINTENANCE_EXIT = "maintenance exit"
-    """Request worker to exit maintenance mode. Once the worker receives this state it will un-pause and fetch new jobs."""
+    """Request Worker is requested to exit maintenance mode. Once the worker receives this state it will un-pause and fetch new tasks."""
     OFFLINE_MAINTENANCE = "offline maintenance"
     """Worker was shut down in maintenance mode. It will be in maintenance mode when restarted."""
 
@@ -264,7 +264,12 @@ def change_maintenance_comment(
     """Write maintenance comment in the db."""
     query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
     worker: EdgeWorkerModel = session.scalar(query)
-    if worker.state in (EdgeWorkerState.MAINTENANCE_MODE, EdgeWorkerState.OFFLINE_MAINTENANCE):
+    if worker.state in (
+        EdgeWorkerState.MAINTENANCE_MODE,
+        EdgeWorkerState.MAINTENANCE_PENDING,
+        EdgeWorkerState.MAINTENANCE_REQUEST,
+        EdgeWorkerState.OFFLINE_MAINTENANCE,
+    ):
         worker.maintenance_comment = maintenance_comment
     else:
         error_message = f"Cannot change maintenance comment as {worker_name} is not in maintenance!"