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

truefoundry/common/storage_provider_utils.py

@@ -0,0 +1,331 @@
+import logging
+import math
+import mmap
+import os
+from concurrent.futures import FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait
+from enum import Enum
+from threading import Event
+from typing import List, NamedTuple, Optional
+
+import requests
+from rich.progress import Progress
+from tqdm.utils import CallbackIOWrapper
+
+from truefoundry.common.constants import ENV_VARS
+from truefoundry.common.exceptions import HttpRequestException
+from truefoundry.common.request_utils import (
+    augmented_raise_for_status,
+    cloud_storage_http_request,
+)
+from truefoundry.pydantic_v1 import BaseModel
+
+logger = logging.getLogger("truefoundry")
+
+
+class MultiPartUploadStorageProvider(str, Enum):
+    S3_COMPATIBLE = "S3_COMPATIBLE"
+    AZURE_BLOB = "AZURE_BLOB"
+
+
+class SignedURL(BaseModel):
+    signed_url: str
+    path: Optional[str] = None
+
+
+class MultiPartUpload(BaseModel):
+    storage_provider: MultiPartUploadStorageProvider
+    part_signed_urls: List[SignedURL]
+    s3_compatible_upload_id: Optional[str] = None
+    azure_blob_block_ids: Optional[List[str]] = None
+    finalize_signed_url: SignedURL
+
+
+class _FileMultiPartInfo(NamedTuple):
+    num_parts: int
+    part_size: int
+    file_size: int
+
+
+class _PartNumberEtag(NamedTuple):
+    part_number: int
+    etag: str
+
+
+_MIN_BYTES_REQUIRED_FOR_MULTIPART = 100 * 1024 * 1024
+# GCP/S3 Maximum number of parts per upload	10,000
+# Maximum number of blocks in a block blob 50,000 blocks
+# TODO: This number is artificially limited now. Later
+# we will ask for parts signed URI in batches rather than in a single
+# API Calls:
+# Create Multipart Upload (Returns maximum number of parts, size limit of
+#                            a single part, upload id for s3 etc )
+#   Get me signed uris for first 500 parts
+#     Upload 500 parts
+#   Get me signed uris for the next 500 parts
+#     Upload 500 parts
+#   ...
+# Finalize the Multipart upload using the finalize signed url returned
+# by Create Multipart Upload or get a new one.
+_MAX_NUM_PARTS_FOR_MULTIPART = 1000
+# Azure Maximum size of a block in a block blob	4000 MiB
+# GCP/S3 Maximum size of an individual part in a multipart upload 5 GiB
+_MAX_PART_SIZE_BYTES_FOR_MULTIPART = 4 * 1024 * 1024 * 1000
+
+
+class _CallbackIOWrapperForMultiPartUpload(CallbackIOWrapper):
+    def __init__(self, callback, stream, method, length: int):
+        self.wrapper_setattr("_length", length)
+        super().__init__(callback, stream, method)
+
+    def __len__(self):
+        return self.wrapper_getattr("_length")
+
+
+def _align_part_size_with_mmap_allocation_granularity(part_size: int) -> int:
+    modulo = part_size % mmap.ALLOCATIONGRANULARITY
+    if modulo == 0:
+        return part_size
+
+    part_size += mmap.ALLOCATIONGRANULARITY - modulo
+    return part_size
+
+
+# Can not be less than 5 * 1024 * 1024
+_PART_SIZE_BYTES_FOR_MULTIPART = _align_part_size_with_mmap_allocation_granularity(
+    10 * 1024 * 1024
+)
+
+
+def decide_file_parts(
+    file_path: str,
+    multipart_upload_allowed: bool = not ENV_VARS.TFY_ARTIFACTS_DISABLE_MULTIPART_UPLOAD,
+    min_file_size_bytes_for_multipart: int = _MIN_BYTES_REQUIRED_FOR_MULTIPART,
+) -> _FileMultiPartInfo:
+    file_size = os.path.getsize(file_path)
+    if not multipart_upload_allowed or file_size < min_file_size_bytes_for_multipart:
+        return _FileMultiPartInfo(1, part_size=file_size, file_size=file_size)
+
+    ideal_num_parts = math.ceil(file_size / _PART_SIZE_BYTES_FOR_MULTIPART)
+    if ideal_num_parts <= _MAX_NUM_PARTS_FOR_MULTIPART:
+        return _FileMultiPartInfo(
+            ideal_num_parts,
+            part_size=_PART_SIZE_BYTES_FOR_MULTIPART,
+            file_size=file_size,
+        )
+
+    part_size_when_using_max_parts = math.ceil(file_size / _MAX_NUM_PARTS_FOR_MULTIPART)
+    part_size_when_using_max_parts = _align_part_size_with_mmap_allocation_granularity(
+        part_size_when_using_max_parts
+    )
+    if part_size_when_using_max_parts > _MAX_PART_SIZE_BYTES_FOR_MULTIPART:
+        raise ValueError(
+            f"file {file_path!r} is too big for upload. Multipart chunk"
+            f" size {part_size_when_using_max_parts} is higher"
+            f" than {_MAX_PART_SIZE_BYTES_FOR_MULTIPART}"
+        )
+    num_parts = math.ceil(file_size / part_size_when_using_max_parts)
+    return _FileMultiPartInfo(
+        num_parts, part_size=part_size_when_using_max_parts, file_size=file_size
+    )
+
+
+def _get_s3_compatible_completion_body(multi_parts: List[_PartNumberEtag]) -> str:
+    body = "<CompleteMultipartUpload>\n"
+    for part in multi_parts:
+        body += "  <Part>\n"
+        body += f"    <PartNumber>{part.part_number}</PartNumber>\n"
+        body += f"    <ETag>{part.etag}</ETag>\n"
+        body += "  </Part>\n"
+    body += "</CompleteMultipartUpload>"
+    return body
+
+
+def _get_azure_blob_completion_body(block_ids: List[str]) -> str:
+    body = "<BlockList>\n"
+    for block_id in block_ids:
+        body += f"<Uncommitted>{block_id}</Uncommitted> "
+    body += "</BlockList>"
+    return body
+
+
+def _file_part_upload(
+    url: str,
+    file_path: str,
+    seek: int,
+    length: int,
+    file_size: int,
+    abort_event: Optional[Event] = None,
+    method: str = "put",
+    exception_class=HttpRequestException,
+):
+    def callback(*_, **__):
+        if abort_event and abort_event.is_set():
+            raise Exception("aborting upload")
+
+    with open(file_path, "rb") as file:
+        with mmap.mmap(
+            file.fileno(),
+            length=min(file_size - seek, length),
+            offset=seek,
+            access=mmap.ACCESS_READ,
+        ) as mapped_file:
+            wrapped_file = _CallbackIOWrapperForMultiPartUpload(
+                callback, mapped_file, "read", len(mapped_file)
+            )
+            with cloud_storage_http_request(
+                method=method,
+                url=url,
+                data=wrapped_file,
+                exception_class=exception_class,
+            ) as response:
+                augmented_raise_for_status(response, exception_class=exception_class)
+                return response
+
+
+def s3_compatible_multipart_upload(
+    multipart_upload: MultiPartUpload,
+    local_file: str,
+    multipart_info: _FileMultiPartInfo,
+    executor: ThreadPoolExecutor,
+    progress_bar: Optional[Progress] = None,
+    abort_event: Optional[Event] = None,
+    exception_class=HttpRequestException,
+):
+    abort_event = abort_event or Event()
+    parts = []
+
+    if progress_bar:
+        multi_part_upload_progress = progress_bar.add_task(
+            f"[green]Uploading {local_file}:", start=True
+        )
+
+    def upload(part_number: int, seek: int) -> None:
+        logger.debug(
+            "Uploading part %d/%d of %s",
+            part_number,
+            multipart_info.num_parts,
+            local_file,
+        )
+        response = _file_part_upload(
+            url=multipart_upload.part_signed_urls[part_number].signed_url,
+            file_path=local_file,
+            seek=seek,
+            length=multipart_info.part_size,
+            file_size=multipart_info.file_size,
+            abort_event=abort_event,
+            exception_class=exception_class,
+        )
+        logger.debug(
+            "Uploaded part %d/%d of %s",
+            part_number,
+            multipart_info.num_parts,
+            local_file,
+        )
+        if progress_bar:
+            progress_bar.update(
+                multi_part_upload_progress,
+                advance=multipart_info.part_size,
+                total=multipart_info.file_size,
+            )
+        etag = response.headers["ETag"]
+        parts.append(_PartNumberEtag(etag=etag, part_number=part_number + 1))
+
+    futures: List[Future] = []
+    for part_number, seek in enumerate(
+        range(0, multipart_info.file_size, multipart_info.part_size)
+    ):
+        future = executor.submit(upload, part_number=part_number, seek=seek)
+        futures.append(future)
+
+    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()
+
+    logger.debug("Finalizing multipart upload of %s", local_file)
+    parts = sorted(parts, key=lambda part: part.part_number)
+    response = requests.post(
+        multipart_upload.finalize_signed_url.signed_url,
+        data=_get_s3_compatible_completion_body(parts),
+        timeout=2 * 60,
+    )
+    response.raise_for_status()
+    logger.debug("Multipart upload of %s completed", local_file)
+
+
+def azure_multi_part_upload(
+    multipart_upload: MultiPartUpload,
+    local_file: str,
+    multipart_info: _FileMultiPartInfo,
+    executor: ThreadPoolExecutor,
+    progress_bar: Optional[Progress] = None,
+    abort_event: Optional[Event] = None,
+    exception_class=HttpRequestException,
+):
+    abort_event = abort_event or Event()
+
+    if progress_bar:
+        multi_part_upload_progress = progress_bar.add_task(
+            f"[green]Uploading {local_file}:", start=True
+        )
+
+    def upload(part_number: int, seek: int):
+        logger.debug(
+            "Uploading part %d/%d of %s",
+            part_number,
+            multipart_info.num_parts,
+            local_file,
+        )
+        _file_part_upload(
+            url=multipart_upload.part_signed_urls[part_number].signed_url,
+            file_path=local_file,
+            seek=seek,
+            length=multipart_info.part_size,
+            file_size=multipart_info.file_size,
+            abort_event=abort_event,
+            exception_class=exception_class,
+        )
+        if progress_bar:
+            progress_bar.update(
+                multi_part_upload_progress,
+                advance=multipart_info.part_size,
+                total=multipart_info.file_size,
+            )
+        logger.debug(
+            "Uploaded part %d/%d of %s",
+            part_number,
+            multipart_info.num_parts,
+            local_file,
+        )
+
+    futures: List[Future] = []
+    for part_number, seek in enumerate(
+        range(0, multipart_info.file_size, multipart_info.part_size)
+    ):
+        future = executor.submit(upload, part_number=part_number, seek=seek)
+        futures.append(future)
+
+    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()
+
+    logger.debug("Finalizing multipart upload of %s", local_file)
+    if multipart_upload.azure_blob_block_ids:
+        response = requests.put(
+            multipart_upload.finalize_signed_url.signed_url,
+            data=_get_azure_blob_completion_body(
+                block_ids=multipart_upload.azure_blob_block_ids
+            ),
+            timeout=2 * 60,
+        )
+        response.raise_for_status()
+    logger.debug("Multipart upload of %s completed", local_file)

