libinephany 1.1.2__py3-none-any.whl → 1.1.3__py3-none-any.whl

libinephany/pydantic_models/schemas/inner_task_profile.py

@@ -7,8 +7,15 @@
 import math
 from typing import Any, Callable
 
-from loguru import logger
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel
+
+# ======================================================================================================================
+#
+# CONSTANTS
+#
+# ======================================================================================================================
+
+VRAM_USAGES_KEY = "vram_usages"
 
 # ======================================================================================================================
 #
@@ -20,26 +27,40 @@ from pydantic import BaseModel, field_validator
 class InnerTaskProfile(BaseModel):
 
     inner_task_name: str
+
     number_of_agents: int
     number_of_layers: int
+    number_of_parameters: int
+
     observation_space_sizes: dict[str, int]
     action_space_sizes: dict[str, int]
-    number_of_parameters: int
-    vram_usage: float
-    idle_vram_usage: float
-    hparam_overrides: dict[str, dict[str, Any]] | None = None
 
-    @field_validator("vram_usage", "idle_vram_usage", mode="before")
-    def replace_non_with_nan(cls, value: None | float) -> float:
+    vram_usages: dict[int, tuple[float, float]] | None = None
+    expected_vram_usage: float | None = None
+    expected_idle_vram_usage: float | None = None
+    max_batch_size_override: int | None = None
+
+    @property
+    def vram_usage(self) -> float:
+        """
+        :return: VRAM usage at the max batch size.
+        """
+
+        if self.expected_vram_usage is None:
+            return float("nan")
+
+        return self.expected_vram_usage
+
+    @property
+    def idle_vram_usage(self) -> float:
         """
-        :param value: Value to replace with NaN if it is None.
-        :return: Either the given float value or NaN.
+        :return: Idle VRAM usage at the max batch size.
         """
 
-        if value is None:
+        if self.expected_idle_vram_usage is None:
             return float("nan")
 
-        return value
+        return self.expected_idle_vram_usage
 
     @property
     def failed_to_profile(self) -> bool:
@@ -49,19 +70,18 @@ class InnerTaskProfile(BaseModel):
 
         return math.isnan(self.vram_usage)
 
-    def model_dump_json(self, **kwargs) -> str:
+    def model_dump(self, **kwargs) -> dict[str, Any]:
         """
         :param kwargs: Standard Pydantic model dump kwargs.
         :return: Dump result of the superclass' method.
         """
 
-        logger.debug(
-            f"Inner task {self.inner_task_name} consumed {self.vram_usage:.3f} MB of VRAM while training and "
-            f"{self.idle_vram_usage:.3f} MB of VRAM while idle. It has {self.number_of_agents} agents across "
-            f"{self.number_of_layers} inner model layers."
-        )
+        super_dump = super().model_dump(**kwargs)
 
-        return super().model_dump_json(**kwargs)
+        if self.vram_usages is not None:
+            super_dump[VRAM_USAGES_KEY] = {k: list(v) for k, v in self.vram_usages.items()}
+
+        return super_dump
 
 
 class InnerTaskProfiles(BaseModel):
@@ -235,42 +255,6 @@ class InnerTaskProfiles(BaseModel):
 
         return inner_task_name in self.profiles
 
