wandb 0.15.3__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/internal/internal_api.py

@@ -39,12 +39,13 @@ from wandb.apis.normalize import normalize_exceptions, parse_backend_error_messa
 from wandb.errors import CommError, UsageError
 from wandb.integration.sagemaker import parse_sm_secrets
 from wandb.old.settings import Settings
+from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings
 from wandb.sdk.lib.gql_request import GraphQLSession
 from wandb.sdk.lib.hashutil import B64MD5, md5_file_b64
 
 from ..lib import retry
 from ..lib.filenames import DIFF_FNAME, METADATA_FNAME
-from ..lib.git import GitRepo
+from ..lib.gitlib import GitRepo
 from . import context
 from .progress import AsyncProgress, Progress
 
@@ -66,6 +67,7 @@ if TYPE_CHECKING:
         md5: str
         mimetype: Optional[str]
         artifactManifestID: Optional[str]  # noqa: N815
+        uploadPartsInput: Optional[List[Dict[str, object]]]  # noqa: N815
 
     class CreateArtifactFilesResponseFile(TypedDict):
         id: str
@@ -73,11 +75,34 @@ if TYPE_CHECKING:
         displayName: str  # noqa: N815
         uploadUrl: Optional[str]  # noqa: N815
         uploadHeaders: Sequence[str]  # noqa: N815
+        uploadMultipartUrls: "UploadPartsResponse"  # noqa: N815
+        storagePath: str  # noqa: N815
         artifact: "CreateArtifactFilesResponseFileNode"
 
     class CreateArtifactFilesResponseFileNode(TypedDict):
         id: str
 
+    class UploadPartsResponse(TypedDict):
+        uploadUrlParts: List["UploadUrlParts"]  # noqa: N815
+        uploadID: str  # noqa: N815
+
+    class UploadUrlParts(TypedDict):
+        partNumber: int  # noqa: N815
+        uploadUrl: str  # noqa: N815
+
+    class CompleteMultipartUploadArtifactInput(TypedDict):
+        """Corresponds to `type CompleteMultipartUploadArtifactInput` in schema.graphql."""
+
+        completeMultipartAction: str  # noqa: N815
+        completedParts: Dict[int, str]  # noqa: N815
+        artifactID: str  # noqa: N815
+        storagePath: str  # noqa: N815
+        uploadID: str  # noqa: N815
+        md5: str
+
+    class CompleteMultipartUploadArtifactResponse(TypedDict):
+        digest: str
+
     class DefaultSettings(TypedDict):
         section: str
         git_remote: str
@@ -205,6 +230,10 @@ class Api:
             self._environ.get("WANDB__EXTRA_HTTP_HEADERS", {})
         )
 
