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

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.3rc5.dist-info}/METADATA

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

{truefoundry-0.5.3rc4.dist-info → truefoundry-0.5.3rc5.dist-info}/RECORD

@@ -1,11 +1,11 @@
-truefoundry/__init__.py,sha256=jynqfY0pDmcaXlxewDnMQC5CYAIW9JFjptBrw7y76rY,81
+truefoundry/__init__.py,sha256=Go9M0PvV9kYX8dhldJzpH0RAYB_Pc1kVOxfNlsps89c,321
 truefoundry/autodeploy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/autodeploy/agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/autodeploy/agents/base.py,sha256=lYw6bKhfCWff5BbtXXn2k4MjMbRIUkOwPBjRT70iXE8,6420
 truefoundry/autodeploy/agents/developer.py,sha256=tO9vmgGjoaqK4rX-DFCyMrViFC2Oqofd89WoXqsEm-A,4142
 truefoundry/autodeploy/agents/project_identifier.py,sha256=KlF7FwrGHDCjF3y3-YaCDxZ8kE34s6XPBagXhG2U1aM,4520
 truefoundry/autodeploy/agents/tester.py,sha256=lqvVBuVcFRvGZnSLAUF2r-1-Yti20o_cbybU19IlSXo,2418
-truefoundry/autodeploy/cli.py,sha256=1VnDAP9ylM6l00v0Ko5KL6_Dr7bL_voeAPdbEvWvU3Y,14215
+truefoundry/autodeploy/cli.py,sha256=9ZxKu_MGIpraMzaW4ZyuQZhlKIQYE3biBrBV4S1h6Fo,14167
 truefoundry/autodeploy/constants.py,sha256=vTh2nA7cjqghqbW2rNh3FbtcIk2scdAWZuuQCmVBO80,1273
 truefoundry/autodeploy/exception.py,sha256=fa_ZyTDUKiMKG2Uayynk1yWcEMsuVluqk2GtV4tfTPU,158
 truefoundry/autodeploy/logger.py,sha256=tkV2UKcOTFl5nz0cn4eRbzxF-2CZd8b7MK9vnhaflYw,325
@@ -24,23 +24,26 @@ truefoundry/autodeploy/utils/client.py,sha256=PvbSkfgAjAogGjisinqmh4mP4svowxAC0I
 truefoundry/autodeploy/utils/diff.py,sha256=Ef8Y-VffDKel_-q-GxRam6gqiv8qTLMcqVg6iifXfcA,5358
 truefoundry/autodeploy/utils/pydantic_compat.py,sha256=hEAUy5kLjhPdzw7yGZ2iXGMXbbMVXVlGzIofmyHafXQ,412
 truefoundry/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/cli/__main__.py,sha256=mYgMZRvI3sBP3_6DA0WwtFUZQTV3r8gJpMFvMJb-OOE,3692
+truefoundry/cli/__main__.py,sha256=0nEwE0UkFagCO5ctGcqEThdmBT2fN8MVcr0ERItnsNQ,3483
 truefoundry/cli/commands/pat.py,sha256=0rm95ny1b3KfeBJSRWh1M4lR46ERnp96ovUZKu-P5Ss,602
 truefoundry/cli/config.py,sha256=tf8w4UfVzcC6eYkENvuuCPYt_V3sqVpO1bclORV9tAk,206
 truefoundry/cli/console.py,sha256=9-dMy4YPisCJQziRKTg8Qa0UJnOGl1soiUnJjsnLDvE,242
 truefoundry/cli/const.py,sha256=dVHPo1uAiDSSMXwXoT2mR5kNQjExT98QNVRz98Hz_Ts,510
 truefoundry/cli/display_util.py,sha256=vypx4FQImmDwl2QAj7GpjwlpvG3VFhqoNPvCkBzGqNg,3120
-truefoundry/cli/util.py,sha256=yA8sie7TBTOu-5QFf2zEFWYEKYcBCaV6o1gtctj1dtI,3390
+truefoundry/cli/util.py,sha256=gLSKts2js8ThKXV19p2ENv5seSf12PrUemUsveNmcGY,3717
 truefoundry/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/common/auth_service_client.py,sha256=RRiGUqITxeVYwKZLc923zJP-61UAvFtVlMaG2HJBvXc,7940
