wandb 0.17.3__py3-none-win_amd64.whl → 0.17.4__py3-none-win_amd64.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()
@@ -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/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/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",
@@ -112,6 +114,7 @@ _Setting = Literal[
     "host",
     "http_proxy",
     "https_proxy",
+    "identity_token_file",
     "ignore_globs",
     "init_timeout",
     "is_local",

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_settings.py

@@ -43,7 +43,7 @@ from wandb.apis.internal import Api
 from wandb.errors import UsageError
 from wandb.proto import wandb_settings_pb2
 from wandb.sdk.internal.system.env_probe_helpers import is_aws_lambda
-from wandb.sdk.lib import filesystem
+from wandb.sdk.lib import credentials, filesystem
 from wandb.sdk.lib._settings_toposort_generated import SETTINGS_TOPOLOGICALLY_SORTED
 from wandb.sdk.lib.run_moment import RunMoment
 from wandb.sdk.wandb_setup import _EarlyLogger
@@ -314,6 +314,7 @@ class SettingsData:
     _disable_machine_info: bool  # Disable automatic machine info collection
     _executable: str
     _extra_http_headers: Mapping[str, str]
+    _file_stream_max_bytes: int  # max size for filestream requests in core
     # file stream retry client configuration
     _file_stream_retry_max: int  # max number of retries
     _file_stream_retry_wait_min_seconds: float  # min wait time between retries
@@ -391,6 +392,7 @@ class SettingsData:
     config_paths: Sequence[str]
     console: str
     console_multipart: bool  # whether to produce multipart console log files
+    credentials_file: str  # file path to write access tokens
     deployment: str
     disable_code: bool
     disable_git: bool
@@ -412,6 +414,7 @@ class SettingsData:
     host: str
     http_proxy: str  # proxy server for the http requests to W&B
     https_proxy: str  # proxy server for the https requests to W&B
+    identity_token_file: str  # file path to supply a jwt for authentication
     ignore_globs: Tuple[str]
     init_timeout: float
     is_local: bool
@@ -659,6 +662,7 @@ class Settings(SettingsData):
             _disable_update_check={"preprocessor": _str_as_bool},
             _disable_viewer={"preprocessor": _str_as_bool},
             _extra_http_headers={"preprocessor": _str_as_json},
+            _file_stream_max_bytes={"preprocessor": int},
             _file_stream_retry_max={"preprocessor": int},
             _file_stream_retry_wait_min_seconds={"preprocessor": float},
             _file_stream_retry_wait_max_seconds={"preprocessor": float},
@@ -783,6 +787,10 @@ class Settings(SettingsData):
                 "auto_hook": True,
             },
             console_multipart={"value": False, "preprocessor": _str_as_bool},
+            credentials_file={
+                "value": str(credentials.DEFAULT_WANDB_CREDENTIALS_FILE),
+                "preprocessor": str,
+            },
             deployment={
                 "hook": lambda _: "local" if self.is_local else "cloud",
                 "auto_hook": True,
@@ -829,6 +837,7 @@ class Settings(SettingsData):
                 "hook": lambda x: self._proxies and self._proxies.get("https") or x,
                 "auto_hook": True,
             },
+            identity_token_file={"value": None, "preprocessor": str},
             ignore_globs={
                 "value": tuple(),
                 "preprocessor": lambda x: tuple(x) if not isinstance(x, tuple) else x,
@@ -873,7 +882,9 @@ class Settings(SettingsData):
             program={
                 "hook": lambda x: self._get_program(x),
             },
-            project={"validator": self._validate_project},
+            project={
+                "validator": self._validate_project,
+            },
             project_url={"hook": lambda _: self._project_url(), "auto_hook": True},
             quiet={"preprocessor": _str_as_bool},
             reinit={"preprocessor": _str_as_bool},

wandb/sdk/wandb_setup.py

@@ -268,20 +268,20 @@ class _WandbSetup__WandbSetup:  # noqa: N801
                     self._config = config_dict
 
     def _teardown(self, exit_code: Optional[int] = None) -> None:
-        exit_code = exit_code or 0
-        self._teardown_manager(exit_code=exit_code)
+        if not self._manager:
+            return
+
+        internal_exit_code = self._manager._teardown(exit_code or 0)
+        self._manager = None
+
+        if internal_exit_code != 0:
+            sys.exit(internal_exit_code)
 
     def _setup_manager(self) -> None:
         if self._settings._disable_service:
             return
         self._manager = wandb_manager._Manager(settings=self._settings)
 
-    def _teardown_manager(self, exit_code: int) -> None:
-        if not self._manager:
-            return
-        self._manager._teardown(exit_code)
-        self._manager = None
-
     def _get_manager(self) -> Optional[wandb_manager._Manager]:
         return self._manager
 
@@ -312,11 +312,9 @@ def _setup(
 ) -> Optional["_WandbSetup"]:
     """Set up library context."""
     if _reset:
-        setup_instance = _WandbSetup._instance
-        if setup_instance:
-            setup_instance._teardown()
-        _WandbSetup._instance = None
+        teardown()
         return None
+
     wl = _WandbSetup(settings=settings)
     return wl
 
@@ -330,6 +328,7 @@ def setup(
 
 def teardown(exit_code: Optional[int] = None) -> None:
     setup_instance = _WandbSetup._instance
+    _WandbSetup._instance = None
+
     if setup_instance:
         setup_instance._teardown(exit_code=exit_code)
-    _WandbSetup._instance = None

{wandb-0.17.3.dist-info → wandb-0.17.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: wandb
-Version: 0.17.3
+Version: 0.17.4
 Summary: A CLI and library for interacting with the Weights & Biases API.
 Project-URL: Source, https://github.com/wandb/wandb
 Project-URL: Bug Reports, https://github.com/wandb/wandb/issues