truefoundry 0.4.0rc2__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/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

@@ -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(

truefoundry/ml/log_types/artifacts/model.py

@@ -146,11 +146,10 @@ class ModelVersion:
         self._metadata = copy.deepcopy(self._model_version.artifact_metadata)
 
     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()
 
@@ -382,13 +381,11 @@ class ModelVersion:
             ```
         """
         self._ensure_not_deleted()
-        kwargs = {}
         _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
@@ -536,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)
@@ -586,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_]",
 ]
 
 

truefoundry/ml/mlfoundry_api.py

@@ -30,6 +30,7 @@ from truefoundry.ml.autogen.client.exceptions import (
     ApiException,
     NotFoundException,
 )
+from truefoundry.ml.clients.servicefoundry_client import ServiceFoundryServiceClient
 from truefoundry.ml.enums import ModelFramework, ViewType
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.internal_namespace import NAMESPACE
@@ -39,7 +40,6 @@ from truefoundry.ml.log_types.artifacts.general_artifact import _log_artifact_ve
 from truefoundry.ml.log_types.artifacts.model import ModelVersion, _log_model_version
 from truefoundry.ml.logger import logger
 from truefoundry.ml.mlfoundry_run import MlFoundryRun
-from truefoundry.ml.services.servicefoundry_service import ServicefoundryService
 from truefoundry.ml.session import (
     Session,
     _get_api_client,
@@ -123,7 +123,6 @@ class MlFoundry:
             )
             raise MlFoundryException(err_msg) from e
 
-        assert ml_repo_instance.experiment_id is not None
         return ml_repo_instance.experiment_id
 
     def list_ml_repos(self) -> List[str]:
@@ -213,7 +212,7 @@ class MlFoundry:
                 "No active session found. Perhaps you are not logged in?\n"
                 "Please log in using `tfy login [--host HOST] --relogin"
             )
-        servicefoundry_service = ServicefoundryService(
+        servicefoundry_client = ServiceFoundryServiceClient(
             tracking_uri=self.get_tracking_uri(),
             token=session.token.access_token,
         )
@@ -221,7 +220,7 @@ class MlFoundry:
         assert existing_ml_repo.storage_integration_id is not None
         try:
             existing_storage_integration = (
-                servicefoundry_service.get_integration_from_id(
+                servicefoundry_client.get_integration_from_id(
                     existing_ml_repo.storage_integration_id
                 )
             )
@@ -335,11 +334,6 @@ class MlFoundry:
             )
         )
         run = _run.run
-
-        assert run is not None
-        assert run.info.run_id is not None
-        assert run.info.fqn is not None
-
         mlf_run_id = run.info.run_id
         kwargs.setdefault("auto_end", True)
         mlf_run = MlFoundryRun(experiment_id=ml_repo_id, run_id=mlf_run_id, **kwargs)
@@ -378,9 +372,6 @@ class MlFoundry:
             return self.get_run_by_fqn(run_id)
         _run = self._runs_api.get_run_get(run_id=run_id)
         run = _run.run
-
-        assert run is not None
-
         mlfoundry_run = MlFoundryRun._from_dto(run)
         logger.info(
             f"Link to the dashboard for the run: {mlfoundry_run.dashboard_link}"
@@ -722,7 +713,6 @@ class MlFoundry:
             version=resolved_version,
         )
         model_version = _model_version.model_version
-        assert model_version.model_id is not None
         _model = self._mlfoundry_artifacts_api.get_model_get(id=model_version.model_id)
         model = _model.model
 
@@ -865,7 +855,6 @@ class MlFoundry:
         model: Optional[ModelDto] = None,
     ) -> Iterator[ModelVersion]:
         if model and not model_id:
-            assert model.id is not None
             model_id = model.id
         elif not model and model_id:
             _model = self._mlfoundry_artifacts_api.get_model_get(id=str(model_id))
@@ -940,7 +929,6 @@ class MlFoundry:
             )
         )
         artifact_version = _artifact_version.artifact_version