-    def add_profile(
-        self,
-        inner_task_name: str,
-        number_of_agents: int,
-        number_of_layers: int,
-        observation_space_sizes: dict[str, int],
-        action_space_sizes: dict[str, int],
-        number_of_parameters: int,
-        vram_usage: float,
-        idle_vram_usage: float,
-        hparam_overrides: dict[str, dict[str, Any]] | None = None,
-    ) -> None:
-        """
-        :param inner_task_name: Name of the inner task to add a profile for.
-        :param number_of_agents: Number of agents active in the inner task's environment.
-        :param number_of_layers: Number of layers in the inner model.
-        :param observation_space_sizes: Dictionary mapping agent IDs to their observation space sizes.
-        :param action_space_sizes: Dictionary mapping agent IDs to their action space sizes.
-        :param vram_usage: VRAM required to perform the inner task. Can be NaN if an OOM was encountered.
-        :param idle_vram_usage: VRAM required for the inner task to sit loaded but not actively being trained. Can be
-        NaN if an OOM was encountered.
-        :param hparam_overrides: Hyperparameter overrides for the inner task.
-        """
-
-        self.profiles[inner_task_name] = InnerTaskProfile(
-            inner_task_name=inner_task_name,
-            number_of_agents=number_of_agents,
-            number_of_layers=number_of_layers,
-            observation_space_sizes=observation_space_sizes,
-            action_space_sizes=action_space_sizes,
-            number_of_parameters=number_of_parameters,
-            vram_usage=vram_usage,
-            idle_vram_usage=idle_vram_usage,
-            hparam_overrides=hparam_overrides,
-        )
-
     def validate_task_profiles(self, policy_mapping_function: Callable[[str, Any, Any], str]) -> None:
         """
         :param policy_mapping_function: Function which maps agent IDs to policy IDs.

{libinephany-1.1.2.dist-info → libinephany-1.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libinephany
-Version: 1.1.2
+Version: 1.1.3
 Summary: Inephany library containing code commonly used by multiple subpackages.
 Author-email: Inephany <info@inephany.com>
 License: Apache 2.0

{libinephany-1.1.2.dist-info → libinephany-1.1.3.dist-info}/RECORD

@@ -28,7 +28,7 @@ libinephany/pydantic_models/configs/observer_config.py,sha256=v_ChzaVXC_rlZ7eDZP
 libinephany/pydantic_models/configs/outer_model_config.py,sha256=GQ0QBSC2Xht8x8X_TEMfYM2GF_x1kErLuFrA_H6Jhs0,1209
 libinephany/pydantic_models/schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 libinephany/pydantic_models/schemas/agent_info.py,sha256=me5gDxvZjP9TNK588mpUvxiiJrPDqy3Z7ZHRzryAYTs,2628
-libinephany/pydantic_models/schemas/inner_task_profile.py,sha256=2ix497uUCHUJzdfXtJRT6MDDn8OFQZ8zpZMBZUlvIHc,11976
+libinephany/pydantic_models/schemas/inner_task_profile.py,sha256=1Q3cDyyW01NOgalmAWGLc-AaLoLum9nBtcpLsfxo_pw,10628
 libinephany/pydantic_models/schemas/observation_models.py,sha256=MLhxqDet9Yol1D5mkQGQsQT23sm37AStRLnPc4sgcZc,2110
 libinephany/pydantic_models/schemas/request_schemas.py,sha256=VED8eAUvBofxeAx9gWU8DyCZOTVD3QsHRq-TO7kyOqk,1260
 libinephany/pydantic_models/schemas/response_schemas.py,sha256=SKFuasdjX5aH_I0vT3SwnpwhyMf9cNPB1ZpDeAGgoO8,2158
@@ -57,8 +57,8 @@ libinephany/utils/typing.py,sha256=rGbaPO3MaUndsWiC_wHzReD_TOLYqb43i01pKN-j7Xs,6
 libinephany/web_apps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 libinephany/web_apps/error_logger.py,sha256=QpspO726Uoyyr6lBEEb3Q9XqhVOXUM4AaYE7vbnk31c,18153
 libinephany/web_apps/web_app_utils.py,sha256=qiq_lasPipgN1RgRudPJc342kYci8O_4RqppxmIX8NY,4095
-libinephany-1.1.2.dist-info/licenses/LICENSE,sha256=pogfDoMBP07ehIOvWymuWIar8pg2YLUhqOHsJQU3wdc,9250
-libinephany-1.1.2.dist-info/METADATA,sha256=4wogF1PgydPU-kNDC5ijjAMbTM7TT_bRRPODhHsNe88,8389
-libinephany-1.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-libinephany-1.1.2.dist-info/top_level.txt,sha256=bYAOXQdJgIoLkO2Ui0kxe7pSYegS_e38u0dMscd7COQ,12
-libinephany-1.1.2.dist-info/RECORD,,
+libinephany-1.1.3.dist-info/licenses/LICENSE,sha256=pogfDoMBP07ehIOvWymuWIar8pg2YLUhqOHsJQU3wdc,9250
+libinephany-1.1.3.dist-info/METADATA,sha256=Z99V_6BvYynyRZkZ2D5Ih6Dk-GWGeNZH0cTiMFEPSvg,8389
+libinephany-1.1.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+libinephany-1.1.3.dist-info/top_level.txt,sha256=bYAOXQdJgIoLkO2Ui0kxe7pSYegS_e38u0dMscd7COQ,12
+libinephany-1.1.3.dist-info/RECORD,,