Skip to content

Commit e9482c6

Browse files
committed
made cores to use configurable
Configurable cores to use.
1 parent 059c487 commit e9482c6

File tree

4 files changed

+17
-5
lines changed

4 files changed

+17
-5
lines changed

src/calculators/qm_xn_lib_calculator.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616

1717
class QMXNLibCalculator(QMCalculator):
18-
def __init__(self, edges_file_path:str, output_file_path:str, polygon_file_path:str=None):
18+
def __init__(self, edges_file_path:str, output_file_path:str, polygon_file_path:str=None, partition_count:int = os.cpu_count()):
1919
"""
2020
Initializes the QMXNLibCalculator class.
2121
@@ -31,6 +31,7 @@ def __init__(self, edges_file_path:str, output_file_path:str, polygon_file_path:
3131
self.default_projection = 'epsg:26910'
3232
self.output_projection = 'epsg:4326'
3333
self.precision = 1e-5
34+
self.partition_count = partition_count
3435

3536
def add_edges_from_linestring(self, graph, linestring, edge_attrs):
3637
points = list(linestring.coords)
@@ -147,7 +148,7 @@ def calculate_quality_metric(self):
147148
gdf = gdf.to_crs(self.default_projection)
148149
tile_gdf = tile_gdf.to_crs(self.default_projection)
149150
tile_gdf = tile_gdf[['geometry']]
150-
no_of_cores = os.cpu_count()
151+
no_of_cores = min(self.partition_count, os.cpu_count())
151152
df_dask = dask_geopandas.from_geopandas(tile_gdf, npartitions=no_of_cores)
152153

153154
output = df_dask.apply(self.qm_func,axis=1, meta=[

src/config.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ class Config(BaseSettings):
1414
storage_container_name: str = os.environ.get('CONTAINER_NAME', 'osw')
1515
algorithm_dictionary: dict = {"fixed":QMFixedCalculator,"ixn":QMXNLibCalculator}
1616
max_concurrent_messages: int = os.environ.get('MAX_CONCURRENT_MESSAGES', 1)
17+
partition_count:int = os.environ.get('PARTITION_COUNT', 2)
1718

1819
def get_download_folder(self) -> str:
1920
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

src/services/osw_qm_calculator_service.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,16 @@ class OswQmCalculator:
2727
2828
"""
2929

30+
def __init__(self, cores_to_use:int):
31+
"""
32+
Initializes the OswQmCalculator class.
33+
34+
Args:
35+
cores_to_use (int): The number of cores to use for calculating quality metrics.
36+
37+
"""
38+
self.cores_to_use = cores_to_use
39+
3040
def calculate_quality_metric(self, input_file, algorithm_names, output_path, ixn_file=None):
3141
"""
3242
Calculates quality metrics for input files using specified algorithms.
@@ -93,7 +103,7 @@ def get_osw_qm_calculator(self, algorithm_name:str, ixn_file:str=None, edges_fil
93103
94104
"""
95105
if algorithm_name == 'ixn':
96-
return QMXNLibCalculator(edges_file, output_file, ixn_file)
106+
return QMXNLibCalculator(edges_file, output_file, ixn_file, self.cores_to_use)
97107
else:
98108
return QMFixedCalculator(edges_file, output_file)
99109

@@ -164,7 +174,6 @@ def parse_and_calculate_quality_metric(self, input_file, algorithm_names, ixn_fi
164174

165175
if 'ixn' in algorithm_names:
166176
# get the edges file from the input
167-
168177
ixn_calculator = QMXNLibCalculator()
169178
input_json = ixn_calculator.calculate_quality_metric(input_file, ixn_file)
170179

src/services/servicebus_service.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,8 @@ def process_message(self, msg: QueueMessage):
7373
output_folder = os.path.join(download_folder,'qm')
7474
os.makedirs(output_folder,exist_ok=True)
7575
output_file_local_path = os.path.join(output_folder,'qm-output.zip')
76-
qm_calculator = OswQmCalculator()
76+
cores_to_use = self.config.partition_count
77+
qm_calculator = OswQmCalculator(cores_to_use=cores_to_use)
7778
algorithm_names = quality_request.data.algorithm.split(',')
7879
qm_calculator.calculate_quality_metric(download_path, algorithm_names,output_file_local_path,ixn_file_path)
7980
# Upload the file

0 commit comments

Comments
 (0)