1#
2# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18
19from __future__ import annotations
20
21import multiprocessing
22import multiprocessing.context
23import typing
24
25from airflow.configuration import conf
26
27if typing.TYPE_CHECKING:
28 from airflow.models.operator import Operator
29 from airflow.utils.context import Context
30
31
32class MultiprocessingStartMethodMixin:
33 """Convenience class to add support for different types of multiprocessing."""
34
35 def _get_multiprocessing_start_method(self) -> str:
36 """
37 Determine method of creating new processes.
38
39 Checks if the mp_start_method is set in configs, else, it uses the OS default.
40 """
41 if conf.has_option("core", "mp_start_method"):
42 return conf.get_mandatory_value("core", "mp_start_method")
43
44 method = multiprocessing.get_start_method()
45 if not method:
46 raise ValueError("Failed to determine start method")
47 return method
48
49 def _get_multiprocessing_context(self) -> multiprocessing.context.DefaultContext:
50 mp_start_method = self._get_multiprocessing_start_method()
51 return multiprocessing.get_context(mp_start_method) # type: ignore
52
53
54class ResolveMixin:
55 """A runtime-resolved value."""
56
57 def iter_references(self) -> typing.Iterable[tuple[Operator, str]]:
58 """Find underlying XCom references this contains.
59
60 This is used by the DAG parser to recursively find task dependencies.
61
62 :meta private:
63 """
64 raise NotImplementedError
65
66 def resolve(self, context: Context) -> typing.Any:
67 """Resolve this value for runtime.
68
69 :meta private:
70 """
71 raise NotImplementedError