apache-airflow-providers-edge3 2.0.0rc1__py3-none-any.whl → 3.0.1rc1__py3-none-any.whl

airflow/providers/edge3/__init__.py

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

airflow/providers/edge3/cli/api_client.py

@@ -16,7 +16,6 @@
 # under the License.
 from __future__ import annotations
 
-import json
 import logging
 import os
 from datetime import datetime
@@ -26,7 +25,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any
 from urllib.parse import quote, urljoin
 
-import requests
+from aiohttp import ClientConnectionError, ClientResponseError, ServerTimeoutError, request
 from retryhttp import retry, wait_retry_after
 from tenacity import before_sleep_log, wait_random_exponential
 
@@ -44,11 +43,11 @@ from airflow.providers.edge3.worker_api.datamodels import (
     WorkerSetStateReturn,
     WorkerStateBody,
 )
-from airflow.utils.state import TaskInstanceState  # noqa: TC001
 
 if TYPE_CHECKING:
-    from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.providers.common.compat.sdk import TaskInstanceKey
     from airflow.providers.edge3.models.edge_worker import EdgeWorkerState
+    from airflow.utils.state import TaskInstanceState
 
 logger = logging.getLogger(__name__)
 
@@ -92,39 +91,42 @@ def jwt_generator() -> JWTGenerator:
     wait_timeouts=_default_wait,
     wait_rate_limited=wait_retry_after(fallback=_default_wait),  # No infinite timeout on HTTP 429
     before_sleep=before_sleep_log(logger, logging.WARNING),
+    network_errors=ClientConnectionError,
+    timeouts=ServerTimeoutError,
 )
-def _make_generic_request(method: str, rest_path: str, data: str | None = None) -> Any:
+async def _make_generic_request(method: str, rest_path: str, data: str | None = None) -> Any:
     authorization = jwt_generator().generate({"method": rest_path})
     api_url = conf.get("edge", "api_url")
+    content_type = {"Content-Type": "application/json"} if data else {}
     headers = {
-        "Content-Type": "application/json",
+        **content_type,
         "Accept": "application/json",
         "Authorization": authorization,
     }
     api_endpoint = urljoin(api_url, rest_path)