-truefoundry/common/constants.py,sha256=VbwTn5nzdUnG8u50_dnq_zWwcLf5PXFinhhSogsOAys,2667
+truefoundry/common/auth_service_client.py,sha256=N3YxKlx63r6cPZqbgb2lqBOPI69ShB7D7RCIq4FSCjc,7949
+truefoundry/common/constants.py,sha256=WQjylPv7xTVYURv4zicJUHYL73Gv3c3i6Iw4oy9Lhg4,2744
 truefoundry/common/credential_file_manager.py,sha256=1yEk1Zm2xS4G0VDFwKSZ4w0VUrcPWQ1nJnoBaz9xyKA,4251
-truefoundry/common/credential_provider.py,sha256=Aht7hFLsnyRgMR34dRbzln7dor0WYSeA8ej8ApNmnKM,4148
+truefoundry/common/credential_provider.py,sha256=R666i1ADKb2ldwEVh9lnaFcPaa4q3vAscIX3MX1UiFA,4097
 truefoundry/common/entities.py,sha256=8O-EGPk4PKqnyoFMKUTxISCU19rz0KBnfRDJU695DhY,3797
-truefoundry/common/exceptions.py,sha256=ePpiQ_zmWe4e94gOgeMiyP_AZnKwjEBfyXsB5ScGYcI,329
-truefoundry/common/request_utils.py,sha256=5xw4YGUcMf71Ncal3OfFCa-PoWDIvG3hYGCDa4Da4OI,2854
-truefoundry/common/servicefoundry_client.py,sha256=2fxmgCM-ckFHpnm6n_mL-5Z8RWN_q-dYVvFC29bkYSg,3120
-truefoundry/common/utils.py,sha256=H87mfvqSE9OSRRdX-G9m3_JvWLVtkNuLmsFsvJHMczc,6098
+truefoundry/common/exceptions.py,sha256=jkU0N7hV_P-EhXeud4I5vuB9glXXZSWPf8LcH04mSbw,459
+truefoundry/common/request_utils.py,sha256=e9qrAQ1MutU7JALDKcucmNd0KQEVBqgW3yx0w1zeHIU,5700
+truefoundry/common/servicefoundry_client.py,sha256=295ZXOEM46lXZP59O04XFTDCsZLtz3MssXRDMk2x9hg,3169
+truefoundry/common/session.py,sha256=BiTJ2WEzLkTvzpqy67xX3fXgiKH_ltzzmQdFG7xUtwI,2904
+truefoundry/common/storage_provider_utils.py,sha256=yURhMw8k0FLFvaviRHDiifhvc6GnuQwGMC9Qd2uM440,10934
+truefoundry/common/utils.py,sha256=sZyg-vpShVXvnrwnINrdsyWJR3FHMUBiOB1CslYdQTo,6072
+truefoundry/common/warnings.py,sha256=rs6BHwk7imQYedo07iwh3TWEOywAR3Lqhj0AY4khByg,504
 truefoundry/deploy/__init__.py,sha256=p14_yIJw2bcaKW3c1DTaoM7UwfHvt2T_MfeFlZIz1g4,2359
 truefoundry/deploy/auto_gen/models.py,sha256=3_GTdaYE95piMzhbM4SvMSQhP4UGWDGpJnLxERjYHGw,82555
 truefoundry/deploy/builder/__init__.py,sha256=1qjHMNBE1poRCZW0WrG46dFM1f1IlivD5352qzsioMU,4953
@@ -49,7 +52,7 @@ truefoundry/deploy/builder/builders/dockerfile.py,sha256=AXXTziCkaqIhuM_bwyD1vT1
 truefoundry/deploy/builder/builders/tfy_notebook_buildpack/__init__.py,sha256=x_GwRFKz-Kb4-ZlxOFjBlr0mTgUDe_hVeG4dsIbHo8c,1796
 truefoundry/deploy/builder/builders/tfy_notebook_buildpack/dockerfile_template.py,sha256=rQgdvKmAT9HArVW4TAG5yd2QTKRs3S5LJ9RQbc_EkHE,2518
 truefoundry/deploy/builder/builders/tfy_python_buildpack/__init__.py,sha256=9r1PYahn-HfzpMth6NkvJoycmmHQpSl0vIVZxWF12xI,1864
-truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py,sha256=lQxEl7VNv9pCSTcaGe-ARicd5481RmRtVX5U6G18ptY,9429
+truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py,sha256=X9mKH7SGZLGZGvY0S2aPB3sgWdSR9TY4eVXtqz_idEA,8973
 truefoundry/deploy/builder/constants.py,sha256=amUkHoHvVKzGv0v_knfiioRuKiJM0V0xW0diERgWiI0,508
 truefoundry/deploy/builder/docker_service.py,sha256=OI8efqK0Gnoii8bcHihpA2StwHVzsMREfBk7NvMR4hY,3950
 truefoundry/deploy/builder/utils.py,sha256=D68-bqM0NQx-Elg-56mtkENyVyg9faZ9tgTmBuo1Sjs,1076
@@ -60,7 +63,7 @@ truefoundry/deploy/cli/commands/build_command.py,sha256=zJBywMatbpUlXx5O2aqpEVmP
 truefoundry/deploy/cli/commands/build_logs_command.py,sha256=DE_yprQQ6PrhOQo3xsYxGWQrXSopzSYDkt3hWB4n728,2790
 truefoundry/deploy/cli/commands/create_command.py,sha256=rCajvQvAfZU10nDZOYpRACbAATH1zj52ihTWrhnLLUc,1829
 truefoundry/deploy/cli/commands/delete_command.py,sha256=JMibxvt1rsZ8sODa0jI1e6vavjN8Pv04ttVFURSTKTw,2362
-truefoundry/deploy/cli/commands/deploy_command.py,sha256=B5taHuy2qRJEwco78kroMEAj6uf2_UMYatqb4Rn2-ko,3986
+truefoundry/deploy/cli/commands/deploy_command.py,sha256=YW4OG-5ZD1BZn82ckrNInNRHWxVOwaUjHhRTgfCEUCU,3890
 truefoundry/deploy/cli/commands/deploy_init_command.py,sha256=g-jBfrEmhZ0TDWsyqPDn4K6q33EqJSGmBTt1eMYig-w,600
 truefoundry/deploy/cli/commands/get_command.py,sha256=HZQGWib-qrS0RdzctRDwDLfMGu0wqWkAM1u26cTZlsc,5944
 truefoundry/deploy/cli/commands/list_command.py,sha256=zKu_JWY35eMIzgSNFYtmwi2uezZ4k-8yk3C1Vqsvshc,4470
@@ -79,40 +82,38 @@ truefoundry/deploy/io/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3h
 truefoundry/deploy/io/output_callback.py,sha256=V2YwUFec4G4a67lM4r-x_64AqdOVNo_9sTdfQWLlvi0,604
 truefoundry/deploy/io/rich_output_callback.py,sha256=TJLiRD-EnFVwgcepxR7WN0koKqW1X2DevETPhNPi_nU,829
 truefoundry/deploy/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/deploy/lib/auth/servicefoundry_session.py,sha256=5TCYPunAygtn5mb0mp_VcWKEalKMKPbyWMWer-Vty2g,1916
 truefoundry/deploy/lib/clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/deploy/lib/clients/servicefoundry_client.py,sha256=zx9c-TjQRzMDKGCa-wDvXEX_8FcWH1QlVif2SjxgLxs,27136
+truefoundry/deploy/lib/clients/servicefoundry_client.py,sha256=l4KaHJmmTC3pP5hbnHuPhDH-MPs7KQBtTtfmmfyyBDI,27020
 truefoundry/deploy/lib/const.py,sha256=Wg0GDnfFu-g1fJr4lU80NH2ULw0R0dYjV7LnW-PbOeM,173
 truefoundry/deploy/lib/dao/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/deploy/lib/dao/application.py,sha256=xUqF34abde71YZo9gkR-WvVJi1C6nBeGBNw4kL5FdIo,9226
