wandb 0.17.2__py3-none-any.whl → 0.17.4__py3-none-any.whl

wandb/sdk/internal/internal_api.py

@@ -11,6 +11,7 @@ import socket
 import sys
 import threading
 from copy import deepcopy
+from pathlib import Path
 from typing import (
     IO,
     TYPE_CHECKING,
@@ -37,14 +38,14 @@ from wandb_gql.client import RetryError
 import wandb
 from wandb import env, util
 from wandb.apis.normalize import normalize_exceptions, parse_backend_error_messages
-from wandb.errors import CommError, UnsupportedError, UsageError
+from wandb.errors import AuthenticationError, CommError, UnsupportedError, 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 import credentials, retry
 from ..lib.filenames import DIFF_FNAME, METADATA_FNAME
 from ..lib.gitlib import GitRepo
 from . import context
@@ -234,14 +235,18 @@ class Api:
         extra_http_headers = self.settings("_extra_http_headers") or json.loads(
             self._environ.get("WANDB__EXTRA_HTTP_HEADERS", "{}")
         )
+        extra_http_headers.update(_thread_local_api_settings.headers or {})
+
+        auth = None
+        if self.access_token is not None:
+            extra_http_headers["Authorization"] = f"Bearer {self.access_token}"
+        elif _thread_local_api_settings.cookies is None:
+            auth = ("api", self.api_key or "")
+
         proxies = self.settings("_proxies") or json.loads(
             self._environ.get("WANDB__PROXIES", "{}")
         )
 
-        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={
@@ -376,6 +381,35 @@ class Api:
         default_key: Optional[str] = self.default_settings.get("api_key")
         return env_key or key or sagemaker_key or default_key
 
+    @property
+    def access_token(self) -> Optional[str]:
+        """Retrieves an access token for authentication.
+
+        This function attempts to exchange an identity token for a temporary
+        access token from the server, and save it to the credentials file.
+        It uses the path to the identity token as defined in the environment
+        variables. If the environment variable is not set, it returns None.
+
+        Returns:
+            Optional[str]: The access token if available, otherwise None if
+            no identity token is supplied.
+        Raises:
+            AuthenticationError: If the path to the identity token is not found.
+        """
+        token_file_str = self._environ.get(env.IDENTITY_TOKEN_FILE)
+        if not token_file_str:
+            return None
+
+        token_file = Path(token_file_str)
+        if not token_file.exists():
+            raise AuthenticationError(f"Identity token file not found: {token_file}")
+
+        base_url = self.settings("base_url")
+        credentials_file = env.get_credentials_file(
+            str(credentials.DEFAULT_WANDB_CREDENTIALS_FILE), self._environ
+        )
+        return credentials.access_token(base_url, token_file, credentials_file)
+
     @property
     def api_url(self) -> str:
         return self.settings("base_url")  # type: ignore
@@ -2674,14 +2708,20 @@ class Api:
             A tuple of the content length and the streaming response
         """
         check_httpclient_logger_handler()
+
+        http_headers = _thread_local_api_settings.headers or {}
+
         auth = None
-        if _thread_local_api_settings.cookies is None:
-            auth = ("user", self.api_key or "")
+        if self.access_token is not None:
+            http_headers["Authorization"] = f"Bearer {self.access_token}"
+        elif _thread_local_api_settings.cookies is None:
+            auth = ("api", self.api_key or "")
+
         response = requests.get(
             url,
             auth=auth,
             cookies=_thread_local_api_settings.cookies or {},
-            headers=_thread_local_api_settings.headers or {},
+            headers=http_headers,
             stream=True,
         )
         response.raise_for_status()
@@ -2780,9 +2820,9 @@ class Api:
         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}")
+            logger.error(f"upload_file request headers: {request_headers!r}")
             response_content = e.response.content if e.response is not None else ""
-            logger.error(f"upload_file response body: {response_content}")
+            logger.error(f"upload_file response body: {response_content!r}")
             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(
@@ -2848,7 +2888,7 @@ class Api:
             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}")
+            logger.error(f"upload_file response body: {response_content!r}")
             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 = (
@@ -3039,6 +3079,7 @@ class Api:
         entity: Optional[str] = None,
         state: Optional[str] = None,
         prior_runs: Optional[List[str]] = None,
+        template_variable_values: Optional[Dict[str, Any]] = None,
     ) -> Tuple[str, List[str]]:
         """Upsert a sweep object.
 
@@ -3052,6 +3093,7 @@ class Api:
             entity (str): entity to use
             state (str): state
             prior_runs (list): IDs of existing runs to add to the sweep
+            template_variable_values (dict): template variable values
         """
         project_query = """
             project {
@@ -3096,7 +3138,17 @@ class Api:
         """
         # TODO(jhr): we need protocol versioning to know schema is not supported
         # for now we will just try both new and old query
-
+        mutation_5 = gql(
+            mutation_str.replace(
+                "$controller: JSONString,",
+                "$controller: JSONString,$launchScheduler: JSONString, $templateVariableValues: JSONString,",
+            )
+            .replace(
+                "controller: $controller,",
+                "controller: $controller,launchScheduler: $launchScheduler,templateVariableValues: $templateVariableValues,",
+            )
+            .replace("_PROJECT_QUERY_", project_query)
+        )
         # launchScheduler was introduced in core v0.14.0
         mutation_4 = gql(
             mutation_str.replace(
@@ -3105,7 +3157,7 @@ class Api:
             )
             .replace(
                 "controller: $controller,",
-                "controller: $controller,launchScheduler: $launchScheduler,",
+                "controller: $controller,launchScheduler: $launchScheduler",
             )
             .replace("_PROJECT_QUERY_", project_query)
         )
@@ -3124,7 +3176,7 @@ class Api:
         )
 
         # TODO(dag): replace this with a query for protocol versioning
-        mutations = [mutation_4, mutation_3, mutation_2, mutation_1]
+        mutations = [mutation_5, mutation_4, mutation_3, mutation_2, mutation_1]
 
         config = self._validate_config_and_fill_distribution(config)
 
@@ -3148,6 +3200,7 @@ class Api:
                     "projectName": project or self.settings("project"),
                     "controller": controller,
                     "launchScheduler": launch_scheduler,
+                    "templateVariableValues": json.dumps(template_variable_values),
                     "scheduler": scheduler,
                     "priorRunsFilters": filters,
                 }

wandb/sdk/launch/_launch_add.py

@@ -61,7 +61,7 @@ def launch_add(
         config: A dictionary containing the configuration for the run. May also contain
             resource specific arguments under the key "resource_args"
         template_variables: A dictionary containing values of template variables for a run queue.
-            Expected format of {"<var-name>": <var-value>}
+            Expected format of {"VAR_NAME": VAR_VALUE}
         project: Target project to send launched run to
         entity: Target entity to send launched run to
         queue: the name of the queue to enqueue the run to

wandb/sdk/launch/agent/agent.py

@@ -308,7 +308,9 @@ class LaunchAgent:
         self._wandb_run = None
 
         if self.gorilla_supports_agents:
-            settings = wandb.Settings(silent=True, disable_git=True)
+            settings = wandb.Settings(
+                silent=True, disable_git=True, disable_job_creation=True
+            )
             self._wandb_run = wandb.init(
                 project=self._project,
                 entity=self._entity,

wandb/sdk/launch/sweeps/scheduler.py

@@ -259,10 +259,12 @@ class Scheduler(ABC):
 
     def _init_wandb_run(self) -> "SdkRun":
         """Controls resume or init logic for a scheduler wandb run."""
+        settings = wandb.Settings(disable_job_creation=True)
         run: SdkRun = wandb.init(  # type: ignore
             name=f"Scheduler.{self._sweep_id}",
             resume="allow",
             config=self._kwargs,  # when run as a job, this sets config
+            settings=settings,
         )
         return run
 

wandb/sdk/launch/utils.py

@@ -60,7 +60,7 @@ AZURE_CONTAINER_REGISTRY_URI_REGEX = re.compile(
 )
 
 ELASTIC_CONTAINER_REGISTRY_URI_REGEX = re.compile(
-    r"^(?:https://)?(?P<account>[\w-]+)\.dkr\.ecr\.(?P<region>[\w-]+)\.amazonaws\.com/(?P<repository>[\w-]+):?(?P<tag>.*)$"
+    r"^(?:https://)?(?P<account>[\w-]+)\.dkr\.ecr\.(?P<region>[\w-]+)\.amazonaws\.com/(?P<repository>[\.\/\w-]+):?(?P<tag>.*)$"
 )
 
 GCP_ARTIFACT_REGISTRY_URI_REGEX = re.compile(

wandb/sdk/lib/_settings_toposort_generated.py

@@ -26,6 +26,7 @@ _Setting = Literal[
     "_disable_machine_info",
     "_executable",
     "_extra_http_headers",
+    "_file_stream_max_bytes",
     "_file_stream_retry_max",
     "_file_stream_retry_wait_min_seconds",
     "_file_stream_retry_wait_max_seconds",
@@ -91,6 +92,7 @@ _Setting = Literal[
     "config_paths",
     "console",
     "console_multipart",
+    "credentials_file",
     "deployment",
     "disable_code",
     "disable_git",
@@ -110,6 +112,9 @@ _Setting = Literal[
     "git_root",
     "heartbeat_seconds",
     "host",
+    "http_proxy",
+    "https_proxy",
+    "identity_token_file",
     "ignore_globs",
     "init_timeout",
     "is_local",
@@ -224,6 +229,9 @@ SETTINGS_TOPOLOGICALLY_SORTED: Final[Tuple[_Setting, ...]] = (
     "disable_git",
     "disable_job_creation",
     "files_dir",
+    "_proxies",
+    "http_proxy",
+    "https_proxy",
     "log_dir",
     "log_internal",
     "log_symlink_internal",

wandb/sdk/lib/apikey.py

@@ -1,6 +1,7 @@
 """apikey util."""
 
 import os
+import platform
 import stat
 import sys
 import textwrap
@@ -15,7 +16,7 @@ else:
     from typing_extensions import Literal
 
 import click
-import requests.utils
+from requests.utils import NETRC_FILES, get_netrc_auth
 
 import wandb
 from wandb.apis import InternalApi
@@ -54,10 +55,22 @@ def _fixup_anon_mode(default: Optional[Mode]) -> Optional[Mode]:
 
 
 def get_netrc_file_path() -> str:
+    """Return the path to the netrc file."""
+    # if the NETRC environment variable is set, use that
     netrc_file = os.environ.get("NETRC")
     if netrc_file:
         return os.path.expanduser(netrc_file)
-    return os.path.join(os.path.expanduser("~"), ".netrc")
+
+    # if either .netrc or _netrc exists in the home directory, use that
+    for netrc_file in NETRC_FILES:
+        home_dir = os.path.expanduser("~")
+        if os.path.exists(os.path.join(home_dir, netrc_file)):
+            return os.path.join(home_dir, netrc_file)
+
+    # otherwise, use .netrc on non-Windows platforms and _netrc on Windows
+    netrc_file = ".netrc" if platform.system() != "Windows" else "_netrc"
+
+    return os.path.join(os.path.expanduser("~"), netrc_file)
 
 
 def prompt_api_key(  # noqa: C901
@@ -254,7 +267,7 @@ def api_key(settings: Optional["Settings"] = None) -> Optional[str]:
         assert settings is not None
     if settings.api_key:
         return settings.api_key
-    auth = requests.utils.get_netrc_auth(settings.base_url)
+    auth = get_netrc_auth(settings.base_url)
     if auth:
         return auth[-1]
     return None

wandb/sdk/lib/credentials.py

@@ -0,0 +1,141 @@
+import json
+import os
+from datetime import datetime, timedelta
+from pathlib import Path
+
+import requests.utils
+
+from wandb.errors import AuthenticationError
+
+DEFAULT_WANDB_CREDENTIALS_FILE = Path(
+    os.path.expanduser("~/.config/wandb/credentials.json")
+)
+
+_expires_at_fmt = "%Y-%m-%d %H:%M:%S"
+
+
+def access_token(base_url: str, token_file: Path, credentials_file: Path) -> str:
+    """Retrieve an access token from the credentials file.
+
+    If no access token exists, create a new one by exchanging the identity
+    token from the token file, and save it to the credentials file.
+
+    Args:
+        base_url (str): The base URL of the server
+        token_file (pathlib.Path): The path to the file containing the
+        identity token
+        credentials_file (pathlib.Path): The path to file used to save
+        temporary access tokens
+
+    Returns:
+        str: The access token
+    """
+    if not credentials_file.exists():
+        _write_credentials_file(base_url, token_file, credentials_file)
+
+    data = _fetch_credentials(base_url, token_file, credentials_file)
+    return data["access_token"]
+
+
+def _write_credentials_file(base_url: str, token_file: Path, credentials_file: Path):
+    """Obtain an access token from the server and write it to the credentials file.
+
+    Args:
+        base_url (str): The base URL of the server
+        token_file (pathlib.Path): The path to the file containing the
+        identity token
+        credentials_file (pathlib.Path): The path to file used to save
+        temporary access tokens
+    """
+    credentials = _create_access_token(base_url, token_file)
+    data = {"credentials": {base_url: credentials}}
+    with open(credentials_file, "w") as file:
+        json.dump(data, file, indent=4)
+
+        # Set file permissions to be read/write by the owner only
+        os.chmod(credentials_file, 0o600)
+
+
+def _fetch_credentials(base_url: str, token_file: Path, credentials_file: Path) -> dict:
+    """Fetch the access token from the credentials file.
+
+    If the access token has expired, fetch a new one from the server and save it
+    to the credentials file.
+
+    Args:
+        base_url (str): The base URL of the server
+        token_file (pathlib.Path): The path to the file containing the
+        identity token
+        credentials_file (pathlib.Path): The path to file used to save
+        temporary access tokens
+
+    Returns:
+        dict: The credentials including the access token.
+    """
+    creds = {}
+    with open(credentials_file) as file:
+        data = json.load(file)
+        if "credentials" not in data:
+            data["credentials"] = {}
+        if base_url in data["credentials"]:
+            creds = data["credentials"][base_url]
+
+    expires_at = datetime.utcnow()
+    if "expires_at" in creds:
+        expires_at = datetime.strptime(creds["expires_at"], _expires_at_fmt)
+
+    if expires_at <= datetime.utcnow():
+        creds = _create_access_token(base_url, token_file)
+        with open(credentials_file, "w") as file:
+            data["credentials"][base_url] = creds
+            json.dump(data, file, indent=4)
+
+    return creds
+
+
+def _create_access_token(base_url: str, token_file: Path) -> dict:
+    """Exchange an identity token for an access token from the server.
+
+    Args:
+        base_url (str): The base URL of the server.
+        token_file (pathlib.Path): The path to the file containing the
+        identity token
+
+    Returns:
+        dict: The access token and its expiration.
+
+    Raises:
+        FileNotFoundError: If the token file is not found.
+        OSError: If there is an issue reading the token file.
+        AuthenticationError: If the server fails to provide an access token.
+    """
+    try:
+        with open(token_file) as file:
+            token = file.read().strip()
+    except FileNotFoundError as e:
+        raise FileNotFoundError(f"Identity token file not found: {token_file}") from e
+    except OSError as e:
+        raise OSError(
+            f"Failed to read the identity token from file: {token_file}"
+        ) from e
+
+    url = f"{base_url}/oidc/token"
+    data = {
+        "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+        "assertion": token,
+    }
+    headers = {"Content-Type": "application/x-www-form-urlencoded"}
+
+    response = requests.post(url, data=data, headers=headers)
+
+    if response.status_code != 200:
+        raise AuthenticationError(
+            f"Failed to retrieve access token: {response.status_code}, {response.text}"
+        )
+
+    resp_json = response.json()
+    expires_at = datetime.utcnow() + timedelta(seconds=float(resp_json["expires_in"]))
+    resp_json["expires_at"] = expires_at.strftime(_expires_at_fmt)
+    del resp_json["expires_in"]
+
+    return resp_json

wandb/sdk/wandb_init.py

@@ -323,6 +323,15 @@ class _WandbInit:
         if save_code_pre_user_settings is False:
             settings.update({"save_code": False}, source=Source.INIT)
 
+        # TODO: remove this once we refactor the client. This is a temporary
+        # fix to make sure that we use the same project name for wandb-core.
+        # The reason this is not going throught the settings object is to
+        # avoid failure cases in other parts of the code that will be
+        # removed with the switch to wandb-core.
+        if settings.project is None:
+            project = wandb.util.auto_project_name(settings.program)
+            settings.update({"project": project}, source=Source.INIT)
+
         # TODO(jhr): should this be moved? probably.
         settings._set_run_start_time(source=Source.INIT)
 
@@ -989,8 +998,9 @@ def init(
 
     Arguments:
         project: (str, optional) The name of the project where you're sending
-            the new run. If the project is not specified, the run is put in an
-            "Uncategorized" project.
+            the new run. If the project is not specified, we will try to infer
+            the project name from git root or the current program file. If we
+            can't infer the project name, we will default to `"uncategorized"`.
         entity: (str, optional) An entity is a username or team name where
             you're sending runs. This entity must exist before you can send runs
             there, so make sure to create your account or team in the UI before

wandb/sdk/wandb_login.py

@@ -156,6 +156,9 @@ class _WandbLogin:
         """Returns whether an API key is set or can be inferred."""
         return apikey.api_key(settings=self._settings) is not None
 
+    def should_use_identity_token(self):
+        return self._settings.identity_token_file is not None
+
     def set_backend(self, backend):
         self._backend = backend
 
@@ -327,6 +330,9 @@ def _login(
         )
         return False
 
+    if wlogin.should_use_identity_token():
+        return True
+
     # perform a login
     logged_in = wlogin.login()
 

wandb/sdk/wandb_manager.py

@@ -127,6 +127,7 @@ class _Manager:
             raise ManagerConnectionError(f"Connection to wandb service failed: {e}")
 
     def __init__(self, settings: "Settings") -> None:
+        """Connects to the internal service, starting it if necessary."""
         from wandb.sdk.service import service
 
         self._settings = settings
@@ -134,6 +135,7 @@ class _Manager:
         self._hooks = None
 
         self._service = service._Service(settings=self._settings)
+
         token = _ManagerToken.from_environment()
         if not token:
             self._service.start()
@@ -144,7 +146,6 @@ class _Manager:
             token = _ManagerToken.from_params(transport=transport, host=host, port=port)
             token.set_environment()
             self._atexit_setup()
-
         self._token = token
 
         try:
@@ -152,6 +153,24 @@ class _Manager:
         except ManagerConnectionError as e:
             wandb._sentry.reraise(e)
 
+    def _teardown(self, exit_code: int) -> int:
+        """Shuts down the internal process and returns its exit code.
+
+        This sends a teardown record to the process. An exception is raised if
+        the process has already been shut down.
+        """
+        unregister_all_post_import_hooks()
+
+        if self._atexit_lambda:
+            atexit.unregister(self._atexit_lambda)
+            self._atexit_lambda = None
+
+        try:
+            self._inform_teardown(exit_code)
+            return self._service.join()
+        finally:
+            self._token.reset_environment()
+
     def _atexit_setup(self) -> None:
         self._atexit_lambda = lambda: self._atexit_teardown()
 
@@ -161,28 +180,18 @@ class _Manager:
 
     def _atexit_teardown(self) -> None:
         trigger.call("on_finished")
-        exit_code = self._hooks.exit_code if self._hooks else 0
-        self._teardown(exit_code)
-
-    def _teardown(self, exit_code: int) -> None:
-        unregister_all_post_import_hooks()
 
-        if self._atexit_lambda:
-            atexit.unregister(self._atexit_lambda)
-            self._atexit_lambda = None
+        # Clear the atexit hook---we're executing it now, after which the
+        # process will exit.
+        self._atexit_lambda = None
 
         try:
-            self._inform_teardown(exit_code)
-            result = self._service.join()
-            if result and not self._settings._notebook:
-                os._exit(result)
+            self._teardown(self._hooks.exit_code if self._hooks else 0)
         except Exception as e:
             wandb.termlog(
-                f"While tearing down the service manager. The following error has occurred: {e}",
+                f"Encountered an error while tearing down the service manager: {e}",
                 repeat=False,
             )
-        finally:
-            self._token.reset_environment()
 
     def _get_service(self) -> "service._Service":
         return self._service

wandb/sdk/wandb_run.py

@@ -2751,12 +2751,14 @@ class Run:
             summary_items = [s.lower() for s in summary.split(",")]
             summary_ops = []
             valid = {"min", "max", "mean", "best", "last", "copy", "none"}
+            # TODO: deprecate copy and best
             for i in summary_items:
                 if i not in valid:
                     raise wandb.Error(f"Unhandled define_metric() arg: summary op: {i}")
                 summary_ops.append(i)
             with telemetry.context(run=self) as tel:
                 tel.feature.metric_summary = True
+        # TODO: deprecate goal
         goal_cleaned: Optional[str] = None
         if goal is not None:
             goal_cleaned = goal[:3].lower()
@@ -2772,6 +2774,9 @@ class Run:
             with telemetry.context(run=self) as tel:
                 tel.feature.metric_step_sync = True
 
+        with telemetry.context(run=self) as tel:
+            tel.feature.metric = True
+
         m = wandb_metric.Metric(
             name=name,
             step_metric=step_metric,
@@ -2783,8 +2788,6 @@ class Run:
         )
         m._set_callback(self._metric_callback)
         m._commit()
-        with telemetry.context(run=self) as tel:
-            tel.feature.metric = True
         return m
 
     # TODO(jhr): annotate this