-
Notifications
You must be signed in to change notification settings - Fork 22
/
setup.py
140 lines (120 loc) · 4.64 KB
/
setup.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
# Copyright 2022 Google LLC
#
# Licensed 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
#
# https://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.
"""Setup arco_era5.
This uses Apache Beam's recommended way to install non-python dependencies. This lets us install ecCodes for cfgrib on
Dataflow workers.
Please see the following documentation and examples for more:
- https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/#nonpython
- https://github.com/apache/beam/blob/master/sdks/python/apache_beam/examples/complete/juliaset/setup.py
"""
import platform
import subprocess
from distutils.command.build import build as _build # type: ignore
from setuptools import setup, find_packages, Command
"""Install the ecCodes binary from ECMWF."""
CUSTOM_COMMANDS = [
cmd.split() for cmd in [
'apt-get update',
'apt-get --assume-yes install libeccodes-dev'
]
]
# This class handles the pip install mechanism.
class build(_build): # pylint: disable=invalid-name
"""A build command class that will be invoked during package install.
The package built using the current setup.py will be staged and later
installed in the worker using `pip install package`. This class will be
instantiated during install for this specific scenario and will trigger
running the custom commands specified.
"""
sub_commands = _build.sub_commands + [('CustomCommands', None)]
class CustomCommands(Command):
"""A setuptools Command class able to run arbitrary commands.
This class allows running arbitrary commands during the package installation
process. It can be used to execute custom build or setup steps required for
the package.
Attributes:
CUSTOM_COMMANDS (list): A list of lists, where each inner list represents
a command and its arguments.
Methods:
initialize_options(): Initialize custom command options.
finalize_options(): Finalize custom command options.
RunCustomCommand(command_list): Run a custom command and print its output.
run(): Run the custom commands specified in CUSTOM_COMMANDS.
Example:
To add custom commands, define them in the CUSTOM_COMMANDS list and
include 'CustomCommands' as a sub-command in the setup.py script.
```python
CUSTOM_COMMANDS = [['command1', 'arg1', 'arg2'], ['command2', 'arg1']]
setup(
# ...
cmdclass={
'build': build,
'CustomCommands': CustomCommands,
}
)
```
"""
def initialize_options(self):
pass
def finalize_options(self):
pass
def RunCustomCommand(self, command_list):
print('Running command: %s' % command_list)
p = subprocess.Popen(
command_list,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT)
# Can use communicate(input='y\n'.encode()) if the command run requires
# some confirmation.
stdout_data, _ = p.communicate()
print('Command output: %s' % stdout_data)
if p.returncode != 0:
raise RuntimeError(
'Command %s failed: exit code: %s' % (command_list, p.returncode))
def run(self):
# Only run on Linux, which is Dataflow's environment.
if 'Linux' not in platform.system():
return
for command in CUSTOM_COMMANDS:
self.RunCustomCommand(command)
setup(
name='arco_era5',
package_dir={'': 'src'},
packages=find_packages('src'),
author_email='anthromet-core+era5@google.com',
description="Analysis-Ready & Cloud-Optimized ERA5.",
version='0.1.0',
platforms=['darwin', 'linux'],
python_requires='>3.8, <3.11',
install_requires=[
"google-cloud-secret-manager==2.0.0",
'apache_beam[gcp]==2.40.0',
'numcodecs==0.11.0',
'pangeo-forge-recipes==0.9.1',
'pandas',
'gcsfs',
'cfgrib',
'immutabledict',
'xarray-beam',
'scipy',
'fsspec==2023.4.0'
],
tests_require=['pytest'],
cmdclass={
# Command class instantiated and run during pip install scenarios.
'build': build,
'CustomCommands': CustomCommands,
}
)