+truefoundry/deploy/lib/dao/application.py,sha256=oMszpueXPUfTUuN_XdKwoRjQyqAgWHhZ-10cbprCVdM,9226
 truefoundry/deploy/lib/dao/apply.py,sha256=Jadfa3DcbIRz-Bb4Tlixed179HLDnmDXOlrvSxfEgQ0,3021
 truefoundry/deploy/lib/dao/version.py,sha256=AtdW_4O1DPUKdfv2qy6iUJsZ_95vM6z0AqeEy3WDKs8,1130
-truefoundry/deploy/lib/dao/workspace.py,sha256=KZA8QrlOTecFLAi4Xt0qJ0hNwQL76UYsTKJgZbsA_5s,2349
+truefoundry/deploy/lib/dao/workspace.py,sha256=6YvfCgWDzAULI3Q6JswyZmP1CwJ5rM-ANsIFkbQia0Q,2349
 truefoundry/deploy/lib/logs_utils.py,sha256=SQxRv3jDDmgHdOUMhlMaAPGYskybnBUMpst7QU_i_sc,1469
 truefoundry/deploy/lib/messages.py,sha256=nhp0bCYf_XpUM68hTq5lBY-__vtEyV2uP7NgnJXJ_Vg,925
 truefoundry/deploy/lib/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/deploy/lib/model/entity.py,sha256=3YRdyvMYI1aUXUu2tB5b5lZWwpxKdAhjpT79gzTCZGo,8373
-truefoundry/deploy/lib/session.py,sha256=Vg6rCA315T0yS0xG4ayJ84Ia_9ZfibH8utOSwPBMAmw,4953
+truefoundry/deploy/lib/session.py,sha256=rNMkdzZejrxq5gTYTqJ-nGm869JzdlhHzaJQvINTQ4U,4953
 truefoundry/deploy/lib/util.py,sha256=J7r8San2wKo48A7-BlH2-OKTlBO67zlPjLEhMsL8os0,1059
 truefoundry/deploy/lib/win32.py,sha256=1RcvPTdlOAJ48rt8rCbE2Ufha2ztRqBAE9dueNXArrY,5009
 truefoundry/deploy/python_deploy_codegen.py,sha256=qJHH1BJQII9e6PhkcRFYiE_3De7_VMMm8nM4AX5Eq1o,6513
 truefoundry/deploy/v2/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/deploy/v2/lib/__init__.py,sha256=WEiVMZXOVljzEE3tpGJil14liIn_PCDoACJ6b3tZ6sI,188
-truefoundry/deploy/v2/lib/deploy.py,sha256=HIcY3SzQ5lWl7avuuKi3J0Z-PBES6Sf4hgMK-m6_53U,11990
-truefoundry/deploy/v2/lib/deploy_workflow.py,sha256=hgAhd1EGwFLz319Vs-WNXHDJmbKjdgkGPzDnBD1Up1k,12579
+truefoundry/deploy/v2/lib/deploy.py,sha256=oK5yTlrhzKNTGZckwuBzu8Q_pxQYgBVcdA-KB-SbzCI,11990
+truefoundry/deploy/v2/lib/deploy_workflow.py,sha256=6Y4PNhd5PCc_4TjpcIZPkr6FCm3BNIQ_SsJTvpUu7a0,12579
 truefoundry/deploy/v2/lib/deployable_patched_models.py,sha256=MROgMxhn9hDEAKwJSWl3iz12tUVvRKzEtqF2QUT6dAk,3343
 truefoundry/deploy/v2/lib/models.py,sha256=pSolLMTArDuYpeNsmeeS5DWliloN_iCDfZSpRllMHUg,1120
