truefoundry 0.4.0rc1__py3-none-any.whl → 0.4.0rc3__py3-none-any.whl

truefoundry/ml/autogen/client/models/model_version_dto.py

@@ -24,7 +24,14 @@ from truefoundry.ml.autogen.client.models.artifact_version_status import (
 )
 from truefoundry.ml.autogen.client.models.metric_dto import MetricDto
 from truefoundry.ml.autogen.client.models.model_schema_dto import ModelSchemaDto
-from truefoundry.pydantic_v1 import BaseModel, StrictBool, StrictInt, StrictStr, conlist
+from truefoundry.pydantic_v1 import (
+    BaseModel,
+    Field,
+    StrictBool,
+    StrictInt,
+    StrictStr,
+    conlist,
+)
 
 
 class ModelVersionDto(BaseModel):
@@ -32,10 +39,11 @@ class ModelVersionDto(BaseModel):
     ModelVersionDto
     """
 
-    id: Optional[StrictStr] = None
-    model_id: Optional[StrictStr] = None
-    version: Optional[StrictInt] = None
-    fqn: Optional[StrictStr] = None
+    id: StrictStr = Field(...)
+    model_id: StrictStr = Field(...)
+    version: StrictInt = Field(...)
+    fqn: StrictStr = Field(...)
+    experiment_id: StrictStr = Field(...)
     artifact_storage_root: Optional[StrictStr] = None
     artifact_metadata: Optional[Dict[str, Any]] = None
     description: Optional[StrictStr] = None
@@ -50,7 +58,6 @@ class ModelVersionDto(BaseModel):
     data_path: Optional[StrictStr] = None
     monitoring_enabled: Optional[StrictBool] = False
     usage_code_snippet: Optional[StrictStr] = None
-    experiment_id: Optional[StrictStr] = None
     run_id: Optional[StrictStr] = None
     metrics: Optional[conlist(MetricDto)] = None
     model_name: Optional[StrictStr] = None
@@ -61,6 +68,7 @@ class ModelVersionDto(BaseModel):
         "model_id",
         "version",
         "fqn",
+        "experiment_id",
         "artifact_storage_root",
         "artifact_metadata",
         "description",
@@ -75,7 +83,6 @@ class ModelVersionDto(BaseModel):
         "data_path",
         "monitoring_enabled",
         "usage_code_snippet",
-        "experiment_id",
         "run_id",
         "metrics",
         "model_name",
@@ -132,6 +139,7 @@ class ModelVersionDto(BaseModel):
                 "model_id": obj.get("model_id"),
                 "version": obj.get("version"),
                 "fqn": obj.get("fqn"),
+                "experiment_id": obj.get("experiment_id"),
                 "artifact_storage_root": obj.get("artifact_storage_root"),
                 "artifact_metadata": obj.get("artifact_metadata"),
                 "description": obj.get("description"),
@@ -150,7 +158,6 @@ class ModelVersionDto(BaseModel):
                 if obj.get("monitoring_enabled") is not None
                 else False,
                 "usage_code_snippet": obj.get("usage_code_snippet"),
-                "experiment_id": obj.get("experiment_id"),
                 "run_id": obj.get("run_id"),
                 "metrics": [MetricDto.from_dict(_item) for _item in obj.get("metrics")]
                 if obj.get("metrics") is not None

truefoundry/ml/autogen/client/models/run_info_dto.py

@@ -18,7 +18,7 @@ import pprint
 import re  # noqa: F401
 from typing import Optional
 
-from truefoundry.pydantic_v1 import BaseModel, StrictInt, StrictStr
+from truefoundry.pydantic_v1 import BaseModel, Field, StrictInt, StrictStr
 
 
 class RunInfoDto(BaseModel):
@@ -26,30 +26,30 @@ class RunInfoDto(BaseModel):
     RunInfoDto
     """
 
-    run_uuid: Optional[StrictStr] = None
-    run_id: Optional[StrictStr] = None
-    experiment_id: Optional[StrictStr] = None
+    run_uuid: StrictStr = Field(...)
+    run_id: StrictStr = Field(...)
+    experiment_id: StrictStr = Field(...)
+    name: StrictStr = Field(...)
+    fqn: StrictStr = Field(...)
     user_id: Optional[StrictStr] = None
     status: Optional[StrictStr] = None
     start_time: Optional[StrictInt] = None
     end_time: Optional[StrictInt] = None
     artifact_uri: Optional[StrictStr] = None
     lifecycle_stage: Optional[StrictStr] = None
