wandb 0.16.4__py3-none-any.whl → 0.16.5__py3-none-any.whl

wandb/sdk/internal/file_stream.py

@@ -1,6 +1,7 @@
 import base64
 import functools
 import itertools
+import json
 import logging
 import os
 import queue
@@ -58,6 +59,7 @@ class Chunk(NamedTuple):
 class DefaultFilePolicy:
     def __init__(self, start_chunk_id: int = 0) -> None:
         self._chunk_id = start_chunk_id
+        self.has_debug_log = False
 
     def process_chunks(
         self, chunks: List[Chunk]
@@ -66,6 +68,21 @@ class DefaultFilePolicy:
         self._chunk_id += len(chunks)
         return {"offset": chunk_id, "content": [c.data for c in chunks]}
 
+    # TODO: this is very inefficient, this is meant for temporary debugging and will be removed in future releases
+    def _debug_log(self, data: Any):
+        if self.has_debug_log or not os.environ.get("WANDB_DEBUG_FILESTREAM_LOG"):
+            return
+
+        loaded = json.loads(data)
+        if not isinstance(loaded, dict):
+            return
+
+        # get key size and convert to MB
+        key_sizes = [(k, len(json.dumps(v))) for k, v in loaded.items()]
+        key_msg = [f"{k}: {v/1048576:.5f} MB" for k, v in key_sizes]
+        wandb.termerror(f"Step: {loaded['_step']} | {key_msg}", repeat=False)
+        self.has_debug_log = True
+
 
 class JsonlFilePolicy(DefaultFilePolicy):
     def process_chunks(self, chunks: List[Chunk]) -> "ProcessedChunk":
@@ -81,6 +98,7 @@ class JsonlFilePolicy(DefaultFilePolicy):
                 )
                 wandb.termerror(msg, repeat=False)
                 wandb._sentry.message(msg, repeat=False)
+                self._debug_log(chunk.data)
             else:
                 chunk_data.append(chunk.data)
 
@@ -99,6 +117,7 @@ class SummaryFilePolicy(DefaultFilePolicy):
             )
             wandb.termerror(msg, repeat=False)
             wandb._sentry.message(msg, repeat=False)
+            self._debug_log(data)
             return False
         return {"offset": 0, "content": [data]}
 

wandb/sdk/internal/handler.py

@@ -689,7 +689,7 @@ class HandleManager:
             self._settings, interface=self._interface, run_proto=run_start.run
         )
 
-        if run_start.run.resumed:
+        if run_start.run.resumed or run_start.run.forked:
             self._step = run_start.run.starting_step
         result = proto_util._result_from_record(record)
         self._respond_result(result)
@@ -862,9 +862,6 @@ class HandleManager:
         self._respond_result(result)
         self._stopped.set()
 
-    def handle_request_job_info(self, record: Record) -> None:
-        self._dispatch_record(record, always_send=True)
-
     def finish(self) -> None:
         logger.info("shutting down handler")
         if self._system_monitor is not None:

wandb/sdk/internal/internal_api.py

@@ -2150,6 +2150,7 @@ class Api:
                             name
                         }
                     }
+                    historyLineCount
                 }
                 inserted
                 _Server_Settings_
@@ -2237,6 +2238,7 @@ class Api:
                 .get("serverSettings", {})
                 .get("serverMessages", [])
             )