-        assert artifact_version.artifact_id is not None
         _artifact = self._mlfoundry_artifacts_api.get_artifact_by_id_get(
             id=artifact_version.artifact_id
         )
@@ -1069,7 +1057,6 @@ class MlFoundry:
         artifact: Optional[ArtifactDto] = None,
     ) -> Iterator[ArtifactVersion]:
         if artifact and not artifact_id:
-            assert artifact.id is not None
             artifact_id = artifact.id
         elif not artifact and artifact_id:
             _artifact = self._mlfoundry_artifacts_api.get_artifact_by_id_get(
@@ -1273,6 +1260,7 @@ class MlFoundry:
             metadata (Optional[Dict[str, Any]], optional): arbitrary json serializable dictionary to store metadata.
                 For example, you can use this to store metrics, params, notes.
                 This field can be updated at any time after logging. Defaults to `None`
+            progress (bool): value to show progress bar, defaults to None.
 
         Returns:
             truefoundry.ml.ModelVersion: an instance of `ModelVersion` that can be used to download the files,
@@ -1382,7 +1370,7 @@ class MlFoundry:
         Args:
             ml_repo (str): Name of the ML Repo in which you want to create data_directory
             name (str): Name of the DataDirectory to be created.
-            description (str): Description of the Datset
+            description (str): Description for the DataDirectory.
             metadata (Dict <str>: Any): Metadata about the data_directory in Dictionary form.
 
         Returns:

truefoundry/ml/mlfoundry_run.py

@@ -110,8 +110,6 @@ class MlFoundryRun:
     @classmethod
     def _from_dto(cls, run_dto: RunDto) -> "MlFoundryRun":
         """classmethod to get MLfoundry run from dto instance"""
-        assert run_dto.info.experiment_id is not None
-        assert run_dto.info.run_id is not None
         run = cls(run_dto.info.experiment_id, run_dto.info.run_id)
         run._run_info = run_dto.info
         run._run_data = run_dto.data
@@ -136,7 +134,6 @@ class MlFoundryRun:
     def run_name(self) -> str:
         """Get run_name for the current `run`"""
         run_info = self._get_run_info()
-        assert run_info.name is not None
         return run_info.name
 
     @property
@@ -144,7 +141,6 @@ class MlFoundryRun:
     def fqn(self) -> str:
         """Get fqn for the current `run`"""
         run_info = self._get_run_info()
-        assert run_info.fqn is not None
         return run_info.fqn
 
     @property
@@ -162,7 +158,6 @@ class MlFoundryRun:
         _experiment = self._experiments_api.get_experiment_get(
             experiment_id=self._experiment_id
         )
-        assert _experiment.experiment.name is not None
         return _experiment.experiment.name
 
     @property

truefoundry/ml/run_utils.py

@@ -7,7 +7,6 @@ from urllib.parse import urljoin, urlsplit
 
 import numpy as np
 
-from truefoundry.ml import env_vars
 from truefoundry.ml.exceptions import MlFoundryException
 
 
@@ -22,21 +21,13 @@ def get_module(
             raise MlFoundryException(msg) from ex
 
 
-def resolve_tracking_uri(tracking_uri: Optional[str]):
-    if not tracking_uri and not os.getenv(env_vars.TRACKING_HOST_GLOBAL):
-        raise ValueError(
-            f"Either `host` should be provided by --host <value>, or `{env_vars.TRACKING_HOST_GLOBAL}` env must be set"
-        )
-    return tracking_uri or os.getenv(env_vars.TRACKING_HOST_GLOBAL)
-
-
 def append_path_to_rest_tracking_uri(tracking_uri: str):
     if urlsplit(tracking_uri).netloc.startswith("localhost"):
         return tracking_uri
     return urljoin(tracking_uri, "/api/ml")
 
 
-def append_servicefoundry_path_to_tracking_ui(tracking_uri: str):
+def append_servicefoundry_path_to_tracking_uri(tracking_uri: str):
     if urlsplit(tracking_uri).netloc.startswith("localhost"):
         return os.getenv("SERVICEFOUNDRY_SERVER_URL")
     return urljoin(tracking_uri, "/api/svc")