+        auth = None
+        if _thread_local_api_settings.cookies is None:
+            auth = ("api", self.api_key or "")
+        extra_http_headers.update(_thread_local_api_settings.headers or {})
         self.client = Client(
             transport=GraphQLSession(
                 headers={
@@ -217,8 +246,9 @@ class Api:
                 # this timeout won't apply when the DNS lookup fails. in that case, it will be 60s
                 # https://bugs.python.org/issue22889
                 timeout=self.HTTP_TIMEOUT,
-                auth=("api", self.api_key or ""),
+                auth=auth,
                 url=f"{self.settings('base_url')}/graphql",
+                cookies=_thread_local_api_settings.cookies,
             )
         )
 
@@ -242,6 +272,11 @@ class Api:
         self.upload_file_retry = normalize_exceptions(
             retry.retriable(retry_timedelta=retry_timedelta)(self.upload_file)
         )
+        self.upload_multipart_file_chunk_retry = normalize_exceptions(
+            retry.retriable(retry_timedelta=retry_timedelta)(
+                self.upload_multipart_file_chunk
+            )
+        )
         self._client_id_mapping: Dict[str, str] = {}
         # Large file uploads to azure can optionally use their SDK
         self._azure_blob_module = util.get_module("azure.storage.blob")
@@ -252,6 +287,7 @@ class Api:
         self.server_use_artifact_input_info: Optional[List[str]] = None
         self._max_cli_version: Optional[str] = None
         self._server_settings_type: Optional[List[str]] = None
+        self.fail_run_queue_item_input_info: Optional[List[str]] = None
 
     def gql(self, *args: Any, **kwargs: Any) -> Any:
         ret = self._retry_gql(
@@ -273,7 +309,7 @@ class Api:
 
     def reauth(self) -> None:
         """Ensure the current api key is set in the transport."""
-        self.client.transport.auth = ("api", self.api_key or "")
+        self.client.transport.session.auth = ("api", self.api_key or "")
 
     def relocate(self) -> None:
         """Ensure the current api points to the right server."""
@@ -307,6 +343,8 @@ class Api:
 
     @property
     def api_key(self) -> Optional[str]:
+        if _thread_local_api_settings.api_key:
+            return _thread_local_api_settings.api_key
         auth = requests.utils.get_netrc_auth(self.api_url)
         key = None
         if auth:
@@ -549,13 +587,103 @@ class Api:
         return "failRunQueueItem" in mutations
 
     @normalize_exceptions
-    def fail_run_queue_item(self, run_queue_item_id: str) -> bool:
+    def fail_run_queue_item_fields_introspection(self) -> List:
+        if self.fail_run_queue_item_input_info:
+            return self.fail_run_queue_item_input_info
+        query_string = """
+           query ProbeServerFailRunQueueItemInput {
+                FailRunQueueItemInputInfoType: __type(name:"FailRunQueueItemInput") {
+                    inputFields{
+                        name
+                    }
+                }
+            }
+        """
+
+        query = gql(query_string)
+        res = self.gql(query)
+
+        self.fail_run_queue_item_input_info = [
+            field.get("name", "")
+            for field in res.get("FailRunQueueItemInputInfoType", {}).get(
+                "inputFields", [{}]
+            )
+        ]
+        return self.fail_run_queue_item_input_info
+
+    @normalize_exceptions
+    def fail_run_queue_item(
+        self,
+        run_queue_item_id: str,
+        message: str,
+        stage: str,
+        file_paths: Optional[List[str]] = None,
+    ) -> bool:
+        if not self.fail_run_queue_item_introspection():
+            return False
+        variable_values: Dict[str, Union[str, Optional[List[str]]]] = {
+            "runQueueItemId": run_queue_item_id,
+        }
+        if "message" in self.fail_run_queue_item_fields_introspection():
+            variable_values.update({"message": message, "stage": stage})
+            if file_paths is not None:
+                variable_values["filePaths"] = file_paths
+            mutation_string = """
+            mutation failRunQueueItem($runQueueItemId: ID!, $message: String!, $stage: String!, $filePaths: [String!]) {
+                failRunQueueItem(
+                    input: {
+                        runQueueItemId: $runQueueItemId
+                        message: $message
+                        stage: $stage
+                        filePaths: $filePaths
+                    }
+                ) {
+                    success
+                }
+            }
+            """
+        else:
+            mutation_string = """
+            mutation failRunQueueItem($runQueueItemId: ID!) {
+                failRunQueueItem(
+                    input: {
+                        runQueueItemId: $runQueueItemId
+                    }
+                ) {
+                    success
+                }
+            }
+            """
+
+        mutation = gql(mutation_string)
+        response = self.gql(mutation, variable_values=variable_values)
+        result: bool = response["failRunQueueItem"]["success"]
+        return result
+
+    @normalize_exceptions
+    def update_run_queue_item_warning_introspection(self) -> bool:
+        _, _, mutations = self.server_info_introspection()
+        return "updateRunQueueItemWarning" in mutations
+
+    @normalize_exceptions
+    def update_run_queue_item_warning(
+        self,
+        run_queue_item_id: str,
+        message: str,
+        stage: str,
+        file_paths: Optional[List[str]] = None,
+    ) -> bool:
+        if not self.update_run_queue_item_warning_introspection():
+            return False
         mutation = gql(
             """
-        mutation failRunQueueItem($runQueueItemId: ID!) {
-            failRunQueueItem(
+        mutation updateRunQueueItemWarning($runQueueItemId: ID!, $message: String!, $stage: String!, $filePaths: [String!]) {
+            updateRunQueueItemWarning(
                 input: {
                     runQueueItemId: $runQueueItemId
+                    message: $message
+                    stage: $stage
+                    filePaths: $filePaths
                 }
             ) {
                 success
@@ -567,9 +695,12 @@ class Api:
             mutation,
             variable_values={
                 "runQueueItemId": run_queue_item_id,
+                "message": message,
+                "stage": stage,
+                "filePaths": file_paths,
             },
         )
-        result: bool = response["failRunQueueItem"]["success"]
+        result: bool = response["updateRunQueueItemWarning"]["success"]
         return result
 
     @normalize_exceptions
@@ -580,6 +711,7 @@ class Api:
             viewer {
                 id
                 entity
+                username
                 flags
                 teams {
                     edges {
@@ -1992,7 +2124,16 @@ class Api:
         Returns:
             A tuple of the content length and the streaming response
         """
-        response = requests.get(url, auth=("user", self.api_key), stream=True)  # type: ignore
+        auth = None
+        if _thread_local_api_settings.cookies is None:
+            auth = ("user", self.api_key or "")
+        response = requests.get(
+            url,
+            auth=auth,
+            cookies=_thread_local_api_settings.cookies or {},
+            headers=_thread_local_api_settings.headers or {},
+            stream=True,
+        )
         response.raise_for_status()
         return int(response.headers.get("content-length", 0)), response
 
@@ -2060,6 +2201,53 @@ class Api:
             else:
                 raise requests.exceptions.ConnectionError(e.message)
 
+    def upload_multipart_file_chunk(
+        self,
+        url: str,
+        upload_chunk: bytes,
+        extra_headers: Optional[Dict[str, str]] = None,
+    ) -> Optional[requests.Response]:
+        """Upload a file chunk to S3 with failure resumption.
+
+        Arguments:
+            url: The url to download
+            upload_chunk: The path to the file you want to upload
+            extra_headers: A dictionary of extra headers to send with the request
+
+        Returns:
+            The `requests` library response object
+        """
+        try:
+            response = self._upload_file_session.put(
+                url, data=upload_chunk, headers=extra_headers
+            )
+            response.raise_for_status()
+        except requests.exceptions.RequestException as e:
+            logger.error(f"upload_file exception {url}: {e}")
+            request_headers = e.request.headers if e.request is not None else ""
+            logger.error(f"upload_file request headers: {request_headers}")
+            response_content = e.response.content if e.response is not None else ""
+            logger.error(f"upload_file response body: {response_content}")
+            status_code = e.response.status_code if e.response is not None else 0
+            # S3 reports retryable request timeouts out-of-band
+            is_aws_retryable = status_code == 400 and "RequestTimeout" in str(
+                response_content
+            )
+            # Retry errors from cloud storage or local network issues
+            if (
+                status_code in (308, 408, 409, 429, 500, 502, 503, 504)
+                or isinstance(
+                    e,
+                    (requests.exceptions.Timeout, requests.exceptions.ConnectionError),
+                )
+                or is_aws_retryable
+            ):
+                _e = retry.TransientError(exc=e)
+                raise _e.with_traceback(sys.exc_info()[2])
+            else:
+                wandb._sentry.reraise(e)
+        return response
+
     def upload_file(
         self,
         url: str,
@@ -2353,7 +2541,8 @@ class Api:
         config = dict(config)
 
         if "parameters" not in config:
-            raise ValueError("sweep config must have a parameters section")
+            # still shows an anaconda warning, but doesn't error
+            return config
 
         for parameter_name in config["parameters"]:
             parameter = config["parameters"][parameter_name]
@@ -3014,6 +3203,50 @@ class Api:
         )
         return response
 
+    def complete_multipart_upload_artifact(
+        self,
+        artifact_id: str,
+        storage_path: str,
+        completed_parts: List[Dict[str, Any]],
+        upload_id: str,
+        complete_multipart_action: str = "Complete",
+    ) -> Optional[str]:
+        mutation = gql(
+            """
+        mutation CompleteMultipartUploadArtifact(
+            $completeMultipartAction: CompleteMultipartAction!,
+            $completedParts: [UploadPartsInput!]!,
+            $artifactID: ID!
+            $storagePath: String!
+            $uploadID: String!
+        ) {
+        completeMultipartUploadArtifact(
+            input: {
+                completeMultipartAction: $completeMultipartAction,
+                completedParts: $completedParts,
+                artifactID: $artifactID,
+                storagePath: $storagePath
+                uploadID: $uploadID
+            }
+            ) {
+                digest
+            }
+        }
+        """
+        )
+        response = self.gql(
+            mutation,
+            variable_values={
+                "completeMultipartAction": complete_multipart_action,
+                "artifactID": artifact_id,
+                "storagePath": storage_path,
+                "completedParts": completed_parts,
+                "uploadID": upload_id,
+            },
+        )
+        digest: Optional[str] = response["completeMultipartUploadArtifact"]["digest"]
+        return digest
+
     def create_artifact_manifest(
         self,
         name: str,
@@ -3171,19 +3404,39 @@ class Api:
                     self._client_id_mapping[client_id] = server_id
         return server_id
 
+    def server_create_artifact_file_spec_input_introspection(self) -> List:
+        query_string = """
+           query ProbeServerCreateArtifactFileSpecInput {
+                CreateArtifactFileSpecInputInfoType: __type(name:"CreateArtifactFileSpecInput") {
+                    inputFields{
+                        name
+                    }
+                }
+            }
+        """
+
+        query = gql(query_string)
+        res = self.gql(query)
+        create_artifact_file_spec_input_info = [
+            field.get("name", "")
+            for field in res.get("CreateArtifactFileSpecInputInfoType", {}).get(
+                "inputFields", [{}]
+            )
+        ]
+        return create_artifact_file_spec_input_info
+
     @normalize_exceptions
     def create_artifact_files(
         self, artifact_files: Iterable["CreateArtifactFileSpecInput"]
     ) -> Mapping[str, "CreateArtifactFilesResponseFile"]:
-        mutation = gql(
-            """
+        query_template = """
         mutation CreateArtifactFiles(
             $storageLayout: ArtifactStorageLayout!
             $artifactFiles: [CreateArtifactFileSpecInput!]!
         ) {
             createArtifactFiles(input: {
                 artifactFiles: $artifactFiles,
-                storageLayout: $storageLayout
+                storageLayout: $storageLayout,
             }) {
                 files {
                     edges {
@@ -3193,6 +3446,7 @@ class Api:
                             displayName
                             uploadUrl
                             uploadHeaders
+                            _MULTIPART_UPLOAD_FIELDS_
                             artifact {
                                 id
                             }
@@ -3202,7 +3456,16 @@ class Api:
             }
         }
         """
-        )
+        multipart_upload_url_query = """
+            storagePath
+            uploadMultipartUrls {
+                uploadID
+                uploadUrlParts {
+                    partNumber
+                    uploadUrl
+                }
+            }
+        """
 
         # TODO: we should use constants here from interface/artifacts.py
         # but probably don't want the dependency. We're going to remove
@@ -3211,6 +3474,17 @@ class Api:
         if env.get_use_v1_artifacts():
             storage_layout = "V1"
 
+        create_artifact_file_spec_input_fields = (
+            self.server_create_artifact_file_spec_input_introspection()
+        )
+        if "uploadPartsInput" in create_artifact_file_spec_input_fields:
+            query_template = query_template.replace(
+                "_MULTIPART_UPLOAD_FIELDS_", multipart_upload_url_query
+            )
+        else:
+            query_template = query_template.replace("_MULTIPART_UPLOAD_FIELDS_", "")
+
+        mutation = gql(query_template)
         response = self.gql(
             mutation,
             variable_values={