truefoundry 0.5.3rc4__py3-none-any.whl → 0.5.3rc5__py3-none-any.whl

truefoundry/ml/clients/servicefoundry_client.py

@@ -1,36 +1,57 @@
-from typing import Optional
+import functools
 
 from truefoundry.common.constants import (
     SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
     VERSION_PREFIX,
 )
+from truefoundry.common.exceptions import HttpRequestException
+from truefoundry.common.request_utils import (
+    http_request,
+    request_handling,
+    requests_retry_session,
+)
 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
 
 
 class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
-    # TODO (chiragjn): Rename tracking_uri to tfy_host
-    def __init__(self, tracking_uri: str, token: Optional[str] = None):
-        super().__init__(base_url=tracking_uri)
-        self.host_creds = HostCreds(host=self._api_server_url, token=token)
+    def __init__(self, tfy_host: str, token: str):
+        super().__init__(tfy_host=tfy_host)
+        self._token = token
+
+    @functools.cached_property
+    def _min_cli_version_required(self) -> str:
+        # TODO (chiragjn): read the mlfoundry min cli version from the config?
+        return self.python_sdk_config.truefoundry_cli_min_version
 
     def get_integration_from_id(self, integration_id: str):
         integration_id = integration_id or ""
-        response = http_request_safe(
-            host_creds=self.host_creds,
-            endpoint=f"{VERSION_PREFIX}/provider-accounts/provider-integrations",
-            params={"id": integration_id, "type": "blob-storage"},
+        session = requests_retry_session(retries=SERVICEFOUNDRY_CLIENT_MAX_RETRIES)
+        response = http_request(
             method="get",
+            url=f"{self._api_server_url}/{VERSION_PREFIX}/provider-accounts/provider-integrations",
+            token=self._token,
             timeout=3,
-            max_retries=SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
+            params={"id": integration_id, "type": "blob-storage"},
+            session=session,
         )
-        data = response.json()
+
+        try:
+            data = request_handling(response)
+            assert isinstance(data, dict)
+        except HttpRequestException as he:
+            raise MlFoundryException(
+                f"Failed to get storage integration from id: {integration_id}. Error: {he.message}",
+                status_code=he.status_code,
+            ) from None
+        except Exception as e:
+            raise MlFoundryException(
+                f"Failed to get storage integration from id: {integration_id}. Error: {str(e)}"
+            ) from None
+
+        # TODO (chiragjn): Parse this using Pydantic
         if (
             data.get("providerIntegrations")
             and len(data["providerIntegrations"]) > 0

truefoundry/ml/exceptions.py

@@ -1,8 +1,9 @@
 from typing import Optional
 
 
+# TODO (chiragjn): We need to establish uniform exception handling across codebase
 class MlFoundryException(Exception):
-    def __init__(self, message, status_code: Optional[int] = None):
+    def __init__(self, message: str, status_code: Optional[int] = None):
         self.message = str(message)
         self.status_code = status_code
         super().__init__(message)

truefoundry/ml/log_types/artifacts/artifact.py

@@ -9,6 +9,7 @@ import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Union
 
+from truefoundry.common.warnings import TrueFoundryDeprecationWarning
 from truefoundry.ml.artifact.truefoundry_artifact_repo import (
     ArtifactIdentifier,
     MlFoundryArtifactsRepository,
@@ -217,7 +218,7 @@ class ArtifactVersion:
         if not self._artifact_version.manifest:
             warnings.warn(
                 message="This model version was created using an older serialization format. tags do not exist, returning empty list",
-                category=DeprecationWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
         return self._tags
@@ -230,7 +231,7 @@ class ArtifactVersion:
         if not self._artifact_version.manifest:
             warnings.warn(
                 message="This model version was created using an older serialization format. Tags will not be updated",
-                category=DeprecationWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
             return

truefoundry/ml/log_types/artifacts/model.py

@@ -10,6 +10,7 @@ import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
+from truefoundry.common.warnings import TrueFoundryDeprecationWarning
 from truefoundry.ml.artifact.truefoundry_artifact_repo import (
     ArtifactIdentifier,
     MlFoundryArtifactsRepository,
@@ -245,7 +246,7 @@ class ModelVersion:
         if not self._model_version.manifest:
             warnings.warn(
                 message="This model version was created using an older serialization format. tags do not exist, returning empty list",
-                category=DeprecationWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
         return self._tags
@@ -258,7 +259,7 @@ class ModelVersion:
         if not self._model_version.manifest:
             warnings.warn(
                 message="This model version was created using an older serialization format. Tags will not be updated",
-                category=DeprecationWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
             return
@@ -270,7 +271,7 @@ class ModelVersion:
         if not self._model_version.manifest:
             warnings.warn(
                 message="This model version was created using an older serialization format. environment does not exist, returning None",
-                category=DeprecationWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
         return self._environment
@@ -281,7 +282,7 @@ class ModelVersion:
         if not self._model_version.manifest:
             warnings.warn(
                 message="This model version was created using an older serialization format. Environment will not be updated",
-                category=DeprecationWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
             return
@@ -300,7 +301,7 @@ class ModelVersion:
         if not self._model_version.manifest:
             warnings.warn(
                 message="This model version was created using an older serialization format. Framework will not be updated",
-                category=DeprecationWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
             return

truefoundry/ml/log_types/artifacts/utils.py

@@ -3,7 +3,7 @@ import logging
 import os
 import posixpath
 from pathlib import Path, PureWindowsPath
-from typing import Any, Dict, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.constants import DESCRIPTION_MAX_LENGTH
@@ -63,7 +63,7 @@ def get_single_file_path_if_only_one_in_directory(path: str) -> Optional[str]:
 
     # If it's a directory, check if it contains a single file
     if is_destination_path_dirlike(path):
-        all_files = []
+        all_files: List[str] = []
         for root, _, files in os.walk(path):
             # Collect all files found in any subdirectory
             all_files.extend(os.path.join(root, f) for f in files)

truefoundry/ml/mlfoundry_api.py

@@ -17,7 +17,7 @@ from typing import (
 
 import coolname
 
-from truefoundry.common.utils import ContextualDirectoryManager, relogin_error_message
+from truefoundry.common.utils import ContextualDirectoryManager
 from truefoundry.ml import constants
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactDto,
@@ -63,10 +63,8 @@ from truefoundry.ml.log_types.artifacts.model import (
 from truefoundry.ml.logger import logger
 from truefoundry.ml.mlfoundry_run import MlFoundryRun
 from truefoundry.ml.session import (
-    Session,
+    MLFoundrySession,
     _get_api_client,
-    get_active_session,
-    init_session,
 )
 from truefoundry.ml.validation_utils import (
     _validate_ml_repo_description,
@@ -112,13 +110,12 @@ class MlFoundry:
     """MlFoundry."""
 
     # TODO (chiragjn): Don't allow session as None here!
-    def __init__(self, session: Session):
+    def __init__(self, session: MLFoundrySession):
         """__init__
 
         Args:
             session (Optional[Session], optional): Session instance to get auth credentials from
         """
-        self._tracking_uri: str = session.tracking_uri
         self._api_client = _get_api_client(session=session)
         self._experiments_api = ExperimentsApi(api_client=self._api_client)
         self._runs_api = RunsApi(api_client=self._api_client)
@@ -236,16 +233,9 @@ class MlFoundry:
                 raise MlFoundryException(err_msg) from e
             return
 
-        session = get_active_session()
-        if session is None:
-            raise MlFoundryException(
-                relogin_error_message(
-                    "No active session found. Perhaps you are not logged in?",
-                )
-            )
         servicefoundry_client = ServiceFoundryServiceClient(
-            tracking_uri=self.get_tracking_uri(),
-            token=session.token.access_token,
+            tfy_host=self._api_client.tfy_host,
+            token=self._api_client.access_token,
         )
 
         assert existing_ml_repo.storage_integration_id is not None
@@ -604,26 +594,6 @@ class MlFoundry:
             if not runs or page_token is None:
                 done = True
 
-    def get_tracking_uri(self) -> str:
-        """
-        Get the current tracking URI.
-
-        Returns:
-            The tracking URI.
-
-        Examples:
-
-            ```python
-            import tempfile
-            from truefoundry.ml import get_client
-
-            client = get_client()
-            tracking_uri = client.get_tracking_uri()
-            print("Current tracking uri: {}".format(tracking_uri))
-            ```
-        """
-        return self._tracking_uri
-
     def _initialize_model_server(
         self,
         name: str,
@@ -1239,7 +1209,6 @@ class MlFoundry:
             raise MlFoundryException(
                 "artifact_paths cannot be empty, atleast one artifact_path must be passed"
             )
-
         ml_repo_id = self._get_ml_repo_id(ml_repo=ml_repo)
         artifact_version = _log_artifact_version(
             run=None,
@@ -1377,7 +1346,6 @@ class MlFoundry:
 
         """
         ml_repo_id = self._get_ml_repo_id(ml_repo=ml_repo)
-
         model_version = _log_model_version(
             run=None,
             mlfoundry_artifacts_api=self._mlfoundry_artifacts_api,
@@ -1648,5 +1616,5 @@ def get_client() -> MlFoundry:
         client = get_client()
         ```
     """
-    session = init_session()
+    session = MLFoundrySession.new()
     return MlFoundry(session=session)

truefoundry/ml/mlfoundry_run.py

@@ -17,7 +17,6 @@ from typing import (
 from urllib.parse import urljoin, urlsplit
 
 from truefoundry import version
-from truefoundry.common.utils import relogin_error_message
 from truefoundry.ml import constants
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
@@ -55,7 +54,7 @@ from truefoundry.ml.log_types.artifacts.model import (
 )
 from truefoundry.ml.logger import logger
 from truefoundry.ml.run_utils import ParamsType, flatten_dict, process_params
-from truefoundry.ml.session import ACTIVE_RUNS, _get_api_client, get_active_session
+from truefoundry.ml.session import ACTIVE_RUNS, _get_api_client
 from truefoundry.ml.validation_utils import (
     MAX_ENTITY_KEY_LENGTH,
     MAX_METRICS_PER_BATCH,
@@ -72,7 +71,7 @@ if TYPE_CHECKING:
 
 def _ensure_not_deleted(method):
     @functools.wraps(method)
-    def _check_deleted_or_not(self, *args, **kwargs):
+    def _check_deleted_or_not(self: "MlFoundryRun", *args, **kwargs):
         if self._deleted:
             raise MlFoundryException("Run was deleted, cannot access a deleted Run")
         else:
@@ -230,18 +229,10 @@ class MlFoundryRun:
     @_ensure_not_deleted
     def dashboard_link(self) -> str:
         """Get Mlfoundry dashboard link for a `run`"""
-        session = get_active_session()
-        if session is None:
-            raise MlFoundryException(
-                relogin_error_message(
-                    "No active session found. Perhaps you are not logged in?",
-                )
-            )
-        base_url = "{uri.scheme}://{uri.netloc}/".format(
-            uri=urlsplit(session.tracking_uri)
+        tfy_host = "{uri.scheme}://{uri.netloc}/".format(
+            uri=urlsplit(self._api_client.tfy_host)
         )
-
-        return urljoin(base_url, f"mlfoundry/{self._experiment_id}/run/{self.run_id}/")
+        return urljoin(tfy_host, f"mlfoundry/{self._experiment_id}/run/{self.run_id}/")
 
     @_ensure_not_deleted
     def end(self, status: RunStatus = RunStatus.FINISHED):
@@ -581,7 +572,7 @@ class MlFoundryRun:
             )
 
         return _log_artifact_version(
-            self,
+            run=self,
             name=name,
             artifact_paths=artifact_paths,
             description=description,

truefoundry/ml/model_framework.py

@@ -20,6 +20,7 @@ from truefoundry.common.utils import (
     get_python_version_major_minor,
     list_pip_packages_installed,
 )
+from truefoundry.common.warnings import TrueFoundryDeprecationWarning
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ModelVersionEnvironment,
     SklearnSerializationFormat,
@@ -260,7 +261,7 @@ class _ModelFramework(BaseModel):
         if isinstance(framework, (str, ModelFramework)):
             warnings.warn(
                 "Passing a string or ModelFramework Enum is deprecated. Please use a ModelFrameworkType object.",
-                DeprecationWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
 

truefoundry/ml/session.py

@@ -3,26 +3,22 @@ import threading
 import weakref
 from typing import TYPE_CHECKING, Dict, Optional
 
-from truefoundry.common.credential_provider import (
-    CredentialProvider,
-    EnvCredentialProvider,
-    FileCredentialProvider,
-)
-from truefoundry.common.entities import Token, UserInfo
 from truefoundry.common.request_utils import urllib3_retry
+from truefoundry.common.session import Session
 from truefoundry.common.utils import get_tfy_servers_config, relogin_error_message
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ApiClient,
     Configuration,
 )
-from truefoundry.ml.clients.entities import HostCreds
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.logger import logger
+from truefoundry.version import __version__
 
 if TYPE_CHECKING:
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
 
 SESSION_LOCK = threading.RLock()
+ACTIVE_SESSION: Optional["MLFoundrySession"] = None
 
 
 class ActiveRuns:
@@ -51,19 +47,7 @@ ACTIVE_RUNS = ActiveRuns()
 atexit.register(ACTIVE_RUNS.close_active_runs)
 
 
-class Session:
-    def __init__(self, cred_provider: CredentialProvider):
-        # Note: Whenever a new session is initialized all the active runs are ended
-        self._closed = False
-        self._cred_provider: Optional[CredentialProvider] = cred_provider
-        self._user_info: Optional[UserInfo] = self._cred_provider.token.to_user_info()
-
-    def close(self):
-        logger.debug("Closing existing session")
-        self._closed = True
-        self._user_info = None
-        self._cred_provider = None
-
+class MLFoundrySession(Session):
     def _assert_not_closed(self):
         if self._closed:
             raise MlFoundryException(
@@ -72,100 +56,88 @@ class Session:
                 "`truefoundry.ml.get_client()` function call) can be used"
             )
 
-    @property
-    def token(self) -> Token:
-        return self._cred_provider.token
+    def close(self):
+        global ACTIVE_RUNS
+        logger.debug("Closing existing session")
+        ACTIVE_RUNS.close_active_runs()
+        super().close()
 
-    @property
-    def user_info(self) -> UserInfo:
-        self._assert_not_closed()
-        return self._user_info
+    @classmethod
+    def new(cls) -> "MLFoundrySession":
+        global ACTIVE_SESSION
+        with SESSION_LOCK:
+            new_session = cls()
+            if ACTIVE_SESSION and ACTIVE_SESSION == new_session:
+                return ACTIVE_SESSION
+
+            if ACTIVE_SESSION:
+                ACTIVE_SESSION.close()
+
+            ACTIVE_SESSION = new_session
+            logger.info(
+                "Logged in to %r as %r (%s)",
+                new_session.tfy_host,
+                new_session.user_info.user_id,
+                new_session.user_info.email or new_session.user_info.user_type.value,
+            )
 
-    # TODO (chiragjn): Rename tracking_uri to tfy_host
-    @property
-    def tracking_uri(self) -> str:
-        return self._cred_provider.base_url
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, Session):
-            return False
-        return (
-            type(self._cred_provider) == type(other._cred_provider)  # noqa: E721
-            and self.user_info == other.user_info
-            and self.tracking_uri == other.tracking_uri
-        )
+            return ACTIVE_SESSION
 
-    def get_host_creds(self) -> HostCreds:
-        tracking_uri = get_tfy_servers_config(self.tracking_uri).mlfoundry_server_url
-        return HostCreds(
-            host=tracking_uri, token=self._cred_provider.token.access_token
-        )
 
+class MLFoundryServerApiClient(ApiClient):
+    def __init__(self, session: Optional[MLFoundrySession] = None, *args, **kwargs):
+        self.session = session
+        super().__init__(*args, **kwargs)
 
-ACTIVE_SESSION: Optional[Session] = None
+    @classmethod
+    def from_session(cls, session: MLFoundrySession) -> "MLFoundryServerApiClient":
+        mlfoundry_server_url = get_tfy_servers_config(
+            session.tfy_host
+        ).mlfoundry_server_url
+        configuration = Configuration(
+            host=mlfoundry_server_url.rstrip("/"),
+            access_token=session.access_token,
+        )
+        configuration.retries = urllib3_retry(retries=2)
+        api_client = cls(session=session, configuration=configuration)
+        api_client.user_agent = f"truefoundry-cli/{__version__}"
+        return api_client
+
+    def _ensure_session(self):
+        if self.session is None:
+            raise MlFoundryException(
+                relogin_error_message(
+                    "No active session found. Perhaps you are not logged in?",
+                )
+            )
 
+    @property
+    def tfy_host(self) -> str:
+        self._ensure_session()
+        assert self.session is not None
+        return self.session.tfy_host
 
-def get_active_session() -> Optional[Session]:
-    return ACTIVE_SESSION
+    @property
+    def access_token(self) -> str:
+        self._ensure_session()
+        assert self.session is not None
+        return self.session.access_token
 
 
 def _get_api_client(
-    session: Optional[Session] = None,
+    session: Optional[MLFoundrySession] = None,
     allow_anonymous: bool = False,
-) -> ApiClient:
-    from truefoundry.version import __version__
+) -> MLFoundryServerApiClient:
+    global ACTIVE_SESSION
 
-    session = session or get_active_session()
+    session = session or ACTIVE_SESSION
     if session is None:
         if allow_anonymous:
-            return ApiClient()
+            return MLFoundryServerApiClient(session=None)
         else:
             raise MlFoundryException(
                 relogin_error_message(
                     "No active session found. Perhaps you are not logged in?",
                 )
             )
-
-    creds = session.get_host_creds()
-    configuration = Configuration(
-        host=creds.host.rstrip("/"),
-        access_token=creds.token,
-    )
-    configuration.retries = urllib3_retry(retries=2)
-    api_client = ApiClient(configuration=configuration)
-    api_client.user_agent = f"truefoundry-cli/{__version__}"
-    return api_client
-
-
-def init_session() -> Session:
-    with SESSION_LOCK:
-        final_cred_provider = None
-        for cred_provider in [EnvCredentialProvider, FileCredentialProvider]:
-            if cred_provider.can_provide():
-                final_cred_provider = cred_provider()
-                break
-        if final_cred_provider is None:
-            raise MlFoundryException(
-                relogin_error_message(
-                    "No active session found. Perhaps you are not logged in?",
-                )
-            )
-        new_session = Session(cred_provider=final_cred_provider)
-
-        global ACTIVE_SESSION
-        if ACTIVE_SESSION and ACTIVE_SESSION == new_session:
-            return ACTIVE_SESSION
-
-        ACTIVE_RUNS.close_active_runs()
-
-        if ACTIVE_SESSION:
-            ACTIVE_SESSION.close()
-        ACTIVE_SESSION = new_session
-
-        logger.info(
-            "Logged in to %r as %r (%s)",
-            ACTIVE_SESSION.tracking_uri,
-            ACTIVE_SESSION.user_info.user_id,
-            ACTIVE_SESSION.user_info.email or ACTIVE_SESSION.user_info.user_type.value,
-        )
-        return ACTIVE_SESSION
+    return MLFoundryServerApiClient.from_session(session)