truefoundry/common/utils.py

@@ -37,25 +37,25 @@ class _TFYServersConfig(BaseSettings):
     mlfoundry_server_url: str
 
     @classmethod
-    def from_base_url(cls, base_url: str) -> "_TFYServersConfig":
-        base_url = base_url.strip("/")
+    def from_tfy_host(cls, tfy_host: str) -> "_TFYServersConfig":
+        tfy_host = tfy_host.strip("/")
         return cls(
-            tenant_host=urlparse(base_url).netloc,
-            servicefoundry_server_url=urljoin(base_url, API_SERVER_RELATIVE_PATH),
-            mlfoundry_server_url=urljoin(base_url, MLFOUNDRY_SERVER_RELATIVE_PATH),
+            tenant_host=urlparse(tfy_host).netloc,
+            servicefoundry_server_url=urljoin(tfy_host, API_SERVER_RELATIVE_PATH),
+            mlfoundry_server_url=urljoin(tfy_host, MLFOUNDRY_SERVER_RELATIVE_PATH),
         )
 
 
 _tfy_servers_config = None
 
 
-def get_tfy_servers_config(base_url: str) -> _TFYServersConfig:
+def get_tfy_servers_config(tfy_host: str) -> _TFYServersConfig:
     global _tfy_servers_config
     if _tfy_servers_config is None:
         if ENV_VARS.TFY_CLI_LOCAL_DEV_MODE:
             _tfy_servers_config = _TFYServersConfig()  # type: ignore[call-arg]
         else:
-            _tfy_servers_config = _TFYServersConfig.from_base_url(base_url)
+            _tfy_servers_config = _TFYServersConfig.from_tfy_host(tfy_host)
     return _tfy_servers_config
 
 
@@ -106,11 +106,11 @@ def validate_tfy_host(tfy_host: str) -> None:
 
 
 def resolve_tfy_host(tfy_host: Optional[str] = None) -> str:
-    if not tfy_host and not ENV_VARS.TFY_HOST:
+    tfy_host = tfy_host or ENV_VARS.TFY_HOST
+    if not tfy_host:
         raise ValueError(
             f"Either `host` should be provided using `--host <value>`, or `{TFY_HOST_ENV_KEY}` env must be set"
         )
-    tfy_host = tfy_host or ENV_VARS.TFY_HOST
     tfy_host = tfy_host.strip("/")
     validate_tfy_host(tfy_host)
     return tfy_host

truefoundry/common/warnings.py

@@ -0,0 +1,21 @@
+import warnings
+
+
+class TrueFoundryDeprecationWarning(DeprecationWarning):
+    pass
+
+
+def surface_truefoundry_deprecation_warnings() -> None:
+    """Unmute TrueFoundry deprecation warnings."""
+    warnings.filterwarnings(
+        "default",
+        category=TrueFoundryDeprecationWarning,
+    )
+
+
+def suppress_truefoundry_deprecation_warnings() -> None:
+    """Mute TrueFoundry deprecation warnings."""
+    warnings.filterwarnings(
+        "ignore",
+        category=TrueFoundryDeprecationWarning,
+    )

truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py

@@ -1,4 +1,3 @@
-import os
 import shlex
 from typing import Dict, List, Optional
 
@@ -12,7 +11,10 @@ from truefoundry.deploy.builder.constants import (
     UV_CONF_BUILDKIT_SECRET_MOUNT,
     UV_CONF_SECRET_MOUNT_AS_ENV,
 )
-from truefoundry.deploy.v2.lib.patched_models import CUDAVersion
+from truefoundry.deploy.v2.lib.patched_models import (
+    CUDAVersion,
+    _resolve_requirements_path,
+)
 
 # TODO (chiragjn): Switch to a non-root user inside the container
 
@@ -80,23 +82,6 @@ CUDA_VERSION_TO_IMAGE_TAG: Dict[str, str] = {
 }
 
 
-def resolve_requirements_txt_path(build_configuration: PythonBuild) -> Optional[str]:
-    if build_configuration.requirements_path:
-        return build_configuration.requirements_path
-
-    # TODO: what if there is a requirements.txt but user does not wants us to use it.
-    possible_requirements_txt_path = os.path.join(
-        build_configuration.build_context_path, "requirements.txt"
-    )
-
-    if os.path.isfile(possible_requirements_txt_path):
-        return os.path.relpath(
-            possible_requirements_txt_path, start=build_configuration.build_context_path
-        )
-
-    return None
-
-
 def generate_apt_install_command(apt_packages: Optional[List[str]]) -> Optional[str]:
     packages_list = None
     if apt_packages:
@@ -182,7 +167,10 @@ def generate_dockerfile_content(
     mount_python_package_manager_conf_secret: bool = False,
 ) -> str:
     # TODO (chiragjn): Handle recursive references to other requirements files e.g. `-r requirements-gpu.txt`
-    requirements_path = resolve_requirements_txt_path(build_configuration)
+    requirements_path = _resolve_requirements_path(
+        build_context_path=build_configuration.build_context_path,
+        requirements_path=build_configuration.requirements_path,
+    )
     requirements_destination_path = (
         "/tmp/requirements.txt" if requirements_path else None
     )

truefoundry/deploy/cli/commands/deploy_command.py

@@ -61,11 +61,11 @@ def deploy_command(
 ):
     if ctx.invoked_subcommand is not None:
         return
-    from truefoundry.deploy.lib.auth.servicefoundry_session import ServiceFoundrySession
+    from truefoundry.common.session import Session
     from truefoundry.deploy.v2.lib.deployable_patched_models import Application
 
     try:
-        _ = ServiceFoundrySession()
+        _ = Session.new()
     except Exception as e:
         raise ClickException(message=str(e)) from e
 
@@ -126,10 +126,10 @@ def deploy_command(
 )
 @handle_exception_wrapper
 def deploy_workflow_command(name: str, file: str, workspace_fqn: str):
-    from truefoundry.deploy.lib.auth.servicefoundry_session import ServiceFoundrySession
+    from truefoundry.common.session import Session
 
     try:
-        _ = ServiceFoundrySession()
+        _ = Session.new()
     except Exception as e:
         raise ClickException(message=str(e)) from e
 

truefoundry/deploy/lib/clients/servicefoundry_client.py

@@ -23,9 +23,9 @@ from truefoundry.common.servicefoundry_client import (
     check_min_cli_version,
     session_with_retries,
 )
+from truefoundry.common.session import Session
 from truefoundry.deploy.auto_gen import models as auto_gen_models
 from truefoundry.deploy.io.output_callback import OutputCallBack
-from truefoundry.deploy.lib.auth.servicefoundry_session import ServiceFoundrySession
 from truefoundry.deploy.lib.model.entity import (
     Application,
     CreateDockerRepositoryResponse,
@@ -78,17 +78,16 @@ def _upload_packaged_code(metadata, package_file):
 
 
 class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
-    # TODO (chiragjn): Rename base_url to tfy_host
-    def __init__(self, init_session: bool = True, base_url: Optional[str] = None):
-        self._session: Optional[ServiceFoundrySession] = None
+    def __init__(self, init_session: bool = True, tfy_host: Optional[str] = None):
+        self._session: Optional[Session] = None
         if init_session:
-            if base_url:
-                logger.warning("Passed base url %r will be ignored", base_url)
-            self._session = ServiceFoundrySession()
-            base_url = self._session.base_url
-        elif not base_url:
-            raise Exception("Neither session, not base_url provided")
-        super().__init__(base_url=base_url)
+            if tfy_host:
+                logger.warning(f"Passed tfy_host {tfy_host!r} will be ignored")
+            self._session = Session.new()
+            tfy_host = self._session.tfy_host
+        elif not tfy_host:
+            raise Exception("Neither session, not tfy_host provided")
+        super().__init__(tfy_host=tfy_host)
 
     def _get_header(self):
         if not self._session:
@@ -250,11 +249,11 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         }
         logger.debug(json.dumps(data))
         url = f"{self._api_server_url}/{VERSION_PREFIX}/deployment"