-    response = requests.request(method, url=api_endpoint, data=data, headers=headers)
-    response.raise_for_status()
-    if response.status_code == HTTPStatus.NO_CONTENT:
-        return None
-    return json.loads(response.content)
+    async with request(method, url=api_endpoint, data=data, headers=headers) as response:
+        response.raise_for_status()
+        if response.status == HTTPStatus.NO_CONTENT:
+            return None
+        return await response.json()
 
 
-def worker_register(
+async def worker_register(
     hostname: str, state: EdgeWorkerState, queues: list[str] | None, sysinfo: dict
 ) -> WorkerRegistrationReturn:
     """Register worker with the Edge API."""
     try:
-        result = _make_generic_request(
+        result = await _make_generic_request(
             "POST",
             f"worker/{quote(hostname)}",
             WorkerStateBody(state=state, jobs_active=0, queues=queues, sysinfo=sysinfo).model_dump_json(
                 exclude_unset=True
             ),
         )
-    except requests.HTTPError as e:
-        if e.response.status_code == 400:
+    except ClientResponseError as e:
+        if e.status == HTTPStatus.BAD_REQUEST:
             raise EdgeWorkerVersionException(str(e))
-        if e.response.status_code == 409:
+        if e.status == HTTPStatus.CONFLICT:
             raise EdgeWorkerDuplicateException(
                 f"A worker with the name '{hostname}' is already active. "
                 "Please ensure worker names are unique, or stop the existing worker before starting a new one."
@@ -133,7 +135,7 @@ def worker_register(
     return WorkerRegistrationReturn(**result)
 
 
-def worker_set_state(
+async def worker_set_state(
     hostname: str,
     state: EdgeWorkerState,
     jobs_active: int,
@@ -143,7 +145,7 @@ def worker_set_state(
 ) -> WorkerSetStateReturn:
     """Update the state of the worker in the central site and thereby implicitly heartbeat."""
     try:
-        result = _make_generic_request(
+        result = await _make_generic_request(
             "PATCH",
             f"worker/{quote(hostname)}",
             WorkerStateBody(
@@ -154,16 +156,16 @@ def worker_set_state(
                 maintenance_comments=maintenance_comments,
             ).model_dump_json(exclude_unset=True),
         )
-    except requests.HTTPError as e:
-        if e.response.status_code == 400:
+    except ClientResponseError as e:
+        if e.status == HTTPStatus.BAD_REQUEST:
             raise EdgeWorkerVersionException(str(e))
         raise e
     return WorkerSetStateReturn(**result)
 
 
-def jobs_fetch(hostname: str, queues: list[str] | None, free_concurrency: int) -> EdgeJobFetched | None:
+async def jobs_fetch(hostname: str, queues: list[str] | None, free_concurrency: int) -> EdgeJobFetched | None:
     """Fetch a job to execute on the edge worker."""
-    result = _make_generic_request(
+    result = await _make_generic_request(
         "POST",
         f"jobs/fetch/{quote(hostname)}",
         WorkerQueuesBody(queues=queues, free_concurrency=free_concurrency).model_dump_json(
@@ -175,17 +177,17 @@ def jobs_fetch(hostname: str, queues: list[str] | None, free_concurrency: int) -
     return None
 
 
-def jobs_set_state(key: TaskInstanceKey, state: TaskInstanceState) -> None:
+async def jobs_set_state(key: TaskInstanceKey, state: TaskInstanceState) -> None:
     """Set the state of a job."""
-    _make_generic_request(
+    await _make_generic_request(
         "PATCH",
         f"jobs/state/{key.dag_id}/{key.task_id}/{key.run_id}/{key.try_number}/{key.map_index}/{state}",
     )
 
 
-def logs_logfile_path(task: TaskInstanceKey) -> Path:
+async def logs_logfile_path(task: TaskInstanceKey) -> Path:
     """Elaborate the path and filename to expect from task execution."""
-    result = _make_generic_request(
+    result = await _make_generic_request(
         "GET",
         f"logs/logfile_path/{task.dag_id}/{task.task_id}/{task.run_id}/{task.try_number}/{task.map_index}",
     )
@@ -193,13 +195,13 @@ def logs_logfile_path(task: TaskInstanceKey) -> Path:
     return Path(base_log_folder, result)
 
 
-def logs_push(
+async def logs_push(
     task: TaskInstanceKey,
     log_chunk_time: datetime,
     log_chunk_data: str,
 ) -> None:
     """Push an incremental log chunk from Edge Worker to central site."""
-    _make_generic_request(
+    await _make_generic_request(
         "POST",
         f"logs/push/{task.dag_id}/{task.task_id}/{task.run_id}/{task.try_number}/{task.map_index}",
         PushLogsBody(log_chunk_time=log_chunk_time, log_chunk_data=log_chunk_data).model_dump_json(

airflow/providers/edge3/cli/dataclasses.py

@@ -23,8 +23,6 @@ from pathlib import Path
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from psutil import Popen
-
     from airflow.providers.edge3.models.edge_worker import EdgeWorkerState
     from airflow.providers.edge3.worker_api.datamodels import EdgeJobFetched
 
@@ -74,22 +72,17 @@ class Job:
     """Holds all information for a task/job to be executed as bundle."""
 
     edge_job: EdgeJobFetched
-    process: Popen | Process
+    process: Process
     logfile: Path
-    logsize: int
+    logsize: int = 0
     """Last size of log file, point of last chunk push."""
 
     @property
     def is_running(self) -> bool:
         """Check if the job is still running."""
-        if hasattr(self.process, "returncode") and hasattr(self.process, "poll"):
-            self.process.poll()
-            return self.process.returncode is None
-        return self.process.exitcode is None
+        return self.process.is_alive()
 
     @property
     def is_success(self) -> bool:
         """Check if the job was successful."""
-        if hasattr(self.process, "returncode"):
-            return self.process.returncode == 0
         return self.process.exitcode == 0

airflow/providers/edge3/cli/definition.py

@@ -0,0 +1,261 @@
+# 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.cli.cli_config import ARG_PID, ARG_VERBOSE, ActionCommand, Arg, GroupCommand, lazy_load_command
+from airflow.configuration import conf
+
+if TYPE_CHECKING:
+    import argparse
+
+
+ARG_CONCURRENCY = Arg(
+    ("-c", "--concurrency"),
+    type=int,
+    help="The number of worker processes",
+    default=conf.getint("edge", "worker_concurrency", fallback=8),
+)
+ARG_QUEUES = Arg(
+    ("-q", "--queues"),
+    help="Comma delimited list of queues to serve, serve all queues if not provided.",
+)
+ARG_EDGE_HOSTNAME = Arg(
+    ("-H", "--edge-hostname"),
+    help="Set the hostname of worker if you have multiple workers on a single machine",
+)
+ARG_REQUIRED_EDGE_HOSTNAME = Arg(
+    ("-H", "--edge-hostname"),
+    help="Set the hostname of worker if you have multiple workers on a single machine",
+    required=True,
+)
+ARG_MAINTENANCE = Arg(("maintenance",), help="Desired maintenance state", choices=("on", "off"))
+ARG_MAINTENANCE_COMMENT = Arg(
+    ("-c", "--comments"),
+    help="Maintenance comments to report reason. Required if maintenance is turned on.",
+)
+ARG_REQUIRED_MAINTENANCE_COMMENT = Arg(
+    ("-c", "--comments"),
+    help="Maintenance comments to report reason. Required if enabling maintenance",
+    required=True,
+)
+ARG_QUEUES_MANAGE = Arg(
+    ("-q", "--queues"),
+    help="Comma delimited list of queues to add or remove.",
+    required=True,
+)
+ARG_WAIT_MAINT = Arg(
+    ("-w", "--wait"),
+    default=False,
+    help="Wait until edge worker has reached desired state.",
+    action="store_true",
+)
+ARG_WAIT_STOP = Arg(
+    ("-w", "--wait"),
+    default=False,
+    help="Wait until edge worker is shut down.",
+    action="store_true",
+)
+ARG_OUTPUT = Arg(
+    (
+        "-o",
+        "--output",
+    ),
+    help="Output format. Allowed values: json, yaml, plain, table (default: table)",
+    metavar="(table, json, yaml, plain)",
+    choices=("table", "json", "yaml", "plain"),
+    default="table",
+)
+ARG_STATE = Arg(
+    (
+        "-s",
+        "--state",
+    ),
+    nargs="+",
+    help="State of the edge worker",
+)
+
+ARG_DAEMON = Arg(
+    ("-D", "--daemon"), help="Daemonize instead of running in the foreground", action="store_true"
+)
+ARG_UMASK = Arg(
+    ("-u", "--umask"),
+    help="Set the umask of edge worker in daemon mode",
+)
+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(
+        name="worker",
+        help="Start Airflow Edge Worker.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.worker"),
+        args=(
+            ARG_CONCURRENCY,
+            ARG_QUEUES,
+            ARG_EDGE_HOSTNAME,
+            ARG_PID,
+            ARG_VERBOSE,
+            ARG_DAEMON,
+            ARG_STDOUT,
+            ARG_STDERR,
+            ARG_LOG_FILE,
+            ARG_UMASK,
+        ),
+    ),
+    ActionCommand(
+        name="status",
+        help="Check for Airflow Local Edge Worker status.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.status"),
+        args=(
+            ARG_PID,
+            ARG_VERBOSE,
+        ),
+    ),
+    ActionCommand(
+        name="maintenance",
+        help="Set or Unset maintenance mode of local edge worker.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.maintenance"),
+        args=(
+            ARG_MAINTENANCE,
+            ARG_MAINTENANCE_COMMENT,
+            ARG_WAIT_MAINT,
+            ARG_PID,
+            ARG_VERBOSE,
+        ),
+    ),
+    ActionCommand(
+        name="stop",
+        help="Stop a running local Airflow Edge Worker.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.stop"),
+        args=(
+            ARG_WAIT_STOP,
+            ARG_PID,
+            ARG_VERBOSE,
+        ),
+    ),
+    ActionCommand(
+        name="list-workers",
+        help="Query the db to list all registered edge workers.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.list_edge_workers"),
+        args=(
+            ARG_OUTPUT,
+            ARG_STATE,
+        ),
+    ),
+    ActionCommand(
+        name="remote-edge-worker-request-maintenance",
+        help="Put remote edge worker on maintenance.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.put_remote_worker_on_maintenance"),
+        args=(
+            ARG_REQUIRED_EDGE_HOSTNAME,
+            ARG_REQUIRED_MAINTENANCE_COMMENT,
+        ),
+    ),
+    ActionCommand(
+        name="remote-edge-worker-exit-maintenance",
+        help="Remove remote edge worker from maintenance.",
+        func=lazy_load_command(
+            "airflow.providers.edge3.cli.edge_command.remove_remote_worker_from_maintenance"
+        ),
+        args=(ARG_REQUIRED_EDGE_HOSTNAME,),
+    ),
+    ActionCommand(
+        name="remote-edge-worker-update-maintenance-comment",
+        help="Update maintenance comments of the remote edge worker.",
+        func=lazy_load_command(
+            "airflow.providers.edge3.cli.edge_command.remote_worker_update_maintenance_comment"
+        ),
+        args=(
+            ARG_REQUIRED_EDGE_HOSTNAME,
+            ARG_REQUIRED_MAINTENANCE_COMMENT,
+        ),
+    ),
+    ActionCommand(
+        name="remove-remote-edge-worker",
+        help="Remove remote edge worker entry from db.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.remove_remote_worker"),
+        args=(ARG_REQUIRED_EDGE_HOSTNAME,),
+    ),
+    ActionCommand(
+        name="shutdown-remote-edge-worker",
+        help="Initiate the shutdown of the remote edge worker.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.remote_worker_request_shutdown"),
+        args=(ARG_REQUIRED_EDGE_HOSTNAME,),
+    ),
+    ActionCommand(
+        name="add-worker-queues",
+        help="Add queues to an edge worker.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.add_worker_queues"),
+        args=(
+            ARG_REQUIRED_EDGE_HOSTNAME,
+            ARG_QUEUES_MANAGE,
+        ),
+    ),
+    ActionCommand(
+        name="remove-worker-queues",
+        help="Remove queues from an edge worker.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.remove_worker_queues"),
+        args=(
+            ARG_REQUIRED_EDGE_HOSTNAME,
+            ARG_QUEUES_MANAGE,
+        ),
+    ),
+    ActionCommand(
+        name="shutdown-all-workers",
+        help="Request graceful shutdown of all edge workers.",
+        func=lazy_load_command("airflow.providers.edge3.cli.edge_command.shutdown_all_workers"),
+        args=(ARG_YES,),
+    ),
+]
+
+
+def get_edge_cli_commands() -> list[GroupCommand]:
+    return [
+        GroupCommand(
+            name="edge",
+            help="Edge Worker components",
+            description=(
+                "Start and manage Edge Worker. Works only when using EdgeExecutor. For more information, "
+                "see https://airflow.apache.org/docs/apache-airflow-providers-edge3/stable/edge_executor.html"
+            ),
+            subcommands=EDGE_COMMANDS,
+        ),
+    ]
+
+
+def get_parser() -> argparse.ArgumentParser:
+    """
+    Generate documentation; used by Sphinx.
+
+    :meta private:
+    """
+    from airflow.cli.cli_parser import AirflowHelpFormatter, DefaultHelpParser, _add_command
+
+    parser = DefaultHelpParser(prog="airflow", formatter_class=AirflowHelpFormatter)
+    subparsers = parser.add_subparsers(dest="subcommand", metavar="GROUP_OR_COMMAND")
+    for group_command in get_edge_cli_commands():
+        _add_command(subparsers, group_command)
+    return parser