-    name: Optional[StrictStr] = None
-    fqn: Optional[StrictStr] = None
     description: Optional[StrictStr] = None
     __properties = [
         "run_uuid",
         "run_id",
         "experiment_id",
+        "name",
+        "fqn",
         "user_id",
         "status",
         "start_time",
         "end_time",
         "artifact_uri",
         "lifecycle_stage",
-        "name",
-        "fqn",
         "description",
     ]
 
@@ -91,14 +91,14 @@ class RunInfoDto(BaseModel):
                 "run_uuid": obj.get("run_uuid"),
                 "run_id": obj.get("run_id"),
                 "experiment_id": obj.get("experiment_id"),
+                "name": obj.get("name"),
+                "fqn": obj.get("fqn"),
                 "user_id": obj.get("user_id"),
                 "status": obj.get("status"),
                 "start_time": obj.get("start_time"),
                 "end_time": obj.get("end_time"),
                 "artifact_uri": obj.get("artifact_uri"),
                 "lifecycle_stage": obj.get("lifecycle_stage"),
-                "name": obj.get("name"),
-                "fqn": obj.get("fqn"),
                 "description": obj.get("description"),
             }
         )

truefoundry/ml/autogen/client/models/update_run_response_dto.py

@@ -16,10 +16,9 @@ from __future__ import annotations
 import json
 import pprint
 import re  # noqa: F401
-from typing import Optional
 
 from truefoundry.ml.autogen.client.models.run_info_dto import RunInfoDto
-from truefoundry.pydantic_v1 import BaseModel
+from truefoundry.pydantic_v1 import BaseModel, Field
 
 
 class UpdateRunResponseDto(BaseModel):
