diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index 3eb900f31..1f4c5cc15 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -43,7 +43,6 @@ import smartsim._core.utils.helpers as _helpers from smartsim._core.schemas.dragonRequests import request_serializer from smartsim._core.schemas.dragonResponses import response_serializer -from smartsim._core.schemas.types import NonEmptyStr from ....error import LauncherError from ....log import get_logger @@ -371,7 +370,7 @@ def _get_managed_step_update(self, step_ids: t.List[str]) -> t.List[StepInfo]: msg += response.error_message raise LauncherError(msg) - status, ret_codes = response.statuses[NonEmptyStr(step_id)] + status, ret_codes = response.statuses[step_id] if ret_codes: grp_ret_code = min(ret_codes) if any(ret_codes): diff --git a/smartsim/_core/schemas/dragonRequests.py b/smartsim/_core/schemas/dragonRequests.py index 7ddfd28b5..01fa7285f 100644 --- a/smartsim/_core/schemas/dragonRequests.py +++ b/smartsim/_core/schemas/dragonRequests.py @@ -26,10 +26,9 @@ import typing as t -from pydantic import BaseModel, PositiveInt +from pydantic import BaseModel, PositiveInt, Field import smartsim._core.schemas.utils as _utils -from smartsim._core.schemas.types import NonEmptyStr # Black and Pylint disagree about where to put the `...` # pylint: disable=multiple-statements @@ -42,16 +41,16 @@ class DragonRequest(BaseModel): ... class DragonRunRequestView(DragonRequest): - exe: NonEmptyStr - exe_args: t.List[NonEmptyStr] = [] - path: NonEmptyStr + exe: t.Annotated[str, Field(min_length=1)] + exe_args: t.List[t.Annotated[str, Field(min_length=1)]] = [] + path: t.Annotated[str, Field(min_length=1)] nodes: PositiveInt = 1 tasks: PositiveInt = 1 tasks_per_node: PositiveInt = 1 - output_file: t.Optional[NonEmptyStr] = None - error_file: t.Optional[NonEmptyStr] = None + output_file: t.Optional[t.Annotated[str, Field(min_length=1)]] = None + error_file: t.Optional[t.Annotated[str, Field(min_length=1)]] = None env: t.Dict[str, t.Optional[str]] = {} - name: t.Optional[NonEmptyStr] + name: t.Optional[t.Annotated[str, Field(min_length=1)]] = None pmi_enabled: bool = True @@ -65,12 +64,12 @@ def __str__(self) -> str: @request_serializer.register("update_status") class DragonUpdateStatusRequest(DragonRequest): - step_ids: t.List[NonEmptyStr] + step_ids: t.List[t.Annotated[str, Field(min_length=1)]] @request_serializer.register("stop") class DragonStopRequest(DragonRequest): - step_id: NonEmptyStr + step_id: t.Annotated[str, Field(min_length=1)] @request_serializer.register("handshake") @@ -79,7 +78,7 @@ class DragonHandshakeRequest(DragonRequest): ... @request_serializer.register("bootstrap") class DragonBootstrapRequest(DragonRequest): - address: NonEmptyStr + address: t.Annotated[str, Field(min_length=1)] @request_serializer.register("shutdown") diff --git a/smartsim/_core/schemas/dragonResponses.py b/smartsim/_core/schemas/dragonResponses.py index 09562c80a..3ef7532b1 100644 --- a/smartsim/_core/schemas/dragonResponses.py +++ b/smartsim/_core/schemas/dragonResponses.py @@ -26,10 +26,9 @@ import typing as t -from pydantic import BaseModel +from pydantic import BaseModel, Field import smartsim._core.schemas.utils as _utils -from smartsim._core.schemas.types import NonEmptyStr # Black and Pylint disagree about where to put the `...` # pylint: disable=multiple-statements @@ -44,13 +43,16 @@ class DragonResponse(BaseModel): @response_serializer.register("run") class DragonRunResponse(DragonResponse): - step_id: NonEmptyStr + step_id: t.Annotated[str, Field(min_length=1)] @response_serializer.register("status_update") class DragonUpdateStatusResponse(DragonResponse): # status is a dict: {step_id: (is_alive, returncode)} - statuses: t.Mapping[NonEmptyStr, t.Tuple[NonEmptyStr, t.Optional[t.List[int]]]] = {} + statuses: t.Mapping[ + t.Annotated[str, Field(min_length=1)], + t.Tuple[t.Annotated[str, Field(min_length=1)], t.Optional[t.List[int]]], + ] = {} @response_serializer.register("stop") diff --git a/smartsim/_core/schemas/types.py b/smartsim/_core/schemas/types.py deleted file mode 100644 index 7136230e8..000000000 --- a/smartsim/_core/schemas/types.py +++ /dev/null @@ -1,31 +0,0 @@ -# BSD 2-Clause License -# -# Copyright (c) 2021-2023, Hewlett Packard Enterprise -# All rights reserved. -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -import pydantic - - -class NonEmptyStr(pydantic.ConstrainedStr): - min_length = 1