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

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

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

{truefoundry-0.5.3rc3.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=Pxksh1MMjz9WAKV6wCcq1wu9uawzZGMe9RU9FrLOZbQ,2626
+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,31 +82,30 @@ 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=9raS6Nq_CNzagEPpbKICwUZ01zMoqmZINKMB9m_gJRE,908
 truefoundry/gateway/lib/entities.py,sha256=1P9YjrMoHbg58M57Pq-HlsT1ISHRyyX4JxK-o4of6cY,695
@@ -111,7 +113,7 @@ truefoundry/gateway/lib/models.py,sha256=TXOur51tuYkq4Xil-icNVaVH-7xZbdOUyqS8N92
 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
@@ -325,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
@@ -350,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
@@ -369,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.3rc3.dist-info/METADATA,sha256=zuoxLW-MzlfmWPLCP8J2NF8eydbOKPu_LNhKouhQi_0,2887
-truefoundry-0.5.3rc3.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-truefoundry-0.5.3rc3.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
-truefoundry-0.5.3rc3.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/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

truefoundry/ml/clients/utils.py

@@ -1,122 +0,0 @@
-import json
-from contextlib import contextmanager
-from typing import Any, Optional
-from urllib.parse import urljoin
-
-import requests
-
-from truefoundry.common.request_utils import requests_retry_session
-from truefoundry.ml.clients.entities import HostCreds
-from truefoundry.ml.exceptions import MlFoundryException
-
-
-# TODO: This will be moved later to truefoundry.common.request_utils
-def _make_url(host: str, endpoint: str) -> str:
-    if endpoint.startswith("/"):
-        raise ValueError("`endpoint` must not start with a leading slash (/)")
-    return urljoin(host, endpoint)
-
-
-def _http_request(
-    *, method: str, url: str, token: Optional[str] = None, session=requests, **kwargs
-) -> requests.Response:
-    headers = kwargs.pop("headers", {}) or {}
-    if token is not None:
-        headers["Authorization"] = f"Bearer {token}"
-    return session.request(method=method, url=url, headers=headers, **kwargs)
-
-
-def http_request(
-    *, method: str, host_creds: HostCreds, endpoint: str, session=requests, **kwargs
-) -> requests.Response:
-    url = _make_url(host=host_creds.host, endpoint=endpoint)
-    return _http_request(
-        method=method, url=url, token=host_creds.token, session=session, **kwargs
-    )
-
-
-def http_request_safe(
-    *, method: str, host_creds: HostCreds, endpoint: str, session=requests, **kwargs
-) -> Any:
-    url = _make_url(host=host_creds.host, endpoint=endpoint)
-    try:
-        response = _http_request(
-            method=method, url=url, token=host_creds.token, session=session, **kwargs
-        )
-        response.raise_for_status()
-        try:
-            return response.json()
-        except json.JSONDecodeError as je:
-            raise MlFoundryException(
-                f"Failed to parse response as json. Response: {response.text}"
-            ) from je
-    except requests.exceptions.ConnectionError as ce:
-        raise MlFoundryException("Failed to connect to TrueFoundry") from ce
-    except requests.exceptions.Timeout as te:
-        raise MlFoundryException(f"Request to {url} timed out") from te
-    except requests.exceptions.HTTPError as he:
-        raise MlFoundryException(
-            f"Request to {url} with status code {he.response.status_code}. Response: {he.response.text}",
-            status_code=he.response.status_code,
-        ) from he
-    except Exception as e:
-        raise MlFoundryException(
-            f"Request to {url} failed with an unknown error"
-        ) from e
-
-
-@contextmanager
-def cloud_storage_http_request(
-    *,
-    method,
-    url,
-    session=None,
-    timeout=None,
-    **kwargs,
-):
-    """
-    Performs an HTTP PUT/GET request using Python's `requests` module with automatic retry.
-    """
-    session = session or requests_retry_session(retries=5, backoff_factor=0.5)
-    kwargs["headers"] = kwargs.get("headers", {}) or {}
-    if "blob.core.windows.net" in url:
-        kwargs["headers"].update({"x-ms-blob-type": "BlockBlob"})
-    if method.lower() not in ("put", "get"):
-        raise ValueError(f"Illegal http method: {method}")
-    try:
-        yield _http_request(
-            method=method, url=url, session=session, timeout=timeout, **kwargs
-        )
-    except Exception as e:
-        raise MlFoundryException(
-            f"API request failed with exception {str(e)}"
-        ) from None
-
-
-def augmented_raise_for_status(response):
-    try:
-        response.raise_for_status()
-    except requests.exceptions.HTTPError as he:
-        raise MlFoundryException(
-            f"Request with status code {he.response.status_code}. Response: {he.response.text}",
-            status_code=he.response.status_code,
-        ) from he
-
-
-def download_file_using_http_uri(http_uri, download_path, chunk_size=100_000_000):
-    """
-    Downloads a file specified using the `http_uri` to a local `download_path`. This function
-    uses a `chunk_size` to ensure an OOM error is not raised a large file is downloaded.
-
-    Note : This function is meant to download files using presigned urls from various cloud
-            providers.
-    """
-    with cloud_storage_http_request(
-        method="get", url=http_uri, stream=True
-    ) as response:
-        augmented_raise_for_status(response)
-        with open(download_path, "wb") as output_file:
-            for chunk in response.iter_content(chunk_size=chunk_size):
-                if not chunk:
-                    break
-                output_file.write(chunk)