@@ -27,7 +26,7 @@ class UpdateRunResponseDto(BaseModel):
     UpdateRunResponseDto
     """
 
-    run_info: Optional[RunInfoDto] = None
+    run_info: RunInfoDto = Field(...)
     __properties = ["run_info"]
 
     class Config:

truefoundry/ml/autogen/client_README.md

@@ -69,7 +69,6 @@ Class | Method | HTTP request | Description
 *DeprecatedApi* | [**get_run_by_name_get**](truefoundry/ml/autogen/client/docs/DeprecatedApi.md#get_run_by_name_get) | **GET** /api/2.0mlflow/runs/get-by-name | Get Run By Name
 *DeprecatedApi* | [**get_search_runs_get**](truefoundry/ml/autogen/client/docs/DeprecatedApi.md#get_search_runs_get) | **GET** /api/2.0/preview/mlflow/runs/search | Get Search Runs
 *DeprecatedApi* | [**get_signed_urls_for_dataset_write_deprecated_get**](truefoundry/ml/autogen/client/docs/DeprecatedApi.md#get_signed_urls_for_dataset_write_deprecated_get) | **GET** /api/2.0/mlflow/mlfoundry-artifacts/datasets/get-signed-urls-for-write | Get Signed Urls For Dataset Write Deprecated
-*ExperimentsApi* | [**backfill_storage_integration_post**](truefoundry/ml/autogen/client/docs/ExperimentsApi.md#backfill_storage_integration_post) | **POST** /api/2.0/mlflow/experiments/storage-integration/backfill | Backfill Storage Integration
 *ExperimentsApi* | [**create_experiment_post**](truefoundry/ml/autogen/client/docs/ExperimentsApi.md#create_experiment_post) | **POST** /api/2.0/mlflow/experiments/create | Create Experiment
 *ExperimentsApi* | [**delete_experiment_post**](truefoundry/ml/autogen/client/docs/ExperimentsApi.md#delete_experiment_post) | **POST** /api/2.0/mlflow/experiments/delete | Delete Experiment
 *ExperimentsApi* | [**get_experiment_by_name_get**](truefoundry/ml/autogen/client/docs/ExperimentsApi.md#get_experiment_by_name_get) | **GET** /api/2.0/mlflow/experiments/get-by-name | Get Experiment By Name
@@ -170,7 +169,6 @@ Class | Method | HTTP request | Description
  - [AssistantMessage](truefoundry/ml/autogen/client/docs/AssistantMessage.md)
  - [AuthorizeUserForModelRequestDto](truefoundry/ml/autogen/client/docs/AuthorizeUserForModelRequestDto.md)
  - [AuthorizeUserForModelVersionRequestDto](truefoundry/ml/autogen/client/docs/AuthorizeUserForModelVersionRequestDto.md)
- - [BackfillDefaultStorageIntegrationIdRequestDto](truefoundry/ml/autogen/client/docs/BackfillDefaultStorageIntegrationIdRequestDto.md)
  - [BlobStorageReference](truefoundry/ml/autogen/client/docs/BlobStorageReference.md)
  - [BodyGetSearchRunsGet](truefoundry/ml/autogen/client/docs/BodyGetSearchRunsGet.md)
  - [ChatPrompt](truefoundry/ml/autogen/client/docs/ChatPrompt.md)

truefoundry/ml/clients/entities.py

@@ -0,0 +1,8 @@
+from typing import Optional
+
+from truefoundry.pydantic_v1 import BaseModel
+
+
+class HostCreds(BaseModel):
+    host: str
+    token: Optional[str] = None

truefoundry/ml/{services/servicefoundry_service.py → clients/servicefoundry_client.py}

@@ -1,28 +1,38 @@
 from typing import Optional
 
-from truefoundry.common.request_utils import requests_retry_session
+from truefoundry.common.constants import (
+    SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
+    VERSION_PREFIX,
+)
+from truefoundry.common.servicefoundry_client import (
+    ServiceFoundryServiceClient as BaseServiceFoundryServiceClient,
+)
+from truefoundry.ml.clients.entities import (
+    HostCreds,
+)
+from truefoundry.ml.clients.utils import http_request_safe
 from truefoundry.ml.exceptions import MlFoundryException
-from truefoundry.ml.run_utils import append_servicefoundry_path_to_tracking_ui
-from truefoundry.ml.services.entities import HostCreds
-from truefoundry.ml.services.utils import http_request_safe
+from truefoundry.ml.run_utils import append_servicefoundry_path_to_tracking_uri
 
 
-class ServicefoundryService:
+class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     def __init__(self, tracking_uri: str, token: Optional[str] = None):
+        super().__init__(base_url=tracking_uri)
         self.host_creds = HostCreds(
-            host=append_servicefoundry_path_to_tracking_ui(tracking_uri), token=token
+            host=append_servicefoundry_path_to_tracking_uri(tracking_uri), token=token
         )
 
     def get_integration_from_id(self, integration_id: str):
         integration_id = integration_id or ""
-        data = http_request_safe(
-            method="get",
+        response = http_request_safe(
             host_creds=self.host_creds,
-            endpoint="v1/provider-accounts/provider-integrations",
-            session=requests_retry_session(retries=1),
+            endpoint=f"{VERSION_PREFIX}/provider-accounts/provider-integrations",
             params={"id": integration_id, "type": "blob-storage"},
+            method="get",
             timeout=3,
+            max_retries=SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
         )
+        data = response.json()
         if (
             data.get("providerIntegrations")
             and len(data["providerIntegrations"]) > 0

truefoundry/ml/{services → clients}/utils.py

@@ -6,8 +6,8 @@ from urllib.parse import urljoin
 import requests
 
 from truefoundry.common.request_utils import requests_retry_session
+from truefoundry.ml.clients.entities import HostCreds
 from truefoundry.ml.exceptions import MlFoundryException
-from truefoundry.ml.services.entities import HostCreds
 
 
 # TODO: This will be moved later to truefoundry.common.request_utils
@@ -77,7 +77,7 @@ def cloud_storage_http_request(
     """
     Performs an HTTP PUT/GET request using Python's `requests` module with automatic retry.
     """
-    session = session or requests_retry_session(retries=5, backoff_factor=2)
+    session = session or requests_retry_session(retries=5, backoff_factor=0.5)
     kwargs["headers"] = kwargs.get("headers", {}) or {}
     if "blob.core.windows.net" in url:
         kwargs["headers"].update({"x-ms-blob-type": "BlockBlob"})

truefoundry/ml/entities.py

@@ -0,0 +1,62 @@
+from typing import Optional, Union
+
+from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+    FileInfoDto,
+    MetricDto,
+)
+from truefoundry.pydantic_v1 import (
+    BaseModel,
+    Field,
+    StrictBool,
+    StrictFloat,
+    StrictInt,
+    StrictStr,
+)
+
+
+class Metric(BaseModel):
+    key: StrictStr = Field(...)
+    value: Optional[Union[StrictFloat, StrictInt]] = None
+    timestamp: Optional[StrictInt] = None
+    step: Optional[StrictInt] = 0
+
+    @classmethod
+    def from_dto(cls, dto: MetricDto) -> "Metric":
+        return cls(
+            key=dto.key,
+            value=dto.value,
+            timestamp=dto.timestamp,
+            step=dto.step,
+        )
+
+    def to_dto(self) -> MetricDto:
+        return MetricDto(
+            key=self.key,
+            value=self.value,
+            timestamp=self.timestamp,
+            step=self.step,
+        )
+
+
+class FileInfo(BaseModel):
+    path: StrictStr
+    is_dir: StrictBool
+    file_size: Optional[StrictInt] = None
+    signed_url: Optional[StrictStr] = None
+
+    @classmethod
+    def from_dto(cls, dto: FileInfoDto) -> "FileInfo":
+        return cls(
+            path=dto.path,
+            is_dir=dto.is_dir,
+            file_size=dto.file_size,
+            signed_url=dto.signed_url,
+        )
+
+    def to_dto(self) -> FileInfoDto:
+        return FileInfoDto(
+            path=self.path,
+            is_dir=self.is_dir,
+            file_size=self.file_size,
+            signed_url=self.signed_url,
+        )

truefoundry/ml/env_vars.py

@@ -1,11 +1,7 @@
 """