+
         return (
             response["upsertBucket"]["bucket"],
             response["upsertBucket"]["inserted"],

wandb/sdk/internal/job_builder.py

@@ -4,7 +4,7 @@ import logging
 import os
 import re
 import sys
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
 import wandb
 from wandb.sdk.artifacts.artifact import Artifact
@@ -28,6 +28,8 @@ FROZEN_REQUIREMENTS_FNAME = "requirements.frozen.txt"
 JOB_FNAME = "wandb-job.json"
 JOB_ARTIFACT_TYPE = "job"
 
+LOG_LEVEL = Literal["log", "warn", "error"]
+
 
 class GitInfo(TypedDict):
     remote: str
@@ -89,8 +91,9 @@ class JobBuilder:
     _job_seq_id: Optional[str]
     _job_version_alias: Optional[str]
     _is_notebook_run: bool
+    _verbose: bool
 
-    def __init__(self, settings: SettingsStatic):
+    def __init__(self, settings: SettingsStatic, verbose: bool = False):
         self._settings = settings
         self._metadatafile_path = None
         self._requirements_path = None
@@ -106,6 +109,7 @@ class JobBuilder:
             Literal["repo", "artifact", "image"]
         ] = settings.job_source  # type: ignore[assignment]
         self._is_notebook_run = self._get_is_notebook_run()
+        self._verbose = verbose
 
     def set_config(self, config: Dict[str, Any]) -> None:
         self._config = config
@@ -197,6 +201,21 @@ class JobBuilder:
 
         return source, name
 
