truefoundry 0.5.2__py3-none-any.whl → 0.5.3__py3-none-any.whl

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,7 +20,9 @@ from truefoundry.common.utils import (
     get_python_version_major_minor,
     list_pip_packages_installed,
 )
-from truefoundry.ml.autogen.client import (
+from truefoundry.common.warnings import TrueFoundryDeprecationWarning
+from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+    ModelVersionEnvironment,
     SklearnSerializationFormat,
     XGBoostSerializationFormat,
 )
@@ -259,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,
             )
 
@@ -344,7 +346,7 @@ def _fetch_framework_specific_pip_packages(
 
 
 def auto_update_environment_details(
-    environment: autogen_artifacts.ModelVersionEnvironment,
+    environment: ModelVersionEnvironment,
     framework: Optional[ModelFrameworkType],
 ):
     """

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)

truefoundry/workflow/remote_filesystem/tfy_signed_url_client.py

@@ -13,6 +13,7 @@ from truefoundry.common.constants import (
     TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN_ENV_KEY,
 )
 from truefoundry.common.request_utils import requests_retry_session
+from truefoundry.common.storage_provider_utils import MultiPartUploadStorageProvider
 from truefoundry.pydantic_v1 import BaseModel, Field
 from truefoundry.workflow.remote_filesystem.logger import log_time, logger
 
@@ -20,6 +21,9 @@ LOG_PREFIX = "[tfy][fs]"
 DEFAULT_TTL = ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL
 MAX_TIMEOUT = ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_MAX_TIMEOUT
 REQUEST_TIMEOUT = ENV_VARS.TFY_INTERNAL_SIGNED_URL_REQUEST_TIMEOUT
+MULTIPART_UPLOAD_FINALIZE_SIGNED_URL_TIMEOUT = (
+    ENV_VARS.TFY_INTERNAL_MULTIPART_UPLOAD_FINALIZE_SIGNED_URL_TIMEOUT
+)
 
 
 class SignedURLAPIResponseDto(BaseModel):
@@ -36,6 +40,18 @@ class SignedURLExistsAPIResponseDto(BaseModel):
     exists: bool
 
 
+class PartSignedUrl(BaseModel):
+    partNumber: int
+    signedUrl: str
+
+
+class SignedURLMultipartUploadAPIResponseDto(BaseModel):
+    uploadId: str
+    partSignedUrls: List[PartSignedUrl]
+    finalizeSignedUrl: str
+    storageProvider: MultiPartUploadStorageProvider
+
+
 class FileInfo(BaseModel):
     path: str
     is_directory: bool = Field(..., alias="isDirectory")
@@ -56,6 +72,7 @@ class SignedURLServerEndpoint(str, Enum):
     EXISTS = "/v1/exists"
     IS_DIRECTORY = "/v1/is-dir"
     LIST_FILES = "/v1/list-files"
+    CREATE_MUTLIPART_UPLOAD = "/v1/multipart-upload"
 
 
 class SignedURLClient:
@@ -98,9 +115,9 @@ class SignedURLClient:
         self,
         endpoint: str,
         method: str = "GET",
-        payload: Optional[Dict] = None,
-        headers: Optional[Dict] = None,
-    ) -> Dict:
+        payload: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
         """Internal method to handle requests to the signed URL server."""
         url = urljoin(self.base_url, endpoint)
         try:
@@ -116,9 +133,9 @@ class SignedURLClient:
     def _make_server_api_call(
         self,
         endpoint: SignedURLServerEndpoint,
-        params: Optional[Dict] = None,
-        headers: Optional[Dict] = None,
-    ) -> Dict:
+        params: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
         """Get a signed URL for the specified operation and URI."""
         query_string = urlencode(params or {})
         endpoint_with_params = f"{endpoint.value}?{query_string}"
@@ -131,7 +148,7 @@ class SignedURLClient:
         self,
         signed_url: str,
         data: Union[bytes, io.BufferedReader],
-        headers: Optional[Dict] = None,
+        headers: Optional[Dict[str, Any]] = None,
     ) -> None:
         """
         Upload data to the specified storage path using a signed URL.
@@ -190,12 +207,28 @@ class SignedURLClient:
             )
         return storage_uri
 
+    def create_multipart_upload(
+        self, storage_uri: str, num_parts: int
+    ) -> SignedURLMultipartUploadAPIResponseDto:
+        response = self._make_server_api_call(
+            endpoint=SignedURLServerEndpoint.CREATE_MUTLIPART_UPLOAD,
+            params={
+                "path": storage_uri,
+                "numParts": num_parts,
+                "partExpiryInSeconds": self.ttl,
+                "finalizationExpiryInSeconds": MULTIPART_UPLOAD_FINALIZE_SIGNED_URL_TIMEOUT,
+            },
+            headers=self.signed_url_server_headers,
+        )
+
+        return SignedURLMultipartUploadAPIResponseDto.parse_obj(response)
+
     @log_time(prefix=LOG_PREFIX)
     def _download_file(
         self,
         signed_url: str,
         local_path: Optional[str] = None,
-        headers: Optional[Dict] = None,
+        headers: Optional[Dict[str, Any]] = None,
     ) -> Optional[bytes]:
         """Common method to download a file using a signed URL."""
         try:
@@ -237,7 +270,7 @@ class SignedURLClient:
         return local_path
 
     @log_time(prefix=LOG_PREFIX)
-    def download_to_bytes(self, storage_uri: str) -> bytes:
+    def download_to_bytes(self, storage_uri: str) -> Optional[bytes]:
         """Download a file from the specified storage path and return it as bytes."""
         response = self._make_server_api_call(
             endpoint=SignedURLServerEndpoint.READ,