-Define all the env variable names that users
-can use
+Define all the env variable names that users can use
 """
 
-# NOTE: These are shared between servicefoundry and mlfoundry
-TRACKING_HOST_GLOBAL = "TFY_HOST"
-API_KEY_GLOBAL = "TFY_API_KEY"
 DISABLE_MULTIPART_UPLOAD = "MLF_DISABLE_MULTIPART_UPLOAD"
 INTERNAL_ENV_VARS = [
     "TFY_INTERNAL_APPLICATION_ID",

truefoundry/ml/internal_namespace.py

@@ -6,39 +6,39 @@ from truefoundry.ml.exceptions import MlFoundryException
 
 class _InternalNamespace:
     NAMESPACE = "mlf"
-    DELIMETER = "."
+    DELIMITER = "."
     NAMESPACE_VIOLATION_MESSAGE = """
         {name} cannot start with {prefix}
     """
 
-    def __call__(self, name: str, delimeter: str = DELIMETER):
+    def __call__(self, name: str, delimiter: str = DELIMITER):
         if not name:
             raise MlFoundryException("name should be a non empty string")
-        return _InternalNamespace.NAMESPACE + delimeter + name
+        return _InternalNamespace.NAMESPACE + delimiter + name
 
     def __truediv__(self, path: str):
         return os.path.join(_InternalNamespace.NAMESPACE, path)
 
     @staticmethod
-    def _validate_name_not_using_namespace(name: typing.Optional[str], delimeter: str):
-        if name and name.startswith(_InternalNamespace.NAMESPACE + delimeter):
+    def _validate_name_not_using_namespace(name: typing.Optional[str], delimiter: str):
+        if name and name.startswith(_InternalNamespace.NAMESPACE + delimiter):
             raise MlFoundryException(
                 _InternalNamespace.NAMESPACE_VIOLATION_MESSAGE.format(
-                    name=name, prefix=_InternalNamespace.NAMESPACE + delimeter
+                    name=name, prefix=_InternalNamespace.NAMESPACE + delimiter
                 )
             )
 
     def validate_namespace_not_used(
         self,
         names: typing.Optional[typing.Union[str, typing.Iterable[str]]] = None,
-        delimeter: str = DELIMETER,
+        delimiter: str = DELIMITER,
         path: typing.Optional[str] = None,
     ):
         if isinstance(names, str):
             names = [names]
         if names is not None:
             for name_ in names:
-                self._validate_name_not_using_namespace(name_, delimeter)
+                self._validate_name_not_using_namespace(name_, delimiter)
         if path:
             prefix = os.path.normpath(os.path.join(_InternalNamespace.NAMESPACE, ""))
             if os.path.normpath(path).startswith(prefix):

truefoundry/ml/log_types/artifacts/artifact.py

@@ -3,6 +3,7 @@ import datetime
 import json
 import os
 import tempfile
+import uuid
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Optional, Tuple, Union
 
@@ -21,7 +22,7 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     NotifyArtifactVersionFailureDto,
     UpdateArtifactVersionRequestDto,
 )
-from truefoundry.ml.autogen.client import (
+from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     InternalMetadata as InternalMetadataDto,
 )
 from truefoundry.ml.exceptions import MlFoundryException
@@ -130,7 +131,7 @@ class ArtifactVersion:
     def _get_artifacts_repo(self):
         return MlFoundryArtifactsRepository(
             artifact_identifier=ArtifactIdentifier(
-                artifact_version_id=self._artifact_version.id
+                artifact_version_id=uuid.UUID(self._artifact_version.id)
             ),
             api_client=self._api_client,
         )
@@ -374,6 +375,7 @@ def _log_artifact_version_helper(
     if run:
         mlfoundry_artifacts_api = run._mlfoundry_artifacts_api
 
+    assert mlfoundry_artifacts_api is not None
     _create_artifact_response = mlfoundry_artifacts_api.create_artifact_version_post(
         create_artifact_version_request_dto=CreateArtifactVersionRequestDto(
             experiment_id=int(run._experiment_id if run else ml_repo_id),
@@ -383,7 +385,9 @@ def _log_artifact_version_helper(
     )
     version_id = _create_artifact_response.id
     artifacts_repo = MlFoundryArtifactsRepository(
-        artifact_identifier=ArtifactIdentifier(artifact_version_id=version_id),
+        artifact_identifier=ArtifactIdentifier(
+            artifact_version_id=uuid.UUID(version_id),
+        ),
         api_client=mlfoundry_artifacts_api.api_client,
     )
     total_size = calculate_local_directory_size(artifact_dir)

truefoundry/ml/log_types/artifacts/dataset.py

@@ -12,10 +12,10 @@ from truefoundry.ml.artifact.truefoundry_artifact_repo import (
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     DatasetDto,
     DeleteDatasetRequestDto,
-    FileInfoDto,
     MlfoundryArtifactsApi,
     UpdateDatasetRequestDto,
 )
+from truefoundry.ml.entities import FileInfo
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.utils import (
     _copy_additional_files,
@@ -136,7 +136,7 @@ class DataDirectory:
 
     @property
     def updated_at(self) -> datetime.datetime:
-        """Get the information about the when the DataDirectory was updated"""
+        """Get the information about when the DataDirectory was updated"""
         return self._dataset.updated_at
 
     def add_files(
@@ -239,7 +239,7 @@ class DataDirectory:
     def list_files(
         self,
         path: Optional[str] = None,
-    ) -> Iterator[FileInfoDto]:
+    ) -> Iterator[FileInfo]:
         """
         List all the files and folders in the data_directory.
 