-truefoundry/deploy/v2/lib/patched_models.py,sha256=NTU8J_CwdvEuF9zNXwFyN3suOp_196Wrm75DDy5qcXE,14184
-truefoundry/deploy/v2/lib/source.py,sha256=VHKuFREiixUP40D3Mrz-TA70spu1M0RbCzl--qwlFaY,9263
+truefoundry/deploy/v2/lib/patched_models.py,sha256=YDM75pdUPdMFgH2G8ZNbztalLIiMcHVHBHt7cLWmiDg,16050
+truefoundry/deploy/v2/lib/source.py,sha256=tzMSOTrTH3dJzICjo8gI8iY6n0-ckktP3Dh7bVK-nTI,9344
 truefoundry/gateway/__init__.py,sha256=HvAMuk60-fuZcMmdOtozHrUaH_JcbEHxpMSlS9ci4ls,41
-truefoundry/gateway/cli/cli.py,sha256=SPEHNP210QtK7G97KOOP6t6QBxgCq2YUezgx65QZToM,1715
-truefoundry/gateway/lib/client.py,sha256=6E9gTl_rOCwcrx5fJkUzbZqacJmgXtSXd1mcmDk9bCA,1717
-truefoundry/gateway/lib/entities.py,sha256=DoYS9QAKjiVao-CH9j9LVUPlVFdeHqJ3xC7-pYhmAyM,918
-truefoundry/gateway/lib/models.py,sha256=0JUsLt2oNf0bBea_03ndDmUkZhNqvMmixh8zgxaH2lk,2600
+truefoundry/gateway/cli/cli.py,sha256=9raS6Nq_CNzagEPpbKICwUZ01zMoqmZINKMB9m_gJRE,908
+truefoundry/gateway/lib/entities.py,sha256=1P9YjrMoHbg58M57Pq-HlsT1ISHRyyX4JxK-o4of6cY,695
+truefoundry/gateway/lib/models.py,sha256=TXOur51tuYkq4Xil-icNVaVH-7xZbdOUyqS8N9295dg,876
 truefoundry/logger.py,sha256=u-YCNjg5HBwE70uQcpjIG64Ghos-K2ulTWaxC03BSj4,714
 truefoundry/ml/__init__.py,sha256=ssUEIs8BixPWxynKoeSh-dkRl6AtLXG0PBGYnUR5Az8,2217
 truefoundry/ml/artifact/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/ml/artifact/truefoundry_artifact_repo.py,sha256=FksxhUpRHb9pgWZmAB16DhXqkAL6UIAPA1S3RJUApQU,46201
+truefoundry/ml/artifact/truefoundry_artifact_repo.py,sha256=myDpr3RZzEkd_UG9C3rZ4UKEFxXopO5HMVJNpNDiuAo,37737
 truefoundry/ml/autogen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/ml/autogen/client/__init__.py,sha256=42Y8v6C1oDlngKCaDODi9k3q-9MPMiqO8jGe39y5QJM,20097
 truefoundry/ml/autogen/client/api/__init__.py,sha256=NyMBxBmIzB1o5LzZZwz9LiydHB3-hPqc_sevsnY6Jrw,746
@@ -326,22 +327,20 @@ truefoundry/ml/cli/commands/download.py,sha256=N9MhsEQ3U24v_OmnMZT8Q4SoAi38Sm7a2
 truefoundry/ml/cli/commands/model_init.py,sha256=9fHCernmQswjEhREzdrnKxwCCWkgmFrpL29H52Sd1gQ,2663
 truefoundry/ml/cli/utils.py,sha256=j6_mZ4Spn114mz3P4QQ8jx0tmorXIuyQnHXVUSDvZi4,1035
 truefoundry/ml/clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/ml/clients/entities.py,sha256=sNP4DnAVdQoMfc06s0r3VTzKHTo7jmxAOuTlQOVmsMs,151
-truefoundry/ml/clients/servicefoundry_client.py,sha256=nmKrWBd7P0VEDGjPHZ8_scX3Xo9tv3wZSWlTIBwGptU,1597
-truefoundry/ml/clients/utils.py,sha256=c0LdC8moejs-Zm30hu1sCqifLEmqhdq4SfZ_m0nUIDk,4402
+truefoundry/ml/clients/servicefoundry_client.py,sha256=Rx-tjiDxNwUgOCDITjtL0vmZcQgJjh_tUdReE23TKIk,2363
 truefoundry/ml/constants.py,sha256=vDq72d4C9FSWqr9MMdjgTF4TuyNFApvo_6RVsSeAjB4,2837
 truefoundry/ml/entities.py,sha256=si5GAqZsWzKu5MPrU4Hk6se7bebHOYhTiNw69ai-Uk8,1485
 truefoundry/ml/enums.py,sha256=arqDkF8duU_oVLFeYMhcfWYbF6Nq5mmjwupJMIheyXM,1790
