truefoundry 0.5.3rc4__py3-none-any.whl → 0.5.4__py3-none-any.whl

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)

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,

truefoundry/workflow/remote_filesystem/tfy_signed_url_fs.py

@@ -2,18 +2,44 @@
 # pylint: disable=W0223
 import io
 import os
+from concurrent.futures import FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait
 from pathlib import Path
-from typing import Optional
+from threading import Event
+from typing import List, Optional, Tuple
 
 from fsspec.spec import DEFAULT_CALLBACK, AbstractBufferedFile, AbstractFileSystem
 
 from truefoundry.common.constants import ENV_VARS
-from truefoundry.workflow.remote_filesystem.logger import log_time
+from truefoundry.common.storage_provider_utils import (
+    MultiPartUpload,
+    SignedURL,
+    _FileMultiPartInfo,
+    decide_file_parts,
+    s3_compatible_multipart_upload,
+)
+from truefoundry.workflow.remote_filesystem.logger import log_time, logger
 from truefoundry.workflow.remote_filesystem.tfy_signed_url_client import (
     LOG_PREFIX,
     SignedURLClient,
+    SignedURLMultipartUploadAPIResponseDto,
 )
 
+MULTIPART_SUPPORTED_PROVIDERS = ["s3"]
+
+
+def _add_file_for_upload(
+    local_file: str,
+    remote_path: str,
+    files_for_normal_upload: List[Tuple[str, str, _FileMultiPartInfo]],
+    files_for_multipart_upload: List[Tuple[str, str, _FileMultiPartInfo]],
+    multipart_upload_allowed: bool,
+):
+    multipart_info = decide_file_parts(local_file, multipart_upload_allowed)
+    if multipart_info.num_parts == 1:
+        files_for_normal_upload.append((remote_path, local_file, multipart_info))
+    else:
+        files_for_multipart_upload.append((remote_path, local_file, multipart_info))
+
 
 class SignedURLFileSystem(AbstractFileSystem):
     def __init__(
@@ -100,7 +126,13 @@ class SignedURLFileSystem(AbstractFileSystem):
         maxdepth=None,
         **kwargs,
     ):
+        files_for_normal_upload: List[Tuple[str, str, _FileMultiPartInfo]] = []
+        files_for_multipart_upload: List[Tuple[str, str, _FileMultiPartInfo]] = []
         local_path = Path(lpath)
+        multipart_upload_allowed = (
+            self.protocol in MULTIPART_SUPPORTED_PROVIDERS
+            and not ENV_VARS.TFY_ARTIFACTS_DISABLE_MULTIPART_UPLOAD
+        )
         if local_path.is_dir():
             if not recursive:
                 raise ValueError(
@@ -127,15 +159,102 @@ class SignedURLFileSystem(AbstractFileSystem):
                 for file in files:
                     local_file_path = Path(root) / file
                     remote_file_path = f"{remote_dir}/{file}"
-                    self.client.upload(
-                        file_path=str(local_file_path),
-                        storage_uri=remote_file_path,
+                    _add_file_for_upload(
+                        local_file=str(local_file_path),
+                        remote_path=str(remote_file_path),
+                        files_for_normal_upload=files_for_normal_upload,
+                        files_for_multipart_upload=files_for_multipart_upload,
+                        multipart_upload_allowed=multipart_upload_allowed,
                     )
-            return None
         else:
             if rpath.endswith("/"):
                 rpath = os.path.join(rpath, local_path.name)
-            return self.client.upload(file_path=lpath, storage_uri=rpath)
+            _add_file_for_upload(
+                local_file=str(lpath),
+                remote_path=str(rpath),
+                files_for_normal_upload=files_for_normal_upload,
+                files_for_multipart_upload=files_for_multipart_upload,
+                multipart_upload_allowed=multipart_upload_allowed,
+            )
+        return self._upload(
+            files_for_normal_upload=files_for_normal_upload,
+            files_for_multipart_upload=files_for_multipart_upload,
+        )
+
+    def _upload(
+        self,
+        files_for_normal_upload: List[Tuple[str, str, _FileMultiPartInfo]],
+        files_for_multipart_upload: List[Tuple[str, str, _FileMultiPartInfo]],
+    ):
+        abort_event = Event()
+        with ThreadPoolExecutor(
+            max_workers=ENV_VARS.TFY_ARTIFACTS_UPLOAD_MAX_WORKERS
+        ) as executor:
+            futures: List[Future] = []
+            for remote_path, local_path, _ in files_for_normal_upload:
+                futures.append(
+                    executor.submit(
+                        self.client.upload,
+                        file_path=local_path,
+                        storage_uri=remote_path,
+                    )
+                )
+
+            done, not_done = wait(futures, return_when=FIRST_EXCEPTION)
+            if len(not_done) > 0:
+                abort_event.set()
+            for future in not_done:
+                future.cancel()
+            for future in done:
+                if future.exception() is not None:
+                    raise future.exception()
+
+            for remote_path, local_path, multipart_info in files_for_multipart_upload:
+                self._multipart_upload(
+                    local_file=local_path,
+                    artifact_path=remote_path,
+                    multipart_info=multipart_info,
+                    executor=executor,
+                    abort_event=abort_event,
+                )
+
+    def _multipart_upload(
+        self,
+        local_file: str,
+        artifact_path: str,
+        multipart_info: _FileMultiPartInfo,
+        executor: ThreadPoolExecutor,
+        abort_event: Optional[Event] = None,
+    ):
+        logger.info(
+            "Uploading %s to %s using multipart upload",
+            local_file,
+            artifact_path,
+        )
+
+        multipart_upload: SignedURLMultipartUploadAPIResponseDto = (
+            self.client.create_multipart_upload(
+                storage_uri=artifact_path,
+                num_parts=multipart_info.num_parts,
+            )
+        )
+        s3_compatible_multipart_upload(
+            multipart_upload=MultiPartUpload(
+                storage_provider=multipart_upload.storageProvider,
+                part_signed_urls=[
+                    SignedURL(signed_url=url.signedUrl)
+                    for url in multipart_upload.partSignedUrls
+                ],
+                s3_compatible_upload_id=multipart_upload.uploadId,
+                finalize_signed_url=SignedURL(
+                    signed_url=multipart_upload.finalizeSignedUrl
+                ),
+            ),
+            local_file=local_file,
+            executor=executor,
+            multipart_info=multipart_info,
+            abort_event=abort_event,
+        )
 
     @log_time(prefix=LOG_PREFIX)
     def isdir(self, path):

{truefoundry-0.5.3rc4.dist-info → truefoundry-0.5.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: truefoundry
-Version: 0.5.3rc4
+Version: 0.5.4
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com