-        deploy_response = session_with_retries().post(
+        response = session_with_retries().post(
             url, json=data, headers=self._get_header()
         )
-        response = request_handling(deploy_response)
-        return Deployment.parse_obj(response["deployment"])
+        response_data = request_handling(response)
+        return Deployment.parse_obj(response_data["deployment"])
 
     def _get_log_print_line(self, log: dict):
         timestamp = int(log["time"]) / 1e6

truefoundry/deploy/lib/dao/application.py

@@ -180,7 +180,7 @@ def trigger_job(
         params=params if params else None,
     )
     jobRunName = result.jobRunName
-    previous_runs_url = f"{client.base_url.strip('/')}/deployments/{application_info.id}?tab=previousRuns"
+    previous_runs_url = f"{client.tfy_host.strip('/')}/deployments/{application_info.id}?tab=previousRuns"
     logger.info(
         f"{message}.\n"
         f"You can check the status of your job run at {previous_runs_url} with jobRunName: {jobRunName}"
@@ -259,6 +259,6 @@ def trigger_workflow(application_fqn: str, inputs: Optional[Dict[str, Any]] = No
     )
     logger.info(f"Started Execution for Workflow: {application_info.name}")
     executions_page = (
-        f"{client.base_url.strip('/')}/deployments/{application_info.id}?tab=executions"
+        f"{client.tfy_host.strip('/')}/deployments/{application_info.id}?tab=executions"
     )
     logger.info(f"You can check the executions at {executions_page}")

truefoundry/deploy/lib/dao/workspace.py

@@ -32,7 +32,7 @@ def create_workspace(
             resources=workspace_resources,
         )
 
-    url = f"{client.base_url.strip('/')}/workspaces"
+    url = f"{client.tfy_host.strip('/')}/workspaces"
     logger.info(
         "You can find your workspace: '%s' on the dashboard: %s", workspace.name, url
     )

truefoundry/deploy/lib/session.py

@@ -68,7 +68,7 @@ def login(
         if api_key:
             token = Token(access_token=api_key, refresh_token=None)
         else:
-            auth_service = AuthServiceClient.from_base_url(base_url=host)
+            auth_service = AuthServiceClient.from_tfy_host(tfy_host=host)
             # interactive login
             token = _login_with_device_code(base_url=host, auth_service=auth_service)
 

truefoundry/deploy/v2/lib/deploy.py

@@ -90,7 +90,7 @@ def _log_application_dashboard_url(deployment: Deployment, log_message: str):
     # TODO: is there any simpler way to get this? :cry
     client = ServiceFoundryServiceClient()
 
-    url = f"{client.base_url.strip('/')}/applications/{application_id}?tab=deployments"
+    url = f"{client.tfy_host.strip('/')}/applications/{application_id}?tab=deployments"
     logger.info(log_message, url)
 
 
@@ -272,7 +272,7 @@ def deploy_component(
             if not last_status_printed or DeploymentTransitionStatus.is_failure_state(
                 last_status_printed
             ):
-                deployment_tab_url = f"{client.base_url.strip('/')}/applications/{response.applicationId}?tab=deployments"
+                deployment_tab_url = f"{client.tfy_host.strip('/')}/applications/{response.applicationId}?tab=deployments"
                 message = f"Deployment Failed. Please refer to the logs for additional details - {deployment_tab_url}"
                 sys.exit(message)
         except KeyboardInterrupt:

truefoundry/deploy/v2/lib/deploy_workflow.py

@@ -297,5 +297,5 @@ def deploy_workflow(
         workflow.name,
         deployment.fqn,
     )
-    deployment_url = f"{client.base_url.strip('/')}/applications/{deployment.applicationId}?tab=deployments"
+    deployment_url = f"{client.tfy_host.strip('/')}/applications/{deployment.applicationId}?tab=deployments"
     logger.info("You can find the application on the dashboard:- '%s'", deployment_url)