-truefoundry/ml/exceptions.py,sha256=8aJm2NYtAWWsRLu4MbzaoOqHsQZ6RjOFwBWQWqb6qrc,244
+truefoundry/ml/exceptions.py,sha256=QpDJSWmF7dIsByS0qOQbQZ_jytdNTzkHDDO3BxhTSo0,332
 truefoundry/ml/git_info.py,sha256=jvAVm9ilqivnGq8qJdUvYdd8Siv0PLtqurB-PXsS5ho,2023
 truefoundry/ml/internal_namespace.py,sha256=QcqMHp6-C2im2H_02hlhi01EIcr1HhNaZprszs13EMU,1790
 truefoundry/ml/log_types/__init__.py,sha256=g4u4D4Jaj0aBK5GtrLV88-qThKZR9pSZ17vFEkN-LmM,125
-truefoundry/ml/log_types/artifacts/artifact.py,sha256=VC0KbcHty1fblybUvGVcfDGFcINfLgT2LByyvPcS1Ro,21565
+truefoundry/ml/log_types/artifacts/artifact.py,sha256=qz657N5WosBNBqqEtZ9hWatvIoa9_0MipudYH9r3Bq0,21657
 truefoundry/ml/log_types/artifacts/constants.py,sha256=qKxQ5mMvJE4j83BvGW3qNTKunxCiBg_EEjTdgbgJtyE,1036
 truefoundry/ml/log_types/artifacts/dataset.py,sha256=a4dxd2EN8p7Ci-cLGGiDOboN3t0395_XhWE1dmTw1Q4,13112
 truefoundry/ml/log_types/artifacts/general_artifact.py,sha256=B4XErLr-m6RmQWtxMTu3wlFRFcqSwPYp6J0OL4Ng6L0,3179
-truefoundry/ml/log_types/artifacts/model.py,sha256=p9cB3_0MSZkVyq8G2pyNWl1F3qICTazQlq6VCSzOV3s,24194
-truefoundry/ml/log_types/artifacts/utils.py,sha256=8mmahc1b6wvrA_E795h9ynRfhYRIdci7LRGN8U6ySm8,7516
+truefoundry/ml/log_types/artifacts/model.py,sha256=J-_v_ertSKQPThx6I9ivLUlCJUHPZJ71X8V0VgaIBAU,24319
+truefoundry/ml/log_types/artifacts/utils.py,sha256=MuPk6by8jIt7LepB1fqUC8emQBSxrTGJmIel7rbMBAc,7533
 truefoundry/ml/log_types/image/__init__.py,sha256=fcOq8yQnNj1rkLcPeIjLXBpdA1WIeiPsXOlAAvMxx7M,76
 truefoundry/ml/log_types/image/constants.py,sha256=wLtGEOA4T5fZHSlOXPuNDLX3lpbCtwlvGKPFk_1fah0,255
 truefoundry/ml/log_types/image/image.py,sha256=qQnAVgErAq4Jn6wXFFpaveOd52zcjUuomUCqNRxO2io,12478
@@ -351,11 +350,11 @@ truefoundry/ml/log_types/plot.py,sha256=HuYvvRA5r8V0xAIuuqMME2IHb9d3SfGHUiuEkOP3
 truefoundry/ml/log_types/pydantic_base.py,sha256=eBlw_AEyAz4iJKDP4zgJOCFWcldwQqpf7FADW1jzIQY,272
 truefoundry/ml/log_types/utils.py,sha256=xjJ21jdPScvFmw3TbVh5NCzbzJwaqiXJyiiT4xxX1EI,335
 truefoundry/ml/logger.py,sha256=VT-BF3BnBYTWVq87O58F0c8uXMu94gYzsiFlGY3_7Ao,458
