From 590684ad8382573cdcabf285259f0b810f93c5e4 Mon Sep 17 00:00:00 2001 From: Linus Seelinger Date: Sun, 17 Mar 2024 21:56:34 +0100 Subject: [PATCH] Add basic Python test model for HPC --- hpc/test/testmodel.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 hpc/test/testmodel.py diff --git a/hpc/test/testmodel.py b/hpc/test/testmodel.py new file mode 100644 index 0000000..ca44d70 --- /dev/null +++ b/hpc/test/testmodel.py @@ -0,0 +1,34 @@ +import umbridge +import time +import os + +class TestModel(umbridge.Model): + + def __init__(self): + super().__init__("forward") + + def get_input_sizes(self, config): + return [1] + + def get_output_sizes(self, config): + return [1] + + def __call__(self, parameters, config): + # Sleep for number of milliseconds defined in env var + time.sleep(int(os.getenv("TEST_DELAY", 0)) / 1000) + + posterior = 2*parameters[0][0] + return [[posterior]] + + def supports_evaluate(self): + return True + + def gradient(self,out_wrt, in_wrt, parameters, sens, config): + # Sleep for number of milliseconds defined in env var + time.sleep(int(os.getenv("TEST_DELAY", 0)) / 1000) + return [2*sens[0]] + + def supports_gradient(self): + return True + +umbridge.serve_models([TestModel()], int(os.getenv("PORT", 4242)))