diff --git a/gempy_engine/API/model/model_api.py b/gempy_engine/API/model/model_api.py index 0404d5e8..ffc8b83e 100644 --- a/gempy_engine/API/model/model_api.py +++ b/gempy_engine/API/model/model_api.py @@ -18,6 +18,7 @@ from ...core.data.solutions import Solutions from ...core.utils import gempy_profiler_decorator from ...core.exceptions import GemPyEngineInputError +from ...core.data.options.temp_interpolation_values import TempInterpolationValues from ...modules.geophysics.fw_gravity import compute_gravity from ...modules.geophysics.fw_magnetic import compute_tmi from ...modules.weights_cache.weights_cache_interface import WeightCache @@ -26,6 +27,10 @@ @gempy_profiler_decorator def compute_model(interpolation_input: InterpolationInput, options: InterpolationOptions, data_descriptor: InputDataDescriptor, *, geophysics_input: Optional[GeophysicsInput] = None) -> Solutions: + # Octree progress and cache timestamps belong to this computation. Keeping + # them on a shared options object lets concurrent requests change each + # other's active octree level. + options = options.model_copy(update={"temp_interpolation_values": TempInterpolationValues()}) try: WeightCache.initialize_cache_dir() options.temp_interpolation_values.start_computation_ts = int(time.time()) diff --git a/gempy_engine/API/server/main_server_pro.py b/gempy_engine/API/server/main_server_pro.py index daaf940b..e1070baf 100644 --- a/gempy_engine/API/server/main_server_pro.py +++ b/gempy_engine/API/server/main_server_pro.py @@ -58,10 +58,15 @@ def compute_gempy_model(gempy_input: GemPyInput) -> Response: logger=logger ) + # Options contain mutable nested configuration, so each request needs its + # own copy. The request schema exposes the desired octree depth on the grid. + options = default_interpolation_options.model_copy(deep=True) + options.evaluation_options.number_octree_levels = gempy_input.interpolation_input.grid.octree_levels + # Compute model solutions = _compute_model( interpolation_input=interpolation_input, - options=default_interpolation_options, + options=options, structure=input_data_descriptor ) logger.info("Finished computing model") @@ -107,4 +112,4 @@ def _compute_model( v_just_points=interpolation_input.surface_points.sp_coords ) - return solutions \ No newline at end of file + return solutions diff --git a/gempy_engine/core/backend_tensor.py b/gempy_engine/core/backend_tensor.py index 458afafd..47f7b555 100644 --- a/gempy_engine/core/backend_tensor.py +++ b/gempy_engine/core/backend_tensor.py @@ -195,8 +195,12 @@ def _sum(tensor, axis=None, dtype=None, keepdims=False): def _repeat(tensor, n_repeats, axis=None): if not isinstance(tensor, torch.Tensor): tensor = torch.as_tensor(tensor, device=cls.device) + elif tensor.device != cls.device: + tensor = tensor.to(cls.device) if not isinstance(n_repeats, torch.Tensor): n_repeats = torch.as_tensor(n_repeats, device=cls.device) + elif n_repeats.device != cls.device: + n_repeats = n_repeats.to(cls.device) return _true_torch_repeat_interleave(tensor, n_repeats, dim=axis) def _array(array_like, dtype=None): diff --git a/gempy_engine/modules/data_preprocess/_input_preparation.py b/gempy_engine/modules/data_preprocess/_input_preparation.py index aa82042c..fe1ae452 100644 --- a/gempy_engine/modules/data_preprocess/_input_preparation.py +++ b/gempy_engine/modules/data_preprocess/_input_preparation.py @@ -38,7 +38,7 @@ def surface_points_preprocess(sp_input: SurfacePoints, tensors_structure: Tensor ref_points_repeated = b.t.repeat(ref_points, number_repetitions, 0) # ref_points shape: (1, 3) ref_nugget_repeated = b.t.repeat(ref_nugget, number_repetitions, 0) surface_ids = b.t.repeat( - b.t.arange(tensors_structure.n_surfaces, dtype=rest_nugget.dtype), + b.t.array(list(range(tensors_structure.n_surfaces)), dtype=rest_nugget.dtype), number_repetitions, 0, ) diff --git a/gempy_engine/modules/evaluator/symbolic_evaluator.py b/gempy_engine/modules/evaluator/symbolic_evaluator.py index 675f5718..835996a0 100644 --- a/gempy_engine/modules/evaluator/symbolic_evaluator.py +++ b/gempy_engine/modules/evaluator/symbolic_evaluator.py @@ -104,6 +104,21 @@ def _build_block_sparse_ranges(M_sizes: list[int], N_sizes: list[int]): return numpy_ranges +def _validate_stacked_dimensions(eval_kernel, weights, M_sizes: list[int], N_sizes: list[int]) -> None: + """Fail before PyKeOps when block ranges do not describe the lazy kernel.""" + expected_i = sum(N_sizes) + expected_j = sum(M_sizes) + kernel_shape = tuple(eval_kernel.shape) + weights_size = weights.shape[0] + + if kernel_shape[0] != expected_i or kernel_shape[1] != expected_j or weights_size != expected_i: + raise ValueError( + "Inconsistent stacked PyKeOps dimensions: " + f"kernel={kernel_shape}, weights={weights_size}, " + f"range dimensions=({expected_i}, {expected_j})." + ) + + def symbolic_evaluator_optimized_stacked( eval_inputs: list[EvaluatorInput], weights_list: list[np.ndarray], @@ -194,6 +209,7 @@ def _run_prep(args): # Concatenate weights all_weights = BackendTensor.t.concatenate(weights_list, axis=0) all_weights = BackendTensor.t.tile(all_weights, tile_factor) + _validate_stacked_dimensions(eval_kernel, all_weights, M_sizes, N_sizes) if BackendTensor.engine_backend == gempy_engine.config.AvailableBackends.numpy: from pykeops.numpy import LazyTensor @@ -234,13 +250,11 @@ def _run_prep(args): except TypeError: raise ValueError("Failed to compute symbolic evaluation with PyKeOps. Ensure that all_weights and eval_kernel are compatible for lazy tensor operations.") - # For torch - # all_results_concat = all_results_concat.to("cpu") - all_results_split = BackendTensor.t.split(all_results_concat, M_sizes) - - # For numpy - # split_indices = np.cumsum(M_sizes)[:-1] - # all_results_split = np.split(all_results_concat, split_indices) + if BackendTensor.engine_backend == gempy_engine.config.AvailableBackends.numpy: + split_indices = np.cumsum(M_sizes)[:-1] + all_results_split = np.split(all_results_concat, split_indices) + else: + all_results_split = BackendTensor.t.split(all_results_concat, M_sizes) original_n_fields = len(eval_inputs) scalar_fields = all_results_split[:original_n_fields] diff --git a/tests/test_common/test_api/test_concurrent_options.py b/tests/test_common/test_api/test_concurrent_options.py new file mode 100644 index 00000000..61a31eab --- /dev/null +++ b/tests/test_common/test_api/test_concurrent_options.py @@ -0,0 +1,60 @@ +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier + +import pytest + +from gempy_engine.API.model import model_api +from gempy_engine.core.data import InterpolationOptions +from gempy_engine.core.data.options.temp_interpolation_values import TempInterpolationValues +from gempy_engine.modules.evaluator.symbolic_evaluator import _validate_stacked_dimensions + + +def test_computations_do_not_share_volatile_option_state(monkeypatch): + options = InterpolationOptions.from_args( + range=1.0, + c_o=1.0, + number_octree_levels=2, + mesh_extraction=False, + ) + barrier = Barrier(2) + + monkeypatch.setattr(model_api.WeightCache, "initialize_cache_dir", lambda: None) + monkeypatch.setattr(model_api.WeightCache, "clear_cache", lambda: None) + monkeypatch.setattr(model_api.BackendTensor, "clear_gpu_memory", lambda: None) + monkeypatch.setattr(model_api, "_check_input_validity", lambda *_: None) + + def observe_options(interpolation_input, options, data_descriptor): + level = interpolation_input + options.temp_interpolation_values.current_octree_level = level + barrier.wait() + return options.temp_interpolation_values.current_octree_level + + monkeypatch.setattr(model_api, "interpolate_n_octree_levels", observe_options) + + class FakeSolutions: + def __init__(self, octrees_output, **_): + self.observed_level = octrees_output + + monkeypatch.setattr(model_api, "Solutions", FakeSolutions) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(model_api.compute_model, level, options, None) + for level in (1, 3) + ] + + assert [future.result().observed_level for future in futures] == [1, 3] + assert options.temp_interpolation_values == TempInterpolationValues() + + +def test_stacked_dimension_validation_reports_range_mismatch(): + class FakeKernel: + shape = (4, 8, 1) + + class FakeWeights: + shape = (4,) + + with pytest.raises(ValueError, match="range dimensions=\\(4, 16\\)") as error: + _validate_stacked_dimensions(FakeKernel(), FakeWeights(), M_sizes=[8, 8], N_sizes=[2, 2]) + + assert "kernel=(4, 8, 1)" in str(error.value) diff --git a/tests/test_common/test_api/test_stack_options_override.py b/tests/test_common/test_api/test_stack_options_override.py index fa2a562d..56057c08 100644 --- a/tests/test_common/test_api/test_stack_options_override.py +++ b/tests/test_common/test_api/test_stack_options_override.py @@ -86,8 +86,10 @@ def test_stack_options_override_serial(override_setup): @pytest.mark.skipif(not PYKEOPS_AVAILABLE, reason="pykeops not installed") -def test_stack_options_override_flat(override_setup): +@pytest.mark.parametrize("number_octree_levels", [1, 2]) +def test_stack_options_override_flat(override_setup, number_octree_levels): ii, global_options, ts = override_setup + global_options.evaluation_options.number_octree_levels = number_octree_levels custom_options = InterpolationOptions.from_args( range=5.0, diff --git a/tests/test_common/test_modules/test_kernel_constructor/test_nugget_propagation.py b/tests/test_common/test_modules/test_kernel_constructor/test_nugget_propagation.py index d3ac5b00..173e0b98 100644 --- a/tests/test_common/test_modules/test_kernel_constructor/test_nugget_propagation.py +++ b/tests/test_common/test_modules/test_kernel_constructor/test_nugget_propagation.py @@ -125,6 +125,36 @@ def test_surface_preprocessing_preserves_nugget_components_and_surface_ids(simpl ) +def test_surface_preprocessing_uses_backend_device_when_torch_default_differs(simple_model_2): + if BackendTensor.engine_backend is not AvailableBackends.PYTORCH or not BackendTensor.use_gpu: + pytest.skip("PyTorch GPU-only device regression test") + import torch + + previous_default_device = torch.get_default_device() + try: + torch.set_default_device("cpu") + surface_points, _, _, descriptor = deepcopy(simple_model_2) + + internal = surface_points_preprocess(surface_points, descriptor.tensors_structure) + + assert internal.surface_ids.device.type == BackendTensor.device.type + finally: + torch.set_default_device(previous_default_device) + + +def test_pytorch_repeat_moves_existing_tensors_to_backend_device(): + if BackendTensor.engine_backend is not AvailableBackends.PYTORCH or not BackendTensor.use_gpu: + pytest.skip("PyTorch GPU-only device regression test") + import torch + + values = torch.tensor([0.0], device="cpu") + repeats = torch.tensor([1], device=BackendTensor.device) + + result = BackendTensor.t.repeat(values, repeats, 0) + + assert result.device.type == BackendTensor.device.type + + def test_all_nuggets_receive_pytorch_gradients(simple_model_2): if BackendTensor.engine_backend is not AvailableBackends.PYTORCH: pytest.skip("PyTorch-only autograd test")