-truefoundry/ml/mlfoundry_api.py,sha256=avuGgjbRnQXk774Vd3Dx7n99bbrOkL9pob9DI-vNvWE,62368
-truefoundry/ml/mlfoundry_run.py,sha256=scHjdjojRGbE4ZZlByZTFrcBTamRbm_JGidaSX9Gi9o,44601
-truefoundry/ml/model_framework.py,sha256=GAGu6bW8L4f8uvukOITB0bJOm1BSWzxC1MJxUGmpOqQ,18882
+truefoundry/ml/mlfoundry_api.py,sha256=ofIbFimPI7u301KFxdbmWbJT0AIinbRMQyCwX7hbPzs,61543
+truefoundry/ml/mlfoundry_run.py,sha256=M-0FkZYDxwv5EZUDtVqiV8JxC_3BWR80N7Kp4Yj7bPY,44291
+truefoundry/ml/model_framework.py,sha256=JGmzdG7o5P6CJr5EYTUOmuly53FfhpoNVqKrhfijAVg,18972
 truefoundry/ml/run_utils.py,sha256=0W208wSLUrbdfk2pjNcZlkUi9bNxG2JORqoe-5rVqHI,2423
-truefoundry/ml/session.py,sha256=F83GTC5WwGBjnJ69Ct8MqMnlutYc56JCc6YhEY1Wl-A,5394
+truefoundry/ml/session.py,sha256=eQo1zG5nd0-8Ok68uEBCQDPAWYeJkpzqHak-DQX1ENU,4553
 truefoundry/ml/validation_utils.py,sha256=J5atNhcJLvKj64ralSV9Y5Fv1Rt4SE237ICdP9-7sP4,12149
 truefoundry/pydantic_v1.py,sha256=jSuhGtz0Mbk1qYu8jJ1AcnIDK4oxUsdhALc4spqstmM,345
 truefoundry/version.py,sha256=bqiT4Q-VWrTC6P4qfK43mez-Ppf-smWfrl6DcwV7mrw,137
@@ -370,11 +369,11 @@ truefoundry/workflow/map_task.py,sha256=2m3qGXQ90k9LdS45q8dqCCECc3qr8t2m_LMCVd1m
 truefoundry/workflow/python_task.py,sha256=SRXRLC4vdBqGjhkwuaY39LEWN6iPCpJAuW17URRdWTY,1128
 truefoundry/workflow/remote_filesystem/__init__.py,sha256=LQ95ViEjJ7Ts4JcCGOxMPs7NZmQdZ4bTiq6qXtsjUhE,206
 truefoundry/workflow/remote_filesystem/logger.py,sha256=em2l7D6sw7xTLDP0kQSLpgfRRCLpN14Qw85TN7ujQcE,1022
-truefoundry/workflow/remote_filesystem/tfy_signed_url_client.py,sha256=5mBCIc-ON7VSTzdyczeADgqSX5oJgso5gq2yT2AQqTY,11085
-truefoundry/workflow/remote_filesystem/tfy_signed_url_fs.py,sha256=Hf6Dk6Fu6P7DqsK5ULgraf9DStjgigf-kjaRAMBW-RU,8680
+truefoundry/workflow/remote_filesystem/tfy_signed_url_client.py,sha256=xcT0wQmQlgzcj0nP3tJopyFSVWT1uv3nhiTIuwfXYeg,12342
+truefoundry/workflow/remote_filesystem/tfy_signed_url_fs.py,sha256=nSGPZu0Gyd_jz0KsEE-7w_BmnTD8CVF1S8cUJoxaCbc,13305
 truefoundry/workflow/task.py,sha256=ToitYiKcNzFCtOVQwz1W8sRjbR97eVS7vQBdbgUQtKg,1779
 truefoundry/workflow/workflow.py,sha256=WaTqUjhwfAXDWu4E5ehuwAxrCbDJkoAf1oWmR2E9Qy0,4575
