wandb 0.17.3__py3-none-any.whl → 0.17.5__py3-none-any.whl

wandb/sdk/launch/builder/kaniko_builder.py

@@ -263,11 +263,17 @@ class KanikoBuilder(AbstractBuilder):
         repo_uri = await self.registry.get_repo_uri()
         image_uri = repo_uri + ":" + image_tag
 
-        if (
-            not launch_project.build_required()
-            and await self.registry.check_image_exists(image_uri)
-        ):
-            return image_uri
+        # The DOCKER_CONFIG_SECRET option is mutually exclusive with the
+        # registry classes, so we must skip the check for image existence in
+        # that case.
+        if not launch_project.build_required():
+            if DOCKER_CONFIG_SECRET:
+                wandb.termlog(
+                    f"Skipping check for existing image {image_uri} due to custom dockerconfig."
+                )
+            else:
+                if await self.registry.check_image_exists(image_uri):
+                    return image_uri
 
         _logger.info(f"Building image {image_uri}...")
         _, api_client = await get_kube_context_and_api_client(
@@ -286,7 +292,12 @@ class KanikoBuilder(AbstractBuilder):
         wandb.termlog(f"{LOG_PREFIX}Created kaniko job {build_job_name}")
 
         try:
-            if isinstance(self.registry, AzureContainerRegistry):
+            # DOCKER_CONFIG_SECRET is a user provided dockerconfigjson. Skip our
+            # dockerconfig handling if it's set.
+            if (
+                isinstance(self.registry, AzureContainerRegistry)
+                and not DOCKER_CONFIG_SECRET
+            ):
                 dockerfile_config_map = client.V1ConfigMap(
                     metadata=client.V1ObjectMeta(
                         name=f"docker-config-{build_job_name}"
@@ -344,7 +355,10 @@ class KanikoBuilder(AbstractBuilder):
         finally:
             wandb.termlog(f"{LOG_PREFIX}Cleaning up resources")
             try:
-                if isinstance(self.registry, AzureContainerRegistry):
+                if (
+                    isinstance(self.registry, AzureContainerRegistry)
+                    and not DOCKER_CONFIG_SECRET
+                ):
                     await core_v1.delete_namespaced_config_map(
                         f"docker-config-{build_job_name}", "wandb"
                     )
@@ -498,7 +512,10 @@ class KanikoBuilder(AbstractBuilder):
                     "readOnly": True,
                 }
             )
-        if isinstance(self.registry, AzureContainerRegistry):
+        if (
+            isinstance(self.registry, AzureContainerRegistry)
+            and not DOCKER_CONFIG_SECRET
+        ):
             # Add the docker config map
             volumes.append(
                 {
@@ -533,7 +550,11 @@ class KanikoBuilder(AbstractBuilder):
         # Apply the rest of our defaults
         pod_labels["wandb"] = "launch"
         # This annotation is required to enable azure workload identity.
-        if isinstance(self.registry, AzureContainerRegistry):
+        # Don't add this label if using a docker config secret for auth.
+        if (
+            isinstance(self.registry, AzureContainerRegistry)
+            and not DOCKER_CONFIG_SECRET
+        ):
             pod_labels["azure.workload.identity/use"] = "true"
         pod_spec["restartPolicy"] = pod_spec.get("restartPolicy", "Never")
         pod_spec["activeDeadlineSeconds"] = pod_spec.get(

wandb/sdk/launch/inputs/internal.py

@@ -11,7 +11,7 @@ import os
 import pathlib
 import shutil
 import tempfile
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
 import wandb
 import wandb.data_types
@@ -62,11 +62,13 @@ class JobInputArguments:
         self,
         include: Optional[List[str]] = None,
         exclude: Optional[List[str]] = None,
+        schema: Optional[dict] = None,
         file_path: Optional[str] = None,
         run_config: Optional[bool] = None,
     ):
         self.include = include
         self.exclude = exclude
+        self.schema = schema
         self.file_path = file_path
         self.run_config = run_config
 
@@ -121,15 +123,66 @@ def _publish_job_input(
         exclude_paths=[_split_on_unesc_dot(path) for path in input.exclude]
         if input.exclude
         else [],
+        input_schema=input.schema,
         run_config=input.run_config,
         file_path=input.file_path or "",
     )
 
 
+def _replace_refs_and_allofs(schema: dict, defs: dict) -> dict:
+    """Recursively fix JSON schemas with common issues.
+
+    1. Replaces any instances of $ref with their associated definition in defs
+    2. Removes any "allOf" lists that only have one item, "lifting" the item up
+    See test_internal.py for examples
+    """
+    ret: Dict[str, Any] = {}
+    if "$ref" in schema:
+        # Reference found, replace it with its definition
+        def_key = schema["$ref"].split("#/$defs/")[1]
+        # Also run recursive replacement in case a ref contains more refs
+        return _replace_refs_and_allofs(defs.pop(def_key), defs)
+    for key, val in schema.items():
+        if isinstance(val, dict):
+            # Step into dicts recursively
+            new_val_dict = _replace_refs_and_allofs(val, defs)
+            ret[key] = new_val_dict
+        elif isinstance(val, list):
+            # Step into each item in the list
+            new_val_list = []
+            for item in val:
+                if isinstance(item, dict):
+                    new_val_list.append(_replace_refs_and_allofs(item, defs))
+                else:
+                    new_val_list.append(item)
+            # Lift up allOf blocks with only one item
+            if (
+                key == "allOf"
+                and len(new_val_list) == 1
+                and isinstance(new_val_list[0], dict)
+            ):
+                ret.update(new_val_list[0])
+            else:
+                ret[key] = new_val_list
+        else:
+            # For anything else (str, int, etc) keep it as-is
+            ret[key] = val
+    return ret
+
+
+def _convert_pydantic_model_to_jsonschema(model: Any) -> dict:
+    schema = model.model_json_schema()
+    defs = schema.pop("$defs")
+    if not defs:
+        return schema
+    return _replace_refs_and_allofs(schema, defs)
+
+
 def handle_config_file_input(
     path: str,
     include: Optional[List[str]] = None,
     exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     """Declare an overridable configuration file for a launch job.
 
@@ -151,9 +204,20 @@ def handle_config_file_input(
         path,
         dest,
     )
+    # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
+    # or the BaseModel class itself (e.g. schema=MySchema)
+    if hasattr(schema, "model_json_schema") and callable(
+        schema.model_json_schema  # type: ignore
+    ):
+        schema = _convert_pydantic_model_to_jsonschema(schema)
+    if schema and not isinstance(schema, dict):
+        raise LaunchError(
+            "schema must be a dict, Pydantic model instance, or Pydantic model class."
+        )
     arguments = JobInputArguments(
         include=include,
         exclude=exclude,
+        schema=schema,
         file_path=path,
         run_config=False,
     )
@@ -165,7 +229,9 @@ def handle_config_file_input(
 
 
 def handle_run_config_input(
-    include: Optional[List[str]] = None, exclude: Optional[List[str]] = None
+    include: Optional[List[str]] = None,
+    exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     """Declare wandb.config as an overridable configuration for a launch job.
 
@@ -175,9 +241,20 @@ def handle_run_config_input(
     If there is no active run, the include and exclude paths are staged and sent
     when a run is created.
     """
+    # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
+    # or the BaseModel class itself (e.g. schema=MySchema)
+    if hasattr(schema, "model_json_schema") and callable(
+        schema.model_json_schema  # type: ignore
+    ):
+        schema = _convert_pydantic_model_to_jsonschema(schema)
+    if schema and not isinstance(schema, dict):
+        raise LaunchError(
+            "schema must be a dict, Pydantic model instance, or Pydantic model class."
+        )
     arguments = JobInputArguments(
         include=include,
         exclude=exclude,
+        schema=schema,
         run_config=True,
         file_path=None,
     )

wandb/sdk/launch/inputs/manage.py

@@ -1,12 +1,13 @@
 """Functions for declaring overridable configuration for launch jobs."""
 
-from typing import List, Optional
+from typing import Any, List, Optional
 
 
 def manage_config_file(
     path: str,
     include: Optional[List[str]] = None,
     exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     r"""Declare an overridable configuration file for a launch job.
 
@@ -43,18 +44,27 @@ def manage_config_file(
             relative and must not contain backwards traversal, i.e. `..`.
         include (List[str]): A list of keys to include in the configuration file.
         exclude (List[str]): A list of keys to exclude from the configuration file.
+        schema (dict | Pydantic model): A JSON Schema or Pydantic model describing
+            describing which attributes will be editable from the Launch drawer.
+            Accepts both an instance of a Pydantic BaseModel class or the BaseModel
+            class itself.
 
     Raises:
         LaunchError: If the path is not valid, or if there is no active run.
     """
+    # note: schema's Any type is because in the case where a BaseModel class is
+    # provided, its type is a pydantic internal type that we don't want our typing
+    # to depend on. schema's type should be considered
+    # "Optional[dict | <something with a .model_json_schema() method>]"
     from .internal import handle_config_file_input
 
-    return handle_config_file_input(path, include, exclude)
+    return handle_config_file_input(path, include, exclude, schema)
 
 
 def manage_wandb_config(
     include: Optional[List[str]] = None,
     exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     r"""Declare wandb.config as an overridable configuration for a launch job.
 
@@ -86,10 +96,18 @@ def manage_wandb_config(
     Args:
         include (List[str]): A list of subtrees to include in the configuration.
         exclude (List[str]): A list of subtrees to exclude from the configuration.
+        schema (dict | Pydantic model): A JSON Schema or Pydantic model describing
+            describing which attributes will be editable from the Launch drawer.
+            Accepts both an instance of a Pydantic BaseModel class or the BaseModel
+            class itself.
 
     Raises:
         LaunchError: If there is no active run.
     """
+    # note: schema's Any type is because in the case where a BaseModel class is
+    # provided, its type is a pydantic internal type that we don't want our typing
+    # to depend on. schema's type should be considered
+    # "Optional[dict | <something with a .model_json_schema() method>]"
     from .internal import handle_run_config_input
 
-    handle_run_config_input(include, exclude)
+    handle_run_config_input(include, exclude, schema)

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/lib/tracelog.py

@@ -45,8 +45,8 @@ logger = logging.getLogger(__name__)
 ANNOTATE_QUEUE_NAME = "_DEBUGLOG_QUEUE_NAME"
 
 # capture stdout and stderr before anyone messes with them
-stdout_write = sys.__stdout__.write
-stderr_write = sys.__stderr__.write
+stdout_write = sys.__stdout__.write  # type: ignore
+stderr_write = sys.__stderr__.write  # type: ignore
 
 
 def _log(

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

@@ -114,19 +114,24 @@ class _Manager:
 
         try:
             svc_iface._svc_connect(port=port)
+
         except ConnectionRefusedError as e:
             if not psutil.pid_exists(self._token.pid):
                 message = (
-                    "Connection to wandb service failed "
-                    "since the process is not available. "
+                    "Connection to wandb service failed"
+                    " because the process is not available."
                 )
             else:
-                message = f"Connection to wandb service failed: {e}. "
-            raise ManagerConnectionRefusedError(message)
+                message = "Connection to wandb service failed."
+            raise ManagerConnectionRefusedError(message) from e
+
         except Exception as e:
-            raise ManagerConnectionError(f"Connection to wandb service failed: {e}")
+            raise ManagerConnectionError(
+                "Connection to wandb service failed.",
+            ) from 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 +139,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 +150,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 +157,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 +184,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