Revert "Remove PriorityWeightStrategy reference in SDK" (#59828)

* Revert "Remove PriorityWeightStrategy reference in SDK (#59780)"

This reverts commit 60b4ed4.

* Tip by TP

Author: jscheffl

airflow-core/newsfragments/59780.significant.rst

@@ -96,7 +96,6 @@
     PriorityWeightStrategy,
     airflow_priority_weight_strategies,
     airflow_priority_weight_strategies_classes,
-    validate_and_load_priority_weight_strategy,
 )
 from airflow.timetables.base import DagRunInfo, Timetable
 from airflow.triggers.base import BaseTrigger, StartTriggerArgs
@@ -250,15 +249,15 @@ def decode_partition_mapper(var: dict[str, Any]) -> PartitionMapper:
     return partition_mapper_class.deserialize(var[Encoding.VAR])
 
 
-def encode_priority_weight_strategy(var: PriorityWeightStrategy | str) -> str:
+def encode_priority_weight_strategy(var: PriorityWeightStrategy) -> str:
     """
     Encode a priority weight strategy instance.
 
     In this version, we only store the importable string, so the class should not wait
     for any parameters to be passed to it. If you need to store the parameters, you
     should store them in the class itself.
     """
-    priority_weight_strategy_class = type(validate_and_load_priority_weight_strategy(var))
+    priority_weight_strategy_class = type(var)
     if priority_weight_strategy_class in airflow_priority_weight_strategies_classes:
         return airflow_priority_weight_strategies_classes[priority_weight_strategy_class]
     importable_string = qualname(priority_weight_strategy_class)

airflow-core/src/airflow/serialization/serialized_objects.py

@@ -20,9 +20,8 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
-from airflow._shared.module_loading import qualname
 from airflow.task.weight_rule import WeightRule
 
 if TYPE_CHECKING:
@@ -42,22 +41,46 @@ class PriorityWeightStrategy(ABC):
     """
 
     @abstractmethod
-    def get_weight(self, ti: TaskInstance) -> int:
+    def get_weight(self, ti: TaskInstance):
         """Get the priority weight of a task."""
-        raise NotImplementedError("must be implemented by a subclass")
+        ...
+
+    @classmethod
+    def deserialize(cls, data: dict[str, Any]) -> PriorityWeightStrategy:
+        """
+        Deserialize a priority weight strategy from data.
+
+        This is called when a serialized DAG is deserialized. ``data`` will be whatever
+        was returned by ``serialize`` during DAG serialization. The default
+        implementation constructs the priority weight strategy without any arguments.
+        """
+        return cls(**data)
+
+    def serialize(self) -> dict[str, Any]:
+        """
+        Serialize the priority weight strategy for JSON encoding.
+
+        This is called during DAG serialization to store priority weight strategy information
+        in the database. This should return a JSON-serializable dict that will be fed into
+        ``deserialize`` when the DAG is deserialized. The default implementation returns
+        an empty dict.
+        """
+        return {}
 
     def __eq__(self, other: object) -> bool:
         """Equality comparison."""
-        return isinstance(other, type(self))
+        if not isinstance(other, type(self)):
+            return False
+        return self.serialize() == other.serialize()
 
-    def __hash__(self) -> int:
-        return hash(None)
+    def __hash__(self):
+        return hash(self.serialize())
 
 
 class _AbsolutePriorityWeightStrategy(PriorityWeightStrategy):
     """Priority weight strategy that uses the task's priority weight directly."""
 
-    def get_weight(self, ti: TaskInstance) -> int:
+    def get_weight(self, ti: TaskInstance):
         if TYPE_CHECKING:
             assert ti.task
         return ti.task.priority_weight
@@ -81,7 +104,7 @@ def get_weight(self, ti: TaskInstance) -> int:
 class _UpstreamPriorityWeightStrategy(PriorityWeightStrategy):
     """Priority weight strategy that uses the sum of the priority weights of all upstream tasks."""
 
-    def get_weight(self, ti: TaskInstance) -> int:
+    def get_weight(self, ti: TaskInstance):
         if TYPE_CHECKING:
             assert ti.task
         dag = ti.task.get_dag()