+    def _log_if_verbose(self, message: str, level: LOG_LEVEL) -> None:
+        log_func: Optional[Union[Callable[[Any], None], Callable[[Any], None]]] = None
+        if level == "log":
+            _logger.info(message)
+            log_func = wandb.termlog
+        elif level == "warn":
+            _logger.warning(message)
+            log_func = wandb.termwarn
+        elif level == "error":
+            _logger.error(message)
+            log_func = wandb.termerror
+
+        if self._verbose and log_func is not None:
+            log_func(message)
+
     def _build_artifact_job_source(
         self,
         program_relpath: str,
@@ -212,8 +231,9 @@ class JobBuilder:
                 # at the directory the notebook is in instead of the jupyter core
                 if not os.path.exists(os.path.basename(program_relpath)):
                     _logger.info("target path does not exist, exiting")
-                    wandb.termwarn(
-                        "No program path found when generating artifact job source for a non-colab notebook run. See https://docs.wandb.ai/guides/launch/create-job"
+                    self._log_if_verbose(
+                        "No program path found when generating artifact job source for a non-colab notebook run. See https://docs.wandb.ai/guides/launch/create-job",
+                        "warn",
                     )
                     return None, None
                 full_program_relpath = os.path.basename(program_relpath)
@@ -299,22 +319,25 @@ class JobBuilder:
         if not os.path.exists(
             os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME)
         ):
-            wandb.termwarn(
-                "No requirements.txt found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job"
+            self._log_if_verbose(
+                "No requirements.txt found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
+                "warn",
             )
             return None
         metadata = self._handle_metadata_file()
         if metadata is None:
-            wandb.termwarn(
-                f"Ensure read and write access to run files dir: {self._settings.files_dir}, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables"
+            self._log_if_verbose(
+                f"Ensure read and write access to run files dir: {self._settings.files_dir}, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables",
+                "warn",
             )
             return None
 
         runtime: Optional[str] = metadata.get("python")
         # can't build a job without a python version
         if runtime is None:
-            wandb.termwarn(
-                "No python version found in metadata, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job"
+            self._log_if_verbose(
+                "No python version found in metadata, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
+                "warn",
             )
             return None
 
@@ -345,13 +368,16 @@ class JobBuilder:
                     or self._settings.job_source
                     or self._source_type
                 ):
-                    wandb.termwarn("No source type found, not creating job artifact")
+                    self._log_if_verbose(
+                        "No source type found, not creating job artifact", "warn"
+                    )
                 return None
 
             program_relpath = self._get_program_relpath(source_type, metadata)
             if source_type != "image" and not program_relpath:
-                wandb.termwarn(
-                    "No program path found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job"
+                self._log_if_verbose(
+                    "No program path found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
+                    "warn",
                 )
                 return None
 
@@ -377,10 +403,11 @@ class JobBuilder:
 
             if source is None:
                 if source_type:
-                    wandb.termwarn(
+                    self._log_if_verbose(
                         f"Source type is set to '{source_type}' but some required information is missing "
                         "from the environment. A job will not be created from this run. See "
-                        "https://docs.wandb.ai/guides/launch/create-job"
+                        "https://docs.wandb.ai/guides/launch/create-job",
+                        "warn",
                     )
                 return None
 
@@ -447,8 +474,9 @@ class JobBuilder:
             program = metadata.get("program")
 
             if not program:
-                wandb.termwarn(
-                    "Notebook 'program' path not found in metadata. See https://docs.wandb.ai/guides/launch/create-job"
+                self._log_if_verbose(
+                    "Notebook 'program' path not found in metadata. See https://docs.wandb.ai/guides/launch/create-job",
+                    "warn",
                 )
 
             return program

wandb/sdk/internal/sender.py

@@ -115,6 +115,7 @@ def _manifest_json_from_proto(manifest: "ArtifactManifest") -> Dict:
                 "ref": content.ref if content.ref else None,
                 "size": content.size if content.size is not None else None,
                 "local_path": content.local_path if content.local_path else None,
+                "skip_cache": content.skip_cache,
                 "extra": {
                     extra.key: json.loads(extra.value_json) for extra in content.extra
                 },
@@ -733,18 +734,7 @@ class SendManager:
             )
         self._respond_result(result)
 
-    def send_request_job_info(self, record: "Record") -> None:
-        """Respond to a request for a job link."""
-        result = proto_util._result_from_record(record)
-        result.response.job_info_response.sequenceId = (
-            self._job_builder._job_seq_id or ""
-        )
-        result.response.job_info_response.version = (
-            self._job_builder._job_version_alias or ""
-        )
-        self._respond_result(result)
-
-    def _maybe_setup_resume(
+    def _setup_resume(
         self, run: "RunRecord"
     ) -> Optional["wandb_internal_pb2.ErrorInfo"]:
         """Queries the backend for a run; fail if the settings are incompatible."""
@@ -890,6 +880,30 @@ class SendManager:
             pass
         # TODO: do something if sync spell is not successful?
 
+    def _setup_fork(self, server_run: dict):
+        assert self._settings.fork_from
+        assert self._settings.fork_from.metric == "_step"
+        assert self._run
+        first_step = int(self._settings.fork_from.value) + 1
+        self._resume_state.step = first_step
+        self._resume_state.history = server_run.get("historyLineCount", 0)
+        self._run.forked = True
+        self._run.starting_step = first_step
+
+    def _handle_error(
+        self,
+        record: "Record",
+        error: "wandb_internal_pb2.ErrorInfo",
+        run: "RunRecord",
+    ) -> None:
+        if record.control.req_resp or record.control.mailbox_slot:
+            result = proto_util._result_from_record(record)
+            result.run_result.run.CopyFrom(run)
+            result.run_result.error.CopyFrom(error)
+            self._respond_result(result)
+        else:
+            logger.error("Got error in async mode: %s", error.message)
+
     def send_run(self, record: "Record", file_dir: Optional[str] = None) -> None:
         run = record.run
         error = None
@@ -911,21 +925,28 @@ class SendManager:
             config_value_dict = self._config_backend_dict()
             self._config_save(config_value_dict)
 
+        do_fork = self._settings.fork_from is not None and is_wandb_init
+        do_resume = bool(self._settings.resume)
+
+        if do_fork and do_resume:
+            error = wandb_internal_pb2.ErrorInfo()
+            error.code = wandb_internal_pb2.ErrorInfo.ErrorCode.USAGE
+            error.message = (
+                "You cannot use `resume` and `fork_from` together. Please choose one."
+            )
+            self._handle_error(record, error, run)
+
         if is_wandb_init:
             # Ensure we have a project to query for status
             if run.project == "":
                 run.project = util.auto_project_name(self._settings.program)
             # Only check resume status on `wandb.init`
-            error = self._maybe_setup_resume(run)
+
+            if do_resume:
+                error = self._setup_resume(run)
 
         if error is not None:
-            if record.control.req_resp or record.control.mailbox_slot:
-                result = proto_util._result_from_record(record)
-                result.run_result.run.CopyFrom(run)
-                result.run_result.error.CopyFrom(error)
-                self._respond_result(result)
-            else:
-                logger.error("Got error in async mode: %s", error.message)
+            self._handle_error(record, error, run)
             return
 
         # Save the resumed config
@@ -945,19 +966,22 @@ class SendManager:
             self._config_save(config_value_dict)
 
         try:
-            self._init_run(run, config_value_dict)
+            server_run = self._init_run(run, config_value_dict)
         except (CommError, UsageError) as e:
             logger.error(e, exc_info=True)
-            if record.control.req_resp or record.control.mailbox_slot:
-                result = proto_util._result_from_record(record)
-                result.run_result.run.CopyFrom(run)
-                error = ProtobufErrorHandler.from_exception(e)
-                result.run_result.error.CopyFrom(error)
-                self._respond_result(result)
+            error = ProtobufErrorHandler.from_exception(e)
+            self._handle_error(record, error, run)
             return
 
         assert self._run  # self._run is configured in _init_run()
 
+        if do_fork:
+            error = self._setup_fork(server_run)
+
+        if error is not None:
+            self._handle_error(record, error, run)
+            return
+
         if record.control.req_resp or record.control.mailbox_slot:
             result = proto_util._result_from_record(record)
             # TODO: we could do self._interface.publish_defer(resp) to notify
@@ -976,7 +1000,7 @@ class SendManager:
         self,
         run: "RunRecord",
         config_dict: Optional[sender_config.BackendConfigDict],
-    ) -> None:
+    ) -> dict:
         # We subtract the previous runs runtime when resuming
         start_time = (
             run.start_time.ToMicroseconds() / 1e6
@@ -1061,6 +1085,7 @@ class SendManager:
             self._run.sweep_id = sweep_id
         if os.getenv("SPELL_RUN_URL"):
             self._sync_spell()
+        return server_run
 
     def _start_run_threads(self, file_dir: Optional[str] = None) -> None:
         assert self._run  # self._run is configured by caller

wandb/sdk/internal/settings_static.py

@@ -2,6 +2,7 @@ from dataclasses import fields
 from typing import Any, Iterable, Sequence, Tuple
 
 from wandb.proto import wandb_settings_pb2
+from wandb.sdk.lib import RunMoment
 from wandb.sdk.wandb_settings import SettingsData
 
 
@@ -38,6 +39,14 @@ class SettingsStatic(SettingsData):
                             unpacked_inner[inner_key] = inner_value
                         unpacked_mapping[outer_key] = unpacked_inner
                     value = unpacked_mapping
+            elif key == "fork_from":
+                value = getattr(proto, key)
+                if value.run:
+                    value = RunMoment(
+                        run=value.run, value=value.value, metric=value.metric
+                    )
+                else:
+                    value = None
             else:
                 if proto.HasField(key):  # type: ignore [arg-type]
                     value = getattr(proto, key).value

wandb/sdk/internal/system/system_info.py

@@ -212,7 +212,10 @@ class SystemInfo:
                 os.path.join(self.settings.files_dir, CONDA_ENVIRONMENTS_FNAME), "w"
             ) as f:
                 subprocess.call(
-                    ["conda", "env", "export"], stdout=f, stderr=subprocess.DEVNULL
+                    ["conda", "env", "export"],
+                    stdout=f,
+                    stderr=subprocess.DEVNULL,
+                    timeout=15,  # add timeout since conda env export could take a really long time
                 )
         except Exception as e:
             logger.exception(f"Error saving conda packages: {e}")

wandb/sdk/launch/create_job.py

@@ -396,6 +396,7 @@ def _configure_job_builder_for_partial(tmpdir: str, job_source: str) -> JobBuild
     settings.update({"files_dir": tmpdir, "job_source": job_source})
     job_builder = JobBuilder(
         settings=settings,  # type: ignore
+        verbose=True,
     )
     # never allow notebook runs
     job_builder._is_notebook_run = False

wandb/sdk/launch/runner/kubernetes_runner.py

@@ -1,6 +1,7 @@
 """Implementation of KubernetesRunner class for wandb launch."""
 import asyncio
 import base64
+import datetime
 import json
 import logging
 import os
@@ -23,6 +24,7 @@ from wandb.sdk.launch.runner.kubernetes_monitor import (
     CustomResource,
     LaunchKubernetesMonitor,
 )
+from wandb.sdk.lib.retry import ExponentialBackoff, retry_async
 from wandb.util import get_module
 
 from .._project_spec import EntryPoint, LaunchProject
@@ -59,6 +61,7 @@ from kubernetes_asyncio.client.models.v1_secret import (  # type: ignore # noqa:
 from kubernetes_asyncio.client.rest import ApiException  # type: ignore # noqa: E402
 
 TIMEOUT = 5
+API_KEY_SECRET_MAX_RETRIES = 5
 
 _logger = logging.getLogger(__name__)
 
@@ -421,8 +424,23 @@ class KubernetesRunner(AbstractRunner):
                     else:
                         secret_name += f"-{launch_project.run_id}"
 
-                    api_key_secret = await ensure_api_key_secret(
-                        core_api, secret_name, namespace, value
+                    def handle_exception(e):
+                        wandb.termwarn(
+                            f"Exception when ensuring Kubernetes API key secret: {e}. Retrying..."
+                        )
+
+                    api_key_secret = await retry_async(
+                        backoff=ExponentialBackoff(
+                            initial_sleep=datetime.timedelta(seconds=1),
+                            max_sleep=datetime.timedelta(minutes=1),
+                            max_retries=API_KEY_SECRET_MAX_RETRIES,
+                        ),
+                        fn=ensure_api_key_secret,
+                        on_exc=handle_exception,
+                        core_api=core_api,
+                        secret_name=secret_name,
+                        namespace=namespace,
+                        api_key=value,
                     )
                     env.append(
                         {

wandb/sdk/launch/utils.py

@@ -222,14 +222,14 @@ def get_default_entity(api: Api, launch_config: Optional[Dict[str, Any]]):
 
 
 def strip_resource_args_and_template_vars(launch_spec: Dict[str, Any]) -> None:
-    wandb.termwarn(
-        "Launch spec contains both resource_args and template_variables, "
-        "only one can be set. Using template_variables."
-    )
     if launch_spec.get("resource_args", None) and launch_spec.get(
         "template_variables", None
     ):
-        launch_spec["resource_args"] = None
+        wandb.termwarn(
+            "Launch spec contains both resource_args and template_variables, "
+            "only one can be set. Using template_variables."
+        )
+        launch_spec.pop("resource_args")
 
 
 def construct_launch_spec(

wandb/sdk/lib/__init__.py

@@ -1,8 +1,5 @@
 from . import lazyloader
 from .disabled import RunDisabled, SummaryDisabled
+from .run_moment import RunMoment
 
-__all__ = (
-    "lazyloader",
-    "RunDisabled",
-    "SummaryDisabled",
-)
+__all__ = ("lazyloader", "RunDisabled", "SummaryDisabled", "RunMoment")

wandb/sdk/lib/_settings_toposort_generated.py

@@ -102,6 +102,7 @@ _Setting = Literal[
     "entity",
     "files_dir",
     "force",
+    "fork_from",
     "git_commit",
     "git_remote",
     "git_remote_url",

wandb/sdk/lib/filesystem.py

@@ -227,7 +227,17 @@ def safe_open(
 
 
 def safe_copy(source_path: StrPath, target_path: StrPath) -> StrPath:
-    """Copy a file, ensuring any changes only apply atomically once finished."""
+    """Copy a file atomically.
+
+    Copying is not usually atomic, and on operating systems that allow multiple
+    writers to the same file, the result can get corrupted. If two writers copy
+    to the same file, the contents can become interleaved.
+
+    We mitigate the issue somewhat by copying to a temporary file first and
+    then renaming. Renaming is atomic: if process 1 renames file A to X and
+    process 2 renames file B to X, then X will either contain the contents
+    of A or the contents of B, not some mixture of both.
+    """
     # TODO (hugh): check that there is enough free space.
     output_path = Path(target_path).resolve()
     output_path.parent.mkdir(parents=True, exist_ok=True)

wandb/sdk/lib/run_moment.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass
+from typing import Literal, Union, cast
+from urllib import parse
+
+_STEP = Literal["_step"]
+
+
+@dataclass
+class RunMoment:
+    """A moment in a run."""
+
+    run: str  # run name
+
+    # currently, the _step value to fork from. in future, this will be optional
+    value: Union[int, float]
+
+    # only step for now, in future this will be relaxed to be any metric
+    metric: _STEP = "_step"
+
+    def __post_init__(self):
+        if self.metric != "_step":
+            raise ValueError(
+                f"Only the metric '_step' is supported, got '{self.metric}'."
+            )
+        if not isinstance(self.value, (int, float)):
+            raise ValueError(
+                f"Only int or float values are supported, got '{self.value}'."
+            )
+        if not isinstance(self.run, str):
+            raise ValueError(f"Only string run names are supported, got '{self.run}'.")
+
+    @classmethod
+    def from_uri(cls, uri: str) -> "RunMoment":
+        parsable = "runmoment://" + uri
+        parse_err = ValueError(
+            f"Could not parse passed run moment string '{uri}', "
+            f"expected format '<run>?<metric>=<numeric_value>'. "
+            f"Currently, only the metric '_step' is supported. "
+            f"Example: 'ans3bsax?_step=123'."
+        )
+
+        try:
+            parsed = parse.urlparse(parsable)
+        except ValueError as e:
+            raise parse_err from e
+
+        if parsed.scheme != "runmoment":
+            raise parse_err
+
+        # extract run, metric, value from parsed
+        if not parsed.netloc:
+            raise parse_err
+
+        run = parsed.netloc
+
+        if parsed.path or parsed.params or parsed.fragment:
+            raise parse_err
+
+        query = parse.parse_qs(parsed.query)
+        if len(query) != 1:
+            raise parse_err
+
+        metric = list(query.keys())[0]
+        if metric != "_step":
+            raise parse_err
+        value: str = query[metric][0]
+        try:
+            num_value = int(value) if value.isdigit() else float(value)
+        except ValueError as e:
+            raise parse_err from e
+
+        return cls(run=run, metric=cast(_STEP, metric), value=num_value)

wandb/sdk/service/streams.py

@@ -5,6 +5,7 @@ StreamRecord: All the external state for the internal thread (queues, etc)
 StreamAction: Lightweight record for stream ops for thread safety
 StreamMux: Container for dictionary of stream threads per runid
 """
+
 import functools
 import multiprocessing
 import queue
@@ -327,7 +328,6 @@ class StreamMux:
             result = internal_messages_handle.wait(timeout=-1)
             assert result
             internal_messages_response = result.response.internal_messages_response
-            job_info_handle = stream.interface.deliver_request_job_info()
 
             # wait for them, it's ok to do this serially but this can be improved
             result = poll_exit_handle.wait(timeout=-1)
@@ -346,17 +346,12 @@ class StreamMux:
             assert result
             final_summary = result.response.get_summary_response
 
-            result = job_info_handle.wait(timeout=-1)
-            assert result
-            job_info = result.response.job_info_response
-
             Run._footer(
                 sampled_history=sampled_history,
                 final_summary=final_summary,
                 poll_exit_response=poll_exit_response,
                 server_info_response=server_info_response,
                 internal_messages_response=internal_messages_response,
-                job_info=job_info,
                 settings=stream._settings,  # type: ignore
                 printer=printer,
             )