@@ -262,7 +262,8 @@ class DataDirectory:
             ```
         """
         artifacts_repo = self._get_artifacts_repo()
-        return artifacts_repo.list_artifacts(path=path)
+        for file_info_dto in artifacts_repo.list_artifacts(path=path):
+            yield FileInfo.from_dto(file_info_dto)
 
     def download(
         self,

truefoundry/ml/log_types/artifacts/model.py

@@ -53,7 +53,6 @@ if TYPE_CHECKING:
 logger = logging.getLogger("truefoundry.ml")
 
 
-# TODO: Add some progress indicators for upload and download
 # TODO: Support async download and upload
 
 
@@ -98,7 +97,6 @@ class ModelVersion:
         self._deleted = False
         self._description: str = ""
         self._metadata: Dict[str, Any] = {}
-        self._model_schema: Optional[ModelSchema] = None
         self._metrics: List[MetricDto] = []
         self._set_metrics_attr()
         self._set_mutable_attrs()
@@ -146,19 +144,17 @@ class ModelVersion:
     def _set_mutable_attrs(self):
         self._description = self._model_version.description or ""
         self._metadata = copy.deepcopy(self._model_version.artifact_metadata)
-        self._model_schema = copy.deepcopy(self._model_version.model_schema)
 
     def _refetch_model_version(self):
-        self._model_version = (
-            self._mlfoundry_artifacts_api.get_artifact_version_by_id_get(
-                id=self._model_version.id
-            )
+        _model_version = self._mlfoundry_artifacts_api.get_model_version_get(
+            id=self._model_version.id
         )
+        self._model_version = _model_version.model_version
         self._set_metrics_attr()
         self._set_mutable_attrs()
 
     def __repr__(self):
-        return f"{self.__class__.__name__}(fqn={self.fqn!r}, model_schema={self._model_schema!r})"
+        return f"{self.__class__.__name__}(fqn={self.fqn!r})"
 
     def _get_artifacts_repo(self):
         return MlFoundryArtifactsRepository(
@@ -215,17 +211,6 @@ class ModelVersion:
         _validate_artifact_metadata(value)
         self._metadata = value
 
-    @property
-    def model_schema(self) -> Optional[ModelSchema]:
-        """get the model schema for current model"""
-        return self._model_schema
-
-    @model_schema.setter
-    def model_schema(self, value: Union[Dict[str, Any], ModelSchema]):
-        if not isinstance(value, ModelSchema):
-            value = ModelSchema.parse_obj(value)
-        self._model_schema = value
-
     @property
     def metrics(self) -> Dict[str, Union[float, int]]:
         """get the metrics for the current version of the model"""
@@ -396,15 +381,11 @@ class ModelVersion:
             ```
         """
         self._ensure_not_deleted()
-        kwargs = {}
-        if self.model_schema is not None:
-            kwargs["model_schema"] = self.model_schema
         _model_version = self._mlfoundry_artifacts_api.update_model_version_post(
             update_model_version_request_dto=UpdateModelVersionRequestDto(
                 id=self._model_version.id,
                 description=self.description,
                 artifact_metadata=self.metadata,
-                **kwargs,
             )
         )
         self._model_version = _model_version.model_version
@@ -552,7 +533,9 @@ def _log_model_version(  # noqa: C901
     )
     version_id = _create_artifact_version_response.id
     artifacts_repo = MlFoundryArtifactsRepository(
-        artifact_identifier=ArtifactIdentifier(artifact_version_id=version_id),
+        artifact_identifier=ArtifactIdentifier(
+            artifact_version_id=uuid.UUID(version_id)
+        ),
         api_client=mlfoundry_artifacts_api.api_client,
     )
     model_size = calculate_model_size(temp_dir)
@@ -602,7 +585,7 @@ def _log_model_version(  # noqa: C901
     model_version = _model_version.model_version
 
     # update model schema at end
-    update_args = {
+    update_args: Dict[str, Any] = {
         "id": version_id,
         "model_framework": framework.value,
     }

truefoundry/ml/log_types/image/image.py

@@ -6,8 +6,6 @@ import tempfile
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-import numpy as np
-
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
 )
@@ -33,17 +31,18 @@ from truefoundry.ml.run_utils import get_module
 
 if TYPE_CHECKING:
     # noinspection PyUnresolvedReferences
+    import numpy  # noqa: F401
     import PIL.Image
-    from numpy.typing import NDArray
+    from numpy.typing import NDArray  # noqa: F401
 
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
 
 ImageLikeNDArray = Union[
-    "NDArray[np.int8]",
-    "NDArray[np.uint8]",
-    "NDArray[np.float32]",
-    "NDArray[np.float64]",
-    "NDArray[np.bool_]",
+    "NDArray[numpy.int8]",
+    "NDArray[numpy.uint8]",
+    "NDArray[numpy.float32]",
+    "NDArray[numpy.float64]",
+    "NDArray[numpy.bool_]",
 ]
 DataOrPathType = Union[str, Path, ImageLikeNDArray, "PIL.Image.Image"]
 ClassGroupsType = Dict[str, Union[str, List[str]]]

truefoundry/ml/log_types/image/image_normalizer.py

@@ -10,15 +10,16 @@ from truefoundry.ml.logger import logger
 from truefoundry.ml.run_utils import get_module
 
 if TYPE_CHECKING:
+    import numpy  # noqa: F401
     import PIL.Image
-    from numpy.typing import NDArray
+    from numpy.typing import NDArray  # noqa: F401
 
 ImageLikeNDArray = Union[
-    "NDArray[np.int8]",
-    "NDArray[np.uint8]",
-    "NDArray[np.float32]",
-    "NDArray[np.float64]",
-    "NDArray[np.bool_]",
+    "NDArray[numpy.int8]",
+    "NDArray[numpy.uint8]",
+    "NDArray[numpy.float32]",
+    "NDArray[numpy.float64]",
+    "NDArray[numpy.bool_]",
 ]