@@ -93,19 +116,14 @@ def get_weight(self, ti: TaskInstance) -> int:
 
 
 airflow_priority_weight_strategies: dict[str, type[PriorityWeightStrategy]] = {
-    qualname(_AbsolutePriorityWeightStrategy): _AbsolutePriorityWeightStrategy,
-    qualname(_DownstreamPriorityWeightStrategy): _DownstreamPriorityWeightStrategy,
-    qualname(_UpstreamPriorityWeightStrategy): _UpstreamPriorityWeightStrategy,
     WeightRule.ABSOLUTE: _AbsolutePriorityWeightStrategy,
     WeightRule.DOWNSTREAM: _DownstreamPriorityWeightStrategy,
     WeightRule.UPSTREAM: _UpstreamPriorityWeightStrategy,
 }
 
 
 airflow_priority_weight_strategies_classes = {
-    _AbsolutePriorityWeightStrategy: WeightRule.ABSOLUTE,
-    _DownstreamPriorityWeightStrategy: WeightRule.DOWNSTREAM,
-    _UpstreamPriorityWeightStrategy: WeightRule.UPSTREAM,
+    cls: name for name, cls in airflow_priority_weight_strategies.items()
 }
 
 
@@ -121,6 +139,7 @@ def validate_and_load_priority_weight_strategy(
 
     :meta private:
     """
+    from airflow._shared.module_loading import qualname
     from airflow.serialization.serialized_objects import _get_registered_priority_weight_strategy
 
     if priority_weight_strategy is None:

airflow-core/src/airflow/task/priority_strategy.py

@@ -130,17 +130,12 @@ def create_trigger_in_db(session, trigger, operator=None):
         operator = BaseOperator(task_id="test_ti", dag=dag)
     session.add(dag_model)
 
-    lazy_serdag = LazyDeserializedDAG.from_dag(dag)
-    SerializedDagModel.write_dag(lazy_serdag, bundle_name=bundle_name)
+    SerializedDagModel.write_dag(LazyDeserializedDAG.from_dag(dag), bundle_name=bundle_name)
     session.add(run)
     session.add(trigger_orm)
     session.flush()
     dag_version = DagVersion.get_latest_version(dag.dag_id)
-    task_instance = TaskInstance(
-        lazy_serdag._real_dag.get_task(operator.task_id),
-        run_id=run.run_id,
-        dag_version_id=dag_version.id,
-    )
+    task_instance = TaskInstance(operator, run_id=run.run_id, dag_version_id=dag_version.id)
     task_instance.trigger_id = trigger_orm.id
     session.add(task_instance)
     session.commit()
@@ -445,8 +440,7 @@ async def test_trigger_kwargs_serialization_cleanup(self, session):
 
 
 @pytest.mark.asyncio
-@pytest.mark.usefixtures("testing_dag_bundle")
-async def test_trigger_create_race_condition_38599(session, supervisor_builder):
+async def test_trigger_create_race_condition_38599(session, supervisor_builder, testing_dag_bundle):
     """
     This verifies the resolution of race condition documented in github issue #38599.
     More details in the issue description.
@@ -471,17 +465,14 @@ async def test_trigger_create_race_condition_38599(session, supervisor_builder):
     session.flush()
 
     bundle_name = "testing"
-    with DAG(dag_id="test-dag") as dag:
-        task = PythonOperator(task_id="dummy-task", python_callable=print)
+    dag = DAG(dag_id="test-dag")
     dm = DagModel(dag_id="test-dag", bundle_name=bundle_name)
     session.add(dm)
-
-    lazy_serdag = LazyDeserializedDAG.from_dag(dag)
-    SerializedDagModel.write_dag(lazy_serdag, bundle_name=bundle_name)
+    SerializedDagModel.write_dag(LazyDeserializedDAG.from_dag(dag), bundle_name=bundle_name)
     dag_run = DagRun(dag.dag_id, run_id="abc", run_type="none", run_after=timezone.utcnow())
     dag_version = DagVersion.get_latest_version(dag.dag_id)
     ti = TaskInstance(
-        lazy_serdag._real_dag.get_task(task.task_id),
+        PythonOperator(task_id="dummy-task", python_callable=print),
         run_id=dag_run.run_id,
         state=TaskInstanceState.DEFERRED,
         dag_version_id=dag_version.id,

airflow-core/tests/unit/jobs/test_triggerer_job.py

@@ -34,6 +34,7 @@
 from airflow.providers.standard.operators.python import PythonOperator
 from airflow.sdk import DAG, BaseOperator, TaskGroup, setup, task, task_group, teardown
 from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
+from airflow.task.priority_strategy import PriorityWeightStrategy
 from airflow.task.trigger_rule import TriggerRule
 from airflow.utils.state import TaskInstanceState
 
@@ -1523,7 +1524,7 @@ def test_properties(
         assert op.pool == SerializedBaseOperator.pool
         assert op.pool_slots == SerializedBaseOperator.pool_slots
         assert op.priority_weight == SerializedBaseOperator.priority_weight
-        assert op.weight_rule == "downstream"
+        assert isinstance(op.weight_rule, PriorityWeightStrategy)
         assert op.email == email
         assert op.execution_timeout == execution_timeout
         assert op.retry_delay == retry_delay

airflow-core/tests/unit/models/test_mappedoperator.py

@@ -82,7 +82,7 @@
 from airflow.serialization.definitions.assets import SerializedAsset
 from airflow.serialization.definitions.dag import SerializedDAG
 from airflow.serialization.encoders import ensure_serialized_asset
-from airflow.serialization.serialized_objects import OperatorSerialization, create_scheduler_operator
+from airflow.serialization.serialized_objects import OperatorSerialization
 from airflow.ti_deps.dep_context import DepContext
 from airflow.ti_deps.dependencies_deps import REQUEUEABLE_DEPS, RUNNING_DEPS
 from airflow.ti_deps.dependencies_states import RUNNABLE_STATES
@@ -2812,7 +2812,7 @@ def mock_policy(task_instance: TaskInstance):
 
         monkeypatch.setattr("airflow.models.taskinstance.task_instance_mutation_hook", mock_policy)
 
-    sdk_task = EmptyOperator(
+    task = EmptyOperator(
         task_id="empty",
         queue=default_queue,
         pool="test_pool1",
@@ -2822,28 +2822,27 @@ def mock_policy(task_instance: TaskInstance):
         retries=30,
         executor_config={"KubernetesExecutor": {"image": "myCustomDockerImage"}},
     )
-    ser_task = create_scheduler_operator(sdk_task)
-    ti = TI(ser_task, run_id=None, dag_version_id=mock.MagicMock())
-    ti.refresh_from_task(ser_task, pool_override=pool_override)
+    ti = TI(task, run_id=None, dag_version_id=mock.MagicMock())
+    ti.refresh_from_task(task, pool_override=pool_override)
 
     assert ti.queue == expected_queue
 
     if pool_override:
         assert ti.pool == pool_override
     else:
-        assert ti.pool == sdk_task.pool
+        assert ti.pool == task.pool
 
-    assert ti.pool_slots == sdk_task.pool_slots
-    assert ti.priority_weight == ser_task.weight_rule.get_weight(ti)
-    assert ti.run_as_user == sdk_task.run_as_user
-    assert ti.max_tries == sdk_task.retries
-    assert ti.executor_config == sdk_task.executor_config
+    assert ti.pool_slots == task.pool_slots
+    assert ti.priority_weight == task.weight_rule.get_weight(ti)
+    assert ti.run_as_user == task.run_as_user
+    assert ti.max_tries == task.retries
+    assert ti.executor_config == task.executor_config
     assert ti.operator == EmptyOperator.__name__
 
     # Test that refresh_from_task does not reset ti.max_tries
-    expected_max_tries = sdk_task.retries + 10
+    expected_max_tries = task.retries + 10
     ti.max_tries = expected_max_tries
-    ti.refresh_from_task(ser_task)
+    ti.refresh_from_task(task)
     assert ti.max_tries == expected_max_tries
 
 

airflow-core/tests/unit/models/test_taskinstance.py

@@ -84,12 +84,7 @@
     OperatorSerialization,
     _XComRef,
 )
-from airflow.task.priority_strategy import (
-    PriorityWeightStrategy,
-    _DownstreamPriorityWeightStrategy,
-    airflow_priority_weight_strategies,
-    validate_and_load_priority_weight_strategy,
-)
+from airflow.task.priority_strategy import _AbsolutePriorityWeightStrategy, _DownstreamPriorityWeightStrategy
 from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep
 from airflow.timetables.simple import NullTimetable, OnceTimetable
 from airflow.triggers.base import StartTriggerArgs
@@ -784,7 +779,6 @@ def validate_deserialized_task(
                 "inlets",
                 "outlets",
                 "task_type",
-                "weight_rule",
                 "_operator_name",
                 # Type is excluded, so don't check it
                 "_log",
@@ -818,7 +812,6 @@ def validate_deserialized_task(
                 "operator_class",
                 "partial_kwargs",
                 "expand_input",
-                "weight_rule",
             }
 
         assert serialized_task.task_type == task.task_type
@@ -853,12 +846,6 @@ def validate_deserialized_task(
         if isinstance(task.params, ParamsDict) and isinstance(serialized_task.params, ParamsDict):
             assert serialized_task.params.dump() == task.params.dump()
 
-        if isinstance(task.weight_rule, PriorityWeightStrategy):
-            assert task.weight_rule == serialized_task.weight_rule
-        else:
-            task_weight_strat = validate_and_load_priority_weight_strategy(task.weight_rule)
-            assert task_weight_strat == serialized_task.weight_rule
-
         if isinstance(task, MappedOperator):
             # MappedOperator.operator_class now stores only minimal type information
             # for memory efficiency (task_type and _operator_name).
@@ -1585,7 +1572,7 @@ def test_no_new_fields_added_to_base_operator(self):
             "ui_fgcolor": "#000",
             "wait_for_downstream": False,
             "wait_for_past_depends_before_skipping": False,
-            "weight_rule": WeightRule.DOWNSTREAM,
+            "weight_rule": _DownstreamPriorityWeightStrategy(),
             "multiple_outputs": False,
         }, """
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
@@ -3999,6 +3986,27 @@ def test_task_callback_backward_compatibility(old_callback_name, new_callback_na
     assert getattr(deserialized_task_empty, new_callback_name) is False
 
 
+def test_weight_rule_absolute_serialization_deserialization():
+    """Test that weight_rule can be serialized and deserialized correctly."""
+    from airflow.sdk import task
+
+    with DAG("test_weight_rule_dag") as dag:
+
+        @task(weight_rule=WeightRule.ABSOLUTE)
+        def test_task():
+            return "test"
+
+        test_task()
+
+    serialized_dag = DagSerialization.to_dict(dag)
+    assert serialized_dag["dag"]["tasks"][0]["__var"]["weight_rule"] == "absolute"
+
+    deserialized_dag = DagSerialization.from_dict(serialized_dag)
+
+    deserialized_task = deserialized_dag.task_dict["test_task"]
+    assert isinstance(deserialized_task.weight_rule, _AbsolutePriorityWeightStrategy)
+
+
 class TestClientDefaultsGeneration:
     """Test client defaults generation functionality."""
 
@@ -4471,50 +4479,3 @@ def test_dag_default_args_callbacks_serialization(callbacks, expected_has_flags,
 
     deserialized_dag = DagSerialization.deserialize_dag(serialized_dag_dict)
     assert deserialized_dag.dag_id == "test_default_args_callbacks"
-
-
-class RegisteredPriorityWeightStrategy(PriorityWeightStrategy):
-    def get_weight(self, ti):
-        return 99
-
-
-class TestWeightRule:
-    def test_default(self):
-        sdkop = BaseOperator(task_id="should_fail")
-        serop = OperatorSerialization.deserialize(OperatorSerialization.serialize(sdkop))
-        assert serop.weight_rule == _DownstreamPriorityWeightStrategy()
-
-    @pytest.mark.parametrize(("value", "expected"), list(airflow_priority_weight_strategies.items()))
-    def test_builtin(self, value, expected):
-        sdkop = BaseOperator(task_id="should_fail", weight_rule=value)
-        serop = OperatorSerialization.deserialize(OperatorSerialization.serialize(sdkop))
-        assert serop.weight_rule == expected()
-
-    def test_custom(self):
-        sdkop = BaseOperator(task_id="should_fail", weight_rule=RegisteredPriorityWeightStrategy())
-        with mock.patch(
-            "airflow.serialization.serialized_objects._get_registered_priority_weight_strategy",
-            return_value=RegisteredPriorityWeightStrategy,
-        ) as mock_get_registered_priority_weight_strategy:
-            serop = OperatorSerialization.deserialize(OperatorSerialization.serialize(sdkop))
-
-        assert serop.weight_rule == RegisteredPriorityWeightStrategy()
-        assert mock_get_registered_priority_weight_strategy.mock_calls == [
-            mock.call("unit.serialization.test_dag_serialization.RegisteredPriorityWeightStrategy"),
-            mock.call("unit.serialization.test_dag_serialization.RegisteredPriorityWeightStrategy"),
-            mock.call("unit.serialization.test_dag_serialization.RegisteredPriorityWeightStrategy"),
-        ]
-
-    def test_invalid(self):
-        op = BaseOperator(task_id="should_fail", weight_rule="no rule")
-        with pytest.raises(ValueError, match="Unknown priority strategy"):
-            OperatorSerialization.serialize(op)
-
-    def test_not_registered_custom(self):
-        class NotRegisteredPriorityWeightStrategy(PriorityWeightStrategy):
-            def get_weight(self, ti):
-                return 99
-
-        op = BaseOperator(task_id="empty_task", weight_rule=NotRegisteredPriorityWeightStrategy())
-        with pytest.raises(ValueError, match="Unknown priority strategy"):
-            OperatorSerialization.serialize(op)