-truefoundry-0.5.3rc4.dist-info/METADATA,sha256=qPTGfmHQVmxxL_jBbf7CJI6fDI1LkOpkSfeBXSqLRpY,2887
-truefoundry-0.5.3rc4.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-truefoundry-0.5.3rc4.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
-truefoundry-0.5.3rc4.dist-info/RECORD,,
+truefoundry-0.5.3rc5.dist-info/METADATA,sha256=pJtlnztVvgBPT53hHtnRD-mp0EzV6iTNPFKlEDvH2oc,2887
+truefoundry-0.5.3rc5.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+truefoundry-0.5.3rc5.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
+truefoundry-0.5.3rc5.dist-info/RECORD,,

truefoundry/deploy/lib/auth/servicefoundry_session.py

@@ -1,61 +0,0 @@
-from __future__ import annotations
-
-from typing import Optional
-
-from truefoundry.common.credential_provider import (
-    CredentialProvider,
-    EnvCredentialProvider,
-    FileCredentialProvider,
-)
-from truefoundry.common.entities import UserInfo
-from truefoundry.common.utils import relogin_error_message
-from truefoundry.logger import logger
-
-ACTIVE_SESSION: Optional[ServiceFoundrySession] = None
-
-
-class ServiceFoundrySession:
-    def __init__(self) -> None:
-        self._cred_provider = self._get_cred_provider()
-        self._user_info: UserInfo = self._cred_provider.token.to_user_info()
-
-        global ACTIVE_SESSION
-        if (ACTIVE_SESSION is None) or (
-            ACTIVE_SESSION
-            and ACTIVE_SESSION.base_url != self.base_url
-            and ACTIVE_SESSION.user_info != self.user_info
-        ):
-            logger.info(
-                "Logged in to %r as %r (%s)",
-                self.base_url,
-                self.user_info.user_id,
-                self.user_info.email or self.user_info.user_type.value,
-            )
-        ACTIVE_SESSION = self
-
-    @staticmethod
-    def _get_cred_provider() -> CredentialProvider:
-        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 Exception(
-                relogin_error_message(
-                    "No active session found. Perhaps you are not logged in?",
-                )
-            )
-        return final_cred_provider
-
-    @property
-    def access_token(self) -> str:
-        return self._cred_provider.token.access_token
-
-    @property
-    def base_url(self) -> str:
-        return self._cred_provider.base_url
-
-    @property
-    def user_info(self) -> UserInfo:
-        return self._user_info

truefoundry/gateway/lib/client.py

@@ -1,51 +0,0 @@
-from typing import Literal, Optional
-
-import requests
-
-from truefoundry.common.constants import GATEWAY_SERVER_RELATIVE_PATH
-from truefoundry.common.credential_provider import (
-    CredentialProvider,
-    EnvCredentialProvider,
-    FileCredentialProvider,
-)
-
-
-class GatewayServiceClient:
-    def __init__(
-        self,
-        credential_provider: Optional[CredentialProvider] = None,
-        host: Optional[str] = None,
-    ) -> None:
-        if credential_provider is None:
-            if EnvCredentialProvider.can_provide():
-                credential_provider = EnvCredentialProvider()
-            elif FileCredentialProvider.can_provide():
-                credential_provider = FileCredentialProvider()
-            else:
-                raise Exception(
-                    "No credentials found. Please login using `tfy login` or set TFY_API_KEY environment variable"
-                )
-
-        self._credential_provider = credential_provider
-        self._host = credential_provider.base_url
-
-    def _get_header(self):
-        return {
-            "Authorization": f"Bearer {self._credential_provider.token.access_token}"
-        }
-
-    @property
-    def _base_url(self) -> str:
-        return f"{self._host}/{GATEWAY_SERVER_RELATIVE_PATH}"
-
-    def generate_code(
-        self, model_id: str, inference_type: Literal["chat", "completion", "embedding"]
-    ) -> str:
-        url = f"{self._base_url}/api/inference/openai/generate-code-snippet"
-        data = {
-            "model": model_id,
-            "playground_endpoint": self._base_url,
-            "inference_type": inference_type,
-        }
-        response = requests.post(url, headers=self._get_header(), json=data)
-        return response.json()

truefoundry/ml/clients/entities.py

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