|
| 1 | +import json |
| 2 | +import random |
| 3 | +from datetime import timedelta |
| 4 | +from pathlib import Path |
| 5 | +from tempfile import TemporaryDirectory |
| 6 | +from typing import Annotated, Any |
| 7 | +from urllib.parse import quote |
| 8 | + |
| 9 | +from locust import HttpUser, run_single_user, task |
| 10 | +from pydantic import BaseModel, Field |
| 11 | +from pydantic_settings import BaseSettings, SettingsConfigDict |
| 12 | +from requests.auth import HTTPBasicAuth |
| 13 | +from tenacity import retry, retry_if_exception_type, stop_after_delay, wait_exponential |
| 14 | +from urllib3 import PoolManager, Retry |
| 15 | + |
| 16 | + |
| 17 | +class UserSettings(BaseSettings): |
| 18 | + model_config = SettingsConfigDict(extra="ignore") |
| 19 | + OSPARC_API_KEY: Annotated[str, Field()] # required, no default |
| 20 | + OSPARC_API_SECRET: Annotated[str, Field()] # required, no default |
| 21 | + |
| 22 | + |
| 23 | +_SOLVER_KEY = "simcore/services/comp/osparc-python-runner" |
| 24 | +_SOLVER_VERSION = "1.2.0" |
| 25 | + |
| 26 | +_PYTHON_SCRIPT = """ |
| 27 | +import numpy as np |
| 28 | +import pathlib as pl |
| 29 | +import json |
| 30 | +import os |
| 31 | +
|
| 32 | +def main(): |
| 33 | +
|
| 34 | + input_json = pl.Path(os.environ["INPUT_FOLDER"]) / "function_inputs.json" |
| 35 | + object = json.loads(input_json.read_text()) |
| 36 | + x = object["x"] |
| 37 | + y = object["y"] |
| 38 | +
|
| 39 | + return np.sinc(x) * np.sinc(y) |
| 40 | +
|
| 41 | +
|
| 42 | +if __name__ == "__main__": |
| 43 | + main() |
| 44 | +
|
| 45 | +""" |
| 46 | + |
| 47 | + |
| 48 | +class Schema(BaseModel): |
| 49 | + schema_content: dict = {} |
| 50 | + schema_class: str = "application/schema+json" |
| 51 | + |
| 52 | + |
| 53 | +class Function(BaseModel): |
| 54 | + function_class: str = "SOLVER" |
| 55 | + title: str |
| 56 | + description: str |
| 57 | + input_schema: Annotated[Schema, Field()] = Schema() |
| 58 | + output_schema: Annotated[Schema, Field()] = Schema() |
| 59 | + default_inputs: Annotated[dict[str, Any], Field()] = dict() |
| 60 | + solver_key: Annotated[str, Field()] = _SOLVER_KEY |
| 61 | + solver_version: Annotated[str, Field()] = _SOLVER_VERSION |
| 62 | + |
| 63 | + |
| 64 | +class MetaModelingUser(HttpUser): |
| 65 | + def __init__(self, *args, **kwargs): |
| 66 | + self._user_settings = UserSettings() |
| 67 | + self._auth = HTTPBasicAuth( |
| 68 | + username=self._user_settings.OSPARC_API_KEY, |
| 69 | + password=self._user_settings.OSPARC_API_SECRET, |
| 70 | + ) |
| 71 | + retry_strategy = Retry( |
| 72 | + total=4, |
| 73 | + backoff_factor=4.0, |
| 74 | + status_forcelist={429, 503, 504}, |
| 75 | + allowed_methods={ |
| 76 | + "DELETE", |
| 77 | + "GET", |
| 78 | + "HEAD", |
| 79 | + "OPTIONS", |
| 80 | + "PUT", |
| 81 | + "TRACE", |
| 82 | + "POST", |
| 83 | + "PATCH", |
| 84 | + "CONNECT", |
| 85 | + }, |
| 86 | + respect_retry_after_header=True, |
| 87 | + raise_on_status=True, |
| 88 | + ) |
| 89 | + self.pool_manager = PoolManager(retries=retry_strategy) |
| 90 | + |
| 91 | + self._function_uid = None |
| 92 | + self._input_json = None |
| 93 | + self._script = None |
| 94 | + self._run_uid = None |
| 95 | + self._solver_job_uid = None |
| 96 | + |
| 97 | + super().__init__(*args, **kwargs) |
| 98 | + |
| 99 | + def on_stop(self) -> None: |
| 100 | + if self._script is not None: |
| 101 | + self.client.delete( |
| 102 | + f"/v0/files/{self._script.get('id')}", |
| 103 | + name="/v0/files/[file_id]", |
| 104 | + auth=self._auth, |
| 105 | + ) |
| 106 | + if self._input_json is not None: |
| 107 | + self.client.delete( |
| 108 | + f"/v0/files/{self._input_json.get('id')}", |
| 109 | + name="/v0/files/[file_id]", |
| 110 | + auth=self._auth, |
| 111 | + ) |
| 112 | + if self._function_uid is not None: |
| 113 | + self.client.delete( |
| 114 | + f"/v0/functions/{self._function_uid}", |
| 115 | + name="/v0/functions/[function_uid]", |
| 116 | + auth=self._auth, |
| 117 | + ) |
| 118 | + if self._run_uid is not None: |
| 119 | + self.client.delete( |
| 120 | + f"/v0/function_jobs/{self._run_uid}", |
| 121 | + name="/v0/function_jobs/[function_run_uid]", |
| 122 | + auth=self._auth, |
| 123 | + ) |
| 124 | + |
| 125 | + @task |
| 126 | + def run_function(self): |
| 127 | + with TemporaryDirectory() as tmpdir_str, Path(tmpdir_str) as tmpdir: |
| 128 | + script = tmpdir / "script.py" |
| 129 | + script.write_text(_PYTHON_SCRIPT) |
| 130 | + self._script = self.upload_file(script) |
| 131 | + |
| 132 | + inputs = {"x": random.uniform(-10, 10), "y": random.uniform(-10, 10)} |
| 133 | + input_json = tmpdir / "function_inputs.json" |
| 134 | + input_json.write_text(json.dumps(inputs)) |
| 135 | + self._input_json = self.upload_file(input_json) |
| 136 | + |
| 137 | + _function = Function( |
| 138 | + title="Test function", |
| 139 | + description="Test function", |
| 140 | + default_inputs={"input_1": self._script}, |
| 141 | + ) |
| 142 | + response = self.client.post( |
| 143 | + "/v0/functions", json=_function.model_dump(), auth=self._auth |
| 144 | + ) |
| 145 | + response.raise_for_status() |
| 146 | + self._function_uid = response.json().get("uid") |
| 147 | + assert self._function_uid is not None |
| 148 | + |
| 149 | + response = self.client.post( |
| 150 | + f"/v0/functions/{self._function_uid}:run", |
| 151 | + json={"input_2": self._input_json}, |
| 152 | + auth=self._auth, |
| 153 | + name="/v0/functions/[function_uid]:run", |
| 154 | + ) |
| 155 | + response.raise_for_status() |
| 156 | + self._run_uid = response.json().get("uid") |
| 157 | + assert self._run_uid is not None |
| 158 | + self._solver_job_uid = response.json().get("solver_job_id") |
| 159 | + assert self._solver_job_uid is not None |
| 160 | + |
| 161 | + self.wait_until_done() |
| 162 | + |
| 163 | + response = self.client.get( |
| 164 | + f"/v0/solvers/{quote(_SOLVER_KEY, safe='')}/releases/{_SOLVER_VERSION}/jobs/{self._solver_job_uid}/outputs", |
| 165 | + auth=self._auth, |
| 166 | + name="/v0/solvers/[solver_key]/releases/[solver_version]/jobs/[solver_job_id]/outputs", |
| 167 | + ) |
| 168 | + response.raise_for_status() |
| 169 | + |
| 170 | + @retry( |
| 171 | + stop=stop_after_delay(timedelta(minutes=10)), |
| 172 | + wait=wait_exponential(multiplier=1, min=1, max=10), |
| 173 | + retry=retry_if_exception_type(AssertionError), |
| 174 | + reraise=False, |
| 175 | + ) |
| 176 | + def wait_until_done(self): |
| 177 | + response = self.client.get( |
| 178 | + f"/v0/function_jobs/{self._run_uid}/status", |
| 179 | + auth=self._auth, |
| 180 | + name="/v0/function_jobs/[function_run_uid]/status", |
| 181 | + ) |
| 182 | + response.raise_for_status() |
| 183 | + status = response.json().get("status") |
| 184 | + assert status in ["SUCCESS", "FAILED"] |
| 185 | + |
| 186 | + def upload_file(self, file: Path) -> dict[str, str]: |
| 187 | + assert file.is_file() |
| 188 | + with file.open(mode="rb") as f: |
| 189 | + files = {"file": f} |
| 190 | + response = self.client.put( |
| 191 | + "/v0/files/content", files=files, auth=self._auth |
| 192 | + ) |
| 193 | + response.raise_for_status() |
| 194 | + assert response.json().get("id") is not None |
| 195 | + return response.json() |
| 196 | + |
| 197 | + |
| 198 | +if __name__ == "__main__": |
| 199 | + run_single_user(MetaModelingUser) |
0 commit comments