diff --git a/lib/python3.12/site-packages/wandb/plot/__init__.py b/lib/python3.12/site-packages/wandb/plot/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c49e40ac3ce6a7ffccd97c553a013bdd97fd305b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/plot/__init__.py @@ -0,0 +1,30 @@ +"""Chart Visualization Utilities + +This module offers a collection of predefined chart types, along with functionality +for creating custom charts, enabling flexible visualization of your data beyond the +built-in options. +""" + +__all__ = [ + "line", + "histogram", + "scatter", + "bar", + "roc_curve", + "pr_curve", + "confusion_matrix", + "line_series", + "plot_table", + "visualize", # doc:exclude +] + +from wandb.plot.bar import bar +from wandb.plot.confusion_matrix import confusion_matrix +from wandb.plot.custom_chart import CustomChart, plot_table +from wandb.plot.histogram import histogram +from wandb.plot.line import line +from wandb.plot.line_series import line_series +from wandb.plot.pr_curve import pr_curve +from wandb.plot.roc_curve import roc_curve +from wandb.plot.scatter import scatter +from wandb.plot.viz import Visualize, visualize diff --git a/lib/python3.12/site-packages/wandb/plot/bar.py b/lib/python3.12/site-packages/wandb/plot/bar.py new file mode 100644 index 0000000000000000000000000000000000000000..21e6496df9e46a00204b78abb4c599ca0430a3c1 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/plot/bar.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from wandb.plot.custom_chart import plot_table + +if TYPE_CHECKING: + import wandb + from wandb.plot.custom_chart import CustomChart + + +def bar( + table: wandb.Table, + label: str, + value: str, + title: str = "", + split_table: bool = False, +) -> CustomChart: + """Constructs a bar chart from a wandb.Table of data. + + Args: + table: A table containing the data for the bar chart. + label: The name of the column to use for the labels of each bar. + value: The name of the column to use for the values of each bar. + title: The title of the bar chart. + split_table: Whether the table should be split into a separate section + in the W&B UI. If `True`, the table will be displayed in a section named + "Custom Chart Tables". Default is `False`. + + Returns: + CustomChart: A custom chart object that can be logged to W&B. To log the + chart, pass it to `wandb.log()`. + + Example: + + ```python + import random + import wandb + + # Generate random data for the table + data = [ + ["car", random.uniform(0, 1)], + ["bus", random.uniform(0, 1)], + ["road", random.uniform(0, 1)], + ["person", random.uniform(0, 1)], + ] + + # Create a table with the data + table = wandb.Table(data=data, columns=["class", "accuracy"]) + + # Initialize a W&B run and log the bar plot + with wandb.init(project="bar_chart") as run: + # Create a bar plot from the table + bar_plot = wandb.plot.bar( + table=table, + label="class", + value="accuracy", + title="Object Classification Accuracy", + ) + + # Log the bar chart to W&B + run.log({"bar_plot": bar_plot}) + ``` + """ + return plot_table( + data_table=table, + vega_spec_name="wandb/bar/v0", + fields={"label": label, "value": value}, + string_fields={"title": title}, + split_table=split_table, + ) diff --git a/lib/python3.12/site-packages/wandb/plot/custom_chart.py b/lib/python3.12/site-packages/wandb/plot/custom_chart.py new file mode 100644 index 0000000000000000000000000000000000000000..69f14a5ab6773e94638b833ccdc750c1725f7d29 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/plot/custom_chart.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import wandb + + +@dataclass +class CustomChartSpec: + spec_name: str + fields: dict[str, Any] + string_fields: dict[str, Any] + key: str = "" + panel_type: str = "Vega2" + split_table: bool = False + + @property + def table_key(self) -> str: + if not self.key: + raise wandb.Error("Key for the custom chart spec is not set.") + if self.split_table: + return f"Custom Chart Tables/{self.key}_table" + return f"{self.key}_table" + + @property + def config_value(self) -> dict[str, Any]: + return { + "panel_type": self.panel_type, + "panel_config": { + "panelDefId": self.spec_name, + "fieldSettings": self.fields, + "stringSettings": self.string_fields, + "transform": {"name": "tableWithLeafColNames"}, + "userQuery": { + "queryFields": [ + { + "name": "runSets", + "args": [{"name": "runSets", "value": "${runSets}"}], + "fields": [ + {"name": "id", "fields": []}, + {"name": "name", "fields": []}, + {"name": "_defaultColorIndex", "fields": []}, + { + "name": "summaryTable", + "args": [ + { + "name": "tableKey", + "value": self.table_key, + } + ], + "fields": [], + }, + ], + } + ], + }, + }, + } + + @property + def config_key(self) -> tuple[str, str, str]: + return ("_wandb", "visualize", self.key) + + +@dataclass +class CustomChart: + table: wandb.Table + spec: CustomChartSpec + + def set_key(self, key: str): + """Sets the key for the spec and updates dependent configurations.""" + self.spec.key = key + + +def plot_table( + vega_spec_name: str, + data_table: wandb.Table, + fields: dict[str, Any], + string_fields: dict[str, Any] | None = None, + split_table: bool = False, +) -> CustomChart: + """Creates a custom charts using a Vega-Lite specification and a `wandb.Table`. + + This function creates a custom chart based on a Vega-Lite specification and + a data table represented by a `wandb.Table` object. The specification needs + to be predefined and stored in the W&B backend. The function returns a custom + chart object that can be logged to W&B using `wandb.log()`. + + Args: + vega_spec_name (str): The name or identifier of the Vega-Lite spec + that defines the visualization structure. + data_table (wandb.Table): A `wandb.Table` object containing the data to be + visualized. + fields (dict[str, Any]): A mapping between the fields in the Vega-Lite spec and the + corresponding columns in the data table to be visualized. + string_fields (dict[str, Any] | None): A dictionary for providing values for any string constants + required by the custom visualization. + split_table (bool): Whether the table should be split into a separate section + in the W&B UI. If `True`, the table will be displayed in a section named + "Custom Chart Tables". Default is `False`. + + Returns: + CustomChart: A custom chart object that can be logged to W&B. To log the + chart, pass it to `wandb.log()`. + + Raises: + wandb.Error: If `data_table` is not a `wandb.Table` object. + """ + + if not isinstance(data_table, wandb.Table): + raise wandb.Error( + f"Expected `data_table` to be `wandb.Table` type, instead got {type(data_table).__name__}" + ) + + return CustomChart( + table=data_table, + spec=CustomChartSpec( + spec_name=vega_spec_name, + fields=fields, + string_fields=string_fields or {}, + split_table=split_table, + ), + ) diff --git a/lib/python3.12/site-packages/wandb/plot/line.py b/lib/python3.12/site-packages/wandb/plot/line.py new file mode 100644 index 0000000000000000000000000000000000000000..34857a4053ac3d699f447f927b7d6d0ea8d6ca14 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/plot/line.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from wandb.plot.custom_chart import plot_table + +if TYPE_CHECKING: + import wandb + from wandb.plot.custom_chart import CustomChart + + +def line( + table: wandb.Table, + x: str, + y: str, + stroke: str | None = None, + title: str = "", + split_table: bool = False, +) -> CustomChart: + """Constructs a customizable line chart. + + Args: + table: The table containing data for the chart. + x: Column name for the x-axis values. + y: Column name for the y-axis values. + stroke: Column name to differentiate line strokes (e.g., for + grouping lines). + title: Title of the chart. + split_table: Whether the table should be split into a separate section + in the W&B UI. If `True`, the table will be displayed in a section named + "Custom Chart Tables". Default is `False`. + + Returns: + CustomChart: A custom chart object that can be logged to W&B. To log the + chart, pass it to `wandb.log()`. + + Example: + + ```python + import math + import random + import wandb + + # Create multiple series of data with different patterns + data = [] + for i in range(100): + # Series 1: Sinusoidal pattern with random noise + data.append([i, math.sin(i / 10) + random.uniform(-0.1, 0.1), "series_1"]) + # Series 2: Cosine pattern with random noise + data.append([i, math.cos(i / 10) + random.uniform(-0.1, 0.1), "series_2"]) + # Series 3: Linear increase with random noise + data.append([i, i / 10 + random.uniform(-0.5, 0.5), "series_3"]) + + # Define the columns for the table + table = wandb.Table(data=data, columns=["step", "value", "series"]) + + # Initialize wandb run and log the line chart + with wandb.init(project="line_chart_example") as run: + line_chart = wandb.plot.line( + table=table, + x="step", + y="value", + stroke="series", # Group by the "series" column + title="Multi-Series Line Plot", + ) + run.log({"line-chart": line_chart}) + ``` + """ + return plot_table( + data_table=table, + vega_spec_name="wandb/line/v0", + fields={"x": x, "y": y, "stroke": stroke}, + string_fields={"title": title}, + split_table=split_table, + ) diff --git a/lib/python3.12/site-packages/wandb/plot/pr_curve.py b/lib/python3.12/site-packages/wandb/plot/pr_curve.py new file mode 100644 index 0000000000000000000000000000000000000000..a8ade060c3dab946002f48d3b3164ed901b8efe3 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/plot/pr_curve.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +import numbers +from typing import TYPE_CHECKING, Iterable, TypeVar + +import wandb +from wandb import util +from wandb.plot.custom_chart import plot_table +from wandb.plot.utils import test_missing, test_types + +if TYPE_CHECKING: + from wandb.plot.custom_chart import CustomChart + + +T = TypeVar("T") + + +def pr_curve( + y_true: Iterable[T] | None = None, + y_probas: Iterable[numbers.Number] | None = None, + labels: list[str] | None = None, + classes_to_plot: list[T] | None = None, + interp_size: int = 21, + title: str = "Precision-Recall Curve", + split_table: bool = False, +) -> CustomChart: + """Constructs a Precision-Recall (PR) curve. + + The Precision-Recall curve is particularly useful for evaluating classifiers + on imbalanced datasets. A high area under the PR curve signifies both high + precision (a low false positive rate) and high recall (a low false negative + rate). The curve provides insights into the balance between false positives + and false negatives at various threshold levels, aiding in the assessment of + a model's performance. + + Args: + y_true: True binary labels. The shape should be (`num_samples`,). + y_probas: Predicted scores or probabilities for each class. + These can be probability estimates, confidence scores, or non-thresholded + decision values. The shape should be (`num_samples`, `num_classes`). + labels: Optional list of class names to replace + numeric values in `y_true` for easier plot interpretation. + For example, `labels = ['dog', 'cat', 'owl']` will replace 0 with + 'dog', 1 with 'cat', and 2 with 'owl' in the plot. If not provided, + numeric values from `y_true` will be used. + classes_to_plot: Optional list of unique class values from + y_true to be included in the plot. If not specified, all unique + classes in y_true will be plotted. + interp_size: Number of points to interpolate recall values. The + recall values will be fixed to `interp_size` uniformly distributed + points in the range [0, 1], and the precision will be interpolated + accordingly. + title: Title of the plot. Defaults to "Precision-Recall Curve". + split_table: Whether the table should be split into a separate section + in the W&B UI. If `True`, the table will be displayed in a section named + "Custom Chart Tables". Default is `False`. + + Returns: + CustomChart: A custom chart object that can be logged to W&B. To log the + chart, pass it to `wandb.log()`. + + Raises: + wandb.Error: If NumPy, pandas, or scikit-learn is not installed. + + + Example: + + ```python + import wandb + + # Example for spam detection (binary classification) + y_true = [0, 1, 1, 0, 1] # 0 = not spam, 1 = spam + y_probas = [ + [0.9, 0.1], # Predicted probabilities for the first sample (not spam) + [0.2, 0.8], # Second sample (spam), and so on + [0.1, 0.9], + [0.8, 0.2], + [0.3, 0.7], + ] + + labels = ["not spam", "spam"] # Optional class names for readability + + with wandb.init(project="spam-detection") as run: + pr_curve = wandb.plot.pr_curve( + y_true=y_true, + y_probas=y_probas, + labels=labels, + title="Precision-Recall Curve for Spam Detection", + ) + run.log({"pr-curve": pr_curve}) + ``` + """ + np = util.get_module( + "numpy", + required="roc requires the numpy library, install with `pip install numpy`", + ) + pd = util.get_module( + "pandas", + required="roc requires the pandas library, install with `pip install pandas`", + ) + sklearn_metrics = util.get_module( + "sklearn.metrics", + "roc requires the scikit library, install with `pip install scikit-learn`", + ) + sklearn_utils = util.get_module( + "sklearn.utils", + "roc requires the scikit library, install with `pip install scikit-learn`", + ) + + def _step(x): + y = np.array(x) + for i in range(1, len(y)): + y[i] = max(y[i], y[i - 1]) + return y + + y_true = np.array(y_true) + y_probas = np.array(y_probas) + + if not test_missing(y_true=y_true, y_probas=y_probas): + return + if not test_types(y_true=y_true, y_probas=y_probas): + return + + classes = np.unique(y_true) + if classes_to_plot is None: + classes_to_plot = classes + + precision = {} + interp_recall = np.linspace(0, 1, interp_size)[::-1] + indices_to_plot = np.where(np.isin(classes, classes_to_plot))[0] + for i in indices_to_plot: + if labels is not None and ( + isinstance(classes[i], int) or isinstance(classes[0], np.integer) + ): + class_label = labels[classes[i]] + else: + class_label = classes[i] + + cur_precision, cur_recall, _ = sklearn_metrics.precision_recall_curve( + y_true, y_probas[:, i], pos_label=classes[i] + ) + # smooth the precision (monotonically increasing) + cur_precision = _step(cur_precision) + + # reverse order so that recall in ascending + cur_precision = cur_precision[::-1] + cur_recall = cur_recall[::-1] + indices = np.searchsorted(cur_recall, interp_recall, side="left") + precision[class_label] = cur_precision[indices] + + df = pd.DataFrame( + { + "class": np.hstack([[k] * len(v) for k, v in precision.items()]), + "precision": np.hstack(list(precision.values())), + "recall": np.tile(interp_recall, len(precision)), + } + ).round(3) + + if len(df) > wandb.Table.MAX_ROWS: + wandb.termwarn( + f"Table has a limit of {wandb.Table.MAX_ROWS} rows. Resampling to fit." + ) + # different sampling could be applied, possibly to ensure endpoints are kept + df = sklearn_utils.resample( + df, + replace=False, + n_samples=wandb.Table.MAX_ROWS, + random_state=42, + stratify=df["class"], + ).sort_values(["precision", "recall", "class"]) + + return plot_table( + data_table=wandb.Table(dataframe=df), + vega_spec_name="wandb/area-under-curve/v0", + fields={ + "x": "recall", + "y": "precision", + "class": "class", + }, + string_fields={ + "title": title, + "x-axis-title": "Recall", + "y-axis-title": "Precision", + }, + split_table=split_table, + ) diff --git a/lib/python3.12/site-packages/wandb/plot/scatter.py b/lib/python3.12/site-packages/wandb/plot/scatter.py new file mode 100644 index 0000000000000000000000000000000000000000..40e3d8033e2d6249554e3de56904add3ef30ffbd --- /dev/null +++ b/lib/python3.12/site-packages/wandb/plot/scatter.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from wandb.plot.custom_chart import plot_table + +if TYPE_CHECKING: + import wandb + from wandb.plot.custom_chart import CustomChart + + +def scatter( + table: wandb.Table, + x: str, + y: str, + title: str = "", + split_table: bool = False, +) -> CustomChart: + """Constructs a scatter plot from a wandb.Table of data. + + Args: + table: The W&B Table containing the data to visualize. + x: The name of the column used for the x-axis. + y: The name of the column used for the y-axis. + title: The title of the scatter chart. + split_table: Whether the table should be split into a separate section + in the W&B UI. If `True`, the table will be displayed in a section named + "Custom Chart Tables". Default is `False`. + + Returns: + CustomChart: A custom chart object that can be logged to W&B. To log the + chart, pass it to `wandb.log()`. + Example: + ```python + import math + import random + import wandb + + # Simulate temperature variations at different altitudes over time + data = [ + [i, random.uniform(-10, 20) - 0.005 * i + 5 * math.sin(i / 50)] + for i in range(300) + ] + + # Create W&B table with altitude (m) and temperature (°C) columns + table = wandb.Table(data=data, columns=["altitude (m)", "temperature (°C)"]) + + # Initialize W&B run and log the scatter plot + with wandb.init(project="temperature-altitude-scatter") as run: + # Create and log the scatter plot + scatter_plot = wandb.plot.scatter( + table=table, + x="altitude (m)", + y="temperature (°C)", + title="Altitude vs Temperature", + ) + run.log({"altitude-temperature-scatter": scatter_plot}) + ``` + """ + return plot_table( + data_table=table, + vega_spec_name="wandb/scatter/v0", + fields={"x": x, "y": y}, + string_fields={"title": title}, + split_table=split_table, + ) diff --git a/lib/python3.12/site-packages/wandb/plot/utils.py b/lib/python3.12/site-packages/wandb/plot/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9436dce8e3394acd731d3efffed23af8188490e6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/plot/utils.py @@ -0,0 +1,184 @@ +from typing import Iterable, Sequence + +import wandb +from wandb import util + + +def test_missing(**kwargs): + np = util.get_module("numpy", required="Logging plots requires numpy") + pd = util.get_module("pandas", required="Logging dataframes requires pandas") + scipy = util.get_module("scipy", required="Logging scipy matrices requires scipy") + + test_passed = True + for k, v in kwargs.items(): + # Missing/empty params/datapoint arrays + if v is None: + wandb.termerror(f"{k} is None. Please try again.") + test_passed = False + if (k == "X") or (k == "X_test"): + if isinstance(v, scipy.sparse.csr.csr_matrix): + v = v.toarray() + elif isinstance(v, (pd.DataFrame, pd.Series)): + v = v.to_numpy() + elif isinstance(v, list): + v = np.asarray(v) + + # Warn the user about missing values + missing = 0 + missing = np.count_nonzero(pd.isnull(v)) + if missing > 0: + wandb.termwarn("%s contains %d missing values. " % (k, missing)) + test_passed = False + # Ensure the dataset contains only integers + non_nums = 0 + if v.ndim == 1: + non_nums = sum( + 1 + for val in v + if ( + not isinstance(val, (int, float, complex)) + and not isinstance(val, np.number) + ) + ) + else: + non_nums = sum( + 1 + for sl in v + for val in sl + if ( + not isinstance(val, (int, float, complex)) + and not isinstance(val, np.number) + ) + ) + if non_nums > 0: + wandb.termerror( + f"{k} contains values that are not numbers. Please vectorize, " + f"label encode or one hot encode {k} and call the plotting function again." + ) + test_passed = False + return test_passed + + +def test_fitted(model): + np = util.get_module("numpy", required="Logging plots requires numpy") + _ = util.get_module("pandas", required="Logging dataframes requires pandas") + _ = util.get_module("scipy", required="Logging scipy matrices requires scipy") + scikit_utils = util.get_module( + "sklearn.utils", + required="roc requires the scikit utils submodule, install with `pip install scikit-learn`", + ) + scikit_exceptions = util.get_module( + "sklearn.exceptions", + "roc requires the scikit preprocessing submodule, install with `pip install scikit-learn`", + ) + + try: + model.predict(np.zeros((7, 3))) + except scikit_exceptions.NotFittedError: + wandb.termerror("Please fit the model before passing it in.") + return False + except AttributeError: + # Some clustering models (LDA, PCA, Agglomerative) don't implement ``predict`` + try: + scikit_utils.validation.check_is_fitted( + model, + [ + "coef_", + "estimator_", + "labels_", + "n_clusters_", + "children_", + "components_", + "n_components_", + "n_iter_", + "n_batch_iter_", + "explained_variance_", + "singular_values_", + "mean_", + ], + all_or_any=any, + ) + except scikit_exceptions.NotFittedError: + wandb.termerror("Please fit the model before passing it in.") + return False + else: + return True + except Exception: + # Assume it's fitted, since ``NotFittedError`` wasn't raised + return True + + +def encode_labels(df): + _ = util.get_module("pandas", required="Logging dataframes requires pandas") + preprocessing = util.get_module( + "sklearn.preprocessing", + "roc requires the scikit preprocessing submodule, install with `pip install scikit-learn`", + ) + + le = preprocessing.LabelEncoder() + # apply le on categorical feature columns + categorical_cols = df.select_dtypes( + exclude=["int", "float", "float64", "float32", "int32", "int64"] + ).columns + df[categorical_cols] = df[categorical_cols].apply(lambda col: le.fit_transform(col)) + + +def test_types(**kwargs): + np = util.get_module("numpy", required="Logging plots requires numpy") + pd = util.get_module("pandas", required="Logging dataframes requires pandas") + _ = util.get_module("scipy", required="Logging scipy matrices requires scipy") + + base = util.get_module( + "sklearn.base", + "roc requires the scikit base submodule, install with `pip install scikit-learn`", + ) + + test_passed = True + for k, v in kwargs.items(): + # check for incorrect types + if ( + (k == "X") + or (k == "X_test") + or (k == "y") + or (k == "y_test") + or (k == "y_true") + or (k == "y_probas") + or (k == "x_labels") + or (k == "y_labels") + or (k == "matrix_values") + ): + # FIXME: do this individually + if not isinstance( + v, + ( + Sequence, + Iterable, + np.ndarray, + np.generic, + pd.DataFrame, + pd.Series, + list, + ), + ): + wandb.termerror(f"{k} is not an array. Please try again.") + test_passed = False + # check for classifier types + if k == "model": + if (not base.is_classifier(v)) and (not base.is_regressor(v)): + wandb.termerror( + f"{k} is not a classifier or regressor. Please try again." + ) + test_passed = False + elif k == "clf" or k == "binary_clf": + if not (base.is_classifier(v)): + wandb.termerror(f"{k} is not a classifier. Please try again.") + test_passed = False + elif k == "regressor": + if not base.is_regressor(v): + wandb.termerror(f"{k} is not a regressor. Please try again.") + test_passed = False + elif k == "clusterer": + if not (getattr(v, "_estimator_type", None) == "clusterer"): + wandb.termerror(f"{k} is not a clusterer. Please try again.") + test_passed = False + return test_passed diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a06caaf7b76317394dbcb0ae457b8ad4a647a8d7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_alerts.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_alerts.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32aeef89d6a82492066dbc5cabb04df38a570901 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_alerts.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_config.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5017ae0f3210f18356d6dcc130b2a64f916f2365 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_helper.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_helper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1893b73a3c781f0cb2deb537d95ef2a689e17cf1 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_helper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_init.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_init.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..653c776bca42d2e717f470b4dd3076510d55261e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_init.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_login.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_login.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b4ecdbc51bfabd03972a70c0bd46997e0986c329 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_login.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_metric.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_metric.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..200957978719354c353c68cf3aae49460e29b0d3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_metric.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_require.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_require.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d8487dcab7a29a197e5bf125839418b0aadb6de Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_require.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_require_helpers.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_require_helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..815fc37f9bc87b1434f8767d1cacd41ec35aede5 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_require_helpers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_settings.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_settings.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..afe2b9c0c196228785355a8a1fa69f367bceb070 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_settings.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_setup.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_setup.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33f5072d4475235ad7af9da229abc78c9fd0042c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_setup.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_summary.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_summary.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b298d6ad62b1b7124bfb2ff3a05ca49dc624dab8 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_summary.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_sweep.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_sweep.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7450446cc3b9b0eb76c3447f2d05ce691e5dee4a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_sweep.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_sync.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_sync.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84ba57f4326328b53ed2c818137b4ddda7c8a665 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_sync.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_watch.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_watch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb432de4af70f72cbb036e46f02fd1d466f508ce Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/__pycache__/wandb_watch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__init__.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f038f0eaf7e45d2d18b91e6427507880a851d35 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/_graphql_fragments.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/_graphql_fragments.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ebee859695642e1f834a9e37776f5f93a485c771 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/_graphql_fragments.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/_internal_artifact.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/_internal_artifact.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e344cd748837c88ac814074586e4267c2d24783 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/_internal_artifact.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/_validators.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/_validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f608679c61700d4d190c9ea1ca54a7c9125032f6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/_validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_download_logger.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_download_logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0015ef6a7ccfca1a2a1cc01613a22553b4e5b8a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_download_logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_file_cache.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_file_cache.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b1bb72e22232b28d2709b684a6ec9bc646b90f3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_file_cache.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_instance_cache.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_instance_cache.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eef5cea2c459402b3bf7c962e900671599764264 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_instance_cache.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_manifest.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_manifest.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79b3e91fab4a73c33a2e7339fbc28b5b719736f2 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_manifest.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_manifest_entry.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_manifest_entry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a5501dfbbbaf574dff7fca917f14a60b578024f Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_manifest_entry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_saver.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_saver.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c33d39b21df880c000c3f39e816ee0080af7709 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_saver.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_state.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_state.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..498c1d0be566a19ca57f1284b9078c694785b1bd Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_state.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_ttl.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_ttl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2624efed06ab0a6399bd8afff4f7dfe92d9616b1 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/artifact_ttl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/exceptions.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/exceptions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4085da2514837c0dda5c7c35617d2565200e3938 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/exceptions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/staging.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/staging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e87b2cf390855b0453236aca4e1b623a0da4a60 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/staging.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/storage_handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/storage_handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..239c755ce5df7fc3c15f389ef361ecdf6fd12265 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/storage_handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/storage_layout.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/storage_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..649f9f8fa9ebb7ddcfe19b4943ef4c7ed1f8b96d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/storage_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/storage_policy.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/storage_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9aae8ec95bf8d5df7f7a88e41960f9bd79393dcf Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/__pycache__/storage_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__init__.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3a844096d71119a33498e2be770093e5fff24fe7 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__init__.py @@ -0,0 +1,300 @@ +# Generated by ariadne-codegen + +from .add_aliases import AddAliases, AddAliasesAddAliases +from .artifact_collection_membership_files import ( + ArtifactCollectionMembershipFiles, + ArtifactCollectionMembershipFilesProject, + ArtifactCollectionMembershipFilesProjectArtifactCollection, + ArtifactCollectionMembershipFilesProjectArtifactCollectionArtifactMembership, +) +from .artifact_version_files import ( + ArtifactVersionFiles, + ArtifactVersionFilesProject, + ArtifactVersionFilesProjectArtifactType, + ArtifactVersionFilesProjectArtifactTypeArtifact, +) +from .create_artifact_collection_tag_assignments import ( + CreateArtifactCollectionTagAssignments, + CreateArtifactCollectionTagAssignmentsCreateArtifactCollectionTagAssignments, + CreateArtifactCollectionTagAssignmentsCreateArtifactCollectionTagAssignmentsTags, +) +from .delete_aliases import DeleteAliases, DeleteAliasesDeleteAliases +from .delete_artifact_collection_tag_assignments import ( + DeleteArtifactCollectionTagAssignments, + DeleteArtifactCollectionTagAssignmentsDeleteArtifactCollectionTagAssignments, +) +from .delete_artifact_portfolio import ( + DeleteArtifactPortfolio, + DeleteArtifactPortfolioDeleteArtifactPortfolio, + DeleteArtifactPortfolioDeleteArtifactPortfolioArtifactCollection, +) +from .delete_artifact_sequence import ( + DeleteArtifactSequence, + DeleteArtifactSequenceDeleteArtifactSequence, + DeleteArtifactSequenceDeleteArtifactSequenceArtifactCollection, +) +from .enums import ArtifactCollectionState, ArtifactState +from .fetch_linked_artifacts import ( + FetchLinkedArtifacts, + FetchLinkedArtifactsArtifact, + FetchLinkedArtifactsArtifactArtifactMemberships, + FetchLinkedArtifactsArtifactArtifactMembershipsEdges, + FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNode, + FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeAliases, + FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeArtifactCollection, + FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeArtifactCollectionProject, +) +from .fragments import ( + ArtifactCollectionsFragment, + ArtifactCollectionsFragmentEdges, + ArtifactCollectionsFragmentEdgesNode, + ArtifactCollectionsFragmentPageInfo, + ArtifactFragment, + ArtifactFragmentAliases, + ArtifactFragmentAliasesArtifactCollection, + ArtifactFragmentAliasesArtifactCollectionProject, + ArtifactFragmentArtifactSequence, + ArtifactFragmentArtifactSequenceProject, + ArtifactFragmentArtifactType, + ArtifactFragmentCurrentManifest, + ArtifactFragmentCurrentManifestFile, + ArtifactFragmentTags, + ArtifactPortfolioTypeFields, + ArtifactSequenceTypeFields, + ArtifactsFragment, + ArtifactsFragmentEdges, + ArtifactsFragmentPageInfo, + ArtifactTypeFragment, + ArtifactTypesFragment, + ArtifactTypesFragmentEdges, + ArtifactTypesFragmentPageInfo, + FilesFragment, + FilesFragmentEdges, + FilesFragmentEdgesNode, + FilesFragmentPageInfo, +) +from .input_types import ( + ArtifactAliasInput, + ArtifactCollectionAliasInput, + LinkArtifactInput, + TagInput, +) +from .link_artifact import LinkArtifact, LinkArtifactLinkArtifact +from .move_artifact_collection import ( + MoveArtifactCollection, + MoveArtifactCollectionMoveArtifactSequence, + MoveArtifactCollectionMoveArtifactSequenceArtifactCollection, +) +from .operations import ( + ADD_ALIASES_GQL, + ARTIFACT_COLLECTION_MEMBERSHIP_FILES_GQL, + ARTIFACT_VERSION_FILES_GQL, + CREATE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL, + DELETE_ALIASES_GQL, + DELETE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL, + DELETE_ARTIFACT_PORTFOLIO_GQL, + DELETE_ARTIFACT_SEQUENCE_GQL, + FETCH_LINKED_ARTIFACTS_GQL, + LINK_ARTIFACT_GQL, + MOVE_ARTIFACT_COLLECTION_GQL, + PROJECT_ARTIFACT_COLLECTION_GQL, + PROJECT_ARTIFACT_COLLECTIONS_GQL, + PROJECT_ARTIFACT_TYPE_GQL, + PROJECT_ARTIFACT_TYPES_GQL, + PROJECT_ARTIFACTS_GQL, + RUN_INPUT_ARTIFACTS_GQL, + RUN_OUTPUT_ARTIFACTS_GQL, + UPDATE_ARTIFACT_GQL, + UPDATE_ARTIFACT_PORTFOLIO_GQL, + UPDATE_ARTIFACT_SEQUENCE_GQL, +) +from .project_artifact_collection import ( + ProjectArtifactCollection, + ProjectArtifactCollectionProject, + ProjectArtifactCollectionProjectArtifactType, + ProjectArtifactCollectionProjectArtifactTypeArtifactCollection, + ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliases, + ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesEdges, + ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesEdgesNode, + ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesPageInfo, + ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTags, + ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTagsEdges, + ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTagsEdgesNode, + ProjectArtifactCollectionProjectArtifactTypeArtifactSequence, +) +from .project_artifact_collections import ( + ProjectArtifactCollections, + ProjectArtifactCollectionsProject, + ProjectArtifactCollectionsProjectArtifactType, +) +from .project_artifact_type import ProjectArtifactType, ProjectArtifactTypeProject +from .project_artifact_types import ProjectArtifactTypes, ProjectArtifactTypesProject +from .project_artifacts import ( + ProjectArtifacts, + ProjectArtifactsProject, + ProjectArtifactsProjectArtifactType, + ProjectArtifactsProjectArtifactTypeArtifactCollection, +) +from .run_input_artifacts import ( + RunInputArtifacts, + RunInputArtifactsProject, + RunInputArtifactsProjectRun, + RunInputArtifactsProjectRunInputArtifacts, + RunInputArtifactsProjectRunInputArtifactsEdges, + RunInputArtifactsProjectRunInputArtifactsPageInfo, +) +from .run_output_artifacts import ( + RunOutputArtifacts, + RunOutputArtifactsProject, + RunOutputArtifactsProjectRun, + RunOutputArtifactsProjectRunOutputArtifacts, + RunOutputArtifactsProjectRunOutputArtifactsEdges, + RunOutputArtifactsProjectRunOutputArtifactsPageInfo, +) +from .update_artifact import UpdateArtifact, UpdateArtifactUpdateArtifact +from .update_artifact_portfolio import ( + UpdateArtifactPortfolio, + UpdateArtifactPortfolioUpdateArtifactPortfolio, + UpdateArtifactPortfolioUpdateArtifactPortfolioArtifactCollection, +) +from .update_artifact_sequence import ( + UpdateArtifactSequence, + UpdateArtifactSequenceUpdateArtifactSequence, + UpdateArtifactSequenceUpdateArtifactSequenceArtifactCollection, +) + +__all__ = [ + "ADD_ALIASES_GQL", + "ARTIFACT_COLLECTION_MEMBERSHIP_FILES_GQL", + "ARTIFACT_VERSION_FILES_GQL", + "CREATE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL", + "DELETE_ALIASES_GQL", + "DELETE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL", + "DELETE_ARTIFACT_PORTFOLIO_GQL", + "DELETE_ARTIFACT_SEQUENCE_GQL", + "FETCH_LINKED_ARTIFACTS_GQL", + "LINK_ARTIFACT_GQL", + "MOVE_ARTIFACT_COLLECTION_GQL", + "PROJECT_ARTIFACTS_GQL", + "PROJECT_ARTIFACT_COLLECTIONS_GQL", + "PROJECT_ARTIFACT_COLLECTION_GQL", + "PROJECT_ARTIFACT_TYPES_GQL", + "PROJECT_ARTIFACT_TYPE_GQL", + "RUN_INPUT_ARTIFACTS_GQL", + "RUN_OUTPUT_ARTIFACTS_GQL", + "UPDATE_ARTIFACT_GQL", + "UPDATE_ARTIFACT_PORTFOLIO_GQL", + "UPDATE_ARTIFACT_SEQUENCE_GQL", + "DeleteArtifactSequence", + "DeleteArtifactSequenceDeleteArtifactSequence", + "DeleteArtifactSequenceDeleteArtifactSequenceArtifactCollection", + "DeleteArtifactPortfolio", + "DeleteArtifactPortfolioDeleteArtifactPortfolio", + "DeleteArtifactPortfolioDeleteArtifactPortfolioArtifactCollection", + "UpdateArtifactSequence", + "UpdateArtifactSequenceUpdateArtifactSequence", + "UpdateArtifactSequenceUpdateArtifactSequenceArtifactCollection", + "UpdateArtifactPortfolio", + "UpdateArtifactPortfolioUpdateArtifactPortfolio", + "UpdateArtifactPortfolioUpdateArtifactPortfolioArtifactCollection", + "MoveArtifactCollection", + "MoveArtifactCollectionMoveArtifactSequence", + "MoveArtifactCollectionMoveArtifactSequenceArtifactCollection", + "CreateArtifactCollectionTagAssignments", + "CreateArtifactCollectionTagAssignmentsCreateArtifactCollectionTagAssignments", + "CreateArtifactCollectionTagAssignmentsCreateArtifactCollectionTagAssignmentsTags", + "DeleteArtifactCollectionTagAssignments", + "DeleteArtifactCollectionTagAssignmentsDeleteArtifactCollectionTagAssignments", + "ProjectArtifactCollections", + "ProjectArtifactCollectionsProject", + "ProjectArtifactCollectionsProjectArtifactType", + "ProjectArtifactCollection", + "ProjectArtifactCollectionProject", + "ProjectArtifactCollectionProjectArtifactType", + "ProjectArtifactCollectionProjectArtifactTypeArtifactCollection", + "ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliases", + "ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesEdges", + "ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesEdgesNode", + "ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesPageInfo", + "ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTags", + "ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTagsEdges", + "ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTagsEdgesNode", + "ProjectArtifactCollectionProjectArtifactTypeArtifactSequence", + "ArtifactVersionFiles", + "ArtifactVersionFilesProject", + "ArtifactVersionFilesProjectArtifactType", + "ArtifactVersionFilesProjectArtifactTypeArtifact", + "ArtifactCollectionMembershipFiles", + "ArtifactCollectionMembershipFilesProject", + "ArtifactCollectionMembershipFilesProjectArtifactCollection", + "ArtifactCollectionMembershipFilesProjectArtifactCollectionArtifactMembership", + "ProjectArtifactTypes", + "ProjectArtifactTypesProject", + "ProjectArtifactType", + "ProjectArtifactTypeProject", + "ProjectArtifacts", + "ProjectArtifactsProject", + "ProjectArtifactsProjectArtifactType", + "ProjectArtifactsProjectArtifactTypeArtifactCollection", + "RunOutputArtifacts", + "RunOutputArtifactsProject", + "RunOutputArtifactsProjectRun", + "RunOutputArtifactsProjectRunOutputArtifacts", + "RunOutputArtifactsProjectRunOutputArtifactsEdges", + "RunOutputArtifactsProjectRunOutputArtifactsPageInfo", + "RunInputArtifacts", + "RunInputArtifactsProject", + "RunInputArtifactsProjectRun", + "RunInputArtifactsProjectRunInputArtifacts", + "RunInputArtifactsProjectRunInputArtifactsEdges", + "RunInputArtifactsProjectRunInputArtifactsPageInfo", + "FetchLinkedArtifacts", + "FetchLinkedArtifactsArtifact", + "FetchLinkedArtifactsArtifactArtifactMemberships", + "FetchLinkedArtifactsArtifactArtifactMembershipsEdges", + "FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNode", + "FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeAliases", + "FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeArtifactCollection", + "FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeArtifactCollectionProject", + "AddAliases", + "AddAliasesAddAliases", + "DeleteAliases", + "DeleteAliasesDeleteAliases", + "UpdateArtifact", + "UpdateArtifactUpdateArtifact", + "LinkArtifact", + "LinkArtifactLinkArtifact", + "ArtifactAliasInput", + "ArtifactCollectionAliasInput", + "LinkArtifactInput", + "TagInput", + "ArtifactCollectionsFragment", + "ArtifactCollectionsFragmentEdges", + "ArtifactCollectionsFragmentEdgesNode", + "ArtifactCollectionsFragmentPageInfo", + "ArtifactFragment", + "ArtifactFragmentAliases", + "ArtifactFragmentAliasesArtifactCollection", + "ArtifactFragmentAliasesArtifactCollectionProject", + "ArtifactFragmentArtifactSequence", + "ArtifactFragmentArtifactSequenceProject", + "ArtifactFragmentArtifactType", + "ArtifactFragmentCurrentManifest", + "ArtifactFragmentCurrentManifestFile", + "ArtifactFragmentTags", + "ArtifactPortfolioTypeFields", + "ArtifactSequenceTypeFields", + "ArtifactTypeFragment", + "ArtifactTypesFragment", + "ArtifactTypesFragmentEdges", + "ArtifactTypesFragmentPageInfo", + "ArtifactsFragment", + "ArtifactsFragmentEdges", + "ArtifactsFragmentPageInfo", + "FilesFragment", + "FilesFragmentEdges", + "FilesFragmentEdgesNode", + "FilesFragmentPageInfo", + "ArtifactCollectionState", + "ArtifactState", +] diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ade8f09c145973931e4049a292d20f42dab10229 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/add_aliases.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/add_aliases.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..932fc40f0e1c385624c7548712cdd83145aa2395 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/add_aliases.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/artifact_collection_membership_files.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/artifact_collection_membership_files.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5c3b27ef7338fc1c2002d04ce3b503c9837464b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/artifact_collection_membership_files.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/artifact_version_files.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/artifact_version_files.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f90af4445823a27bc34ee51aedd0ff82484900b1 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/artifact_version_files.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/create_artifact_collection_tag_assignments.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/create_artifact_collection_tag_assignments.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ef4f6cf49966fa75615e714814f0d8d703b8964d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/create_artifact_collection_tag_assignments.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_aliases.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_aliases.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59969220232311ec2ac27f812665996d46357dc6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_aliases.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_artifact_collection_tag_assignments.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_artifact_collection_tag_assignments.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea6a3aa677ea1b2bd292538ee2eec49100173c48 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_artifact_collection_tag_assignments.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_artifact_portfolio.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_artifact_portfolio.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60e0ede3cb9c82570b1ebc6ea4c1e834c6473640 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_artifact_portfolio.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_artifact_sequence.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_artifact_sequence.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a4eb421f6cd369ef3bfff142436e2ab5c334f60 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/delete_artifact_sequence.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/enums.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/enums.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..914423d05ce9fe0db4ce471c0753c47b00893083 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/enums.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/fetch_linked_artifacts.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/fetch_linked_artifacts.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ceed3018b490e2445c453338614554daf2252093 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/fetch_linked_artifacts.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/fragments.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/fragments.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f90ce4310d14e92ae03a11bd515ade7e448500f4 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/fragments.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/input_types.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/input_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14917c010f22bd224609024b2e516770dc8bc68b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/input_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/link_artifact.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/link_artifact.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..29e0ad2d7eeea2d7b59aa2e1c9dc903ca6a97f03 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/link_artifact.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/move_artifact_collection.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/move_artifact_collection.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d4fd6e5c57033efd0647b0117f80ab28af2bbd16 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/move_artifact_collection.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/operations.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/operations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4caa1682b675e485c7006dcc3db0709e47be9e68 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/operations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_collection.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_collection.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58a8a7087c2222eb1d2bb6e9a45510a5273d5430 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_collection.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_collections.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_collections.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e9748f2ffe7fc788437909cf86afa050cf73912 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_collections.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_type.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_type.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db110b57cbde210ece175d78cd02d1ca94dd5e01 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_type.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_types.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21b551fe7ad4d228769c1f330169d0e1e02bf568 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifact_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifacts.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifacts.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c184f50f29a47b28ad87e85b9070c4ae95b1fc71 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/project_artifacts.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/run_input_artifacts.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/run_input_artifacts.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..68f253012ed8c8b562cf253e26c87f7f0da83d57 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/run_input_artifacts.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/run_output_artifacts.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/run_output_artifacts.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf599a5dd0980e3db9360f0fd674c756d18d85c7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/run_output_artifacts.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/update_artifact.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/update_artifact.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f521f9758ca0e4ddaac865dfac849d2bda5d01ba Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/update_artifact.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/update_artifact_portfolio.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/update_artifact_portfolio.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dc24b50aaae86b2c753006de8e87a03fbd661066 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/update_artifact_portfolio.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/update_artifact_sequence.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/update_artifact_sequence.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76005188b0e66d1f148d68dd198a8cb8fa118def Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/__pycache__/update_artifact_sequence.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/add_aliases.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/add_aliases.py new file mode 100644 index 0000000000000000000000000000000000000000..14b288fdde8f9057dd43263a952d777568da7aa4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/add_aliases.py @@ -0,0 +1,21 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + + +class AddAliases(GQLBase): + add_aliases: Optional[AddAliasesAddAliases] = Field(alias="addAliases") + + +class AddAliasesAddAliases(GQLBase): + success: bool + + +AddAliases.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/artifact_collection_membership_files.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/artifact_collection_membership_files.py new file mode 100644 index 0000000000000000000000000000000000000000..87dc90d69600c9d4e06c3254412ae72aaaca3d8f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/artifact_collection_membership_files.py @@ -0,0 +1,43 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, Typename + +from .fragments import FilesFragment + + +class ArtifactCollectionMembershipFiles(GQLBase): + project: Optional[ArtifactCollectionMembershipFilesProject] + + +class ArtifactCollectionMembershipFilesProject(GQLBase): + artifact_collection: Optional[ + ArtifactCollectionMembershipFilesProjectArtifactCollection + ] = Field(alias="artifactCollection") + + +class ArtifactCollectionMembershipFilesProjectArtifactCollection(GQLBase): + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + artifact_membership: Optional[ + ArtifactCollectionMembershipFilesProjectArtifactCollectionArtifactMembership + ] = Field(alias="artifactMembership") + + +class ArtifactCollectionMembershipFilesProjectArtifactCollectionArtifactMembership( + GQLBase +): + files: Optional[FilesFragment] + + +ArtifactCollectionMembershipFiles.model_rebuild() +ArtifactCollectionMembershipFilesProject.model_rebuild() +ArtifactCollectionMembershipFilesProjectArtifactCollection.model_rebuild() +ArtifactCollectionMembershipFilesProjectArtifactCollectionArtifactMembership.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/artifact_version_files.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/artifact_version_files.py new file mode 100644 index 0000000000000000000000000000000000000000..c35018dccf3ae1ab003644345ffcebfc8530526d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/artifact_version_files.py @@ -0,0 +1,36 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + +from .fragments import FilesFragment + + +class ArtifactVersionFiles(GQLBase): + project: Optional[ArtifactVersionFilesProject] + + +class ArtifactVersionFilesProject(GQLBase): + artifact_type: Optional[ArtifactVersionFilesProjectArtifactType] = Field( + alias="artifactType" + ) + + +class ArtifactVersionFilesProjectArtifactType(GQLBase): + artifact: Optional[ArtifactVersionFilesProjectArtifactTypeArtifact] + + +class ArtifactVersionFilesProjectArtifactTypeArtifact(GQLBase): + files: Optional[FilesFragment] + + +ArtifactVersionFiles.model_rebuild() +ArtifactVersionFilesProject.model_rebuild() +ArtifactVersionFilesProjectArtifactType.model_rebuild() +ArtifactVersionFilesProjectArtifactTypeArtifact.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/create_artifact_collection_tag_assignments.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/create_artifact_collection_tag_assignments.py new file mode 100644 index 0000000000000000000000000000000000000000..5c625cc843e0c16934d8b59f95c646e79ea09756 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/create_artifact_collection_tag_assignments.py @@ -0,0 +1,36 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import List, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, GQLId + + +class CreateArtifactCollectionTagAssignments(GQLBase): + create_artifact_collection_tag_assignments: Optional[ + CreateArtifactCollectionTagAssignmentsCreateArtifactCollectionTagAssignments + ] = Field(alias="createArtifactCollectionTagAssignments") + + +class CreateArtifactCollectionTagAssignmentsCreateArtifactCollectionTagAssignments( + GQLBase +): + tags: List[ + CreateArtifactCollectionTagAssignmentsCreateArtifactCollectionTagAssignmentsTags + ] + + +class CreateArtifactCollectionTagAssignmentsCreateArtifactCollectionTagAssignmentsTags( + GQLBase +): + id: GQLId + name: str + tag_category_name: str = Field(alias="tagCategoryName") + + +CreateArtifactCollectionTagAssignments.model_rebuild() +CreateArtifactCollectionTagAssignmentsCreateArtifactCollectionTagAssignments.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_aliases.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_aliases.py new file mode 100644 index 0000000000000000000000000000000000000000..d1aa3e1973e943f254b02907a34da58f951eadb4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_aliases.py @@ -0,0 +1,21 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + + +class DeleteAliases(GQLBase): + delete_aliases: Optional[DeleteAliasesDeleteAliases] = Field(alias="deleteAliases") + + +class DeleteAliasesDeleteAliases(GQLBase): + success: bool + + +DeleteAliases.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_artifact_collection_tag_assignments.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_artifact_collection_tag_assignments.py new file mode 100644 index 0000000000000000000000000000000000000000..d89d18340cd4a4cb17d09af8edef4e72c278151c --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_artifact_collection_tag_assignments.py @@ -0,0 +1,25 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + + +class DeleteArtifactCollectionTagAssignments(GQLBase): + delete_artifact_collection_tag_assignments: Optional[ + DeleteArtifactCollectionTagAssignmentsDeleteArtifactCollectionTagAssignments + ] = Field(alias="deleteArtifactCollectionTagAssignments") + + +class DeleteArtifactCollectionTagAssignmentsDeleteArtifactCollectionTagAssignments( + GQLBase +): + success: bool + + +DeleteArtifactCollectionTagAssignments.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_artifact_portfolio.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_artifact_portfolio.py new file mode 100644 index 0000000000000000000000000000000000000000..91b88aed6108959893e101955baf75840f9eaba0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_artifact_portfolio.py @@ -0,0 +1,35 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, Typename + +from .enums import ArtifactCollectionState + + +class DeleteArtifactPortfolio(GQLBase): + delete_artifact_portfolio: Optional[ + DeleteArtifactPortfolioDeleteArtifactPortfolio + ] = Field(alias="deleteArtifactPortfolio") + + +class DeleteArtifactPortfolioDeleteArtifactPortfolio(GQLBase): + artifact_collection: DeleteArtifactPortfolioDeleteArtifactPortfolioArtifactCollection = Field( + alias="artifactCollection" + ) + + +class DeleteArtifactPortfolioDeleteArtifactPortfolioArtifactCollection(GQLBase): + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + state: ArtifactCollectionState + + +DeleteArtifactPortfolio.model_rebuild() +DeleteArtifactPortfolioDeleteArtifactPortfolio.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_artifact_sequence.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_artifact_sequence.py new file mode 100644 index 0000000000000000000000000000000000000000..85eeca9bfb32f452b6281733fbfcf5ef5b4f0700 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/delete_artifact_sequence.py @@ -0,0 +1,35 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, Typename + +from .enums import ArtifactCollectionState + + +class DeleteArtifactSequence(GQLBase): + delete_artifact_sequence: Optional[DeleteArtifactSequenceDeleteArtifactSequence] = ( + Field(alias="deleteArtifactSequence") + ) + + +class DeleteArtifactSequenceDeleteArtifactSequence(GQLBase): + artifact_collection: DeleteArtifactSequenceDeleteArtifactSequenceArtifactCollection = Field( + alias="artifactCollection" + ) + + +class DeleteArtifactSequenceDeleteArtifactSequenceArtifactCollection(GQLBase): + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + state: ArtifactCollectionState + + +DeleteArtifactSequence.model_rebuild() +DeleteArtifactSequenceDeleteArtifactSequence.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/enums.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/enums.py new file mode 100644 index 0000000000000000000000000000000000000000..0d102f8a72f024db4de1b5311b02f331ebe6a4f7 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/enums.py @@ -0,0 +1,17 @@ +# Generated by ariadne-codegen +# Source: core/api/graphql/schemas/schema-latest.graphql + +from __future__ import annotations + +from enum import Enum + + +class ArtifactState(str, Enum): + PENDING = "PENDING" + COMMITTED = "COMMITTED" + DELETED = "DELETED" + + +class ArtifactCollectionState(str, Enum): + READY = "READY" + DELETED = "DELETED" diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/fetch_linked_artifacts.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/fetch_linked_artifacts.py new file mode 100644 index 0000000000000000000000000000000000000000..2617735522a65f6b37c76d80952fe3d3b1065495 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/fetch_linked_artifacts.py @@ -0,0 +1,67 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import List, Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, Typename + + +class FetchLinkedArtifacts(GQLBase): + artifact: Optional[FetchLinkedArtifactsArtifact] + + +class FetchLinkedArtifactsArtifact(GQLBase): + artifact_memberships: FetchLinkedArtifactsArtifactArtifactMemberships = Field( + alias="artifactMemberships" + ) + + +class FetchLinkedArtifactsArtifactArtifactMemberships(GQLBase): + edges: List[FetchLinkedArtifactsArtifactArtifactMembershipsEdges] + + +class FetchLinkedArtifactsArtifactArtifactMembershipsEdges(GQLBase): + node: Optional[FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNode] + + +class FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNode(GQLBase): + aliases: List[FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeAliases] + version_index: Optional[int] = Field(alias="versionIndex") + artifact_collection: Optional[ + FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeArtifactCollection + ] = Field(alias="artifactCollection") + + +class FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeAliases(GQLBase): + alias: str + + +class FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeArtifactCollection( + GQLBase +): + project: Optional[ + FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeArtifactCollectionProject + ] + name: str + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + + +class FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeArtifactCollectionProject( + GQLBase +): + entity_name: str = Field(alias="entityName") + name: str + + +FetchLinkedArtifacts.model_rebuild() +FetchLinkedArtifactsArtifact.model_rebuild() +FetchLinkedArtifactsArtifactArtifactMemberships.model_rebuild() +FetchLinkedArtifactsArtifactArtifactMembershipsEdges.model_rebuild() +FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNode.model_rebuild() +FetchLinkedArtifactsArtifactArtifactMembershipsEdgesNodeArtifactCollection.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/fragments.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/fragments.py new file mode 100644 index 0000000000000000000000000000000000000000..46ebe35435b489409fc974b5025c43b5b10f1cca --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/fragments.py @@ -0,0 +1,221 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Any, List, Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, GQLId, Typename + +from .enums import ArtifactState + + +class ArtifactCollectionsFragment(GQLBase): + page_info: ArtifactCollectionsFragmentPageInfo = Field(alias="pageInfo") + total_count: int = Field(alias="totalCount") + edges: List[ArtifactCollectionsFragmentEdges] + + +class ArtifactCollectionsFragmentEdges(GQLBase): + node: Optional[ArtifactCollectionsFragmentEdgesNode] + cursor: str + + +class ArtifactCollectionsFragmentEdgesNode(GQLBase): + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + id: GQLId + name: str + description: Optional[str] + created_at: str = Field(alias="createdAt") + + +class ArtifactCollectionsFragmentPageInfo(GQLBase): + end_cursor: Optional[str] = Field(alias="endCursor") + has_next_page: bool = Field(alias="hasNextPage") + + +class ArtifactFragment(GQLBase): + id: GQLId + artifact_sequence: ArtifactFragmentArtifactSequence = Field( + alias="artifactSequence" + ) + version_index: Optional[int] = Field(alias="versionIndex") + artifact_type: ArtifactFragmentArtifactType = Field(alias="artifactType") + description: Optional[str] + metadata: Optional[Any] + ttl_duration_seconds: Optional[Any] = Field( + alias="ttlDurationSeconds", default=None + ) + ttl_is_inherited: Optional[bool] = Field(alias="ttlIsInherited", default=None) + aliases: Optional[List[ArtifactFragmentAliases]] = None + tags: Optional[List[ArtifactFragmentTags]] = None + history_step: Optional[Any] = Field(alias="historyStep", default=None) + state: ArtifactState + current_manifest: Optional[ArtifactFragmentCurrentManifest] = Field( + alias="currentManifest" + ) + commit_hash: Optional[str] = Field(alias="commitHash") + file_count: Any = Field(alias="fileCount") + created_at: str = Field(alias="createdAt") + updated_at: Optional[str] = Field(alias="updatedAt") + + +class ArtifactFragmentAliases(GQLBase): + artifact_collection: Optional[ArtifactFragmentAliasesArtifactCollection] = Field( + alias="artifactCollection" + ) + alias: str + + +class ArtifactFragmentAliasesArtifactCollection(GQLBase): + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + project: Optional[ArtifactFragmentAliasesArtifactCollectionProject] + name: str + + +class ArtifactFragmentAliasesArtifactCollectionProject(GQLBase): + entity_name: str = Field(alias="entityName") + name: str + + +class ArtifactFragmentArtifactSequence(GQLBase): + project: Optional[ArtifactFragmentArtifactSequenceProject] + name: str + + +class ArtifactFragmentArtifactSequenceProject(GQLBase): + entity_name: str = Field(alias="entityName") + name: str + + +class ArtifactFragmentArtifactType(GQLBase): + name: str + + +class ArtifactFragmentCurrentManifest(GQLBase): + file: ArtifactFragmentCurrentManifestFile + + +class ArtifactFragmentCurrentManifestFile(GQLBase): + direct_url: str = Field(alias="directUrl") + + +class ArtifactFragmentTags(GQLBase): + name: str + + +class ArtifactPortfolioTypeFields(GQLBase): + typename__: Typename[Literal["ArtifactPortfolio"]] = "ArtifactPortfolio" + id: GQLId + name: str + + +class ArtifactSequenceTypeFields(GQLBase): + typename__: Typename[Literal["ArtifactSequence"]] = "ArtifactSequence" + id: GQLId + name: str + + +class ArtifactTypeFragment(GQLBase): + typename__: Typename[Literal["ArtifactType"]] = "ArtifactType" + id: GQLId + name: str + description: Optional[str] + created_at: str = Field(alias="createdAt") + + +class ArtifactTypesFragment(GQLBase): + edges: List[ArtifactTypesFragmentEdges] + page_info: ArtifactTypesFragmentPageInfo = Field(alias="pageInfo") + + +class ArtifactTypesFragmentEdges(GQLBase): + node: Optional[ArtifactTypeFragment] + cursor: str + + +class ArtifactTypesFragmentPageInfo(GQLBase): + end_cursor: Optional[str] = Field(alias="endCursor") + has_next_page: bool = Field(alias="hasNextPage") + + +class ArtifactsFragment(GQLBase): + total_count: int = Field(alias="totalCount") + edges: List[ArtifactsFragmentEdges] + page_info: ArtifactsFragmentPageInfo = Field(alias="pageInfo") + + +class ArtifactsFragmentEdges(GQLBase): + node: ArtifactFragment + version: str + cursor: str + + +class ArtifactsFragmentPageInfo(GQLBase): + end_cursor: Optional[str] = Field(alias="endCursor") + has_next_page: bool = Field(alias="hasNextPage") + + +class FilesFragment(GQLBase): + edges: List[FilesFragmentEdges] + page_info: FilesFragmentPageInfo = Field(alias="pageInfo") + + +class FilesFragmentEdges(GQLBase): + node: Optional[FilesFragmentEdgesNode] + cursor: str + + +class FilesFragmentEdgesNode(GQLBase): + id: GQLId + name: str + url: Optional[str] + size_bytes: Any = Field(alias="sizeBytes") + storage_path: Optional[str] = Field(alias="storagePath", default=None) + mimetype: Optional[str] + updated_at: Optional[str] = Field(alias="updatedAt") + digest: Optional[str] + md_5: Optional[str] = Field(alias="md5") + direct_url: str = Field(alias="directUrl") + + +class FilesFragmentPageInfo(GQLBase): + end_cursor: Optional[str] = Field(alias="endCursor") + has_next_page: bool = Field(alias="hasNextPage") + + +ArtifactCollectionsFragment.model_rebuild() +ArtifactCollectionsFragmentEdges.model_rebuild() +ArtifactCollectionsFragmentEdgesNode.model_rebuild() +ArtifactCollectionsFragmentPageInfo.model_rebuild() +ArtifactFragment.model_rebuild() +ArtifactFragmentAliases.model_rebuild() +ArtifactFragmentAliasesArtifactCollection.model_rebuild() +ArtifactFragmentAliasesArtifactCollectionProject.model_rebuild() +ArtifactFragmentArtifactSequence.model_rebuild() +ArtifactFragmentArtifactSequenceProject.model_rebuild() +ArtifactFragmentArtifactType.model_rebuild() +ArtifactFragmentCurrentManifest.model_rebuild() +ArtifactFragmentCurrentManifestFile.model_rebuild() +ArtifactFragmentTags.model_rebuild() +ArtifactPortfolioTypeFields.model_rebuild() +ArtifactSequenceTypeFields.model_rebuild() +ArtifactTypeFragment.model_rebuild() +ArtifactTypesFragment.model_rebuild() +ArtifactTypesFragmentEdges.model_rebuild() +ArtifactTypesFragmentPageInfo.model_rebuild() +ArtifactsFragment.model_rebuild() +ArtifactsFragmentEdges.model_rebuild() +ArtifactsFragmentPageInfo.model_rebuild() +FilesFragment.model_rebuild() +FilesFragmentEdges.model_rebuild() +FilesFragmentEdgesNode.model_rebuild() +FilesFragmentPageInfo.model_rebuild() +ArtifactFragment.model_rebuild() +ArtifactTypeFragment.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/input_types.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/input_types.py new file mode 100644 index 0000000000000000000000000000000000000000..2b3f1b971142127810c96b52b222453b00382935 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/input_types.py @@ -0,0 +1,46 @@ +# Generated by ariadne-codegen +# Source: core/api/graphql/schemas/schema-latest.graphql + +from __future__ import annotations + +from typing import Any, List, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, GQLId + + +class ArtifactAliasInput(GQLBase): + artifact_collection_name: str = Field(alias="artifactCollectionName") + alias: str + + +class LinkArtifactInput(GQLBase): + artifact_id: Optional[GQLId] = Field(alias="artifactID", default=None) + artifact_portfolio_id: Optional[GQLId] = Field( + alias="artifactPortfolioID", default=None + ) + artifact_portfolio_name: Optional[str] = Field( + alias="artifactPortfolioName", default=None + ) + entity_name: Optional[str] = Field(alias="entityName", default=None) + project_name: Optional[str] = Field(alias="projectName", default=None) + aliases: Optional[List[ArtifactAliasInput]] = None + client_id: Optional[GQLId] = Field(alias="clientID", default=None) + client_mutation_id: Optional[str] = Field(alias="clientMutationId", default=None) + + +class ArtifactCollectionAliasInput(GQLBase): + alias: str + entity_name: str = Field(alias="entityName") + project_name: str = Field(alias="projectName") + artifact_collection_name: str = Field(alias="artifactCollectionName") + + +class TagInput(GQLBase): + tag_category_name: Optional[str] = Field(alias="tagCategoryName", default=None) + tag_name: str = Field(alias="tagName") + attributes: Optional[Any] = None + + +LinkArtifactInput.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/link_artifact.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/link_artifact.py new file mode 100644 index 0000000000000000000000000000000000000000..086f8b5a8a57d44d81770de9471a3eef66ca1ac6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/link_artifact.py @@ -0,0 +1,21 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + + +class LinkArtifact(GQLBase): + link_artifact: Optional[LinkArtifactLinkArtifact] = Field(alias="linkArtifact") + + +class LinkArtifactLinkArtifact(GQLBase): + version_index: Optional[int] = Field(alias="versionIndex") + + +LinkArtifact.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/move_artifact_collection.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/move_artifact_collection.py new file mode 100644 index 0000000000000000000000000000000000000000..31c142e24ace6a111ffc6c53aefe192d9a3c293b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/move_artifact_collection.py @@ -0,0 +1,35 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, GQLId, Typename + + +class MoveArtifactCollection(GQLBase): + move_artifact_sequence: Optional[MoveArtifactCollectionMoveArtifactSequence] = ( + Field(alias="moveArtifactSequence") + ) + + +class MoveArtifactCollectionMoveArtifactSequence(GQLBase): + artifact_collection: Optional[ + MoveArtifactCollectionMoveArtifactSequenceArtifactCollection + ] = Field(alias="artifactCollection") + + +class MoveArtifactCollectionMoveArtifactSequenceArtifactCollection(GQLBase): + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + id: GQLId + name: str + description: Optional[str] + + +MoveArtifactCollection.model_rebuild() +MoveArtifactCollectionMoveArtifactSequence.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/operations.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/operations.py new file mode 100644 index 0000000000000000000000000000000000000000..ed6d6061b9fa574d7e654fb2d9ab32d17f0408a6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/operations.py @@ -0,0 +1,620 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +__all__ = [ + "ADD_ALIASES_GQL", + "ARTIFACT_COLLECTION_MEMBERSHIP_FILES_GQL", + "ARTIFACT_VERSION_FILES_GQL", + "CREATE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL", + "DELETE_ALIASES_GQL", + "DELETE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL", + "DELETE_ARTIFACT_PORTFOLIO_GQL", + "DELETE_ARTIFACT_SEQUENCE_GQL", + "FETCH_LINKED_ARTIFACTS_GQL", + "LINK_ARTIFACT_GQL", + "MOVE_ARTIFACT_COLLECTION_GQL", + "PROJECT_ARTIFACTS_GQL", + "PROJECT_ARTIFACT_COLLECTIONS_GQL", + "PROJECT_ARTIFACT_COLLECTION_GQL", + "PROJECT_ARTIFACT_TYPES_GQL", + "PROJECT_ARTIFACT_TYPE_GQL", + "RUN_INPUT_ARTIFACTS_GQL", + "RUN_OUTPUT_ARTIFACTS_GQL", + "UPDATE_ARTIFACT_GQL", + "UPDATE_ARTIFACT_PORTFOLIO_GQL", + "UPDATE_ARTIFACT_SEQUENCE_GQL", +] + +DELETE_ARTIFACT_SEQUENCE_GQL = """ +mutation DeleteArtifactSequence($id: ID!) { + deleteArtifactSequence(input: {artifactSequenceID: $id}) { + artifactCollection { + __typename + state + } + } +} +""" + +DELETE_ARTIFACT_PORTFOLIO_GQL = """ +mutation DeleteArtifactPortfolio($id: ID!) { + deleteArtifactPortfolio(input: {artifactPortfolioID: $id}) { + artifactCollection { + __typename + state + } + } +} +""" + +UPDATE_ARTIFACT_SEQUENCE_GQL = """ +mutation UpdateArtifactSequence($id: ID!, $name: String, $description: String) { + updateArtifactSequence( + input: {artifactSequenceID: $id, name: $name, description: $description} + ) { + artifactCollection { + __typename + id + name + description + } + } +} +""" + +UPDATE_ARTIFACT_PORTFOLIO_GQL = """ +mutation UpdateArtifactPortfolio($id: ID!, $name: String, $description: String) { + updateArtifactPortfolio( + input: {artifactPortfolioID: $id, name: $name, description: $description} + ) { + artifactCollection { + __typename + id + name + description + } + } +} +""" + +MOVE_ARTIFACT_COLLECTION_GQL = """ +mutation MoveArtifactCollection($artifactSequenceID: ID!, $destinationArtifactTypeName: String!) { + moveArtifactSequence( + input: {artifactSequenceID: $artifactSequenceID, destinationArtifactTypeName: $destinationArtifactTypeName} + ) { + artifactCollection { + __typename + id + name + description + } + } +} +""" + +CREATE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL = """ +mutation CreateArtifactCollectionTagAssignments($entityName: String!, $projectName: String!, $artifactCollectionName: String!, $tags: [TagInput!]!) { + createArtifactCollectionTagAssignments( + input: {entityName: $entityName, projectName: $projectName, artifactCollectionName: $artifactCollectionName, tags: $tags} + ) { + tags { + id + name + tagCategoryName + } + } +} +""" + +DELETE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL = """ +mutation DeleteArtifactCollectionTagAssignments($entityName: String!, $projectName: String!, $artifactCollectionName: String!, $tags: [TagInput!]!) { + deleteArtifactCollectionTagAssignments( + input: {entityName: $entityName, projectName: $projectName, artifactCollectionName: $artifactCollectionName, tags: $tags} + ) { + success + } +} +""" + +PROJECT_ARTIFACT_COLLECTIONS_GQL = """ +query ProjectArtifactCollections($entityName: String!, $projectName: String!, $artifactTypeName: String!, $cursor: String) { + project(name: $projectName, entityName: $entityName) { + artifactType(name: $artifactTypeName) { + artifactCollections: artifactCollections(after: $cursor) { + ...ArtifactCollectionsFragment + } + } + } +} + +fragment ArtifactCollectionsFragment on ArtifactCollectionConnection { + pageInfo { + endCursor + hasNextPage + } + totalCount + edges { + node { + __typename + id + name + description + createdAt + } + cursor + } +} +""" + +PROJECT_ARTIFACT_COLLECTION_GQL = """ +query ProjectArtifactCollection($entityName: String!, $projectName: String!, $artifactTypeName: String!, $artifactCollectionName: String!, $cursor: String, $perPage: Int = 1000) { + project(name: $projectName, entityName: $entityName) { + artifactType(name: $artifactTypeName) { + artifactCollection: artifactCollection(name: $artifactCollectionName) { + __typename + id + name + description + createdAt + tags { + edges { + node { + id + name + } + } + } + aliases(after: $cursor, first: $perPage) { + edges { + node { + alias + } + cursor + } + pageInfo { + endCursor + hasNextPage + } + } + } + artifactSequence(name: $artifactCollectionName) { + __typename + } + } + } +} +""" + +ARTIFACT_VERSION_FILES_GQL = """ +query ArtifactVersionFiles($entityName: String!, $projectName: String!, $artifactTypeName: String!, $artifactName: String!, $fileNames: [String!], $fileCursor: String, $fileLimit: Int = 50) { + project(name: $projectName, entityName: $entityName) { + artifactType(name: $artifactTypeName) { + artifact(name: $artifactName) { + files(names: $fileNames, after: $fileCursor, first: $fileLimit) { + ...FilesFragment + } + } + } + } +} + +fragment FilesFragment on FileConnection { + edges { + node { + id + name: displayName + url + sizeBytes + storagePath @include(if: true) + mimetype + updatedAt + digest + md5 + directUrl + } + cursor + } + pageInfo { + endCursor + hasNextPage + } +} +""" + +ARTIFACT_COLLECTION_MEMBERSHIP_FILES_GQL = """ +query ArtifactCollectionMembershipFiles($entityName: String!, $projectName: String!, $artifactName: String!, $artifactVersionIndex: String!, $fileNames: [String!], $fileCursor: String, $fileLimit: Int = 50) { + project(name: $projectName, entityName: $entityName) { + artifactCollection(name: $artifactName) { + __typename + artifactMembership(aliasName: $artifactVersionIndex) { + files(names: $fileNames, after: $fileCursor, first: $fileLimit) { + ...FilesFragment + } + } + } + } +} + +fragment FilesFragment on FileConnection { + edges { + node { + id + name: displayName + url + sizeBytes + storagePath @include(if: true) + mimetype + updatedAt + digest + md5 + directUrl + } + cursor + } + pageInfo { + endCursor + hasNextPage + } +} +""" + +PROJECT_ARTIFACT_TYPES_GQL = """ +query ProjectArtifactTypes($entityName: String!, $projectName: String!, $cursor: String) { + project(name: $projectName, entityName: $entityName) { + artifactTypes(after: $cursor) { + ...ArtifactTypesFragment + } + } +} + +fragment ArtifactTypeFragment on ArtifactType { + __typename + id + name + description + createdAt +} + +fragment ArtifactTypesFragment on ArtifactTypeConnection { + edges { + node { + ...ArtifactTypeFragment + } + cursor + } + pageInfo { + endCursor + hasNextPage + } +} +""" + +PROJECT_ARTIFACT_TYPE_GQL = """ +query ProjectArtifactType($entityName: String!, $projectName: String!, $artifactTypeName: String!) { + project(name: $projectName, entityName: $entityName) { + artifactType(name: $artifactTypeName) { + ...ArtifactTypeFragment + } + } +} + +fragment ArtifactTypeFragment on ArtifactType { + __typename + id + name + description + createdAt +} +""" + +PROJECT_ARTIFACTS_GQL = """ +query ProjectArtifacts($project: String!, $entity: String!, $type: String!, $collection: String!, $cursor: String, $perPage: Int = 50, $order: String, $filters: JSONString) { + project(name: $project, entityName: $entity) { + artifactType(name: $type) { + artifactCollection: artifactCollection(name: $collection) { + __typename + name + artifacts(filters: $filters, after: $cursor, first: $perPage, order: $order) { + ...ArtifactsFragment + } + } + } + } +} + +fragment ArtifactFragment on Artifact { + id + artifactSequence { + project { + entityName + name + } + name + } + versionIndex + artifactType { + name + } + description + metadata + ttlDurationSeconds @include(if: true) + ttlIsInherited @include(if: true) + aliases @include(if: true) { + artifactCollection { + __typename + project { + entityName + name + } + name + } + alias + } + tags @include(if: true) { + name + } + historyStep @include(if: true) + state + currentManifest { + file { + directUrl + } + } + commitHash + fileCount + createdAt + updatedAt +} + +fragment ArtifactsFragment on VersionedArtifactConnection { + totalCount + edges { + node { + ...ArtifactFragment + } + version + cursor + } + pageInfo { + endCursor + hasNextPage + } +} +""" + +RUN_OUTPUT_ARTIFACTS_GQL = """ +query RunOutputArtifacts($entity: String!, $project: String!, $runName: String!, $cursor: String, $perPage: Int) { + project(name: $project, entityName: $entity) { + run(name: $runName) { + outputArtifacts(after: $cursor, first: $perPage) { + totalCount + edges { + node { + ...ArtifactFragment + } + cursor + } + pageInfo { + endCursor + hasNextPage + } + } + } + } +} + +fragment ArtifactFragment on Artifact { + id + artifactSequence { + project { + entityName + name + } + name + } + versionIndex + artifactType { + name + } + description + metadata + ttlDurationSeconds @include(if: true) + ttlIsInherited @include(if: true) + aliases @include(if: true) { + artifactCollection { + __typename + project { + entityName + name + } + name + } + alias + } + tags @include(if: true) { + name + } + historyStep @include(if: true) + state + currentManifest { + file { + directUrl + } + } + commitHash + fileCount + createdAt + updatedAt +} +""" + +RUN_INPUT_ARTIFACTS_GQL = """ +query RunInputArtifacts($entity: String!, $project: String!, $runName: String!, $cursor: String, $perPage: Int) { + project(name: $project, entityName: $entity) { + run(name: $runName) { + inputArtifacts(after: $cursor, first: $perPage) { + totalCount + edges { + node { + ...ArtifactFragment + } + cursor + } + pageInfo { + endCursor + hasNextPage + } + } + } + } +} + +fragment ArtifactFragment on Artifact { + id + artifactSequence { + project { + entityName + name + } + name + } + versionIndex + artifactType { + name + } + description + metadata + ttlDurationSeconds @include(if: true) + ttlIsInherited @include(if: true) + aliases @include(if: true) { + artifactCollection { + __typename + project { + entityName + name + } + name + } + alias + } + tags @include(if: true) { + name + } + historyStep @include(if: true) + state + currentManifest { + file { + directUrl + } + } + commitHash + fileCount + createdAt + updatedAt +} +""" + +FETCH_LINKED_ARTIFACTS_GQL = """ +query FetchLinkedArtifacts($artifactID: ID!) { + artifact(id: $artifactID) { + artifactMemberships { + edges { + node { + aliases { + alias + } + versionIndex + artifactCollection { + project { + entityName + name + } + name + __typename + } + } + } + } + } +} +""" + +ADD_ALIASES_GQL = """ +mutation AddAliases($artifactID: ID!, $aliases: [ArtifactCollectionAliasInput!]!) { + addAliases(input: {artifactID: $artifactID, aliases: $aliases}) { + success + } +} +""" + +DELETE_ALIASES_GQL = """ +mutation DeleteAliases($artifactID: ID!, $aliases: [ArtifactCollectionAliasInput!]!) { + deleteAliases(input: {artifactID: $artifactID, aliases: $aliases}) { + success + } +} +""" + +UPDATE_ARTIFACT_GQL = """ +mutation UpdateArtifact($artifactID: ID!, $description: String, $metadata: JSONString, $ttlDurationSeconds: Int64, $tagsToAdd: [TagInput!], $tagsToDelete: [TagInput!], $aliases: [ArtifactAliasInput!]) { + updateArtifact( + input: {artifactID: $artifactID, description: $description, metadata: $metadata, ttlDurationSeconds: $ttlDurationSeconds, tagsToAdd: $tagsToAdd, tagsToDelete: $tagsToDelete, aliases: $aliases} + ) { + artifact { + ...ArtifactFragment + } + } +} + +fragment ArtifactFragment on Artifact { + id + artifactSequence { + project { + entityName + name + } + name + } + versionIndex + artifactType { + name + } + description + metadata + ttlDurationSeconds @include(if: true) + ttlIsInherited @include(if: true) + aliases @include(if: true) { + artifactCollection { + __typename + project { + entityName + name + } + name + } + alias + } + tags @include(if: true) { + name + } + historyStep @include(if: true) + state + currentManifest { + file { + directUrl + } + } + commitHash + fileCount + createdAt + updatedAt +} +""" + +LINK_ARTIFACT_GQL = """ +mutation LinkArtifact($input: LinkArtifactInput!) { + linkArtifact(input: $input) { + versionIndex + } +} +""" diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_collection.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_collection.py new file mode 100644 index 0000000000000000000000000000000000000000..12e2342db24aae84e79b87e287f0abb87d46453f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_collection.py @@ -0,0 +1,101 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import List, Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, GQLId, Typename + + +class ProjectArtifactCollection(GQLBase): + project: Optional[ProjectArtifactCollectionProject] + + +class ProjectArtifactCollectionProject(GQLBase): + artifact_type: Optional[ProjectArtifactCollectionProjectArtifactType] = Field( + alias="artifactType" + ) + + +class ProjectArtifactCollectionProjectArtifactType(GQLBase): + artifact_collection: Optional[ + ProjectArtifactCollectionProjectArtifactTypeArtifactCollection + ] = Field(alias="artifactCollection") + artifact_sequence: Optional[ + ProjectArtifactCollectionProjectArtifactTypeArtifactSequence + ] = Field(alias="artifactSequence") + + +class ProjectArtifactCollectionProjectArtifactTypeArtifactCollection(GQLBase): + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + id: GQLId + name: str + description: Optional[str] + created_at: str = Field(alias="createdAt") + tags: ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTags + aliases: ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliases + + +class ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTags(GQLBase): + edges: List[ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTagsEdges] + + +class ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTagsEdges(GQLBase): + node: ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTagsEdgesNode + + +class ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTagsEdgesNode( + GQLBase +): + id: GQLId + name: str + + +class ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliases(GQLBase): + edges: List[ + ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesEdges + ] + page_info: ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesPageInfo = Field( + alias="pageInfo" + ) + + +class ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesEdges( + GQLBase +): + node: Optional[ + ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesEdgesNode + ] + cursor: str + + +class ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesEdgesNode( + GQLBase +): + alias: str + + +class ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesPageInfo( + GQLBase +): + end_cursor: Optional[str] = Field(alias="endCursor") + has_next_page: bool = Field(alias="hasNextPage") + + +class ProjectArtifactCollectionProjectArtifactTypeArtifactSequence(GQLBase): + typename__: Typename[Literal["ArtifactSequence"]] + + +ProjectArtifactCollection.model_rebuild() +ProjectArtifactCollectionProject.model_rebuild() +ProjectArtifactCollectionProjectArtifactType.model_rebuild() +ProjectArtifactCollectionProjectArtifactTypeArtifactCollection.model_rebuild() +ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTags.model_rebuild() +ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionTagsEdges.model_rebuild() +ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliases.model_rebuild() +ProjectArtifactCollectionProjectArtifactTypeArtifactCollectionAliasesEdges.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_collections.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_collections.py new file mode 100644 index 0000000000000000000000000000000000000000..d6d113d0d0ce6f0f34f2c6689ca0309e373559e3 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_collections.py @@ -0,0 +1,33 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + +from .fragments import ArtifactCollectionsFragment + + +class ProjectArtifactCollections(GQLBase): + project: Optional[ProjectArtifactCollectionsProject] + + +class ProjectArtifactCollectionsProject(GQLBase): + artifact_type: Optional[ProjectArtifactCollectionsProjectArtifactType] = Field( + alias="artifactType" + ) + + +class ProjectArtifactCollectionsProjectArtifactType(GQLBase): + artifact_collections: Optional[ArtifactCollectionsFragment] = Field( + alias="artifactCollections" + ) + + +ProjectArtifactCollections.model_rebuild() +ProjectArtifactCollectionsProject.model_rebuild() +ProjectArtifactCollectionsProjectArtifactType.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_type.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_type.py new file mode 100644 index 0000000000000000000000000000000000000000..34c3ed067fba86c70f0e8b1f579eac564675eeaa --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_type.py @@ -0,0 +1,24 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + +from .fragments import ArtifactTypeFragment + + +class ProjectArtifactType(GQLBase): + project: Optional[ProjectArtifactTypeProject] + + +class ProjectArtifactTypeProject(GQLBase): + artifact_type: Optional[ArtifactTypeFragment] = Field(alias="artifactType") + + +ProjectArtifactType.model_rebuild() +ProjectArtifactTypeProject.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_types.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_types.py new file mode 100644 index 0000000000000000000000000000000000000000..e722d692620f0c496239e7e2d4324bb7c1c38cf8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifact_types.py @@ -0,0 +1,24 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + +from .fragments import ArtifactTypesFragment + + +class ProjectArtifactTypes(GQLBase): + project: Optional[ProjectArtifactTypesProject] + + +class ProjectArtifactTypesProject(GQLBase): + artifact_types: ArtifactTypesFragment = Field(alias="artifactTypes") + + +ProjectArtifactTypes.model_rebuild() +ProjectArtifactTypesProject.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifacts.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifacts.py new file mode 100644 index 0000000000000000000000000000000000000000..4ad46c3196e3bbf2de696e798debd23e8d9e3d72 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/project_artifacts.py @@ -0,0 +1,42 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, Typename + +from .fragments import ArtifactsFragment + + +class ProjectArtifacts(GQLBase): + project: Optional[ProjectArtifactsProject] + + +class ProjectArtifactsProject(GQLBase): + artifact_type: Optional[ProjectArtifactsProjectArtifactType] = Field( + alias="artifactType" + ) + + +class ProjectArtifactsProjectArtifactType(GQLBase): + artifact_collection: Optional[ + ProjectArtifactsProjectArtifactTypeArtifactCollection + ] = Field(alias="artifactCollection") + + +class ProjectArtifactsProjectArtifactTypeArtifactCollection(GQLBase): + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + name: str + artifacts: Optional[ArtifactsFragment] + + +ProjectArtifacts.model_rebuild() +ProjectArtifactsProject.model_rebuild() +ProjectArtifactsProjectArtifactType.model_rebuild() +ProjectArtifactsProjectArtifactTypeArtifactCollection.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/run_input_artifacts.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/run_input_artifacts.py new file mode 100644 index 0000000000000000000000000000000000000000..7ac3dca13e97d5ce67c3d18394db35107b825f24 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/run_input_artifacts.py @@ -0,0 +1,51 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import List, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + +from .fragments import ArtifactFragment + + +class RunInputArtifacts(GQLBase): + project: Optional[RunInputArtifactsProject] + + +class RunInputArtifactsProject(GQLBase): + run: Optional[RunInputArtifactsProjectRun] + + +class RunInputArtifactsProjectRun(GQLBase): + input_artifacts: Optional[RunInputArtifactsProjectRunInputArtifacts] = Field( + alias="inputArtifacts" + ) + + +class RunInputArtifactsProjectRunInputArtifacts(GQLBase): + total_count: int = Field(alias="totalCount") + edges: List[RunInputArtifactsProjectRunInputArtifactsEdges] + page_info: RunInputArtifactsProjectRunInputArtifactsPageInfo = Field( + alias="pageInfo" + ) + + +class RunInputArtifactsProjectRunInputArtifactsEdges(GQLBase): + node: Optional[ArtifactFragment] + cursor: str + + +class RunInputArtifactsProjectRunInputArtifactsPageInfo(GQLBase): + end_cursor: Optional[str] = Field(alias="endCursor") + has_next_page: bool = Field(alias="hasNextPage") + + +RunInputArtifacts.model_rebuild() +RunInputArtifactsProject.model_rebuild() +RunInputArtifactsProjectRun.model_rebuild() +RunInputArtifactsProjectRunInputArtifacts.model_rebuild() +RunInputArtifactsProjectRunInputArtifactsEdges.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/run_output_artifacts.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/run_output_artifacts.py new file mode 100644 index 0000000000000000000000000000000000000000..80386993a2095586b2a6ab29ea0865657c076cf0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/run_output_artifacts.py @@ -0,0 +1,51 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import List, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + +from .fragments import ArtifactFragment + + +class RunOutputArtifacts(GQLBase): + project: Optional[RunOutputArtifactsProject] + + +class RunOutputArtifactsProject(GQLBase): + run: Optional[RunOutputArtifactsProjectRun] + + +class RunOutputArtifactsProjectRun(GQLBase): + output_artifacts: Optional[RunOutputArtifactsProjectRunOutputArtifacts] = Field( + alias="outputArtifacts" + ) + + +class RunOutputArtifactsProjectRunOutputArtifacts(GQLBase): + total_count: int = Field(alias="totalCount") + edges: List[RunOutputArtifactsProjectRunOutputArtifactsEdges] + page_info: RunOutputArtifactsProjectRunOutputArtifactsPageInfo = Field( + alias="pageInfo" + ) + + +class RunOutputArtifactsProjectRunOutputArtifactsEdges(GQLBase): + node: Optional[ArtifactFragment] + cursor: str + + +class RunOutputArtifactsProjectRunOutputArtifactsPageInfo(GQLBase): + end_cursor: Optional[str] = Field(alias="endCursor") + has_next_page: bool = Field(alias="hasNextPage") + + +RunOutputArtifacts.model_rebuild() +RunOutputArtifactsProject.model_rebuild() +RunOutputArtifactsProjectRun.model_rebuild() +RunOutputArtifactsProjectRunOutputArtifacts.model_rebuild() +RunOutputArtifactsProjectRunOutputArtifactsEdges.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/update_artifact.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/update_artifact.py new file mode 100644 index 0000000000000000000000000000000000000000..7eaf59e7e9dde4e818449319e77f52912f6d322d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/update_artifact.py @@ -0,0 +1,26 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + +from .fragments import ArtifactFragment + + +class UpdateArtifact(GQLBase): + update_artifact: Optional[UpdateArtifactUpdateArtifact] = Field( + alias="updateArtifact" + ) + + +class UpdateArtifactUpdateArtifact(GQLBase): + artifact: ArtifactFragment + + +UpdateArtifact.model_rebuild() +UpdateArtifactUpdateArtifact.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/update_artifact_portfolio.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/update_artifact_portfolio.py new file mode 100644 index 0000000000000000000000000000000000000000..adae2799116e5adeb32599803680d89d1e8682bd --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/update_artifact_portfolio.py @@ -0,0 +1,35 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, GQLId, Typename + + +class UpdateArtifactPortfolio(GQLBase): + update_artifact_portfolio: Optional[ + UpdateArtifactPortfolioUpdateArtifactPortfolio + ] = Field(alias="updateArtifactPortfolio") + + +class UpdateArtifactPortfolioUpdateArtifactPortfolio(GQLBase): + artifact_collection: UpdateArtifactPortfolioUpdateArtifactPortfolioArtifactCollection = Field( + alias="artifactCollection" + ) + + +class UpdateArtifactPortfolioUpdateArtifactPortfolioArtifactCollection(GQLBase): + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + id: GQLId + name: str + description: Optional[str] + + +UpdateArtifactPortfolio.model_rebuild() +UpdateArtifactPortfolioUpdateArtifactPortfolio.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/update_artifact_sequence.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/update_artifact_sequence.py new file mode 100644 index 0000000000000000000000000000000000000000..4e875f1f732a0a33a989701ea0d37fdc43f1f34b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_generated/update_artifact_sequence.py @@ -0,0 +1,35 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/artifacts/ + +from __future__ import annotations + +from typing import Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, GQLId, Typename + + +class UpdateArtifactSequence(GQLBase): + update_artifact_sequence: Optional[UpdateArtifactSequenceUpdateArtifactSequence] = ( + Field(alias="updateArtifactSequence") + ) + + +class UpdateArtifactSequenceUpdateArtifactSequence(GQLBase): + artifact_collection: UpdateArtifactSequenceUpdateArtifactSequenceArtifactCollection = Field( + alias="artifactCollection" + ) + + +class UpdateArtifactSequenceUpdateArtifactSequenceArtifactCollection(GQLBase): + typename__: Typename[ + Literal["ArtifactCollection", "ArtifactPortfolio", "ArtifactSequence"] + ] + id: GQLId + name: str + description: Optional[str] + + +UpdateArtifactSequence.model_rebuild() +UpdateArtifactSequenceUpdateArtifactSequence.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_graphql_fragments.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_graphql_fragments.py new file mode 100644 index 0000000000000000000000000000000000000000..936bd576315f7a9e1512993a5c4e0cca04f62dbe --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_graphql_fragments.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from textwrap import dedent + +from wandb_graphql.language.printer import print_ast + +from wandb.apis.public.utils import gql_compat +from wandb.sdk.internal.internal_api import Api as InternalApi + +OMITTABLE_ARTIFACT_FIELDS = frozenset( + { + "ttlDurationSeconds", + "ttlIsInherited", + "aliases", + "tags", + "historyStep", + } +) + + +def omit_artifact_fields(api: InternalApi) -> set[str]: + """Return names of Artifact fields to remove from GraphQL requests (for server compatibility).""" + allowed_fields = set(api.server_artifact_introspection()) + return set(OMITTABLE_ARTIFACT_FIELDS - allowed_fields) + + +def _gql_artifact_fragment(include_aliases: bool = True) -> str: + """Return a GraphQL query fragment with all parseable Artifact attributes.""" + omit_fields = omit_artifact_fields(api=InternalApi()) + + # Respect the `include_aliases` flag + if not include_aliases: + omit_fields.add("aliases") + + artifact_fragment_str = dedent( + """\ + fragment ArtifactFragment on Artifact { + id + artifactSequence { + project { + entityName + name + } + name + } + versionIndex + artifactType { + name + } + description + metadata + ttlDurationSeconds + ttlIsInherited + aliases { + artifactCollection { + project { + entityName + name + } + name + } + alias + } + tags { + name + } + historyStep + state + currentManifest { + file { + directUrl + } + } + commitHash + fileCount + createdAt + updatedAt + }""" + ) + compat_doc = gql_compat(artifact_fragment_str, omit_fields=omit_fields) + return print_ast(compat_doc) + + +def _gql_registry_fragment() -> str: + return """ + fragment RegistryFragment on Project { + id + allowAllArtifactTypesInRegistry + artifactTypes(includeAll: true) { + edges { + node { + name + } + } + } + name + description + createdAt + updatedAt + access + } + """ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_internal_artifact.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_internal_artifact.py new file mode 100644 index 0000000000000000000000000000000000000000..7f65932728257582aed9b0653187cd26fb5e3853 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_internal_artifact.py @@ -0,0 +1,43 @@ +import re +from typing import Any, Dict, Final, Optional + +from wandb.sdk.artifacts.artifact import Artifact + +PLACEHOLDER: Final[str] = "PLACEHOLDER" + + +def sanitize_artifact_name(name: str) -> str: + """Sanitize the string to satisfy constraints on artifact names.""" + from wandb.sdk.lib.hashutil import _md5 + + # If the name is already sanitized, don't change it. + if (sanitized := re.sub(r"[^a-zA-Z0-9_\-.]+", "", name)) == name: + return name + + # Append an alphanumeric suffix to maintain uniqueness of the name. + suffix = _md5(name.encode("utf-8")).hexdigest() + return f"{sanitized}-{suffix}" + + +class InternalArtifact(Artifact): + """InternalArtifact is used to create artifacts that are intended for internal use. + + This includes artifacts of type: `job`, `code`(with `source-` prefix in the collection name), + `run_table` (with `run-` prefix in the collection name), and artifacts that start with `wandb-`. + Users should not use this class directly. + """ + + def __init__( + self, + name: str, + type: str, + description: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + incremental: bool = False, + use_as: Optional[str] = None, + ) -> None: + sanitized_name = sanitize_artifact_name(name) + super().__init__( + sanitized_name, PLACEHOLDER, description, metadata, incremental, use_as + ) + self._type = type diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/_validators.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/_validators.py new file mode 100644 index 0000000000000000000000000000000000000000..49fd09ebcc157cf857a940910940b42c623a1493 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/_validators.py @@ -0,0 +1,299 @@ +"""Internal validation utilities that are specific to artifacts.""" + +from __future__ import annotations + +import json +import re +from dataclasses import astuple, dataclass, field, replace +from functools import wraps +from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, TypeVar, cast + +from pydantic.dataclasses import dataclass as pydantic_dataclass +from typing_extensions import Self + +from wandb._iterutils import always_list +from wandb._pydantic import gql_typename +from wandb.util import json_friendly_val + +from ._generated import ArtifactPortfolioTypeFields, ArtifactSequenceTypeFields +from .exceptions import ArtifactFinalizedError, ArtifactNotLoggedError + +if TYPE_CHECKING: + from typing import Collection, Final + + from wandb.sdk.artifacts.artifact import Artifact + + ArtifactT = TypeVar("ArtifactT", bound=Artifact) + T = TypeVar("T") + + +REGISTRY_PREFIX: Final[str] = "wandb-registry-" +MAX_ARTIFACT_METADATA_KEYS: Final[int] = 100 + +ARTIFACT_NAME_MAXLEN: Final[int] = 128 +ARTIFACT_NAME_INVALID_CHARS: Final[frozenset[str]] = frozenset({"/"}) + +LINKED_ARTIFACT_COLLECTION_TYPE: Final[str] = gql_typename(ArtifactPortfolioTypeFields) +SOURCE_ARTIFACT_COLLECTION_TYPE: Final[str] = gql_typename(ArtifactSequenceTypeFields) + + +@dataclass +class _LinkArtifactFields: + """Keep this list updated with fields where the linked artifact and the source artifact differ.""" + + entity_name: str + project_name: str + name: str + version: str + aliases: list[str] + + # These fields shouldn't be set as they should always be + # these values for a linked artifact + # These fields shouldn't be set by the user as they should always be these values for a linked artifact + _is_link: Literal[True] = field(init=False, default=True) + _linked_artifacts: list[Artifact] = field(init=False, default_factory=list) + + @property + def is_link(self) -> bool: + return self._is_link + + @property + def linked_artifacts(self) -> list[Artifact]: + return self._linked_artifacts + + +def validate_artifact_name(name: str) -> str: + """Validate the artifact name, returning it if successful. + + Raises: + ValueError: If the artifact name is invalid. + """ + if len(name) > ARTIFACT_NAME_MAXLEN: + short_name = f"{name[:ARTIFACT_NAME_MAXLEN]} ..." + raise ValueError( + f"Artifact name is longer than {ARTIFACT_NAME_MAXLEN} characters: {short_name!r}" + ) + + if ARTIFACT_NAME_INVALID_CHARS.intersection(name): + raise ValueError( + "Artifact names must not contain any of the following characters: " + f"{', '.join(sorted(ARTIFACT_NAME_INVALID_CHARS))}. Got: {name!r}" + ) + + return name + + +INVALID_URL_CHARACTERS = ("/", "\\", "#", "?", "%", ":", "\r", "\n") + + +def validate_project_name(name: str) -> None: + """Validates a project name according to W&B rules. + + Args: + name: The project name string. + + Raises: + ValueError: If the name is invalid (too long or contains invalid characters). + """ + max_len = 128 + + if len(name) == 0: + raise ValueError("Project name cannot be empty") + + registry_name = "" + if name.startswith(REGISTRY_PREFIX): + registry_name = name[len(REGISTRY_PREFIX) :] + if len(registry_name) == 0: + raise ValueError("Registry name cannot be empty") + + if len(name) > max_len: + if registry_name: + raise ValueError( + f"Invalid registry name {registry_name!r}, must be {max_len - len(REGISTRY_PREFIX)} characters or less" + ) + else: + raise ValueError( + f"Invalid project name {name!r}, must be {max_len} characters or less" + ) + + # Find the first occurrence of any invalid character + if invalid_chars := set(INVALID_URL_CHARACTERS).intersection(name): + error_name = registry_name or name + invalid_chars_repr = ", ".join(sorted(map(repr, invalid_chars))) + raise ValueError( + f"Invalid project/registry name {error_name!r}, cannot contain characters: {invalid_chars_repr!s}" + ) + + +def validate_aliases(aliases: Collection[str] | str) -> list[str]: + """Validate the artifact aliases and return them as a list. + + Raises: + ValueError: If any of the aliases contain invalid characters. + """ + aliases_list = always_list(aliases) + + invalid_chars = ("/", ":") + if any(char in alias for alias in aliases_list for char in invalid_chars): + raise ValueError( + f"Aliases must not contain any of the following characters: {', '.join(invalid_chars)}" + ) + return aliases_list + + +def validate_artifact_types_list(artifact_types: list[str]) -> list[str]: + """Return True if the artifact types list is valid, False otherwise.""" + artifact_types = always_list(artifact_types) + invalid_chars = ("/", ":") + if any( + char in type or len(type) > 128 + for type in artifact_types + for char in invalid_chars + ): + raise ValueError( + f"""Artifact types must not contain any of the following characters: {", ".join(invalid_chars)} + and must be less than equal to 128 characters""" + ) + return artifact_types + + +_VALID_TAG_PATTERN: re.Pattern[str] = re.compile(r"^[-\w]+( +[-\w]+)*$") + + +def validate_tags(tags: Collection[str] | str) -> list[str]: + """Validate the artifact tag names and return them as a deduped list. + + In the case of duplicates, only keep the first tag, and otherwise maintain the order of appearance. + + Raises: + ValueError: If any of the tags contain invalid characters. + """ + tags_list = always_list(tags) + + if any(not _VALID_TAG_PATTERN.match(tag) for tag in tags_list): + raise ValueError( + "Invalid tag(s). " + "Tags must only contain alphanumeric characters separated by hyphens, underscores, and/or spaces." + ) + return list(dict.fromkeys(tags_list)) + + +RESERVED_ARTIFACT_TYPE_PREFIX: Final[str] = "wandb-" +RESERVED_ARTIFACT_TYPES: Final[tuple[str, ...]] = ("job", "run_table", "code") +RESERVED_ARTIFACT_NAME_PREFIXES: Final[dict[str, str]] = { + "run_table": "run-", + "code": "source-", +} + + +def validate_artifact_type(typ: str, name: str) -> str: + """Validate the artifact type and return it as a string.""" + if typ in RESERVED_ARTIFACT_TYPES: + reserved_name_prefix = RESERVED_ARTIFACT_NAME_PREFIXES.get(typ) + if not reserved_name_prefix or name.startswith(reserved_name_prefix): + raise ValueError( + f"Artifact type {typ!r} is reserved for internal use. " + "Please use a different type." + ) + elif typ.startswith(RESERVED_ARTIFACT_TYPE_PREFIX): + raise ValueError( + f"Artifact type {typ!r} is reserved for internal use. " + "Please use a different type." + ) + return typ + + +def validate_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]: + """Validate the artifact metadata and return it as a dict.""" + if metadata is None: + return {} + if not isinstance(metadata, dict): + raise TypeError(f"metadata must be dict, not {type(metadata)}") + return cast(Dict[str, Any], json.loads(json.dumps(json_friendly_val(metadata)))) + + +def validate_ttl_duration_seconds(gql_ttl_duration_seconds: int | None) -> int | None: + """Validate the `ttlDurationSeconds` value (if any) from a GraphQL response.""" + # If gql_ttl_duration_seconds is not positive, its indicating that TTL is DISABLED(-2) + # gql_ttl_duration_seconds only returns None if the server is not compatible with setting Artifact TTLs + if gql_ttl_duration_seconds and gql_ttl_duration_seconds > 0: + return gql_ttl_duration_seconds + return None + + +# ---------------------------------------------------------------------------- +DecoratedF = TypeVar("DecoratedF", bound=Callable[..., Any]) +"""Type hint for a decorated function that'll preserve its signature (e.g. for arg autocompletion in IDEs).""" + + +def ensure_logged(method: DecoratedF) -> DecoratedF: + """Decorator to ensure that an Artifact method can only be called if the artifact has been logged. + + If the method is called on an artifact that's not logged, `ArtifactNotLoggedError` is raised. + """ + # For clarity, use the qualified (full) name of the method + method_fullname = method.__qualname__ + + @wraps(method) + def wrapper(self: ArtifactT, *args: Any, **kwargs: Any) -> Any: + if self.is_draft(): + raise ArtifactNotLoggedError(fullname=method_fullname, obj=self) + return method(self, *args, **kwargs) + + return cast(DecoratedF, wrapper) + + +def ensure_not_finalized(method: DecoratedF) -> DecoratedF: + """Decorator to ensure that an `Artifact` method can only be called if the artifact isn't finalized. + + If the method is called on an artifact that's not logged, `ArtifactFinalizedError` is raised. + """ + # For clarity, use the qualified (full) name of the method + method_fullname = method.__qualname__ + + @wraps(method) + def wrapper(self: ArtifactT, *args: Any, **kwargs: Any) -> Any: + if self._final: + raise ArtifactFinalizedError(fullname=method_fullname, obj=self) + return method(self, *args, **kwargs) + + return cast(DecoratedF, wrapper) + + +def is_artifact_registry_project(project: str) -> bool: + return project.startswith(REGISTRY_PREFIX) + + +@pydantic_dataclass +class ArtifactPath: + #: The collection name. + name: str + #: The project name, which can also be a registry name. + project: Optional[str] = None # noqa: UP045 + #: Prefix is often an org or entity name. + prefix: Optional[str] = None # noqa: UP045 + + @classmethod + def from_str(cls, path: str) -> Self: + """Instantiate by parsing an artifact path.""" + if len(parts := path.split("/")) <= 3: + return cls(*reversed(parts)) + raise ValueError( + f"Expected a valid path like `name`, `project/name`, or `prefix/project/name`. Got: {path!r}" + ) + + def to_str(self) -> str: + """Returns the slash-separated string representation of the path.""" + return "/".join(filter(bool, reversed(astuple(self)))) + + def with_defaults( + self, + prefix: str | None = None, + project: str | None = None, + ) -> Self: + """Returns this path with missing values set to the given defaults.""" + return replace( + self, + prefix=self.prefix or prefix, + project=self.project or project, + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact.py new file mode 100644 index 0000000000000000000000000000000000000000..2624facd9c0f83043da8fe627ce5b4a7ba98e313 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact.py @@ -0,0 +1,2808 @@ +"""Artifact class.""" + +from __future__ import annotations + +import atexit +import concurrent.futures +import contextlib +import json +import logging +import multiprocessing.dummy +import os +import re +import shutil +import stat +import tempfile +import time +from collections import deque +from copy import copy +from dataclasses import dataclass +from datetime import datetime, timedelta +from functools import partial +from itertools import filterfalse +from pathlib import PurePosixPath +from typing import IO, TYPE_CHECKING, Any, Iterator, Literal, Sequence, Type, final +from urllib.parse import quote, urljoin, urlparse + +import requests + +import wandb +from wandb import data_types, env +from wandb._iterutils import one +from wandb.apis.normalize import normalize_exceptions +from wandb.apis.public import ArtifactCollection, ArtifactFiles, Run +from wandb.apis.public.utils import gql_compat +from wandb.data_types import WBValue +from wandb.errors import CommError +from wandb.errors.errors import UnsupportedError +from wandb.errors.term import termerror, termlog, termwarn +from wandb.proto import wandb_internal_pb2 as pb +from wandb.proto.wandb_deprecated import Deprecated +from wandb.sdk import wandb_setup +from wandb.sdk.data_types._dtypes import Type as WBType +from wandb.sdk.data_types._dtypes import TypeRegistry +from wandb.sdk.internal.internal_api import Api as InternalApi +from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings +from wandb.sdk.lib import filesystem, retry, runid, telemetry +from wandb.sdk.lib.deprecate import deprecate +from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, md5_file_b64 +from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr +from wandb.sdk.lib.runid import generate_id +from wandb.sdk.mailbox import MailboxHandle +from wandb.util import ( + alias_is_version_index, + artifact_to_json, + fsync_open, + json_dumps_safer, + uri_from_path, + vendor_setup, +) + +from ._generated import ( + ADD_ALIASES_GQL, + DELETE_ALIASES_GQL, + FETCH_LINKED_ARTIFACTS_GQL, + LINK_ARTIFACT_GQL, + UPDATE_ARTIFACT_GQL, + ArtifactAliasInput, + ArtifactCollectionAliasInput, + FetchLinkedArtifacts, + LinkArtifact, + LinkArtifactInput, + TagInput, + UpdateArtifact, +) +from ._graphql_fragments import _gql_artifact_fragment, omit_artifact_fields +from ._validators import ( + LINKED_ARTIFACT_COLLECTION_TYPE, + ArtifactPath, + _LinkArtifactFields, + ensure_logged, + ensure_not_finalized, + is_artifact_registry_project, + validate_aliases, + validate_artifact_name, + validate_artifact_type, + validate_metadata, + validate_tags, + validate_ttl_duration_seconds, +) +from .artifact_download_logger import ArtifactDownloadLogger +from .artifact_instance_cache import artifact_instance_cache +from .artifact_manifest import ArtifactManifest +from .artifact_manifest_entry import ArtifactManifestEntry +from .artifact_manifests.artifact_manifest_v1 import ArtifactManifestV1 +from .artifact_state import ArtifactState +from .artifact_ttl import ArtifactTTL +from .exceptions import ( + ArtifactNotLoggedError, + TooFewItemsError, + TooManyItemsError, + WaitTimeoutError, +) +from .staging import get_staging_dir +from .storage_handlers.gcs_handler import _GCSIsADirectoryError +from .storage_layout import StorageLayout +from .storage_policies import WANDB_STORAGE_POLICY +from .storage_policy import StoragePolicy + +reset_path = vendor_setup() + +from wandb_gql import gql # noqa: E402 + +reset_path() + +if TYPE_CHECKING: + from wandb.apis.public import RetryingClient + +logger = logging.getLogger(__name__) + + +@final +@dataclass +class _DeferredArtifactManifest: + """A lightweight wrapper around the manifest URL, used to indicate deferred loading of the actual manifest.""" + + url: str + + +class Artifact: + """Flexible and lightweight building block for dataset and model versioning. + + Construct an empty W&B Artifact. Populate an artifacts contents with methods that + begin with `add`. Once the artifact has all the desired files, you can call + `wandb.log_artifact()` to log it. + + Args: + name: A human-readable name for the artifact. Use the name to identify + a specific artifact in the W&B App UI or programmatically. You can + interactively reference an artifact with the `use_artifact` Public API. + A name can contain letters, numbers, underscores, hyphens, and dots. + The name must be unique across a project. + type: The artifact's type. Use the type of an artifact to both organize + and differentiate artifacts. You can use any string that contains letters, + numbers, underscores, hyphens, and dots. Common types include `dataset` or `model`. + Note: Some types are reserved for internal use and cannot be set by users. + Such types include `job` and types that start with `wandb-`. + description: A description of the artifact. For Model or Dataset Artifacts, + add documentation for your standardized team model or dataset card. View + an artifact's description programmatically with the `Artifact.description` + attribute or programmatically with the W&B App UI. W&B renders the + description as markdown in the W&B App. + metadata: Additional information about an artifact. Specify metadata as a + dictionary of key-value pairs. You can specify no more than 100 total keys. + incremental: Use `Artifact.new_draft()` method instead to modify an + existing artifact. + use_as: Deprecated. + is_link: Boolean indication of if the artifact is a linked artifact(`True`) or source artifact(`False`). + + Returns: + An `Artifact` object. + """ + + _TMP_DIR = tempfile.TemporaryDirectory("wandb-artifacts") + atexit.register(_TMP_DIR.cleanup) + + def __init__( + self, + name: str, + type: str, + description: str | None = None, + metadata: dict[str, Any] | None = None, + incremental: bool = False, + use_as: str | None = None, + ) -> None: + if not re.match(r"^[a-zA-Z0-9_\-.]+$", name): + raise ValueError( + f"Artifact name may only contain alphanumeric characters, dashes, " + f"underscores, and dots. Invalid name: {name}" + ) + + from wandb.sdk.artifacts._internal_artifact import InternalArtifact + + if incremental and not isinstance(self, InternalArtifact): + termwarn("Using experimental arg `incremental`") + + # Internal. + self._client: RetryingClient | None = None + + storage_policy_cls = StoragePolicy.lookup_by_name(WANDB_STORAGE_POLICY) + layout = StorageLayout.V1 if env.get_use_v1_artifacts() else StorageLayout.V2 + policy_config = {"storageLayout": layout} + self._storage_policy = storage_policy_cls.from_config(config=policy_config) + + self._tmp_dir: tempfile.TemporaryDirectory | None = None + self._added_objs: dict[int, tuple[WBValue, ArtifactManifestEntry]] = {} + self._added_local_paths: dict[str, ArtifactManifestEntry] = {} + self._save_handle: MailboxHandle[pb.Result] | None = None + self._download_roots: set[str] = set() + # Set by new_draft(), otherwise the latest artifact will be used as the base. + self._base_id: str | None = None + # Properties. + self._id: str | None = None + self._client_id: str = runid.generate_id(128) + self._sequence_client_id: str = runid.generate_id(128) + self._entity: str | None = None + self._project: str | None = None + self._name: str = validate_artifact_name(name) # includes version after saving + self._version: str | None = None + self._source_entity: str | None = None + self._source_project: str | None = None + self._source_name: str = name # includes version after saving + self._source_version: str | None = None + self._source_artifact: Artifact | None = None + self._is_link: bool = False + self._type: str = validate_artifact_type(type, name) + self._description: str | None = description + self._metadata: dict[str, Any] = validate_metadata(metadata) + self._ttl_duration_seconds: int | None = None + self._ttl_is_inherited: bool = True + self._ttl_changed: bool = False + self._aliases: list[str] = [] + self._saved_aliases: list[str] = [] + self._tags: list[str] = [] + self._saved_tags: list[str] = [] + self._distributed_id: str | None = None + self._incremental: bool = incremental + if use_as is not None: + deprecate( + field_name=Deprecated.artifact__init_use_as, + warning_message=( + "`use_as` argument is deprecated and does not affect the behaviour of `wandb.Artifact()`" + ), + ) + self._use_as: str | None = None + self._state: ArtifactState = ArtifactState.PENDING + self._manifest: ArtifactManifest | _DeferredArtifactManifest | None = ( + ArtifactManifestV1(self._storage_policy) + ) + self._commit_hash: str | None = None + self._file_count: int | None = None + self._created_at: str | None = None + self._updated_at: str | None = None + self._final: bool = False + self._history_step: int | None = None + self._linked_artifacts: list[Artifact] = [] + + # Cache. + artifact_instance_cache[self._client_id] = self + + def __repr__(self) -> str: + return f"" + + @classmethod + def _from_id(cls, artifact_id: str, client: RetryingClient) -> Artifact | None: + if (artifact := artifact_instance_cache.get(artifact_id)) is not None: + return artifact + + query = gql( + """ + query ArtifactByID($id: ID!) { + artifact(id: $id) { + ...ArtifactFragment + } + } + """ + + _gql_artifact_fragment() + ) + response = client.execute( + query, + variable_values={"id": artifact_id}, + ) + attrs = response.get("artifact") + if attrs is None: + return None + + src_collection = attrs["artifactSequence"] + src_project = src_collection["project"] + + entity_name = src_project["entityName"] if src_project else "" + project_name = src_project["name"] if src_project else "" + + name = "{}:v{}".format(src_collection["name"], attrs["versionIndex"]) + return cls._from_attrs(entity_name, project_name, name, attrs, client) + + @classmethod + def _membership_from_name( + cls, + *, + entity: str, + project: str, + name: str, + client: RetryingClient, + ) -> Artifact: + if not InternalApi()._server_supports( + pb.ServerFeature.PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP + ): + raise UnsupportedError( + "querying for the artifact collection membership is not supported " + "by this version of wandb server. Consider updating to the latest version." + ) + + query = gql( + f""" + query ArtifactByName($entityName: String!, $projectName: String!, $name: String!) {{ + project(name: $projectName, entityName: $entityName) {{ + artifactCollectionMembership(name: $name) {{ + id + artifactCollection {{ + id + name + project {{ + id + entityName + name + }} + }} + artifact {{ + ...ArtifactFragment + }} + }} + }} + }} + {_gql_artifact_fragment()} + """ + ) + + query_variable_values: dict[str, Any] = { + "entityName": entity, + "projectName": project, + "name": name, + } + response = client.execute( + query, + variable_values=query_variable_values, + ) + if not (project_attrs := response.get("project")): + raise ValueError(f"project {project!r} not found under entity {entity!r}") + if not (acm_attrs := project_attrs.get("artifactCollectionMembership")): + entity_project = f"{entity}/{project}" + raise ValueError( + f"artifact membership {name!r} not found in {entity_project!r}" + ) + if not (ac_attrs := acm_attrs.get("artifactCollection")): + raise ValueError("artifact collection not found") + if not ( + (ac_name := ac_attrs.get("name")) + and (ac_project_attrs := ac_attrs.get("project")) + ): + raise ValueError("artifact collection project not found") + ac_project = ac_project_attrs.get("name") + ac_entity = ac_project_attrs.get("entityName") + if is_artifact_registry_project(ac_project) and project == "model-registry": + wandb.termwarn( + "This model registry has been migrated and will be discontinued. " + f"Your request was redirected to the corresponding artifact `{ac_name}` in the new registry. " + f"Please update your paths to point to the migrated registry directly, `{ac_project}/{ac_name}`." + ) + entity = ac_entity + project = ac_project + if not (attrs := acm_attrs.get("artifact")): + entity_project = f"{entity}/{project}" + raise ValueError(f"artifact {name!r} not found in {entity_project!r}") + + return cls._from_attrs(entity, project, name, attrs, client) + + @classmethod + def _from_name( + cls, + *, + entity: str, + project: str, + name: str, + client: RetryingClient, + enable_tracking: bool = False, + ) -> Artifact: + if InternalApi()._server_supports( + pb.ServerFeature.PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP + ): + return cls._membership_from_name( + entity=entity, + project=project, + name=name, + client=client, + ) + + query_variable_values: dict[str, Any] = { + "entityName": entity, + "projectName": project, + "name": name, + } + query_vars = ["$entityName: String!", "$projectName: String!", "$name: String!"] + query_args = ["name: $name"] + + server_supports_enabling_artifact_usage_tracking = ( + InternalApi().server_project_type_introspection() + ) + if server_supports_enabling_artifact_usage_tracking: + query_vars.append("$enableTracking: Boolean") + query_args.append("enableTracking: $enableTracking") + query_variable_values["enableTracking"] = enable_tracking + + vars_str = ", ".join(query_vars) + args_str = ", ".join(query_args) + + query = gql( + f""" + query ArtifactByName({vars_str}) {{ + project(name: $projectName, entityName: $entityName) {{ + artifact({args_str}) {{ + ...ArtifactFragment + }} + }} + }} + {_gql_artifact_fragment()} + """ + ) + response = client.execute( + query, + variable_values=query_variable_values, + ) + project_attrs = response.get("project") + if not project_attrs: + raise ValueError(f"project '{project}' not found under entity '{entity}'") + attrs = project_attrs.get("artifact") + if not attrs: + raise ValueError(f"artifact '{name}' not found in '{entity}/{project}'") + + return cls._from_attrs(entity, project, name, attrs, client) + + @classmethod + def _from_attrs( + cls, + entity: str, + project: str, + name: str, + attrs: dict[str, Any], + client: RetryingClient, + aliases: list[str] | None = None, + ) -> Artifact: + # Placeholder is required to skip validation. + artifact = cls("placeholder", type="placeholder") + artifact._client = client + artifact._entity = entity + artifact._project = project + artifact._name = name + artifact._assign_attrs(attrs, aliases) + + artifact.finalize() + + # Cache. + assert artifact.id is not None + artifact_instance_cache[artifact.id] = artifact + return artifact + + # TODO: Eventually factor out is_link. Have to currently use it since some forms of fetching the artifact + # doesn't make it clear if the artifact is a link or not and have to manually set it. + def _assign_attrs( + self, + attrs: dict[str, Any], + aliases: list[str] | None = None, + is_link: bool | None = None, + ) -> None: + """Update this Artifact's attributes using the server response.""" + self._id = attrs["id"] + + src_version = f"v{attrs['versionIndex']}" + src_collection = attrs["artifactSequence"] + src_project = src_collection["project"] + + self._source_entity = src_project["entityName"] if src_project else "" + self._source_project = src_project["name"] if src_project else "" + self._source_name = f"{src_collection['name']}:{src_version}" + self._source_version = src_version + + if self._entity is None: + self._entity = self._source_entity + if self._project is None: + self._project = self._source_project + + if self._name is None: + self._name = self._source_name + + # TODO: Refactor artifact query to fetch artifact via membership instead + # and get the collection type + if is_link is None: + self._is_link = ( + self._entity != self._source_entity + or self._project != self._source_project + or self._name != self._source_name + ) + else: + self._is_link = is_link + + self._type = attrs["artifactType"]["name"] + self._description = attrs["description"] + + entity = self._entity + project = self._project + collection, *_ = self._name.split(":") + + processed_aliases = [] + # The future of aliases is to move all alias fetches to the membership level + # so we don't have to do the collection fetches below + if aliases: + processed_aliases = aliases + else: + processed_aliases = [ + obj["alias"] + for obj in attrs["aliases"] + if obj["artifactCollection"] + and obj["artifactCollection"]["project"] + and obj["artifactCollection"]["project"]["entityName"] == entity + and obj["artifactCollection"]["project"]["name"] == project + and obj["artifactCollection"]["name"] == collection + ] + + version_aliases = list(filter(alias_is_version_index, processed_aliases)) + other_aliases = list(filterfalse(alias_is_version_index, processed_aliases)) + + try: + version = one( + version_aliases, too_short=TooFewItemsError, too_long=TooManyItemsError + ) + except TooFewItemsError: + version = src_version # default to the source version + except TooManyItemsError: + msg = f"Expected at most one version alias, got {len(version_aliases)}: {version_aliases!r}" + raise ValueError(msg) from None + + self._version = version + + if ":" not in self._name: + self._name = f"{self._name}:{version}" + + self._aliases = other_aliases + self._saved_aliases = copy(other_aliases) + + tags = [obj["name"] for obj in (attrs.get("tags") or [])] + self._tags = tags + self._saved_tags = copy(tags) + + metadata_str = attrs["metadata"] + self._metadata = validate_metadata( + json.loads(metadata_str) if metadata_str else {} + ) + + self._ttl_duration_seconds = validate_ttl_duration_seconds( + attrs.get("ttlDurationSeconds") + ) + self._ttl_is_inherited = ( + True if (attrs.get("ttlIsInherited") is None) else attrs["ttlIsInherited"] + ) + + self._state = ArtifactState(attrs["state"]) + + try: + manifest_url = attrs["currentManifest"]["file"]["directUrl"] + except (LookupError, TypeError): + self._manifest = None + else: + self._manifest = _DeferredArtifactManifest(manifest_url) + + self._commit_hash = attrs["commitHash"] + self._file_count = attrs["fileCount"] + self._created_at = attrs["createdAt"] + self._updated_at = attrs["updatedAt"] + self._history_step = attrs.get("historyStep", None) + + @ensure_logged + def new_draft(self) -> Artifact: + """Create a new draft artifact with the same content as this committed artifact. + + Modifying an existing artifact creates a new artifact version known + as an "incremental artifact". The artifact returned can be extended or + modified and logged as a new version. + + Returns: + An `Artifact` object. + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + """ + # Name, _entity and _project are set to the *source* name/entity/project: + # if this artifact is saved it must be saved to the source sequence. + artifact = Artifact(self.source_name.split(":")[0], self.type) + artifact._entity = self._source_entity + artifact._project = self._source_project + artifact._source_entity = self._source_entity + artifact._source_project = self._source_project + + # This artifact's parent is the one we are making a draft from. + artifact._base_id = self.id + + # We can reuse the client, and copy over all the attributes that aren't + # version-dependent and don't depend on having been logged. + artifact._client = self._client + artifact._description = self.description + artifact._metadata = self.metadata + artifact._manifest = ArtifactManifest.from_manifest_json( + self.manifest.to_manifest_json() + ) + return artifact + + # Properties (Python Class managed attributes). + + @property + def id(self) -> str | None: + """The artifact's ID.""" + if self.is_draft(): + return None + assert self._id is not None + return self._id + + @property + @ensure_logged + def entity(self) -> str: + """The name of the entity that the artifact collection belongs to. + + If the artifact is a link, the entity will be the entity of the linked artifact. + """ + assert self._entity is not None + return self._entity + + @property + @ensure_logged + def project(self) -> str: + """The name of the project that the artifact collection belongs to. + + If the artifact is a link, the project will be the project of the linked artifact. + """ + assert self._project is not None + return self._project + + @property + def name(self) -> str: + """The artifact name and version of the artifact. + + A string with the format `{collection}:{alias}`. If fetched before an artifact is logged/saved, the name won't contain the alias. + If the artifact is a link, the name will be the name of the linked artifact. + """ + return self._name + + @property + def qualified_name(self) -> str: + """The entity/project/name of the artifact. + + If the artifact is a link, the qualified name will be the qualified name of the linked artifact path. + """ + return f"{self.entity}/{self.project}/{self.name}" + + @property + @ensure_logged + def version(self) -> str: + """The artifact's version. + + A string with the format `v{number}`. + If the artifact is a link artifact, the version will be from the linked collection. + """ + assert self._version is not None + return self._version + + @property + @ensure_logged + def collection(self) -> ArtifactCollection: + """The collection this artifact was retrieved from. + + A collection is an ordered group of artifact versions. + If this artifact was retrieved from a portfolio / linked collection, that + collection will be returned rather than the collection + that an artifact version originated from. The collection + that an artifact originates from is known as the source sequence. + """ + base_name = self.name.split(":")[0] + return ArtifactCollection( + self._client, self.entity, self.project, base_name, self.type + ) + + @property + @ensure_logged + def source_entity(self) -> str: + """The name of the entity of the source artifact.""" + assert self._source_entity is not None + return self._source_entity + + @property + @ensure_logged + def source_project(self) -> str: + """The name of the project of the source artifact.""" + assert self._source_project is not None + return self._source_project + + @property + def source_name(self) -> str: + """The artifact name and version of the source artifact. + + A string with the format `{source_collection}:{alias}`. Before the artifact is saved, + contains only the name since the version is not yet known. + """ + return self._source_name + + @property + def source_qualified_name(self) -> str: + """The source_entity/source_project/source_name of the source artifact.""" + return f"{self.source_entity}/{self.source_project}/{self.source_name}" + + @property + @ensure_logged + def source_version(self) -> str: + """The source artifact's version. + + A string with the format `v{number}`. + """ + assert self._source_version is not None + return self._source_version + + @property + @ensure_logged + def source_collection(self) -> ArtifactCollection: + """The artifact's source collection. + + The source collection is the collection that the artifact was logged from. + """ + base_name = self.source_name.split(":")[0] + return ArtifactCollection( + self._client, self.source_entity, self.source_project, base_name, self.type + ) + + @property + def is_link(self) -> bool: + """Boolean flag indicating if the artifact is a link artifact. + + True: The artifact is a link artifact to a source artifact. + False: The artifact is a source artifact. + """ + return self._is_link + + @property + @ensure_logged + def linked_artifacts(self) -> list[Artifact]: + """Returns a list of all the linked artifacts of a source artifact. + + If the artifact is a link artifact (`artifact.is_link == True`), it will return an empty list. + Limited to 500 results.""" + if not self.is_link: + self._linked_artifacts = self._fetch_linked_artifacts() + return self._linked_artifacts + + @property + @ensure_logged + def source_artifact(self) -> Artifact: + """Returns the source artifact. The source artifact is the original logged artifact. + + If the artifact itself is a source artifact (`artifact.is_link == False`), it will return itself.""" + if not self.is_link: + return self + if self._source_artifact is None: + if self._client is None: + raise ValueError("Client is not initialized") + + try: + artifact = self._from_name( + entity=self.source_entity, + project=self.source_project, + name=self.source_name, + client=self._client, + ) + self._source_artifact = artifact + except Exception as e: + raise ValueError( + f"Unable to fetch source artifact for linked artifact {self.name}" + ) from e + return self._source_artifact + + @property + def type(self) -> str: + """The artifact's type. Common types include `dataset` or `model`.""" + return self._type + + @property + @ensure_logged + def url(self) -> str: + """ + Constructs the URL of the artifact. + + Returns: + str: The URL of the artifact. + """ + try: + base_url = self._client.app_url # type: ignore[union-attr] + except AttributeError: + return "" + + if not self.is_link: + return self._construct_standard_url(base_url) + if is_artifact_registry_project(self.project): + return self._construct_registry_url(base_url) + if self._type == "model" or self.project == "model-registry": + return self._construct_model_registry_url(base_url) + return self._construct_standard_url(base_url) + + def _construct_standard_url(self, base_url: str) -> str: + if not all( + [ + base_url, + self.entity, + self.project, + self._type, + self.collection.name, + self._version, + ] + ): + return "" + return urljoin( + base_url, + f"{self.entity}/{self.project}/artifacts/{quote(self._type)}/{quote(self.collection.name)}/{self._version}", + ) + + def _construct_registry_url(self, base_url: str) -> str: + if not all( + [ + base_url, + self.entity, + self.project, + self.collection.name, + self._version, + ] + ): + return "" + + try: + org, *_ = InternalApi()._fetch_orgs_and_org_entities_from_entity( + self.entity + ) + except ValueError: + return "" + + selection_path = quote( + f"{self.entity}/{self.project}/{self.collection.name}", safe="" + ) + return urljoin( + base_url, + f"orgs/{org.display_name}/registry/{self._type}?selectionPath={selection_path}&view=membership&version={self._version}", + ) + + def _construct_model_registry_url(self, base_url: str) -> str: + if not all( + [ + base_url, + self.entity, + self.project, + self.collection.name, + self._version, + ] + ): + return "" + selection_path = quote( + f"{self.entity}/{self.project}/{self.collection.name}", safe="" + ) + return urljoin( + base_url, + f"{self.entity}/registry/model?selectionPath={selection_path}&view=membership&version={self._version}", + ) + + @property + def description(self) -> str | None: + """A description of the artifact.""" + return self._description + + @description.setter + def description(self, description: str | None) -> None: + """Set the description of the artifact. + + For model or dataset Artifacts, add documentation for your + standardized team model or dataset card. In the W&B UI the + description is rendered as markdown. + + Editing the description will apply the changes to the source artifact and all linked artifacts associated with it. + + Args: + description: Free text that offers a description of the artifact. + """ + if self.is_link: + wandb.termwarn( + "Editing the description of this linked artifact will edit the description for the source artifact and it's linked artifacts as well." + ) + self._description = description + + @property + def metadata(self) -> dict: + """User-defined artifact metadata. + + Structured data associated with the artifact. + """ + return self._metadata + + @metadata.setter + def metadata(self, metadata: dict) -> None: + """User-defined artifact metadata. + + Metadata set this way will eventually be queryable and plottable in the UI; e.g. + the class distribution of a dataset. + + Note: There is currently a limit of 100 total keys. + Editing the metadata will apply the changes to the source artifact and all linked artifacts associated with it. + + Args: + metadata: Structured data associated with the artifact. + """ + if self.is_link: + wandb.termwarn( + "Editing the metadata of this linked artifact will edit the metadata for the source artifact and it's linked artifacts as well." + ) + self._metadata = validate_metadata(metadata) + + @property + def ttl(self) -> timedelta | None: + """The time-to-live (TTL) policy of an artifact. + + Artifacts are deleted shortly after a TTL policy's duration passes. + If set to `None`, the artifact deactivates TTL policies and will be not + scheduled for deletion, even if there is a team default TTL. + An artifact inherits a TTL policy from + the team default if the team administrator defines a default + TTL and there is no custom policy set on an artifact. + + Raises: + ArtifactNotLoggedError: Unable to fetch inherited TTL if the artifact has not been logged or saved + """ + if self._ttl_is_inherited and (self.is_draft() or self._ttl_changed): + raise ArtifactNotLoggedError(f"{type(self).__name__}.ttl", self) + if self._ttl_duration_seconds is None: + return None + return timedelta(seconds=self._ttl_duration_seconds) + + @ttl.setter + def ttl(self, ttl: timedelta | ArtifactTTL | None) -> None: + """The time-to-live (TTL) policy of an artifact. + + Artifacts are deleted shortly after a TTL policy's duration passes. + If set to `None`, the artifact has no TTL policy set and it is not + scheduled for deletion. An artifact inherits a TTL policy from + the team default if the team administrator defines a default + TTL and there is no custom policy set on an artifact. + + Args: + ttl: The duration as a positive Python `datetime.timedelta` Type + that represents how long the artifact will remain active from its creation. + + """ + if self.type == "wandb-history": + raise ValueError("Cannot set artifact TTL for type wandb-history") + + if self.is_link: + raise ValueError( + "Cannot set TTL for link artifact. " + "Unlink the artifact first then set the TTL for the source artifact" + ) + + self._ttl_changed = True + if isinstance(ttl, ArtifactTTL): + if ttl == ArtifactTTL.INHERIT: + self._ttl_is_inherited = True + else: + raise ValueError(f"Unhandled ArtifactTTL enum {ttl}") + else: + self._ttl_is_inherited = False + if ttl is None: + self._ttl_duration_seconds = None + else: + if ttl.total_seconds() <= 0: + raise ValueError( + f"Artifact TTL Duration has to be positive. ttl: {ttl.total_seconds()}" + ) + self._ttl_duration_seconds = int(ttl.total_seconds()) + + @property + @ensure_logged + def aliases(self) -> list[str]: + """List of one or more semantically-friendly references or identifying "nicknames" assigned to an artifact version. + + Aliases are mutable references that you can programmatically reference. + Change an artifact's alias with the W&B App UI or programmatically. + See [Create new artifact versions](https://docs.wandb.ai/guides/artifacts/create-a-new-artifact-version) + for more information. + """ + return self._aliases + + @aliases.setter + @ensure_logged + def aliases(self, aliases: list[str]) -> None: + """Set the aliases associated with this artifact.""" + self._aliases = validate_aliases(aliases) + + @property + @ensure_logged + def tags(self) -> list[str]: + """List of one or more tags assigned to this artifact version.""" + return self._tags + + @tags.setter + @ensure_logged + def tags(self, tags: list[str]) -> None: + """Set the tags associated with this artifact. + + Editing tags will apply the changes to the source artifact and all linked artifacts associated with it. + """ + if self.is_link: + wandb.termwarn( + "Editing tags will apply the changes to the source artifact and all linked artifacts associated with it." + ) + self._tags = validate_tags(tags) + + @property + def distributed_id(self) -> str | None: + return self._distributed_id + + @distributed_id.setter + def distributed_id(self, distributed_id: str | None) -> None: + self._distributed_id = distributed_id + + @property + def incremental(self) -> bool: + return self._incremental + + @property + def use_as(self) -> str | None: + """Deprecated.""" + deprecate( + field_name=Deprecated.artifact__use_as, + warning_message=("The use_as property of Artifact is deprecated."), + ) + return self._use_as + + @property + def state(self) -> str: + """The status of the artifact. One of: "PENDING", "COMMITTED", or "DELETED".""" + return self._state.value + + @property + def manifest(self) -> ArtifactManifest: + """The artifact's manifest. + + The manifest lists all of its contents, and can't be changed once the artifact + has been logged. + """ + if isinstance(self._manifest, _DeferredArtifactManifest): + # A deferred manifest URL flags a deferred download request, + # so fetch the manifest to override the placeholder object + self._manifest = self._load_manifest(self._manifest.url) + return self._manifest + + if self._manifest is None: + query = gql( + """ + query ArtifactManifest( + $entityName: String!, + $projectName: String!, + $name: String! + ) { + project(entityName: $entityName, name: $projectName) { + artifact(name: $name) { + currentManifest { + file { + directUrl + } + } + } + } + } + """ + ) + assert self._client is not None + response = self._client.execute( + query, + variable_values={ + "entityName": self._entity, + "projectName": self._project, + "name": self._name, + }, + ) + attrs = response["project"]["artifact"] + manifest_url = attrs["currentManifest"]["file"]["directUrl"] + self._manifest = self._load_manifest(manifest_url) + + return self._manifest + + @property + def digest(self) -> str: + """The logical digest of the artifact. + + The digest is the checksum of the artifact's contents. If an artifact has the + same digest as the current `latest` version, then `log_artifact` is a no-op. + """ + return self.manifest.digest() + + @property + def size(self) -> int: + """The total size of the artifact in bytes. + + Includes any references tracked by this artifact. + """ + total_size: int = 0 + for entry in self.manifest.entries.values(): + if entry.size is not None: + total_size += entry.size + return total_size + + @property + @ensure_logged + def commit_hash(self) -> str: + """The hash returned when this artifact was committed.""" + assert self._commit_hash is not None + return self._commit_hash + + @property + @ensure_logged + def file_count(self) -> int: + """The number of files (including references).""" + assert self._file_count is not None + return self._file_count + + @property + @ensure_logged + def created_at(self) -> str: + """Timestamp when the artifact was created.""" + assert self._created_at is not None + return self._created_at + + @property + @ensure_logged + def updated_at(self) -> str: + """The time when the artifact was last updated.""" + assert self._created_at is not None + return self._updated_at or self._created_at + + @property + @ensure_logged + def history_step(self) -> int | None: + """The nearest step at which history metrics were logged for the source run of the artifact. + + Examples: + ```python + run = artifact.logged_by() + if run and (artifact.history_step is not None): + history = run.sample_history( + min_step=artifact.history_step, + max_step=artifact.history_step + 1, + keys=["my_metric"], + ) + ``` + """ + if self._history_step is None: + return None + return max(0, self._history_step - 1) + + # State management. + + def finalize(self) -> None: + """Finalize the artifact version. + + You cannot modify an artifact version once it is finalized because the artifact + is logged as a specific artifact version. Create a new artifact version + to log more data to an artifact. An artifact is automatically finalized + when you log the artifact with `log_artifact`. + """ + self._final = True + + def is_draft(self) -> bool: + """Check if artifact is not saved. + + Returns: Boolean. `False` if artifact is saved. `True` if artifact is not saved. + """ + return self._state == ArtifactState.PENDING + + def _is_draft_save_started(self) -> bool: + return self._save_handle is not None + + def save( + self, + project: str | None = None, + settings: wandb.Settings | None = None, + ) -> None: + """Persist any changes made to the artifact. + + If currently in a run, that run will log this artifact. If not currently in a + run, a run of type "auto" is created to track this artifact. + + Args: + project: A project to use for the artifact in the case that a run is not + already in context. + settings: A settings object to use when initializing an automatic run. Most + commonly used in testing harness. + """ + if self._state != ArtifactState.PENDING: + return self._update() + + if self._incremental: + with telemetry.context() as tel: + tel.feature.artifact_incremental = True + + if run := wandb_setup.singleton().most_recent_active_run: + # TODO: Deprecate and encourage explicit log_artifact(). + run.log_artifact(self) + else: + if settings is None: + settings = wandb.Settings(silent="true") + with wandb.init( # type: ignore + entity=self._source_entity, + project=project or self._source_project, + job_type="auto", + settings=settings, + ) as run: + # redoing this here because in this branch we know we didn't + # have the run at the beginning of the method + if self._incremental: + with telemetry.context(run=run) as tel: + tel.feature.artifact_incremental = True + run.log_artifact(self) + + def _set_save_handle( + self, + save_handle: MailboxHandle[pb.Result], + client: RetryingClient, + ) -> None: + self._save_handle = save_handle + self._client = client + + def wait(self, timeout: int | None = None) -> Artifact: + """If needed, wait for this artifact to finish logging. + + Args: + timeout: The time, in seconds, to wait. + + Returns: + An `Artifact` object. + """ + if self.is_draft(): + if self._save_handle is None: + raise ArtifactNotLoggedError(type(self).wait.__qualname__, self) + + try: + result = self._save_handle.wait_or(timeout=timeout) + except TimeoutError as e: + raise WaitTimeoutError( + "Artifact upload wait timed out, failed to fetch Artifact response" + ) from e + + response = result.response.log_artifact_response + if response.error_message: + raise ValueError(response.error_message) + self._populate_after_save(response.artifact_id) + return self + + def _populate_after_save(self, artifact_id: str) -> None: + query_template = """ + query ArtifactByIDShort($id: ID!) { + artifact(id: $id) { + ...ArtifactFragment + } + } + """ + _gql_artifact_fragment() + + query = gql(query_template) + + assert self._client is not None + response = self._client.execute( + query, + variable_values={"id": artifact_id}, + ) + + try: + attrs = response["artifact"] + except LookupError: + raise ValueError(f"Unable to fetch artifact with id: {artifact_id!r}") + else: + # _populate_after_save is only called on source artifacts, not linked artifacts + # We have to manually set is_link because we aren't fetching the collection the artifact. + # That requires greater refactoring for commitArtifact to return the artifact collection type. + self._assign_attrs(attrs, is_link=False) + + @normalize_exceptions + def _update(self) -> None: + """Persists artifact changes to the wandb backend.""" + if self._client is None: + raise RuntimeError("Client not initialized for artifact mutations") + + entity = self.entity + project = self.project + collection = self.name.split(":")[0] + + aliases = None + introspect_query = gql( + """ + query ProbeServerAddAliasesInput { + AddAliasesInputInfoType: __type(name: "AddAliasesInput") { + name + inputFields { + name + } + } + } + """ + ) + + data = self._client.execute(introspect_query) + if data.get("AddAliasesInputInfoType"): # wandb backend version >= 0.13.0 + alias_props = { + "entity_name": entity, + "project_name": project, + "artifact_collection_name": collection, + } + if aliases_to_add := (set(self.aliases) - set(self._saved_aliases)): + add_mutation = gql(ADD_ALIASES_GQL) + add_alias_inputs = [ + ArtifactCollectionAliasInput(**alias_props, alias=alias) + for alias in aliases_to_add + ] + try: + self._client.execute( + add_mutation, + variable_values={ + "artifactID": self.id, + "aliases": [a.model_dump() for a in add_alias_inputs], + }, + ) + except CommError as e: + raise CommError( + "You do not have permission to add" + f" {'at least one of the following aliases' if len(aliases_to_add) > 1 else 'the following alias'}" + f" to this artifact: {aliases_to_add}" + ) from e + + if aliases_to_delete := (set(self._saved_aliases) - set(self.aliases)): + delete_mutation = gql(DELETE_ALIASES_GQL) + delete_alias_inputs = [ + ArtifactCollectionAliasInput(**alias_props, alias=alias) + for alias in aliases_to_delete + ] + try: + self._client.execute( + delete_mutation, + variable_values={ + "artifactID": self.id, + "aliases": [a.model_dump() for a in delete_alias_inputs], + }, + ) + except CommError as e: + raise CommError( + f"You do not have permission to delete" + f" {'at least one of the following aliases' if len(aliases_to_delete) > 1 else 'the following alias'}" + f" from this artifact: {aliases_to_delete}" + ) from e + + self._saved_aliases = copy(self.aliases) + + else: # wandb backend version < 0.13.0 + aliases = [ + ArtifactAliasInput( + artifact_collection_name=collection, alias=alias + ).model_dump() + for alias in self.aliases + ] + + omit_fields = omit_artifact_fields(api=InternalApi()) + omit_variables = set() + + if {"ttlIsInherited", "ttlDurationSeconds"} & omit_fields: + if self._ttl_changed: + termwarn( + "Server not compatible with setting Artifact TTLs, please upgrade the server to use Artifact TTL" + ) + + omit_variables |= {"ttlDurationSeconds"} + + tags_to_add = validate_tags(set(self.tags) - set(self._saved_tags)) + tags_to_del = validate_tags(set(self._saved_tags) - set(self.tags)) + + if {"tags"} & omit_fields: + if tags_to_add or tags_to_del: + termwarn( + "Server not compatible with Artifact tags. " + "To use Artifact tags, please upgrade the server to v0.85 or higher." + ) + + omit_variables |= {"tagsToAdd", "tagsToDelete"} + + mutation = gql_compat( + UPDATE_ARTIFACT_GQL, omit_variables=omit_variables, omit_fields=omit_fields + ) + + gql_vars = { + "artifactID": self.id, + "description": self.description, + "metadata": json_dumps_safer(self.metadata), + "ttlDurationSeconds": self._ttl_duration_seconds_to_gql(), + "aliases": aliases, + "tagsToAdd": [TagInput(tag_name=t).model_dump() for t in tags_to_add], + "tagsToDelete": [TagInput(tag_name=t).model_dump() for t in tags_to_del], + } + + data = self._client.execute(mutation, variable_values=gql_vars) + + result = UpdateArtifact.model_validate(data).update_artifact + if not (result and (artifact := result.artifact)): + raise ValueError("Unable to parse updateArtifact response") + self._assign_attrs(artifact.model_dump()) + + self._ttl_changed = False # Reset after updating artifact + + # Adding, removing, getting entries. + + def __getitem__(self, name: str) -> WBValue | None: + """Get the WBValue object located at the artifact relative `name`. + + Args: + name: The artifact relative name to get. + + Returns: + W&B object that can be logged with `wandb.log()` and visualized in the W&B UI. + + Raises: + ArtifactNotLoggedError: If the artifact isn't logged or the run is offline. + """ + return self.get(name) + + def __setitem__(self, name: str, item: WBValue) -> ArtifactManifestEntry: + """Add `item` to the artifact at path `name`. + + Args: + name: The path within the artifact to add the object. + item: The object to add. + + Returns: + The added manifest entry + + Raises: + ArtifactFinalizedError: You cannot make changes to the current artifact + version because it is finalized. Log a new artifact version instead. + """ + return self.add(item, name) + + @contextlib.contextmanager + @ensure_not_finalized + def new_file( + self, name: str, mode: str = "x", encoding: str | None = None + ) -> Iterator[IO]: + """Open a new temporary file and add it to the artifact. + + Args: + name: The name of the new file to add to the artifact. + mode: The file access mode to use to open the new file. + encoding: The encoding used to open the new file. + + Returns: + A new file object that can be written to. Upon closing, the file will be + automatically added to the artifact. + + Raises: + ArtifactFinalizedError: You cannot make changes to the current artifact + version because it is finalized. Log a new artifact version instead. + """ + overwrite: bool = "x" not in mode + + if self._tmp_dir is None: + self._tmp_dir = tempfile.TemporaryDirectory() + path = os.path.join(self._tmp_dir.name, name.lstrip("/")) + + filesystem.mkdir_exists_ok(os.path.dirname(path)) + try: + with fsync_open(path, mode, encoding) as f: + yield f + except FileExistsError: + raise ValueError(f"File with name {name!r} already exists at {path!r}") + except UnicodeEncodeError as e: + termerror( + f"Failed to open the provided file ({type(e).__name__}: {e}). Please " + f"provide the proper encoding." + ) + raise + + self.add_file( + path, name=name, policy="immutable", skip_cache=True, overwrite=overwrite + ) + + @ensure_not_finalized + def add_file( + self, + local_path: str, + name: str | None = None, + is_tmp: bool | None = False, + skip_cache: bool | None = False, + policy: Literal["mutable", "immutable"] | None = "mutable", + overwrite: bool = False, + ) -> ArtifactManifestEntry: + """Add a local file to the artifact. + + Args: + local_path: The path to the file being added. + name: The path within the artifact to use for the file being added. Defaults + to the basename of the file. + is_tmp: If true, then the file is renamed deterministically to avoid + collisions. + skip_cache: If `True`, W&B will not copy files to the cache after uploading. + policy: By default, set to "mutable". If set to "mutable", create a temporary copy of the + file to prevent corruption during upload. If set to "immutable", disable + protection and rely on the user not to delete or change the file. + overwrite: If `True`, overwrite the file if it already exists. + + Returns: + The added manifest entry. + + Raises: + ArtifactFinalizedError: You cannot make changes to the current artifact + version because it is finalized. Log a new artifact version instead. + ValueError: Policy must be "mutable" or "immutable" + """ + if not os.path.isfile(local_path): + raise ValueError(f"Path is not a file: {local_path!r}") + + name = LogicalPath(name or os.path.basename(local_path)) + digest = md5_file_b64(local_path) + + if is_tmp: + file_path, file_name = os.path.split(name) + file_name_parts = file_name.split(".") + file_name_parts[0] = b64_to_hex_id(digest)[:20] + name = os.path.join(file_path, ".".join(file_name_parts)) + + return self._add_local_file( + name, + local_path, + digest=digest, + skip_cache=skip_cache, + policy=policy, + overwrite=overwrite, + ) + + @ensure_not_finalized + def add_dir( + self, + local_path: str, + name: str | None = None, + skip_cache: bool | None = False, + policy: Literal["mutable", "immutable"] | None = "mutable", + merge: bool = False, + ) -> None: + """Add a local directory to the artifact. + + Args: + local_path: The path of the local directory. + name: The subdirectory name within an artifact. The name you specify appears + in the W&B App UI nested by artifact's `type`. + Defaults to the root of the artifact. + skip_cache: If set to `True`, W&B will not copy/move files to the cache while uploading + policy: "mutable" | "immutable". By default, "mutable" + "mutable": Create a temporary copy of the file to prevent corruption during upload. + "immutable": Disable protection, rely on the user not to delete or change the file. + merge: If `False` (default), throws ValueError if a file was already added in a previous add_dir call + and its content has changed. If `True`, overwrites existing files with changed content. + Always adds new files and never removes files. To replace an entire directory, pass a name when adding the directory + using `add_dir(local_path, name=my_prefix)` and call `remove(my_prefix)` to remove the directory, then add it again. + + Raises: + ArtifactFinalizedError: You cannot make changes to the current artifact + version because it is finalized. Log a new artifact version instead. + ValueError: Policy must be "mutable" or "immutable" + """ + if not os.path.isdir(local_path): + raise ValueError(f"Path is not a directory: {local_path}") + + termlog( + "Adding directory to artifact ({})... ".format( + os.path.join(".", os.path.normpath(local_path)) + ), + newline=False, + ) + start_time = time.time() + + paths = [] + for dirpath, _, filenames in os.walk(local_path, followlinks=True): + for fname in filenames: + physical_path = os.path.join(dirpath, fname) + logical_path = os.path.relpath(physical_path, start=local_path) + if name is not None: + logical_path = os.path.join(name, logical_path) + paths.append((logical_path, physical_path)) + + def add_manifest_file(log_phy_path: tuple[str, str]) -> None: + logical_path, physical_path = log_phy_path + self._add_local_file( + name=logical_path, + path=physical_path, + skip_cache=skip_cache, + policy=policy, + overwrite=merge, + ) + + num_threads = 8 + pool = multiprocessing.dummy.Pool(num_threads) + pool.map(add_manifest_file, paths) + pool.close() + pool.join() + + termlog("Done. %.1fs" % (time.time() - start_time), prefix=False) + + @ensure_not_finalized + def add_reference( + self, + uri: ArtifactManifestEntry | str, + name: StrPath | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + """Add a reference denoted by a URI to the artifact. + + Unlike files or directories that you add to an artifact, references are not + uploaded to W&B. For more information, + see [Track external files](https://docs.wandb.ai/guides/artifacts/track-external-files). + + By default, the following schemes are supported: + + - http(s): The size and digest of the file will be inferred by the + `Content-Length` and the `ETag` response headers returned by the server. + - s3: The checksum and size are pulled from the object metadata. If bucket + versioning is enabled, then the version ID is also tracked. + - gs: The checksum and size are pulled from the object metadata. If bucket + versioning is enabled, then the version ID is also tracked. + - https, domain matching `*.blob.core.windows.net` (Azure): The checksum and size + are be pulled from the blob metadata. If storage account versioning is + enabled, then the version ID is also tracked. + - file: The checksum and size are pulled from the file system. This scheme + is useful if you have an NFS share or other externally mounted volume + containing files you wish to track but not necessarily upload. + + For any other scheme, the digest is just a hash of the URI and the size is left + blank. + + Args: + uri: The URI path of the reference to add. The URI path can be an object + returned from `Artifact.get_entry` to store a reference to another + artifact's entry. + name: The path within the artifact to place the contents of this reference. + checksum: Whether or not to checksum the resource(s) located at the + reference URI. Checksumming is strongly recommended as it enables + automatic integrity validation. Disabling checksumming will speed up + artifact creation but reference directories will not iterated through so the + objects in the directory will not be saved to the artifact. We recommend + setting `checksum=False` when adding reference objects, in which case + a new version will only be created if the reference URI changes. + max_objects: The maximum number of objects to consider when adding a + reference that points to directory or bucket store prefix. By default, + the maximum number of objects allowed for Amazon S3, + GCS, Azure, and local files is 10,000,000. Other URI schemas do not have a maximum. + + Returns: + The added manifest entries. + + Raises: + ArtifactFinalizedError: You cannot make changes to the current artifact + version because it is finalized. Log a new artifact version instead. + """ + if name is not None: + name = LogicalPath(name) + + # This is a bit of a hack, we want to check if the uri is a of the type + # ArtifactManifestEntry. If so, then recover the reference URL. + if isinstance(uri, ArtifactManifestEntry): + uri_str = uri.ref_url() + elif isinstance(uri, str): + uri_str = uri + url = urlparse(str(uri_str)) + if not url.scheme: + raise ValueError( + "References must be URIs. To reference a local file, use file://" + ) + + manifest_entries = self._storage_policy.store_reference( + self, + URIStr(uri_str), + name=name, + checksum=checksum, + max_objects=max_objects, + ) + for entry in manifest_entries: + self.manifest.add_entry(entry) + + return manifest_entries + + @ensure_not_finalized + def add( + self, obj: WBValue, name: StrPath, overwrite: bool = False + ) -> ArtifactManifestEntry: + """Add wandb.WBValue `obj` to the artifact. + + Args: + obj: The object to add. Currently support one of Bokeh, JoinedTable, + PartitionedTable, Table, Classes, ImageMask, BoundingBoxes2D, Audio, + Image, Video, Html, Object3D + name: The path within the artifact to add the object. + overwrite: If True, overwrite existing objects with the same file path (if applicable). + + Returns: + The added manifest entry + + Raises: + ArtifactFinalizedError: You cannot make changes to the current artifact + version because it is finalized. Log a new artifact version instead. + """ + name = LogicalPath(name) + + # This is a "hack" to automatically rename tables added to + # the wandb /media/tables directory to their sha-based name. + # TODO: figure out a more appropriate convention. + is_tmp_name = name.startswith("media/tables") + + # Validate that the object is one of the correct wandb.Media types + # TODO: move this to checking subclass of wandb.Media once all are + # generally supported + allowed_types = ( + data_types.Bokeh, + data_types.JoinedTable, + data_types.PartitionedTable, + data_types.Table, + data_types.Classes, + data_types.ImageMask, + data_types.BoundingBoxes2D, + data_types.Audio, + data_types.Image, + data_types.Video, + data_types.Html, + data_types.Object3D, + data_types.Molecule, + data_types._SavedModel, + ) + if not isinstance(obj, allowed_types): + raise TypeError( + f"Found object of type {obj.__class__}, expected one of:" + f" {allowed_types}" + ) + + obj_id = id(obj) + if obj_id in self._added_objs: + return self._added_objs[obj_id][1] + + # If the object is coming from another artifact, save it as a reference + ref_path = obj._get_artifact_entry_ref_url() + if ref_path is not None: + return self.add_reference(ref_path, type(obj).with_suffix(name))[0] + + val = obj.to_json(self) + name = obj.with_suffix(name) + entry = self.manifest.get_entry_by_path(name) + if (not overwrite) and (entry is not None): + return entry + + if is_tmp_name: + file_path = os.path.join(self._TMP_DIR.name, str(id(self)), name) + folder_path, _ = os.path.split(file_path) + os.makedirs(folder_path, exist_ok=True) + with open(file_path, "w", encoding="utf-8") as tmp_f: + json.dump(val, tmp_f, sort_keys=True) + else: + filemode = "w" if overwrite else "x" + with self.new_file(name, mode=filemode, encoding="utf-8") as f: + json.dump(val, f, sort_keys=True) + file_path = f.name + + # Note, we add the file from our temp directory. + # It will be added again later on finalize, but succeed since + # the checksum should match + entry = self.add_file(file_path, name, is_tmp_name) + # We store a reference to the obj so that its id doesn't get reused. + self._added_objs[obj_id] = (obj, entry) + if obj._artifact_target is None: + obj._set_artifact_target(self, entry.path) + + if is_tmp_name: + with contextlib.suppress(FileNotFoundError): + os.remove(file_path) + + return entry + + def _add_local_file( + self, + name: StrPath, + path: StrPath, + digest: B64MD5 | None = None, + skip_cache: bool | None = False, + policy: Literal["mutable", "immutable"] | None = "mutable", + overwrite: bool = False, + ) -> ArtifactManifestEntry: + policy = policy or "mutable" + if policy not in ["mutable", "immutable"]: + raise ValueError( + f"Invalid policy {policy!r}. Policy may only be `mutable` or `immutable`." + ) + upload_path = path + if policy == "mutable": + with tempfile.NamedTemporaryFile(dir=get_staging_dir(), delete=False) as f: + staging_path = f.name + shutil.copyfile(path, staging_path) + # Set as read-only to prevent changes to the file during upload process + os.chmod(staging_path, stat.S_IRUSR) + upload_path = staging_path + + entry = ArtifactManifestEntry( + path=name, + digest=digest or md5_file_b64(upload_path), + size=os.path.getsize(upload_path), + local_path=upload_path, + skip_cache=skip_cache, + ) + self.manifest.add_entry(entry, overwrite=overwrite) + self._added_local_paths[os.fspath(path)] = entry + return entry + + @ensure_not_finalized + def remove(self, item: StrPath | ArtifactManifestEntry) -> None: + """Remove an item from the artifact. + + Args: + item: The item to remove. Can be a specific manifest entry or the name of an + artifact-relative path. If the item matches a directory all items in + that directory will be removed. + + Raises: + ArtifactFinalizedError: You cannot make changes to the current artifact + version because it is finalized. Log a new artifact version instead. + FileNotFoundError: If the item isn't found in the artifact. + """ + if isinstance(item, ArtifactManifestEntry): + self.manifest.remove_entry(item) + return + + path = str(PurePosixPath(item)) + entry = self.manifest.get_entry_by_path(path) + if entry: + self.manifest.remove_entry(entry) + return + + entries = self.manifest.get_entries_in_directory(path) + if not entries: + raise FileNotFoundError(f"No such file or directory: {path}") + for entry in entries: + self.manifest.remove_entry(entry) + + def get_path(self, name: StrPath) -> ArtifactManifestEntry: + """Deprecated. Use `get_entry(name)`.""" + deprecate( + field_name=Deprecated.artifact__get_path, + warning_message="Artifact.get_path(name) is deprecated, use Artifact.get_entry(name) instead.", + ) + return self.get_entry(name) + + @ensure_logged + def get_entry(self, name: StrPath) -> ArtifactManifestEntry: + """Get the entry with the given name. + + Args: + name: The artifact relative name to get + + Returns: + A `W&B` object. + + Raises: + ArtifactNotLoggedError: if the artifact isn't logged or the run is offline. + KeyError: if the artifact doesn't contain an entry with the given name. + """ + name = LogicalPath(name) + entry = self.manifest.entries.get(name) or self._get_obj_entry(name)[0] + if entry is None: + raise KeyError(f"Path not contained in artifact: {name}") + entry._parent_artifact = self + return entry + + @ensure_logged + def get(self, name: str) -> WBValue | None: + """Get the WBValue object located at the artifact relative `name`. + + Args: + name: The artifact relative name to retrieve. + + Returns: + W&B object that can be logged with `wandb.log()` and visualized in the W&B UI. + + Raises: + ArtifactNotLoggedError: if the artifact isn't logged or the run is offline + """ + entry, wb_class = self._get_obj_entry(name) + if entry is None or wb_class is None: + return None + + # If the entry is a reference from another artifact, then get it directly from + # that artifact. + referenced_id = entry._referenced_artifact_id() + if referenced_id: + assert self._client is not None + artifact = self._from_id(referenced_id, client=self._client) + assert artifact is not None + return artifact.get(uri_from_path(entry.ref)) + + # Special case for wandb.Table. This is intended to be a short term + # optimization. Since tables are likely to download many other assets in + # artifact(s), we eagerly download the artifact using the parallelized + # `artifact.download`. In the future, we should refactor the deserialization + # pattern such that this special case is not needed. + if wb_class == wandb.Table: + self.download() + + # Get the ArtifactManifestEntry + item = self.get_entry(entry.path) + item_path = item.download() + + # Load the object from the JSON blob + result = None + json_obj = {} + with open(item_path) as file: + json_obj = json.load(file) + result = wb_class.from_json(json_obj, self) + result._set_artifact_source(self, name) + return result + + def get_added_local_path_name(self, local_path: str) -> str | None: + """Get the artifact relative name of a file added by a local filesystem path. + + Args: + local_path: The local path to resolve into an artifact relative name. + + Returns: + The artifact relative name. + """ + entry = self._added_local_paths.get(local_path, None) + if entry is None: + return None + return entry.path + + def _get_obj_entry( + self, name: str + ) -> tuple[ArtifactManifestEntry, Type[WBValue]] | tuple[None, None]: # noqa: UP006 # `type` shadows `Artifact.type` + """Return an object entry by name, handling any type suffixes. + + When objects are added with `.add(obj, name)`, the name is typically changed to + include the suffix of the object type when serializing to JSON. So we need to be + able to resolve a name, without tasking the user with appending .THING.json. + This method returns an entry if it exists by a suffixed name. + + Args: + name: name used when adding + """ + for wb_class in WBValue.type_mapping().values(): + wandb_file_name = wb_class.with_suffix(name) + entry = self.manifest.entries.get(wandb_file_name) + if entry is not None: + return entry, wb_class + return None, None + + # Downloading. + + @ensure_logged + def download( + self, + root: StrPath | None = None, + allow_missing_references: bool = False, + skip_cache: bool | None = None, + path_prefix: StrPath | None = None, + multipart: bool | None = None, + ) -> FilePathStr: + """Download the contents of the artifact to the specified root directory. + + Existing files located within `root` are not modified. Explicitly delete `root` + before you call `download` if you want the contents of `root` to exactly match + the artifact. + + Args: + root: The directory W&B stores the artifact's files. + allow_missing_references: If set to `True`, any invalid reference paths + will be ignored while downloading referenced files. + skip_cache: If set to `True`, the artifact cache will be skipped when + downloading and W&B will download each file into the default root or + specified download directory. + path_prefix: If specified, only files with a path that starts with the given + prefix will be downloaded. Uses unix format (forward slashes). + multipart: If set to `None` (default), the artifact will be downloaded + in parallel using multipart download if individual file size is greater than + 2GB. If set to `True` or `False`, the artifact will be downloaded in + parallel or serially regardless of the file size. + + Returns: + The path to the downloaded contents. + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + """ + root = FilePathStr(str(root or self._default_root())) + self._add_download_root(root) + + # TODO: download artifacts using core when implemented + # if is_require_core(): + # return self._download_using_core( + # root=root, + # allow_missing_references=allow_missing_references, + # skip_cache=bool(skip_cache), + # path_prefix=path_prefix, + # ) + return self._download( + root=root, + allow_missing_references=allow_missing_references, + skip_cache=skip_cache, + path_prefix=path_prefix, + multipart=multipart, + ) + + def _download_using_core( + self, + root: str, + allow_missing_references: bool = False, + skip_cache: bool = False, + path_prefix: StrPath | None = None, + ) -> FilePathStr: + import pathlib + + from wandb.sdk.backend.backend import Backend + + # TODO: Create a special stream instead of relying on an existing run. + if wandb.run is None: + wl = wandb_setup.singleton() + + stream_id = generate_id() + + settings = wl.settings.to_proto() + # TODO: remove this + tmp_dir = pathlib.Path(tempfile.mkdtemp()) + + settings.sync_dir.value = str(tmp_dir) + settings.sync_file.value = str(tmp_dir / f"{stream_id}.wandb") + settings.files_dir.value = str(tmp_dir / "files") + settings.run_id.value = stream_id + + service = wl.ensure_service() + service.inform_init(settings=settings, run_id=stream_id) + + backend = Backend(settings=wl.settings, service=service) + backend.ensure_launched() + + assert backend.interface + backend.interface._stream_id = stream_id # type: ignore + else: + assert wandb.run._backend + backend = wandb.run._backend + + assert backend.interface + handle = backend.interface.deliver_download_artifact( + self.id, # type: ignore + root, + allow_missing_references, + skip_cache, + path_prefix, # type: ignore + ) + # TODO: Start the download process in the user process too, to handle reference downloads + self._download( + root=root, + allow_missing_references=allow_missing_references, + skip_cache=skip_cache, + path_prefix=path_prefix, + ) + result = handle.wait_or(timeout=None) + + response = result.response.download_artifact_response + if response.error_message: + raise ValueError(f"Error downloading artifact: {response.error_message}") + + return FilePathStr(root) + + def _download( + self, + root: str, + allow_missing_references: bool = False, + skip_cache: bool | None = None, + path_prefix: StrPath | None = None, + multipart: bool | None = None, + ) -> FilePathStr: + nfiles = len(self.manifest.entries) + size = sum(e.size or 0 for e in self.manifest.entries.values()) + log = False + if nfiles > 5000 or size > 50 * 1024 * 1024: + log = True + termlog( + f"Downloading large artifact {self.name}, {size / (1024 * 1024):.2f}MB. {nfiles} files... ", + ) + start_time = datetime.now() + download_logger = ArtifactDownloadLogger(nfiles=nfiles) + + def _download_entry( + entry: ArtifactManifestEntry, + executor: concurrent.futures.Executor, + api_key: str | None, + cookies: dict | None, + headers: dict | None, + ) -> None: + _thread_local_api_settings.api_key = api_key + _thread_local_api_settings.cookies = cookies + _thread_local_api_settings.headers = headers + + try: + entry.download( + root, + skip_cache=skip_cache, + executor=executor, + multipart=multipart, + ) + except FileNotFoundError as e: + if allow_missing_references: + wandb.termwarn(str(e)) + return + raise + except _GCSIsADirectoryError as e: + logger.debug(str(e)) + return + download_logger.notify_downloaded() + + with concurrent.futures.ThreadPoolExecutor(64) as executor: + download_entry = partial( + _download_entry, + executor=executor, + api_key=_thread_local_api_settings.api_key, + cookies=_thread_local_api_settings.cookies, + headers=_thread_local_api_settings.headers, + ) + active_futures = set() + has_next_page = True + cursor = None + while has_next_page: + fetch_url_batch_size = env.get_artifact_fetch_file_url_batch_size() + attrs = self._fetch_file_urls(cursor, fetch_url_batch_size) + has_next_page = attrs["pageInfo"]["hasNextPage"] + cursor = attrs["pageInfo"]["endCursor"] + for edge in attrs["edges"]: + entry = self.get_entry(edge["node"]["name"]) + # TODO: uncomment once artifact downloads are supported in core + # if require_core and entry.ref is None: + # # Handled by core + # continue + entry._download_url = edge["node"]["directUrl"] + if (not path_prefix) or entry.path.startswith(str(path_prefix)): + active_futures.add(executor.submit(download_entry, entry)) + # Wait for download threads to catch up. + max_backlog = fetch_url_batch_size + if len(active_futures) > max_backlog: + for future in concurrent.futures.as_completed(active_futures): + future.result() # check for errors + active_futures.remove(future) + if len(active_futures) <= max_backlog: + break + # Check for errors. + for future in concurrent.futures.as_completed(active_futures): + future.result() + + if log: + now = datetime.now() + delta = abs((now - start_time).total_seconds()) + hours = int(delta // 3600) + minutes = int((delta - hours * 3600) // 60) + seconds = delta - hours * 3600 - minutes * 60 + speed = size / 1024 / 1024 / delta + termlog( + f"Done. {hours}:{minutes}:{seconds:.1f} ({speed:.1f}MB/s)", + prefix=False, + ) + return FilePathStr(root) + + @retry.retriable( + retry_timedelta=timedelta(minutes=3), + retryable_exceptions=(requests.RequestException), + ) + def _fetch_file_urls(self, cursor: str | None, per_page: int | None = 5000) -> Any: + if InternalApi()._server_supports( + pb.ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILES + ): + query = gql( + """ + query ArtifactCollectionMembershipFileURLs($entityName: String!, $projectName: String!, \ + $artifactName: String!, $artifactVersionIndex: String!, $cursor: String, $perPage: Int) { + project(name: $projectName, entityName: $entityName) { + artifactCollection(name: $artifactName) { + artifactMembership(aliasName: $artifactVersionIndex) { + files(after: $cursor, first: $perPage) { + pageInfo { + hasNextPage + endCursor + } + edges { + node { + name + directUrl + } + } + } + } + } + } + } + """ + ) + assert self._client is not None + response = self._client.execute( + query, + variable_values={ + "entityName": self.entity, + "projectName": self.project, + "artifactName": self.name.split(":")[0], + "artifactVersionIndex": self.version, + "cursor": cursor, + "perPage": per_page, + }, + timeout=60, + ) + return response["project"]["artifactCollection"]["artifactMembership"][ + "files" + ] + else: + query = gql( + """ + query ArtifactFileURLs($id: ID!, $cursor: String, $perPage: Int) { + artifact(id: $id) { + files(after: $cursor, first: $perPage) { + pageInfo { + hasNextPage + endCursor + } + edges { + node { + name + directUrl + } + } + } + } + } + """ + ) + assert self._client is not None + response = self._client.execute( + query, + variable_values={"id": self.id, "cursor": cursor, "perPage": per_page}, + timeout=60, + ) + return response["artifact"]["files"] + + @ensure_logged + def checkout(self, root: str | None = None) -> str: + """Replace the specified root directory with the contents of the artifact. + + WARNING: This will delete all files in `root` that are not included in the + artifact. + + Args: + root: The directory to replace with this artifact's files. + + Returns: + The path of the checked out contents. + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + """ + root = root or self._default_root(include_version=False) + + for dirpath, _, files in os.walk(root): + for file in files: + full_path = os.path.join(dirpath, file) + artifact_path = os.path.relpath(full_path, start=root) + try: + self.get_entry(artifact_path) + except KeyError: + # File is not part of the artifact, remove it. + os.remove(full_path) + + return self.download(root=root) + + @ensure_logged + def verify(self, root: str | None = None) -> None: + """Verify that the contents of an artifact match the manifest. + + All files in the directory are checksummed and the checksums are then + cross-referenced against the artifact's manifest. References are not verified. + + Args: + root: The directory to verify. If None artifact will be downloaded to + './artifacts/self.name/'. + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + ValueError: If the verification fails. + """ + root = root or self._default_root() + + for dirpath, _, files in os.walk(root): + for file in files: + full_path = os.path.join(dirpath, file) + artifact_path = os.path.relpath(full_path, start=root) + try: + self.get_entry(artifact_path) + except KeyError: + raise ValueError( + f"Found file {full_path} which is not a member of artifact {self.name}" + ) + + ref_count = 0 + for entry in self.manifest.entries.values(): + if entry.ref is None: + if md5_file_b64(os.path.join(root, entry.path)) != entry.digest: + raise ValueError(f"Digest mismatch for file: {entry.path}") + else: + ref_count += 1 + if ref_count > 0: + termwarn(f"skipped verification of {ref_count} refs") + + @ensure_logged + def file(self, root: str | None = None) -> StrPath: + """Download a single file artifact to the directory you specify with `root`. + + Args: + root: The root directory to store the file. Defaults to + './artifacts/self.name/'. + + Returns: + The full path of the downloaded file. + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + ValueError: If the artifact contains more than one file. + """ + if root is None: + root = os.path.join(".", "artifacts", self.name) + + if len(self.manifest.entries) > 1: + raise ValueError( + "This artifact contains more than one file, call `.download()` to get " + 'all files or call .get_entry("filename").download()' + ) + + return self.get_entry(list(self.manifest.entries)[0]).download(root) + + @ensure_logged + def files( + self, names: list[str] | None = None, per_page: int = 50 + ) -> ArtifactFiles: + """Iterate over all files stored in this artifact. + + Args: + names: The filename paths relative to the root of the artifact you wish to + list. + per_page: The number of files to return per request. + + Returns: + An iterator containing `File` objects. + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + """ + return ArtifactFiles(self._client, self, names, per_page) + + def _default_root(self, include_version: bool = True) -> FilePathStr: + name = self.source_name if include_version else self.source_name.split(":")[0] + root = os.path.join(env.get_artifact_dir(), name) + # In case we're on a system where the artifact dir has a name corresponding to + # an unexpected filesystem, we'll check for alternate roots. If one exists we'll + # use that, otherwise we'll fall back to the system-preferred path. + path = filesystem.check_exists(root) or filesystem.system_preferred_path(root) + return FilePathStr(str(path)) + + def _add_download_root(self, dir_path: str) -> None: + self._download_roots.add(os.path.abspath(dir_path)) + + def _local_path_to_name(self, file_path: str) -> str | None: + """Convert a local file path to a path entry in the artifact.""" + abs_file_path = os.path.abspath(file_path) + abs_file_parts = abs_file_path.split(os.sep) + for i in range(len(abs_file_parts) + 1): + if os.path.join(os.sep, *abs_file_parts[:i]) in self._download_roots: + return os.path.join(*abs_file_parts[i:]) + return None + + # Others. + + @ensure_logged + def delete(self, delete_aliases: bool = False) -> None: + """Delete an artifact and its files. + + If called on a linked artifact (i.e. a member of a portfolio collection): only the link is deleted, and the + source artifact is unaffected. + + Use `artifact.unlink()` instead of `artifact.delete()` to remove a link between a source artifact and a linked artifact. + + Args: + delete_aliases: If set to `True`, deletes all aliases associated with the artifact. + Otherwise, this raises an exception if the artifact has existing + aliases. + This parameter is ignored if the artifact is linked (i.e. a member of a portfolio collection). + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + """ + if self.is_link: + wandb.termwarn( + "Deleting a link artifact will only unlink the artifact from the source artifact and not delete the source artifact and the data of the source artifact." + ) + self._unlink() + else: + self._delete(delete_aliases) + + @normalize_exceptions + def _delete(self, delete_aliases: bool = False) -> None: + mutation = gql( + """ + mutation DeleteArtifact($artifactID: ID!, $deleteAliases: Boolean) { + deleteArtifact(input: { + artifactID: $artifactID + deleteAliases: $deleteAliases + }) { + artifact { + id + } + } + } + """ + ) + assert self._client is not None + self._client.execute( + mutation, + variable_values={ + "artifactID": self.id, + "deleteAliases": delete_aliases, + }, + ) + + @normalize_exceptions + def link( + self, target_path: str, aliases: list[str] | None = None + ) -> Artifact | None: + """Link this artifact to a portfolio (a promoted collection of artifacts). + + Args: + target_path: The path to the portfolio inside a project. + The target path must adhere to one of the following + schemas `{portfolio}`, `{project}/{portfolio}` or + `{entity}/{project}/{portfolio}`. + To link the artifact to the Model Registry, rather than to a generic + portfolio inside a project, set `target_path` to the following + schema `{"model-registry"}/{Registered Model Name}` or + `{entity}/{"model-registry"}/{Registered Model Name}`. + aliases: A list of strings that uniquely identifies the artifact inside the + specified portfolio. + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + + Returns: + The linked artifact if linking was successful, otherwise None. + """ + from wandb import Api + + if self.is_link: + wandb.termwarn( + "Linking to a link artifact will result in directly linking to the source artifact of that link artifact." + ) + + if self._client is None: + raise ValueError("Client not initialized for artifact mutations") + + # Save the artifact first if necessary + if self.is_draft(): + if not self._is_draft_save_started(): + self.save(project=self.source_project) + + # Wait until the artifact is committed before trying to link it. + self.wait() + + api = Api(overrides={"entity": self.source_entity}) + + target = ArtifactPath.from_str(target_path).with_defaults( + project=api.settings.get("project") or "uncategorized", + ) + + # Parse the entity (first part of the path) appropriately, + # depending on whether we're linking to a registry + if target.project and is_artifact_registry_project(target.project): + # In a Registry linking, the entity is used to fetch the organization of the artifact + # therefore the source artifact's entity is passed to the backend + organization = target.prefix or api.settings.get("organization") or "" + + target.prefix = InternalApi()._resolve_org_entity_name( + self.source_entity, organization + ) + else: + target = target.with_defaults(prefix=self.source_entity) + + # Prepare the validated GQL input, send it + alias_inputs = [ + ArtifactAliasInput(artifact_collection_name=target.name, alias=a) + for a in (aliases or []) + ] + gql_input = LinkArtifactInput( + artifact_id=self.id, + artifact_portfolio_name=target.name, + entity_name=target.prefix, + project_name=target.project, + aliases=alias_inputs, + ) + gql_vars = {"input": gql_input.model_dump(exclude_none=True)} + gql_op = gql(LINK_ARTIFACT_GQL) + data = self._client.execute(gql_op, variable_values=gql_vars) + + result = LinkArtifact.model_validate(data).link_artifact + if not (result and (version_idx := result.version_index) is not None): + raise ValueError("Unable to parse linked artifact version from response") + + # Fetch the linked artifact to return it + linked_path = f"{target.to_str()}:v{version_idx}" + + try: + return api._artifact(linked_path) + except Exception as e: + wandb.termerror(f"Error fetching link artifact after linking: {e}") + return None + + @ensure_logged + def unlink(self) -> None: + """Unlink this artifact if it is currently a member of a portfolio (a promoted collection of artifacts). + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + ValueError: If the artifact is not linked, i.e. it is not a member of a portfolio collection. + """ + # Fail early if this isn't a linked artifact to begin with + if not self.is_link: + raise ValueError( + f"Artifact {self.qualified_name!r} is not a linked artifact and cannot be unlinked. " + f"To delete it, use {self.delete.__qualname__!r} instead." + ) + + self._unlink() + + @normalize_exceptions + def _unlink(self) -> None: + mutation = gql( + """ + mutation UnlinkArtifact($artifactID: ID!, $artifactPortfolioID: ID!) { + unlinkArtifact( + input: { artifactID: $artifactID, artifactPortfolioID: $artifactPortfolioID } + ) { + artifactID + success + clientMutationId + } + } + """ + ) + assert self._client is not None + try: + self._client.execute( + mutation, + variable_values={ + "artifactID": self.id, + "artifactPortfolioID": self.collection.id, + }, + ) + except CommError as e: + raise CommError( + f"You do not have permission to unlink the artifact {self.qualified_name}" + ) from e + + @ensure_logged + def used_by(self) -> list[Run]: + """Get a list of the runs that have used this artifact and its linked artifacts. + + Returns: + A list of `Run` objects. + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + """ + query = gql( + """ + query ArtifactUsedBy( + $id: ID!, + ) { + artifact(id: $id) { + usedBy { + edges { + node { + name + project { + name + entityName + } + } + } + } + } + } + """ + ) + assert self._client is not None + response = self._client.execute( + query, + variable_values={"id": self.id}, + ) + return [ + Run( + self._client, + edge["node"]["project"]["entityName"], + edge["node"]["project"]["name"], + edge["node"]["name"], + ) + for edge in response.get("artifact", {}).get("usedBy", {}).get("edges", []) + ] + + @ensure_logged + def logged_by(self) -> Run | None: + """Get the W&B run that originally logged the artifact. + + Returns: + The name of the W&B run that originally logged the artifact. + + Raises: + ArtifactNotLoggedError: If the artifact is not logged. + """ + query = gql( + """ + query ArtifactCreatedBy( + $id: ID! + ) { + artifact(id: $id) { + createdBy { + ... on Run { + name + project { + name + entityName + } + } + } + } + } + """ + ) + assert self._client is not None + response = self._client.execute( + query, + variable_values={"id": self.id}, + ) + creator = response.get("artifact", {}).get("createdBy", {}) + if creator.get("name") is None: + return None + return Run( + self._client, + creator["project"]["entityName"], + creator["project"]["name"], + creator["name"], + ) + + @ensure_logged + def json_encode(self) -> dict[str, Any]: + """Returns the artifact encoded to the JSON format. + + Returns: + A `dict` with `string` keys representing attributes of the artifact. + """ + return artifact_to_json(self) + + @staticmethod + def _expected_type( + entity_name: str, project_name: str, name: str, client: RetryingClient + ) -> str | None: + """Returns the expected type for a given artifact name and project.""" + query = gql( + """ + query ArtifactType( + $entityName: String, + $projectName: String, + $name: String! + ) { + project(name: $projectName, entityName: $entityName) { + artifact(name: $name) { + artifactType { + name + } + } + } + } + """ + ) + if ":" not in name: + name += ":latest" + response = client.execute( + query, + variable_values={ + "entityName": entity_name, + "projectName": project_name, + "name": name, + }, + ) + return ( + ((response.get("project") or {}).get("artifact") or {}).get("artifactType") + or {} + ).get("name") + + def _load_manifest(self, url: str) -> ArtifactManifest: + with requests.get(url) as response: + response.raise_for_status() + return ArtifactManifest.from_manifest_json(response.json()) + + def _ttl_duration_seconds_to_gql(self) -> int | None: + # Set artifact ttl value to ttl_duration_seconds if the user set a value + # otherwise use ttl_status to indicate the backend INHERIT(-1) or DISABLED(-2) when the TTL is None + # When ttl_change = None its a no op since nothing changed + INHERIT = -1 # noqa: N806 + DISABLED = -2 # noqa: N806 + + if not self._ttl_changed: + return None + if self._ttl_is_inherited: + return INHERIT + return self._ttl_duration_seconds or DISABLED + + def _fetch_linked_artifacts(self) -> list[Artifact]: + """Fetches all linked artifacts from the server.""" + if self.id is None: + raise ValueError( + "Unable to find any artifact memberships for artifact without an ID" + ) + if self._client is None: + raise ValueError("Client is not initialized") + response = self._client.execute( + gql_compat(FETCH_LINKED_ARTIFACTS_GQL), + variable_values={"artifactID": self.id}, + ) + result = FetchLinkedArtifacts.model_validate(response) + + if not ( + (artifact := result.artifact) + and (memberships := artifact.artifact_memberships) + and (membership_edges := memberships.edges) + ): + raise ValueError("Unable to find any artifact memberships for artifact") + + linked_artifacts: deque[Artifact] = deque() + linked_nodes = ( + node + for edge in membership_edges + if ( + (node := edge.node) + and (col := node.artifact_collection) + and (col.typename__ == LINKED_ARTIFACT_COLLECTION_TYPE) + ) + ) + for node in linked_nodes: + # Trick for O(1) membership check that maintains order + alias_names = dict.fromkeys(a.alias for a in node.aliases) + version = f"v{node.version_index}" + aliases = ( + [*alias_names, version] + if version not in alias_names + else [*alias_names] + ) + + if not ( + node + and (col := node.artifact_collection) + and (proj := col.project) + and (proj.entity_name and proj.name) + ): + raise ValueError("Unable to fetch fields for linked artifact") + + link_fields = _LinkArtifactFields( + entity_name=proj.entity_name, + project_name=proj.name, + name=f"{col.name}:{version}", + version=version, + aliases=aliases, + ) + link = self._create_linked_artifact_using_source_artifact(link_fields) + linked_artifacts.append(link) + return list(linked_artifacts) + + def _create_linked_artifact_using_source_artifact( + self, + link_fields: _LinkArtifactFields, + ) -> Artifact: + """Copies the source artifact to a linked artifact.""" + linked_artifact = copy(self) + linked_artifact._version = link_fields.version + linked_artifact._aliases = link_fields.aliases + linked_artifact._saved_aliases = copy(link_fields.aliases) + linked_artifact._name = link_fields.name + linked_artifact._entity = link_fields.entity_name + linked_artifact._project = link_fields.project_name + linked_artifact._is_link = link_fields.is_link + linked_artifact._linked_artifacts = link_fields.linked_artifacts + return linked_artifact + + +class _ArtifactVersionType(WBType): + name = "artifactVersion" + types = [Artifact] + + +TypeRegistry.add(_ArtifactVersionType) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_download_logger.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_download_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..f4370d68bc74b9252827b6b18aeac87a4b06ef55 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_download_logger.py @@ -0,0 +1,45 @@ +"""Artifact download logger.""" + +from __future__ import annotations + +import multiprocessing.dummy +import time +from typing import Callable + +from wandb.errors.term import termlog + + +class ArtifactDownloadLogger: + def __init__( + self, + nfiles: int, + clock_for_testing: Callable[[], float] = time.monotonic, + termlog_for_testing: Callable[..., None] = termlog, + ) -> None: + self._nfiles = nfiles + self._clock = clock_for_testing + self._termlog = termlog_for_testing + + self._n_files_downloaded = 0 + self._spinner_index = 0 + self._last_log_time = self._clock() + self._lock = multiprocessing.dummy.Lock() + + def notify_downloaded(self) -> None: + with self._lock: + self._n_files_downloaded += 1 + if self._n_files_downloaded == self._nfiles: + self._termlog( + f" {self._nfiles} of {self._nfiles} files downloaded. ", + # ^ trailing spaces to wipe out ellipsis from previous logs + newline=True, + ) + self._last_log_time = self._clock() + elif self._clock() - self._last_log_time > 0.1: + self._spinner_index += 1 + spinner = r"-\|/"[self._spinner_index % 4] + self._termlog( + f"{spinner} {self._n_files_downloaded} of {self._nfiles} files downloaded...\r", + newline=False, + ) + self._last_log_time = self._clock() diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_file_cache.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_file_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..aa4dd0d8e97466a24abf34b52ff6ad625833cf66 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_file_cache.py @@ -0,0 +1,247 @@ +"""Artifact cache.""" + +from __future__ import annotations + +import contextlib +import errno +import hashlib +import os +import shutil +import subprocess +import sys +from pathlib import Path +from tempfile import NamedTemporaryFile +from typing import IO, ContextManager, Iterator, Protocol + +import wandb +from wandb import env, util +from wandb.sdk.lib.filesystem import files_in +from wandb.sdk.lib.hashutil import B64MD5, ETag, b64_to_hex_id +from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr + + +class Opener(Protocol): + def __call__(self, mode: str = ...) -> ContextManager[IO]: + pass + + +def _get_sys_umask_threadsafe() -> int: + # Workaround to get the current system umask, since + # - `os.umask()` isn't thread-safe + # - we don't want to inadvertently change the umask of the current process + # See: https://stackoverflow.com/questions/53227072/reading-umask-thread-safe + umask_cmd = (sys.executable, "-c", "import os; print(os.umask(22))") + return int(subprocess.check_output(umask_cmd)) + + +class ArtifactFileCache: + def __init__(self, cache_dir: StrPath) -> None: + self._cache_dir = Path(cache_dir) + self._obj_dir = self._cache_dir / "obj" + self._temp_dir = self._cache_dir / "tmp" + self._ensure_write_permissions() + + # NamedTemporaryFile sets the file mode to 600 [1], we reset to the default. + # [1] https://stackoverflow.com/questions/10541760/can-i-set-the-umask-for-tempfile-namedtemporaryfile-in-python + self._sys_umask = _get_sys_umask_threadsafe() + + self._override_cache_path: StrPath | None = None + + def check_md5_obj_path( + self, b64_md5: B64MD5, size: int + ) -> tuple[FilePathStr, bool, Opener]: + # Check if we're using vs skipping the cache + if self._override_cache_path is not None: + skip_cache = True + path = Path(self._override_cache_path) + else: + skip_cache = False + hex_md5 = b64_to_hex_id(b64_md5) + path = self._obj_dir / "md5" / hex_md5[:2] / hex_md5[2:] + return self._check_or_create(path, size, skip_cache=skip_cache) + + # TODO(spencerpearson): this method at least needs its signature changed. + # An ETag is not (necessarily) a checksum. + def check_etag_obj_path( + self, + url: URIStr, + etag: ETag, + size: int, + ) -> tuple[FilePathStr, bool, Opener]: + # Check if we're using vs skipping the cache + if self._override_cache_path is not None: + skip_cache = True + path = Path(self._override_cache_path) + else: + skip_cache = False + hexhash = hashlib.sha256( + hashlib.sha256(url.encode("utf-8")).digest() + + hashlib.sha256(etag.encode("utf-8")).digest() + ).hexdigest() + path = self._obj_dir / "etag" / hexhash[:2] / hexhash[2:] + return self._check_or_create(path, size, skip_cache=skip_cache) + + def _check_or_create( + self, path: Path, size: int, skip_cache: bool = False + ) -> tuple[FilePathStr, bool, Opener]: + opener = self._opener(path, size, skip_cache=skip_cache) + hit = path.is_file() and path.stat().st_size == size + return FilePathStr(str(path)), hit, opener + + def cleanup( + self, + target_size: int | None = None, + remove_temp: bool = False, + target_fraction: float | None = None, + ) -> int: + """Clean up the cache, removing the least recently used files first. + + Args: + target_size: The target size of the cache in bytes. If the cache is larger + than this, we will remove the least recently used files until the cache + is smaller than this size. + remove_temp: Whether to remove temporary files. Temporary files are files + that are currently being written to the cache. If remove_temp is True, + all temp files will be removed, regardless of the target_size or + target_fraction. + target_fraction: The target fraction of the cache to reclaim. If the cache + is larger than this, we will remove the least recently used files until + the cache is smaller than this fraction of its current size. It is an + error to specify both target_size and target_fraction. + + Returns: + The number of bytes reclaimed. + """ + if target_size is None and target_fraction is None: + # Default to clearing the entire cache. + target_size = 0 + if target_size is not None and target_fraction is not None: + raise ValueError("Cannot specify both target_size and target_fraction") + if target_size is not None and target_size < 0: + raise ValueError("target_size must be non-negative") + if target_fraction is not None and (target_fraction < 0 or target_fraction > 1): + raise ValueError("target_fraction must be between 0 and 1") + + bytes_reclaimed = 0 + total_size = 0 + temp_size = 0 + + # Remove all temporary files if requested. Otherwise sum their size. + for entry in files_in(self._temp_dir): + size = entry.stat().st_size + total_size += size + if remove_temp: + try: + os.remove(entry.path) + bytes_reclaimed += size + except OSError: + pass + else: + temp_size += size + if temp_size: + wandb.termwarn( + f"Cache contains {util.to_human_size(temp_size)} of temporary files. " + "Run `wandb artifact cleanup --remove-temp` to remove them." + ) + + entries = [] + for file_entry in files_in(self._obj_dir): + total_size += file_entry.stat().st_size + entries.append(file_entry) + + if target_fraction is not None: + target_size = int(total_size * target_fraction) + assert target_size is not None + + for entry in sorted(entries, key=lambda x: x.stat().st_atime): + if total_size <= target_size: + return bytes_reclaimed + try: + os.remove(entry.path) + except OSError: + pass + total_size -= entry.stat().st_size + bytes_reclaimed += entry.stat().st_size + + if total_size > target_size: + wandb.termerror( + f"Failed to reclaim enough space in {self._cache_dir}. Try running" + " `wandb artifact cache cleanup --remove-temp` to remove temporary files." + ) + + return bytes_reclaimed + + def _free_space(self) -> int: + """Return the number of bytes of free space in the cache directory.""" + return shutil.disk_usage(self._cache_dir)[2] + + def _reserve_space(self, size: int) -> None: + """If a `size` write would exceed disk space, remove cached items to make space. + + Raises: + OSError: If there is not enough space to write `size` bytes, even after + removing cached items. + """ + if size <= self._free_space(): + return + + wandb.termwarn("Cache size exceeded. Attempting to reclaim space...") + self.cleanup(target_fraction=0.5) + if size <= self._free_space(): + return + + self.cleanup(target_size=0) + if size > self._free_space(): + raise OSError(errno.ENOSPC, f"Insufficient free space in {self._cache_dir}") + + def _opener(self, path: Path, size: int, skip_cache: bool = False) -> Opener: + @contextlib.contextmanager + def atomic_open(mode: str = "w") -> Iterator[IO]: + if "a" in mode: + raise ValueError("Appending to cache files is not supported") + + if skip_cache: + # We skip the cache, but we'll still need an intermediate, temporary file to ensure atomicity. + # Put the temp file in the same root as the destination file in an attempt to avoid moving/copying + # across filesystems. + temp_dir = path.parent + else: + self._reserve_space(size) + temp_dir = self._temp_dir + + temp_dir.mkdir(parents=True, exist_ok=True) + temp_file = NamedTemporaryFile(dir=temp_dir, mode=mode, delete=False) + try: + yield temp_file + temp_file.close() + os.chmod(temp_file.name, 0o666 & ~self._sys_umask) + path.parent.mkdir(parents=True, exist_ok=True) + os.replace(temp_file.name, path) + except Exception: + os.remove(temp_file.name) + raise + + return atomic_open + + def _ensure_write_permissions(self) -> None: + """Raise an error if we cannot write to the cache directory.""" + try: + self._temp_dir.mkdir(parents=True, exist_ok=True) + with NamedTemporaryFile(dir=self._temp_dir) as f: + f.write(b"wandb") + except PermissionError as e: + raise PermissionError( + f"Unable to write to {self._cache_dir}. " + "Ensure that the current user has write permissions." + ) from e + + +_artifact_file_cache: ArtifactFileCache | None = None + + +def get_artifact_file_cache() -> ArtifactFileCache: + global _artifact_file_cache + cache_dir = env.get_cache_dir() / "artifacts" + if _artifact_file_cache is None or _artifact_file_cache._cache_dir != cache_dir: + _artifact_file_cache = ArtifactFileCache(cache_dir) + return _artifact_file_cache diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_instance_cache.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_instance_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..e6fd6e9676d048a0a7fc0715bb7ae107ae01e7cc --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_instance_cache.py @@ -0,0 +1,17 @@ +"""Recent Artifact storage. + +Artifacts are registered in the cache to ensure they won't be immediately garbage +collected and can be retrieved by their ID. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from wandb.sdk.lib.capped_dict import CappedDict + +if TYPE_CHECKING: + from wandb.sdk.artifacts.artifact import Artifact + +# There is nothing special about the artifact cache, it's just a global capped dict. +artifact_instance_cache: dict[str, Artifact] = CappedDict(100) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifest.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifest.py new file mode 100644 index 0000000000000000000000000000000000000000..74df51a7fd697374859f6eec3bed1d9f3b636c61 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifest.py @@ -0,0 +1,75 @@ +"""Artifact manifest.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Mapping + +from wandb.sdk.internal.internal_api import Api as InternalApi +from wandb.sdk.lib.hashutil import HexMD5 + +if TYPE_CHECKING: + from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry + from wandb.sdk.artifacts.storage_policy import StoragePolicy + + +class ArtifactManifest: + entries: dict[str, ArtifactManifestEntry] + + @classmethod + def from_manifest_json( + cls, manifest_json: dict, api: InternalApi | None = None + ) -> ArtifactManifest: + if "version" not in manifest_json: + raise ValueError("Invalid manifest format. Must contain version field.") + version = manifest_json["version"] + for sub in cls.__subclasses__(): + if sub.version() == version: + return sub.from_manifest_json(manifest_json, api=api) + raise ValueError("Invalid manifest version.") + + @classmethod + def version(cls) -> int: + raise NotImplementedError + + def __init__( + self, + storage_policy: StoragePolicy, + entries: Mapping[str, ArtifactManifestEntry] | None = None, + ) -> None: + self.storage_policy = storage_policy + self.entries = dict(entries) if entries else {} + + def __len__(self) -> int: + return len(self.entries) + + def to_manifest_json(self) -> dict: + raise NotImplementedError + + def digest(self) -> HexMD5: + raise NotImplementedError + + def add_entry(self, entry: ArtifactManifestEntry, overwrite: bool = False) -> None: + path = entry.path + if not overwrite: + prev_entry = self.entries.get(path) + if prev_entry and (entry.digest != prev_entry.digest): + raise ValueError(f"Cannot add the same path twice: {path!r}") + self.entries[path] = entry + + def remove_entry(self, entry: ArtifactManifestEntry) -> None: + try: + del self.entries[entry.path] + except LookupError: + raise FileNotFoundError(f"Cannot remove missing entry: '{entry.path}'") + + def get_entry_by_path(self, path: str) -> ArtifactManifestEntry | None: + return self.entries.get(path) + + def get_entries_in_directory(self, directory: str) -> list[ArtifactManifestEntry]: + return [ + self.entries[entry_key] + for entry_key in self.entries + if entry_key.startswith( + directory + "/" + ) # entries use forward slash even for windows + ] diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifest_entry.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifest_entry.py new file mode 100644 index 0000000000000000000000000000000000000000..c4163d61f6f67e90237c2039f37dd5ca2a3bec21 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifest_entry.py @@ -0,0 +1,259 @@ +"""Artifact manifest entry.""" + +from __future__ import annotations + +import concurrent.futures +import json +import logging +import os +from pathlib import Path +from typing import TYPE_CHECKING +from urllib.parse import urlparse + +from wandb.proto.wandb_deprecated import Deprecated +from wandb.sdk.lib import filesystem +from wandb.sdk.lib.deprecate import deprecate +from wandb.sdk.lib.hashutil import ( + B64MD5, + ETag, + b64_to_hex_id, + hex_to_b64_id, + md5_file_b64, +) +from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from typing_extensions import TypedDict + + from wandb.sdk.artifacts.artifact import Artifact + + class ArtifactManifestEntryDict(TypedDict, total=False): + path: str + digest: str + skip_cache: bool + ref: str + birthArtifactID: str + size: int + extra: dict + local_path: str + + +class ArtifactManifestEntry: + """A single entry in an artifact manifest.""" + + path: LogicalPath + digest: B64MD5 | URIStr | FilePathStr | ETag + skip_cache: bool + ref: FilePathStr | URIStr | None + birth_artifact_id: str | None + size: int | None + extra: dict + local_path: str | None + + _parent_artifact: Artifact | None = None + _download_url: str | None = None + + def __init__( + self, + path: StrPath, + digest: B64MD5 | URIStr | FilePathStr | ETag, + skip_cache: bool | None = False, + ref: FilePathStr | URIStr | None = None, + birth_artifact_id: str | None = None, + size: int | None = None, + extra: dict | None = None, + local_path: StrPath | None = None, + ) -> None: + self.path = LogicalPath(path) + self.digest = digest + self.ref = ref + self.birth_artifact_id = birth_artifact_id + self.size = size + self.extra = extra or {} + self.local_path = str(local_path) if local_path else None + if self.local_path and self.size is None: + self.size = Path(self.local_path).stat().st_size + self.skip_cache = skip_cache or False + + def __repr__(self) -> str: + cls = self.__class__.__name__ + ref = f", ref={self.ref!r}" if self.ref is not None else "" + birth_artifact_id = ( + f", birth_artifact_id={self.birth_artifact_id!r}" + if self.birth_artifact_id is not None + else "" + ) + size = f", size={self.size}" if self.size is not None else "" + extra = f", extra={json.dumps(self.extra)}" if self.extra else "" + local_path = f", local_path={self.local_path!r}" if self.local_path else "" + skip_cache = f", skip_cache={self.skip_cache}" + others = ref + birth_artifact_id + size + extra + local_path + skip_cache + return f"{cls}(path={self.path!r}, digest={self.digest!r}{others})" + + def __eq__(self, other: object) -> bool: + """Strict equality, comparing all public fields. + + ArtifactManifestEntries for the same file may not compare equal if they were + added in different ways or created for different parent artifacts. + """ + if not isinstance(other, ArtifactManifestEntry): + return False + return ( + self.path == other.path + and self.digest == other.digest + and self.ref == other.ref + and self.birth_artifact_id == other.birth_artifact_id + and self.size == other.size + and self.extra == other.extra + and self.local_path == other.local_path + and self.skip_cache == other.skip_cache + ) + + @property + def name(self) -> LogicalPath: + """Deprecated; use `path` instead.""" + deprecate( + field_name=Deprecated.artifactmanifestentry__name, + warning_message="ArtifactManifestEntry.name is deprecated, use .path instead.", + ) + return self.path + + def parent_artifact(self) -> Artifact: + """Get the artifact to which this artifact entry belongs. + + Returns: + (PublicArtifact): The parent artifact + """ + if self._parent_artifact is None: + raise NotImplementedError + return self._parent_artifact + + def download( + self, + root: str | None = None, + skip_cache: bool | None = None, + executor: concurrent.futures.Executor | None = None, + multipart: bool | None = None, + ) -> FilePathStr: + """Download this artifact entry to the specified root path. + + Args: + root: (str, optional) The root path in which to download this + artifact entry. Defaults to the artifact's root. + + Returns: + (str): The path of the downloaded artifact entry. + """ + if self._parent_artifact is None: + raise NotImplementedError + + root = root or self._parent_artifact._default_root() + self._parent_artifact._add_download_root(root) + path = str(Path(self.path)) + dest_path = os.path.join(root, path) + + if skip_cache: + override_cache_path = dest_path + else: + override_cache_path = None + + # Skip checking the cache (and possibly downloading) if the file already exists + # and has the digest we're expecting. + try: + md5_hash = md5_file_b64(dest_path) + except (FileNotFoundError, IsADirectoryError): + logger.debug(f"unable to find {dest_path}, skip searching for file") + else: + if self.digest == md5_hash: + return FilePathStr(dest_path) + + if self.ref is not None: + cache_path = self._parent_artifact.manifest.storage_policy.load_reference( + self, local=True, dest_path=override_cache_path + ) + else: + cache_path = self._parent_artifact.manifest.storage_policy.load_file( + self._parent_artifact, + self, + dest_path=override_cache_path, + executor=executor, + multipart=multipart, + ) + + if skip_cache: + return FilePathStr(dest_path) + else: + return FilePathStr( + str(filesystem.copy_or_overwrite_changed(cache_path, dest_path)) + ) + + def ref_target(self) -> FilePathStr | URIStr: + """Get the reference URL that is targeted by this artifact entry. + + Returns: + (str): The reference URL of this artifact entry. + + Raises: + ValueError: If this artifact entry was not a reference. + """ + if self.ref is None: + raise ValueError("Only reference entries support ref_target().") + if self._parent_artifact is None: + return self.ref + return self._parent_artifact.manifest.storage_policy.load_reference( + self._parent_artifact.manifest.entries[self.path], local=False + ) + + def ref_url(self) -> str: + """Get a URL to this artifact entry. + + These URLs can be referenced by another artifact. + + Returns: + (str): A URL representing this artifact entry. + + Examples: + Basic usage + ``` + ref_url = source_artifact.get_entry("file.txt").ref_url() + derived_artifact.add_reference(ref_url) + ``` + """ + if self._parent_artifact is None: + raise NotImplementedError + assert self._parent_artifact.id is not None + return ( + "wandb-artifact://" + + b64_to_hex_id(B64MD5(self._parent_artifact.id)) + + "/" + + self.path + ) + + def to_json(self) -> ArtifactManifestEntryDict: + contents: ArtifactManifestEntryDict = { + "path": self.path, + "digest": self.digest, + } + if self.size is not None: + contents["size"] = self.size + if self.ref: + contents["ref"] = self.ref + if self.birth_artifact_id: + contents["birthArtifactID"] = self.birth_artifact_id + if self.local_path: + contents["local_path"] = self.local_path + if self.skip_cache: + contents["skip_cache"] = self.skip_cache + if self.extra: + contents["extra"] = self.extra + return contents + + def _is_artifact_reference(self) -> bool: + return self.ref is not None and urlparse(self.ref).scheme == "wandb-artifact" + + def _referenced_artifact_id(self) -> str | None: + if not self._is_artifact_reference(): + return None + return hex_to_b64_id(urlparse(self.ref).netloc) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/__init__.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2078fd8de6b827922b346da0828fab0d26915974 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/__pycache__/artifact_manifest_v1.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/__pycache__/artifact_manifest_v1.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..555545fb458a5eac8689a5e14c9933ba80f3b64e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/__pycache__/artifact_manifest_v1.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py new file mode 100644 index 0000000000000000000000000000000000000000..613871ff6c137385ad026b7f397a00904637453e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py @@ -0,0 +1,92 @@ +"""Artifact manifest v1.""" + +from __future__ import annotations + +from typing import Any, Mapping + +from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest +from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry +from wandb.sdk.artifacts.storage_policy import StoragePolicy +from wandb.sdk.internal.internal_api import Api as InternalApi +from wandb.sdk.lib.hashutil import HexMD5, _md5 + + +class ArtifactManifestV1(ArtifactManifest): + @classmethod + def version(cls) -> int: + return 1 + + @classmethod + def from_manifest_json( + cls, manifest_json: dict, api: InternalApi | None = None + ) -> ArtifactManifestV1: + if manifest_json["version"] != cls.version(): + raise ValueError( + "Expected manifest version 1, got {}".format(manifest_json["version"]) + ) + + storage_policy_name = manifest_json["storagePolicy"] + storage_policy_config = manifest_json.get("storagePolicyConfig", {}) + storage_policy_cls = StoragePolicy.lookup_by_name(storage_policy_name) + + entries: Mapping[str, ArtifactManifestEntry] + entries = { + name: ArtifactManifestEntry( + path=name, + digest=val["digest"], + birth_artifact_id=val.get("birthArtifactID"), + ref=val.get("ref"), + size=val.get("size"), + extra=val.get("extra"), + local_path=val.get("local_path"), + skip_cache=val.get("skip_cache"), + ) + for name, val in manifest_json["contents"].items() + } + + return cls( + storage_policy_cls.from_config(storage_policy_config, api=api), entries + ) + + def __init__( + self, + storage_policy: StoragePolicy, + entries: Mapping[str, ArtifactManifestEntry] | None = None, + ) -> None: + super().__init__(storage_policy, entries=entries) + + def to_manifest_json(self) -> dict: + """This is the JSON that's stored in wandb_manifest.json. + + If include_local is True we also include the local paths to files. This is + used to represent an artifact that's waiting to be saved on the current + system. We don't need to include the local paths in the artifact manifest + contents. + """ + contents = {} + for entry in sorted(self.entries.values(), key=lambda k: k.path): + json_entry: dict[str, Any] = { + "digest": entry.digest, + } + if entry.birth_artifact_id: + json_entry["birthArtifactID"] = entry.birth_artifact_id + if entry.ref: + json_entry["ref"] = entry.ref + if entry.extra: + json_entry["extra"] = entry.extra + if entry.size is not None: + json_entry["size"] = entry.size + contents[entry.path] = json_entry + return { + "version": self.__class__.version(), + "storagePolicy": self.storage_policy.name(), + "storagePolicyConfig": self.storage_policy.config() or {}, + "contents": contents, + } + + def digest(self) -> HexMD5: + hasher = _md5() + hasher.update(b"wandb-artifact-manifest-v1\n") + for name, entry in sorted(self.entries.items(), key=lambda kv: kv[0]): + hasher.update(f"{name}:{entry.digest}\n".encode()) + return HexMD5(hasher.hexdigest()) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_saver.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_saver.py new file mode 100644 index 0000000000000000000000000000000000000000..e163f150060f553cf22ce5dd9916dcc3868d4830 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_saver.py @@ -0,0 +1,277 @@ +"""Artifact saver.""" + +from __future__ import annotations + +import concurrent.futures +import json +import os +import tempfile +from typing import TYPE_CHECKING, Awaitable, Sequence + +import wandb +import wandb.filesync.step_prepare +from wandb import util +from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest +from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, md5_file_b64 +from wandb.sdk.lib.paths import URIStr + +if TYPE_CHECKING: + from typing import Protocol + + from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry + from wandb.sdk.internal.file_pusher import FilePusher + from wandb.sdk.internal.internal_api import Api as InternalApi + from wandb.sdk.internal.progress import ProgressFn + + class SaveFn(Protocol): + def __call__( + self, entry: ArtifactManifestEntry, progress_callback: ProgressFn + ) -> bool: + pass + + class SaveFnAsync(Protocol): + def __call__( + self, entry: ArtifactManifestEntry, progress_callback: ProgressFn + ) -> Awaitable[bool]: + pass + + +class ArtifactSaver: + _server_artifact: dict | None # TODO better define this dict + + def __init__( + self, + api: InternalApi, + digest: str, + manifest_json: dict, + file_pusher: FilePusher, + is_user_created: bool = False, + ) -> None: + self._api = api + self._file_pusher = file_pusher + self._digest = digest + self._manifest = ArtifactManifest.from_manifest_json( + manifest_json, + api=self._api, + ) + self._is_user_created = is_user_created + self._server_artifact = None + + def save( + self, + entity: str, + project: str, + type: str, + name: str, + client_id: str, + sequence_client_id: str, + distributed_id: str | None = None, + finalize: bool = True, + metadata: dict | None = None, + ttl_duration_seconds: int | None = None, + description: str | None = None, + aliases: Sequence[str] | None = None, + tags: Sequence[str] | None = None, + use_after_commit: bool = False, + incremental: bool = False, + history_step: int | None = None, + base_id: str | None = None, + ) -> dict | None: + return self._save_internal( + entity, + project, + type, + name, + client_id, + sequence_client_id, + distributed_id, + finalize, + metadata, + ttl_duration_seconds, + description, + aliases, + tags, + use_after_commit, + incremental, + history_step, + base_id, + ) + + def _save_internal( + self, + entity: str, + project: str, + type: str, + name: str, + client_id: str, + sequence_client_id: str, + distributed_id: str | None = None, + finalize: bool = True, + metadata: dict | None = None, + ttl_duration_seconds: int | None = None, + description: str | None = None, + aliases: Sequence[str] | None = None, + tags: Sequence[str] | None = None, + use_after_commit: bool = False, + incremental: bool = False, + history_step: int | None = None, + base_id: str | None = None, + ) -> dict | None: + alias_specs = [] + for alias in aliases or []: + alias_specs.append({"artifactCollectionName": name, "alias": alias}) + + tag_specs = [{"tagName": tag} for tag in tags or []] + + """Returns the server artifact.""" + self._server_artifact, latest = self._api.create_artifact( + type, + name, + self._digest, + metadata=metadata, + ttl_duration_seconds=ttl_duration_seconds, + aliases=alias_specs, + tags=tag_specs, + description=description, + is_user_created=self._is_user_created, + distributed_id=distributed_id, + client_id=client_id, + sequence_client_id=sequence_client_id, + history_step=history_step, + ) + + assert self._server_artifact is not None # mypy optionality unwrapper + artifact_id = self._server_artifact["id"] + if base_id is None and latest: + base_id = latest["id"] + if self._server_artifact["state"] == "COMMITTED": + if use_after_commit: + self._api.use_artifact( + artifact_id, + artifact_entity_name=entity, + artifact_project_name=project, + ) + return self._server_artifact + if ( + self._server_artifact["state"] != "PENDING" + # For old servers, see https://github.com/wandb/wandb/pull/6190 + and self._server_artifact["state"] != "DELETED" + ): + raise Exception( + 'Unknown artifact state "{}"'.format(self._server_artifact["state"]) + ) + + manifest_type = "FULL" + manifest_filename = "wandb_manifest.json" + if incremental: + manifest_type = "INCREMENTAL" + manifest_filename = "wandb_manifest.incremental.json" + elif distributed_id: + manifest_type = "PATCH" + manifest_filename = "wandb_manifest.patch.json" + artifact_manifest_id, _ = self._api.create_artifact_manifest( + manifest_filename, + "", + artifact_id, + base_artifact_id=base_id, + include_upload=False, + type=manifest_type, + ) + + step_prepare = wandb.filesync.step_prepare.StepPrepare( + self._api, 0.1, 0.01, 1000 + ) # TODO: params + step_prepare.start() + + # Upload Artifact "L1" files, the actual artifact contents + self._file_pusher.store_manifest_files( + self._manifest, + artifact_id, + lambda entry, progress_callback: self._manifest.storage_policy.store_file( + artifact_id, + artifact_manifest_id, + entry, + step_prepare, + progress_callback=progress_callback, + ), + ) + + def before_commit() -> None: + self._resolve_client_id_manifest_references() + with tempfile.NamedTemporaryFile("w+", suffix=".json", delete=False) as fp: + path = os.path.abspath(fp.name) + json.dump(self._manifest.to_manifest_json(), fp, indent=4) + digest = md5_file_b64(path) + if distributed_id or incremental: + # If we're in the distributed flow, we want to update the + # patch manifest we created with our finalized digest. + _, resp = self._api.update_artifact_manifest( + artifact_manifest_id, + digest=digest, + ) + else: + # In the regular flow, we can recreate the full manifest with the + # updated digest. + # + # NOTE: We do this for backwards compatibility with older backends + # that don't support the 'updateArtifactManifest' API. + _, resp = self._api.create_artifact_manifest( + manifest_filename, + digest, + artifact_id, + base_artifact_id=base_id, + ) + + # We're duplicating the file upload logic a little, which isn't great. + upload_url = resp["uploadUrl"] + upload_headers = resp["uploadHeaders"] + extra_headers = {} + for upload_header in upload_headers: + key, val = upload_header.split(":", 1) + extra_headers[key] = val + with open(path, "rb") as fp2: + self._api.upload_file_retry( + upload_url, + fp2, + extra_headers=extra_headers, + ) + + commit_result: concurrent.futures.Future[None] = concurrent.futures.Future() + + # This will queue the commit. It will only happen after all the file uploads are done + self._file_pusher.commit_artifact( + artifact_id, + finalize=finalize, + before_commit=before_commit, + result_future=commit_result, + ) + + # Block until all artifact files are uploaded and the + # artifact is committed. + try: + commit_result.result() + finally: + step_prepare.shutdown() + + if finalize and use_after_commit: + self._api.use_artifact( + artifact_id, + artifact_entity_name=entity, + artifact_project_name=project, + ) + + return self._server_artifact + + def _resolve_client_id_manifest_references(self) -> None: + for entry_path in self._manifest.entries: + entry = self._manifest.entries[entry_path] + if entry.ref is not None: + if entry.ref.startswith("wandb-client-artifact:"): + client_id = util.host_from_path(entry.ref) + artifact_file_path = util.uri_from_path(entry.ref) + artifact_id = self._api._resolve_client_id(client_id) + if artifact_id is None: + raise RuntimeError(f"Could not resolve client id {client_id}") + entry.ref = URIStr( + f"wandb-artifact://{b64_to_hex_id(B64MD5(artifact_id))}/{artifact_file_path}" + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_state.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_state.py new file mode 100644 index 0000000000000000000000000000000000000000..4df836c8b60e482e3de8a103b4b0e646800922a7 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_state.py @@ -0,0 +1,13 @@ +"""Artifact state.""" + +from __future__ import annotations + +from enum import Enum + + +class ArtifactState(Enum): + PENDING = "PENDING" + COMMITTED = "COMMITTED" + DELETED = "DELETED" + GARBAGE_COLLECTED = "GARBAGE_COLLECTED" + PENDING_DELETION = "PENDING_DELETION" diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_ttl.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_ttl.py new file mode 100644 index 0000000000000000000000000000000000000000..e63c69031923aa36a2214526443861497f3575d6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/artifact_ttl.py @@ -0,0 +1,9 @@ +"""Artifact TTL.""" + +from __future__ import annotations + +from enum import Enum + + +class ArtifactTTL(Enum): + INHERIT = 0 diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/exceptions.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..541d9c7d47b813edf899d7cac9b7d001cabc631f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/exceptions.py @@ -0,0 +1,71 @@ +"""Artifact exceptions.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, TypeVar + +from wandb import errors + +if TYPE_CHECKING: + from wandb.sdk.artifacts.artifact import Artifact + + ArtifactT = TypeVar("ArtifactT", bound=Artifact) + + +class ArtifactStatusError(AttributeError): + """Raised when an artifact is in an invalid state for the requested operation.""" + + def __init__( + self, + msg: str = "Artifact is in an invalid state for the requested operation.", + name: str | None = None, + obj: ArtifactT | None = None, + ): + # Follow the same pattern as AttributeError in python 3.10+ by `name/obj` attributes + # See: https://docs.python.org/3/library/exceptions.html#AttributeError + try: + super().__init__(msg, name=name, obj=obj) + except TypeError: + # The `name`/`obj` keyword args and attributes were only added in python >= 3.10 + super().__init__(msg) + self.name = name or "" + self.obj = obj + + +class ArtifactNotLoggedError(ArtifactStatusError): + """Raised for Artifact methods or attributes only available after logging.""" + + def __init__(self, fullname: str, obj: ArtifactT): + *_, name = fullname.split(".") + msg = ( + f"{fullname!r} used prior to logging artifact or while in offline mode. " + f"Call {type(obj).wait.__qualname__}() before accessing logged artifact properties." + ) + super().__init__(msg=msg, name=name, obj=obj) + + +class ArtifactFinalizedError(ArtifactStatusError): + """Raised for Artifact methods or attributes that can't be changed after logging.""" + + def __init__(self, fullname: str, obj: ArtifactT): + *_, name = fullname.split(".") + msg = f"{fullname!r} used on logged artifact. Can't modify finalized artifact." + super().__init__(msg=msg, name=name, obj=obj) + + +class WaitTimeoutError(errors.Error): + """Raised when wait() timeout occurs before process is finished.""" + + +class TooFewItemsError(ValueError): + """Raised when there are fewer items than expected in a collection. + + Intended for internal use only. + """ + + +class TooManyItemsError(ValueError): + """Raised when there are more items than expected in a collection. + + Intended for internal use only. + """ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/staging.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/staging.py new file mode 100644 index 0000000000000000000000000000000000000000..24a6313ca0507d1f4ee110570c80062644c08cb9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/staging.py @@ -0,0 +1,27 @@ +"""Manages artifact file staging. + +Artifact files are copied to the staging area as soon as they are added to an artifact +in order to avoid file changes corrupting the artifact. Once the upload is complete, the +file should be moved to the artifact cache. +""" + +from __future__ import annotations + +import os + +from wandb import env +from wandb.sdk.lib.filesystem import mkdir_exists_ok +from wandb.sdk.lib.paths import FilePathStr + + +def get_staging_dir() -> FilePathStr: + path = os.path.join(env.get_data_dir(), "artifacts", "staging") + try: + mkdir_exists_ok(path) + except OSError as e: + raise PermissionError( + f"Unable to write staging files to {path}. To fix this problem, please set " + f"{env.DATA_DIR} to a directory where you have the necessary write access." + ) from e + + return FilePathStr(os.path.abspath(os.path.expanduser(path))) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handler.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..0f3abd02c93dd87497fbbfc413dbc9695a259e06 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handler.py @@ -0,0 +1,62 @@ +"""Storage handler.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Sequence + +from wandb.sdk.lib.paths import FilePathStr, URIStr + +if TYPE_CHECKING: + from urllib.parse import ParseResult + + from wandb.sdk.artifacts.artifact import Artifact + from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry + +DEFAULT_MAX_OBJECTS = 10**7 + + +class StorageHandler: + def can_handle(self, parsed_url: ParseResult) -> bool: + """Checks whether this handler can handle the given url. + + Returns: + Whether this handler can handle the given url. + """ + raise NotImplementedError + + def load_path( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + ) -> URIStr | FilePathStr: + """Load a file or directory given the corresponding index entry. + + Args: + manifest_entry: The index entry to load + local: Whether to load the file locally or not + + Returns: + A path to the file represented by `index_entry` + """ + raise NotImplementedError + + def store_path( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: str | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + """Store the file or directory at the given path to the specified artifact. + + Args: + path: The path to store + name: If specified, the logical name that should map to `path` + checksum: Whether to compute the checksum of the file + max_objects: The maximum number of objects to store + + Returns: + A list of manifest entries to store within the artifact + """ + raise NotImplementedError diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__init__.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1cb1e13eabb0fbdab0ae804c1d4069acf8d55866 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/azure_handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/azure_handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0fcf0fdc58e6ecdbfdc55edf46589082dec8259 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/azure_handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/gcs_handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/gcs_handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ec19d0b87158b4e4058775e2810508fb92856fbc Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/gcs_handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/http_handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/http_handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c1ac2af2bd00d00c674f421538d00a2f540f27de Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/http_handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/local_file_handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/local_file_handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0c1e41666c2105028312bf5cce2f53b562e2a5e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/local_file_handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/multi_handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/multi_handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..532e6bdfff3b545f00124e2c198c3825c648ec19 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/multi_handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/s3_handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/s3_handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fe3640944537a644ea4daa05eb4a603c3cc4c1d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/s3_handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/tracking_handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/tracking_handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9818c83b7e9f38ee6ff540efef771211b30efb5e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/tracking_handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/wb_artifact_handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/wb_artifact_handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b573e985c7498f7d18e6598c1e000f07b0ea8a41 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/wb_artifact_handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/wb_local_artifact_handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/wb_local_artifact_handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c0fa44c796c3296c45ce5fbcc1052eaf8d3d2696 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/__pycache__/wb_local_artifact_handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/azure_handler.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/azure_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..6a933a7e71829afd5e2e4e719b51b4160ec2c26d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/azure_handler.py @@ -0,0 +1,214 @@ +"""Azure storage handler.""" + +from __future__ import annotations + +from pathlib import PurePosixPath +from types import ModuleType +from typing import TYPE_CHECKING, Sequence +from urllib.parse import ParseResult, parse_qsl, urlparse + +import wandb +from wandb import util +from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache +from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry +from wandb.sdk.artifacts.storage_handler import DEFAULT_MAX_OBJECTS, StorageHandler +from wandb.sdk.lib.hashutil import ETag +from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr + +if TYPE_CHECKING: + import azure.identity # type: ignore + import azure.storage.blob # type: ignore + + from wandb.sdk.artifacts.artifact import Artifact + + +class AzureHandler(StorageHandler): + def can_handle(self, parsed_url: ParseResult) -> bool: + return parsed_url.scheme == "https" and parsed_url.netloc.endswith( + ".blob.core.windows.net" + ) + + def __init__(self, scheme: str | None = None) -> None: + self._cache = get_artifact_file_cache() + + def load_path( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + ) -> URIStr | FilePathStr: + assert manifest_entry.ref is not None + if not local: + return manifest_entry.ref + + path, hit, cache_open = self._cache.check_etag_obj_path( + URIStr(manifest_entry.ref), + ETag(manifest_entry.digest), + manifest_entry.size or 0, + ) + if hit: + return path + + account_url, container_name, blob_name, query = self._parse_uri( + manifest_entry.ref + ) + version_id = manifest_entry.extra.get("versionID") + blob_service_client = self._get_module("azure.storage.blob").BlobServiceClient( + account_url, credential=self._get_credential(account_url) + ) + blob_client = blob_service_client.get_blob_client( + container=container_name, blob=blob_name + ) + if version_id is None: + # Try current version, then all versions. + try: + downloader = blob_client.download_blob( + etag=manifest_entry.digest, + match_condition=self._get_module( + "azure.core" + ).MatchConditions.IfNotModified, + ) + except self._get_module("azure.core.exceptions").ResourceModifiedError: + container_client = blob_service_client.get_container_client( + container_name + ) + for blob_properties in container_client.walk_blobs( + name_starts_with=blob_name, include=["versions"] + ): + if ( + blob_properties.name == blob_name + and blob_properties.etag == manifest_entry.digest + and blob_properties.version_id is not None + ): + downloader = blob_client.download_blob( + version_id=blob_properties.version_id + ) + break + else: # didn't break + raise ValueError( + f"Couldn't find blob version for {manifest_entry.ref} matching " + f"etag {manifest_entry.digest}." + ) + else: + downloader = blob_client.download_blob(version_id=version_id) + with cache_open(mode="wb") as f: + downloader.readinto(f) + return path + + def store_path( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: StrPath | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + account_url, container_name, blob_name, query = self._parse_uri(path) + path = URIStr(f"{account_url}/{container_name}/{blob_name}") + + if not checksum: + return [ + ArtifactManifestEntry(path=name or blob_name, digest=path, ref=path) + ] + + blob_service_client = self._get_module("azure.storage.blob").BlobServiceClient( + account_url, credential=self._get_credential(account_url) + ) + blob_client = blob_service_client.get_blob_client( + container=container_name, blob=blob_name + ) + if blob_client.exists(version_id=query.get("versionId")): + blob_properties = blob_client.get_blob_properties( + version_id=query.get("versionId") + ) + + if not self._is_directory_stub(blob_properties): + return [ + self._create_entry( + blob_properties, + path=name or PurePosixPath(blob_name).name, + ref=URIStr( + f"{account_url}/{container_name}/{blob_properties.name}" + ), + ) + ] + + entries: list[ArtifactManifestEntry] = [] + container_client = blob_service_client.get_container_client(container_name) + max_objects = max_objects or DEFAULT_MAX_OBJECTS + for blob_properties in container_client.list_blobs( + name_starts_with=f"{blob_name}/" + ): + if len(entries) >= max_objects: + wandb.termwarn( + f"Found more than {max_objects} objects under path, limiting upload " + f"to {max_objects} objects. Increase max_objects to upload more" + ) + break + if not self._is_directory_stub(blob_properties): + suffix = PurePosixPath(blob_properties.name).relative_to(blob_name) + entries.append( + self._create_entry( + blob_properties, + path=LogicalPath(name) / suffix if name else suffix, + ref=URIStr( + f"{account_url}/{container_name}/{blob_properties.name}" + ), + ) + ) + + return entries + + def _get_module(self, name: str) -> ModuleType: + module = util.get_module( + name, + lazy=False, + required="Azure references require the azure library, run " + "pip install wandb[azure]", + ) + assert isinstance(module, ModuleType) + return module + + def _get_credential( + self, account_url: str + ) -> azure.identity.DefaultAzureCredential | str: + # NOTE: Always returns default credential for reinit="create_new" runs. + if ( + wandb.run + and wandb.run.settings.azure_account_url_to_access_key is not None + and account_url in wandb.run.settings.azure_account_url_to_access_key + ): + return wandb.run.settings.azure_account_url_to_access_key[account_url] + return self._get_module("azure.identity").DefaultAzureCredential() + + def _parse_uri(self, uri: str) -> tuple[str, str, str, dict[str, str]]: + parsed_url = urlparse(uri) + query = dict(parse_qsl(parsed_url.query)) + account_url = f"{parsed_url.scheme}://{parsed_url.netloc}" + _, container_name, blob_name = parsed_url.path.split("/", 2) + return account_url, container_name, blob_name, query + + def _create_entry( + self, + blob_properties: azure.storage.blob.BlobProperties, + path: StrPath, + ref: URIStr, + ) -> ArtifactManifestEntry: + extra = {"etag": blob_properties.etag.strip('"')} + if blob_properties.version_id: + extra["versionID"] = blob_properties.version_id + return ArtifactManifestEntry( + path=path, + ref=ref, + digest=blob_properties.etag.strip('"'), + size=blob_properties.size, + extra=extra, + ) + + def _is_directory_stub( + self, blob_properties: azure.storage.blob.BlobProperties + ) -> bool: + return ( + blob_properties.has_key("metadata") + and "hdi_isfolder" in blob_properties.metadata + and blob_properties.metadata["hdi_isfolder"] == "true" + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/gcs_handler.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/gcs_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..b0600bb6a2f8277e48224525099c02aecb01ab8a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/gcs_handler.py @@ -0,0 +1,224 @@ +"""GCS storage handler.""" + +from __future__ import annotations + +import time +from pathlib import PurePosixPath +from typing import TYPE_CHECKING, Sequence +from urllib.parse import ParseResult, urlparse + +from wandb import util +from wandb.errors.term import termlog +from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache +from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry +from wandb.sdk.artifacts.storage_handler import DEFAULT_MAX_OBJECTS, StorageHandler +from wandb.sdk.lib.hashutil import ETag +from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr + +if TYPE_CHECKING: + import google.cloud.storage as gcs_module # type: ignore + + from wandb.sdk.artifacts.artifact import Artifact + + +class _GCSIsADirectoryError(Exception): + """Raised when we try to download a GCS folder.""" + + +class GCSHandler(StorageHandler): + _client: gcs_module.client.Client | None + + def __init__(self, scheme: str | None = None) -> None: + self._scheme = scheme or "gs" + self._client = None + self._cache = get_artifact_file_cache() + + def can_handle(self, parsed_url: ParseResult) -> bool: + return parsed_url.scheme == self._scheme + + def init_gcs(self) -> gcs_module.client.Client: + if self._client is not None: + return self._client + storage = util.get_module( + "google.cloud.storage", + required="gs:// references requires the google-cloud-storage library, run pip install wandb[gcp]", + ) + self._client = storage.Client() + return self._client + + def _parse_uri(self, uri: str) -> tuple[str, str, str | None]: + url = urlparse(uri) + bucket = url.netloc + key = url.path[1:] + version = url.fragment if url.fragment else None + return bucket, key, version + + def load_path( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + ) -> URIStr | FilePathStr: + assert manifest_entry.ref is not None + if not local: + return manifest_entry.ref + + path, hit, cache_open = self._cache.check_etag_obj_path( + url=URIStr(manifest_entry.ref), + etag=ETag(manifest_entry.digest), + size=manifest_entry.size if manifest_entry.size is not None else 0, + ) + if hit: + return path + + self.init_gcs() + assert self._client is not None # mypy: unwraps optionality + assert manifest_entry.ref is not None + bucket, key, _ = self._parse_uri(manifest_entry.ref) + version = manifest_entry.extra.get("versionID") + + if self._is_dir(manifest_entry): + raise _GCSIsADirectoryError( + f"Unable to download GCS folder {manifest_entry.ref!r}, skipping" + ) + + obj = None + # First attempt to get the generation specified, this will return None if versioning is not enabled + if version is not None: + obj = self._client.bucket(bucket).get_blob(key, generation=version) + + if obj is None: + # Object versioning is disabled on the bucket, so just get + # the latest version and make sure the MD5 matches. + obj = self._client.bucket(bucket).get_blob(key) + if obj is None: + raise ValueError( + f"Unable to download object {manifest_entry.ref} with generation {version}" + ) + if obj.etag != manifest_entry.digest: + raise ValueError( + f"Digest mismatch for object {manifest_entry.ref}: " + f"expected {manifest_entry.digest} but found {obj.etag}" + ) + + with cache_open(mode="wb") as f: + obj.download_to_file(f) + return path + + def store_path( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: StrPath | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + self.init_gcs() + assert self._client is not None # mypy: unwraps optionality + + # After parsing any query params / fragments for additional context, + # such as version identifiers, pare down the path to just the bucket + # and key. + bucket, key, version = self._parse_uri(path) + path = URIStr(f"{self._scheme}://{bucket}/{key}") + max_objects = max_objects or DEFAULT_MAX_OBJECTS + + if not checksum: + return [ArtifactManifestEntry(path=name or key, ref=path, digest=path)] + + start_time = None + obj = self._client.bucket(bucket).get_blob(key, generation=version) + if obj is None and version is not None: + raise ValueError(f"Object does not exist: {path}#{version}") + multi = obj is None + if multi: + start_time = time.time() + termlog( + f'Generating checksum for up to {max_objects} objects with prefix "{key}"... ', + newline=False, + ) + objects = self._client.bucket(bucket).list_blobs( + prefix=key, max_results=max_objects + ) + else: + objects = [obj] + + entries = [ + self._entry_from_obj(obj, path, name, prefix=key, multi=multi) + for obj in objects + if not obj.name.endswith("/") + ] + if start_time is not None: + termlog("Done. %.1fs" % (time.time() - start_time), prefix=False) + if len(entries) > max_objects: + raise ValueError( + f"Exceeded {max_objects} objects tracked, pass max_objects to add_reference" + ) + return entries + + def _entry_from_obj( + self, + obj: gcs_module.blob.Blob, + path: str, + name: StrPath | None = None, + prefix: str = "", + multi: bool = False, + ) -> ArtifactManifestEntry: + """Create an ArtifactManifestEntry from a GCS object. + + Args: + obj: The GCS object + path: The GCS-style path (e.g.: "gs://bucket/file.txt") + name: The user assigned name, or None if not specified + prefix: The prefix to add (will be the same as `path` for directories) + multi: Whether or not this is a multi-object add. + """ + bucket, key, _ = self._parse_uri(path) + + # Always use posix paths, since that's what S3 uses. + posix_key = PurePosixPath(obj.name) # the bucket key + posix_path = PurePosixPath(bucket) / PurePosixPath( + key + ) # the path, with the scheme stripped + posix_prefix = PurePosixPath(prefix) # the prefix, if adding a prefix + posix_name = PurePosixPath(name or "") + posix_ref = posix_path + + if name is None: + # We're adding a directory (prefix), so calculate a relative path. + if str(posix_prefix) in str(posix_key) and posix_prefix != posix_key: + posix_name = posix_key.relative_to(posix_prefix) + posix_ref = posix_path / posix_name + else: + posix_name = PurePosixPath(posix_key.name) + posix_ref = posix_path + elif multi: + # We're adding a directory with a name override. + relpath = posix_key.relative_to(posix_prefix) + posix_name = posix_name / relpath + posix_ref = posix_path / relpath + return ArtifactManifestEntry( + path=posix_name, + ref=URIStr(f"{self._scheme}://{str(posix_ref)}"), + digest=obj.etag, + size=obj.size, + extra={"versionID": obj.generation}, + ) + + def _is_dir( + self, + manifest_entry: ArtifactManifestEntry, + ) -> bool: + assert self._client is not None + assert manifest_entry.ref is not None + bucket, key, _ = self._parse_uri(manifest_entry.ref) + bucket_obj = self._client.bucket(bucket) + # A gcs bucket key should end with a forward slash on gcloud, but + # we save these refs without the forward slash in the manifest entry + # so we check the size and extension, make sure its not referring to + # an actual file with this reference, and that the ref with the slash + # exists on gcloud + return key.endswith("/") or ( + not (manifest_entry.size or PurePosixPath(key).suffix) + and bucket_obj.get_blob(key) is None + and bucket_obj.get_blob(f"{key}/") is not None + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/http_handler.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/http_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..a06af2683439925ded44e0ebf4b4bcd0dca1a018 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/http_handler.py @@ -0,0 +1,114 @@ +"""HTTP storage handler.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Sequence +from urllib.parse import ParseResult + +from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache +from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry +from wandb.sdk.artifacts.storage_handler import StorageHandler +from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings +from wandb.sdk.lib.hashutil import ETag +from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr + +if TYPE_CHECKING: + import requests + from requests.structures import CaseInsensitiveDict + + from wandb.sdk.artifacts.artifact import Artifact + + +class HTTPHandler(StorageHandler): + def __init__(self, session: requests.Session, scheme: str | None = None) -> None: + self._scheme = scheme or "http" + self._cache = get_artifact_file_cache() + self._session = session + + def can_handle(self, parsed_url: ParseResult) -> bool: + return parsed_url.scheme == self._scheme + + def load_path( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + ) -> URIStr | FilePathStr: + if not local: + assert manifest_entry.ref is not None + return manifest_entry.ref + + assert manifest_entry.ref is not None + + path, hit, cache_open = self._cache.check_etag_obj_path( + URIStr(manifest_entry.ref), + ETag(manifest_entry.digest), # TODO(spencerpearson): unsafe cast + manifest_entry.size if manifest_entry.size is not None else 0, + ) + if hit: + return path + + response = self._session.get( + manifest_entry.ref, + stream=True, + cookies=_thread_local_api_settings.cookies, + headers=_thread_local_api_settings.headers, + ) + response.raise_for_status() + + digest: ETag | FilePathStr | URIStr | None + digest, size, extra = self._entry_from_headers(response.headers) + digest = digest or manifest_entry.ref + if manifest_entry.digest != digest: + raise ValueError( + f"Digest mismatch for url {manifest_entry.ref}: expected {manifest_entry.digest} but found {digest}" + ) + + with cache_open(mode="wb") as file: + for data in response.iter_content(chunk_size=16 * 1024): + file.write(data) + return path + + def store_path( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: StrPath | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + name = name or os.path.basename(path) + if not checksum: + return [ArtifactManifestEntry(path=name, ref=path, digest=path)] + + with self._session.get( + path, + stream=True, + cookies=_thread_local_api_settings.cookies, + headers=_thread_local_api_settings.headers, + ) as response: + response.raise_for_status() + digest: ETag | FilePathStr | URIStr | None + digest, size, extra = self._entry_from_headers(response.headers) + digest = digest or path + return [ + ArtifactManifestEntry( + path=name, ref=path, digest=digest, size=size, extra=extra + ) + ] + + def _entry_from_headers( + self, headers: CaseInsensitiveDict + ) -> tuple[ETag | None, int | None, dict[str, str]]: + response_headers = {k.lower(): v for k, v in headers.items()} + size = None + if response_headers.get("content-length", None): + size = int(response_headers["content-length"]) + + digest = response_headers.get("etag", None) + extra = {} + if digest: + extra["etag"] = digest + if digest and digest[:1] == '"' and digest[-1:] == '"': + digest = digest[1:-1] # trim leading and trailing quotes around etag + return digest, size, extra diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/local_file_handler.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/local_file_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..2e6e508fc5cf8abdf96ea5b2ea3cd5801196d1e3 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/local_file_handler.py @@ -0,0 +1,142 @@ +"""Local file storage handler.""" + +from __future__ import annotations + +import os +import shutil +import time +from pathlib import Path +from typing import TYPE_CHECKING, Sequence +from urllib.parse import ParseResult + +from wandb import util +from wandb.errors.term import termlog +from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache +from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry +from wandb.sdk.artifacts.storage_handler import DEFAULT_MAX_OBJECTS, StorageHandler +from wandb.sdk.lib import filesystem +from wandb.sdk.lib.hashutil import B64MD5, md5_file_b64, md5_string +from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr + +if TYPE_CHECKING: + from wandb.sdk.artifacts.artifact import Artifact + + +class LocalFileHandler(StorageHandler): + """Handles file:// references.""" + + def __init__(self, scheme: str | None = None) -> None: + """Track files or directories on a local filesystem. + + Expand directories to create an entry for each file contained. + """ + self._scheme = scheme or "file" + self._cache = get_artifact_file_cache() + + def can_handle(self, parsed_url: ParseResult) -> bool: + return parsed_url.scheme == self._scheme + + def load_path( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + ) -> URIStr | FilePathStr: + if manifest_entry.ref is None: + raise ValueError(f"Cannot add path with no ref: {manifest_entry.path}") + local_path = util.local_file_uri_to_path(str(manifest_entry.ref)) + if not os.path.exists(local_path): + raise ValueError( + f"Local file reference: Failed to find file at path {local_path}" + ) + + path, hit, cache_open = self._cache.check_md5_obj_path( + B64MD5(manifest_entry.digest), # TODO(spencerpearson): unsafe cast + manifest_entry.size if manifest_entry.size is not None else 0, + ) + if hit: + return path + + md5 = md5_file_b64(local_path) + if md5 != manifest_entry.digest: + raise ValueError( + f"Local file reference: Digest mismatch for path {local_path}: expected {manifest_entry.digest} but found {md5}" + ) + + filesystem.mkdir_exists_ok(os.path.dirname(path)) + + with cache_open() as f: + shutil.copy(local_path, f.name) + return path + + def store_path( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: StrPath | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + local_path = util.local_file_uri_to_path(path) + max_objects = max_objects or DEFAULT_MAX_OBJECTS + # We have a single file or directory + # Note, we follow symlinks for files contained within the directory + entries = [] + + # If checksum=False, the file's hash should only + # depend on its absolute path/URI, not its contents + + # Closure func for calculating the file hash from its path + def md5(path: str) -> B64MD5: + return ( + md5_file_b64(path) + if checksum + else md5_string(Path(path).resolve().as_uri()) + ) + + if os.path.isdir(local_path): + i = 0 + start_time = time.time() + if checksum: + termlog( + f'Generating checksum for up to {max_objects} files in "{local_path}"... ', + newline=False, + ) + for root, _, files in os.walk(local_path): + for sub_path in files: + i += 1 + if i > max_objects: + raise ValueError( + f"Exceeded {max_objects} objects tracked, pass max_objects to add_reference" + ) + physical_path = os.path.join(root, sub_path) + # TODO(spencerpearson): this is not a "logical path" in the sense that + # `LogicalPath` returns a "logical path"; it's a relative path + # **on the local filesystem**. + file_path = os.path.relpath(physical_path, start=local_path) + if name is not None: + artifact_path = os.path.join(name, file_path) + else: + artifact_path = file_path + + entry = ArtifactManifestEntry( + path=artifact_path, + ref=FilePathStr(os.path.join(path, file_path)), + size=os.path.getsize(physical_path), + digest=md5(physical_path), + ) + entries.append(entry) + if checksum: + termlog("Done. %.1fs" % (time.time() - start_time), prefix=False) + elif os.path.isfile(local_path): + name = name or os.path.basename(local_path) + entry = ArtifactManifestEntry( + path=name, + ref=path, + size=os.path.getsize(local_path), + digest=md5(local_path), + ) + entries.append(entry) + else: + # TODO: update error message if we don't allow directories. + raise ValueError(f'Path "{path}" must be a valid file or directory path') + return entries diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/multi_handler.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/multi_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..a44f102ff94cc6b3b702e5154f3df6250dd4e5c8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/multi_handler.py @@ -0,0 +1,56 @@ +"""Multi storage handler.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Sequence +from urllib.parse import urlparse + +from wandb.sdk.artifacts.storage_handler import StorageHandler +from wandb.sdk.lib.paths import FilePathStr, URIStr + +if TYPE_CHECKING: + from wandb.sdk.artifacts.artifact import Artifact + from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry + + +class MultiHandler(StorageHandler): + _handlers: list[StorageHandler] + + def __init__( + self, + handlers: list[StorageHandler] | None = None, + default_handler: StorageHandler | None = None, + ) -> None: + self._handlers = handlers or [] + self._default_handler = default_handler + + def _get_handler(self, url: FilePathStr | URIStr) -> StorageHandler: + parsed_url = urlparse(url) + for handler in self._handlers: + if handler.can_handle(parsed_url): + return handler + if self._default_handler is not None: + return self._default_handler + raise ValueError(f'No storage handler registered for url "{url!s}"') + + def load_path( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + ) -> URIStr | FilePathStr: + assert manifest_entry.ref is not None + handler = self._get_handler(manifest_entry.ref) + return handler.load_path(manifest_entry, local=local) + + def store_path( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: str | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + handler = self._get_handler(path) + return handler.store_path( + artifact, path, name=name, checksum=checksum, max_objects=max_objects + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/s3_handler.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/s3_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..a9ea22225e477b449556e45c435ccad16dfd0318 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/s3_handler.py @@ -0,0 +1,339 @@ +"""S3 storage handler.""" + +from __future__ import annotations + +import os +import re +import time +from pathlib import PurePosixPath +from typing import TYPE_CHECKING, Sequence +from urllib.parse import parse_qsl, urlparse + +from wandb import util +from wandb.errors import CommError +from wandb.errors.term import termlog +from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache +from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry +from wandb.sdk.artifacts.storage_handler import DEFAULT_MAX_OBJECTS, StorageHandler +from wandb.sdk.lib.hashutil import ETag +from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr + +if TYPE_CHECKING: + from urllib.parse import ParseResult + + # We could probably use https://pypi.org/project/boto3-stubs/ or something + # instead of `type:ignore`ing these boto imports, but it's nontrivial: + # for some reason, despite being actively maintained as of 2022-09-30, + # the latest release of boto3-stubs doesn't include all the features we use. + import boto3 # type: ignore + import boto3.resources.base # type: ignore + import boto3.s3 # type: ignore + import boto3.session # type: ignore + + from wandb.sdk.artifacts.artifact import Artifact + + +class S3Handler(StorageHandler): + _s3: boto3.resources.base.ServiceResource | None + _scheme: str + + def __init__(self, scheme: str | None = None) -> None: + self._scheme = scheme or "s3" + self._s3 = None + self._cache = get_artifact_file_cache() + + def can_handle(self, parsed_url: ParseResult) -> bool: + return parsed_url.scheme == self._scheme + + def init_boto(self) -> boto3.resources.base.ServiceResource: + if self._s3 is not None: + return self._s3 + boto: boto3 = util.get_module( + "boto3", + required="s3:// references requires the boto3 library, run pip install wandb[aws]", + lazy=False, + ) + + from botocore.client import Config # type: ignore + + s3_endpoint = os.getenv("AWS_S3_ENDPOINT_URL") + config = ( + Config(s3={"addressing_style": "virtual"}) + if s3_endpoint and self._is_coreweave_endpoint(s3_endpoint) + else None + ) + self._s3 = boto.session.Session().resource( + "s3", + endpoint_url=s3_endpoint, + region_name=os.getenv("AWS_REGION"), + config=config, + ) + self._botocore = util.get_module("botocore") + return self._s3 + + def _parse_uri(self, uri: str) -> tuple[str, str, str | None]: + url = urlparse(uri) + query = dict(parse_qsl(url.query)) + + bucket = url.netloc + key = url.path[1:] # strip leading slash + version = query.get("versionId") + + return bucket, key, version + + def load_path( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + ) -> URIStr | FilePathStr: + if not local: + assert manifest_entry.ref is not None + return manifest_entry.ref + + assert manifest_entry.ref is not None + + path, hit, cache_open = self._cache.check_etag_obj_path( + URIStr(manifest_entry.ref), + ETag(manifest_entry.digest), # TODO(spencerpearson): unsafe cast + manifest_entry.size if manifest_entry.size is not None else 0, + ) + if hit: + return path + + self.init_boto() + assert self._s3 is not None # mypy: unwraps optionality + bucket, key, _ = self._parse_uri(manifest_entry.ref) + version = manifest_entry.extra.get("versionID") + + extra_args = {} + if version: + obj_version = self._s3.ObjectVersion(bucket, key, version) + extra_args["VersionId"] = version + obj = obj_version.Object() + else: + obj = self._s3.Object(bucket, key) + + try: + etag = ( + obj_version.head()["ETag"][1:-1] # escape leading and trailing + if version + else self._etag_from_obj(obj) + ) + except self._botocore.exceptions.ClientError as e: + if e.response["Error"]["Code"] == "404": + raise FileNotFoundError( + f"Unable to find {manifest_entry.path} at s3://{bucket}/{key}" + ) from e + raise + + if etag != manifest_entry.digest: + # Try to match the etag with some other version. + if version: + raise ValueError( + f"Digest mismatch for object {manifest_entry.ref} with version {version}: expected {manifest_entry.digest} but found {etag}" + ) + obj = None + object_versions = self._s3.Bucket(bucket).object_versions.filter(Prefix=key) + for object_version in object_versions: + if manifest_entry.extra.get("etag") == self._etag_from_obj( + object_version + ): + obj = object_version.Object() + extra_args["VersionId"] = object_version.version_id + break + if obj is None: + raise FileNotFoundError( + "Couldn't find object version for {}/{} matching etag {}".format( + bucket, key, manifest_entry.extra.get("etag") + ) + ) + + with cache_open(mode="wb") as f: + obj.download_fileobj(f, ExtraArgs=extra_args) + return path + + def store_path( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: StrPath | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + self.init_boto() + assert self._s3 is not None # mypy: unwraps optionality + + # The passed in path might have query string parameters. + # We only need to care about a subset, like version, when + # parsing. Once we have that, we can store the rest of the + # metadata in the artifact entry itself. + bucket, key, version = self._parse_uri(path) + path = URIStr(f"{self._scheme}://{bucket}/{key}") + + max_objects = max_objects or DEFAULT_MAX_OBJECTS + if not checksum: + entry_path = name or (key if key != "" else bucket) + return [ArtifactManifestEntry(path=entry_path, ref=path, digest=path)] + + # If an explicit version is specified, use that. Otherwise, use the head version. + objs = ( + [self._s3.ObjectVersion(bucket, key, version).Object()] + if version + else [self._s3.Object(bucket, key)] + ) + start_time = None + multi = False + if key != "": + try: + objs[0].load() + # S3 doesn't have real folders, however there are cases where the folder key has a valid file which will not + # trigger a recursive upload. + # we should check the object's metadata says it is a directory and do a multi file upload if it is + if "x-directory" in objs[0].content_type: + multi = True + except self._botocore.exceptions.ClientError as e: + if e.response["Error"]["Code"] == "404": + multi = True + else: + raise CommError( + f"Unable to connect to S3 ({e.response['Error']['Code']}): " + f"{e.response['Error']['Message']}. Check that your " + "authentication credentials are valid and that your region is " + "set correctly." + ) + else: + multi = True + + if multi: + start_time = time.time() + termlog( + f'Generating checksum for up to {max_objects} objects in "{bucket}/{key}"... ', + newline=False, + ) + if key != "": + objs = ( + self._s3.Bucket(bucket) + .objects.filter(Prefix=key) + .limit(max_objects) + ) + else: + objs = self._s3.Bucket(bucket).objects.limit(max_objects) + # Weird iterator scoping makes us assign this to a local function + size = self._size_from_obj + entries = [ + self._entry_from_obj(obj, path, name, prefix=key, multi=multi) + for obj in objs + if size(obj) > 0 + ] + if start_time is not None: + termlog("Done. %.1fs" % (time.time() - start_time), prefix=False) + if len(entries) > max_objects: + raise ValueError( + f"Exceeded {max_objects} objects tracked, pass max_objects to add_reference" + ) + return entries + + def _size_from_obj(self, obj: boto3.s3.Object | boto3.s3.ObjectSummary) -> int: + # ObjectSummary has size, Object has content_length + size: int + if hasattr(obj, "size"): + size = obj.size + else: + size = obj.content_length + return size + + def _entry_from_obj( + self, + obj: boto3.s3.Object | boto3.s3.ObjectSummary, + path: str, + name: StrPath | None = None, + prefix: str = "", + multi: bool = False, + ) -> ArtifactManifestEntry: + """Create an ArtifactManifestEntry from an S3 object. + + Args: + obj: The S3 object + path: The S3-style path (e.g.: "s3://bucket/file.txt") + name: The user assigned name, or None if not specified + prefix: The prefix to add (will be the same as `path` for directories) + multi: Whether or not this is a multi-object add. + """ + bucket, key, _ = self._parse_uri(path) + + # Always use posix paths, since that's what S3 uses. + posix_key = PurePosixPath(obj.key) # the bucket key + posix_path = PurePosixPath(bucket) / key # the path, with the scheme stripped + posix_prefix = PurePosixPath(prefix) # the prefix, if adding a prefix + posix_name = PurePosixPath(name or "") + posix_ref = posix_path + + if name is None: + # We're adding a directory (prefix), so calculate a relative path. + if str(posix_prefix) in str(posix_key) and posix_prefix != posix_key: + posix_name = posix_key.relative_to(posix_prefix) + posix_ref = posix_path / posix_name + else: + posix_name = PurePosixPath(posix_key.name) + posix_ref = posix_path + elif multi: + # We're adding a directory with a name override. + relpath = posix_key.relative_to(posix_prefix) + posix_name = posix_name / relpath + posix_ref = posix_path / relpath + return ArtifactManifestEntry( + path=posix_name, + ref=URIStr(f"{self._scheme}://{str(posix_ref)}"), + digest=ETag(self._etag_from_obj(obj)), + size=self._size_from_obj(obj), + extra=self._extra_from_obj(obj), + ) + + @staticmethod + def _etag_from_obj(obj: boto3.s3.Object | boto3.s3.ObjectSummary) -> ETag: + etag: ETag + etag = obj.e_tag[1:-1] # escape leading and trailing quote + return etag + + def _extra_from_obj( + self, obj: boto3.s3.Object | boto3.s3.ObjectSummary + ) -> dict[str, str]: + extra = { + "etag": obj.e_tag[1:-1], # escape leading and trailing quote + } + if not hasattr(obj, "version_id"): + # Convert ObjectSummary to Object to get the version_id. + obj = self._s3.Object(obj.bucket_name, obj.key) # type: ignore[union-attr] + if hasattr(obj, "version_id") and obj.version_id and obj.version_id != "null": + extra["versionID"] = obj.version_id + return extra + + _CW_LEGACY_NETLOC_REGEX: re.Pattern[str] = re.compile( + r""" + # accelerated endpoints like "accel-object..coreweave.com" + accel-object\.[a-z0-9-]+\.coreweave\.com + | + # URLs like "object..coreweave.com" + object\.[a-z0-9-]+\.coreweave\.com + """, + flags=re.VERBOSE, + ) + + def _is_coreweave_endpoint(self, endpoint_url: str) -> bool: + if not (url := endpoint_url.strip().rstrip("/")): + return False + + # Only http://cwlota.com is supported using HTTP + if url == "http://cwlota.com": + return True + + # Enforce HTTPS otherwise + https_url = url if url.startswith("https://") else f"https://{url}" + netloc = urlparse(https_url).netloc + return bool( + # Match for https://cwobject.com + (netloc == "cwobject.com") + or + # Check for legacy endpoints + self._CW_LEGACY_NETLOC_REGEX.fullmatch(netloc) + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/tracking_handler.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/tracking_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..d8efedd48f45f9dd729c871bb3a3f409bc7a8eb4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/tracking_handler.py @@ -0,0 +1,68 @@ +"""Tracking storage handler.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Sequence +from urllib.parse import urlparse + +from wandb.errors.term import termwarn +from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry +from wandb.sdk.artifacts.storage_handler import StorageHandler +from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr + +if TYPE_CHECKING: + from urllib.parse import ParseResult + + from wandb.sdk.artifacts.artifact import Artifact + + +class TrackingHandler(StorageHandler): + def __init__(self, scheme: str | None = None) -> None: + """Track paths with no modification or special processing. + + Useful when paths being tracked are on file systems mounted at a standardized + location. + + For example, if the data to track is located on an NFS share mounted on + `/data`, then it is sufficient to just track the paths. + """ + self._scheme = scheme or "" + + def can_handle(self, parsed_url: ParseResult) -> bool: + return parsed_url.scheme == self._scheme + + def load_path( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + ) -> URIStr | FilePathStr: + if local: + # Likely a user error. The tracking handler is + # oblivious to the underlying paths, so it has + # no way of actually loading it. + url = urlparse(manifest_entry.ref) + raise ValueError( + f"Cannot download file at path {str(manifest_entry.ref)}, scheme {str(url.scheme)} not recognized" + ) + # TODO(spencerpearson): should this go through util.to_native_slash_path + # instead of just getting typecast? + return FilePathStr(manifest_entry.path) + + def store_path( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: StrPath | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + url = urlparse(path) + if name is None: + raise ValueError( + f'You must pass name="" when tracking references with unknown schemes. ref: {path}' + ) + termwarn( + f"Artifact references with unsupported schemes cannot be checksummed: {path}" + ) + name = name or url.path[1:] # strip leading slash + return [ArtifactManifestEntry(path=name, ref=path, digest=path)] diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/wb_artifact_handler.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/wb_artifact_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..16e490dbc54b7df63e326e99f2acaa2de415621d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/wb_artifact_handler.py @@ -0,0 +1,131 @@ +"""WB artifact storage handler.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Sequence +from urllib.parse import urlparse + +import wandb +from wandb import util +from wandb.apis import PublicApi +from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache +from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry +from wandb.sdk.artifacts.storage_handler import StorageHandler +from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, hex_to_b64_id +from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr + +if TYPE_CHECKING: + from urllib.parse import ParseResult + + from wandb.sdk.artifacts.artifact import Artifact + + +class WBArtifactHandler(StorageHandler): + """Handles loading and storing Artifact reference-type files.""" + + _client: PublicApi | None + + def __init__(self) -> None: + self._scheme = "wandb-artifact" + self._cache = get_artifact_file_cache() + self._client = None + + def can_handle(self, parsed_url: ParseResult) -> bool: + return parsed_url.scheme == self._scheme + + @property + def client(self) -> PublicApi: + if self._client is None: + self._client = PublicApi() + return self._client + + def load_path( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + ) -> URIStr | FilePathStr: + """Load the file in the specified artifact given its corresponding entry. + + Download the referenced artifact; create and return a new symlink to the caller. + + Args: + manifest_entry (ArtifactManifestEntry): The index entry to load + + Returns: + (os.PathLike): A path to the file represented by `index_entry` + """ + # We don't check for cache hits here. Since we have 0 for size (since this + # is a cross-artifact reference which and we've made the choice to store 0 + # in the size field), we can't confirm if the file is complete. So we just + # rely on the dep_artifact entry's download() method to do its own cache + # check. + + # Parse the reference path and download the artifact if needed + artifact_id = util.host_from_path(manifest_entry.ref) + artifact_file_path = util.uri_from_path(manifest_entry.ref) + + dep_artifact = wandb.Artifact._from_id( + hex_to_b64_id(artifact_id), self.client.client + ) + assert dep_artifact is not None + link_target_path: URIStr | FilePathStr + if local: + link_target_path = dep_artifact.get_entry(artifact_file_path).download() + else: + link_target_path = dep_artifact.get_entry(artifact_file_path).ref_target() + + return link_target_path + + def store_path( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: StrPath | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + """Store the file or directory at the given path into the specified artifact. + + Recursively resolves the reference until the result is a concrete asset. + + Args: + artifact: The artifact doing the storing path (str): The path to store name + (str): If specified, the logical name that should map to `path` + + Returns: + (list[ArtifactManifestEntry]): A list of manifest entries to store within + the artifact + """ + # Recursively resolve the reference until a concrete asset is found + # TODO: Consider resolving server-side for performance improvements. + iter_path: URIStr | FilePathStr | None = path + while iter_path is not None and urlparse(iter_path).scheme == self._scheme: + artifact_id = util.host_from_path(iter_path) + artifact_file_path = util.uri_from_path(iter_path) + target_artifact = wandb.Artifact._from_id( + hex_to_b64_id(artifact_id), self.client.client + ) + assert target_artifact is not None + + entry = target_artifact.manifest.get_entry_by_path(artifact_file_path) + assert entry is not None + iter_path = entry.ref + + # Create the path reference + assert target_artifact is not None + assert target_artifact.id is not None + path = URIStr( + f"{self._scheme}://{b64_to_hex_id(B64MD5(target_artifact.id))}/{artifact_file_path}" + ) + + # Return the new entry + assert entry is not None + return [ + ArtifactManifestEntry( + path=name or os.path.basename(path), + ref=path, + size=0, + digest=entry.digest, + ) + ] diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/wb_local_artifact_handler.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/wb_local_artifact_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..6a2e5a97f3e23a557fcc178df0062afaad1cfdb8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_handlers/wb_local_artifact_handler.py @@ -0,0 +1,74 @@ +"""WB local artifact storage handler.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Sequence + +import wandb +from wandb import util +from wandb.sdk.artifacts.artifact_instance_cache import artifact_instance_cache +from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry +from wandb.sdk.artifacts.storage_handler import StorageHandler +from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr + +if TYPE_CHECKING: + from urllib.parse import ParseResult + + from wandb.sdk.artifacts.artifact import Artifact + + +class WBLocalArtifactHandler(StorageHandler): + """Handles loading and storing Artifact reference-type files.""" + + def __init__(self) -> None: + self._scheme = "wandb-client-artifact" + + def can_handle(self, parsed_url: ParseResult) -> bool: + return parsed_url.scheme == self._scheme + + def load_path( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + ) -> URIStr | FilePathStr: + raise NotImplementedError( + "Should not be loading a path for an artifact entry with unresolved client id." + ) + + def store_path( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: StrPath | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + """Store the file or directory at the given path within the specified artifact. + + Args: + artifact: The artifact doing the storing + path (str): The path to store + name (str): If specified, the logical name that should map to `path` + + Returns: + (list[ArtifactManifestEntry]): A list of manifest entries to store within the artifact + """ + client_id = util.host_from_path(path) + target_path = util.uri_from_path(path) + target_artifact = artifact_instance_cache.get(client_id) + if not isinstance(target_artifact, wandb.Artifact): + raise TypeError("Artifact passed to store_path() must be a wandb.Artifact.") + target_entry = target_artifact.manifest.entries[target_path] # type: ignore + if target_entry is None: + raise RuntimeError("Local entry not found - invalid reference") + + # Return the new entry + return [ + ArtifactManifestEntry( + path=name or os.path.basename(path), + ref=path, + size=0, + digest=target_entry.digest, + ) + ] diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_layout.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..7fb20c522ac667a83be3ad55abc74b32a453650f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_layout.py @@ -0,0 +1,8 @@ +"""Storage layout.""" + +from __future__ import annotations + + +class StorageLayout: + V1 = "V1" + V2 = "V2" diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__init__.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4d409bcf913dfde9b84f6cc53d3c980a2cbe02e2 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__init__.py @@ -0,0 +1,4 @@ +from wandb.sdk.artifacts.storage_policies.register import WANDB_STORAGE_POLICY +from wandb.sdk.artifacts.storage_policies.wandb_storage_policy import WandbStoragePolicy + +__all__ = ["WANDB_STORAGE_POLICY", "WandbStoragePolicy"] diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e5a671b0d5189a558b78a45d15abcfe2baf18195 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__pycache__/register.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__pycache__/register.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c6b84c0eb058c5843e456ba3f1d691298edf9a9b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__pycache__/register.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__pycache__/wandb_storage_policy.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__pycache__/wandb_storage_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..beba0c07ff1308cfe113b9a064bd00371a3b214e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/__pycache__/wandb_storage_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/register.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/register.py new file mode 100644 index 0000000000000000000000000000000000000000..84542d8a70fca03538985fa8aa9c6ef7ea86c2b5 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/register.py @@ -0,0 +1 @@ +WANDB_STORAGE_POLICY = "wandb-storage-policy-v1" diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..6bfd740e48bfe2f1eb9023ed3998e55fa3508ddf --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py @@ -0,0 +1,580 @@ +"""WandB storage policy.""" + +from __future__ import annotations + +import concurrent.futures +import functools +import hashlib +import logging +import math +import os +import queue +import shutil +import threading +from collections import deque +from typing import IO, TYPE_CHECKING, Any, NamedTuple, Sequence +from urllib.parse import quote + +import requests +import urllib3 + +from wandb import env +from wandb.errors.term import termwarn +from wandb.proto.wandb_internal_pb2 import ServerFeature +from wandb.sdk.artifacts.artifact_file_cache import ( + ArtifactFileCache, + Opener, + get_artifact_file_cache, +) +from wandb.sdk.artifacts.staging import get_staging_dir +from wandb.sdk.artifacts.storage_handlers.azure_handler import AzureHandler +from wandb.sdk.artifacts.storage_handlers.gcs_handler import GCSHandler +from wandb.sdk.artifacts.storage_handlers.http_handler import HTTPHandler +from wandb.sdk.artifacts.storage_handlers.local_file_handler import LocalFileHandler +from wandb.sdk.artifacts.storage_handlers.multi_handler import MultiHandler +from wandb.sdk.artifacts.storage_handlers.s3_handler import S3Handler +from wandb.sdk.artifacts.storage_handlers.tracking_handler import TrackingHandler +from wandb.sdk.artifacts.storage_handlers.wb_artifact_handler import WBArtifactHandler +from wandb.sdk.artifacts.storage_handlers.wb_local_artifact_handler import ( + WBLocalArtifactHandler, +) +from wandb.sdk.artifacts.storage_layout import StorageLayout +from wandb.sdk.artifacts.storage_policies.register import WANDB_STORAGE_POLICY +from wandb.sdk.artifacts.storage_policy import StoragePolicy +from wandb.sdk.internal.internal_api import Api as InternalApi +from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings +from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, hex_to_b64_id +from wandb.sdk.lib.paths import FilePathStr, URIStr + +if TYPE_CHECKING: + from wandb.filesync.step_prepare import StepPrepare + from wandb.sdk.artifacts.artifact import Artifact + from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry + from wandb.sdk.internal import progress + +# Sleep length: 0, 2, 4, 8, 16, 32, 64, 120, 120, 120, 120, 120, 120, 120, 120, 120 +# seconds, i.e. a total of 20min 6s. +_REQUEST_RETRY_STRATEGY = urllib3.util.retry.Retry( + backoff_factor=1, + total=16, + status_forcelist=(308, 408, 409, 429, 500, 502, 503, 504), +) +_REQUEST_POOL_CONNECTIONS = 64 +_REQUEST_POOL_MAXSIZE = 64 + +# AWS S3 max upload parts without having to make additional requests for extra parts +S3_MAX_PART_NUMBERS = 1000 +S3_MIN_MULTI_UPLOAD_SIZE = 2 * 1024**3 +S3_MAX_MULTI_UPLOAD_SIZE = 5 * 1024**4 + + +# Minimum size to switch to multipart download, same as upload, 2GB. +_MULTIPART_DOWNLOAD_SIZE = S3_MIN_MULTI_UPLOAD_SIZE +# Multipart download part size is same as multpart upload size, which is hard coded to 100MB. +# https://github.com/wandb/wandb/blob/7b2a13cb8efcd553317167b823c8e52d8c3f7c4e/core/pkg/artifacts/saver.go#L496 +# https://docs.aws.amazon.com/AmazonS3/latest/userguide/optimizing-performance-guidelines.html#optimizing-performance-guidelines-get-range +_DOWNLOAD_PART_SIZE_BYTES = 100 * 1024 * 1024 +# Chunk size for reading http response and writing to disk. 1MB. +_HTTP_RES_CHUNK_SIZE_BYTES = 1 * 1024 * 1024 +# Signal end of _ChunkQueue, consumer (file writer) should stop after getting this item. +# NOTE: it should only be used for multithread executor, it does notwork for multiprocess executor. +# multipart download is using the executor from artifact.download() which is a multithread executor. +_CHUNK_QUEUE_SENTINEL = object() + +logger = logging.getLogger(__name__) + + +class _ChunkContent(NamedTuple): + offset: int + data: bytes + + +class WandbStoragePolicy(StoragePolicy): + @classmethod + def name(cls) -> str: + return WANDB_STORAGE_POLICY + + @classmethod + def from_config( + cls, config: dict, api: InternalApi | None = None + ) -> WandbStoragePolicy: + return cls(config=config, api=api) + + def __init__( + self, + config: dict | None = None, + cache: ArtifactFileCache | None = None, + api: InternalApi | None = None, + ) -> None: + self._cache = cache or get_artifact_file_cache() + self._config = config or {} + self._session = requests.Session() + adapter = requests.adapters.HTTPAdapter( + max_retries=_REQUEST_RETRY_STRATEGY, + pool_connections=_REQUEST_POOL_CONNECTIONS, + pool_maxsize=_REQUEST_POOL_MAXSIZE, + ) + self._session.mount("http://", adapter) + self._session.mount("https://", adapter) + + s3 = S3Handler() + gcs = GCSHandler() + azure = AzureHandler() + http = HTTPHandler(self._session) + https = HTTPHandler(self._session, scheme="https") + artifact = WBArtifactHandler() + local_artifact = WBLocalArtifactHandler() + file_handler = LocalFileHandler() + + self._api = api or InternalApi() + self._handler = MultiHandler( + handlers=[ + s3, + gcs, + azure, + http, + https, + artifact, + local_artifact, + file_handler, + ], + default_handler=TrackingHandler(), + ) + + def config(self) -> dict: + return self._config + + def load_file( + self, + artifact: Artifact, + manifest_entry: ArtifactManifestEntry, + dest_path: str | None = None, + executor: concurrent.futures.Executor | None = None, + multipart: bool | None = None, + ) -> FilePathStr: + """Use cache or download the file using signed url. + + Args: + executor: Passed from caller, artifact has a thread pool for multi file download. + Reuse the thread pool for multi part download. The thread pool is closed when + artifact download is done. + multipart: If set to `None` (default), the artifact will be downloaded + in parallel using multipart download if individual file size is greater than + 2GB. If set to `True` or `False`, the artifact will be downloaded in + parallel or serially regardless of the file size. + """ + if dest_path is not None: + self._cache._override_cache_path = dest_path + + path, hit, cache_open = self._cache.check_md5_obj_path( + B64MD5(manifest_entry.digest), # TODO(spencerpearson): unsafe cast + manifest_entry.size if manifest_entry.size is not None else 0, + ) + if hit: + return path + + if manifest_entry._download_url is not None: + # Use multipart parallel download for large file + if ( + executor is not None + and manifest_entry.size is not None + and self._should_multipart_download(manifest_entry.size, multipart) + ): + self._multipart_file_download( + executor, + manifest_entry._download_url, + manifest_entry.size, + cache_open, + ) + return path + # Serial download + response = self._session.get(manifest_entry._download_url, stream=True) + try: + response.raise_for_status() + except Exception: + # Signed URL might have expired, fall back to fetching it one by one. + manifest_entry._download_url = None + if manifest_entry._download_url is None: + auth = None + http_headers = _thread_local_api_settings.headers or {} + if self._api.access_token is not None: + http_headers["Authorization"] = f"Bearer {self._api.access_token}" + elif _thread_local_api_settings.cookies is None: + auth = ("api", self._api.api_key or "") + response = self._session.get( + self._file_url( + self._api, + artifact.entity, + artifact.project, + artifact.name.split(":")[0], + manifest_entry, + ), + auth=auth, + cookies=_thread_local_api_settings.cookies, + headers=http_headers, + stream=True, + ) + response.raise_for_status() + + with cache_open(mode="wb") as file: + for data in response.iter_content(chunk_size=16 * 1024): + file.write(data) + return path + + def _should_multipart_download( + self, + file_size: int, + multipart: bool | None, + ) -> bool: + if multipart is not None: + return multipart + return file_size >= _MULTIPART_DOWNLOAD_SIZE + + def _write_chunks_to_file( + self, + f: IO, + q: queue.Queue, + download_has_error: threading.Event, + ): + while not download_has_error.is_set(): + item = q.get() + if item is _CHUNK_QUEUE_SENTINEL: + # Normal shutdown, all the chunks are written + return + elif isinstance(item, _ChunkContent): + try: + # NOTE: Seek works without pre allocating the file on disk. + # It automatically creates a sparse file, e.g. ls -hl would show + # a bigger size compared to du -sh * because downloading different + # chunks is not a sequential write. + # See https://man7.org/linux/man-pages/man2/lseek.2.html + f.seek(item.offset) + f.write(item.data) + except Exception as e: + if env.is_debug(): + logger.debug(f"Error writing chunk to file: {e}") + download_has_error.set() + raise + else: + raise ValueError(f"Unknown queue item type: {type(item)}") + + def _download_part( + self, + download_url: str, + headers: dict, + start: int, + q: queue.Queue, + download_has_error: threading.Event, + ): + # Other threads has error, no need to start + if download_has_error.is_set(): + return + response = self._session.get( + url=download_url, + headers=headers, + stream=True, + ) + response.raise_for_status() + + file_offset = start + for content in response.iter_content(chunk_size=_HTTP_RES_CHUNK_SIZE_BYTES): + if download_has_error.is_set(): + return + q.put(_ChunkContent(offset=file_offset, data=content)) + file_offset += len(content) + + def _multipart_file_download( + self, + executor: concurrent.futures.Executor, + download_url: str, + file_size_bytes: int, + cache_open: Opener, + ): + """Download file as multiple parts in parallel. + + Only one thread for writing to file. Each part run one http request in one thread. + HTTP response chunk of a file part is sent to the writer thread via a queue. + """ + q: queue.Queue[_ChunkContent | object] = queue.Queue(maxsize=500) + download_has_error = threading.Event() + + # Put cache_open at top so we remove the tmp file when there is network error. + with cache_open("wb") as f: + # Start writer thread first. + write_handler = functools.partial( + self._write_chunks_to_file, f, q, download_has_error + ) + write_future = executor.submit(write_handler) + + # Start download threads for each part. + download_futures: deque[concurrent.futures.Future] = deque() + part_size = _DOWNLOAD_PART_SIZE_BYTES + num_parts = int(math.ceil(file_size_bytes / float(part_size))) + for i in range(num_parts): + # https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Range + # Start and end are both inclusive, empty end means use the actual end of the file. + start = i * part_size + bytes_range = f"bytes={start}-" + if i != (num_parts - 1): + # bytes=0-499 + bytes_range += f"{start + part_size - 1}" + headers = {"Range": bytes_range} + download_handler = functools.partial( + self._download_part, + download_url, + headers, + start, + q, + download_has_error, + ) + download_futures.append(executor.submit(download_handler)) + + # Wait for download + done, not_done = concurrent.futures.wait( + download_futures, return_when=concurrent.futures.FIRST_EXCEPTION + ) + try: + for fut in done: + fut.result() + except Exception as e: + if env.is_debug(): + logger.debug(f"Error downloading file: {e}") + download_has_error.set() + raise + finally: + # Always signal the writer to stop + q.put(_CHUNK_QUEUE_SENTINEL) + write_future.result() + + def store_reference( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: str | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + return self._handler.store_path( + artifact, path, name=name, checksum=checksum, max_objects=max_objects + ) + + def load_reference( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + dest_path: str | None = None, + ) -> FilePathStr | URIStr: + assert manifest_entry.ref is not None + used_handler = self._handler._get_handler(manifest_entry.ref) + if hasattr(used_handler, "_cache") and (dest_path is not None): + used_handler._cache._override_cache_path = dest_path + return self._handler.load_path(manifest_entry, local) + + def _file_url( + self, + api: InternalApi, + entity_name: str, + project_name: str, + artifact_name: str, + manifest_entry: ArtifactManifestEntry, + ) -> str: + storage_layout = self._config.get("storageLayout", StorageLayout.V1) + storage_region = self._config.get("storageRegion", "default") + md5_hex = b64_to_hex_id(B64MD5(manifest_entry.digest)) + + if storage_layout == StorageLayout.V1: + return "{}/artifacts/{}/{}".format( + api.settings("base_url"), entity_name, md5_hex + ) + elif storage_layout == StorageLayout.V2: + if api._server_supports( + ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER # type: ignore + ): + return "{}/artifactsV2/{}/{}/{}/{}/{}/{}/{}".format( + api.settings("base_url"), + storage_region, + quote(entity_name), + quote(project_name), + quote(artifact_name), + quote(manifest_entry.birth_artifact_id or ""), + md5_hex, + manifest_entry.path.name, + ) + return "{}/artifactsV2/{}/{}/{}/{}".format( + api.settings("base_url"), + storage_region, + entity_name, + quote( + manifest_entry.birth_artifact_id + if manifest_entry.birth_artifact_id is not None + else "" + ), + md5_hex, + ) + else: + raise Exception(f"unrecognized storage layout: {storage_layout}") + + def s3_multipart_file_upload( + self, + file_path: str, + chunk_size: int, + hex_digests: dict[int, str], + multipart_urls: dict[int, str], + extra_headers: dict[str, str], + ) -> list[dict[str, Any]]: + etags = [] + part_number = 1 + + with open(file_path, "rb") as f: + while True: + data = f.read(chunk_size) + if not data: + break + md5_b64_str = str(hex_to_b64_id(hex_digests[part_number])) + upload_resp = self._api.upload_multipart_file_chunk_retry( + multipart_urls[part_number], + data, + extra_headers={ + "content-md5": md5_b64_str, + "content-length": str(len(data)), + "content-type": extra_headers.get("Content-Type", ""), + }, + ) + assert upload_resp is not None + etags.append( + {"partNumber": part_number, "hexMD5": upload_resp.headers["ETag"]} + ) + part_number += 1 + return etags + + def default_file_upload( + self, + upload_url: str, + file_path: str, + extra_headers: dict[str, Any], + progress_callback: progress.ProgressFn | None = None, + ) -> None: + """Upload a file to the artifact store and write to cache.""" + with open(file_path, "rb") as file: + # This fails if we don't send the first byte before the signed URL expires. + self._api.upload_file_retry( + upload_url, + file, + progress_callback, + extra_headers=extra_headers, + ) + + def calc_chunk_size(self, file_size: int) -> int: + # Default to chunk size of 100MiB. S3 has cap of 10,000 upload parts. + # If file size exceeds the default chunk size, recalculate chunk size. + default_chunk_size = 100 * 1024**2 + if default_chunk_size * S3_MAX_PART_NUMBERS < file_size: + return math.ceil(file_size / S3_MAX_PART_NUMBERS) + return default_chunk_size + + def store_file( + self, + artifact_id: str, + artifact_manifest_id: str, + entry: ArtifactManifestEntry, + preparer: StepPrepare, + progress_callback: progress.ProgressFn | None = None, + ) -> bool: + """Upload a file to the artifact store. + + Returns: + True if the file was a duplicate (did not need to be uploaded), + False if it needed to be uploaded or was a reference (nothing to dedupe). + """ + file_size = entry.size if entry.size is not None else 0 + chunk_size = self.calc_chunk_size(file_size) + upload_parts = [] + hex_digests = {} + file_path = entry.local_path if entry.local_path is not None else "" + # Logic for AWS s3 multipart upload. + # Only chunk files if larger than 2 GiB. Currently can only support up to 5TiB. + if ( + file_size >= S3_MIN_MULTI_UPLOAD_SIZE + and file_size <= S3_MAX_MULTI_UPLOAD_SIZE + ): + part_number = 1 + with open(file_path, "rb") as f: + while True: + data = f.read(chunk_size) + if not data: + break + hex_digest = hashlib.md5(data).hexdigest() + upload_parts.append( + {"hexMD5": hex_digest, "partNumber": part_number} + ) + hex_digests[part_number] = hex_digest + part_number += 1 + + resp = preparer.prepare( + { + "artifactID": artifact_id, + "artifactManifestID": artifact_manifest_id, + "name": entry.path, + "md5": entry.digest, + "uploadPartsInput": upload_parts, + } + ).get() + + entry.birth_artifact_id = resp.birth_artifact_id + + multipart_urls = resp.multipart_upload_urls + if resp.upload_url is None: + return True + if entry.local_path is None: + return False + extra_headers = { + header.split(":", 1)[0]: header.split(":", 1)[1] + for header in (resp.upload_headers or {}) + } + + # This multipart upload isn't available, do a regular single url upload + if multipart_urls is None and resp.upload_url: + self.default_file_upload( + resp.upload_url, file_path, extra_headers, progress_callback + ) + else: + if multipart_urls is None: + raise ValueError(f"No multipart urls to upload for file: {file_path}") + # Upload files using s3 multipart upload urls + etags = self.s3_multipart_file_upload( + file_path, + chunk_size, + hex_digests, + multipart_urls, + extra_headers, + ) + assert resp.storage_path is not None + self._api.complete_multipart_upload_artifact( + artifact_id, resp.storage_path, etags, resp.upload_id + ) + self._write_cache(entry) + + return False + + def _write_cache(self, entry: ArtifactManifestEntry) -> None: + if entry.local_path is None: + return + + # Cache upon successful upload. + _, hit, cache_open = self._cache.check_md5_obj_path( + B64MD5(entry.digest), + entry.size if entry.size is not None else 0, + ) + + staging_dir = get_staging_dir() + try: + if not entry.skip_cache and not hit: + with cache_open("wb") as f, open(entry.local_path, "rb") as src: + shutil.copyfileobj(src, f) + if entry.local_path.startswith(staging_dir): + # Delete staged files here instead of waiting till + # all the files are uploaded + os.chmod(entry.local_path, 0o600) + os.remove(entry.local_path) + except OSError as e: + termwarn(f"Failed to cache {entry.local_path}, ignoring {e}") diff --git a/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policy.py b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..b33beabd42b98184ee39d350458dd1f3efb77973 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/artifacts/storage_policy.py @@ -0,0 +1,75 @@ +"""Storage policy.""" + +from __future__ import annotations + +import concurrent.futures +from typing import TYPE_CHECKING, Sequence + +from wandb.sdk.internal.internal_api import Api as InternalApi +from wandb.sdk.lib.paths import FilePathStr, URIStr + +if TYPE_CHECKING: + from wandb.filesync.step_prepare import StepPrepare + from wandb.sdk.artifacts.artifact import Artifact + from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry + from wandb.sdk.internal.progress import ProgressFn + + +class StoragePolicy: + @classmethod + def lookup_by_name(cls, name: str) -> type[StoragePolicy]: + import wandb.sdk.artifacts.storage_policies # noqa: F401 + + for sub in cls.__subclasses__(): + if sub.name() == name: + return sub + raise NotImplementedError(f"Failed to find storage policy '{name}'") + + @classmethod + def name(cls) -> str: + raise NotImplementedError + + @classmethod + def from_config(cls, config: dict, api: InternalApi | None = None) -> StoragePolicy: + raise NotImplementedError + + def config(self) -> dict: + raise NotImplementedError + + def load_file( + self, + artifact: Artifact, + manifest_entry: ArtifactManifestEntry, + dest_path: str | None = None, + executor: concurrent.futures.Executor | None = None, + multipart: bool | None = None, + ) -> FilePathStr: + raise NotImplementedError + + def store_file( + self, + artifact_id: str, + artifact_manifest_id: str, + entry: ArtifactManifestEntry, + preparer: StepPrepare, + progress_callback: ProgressFn | None = None, + ) -> bool: + raise NotImplementedError + + def store_reference( + self, + artifact: Artifact, + path: URIStr | FilePathStr, + name: str | None = None, + checksum: bool = True, + max_objects: int | None = None, + ) -> Sequence[ArtifactManifestEntry]: + raise NotImplementedError + + def load_reference( + self, + manifest_entry: ArtifactManifestEntry, + local: bool = False, + dest_path: str | None = None, + ) -> FilePathStr | URIStr: + raise NotImplementedError diff --git a/lib/python3.12/site-packages/wandb/sdk/backend/__init__.py b/lib/python3.12/site-packages/wandb/sdk/backend/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/backend/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/backend/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..20148446e5c85b6db1dfe300bf9069d168cf2b6c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/backend/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/backend/__pycache__/backend.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/backend/__pycache__/backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6bf0bd15fd16d6663fffacae7fdd3d105b838e7e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/backend/__pycache__/backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/backend/backend.py b/lib/python3.12/site-packages/wandb/sdk/backend/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..87079251b0a18de4be064cddd6d92069d4c9d052 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/backend/backend.py @@ -0,0 +1,57 @@ +"""Backend - Send to internal process. + +Manage backend. + +""" + +import logging +from typing import TYPE_CHECKING, Optional + +from wandb.sdk.interface.interface import InterfaceBase +from wandb.sdk.wandb_settings import Settings + +if TYPE_CHECKING: + from wandb.sdk.lib.service import service_connection + +logger = logging.getLogger("wandb") + + +class Backend: + interface: Optional[InterfaceBase] + + _settings: Settings + + _done: bool + + _service: Optional["service_connection.ServiceConnection"] + + def __init__( + self, + settings: Settings, + service: Optional["service_connection.ServiceConnection"] = None, + ) -> None: + self._done = False + + self.interface = None + + self._settings = settings + self._service = service + + def ensure_launched(self) -> None: + """Launch backend worker if not running.""" + assert self._settings.run_id + assert self._service + self.interface = self._service.make_interface( + stream_id=self._settings.run_id, + ) + + def server_status(self) -> None: + """Report server status.""" + + def cleanup(self) -> None: + # TODO: make _done atomic + if self._done: + return + self._done = True + if self.interface: + self.interface.join() diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__init__.py b/lib/python3.12/site-packages/wandb/sdk/data_types/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf387c9c651eccc186f48b5982b9ccabc6e28086 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/_dtypes.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/_dtypes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee85bdeeafb04577e1b198cd50de0062dfd50748 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/_dtypes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/_private.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/_private.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e8fad65d8fbdceb757c09c0c356f31c23275cb68 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/_private.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/audio.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/audio.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d08e3ac2908fd950bde4f5a3904723e122a3c9c6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/audio.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/bokeh.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/bokeh.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..26a369c1e99380a27002d1fd83fa2bf0c6611879 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/bokeh.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/graph.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/graph.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb0bf8c698a91ca15fc059a3b6a703a7f0f55601 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/graph.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/histogram.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/histogram.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45c9540e134396a93f58986f8c9325a9f7bbadc6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/histogram.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/html.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/html.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..38ca59bc36dc4df4d41e4910c0f31629825f65f3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/html.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/image.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/image.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f53b0f3a0796b6ee139106e3601654df84658d4b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/image.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/molecule.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/molecule.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac39d65a4f61d426f7528eaa1b7322d364fddc86 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/molecule.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/object_3d.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/object_3d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4066d53c79afcf3f0c31ba04a3cfb7b249877d0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/object_3d.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/plotly.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/plotly.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb860ef37002a98e8b27f103faa86d7de76334c8 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/plotly.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/saved_model.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/saved_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d377657553a3ea99a1c6d3b2f7965de04ab59e94 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/saved_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/table.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/table.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6bcf21ccfca552860daf097cbeeb94f8a8407c17 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/table.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/table_decorators.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/table_decorators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4dd0b91523ae6ecaa019409455f2d2def92a3ab Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/table_decorators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/trace_tree.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/trace_tree.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee24ec0b500eaa9886e6bed70ad9911c17e5e77a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/trace_tree.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd6a2fca807a30394c06a038ed13033e7d90c2a0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/video.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/video.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..55cc0b78a6a13d5b18b82def166ca3c2c8598c39 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/__pycache__/video.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/_dtypes.py b/lib/python3.12/site-packages/wandb/sdk/data_types/_dtypes.py new file mode 100644 index 0000000000000000000000000000000000000000..015b6af0501266264dcf710ce6f883a3c0f7dad9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/_dtypes.py @@ -0,0 +1,914 @@ +import datetime +import math +import typing as t + +from wandb.util import ( + _is_artifact_string, + _is_artifact_version_weave_dict, + get_module, + is_numpy_array, +) + +np = get_module("numpy") # intentionally not required + +if t.TYPE_CHECKING: + from wandb.sdk.artifacts.artifact import Artifact + +ConvertibleToType = t.Union["Type", t.Type["Type"], type, t.Any] + + +class TypeRegistry: + """A resolver for python objects that can deserialize JSON dicts. + + Additional types can be registered via the .add call. + """ + + _types_by_name = None + _types_by_class = None + + @staticmethod + def types_by_name(): + if TypeRegistry._types_by_name is None: + TypeRegistry._types_by_name = {} + return TypeRegistry._types_by_name + + @staticmethod + def types_by_class(): + if TypeRegistry._types_by_class is None: + TypeRegistry._types_by_class = {} + return TypeRegistry._types_by_class + + @staticmethod + def add(wb_type: t.Type["Type"]) -> None: + assert issubclass(wb_type, Type) + TypeRegistry.types_by_name().update({wb_type.name: wb_type}) + for name in wb_type.legacy_names: + TypeRegistry.types_by_name().update({name: wb_type}) + TypeRegistry.types_by_class().update( + {_type: wb_type for _type in wb_type.types} + ) + + @staticmethod + def type_of(py_obj: t.Optional[t.Any]) -> "Type": + # Special case handler for common case of np.nans. np.nan + # is of type 'float', but should be treated as a None. This is + # because np.nan can co-exist with other types in dataframes, + # but will be ultimately treated as a None. Ignoring type since + # mypy does not trust that py_obj is a float by the time it is + # passed to isnan. + if py_obj.__class__ is float and math.isnan(py_obj): # type: ignore + return NoneType() + + # TODO: generalize this to handle other config input types + if _is_artifact_string(py_obj) or _is_artifact_version_weave_dict(py_obj): + return TypeRegistry.types_by_name().get("artifactVersion")() + + class_handler = TypeRegistry.types_by_class().get(py_obj.__class__) + _type = None + if class_handler: + _type = class_handler.from_obj(py_obj) + else: + _type = PythonObjectType.from_obj(py_obj) + return _type + + @staticmethod + def type_from_dict( + json_dict: t.Dict[str, t.Any], artifact: t.Optional["Artifact"] = None + ) -> "Type": + wb_type = json_dict.get("wb_type") + if wb_type is None: + TypeError("json_dict must contain `wb_type` key") + _type = TypeRegistry.types_by_name().get(wb_type) + if _type is None: + TypeError(f"missing type handler for {wb_type}") + return _type.from_json(json_dict, artifact) + + @staticmethod + def type_from_dtype(dtype: ConvertibleToType) -> "Type": + # The dtype is already an instance of Type + if isinstance(dtype, Type): + wbtype: Type = dtype + + # The dtype is a subclass of Type + elif isinstance(dtype, type) and issubclass(dtype, Type): + wbtype = dtype() + + # The dtype is a subclass of generic python type + elif isinstance(dtype, type): + handler = TypeRegistry.types_by_class().get(dtype) + + # and we have a registered handler + if handler: + wbtype = handler() + + # else, fallback to object type + else: + wbtype = PythonObjectType.from_obj(dtype) + + # The dtype is a list, then we resolve the list notation + elif isinstance(dtype, list): + if len(dtype) == 0: + wbtype = ListType() + elif len(dtype) == 1: + wbtype = ListType(TypeRegistry.type_from_dtype(dtype[0])) + + # lists of more than 1 are treated as unions + else: + wbtype = UnionType([TypeRegistry.type_from_dtype(dt) for dt in dtype]) + + # The dtype is a dict, then we resolve the dict notation + elif isinstance(dtype, dict): + wbtype = TypedDictType( + {key: TypeRegistry.type_from_dtype(dtype[key]) for key in dtype} + ) + + # The dtype is a concrete instance, which we will treat as a constant + else: + wbtype = ConstType(dtype) + + return wbtype + + +def _params_obj_to_json_obj( + params_obj: t.Any, + artifact: t.Optional["Artifact"] = None, +) -> t.Any: + """Helper method.""" + if params_obj.__class__ is dict: + return { + key: _params_obj_to_json_obj(params_obj[key], artifact) + for key in params_obj + } + elif params_obj.__class__ in [list, set, tuple, frozenset]: + return [_params_obj_to_json_obj(item, artifact) for item in list(params_obj)] + elif isinstance(params_obj, Type): + return params_obj.to_json(artifact) + else: + return params_obj + + +def _json_obj_to_params_obj( + json_obj: t.Any, artifact: t.Optional["Artifact"] = None +) -> t.Any: + """Helper method.""" + if json_obj.__class__ is dict: + if "wb_type" in json_obj: + return TypeRegistry.type_from_dict(json_obj, artifact) + else: + return { + key: _json_obj_to_params_obj(json_obj[key], artifact) + for key in json_obj + } + elif json_obj.__class__ is list: + return [_json_obj_to_params_obj(item, artifact) for item in json_obj] + else: + return json_obj + + +class Type: + """The most generic type that all types subclass. + + It provides simple serialization and deserialization as well as equality checks. + A name class-level property must be uniquely set by subclasses. + """ + + # Subclasses must override with a unique name. This is used to identify the + # class during serializations and deserializations + name: t.ClassVar[str] = "" + + # List of names by which this class can deserialize + legacy_names: t.ClassVar[t.List[str]] = [] + + # Subclasses may override with a list of `types` which this Type is capable + # of being initialized. This is used by the Type Registry when calling `TypeRegistry.type_of`. + # Some types will have an empty list - for example `Union`. There is no raw python type which + # inherently maps to a Union and therefore the list should be empty. + types: t.ClassVar[t.List[type]] = [] + + # Contains the further specification of the Type + _params: t.Dict[str, t.Any] + + def __init__(*args, **kwargs): + pass + + @property + def params(self): + if not hasattr(self, "_params") or self._params is None: + self._params = {} + return self._params + + def assign(self, py_obj: t.Optional[t.Any] = None) -> "Type": + """Assign a python object to the type. + + May to be overridden by subclasses + + Args: + py_obj (any, optional): Any python object which the user wishes to assign to + this type + + Returns: + Type: a new type representing the result of the assignment. + """ + return self.assign_type(TypeRegistry.type_of(py_obj)) + + def assign_type(self, wb_type: "Type") -> "Type": + # Default - should be overridden + if isinstance(wb_type, self.__class__) and self.params == wb_type.params: + return self + else: + return InvalidType() + + def to_json(self, artifact: t.Optional["Artifact"] = None) -> t.Dict[str, t.Any]: + """Generate a jsonable dictionary serialization the type. + + If overridden by subclass, ensure that `from_json` is equivalently overridden. + + Args: + artifact (wandb.Artifact, optional): If the serialization is being performed + for a particular artifact, pass that artifact. Defaults to None. + + Returns: + dict: Representation of the type + """ + res = { + "wb_type": self.name, + "params": _params_obj_to_json_obj(self.params, artifact), + } + if res["params"] is None or res["params"] == {}: + del res["params"] + + return res + + @classmethod + def from_json( + cls, + json_dict: t.Dict[str, t.Any], + artifact: t.Optional["Artifact"] = None, + ) -> "Type": + """Construct a new instance of the type using a JSON dictionary. + + The mirror function of `to_json`. If overridden by subclass, ensure that + `to_json` is equivalently overridden. + + Returns: + _Type: an instance of a subclass of the _Type class. + """ + return cls(**_json_obj_to_params_obj(json_dict.get("params", {}), artifact)) + + @classmethod + def from_obj(cls, py_obj: t.Optional[t.Any] = None) -> "Type": + return cls() + + def explain(self, other: t.Any, depth=0) -> str: + """Explain why an item is not assignable to a type. + + Assumes that the caller has already validated that the assignment fails. + + Args: + other (any): Any object depth (int, optional): depth of the type checking. + Defaults to 0. + + Returns: + str: human-readable explanation + """ + wbtype = TypeRegistry.type_of(other) + gap = "".join(["\t"] * depth) + if depth > 0: + return f"{gap}{wbtype} not assignable to {self}" + else: + return f"{gap}{other} of type {wbtype} is not assignable to {self}" + + def __repr__(self): + rep = self.name.capitalize() + if len(self.params.keys()) > 0: + rep += "(" + for ndx, key in enumerate(self.params.keys()): + if ndx > 0: + rep += ", " + rep += key + ":" + str(self.params[key]) + rep += ")" + return rep + + def __eq__(self, other): + return self is other or ( + isinstance(self, Type) + and isinstance(other, Type) + and self.name == other.name + and self.params.keys() == other.params.keys() + and all([self.params[k] == other.params[k] for k in self.params]) + ) + + +class InvalidType(Type): + """A disallowed type. + + Assignments to a InvalidType result in a Never Type. InvalidType is basically the + invalid case. + """ + + name = "invalid" + types: t.ClassVar[t.List[type]] = [] + + def assign_type(self, wb_type: "Type") -> "InvalidType": + return self + + +class AnyType(Type): + """An object that can be any type. + + Assignments to an AnyType result in the AnyType except None which results in an + InvalidType. + """ + + name = "any" + types: t.ClassVar[t.List[type]] = [] + + def assign_type(self, wb_type: "Type") -> t.Union["AnyType", InvalidType]: + return ( + self + if not (isinstance(wb_type, NoneType) or isinstance(wb_type, InvalidType)) + else InvalidType() + ) + + +class UnknownType(Type): + """An object with an unknown type. + + All assignments to an UnknownType result in the type of the assigned object except + `None` which results in a InvalidType. + """ + + name = "unknown" + types: t.ClassVar[t.List[type]] = [] + + def assign_type(self, wb_type: "Type") -> "Type": + return wb_type if not isinstance(wb_type, NoneType) else InvalidType() + + +class NoneType(Type): + name = "none" + types: t.ClassVar[t.List[type]] = [None.__class__] + + +class StringType(Type): + name = "string" + types: t.ClassVar[t.List[type]] = [str] + + +class NumberType(Type): + name = "number" + types: t.ClassVar[t.List[type]] = [int, float] + + +if np: + NumberType.types.append(np.byte) + NumberType.types.append(np.short) + NumberType.types.append(np.ushort) + NumberType.types.append(np.intc) + NumberType.types.append(np.uintc) + NumberType.types.append(np.int_) + NumberType.types.append(np.uint) + NumberType.types.append(np.longlong) + NumberType.types.append(np.ulonglong) + NumberType.types.append(np.half) + NumberType.types.append(np.float16) + NumberType.types.append(np.single) + NumberType.types.append(np.double) + NumberType.types.append(np.longdouble) + NumberType.types.append(np.csingle) + NumberType.types.append(np.cdouble) + NumberType.types.append(np.clongdouble) + NumberType.types.append(np.int8) + NumberType.types.append(np.int16) + NumberType.types.append(np.int32) + NumberType.types.append(np.int64) + NumberType.types.append(np.uint8) + NumberType.types.append(np.uint16) + NumberType.types.append(np.uint32) + NumberType.types.append(np.uint64) + NumberType.types.append(np.intp) + NumberType.types.append(np.uintp) + NumberType.types.append(np.float32) + NumberType.types.append(np.float64) + NumberType.types.append(np.complex64) + NumberType.types.append(np.complex128) + + numpy_major_version = np.__version__.split(".")[0] + if int(numpy_major_version) < 2: + NumberType.types.append(np.float_) + NumberType.types.append(np.complex_) + + +class TimestampType(Type): + name = "timestamp" + types: t.ClassVar[t.List[type]] = [datetime.datetime, datetime.date] + + +if np: + TimestampType.types.append(np.datetime64) + + +class BooleanType(Type): + name = "boolean" + types: t.ClassVar[t.List[type]] = [bool] + + +if np: + BooleanType.types.append(np.bool_) + + +class PythonObjectType(Type): + """A backup type that keeps track of the python object name.""" + + name = "pythonObject" + legacy_names = ["object"] + types: t.ClassVar[t.List[type]] = [] + + def __init__(self, class_name: str): + self.params.update({"class_name": class_name}) + + @classmethod + def from_obj(cls, py_obj: t.Optional[t.Any] = None) -> "PythonObjectType": + return cls(py_obj.__class__.__name__) + + +class ConstType(Type): + """A constant value (currently only primitives supported).""" + + name = "const" + types: t.ClassVar[t.List[type]] = [] + + def __init__(self, val: t.Optional[t.Any] = None, is_set: t.Optional[bool] = False): + if val.__class__ not in [str, int, float, bool, set, list, None.__class__]: + TypeError( + f"ConstType only supports str, int, float, bool, set, list, and None types. Found {val}" + ) + if is_set or isinstance(val, set): + is_set = True + assert isinstance(val, set) or isinstance(val, list) + val = set(val) + + self.params.update({"val": val, "is_set": is_set}) + + def assign(self, py_obj: t.Optional[t.Any] = None) -> "Type": + return self.assign_type(ConstType(py_obj)) + + @classmethod + def from_obj(cls, py_obj: t.Optional[t.Any] = None) -> "ConstType": + return cls(py_obj) + + def __repr__(self): + return str(self.params["val"]) + + +def _flatten_union_types(wb_types: t.List[Type]) -> t.List[Type]: + final_types = [] + for allowed_type in wb_types: + if isinstance(allowed_type, UnionType): + internal_types = _flatten_union_types(allowed_type.params["allowed_types"]) + for internal_type in internal_types: + final_types.append(internal_type) + else: + final_types.append(allowed_type) + return final_types + + +def _union_assigner( + allowed_types: t.List[Type], + obj_or_type: t.Union[Type, t.Optional[t.Any]], + type_mode=False, +) -> t.Union[t.List[Type], InvalidType]: + resolved_types = [] + valid = False + unknown_count = 0 + + for allowed_type in allowed_types: + if valid: + resolved_types.append(allowed_type) + else: + if isinstance(allowed_type, UnknownType): + unknown_count += 1 + else: + if type_mode: + assert isinstance(obj_or_type, Type) + assigned_type = allowed_type.assign_type(obj_or_type) + else: + assigned_type = allowed_type.assign(obj_or_type) + if isinstance(assigned_type, InvalidType): + resolved_types.append(allowed_type) + else: + resolved_types.append(assigned_type) + valid = True + + if not valid: + if unknown_count == 0: + return InvalidType() + else: + if type_mode: + assert isinstance(obj_or_type, Type) + new_type = obj_or_type + else: + new_type = UnknownType().assign(obj_or_type) + if isinstance(new_type, InvalidType): + return InvalidType() + else: + resolved_types.append(new_type) + unknown_count -= 1 + + for _ in range(unknown_count): + resolved_types.append(UnknownType()) + + resolved_types = _flatten_union_types(resolved_types) + resolved_types.sort(key=str) + return resolved_types + + +class UnionType(Type): + """An "or" of types.""" + + name = "union" + types: t.ClassVar[t.List[type]] = [] + + def __init__( + self, + allowed_types: t.Optional[t.Sequence[ConvertibleToType]] = None, + ): + assert allowed_types is None or (allowed_types.__class__ is list) + if allowed_types is None: + wb_types = [] + else: + wb_types = [TypeRegistry.type_from_dtype(dt) for dt in allowed_types] + + wb_types = _flatten_union_types(wb_types) + wb_types.sort(key=str) + self.params.update({"allowed_types": wb_types}) + + def assign( + self, py_obj: t.Optional[t.Any] = None + ) -> t.Union["UnionType", InvalidType]: + resolved_types = _union_assigner( + self.params["allowed_types"], py_obj, type_mode=False + ) + if isinstance(resolved_types, InvalidType): + return InvalidType() + return self.__class__(resolved_types) + + def assign_type(self, wb_type: "Type") -> t.Union["UnionType", InvalidType]: + if isinstance(wb_type, UnionType): + assignees = wb_type.params["allowed_types"] + else: + assignees = [wb_type] + + resolved_types = self.params["allowed_types"] + for assignee in assignees: + resolved_types = _union_assigner(resolved_types, assignee, type_mode=True) + if isinstance(resolved_types, InvalidType): + return InvalidType() + + return self.__class__(resolved_types) + + def explain(self, other: t.Any, depth=0) -> str: + exp = super().explain(other, depth) + for ndx, subtype in enumerate(self.params["allowed_types"]): + if ndx > 0: + exp += "\n{}and".format("".join(["\t"] * depth)) + exp += "\n" + subtype.explain(other, depth=depth + 1) + return exp + + def __repr__(self): + return "{}".format(" or ".join([str(t) for t in self.params["allowed_types"]])) + + +def OptionalType(dtype: ConvertibleToType) -> UnionType: # noqa: N802 + """Function that mimics the Type class API for constructing an "Optional Type". + + This is just a Union[wb_type, NoneType]. + + Args: + dtype (Type): type to be optional + + Returns: + Type: Optional version of the type. + """ + return UnionType([TypeRegistry.type_from_dtype(dtype), NoneType()]) + + +class ListType(Type): + """A list of homogeneous types.""" + + name = "list" + types: t.ClassVar[t.List[type]] = [list, tuple, set, frozenset] + + def __init__( + self, + element_type: t.Optional[ConvertibleToType] = None, + length: t.Optional[int] = None, + ): + if element_type is None: + wb_type: Type = UnknownType() + else: + wb_type = TypeRegistry.type_from_dtype(element_type) + + self.params.update({"element_type": wb_type, "length": length}) + + @classmethod + def from_obj(cls, py_obj: t.Optional[t.Any] = None) -> "ListType": + if py_obj is None or not hasattr(py_obj, "__iter__"): + raise TypeError("ListType.from_obj expects py_obj to by list-like") + else: + if hasattr(py_obj, "tolist"): + py_list = py_obj.tolist() + else: + py_list = list(py_obj) + + elm_type: Type + if None not in py_list: + elm_type = UnknownType() + else: + elm_type = OptionalType(UnknownType()) + + for item in py_list: + _elm_type = elm_type.assign(item) + # Commenting this out since we don't want to crash user code at this point, but rather + # retain an invalid internal list type. + # if isinstance(_elm_type, InvalidType): + # raise TypeError( + # "List contained incompatible types. Item at index {}: \n{}".format( + # ndx, elm_type.explain(item, 1) + # ) + # ) + + elm_type = _elm_type + + return cls(elm_type, len(py_list)) + + def assign_type(self, wb_type: "Type") -> t.Union["ListType", InvalidType]: + if isinstance(wb_type, ListType): + assigned_type = self.params["element_type"].assign_type( + wb_type.params["element_type"] + ) + if not isinstance(assigned_type, InvalidType): + return ListType( + assigned_type, + None + if self.params["length"] != wb_type.params["length"] + else self.params["length"], + ) + + return InvalidType() + + def assign( + self, py_obj: t.Optional[t.Any] = None + ) -> t.Union["ListType", InvalidType]: + if hasattr(py_obj, "__iter__"): + new_element_type = self.params["element_type"] + # The following ignore is needed since the above hasattr(py_obj, "__iter__") enforces iteration + # error: Argument 1 to "list" has incompatible type "Optional[Any]"; expected "Iterable[Any]" + py_list = list(py_obj) # type: ignore + for obj in py_list: + new_element_type = new_element_type.assign(obj) + if isinstance(new_element_type, InvalidType): + return InvalidType() + return ListType(new_element_type, len(py_list)) + + return InvalidType() + + def explain(self, other: t.Any, depth=0) -> str: + exp = super().explain(other, depth) + gap = "".join(["\t"] * depth) + if ( # yes, this is a bit verbose, but the mypy typechecker likes it this way + isinstance(other, list) + or isinstance(other, tuple) + or isinstance(other, set) + or isinstance(other, frozenset) + ): + new_element_type = self.params["element_type"] + for ndx, obj in enumerate(list(other)): + _new_element_type = new_element_type.assign(obj) + if isinstance(_new_element_type, InvalidType): + exp += f"\n{gap}Index {ndx}:\n{new_element_type.explain(obj, depth + 1)}" + break + new_element_type = _new_element_type + return exp + + def __repr__(self): + return "{}[]".format(self.params["element_type"]) + + +class NDArrayType(Type): + """Represents a list of homogeneous types.""" + + name = "ndarray" + types: t.ClassVar[t.List[type]] = [] # will manually add type if np is available + _serialization_path: t.Optional[t.Dict[str, str]] + + def __init__( + self, + shape: t.Sequence[int], + serialization_path: t.Optional[t.Dict[str, str]] = None, + ): + self.params.update({"shape": list(shape)}) + self._serialization_path = serialization_path + + @classmethod + def from_obj(cls, py_obj: t.Optional[t.Any] = None) -> "NDArrayType": + if is_numpy_array(py_obj): + return cls(py_obj.shape) # type: ignore + elif isinstance(py_obj, list): + shape = [] + target = py_obj + while isinstance(target, list): + dim = len(target) + shape.append(dim) + if dim > 0: + target = target[0] + return cls(shape) + else: + raise TypeError( + f"NDArrayType.from_obj expects py_obj to be ndarray or list, found {py_obj.__class__}" + ) + + def assign_type(self, wb_type: "Type") -> t.Union["NDArrayType", InvalidType]: + if ( + isinstance(wb_type, NDArrayType) + and self.params["shape"] == wb_type.params["shape"] + ): + return self + elif isinstance(wb_type, ListType): + # Should we return error here? + return self + + return InvalidType() + + def assign( + self, py_obj: t.Optional[t.Any] = None + ) -> t.Union["NDArrayType", InvalidType]: + if is_numpy_array(py_obj) or isinstance(py_obj, list): + py_type = self.from_obj(py_obj) + return self.assign_type(py_type) + + return InvalidType() + + def to_json(self, artifact: t.Optional["Artifact"] = None) -> t.Dict[str, t.Any]: + # custom override to support serialization path outside of params internal dict + res = { + "wb_type": self.name, + "params": { + "shape": self.params["shape"], + "serialization_path": self._serialization_path, + }, + } + + return res + + def _get_serialization_path(self) -> t.Optional[t.Dict[str, str]]: + return self._serialization_path + + def _set_serialization_path(self, path: str, key: str) -> None: + self._serialization_path = {"path": path, "key": key} + + def _clear_serialization_path(self) -> None: + self._serialization_path = None + + +if np: + NDArrayType.types.append(np.ndarray) + +# class KeyPolicy: +# EXACT = "E" # require exact key match +# SUBSET = "S" # all known keys are optional and unknown keys are disallowed +# UNRESTRICTED = "U" # all known keys are optional and unknown keys are Unknown + + +class TypedDictType(Type): + """Represents a dictionary object where each key can have a type.""" + + name = "typedDict" + legacy_names = ["dictionary"] + types: t.ClassVar[t.List[type]] = [dict] + + def __init__( + self, + type_map: t.Optional[t.Dict[str, ConvertibleToType]] = None, + ): + if type_map is None: + type_map = {} + self.params.update( + { + "type_map": { + key: TypeRegistry.type_from_dtype(type_map[key]) for key in type_map + } + } + ) + + @classmethod + def from_obj(cls, py_obj: t.Optional[t.Any] = None) -> "TypedDictType": + if not isinstance(py_obj, dict): + TypeError("TypedDictType.from_obj expects a dictionary") + + assert isinstance(py_obj, dict) # helps mypy type checker + return cls({key: TypeRegistry.type_of(py_obj[key]) for key in py_obj}) + + def assign_type(self, wb_type: "Type") -> t.Union["TypedDictType", InvalidType]: + if ( + isinstance(wb_type, TypedDictType) + and len( + set(wb_type.params["type_map"].keys()) + - set(self.params["type_map"].keys()) + ) + == 0 + ): + type_map = {} + for key in self.params["type_map"]: + type_map[key] = self.params["type_map"][key].assign_type( + wb_type.params["type_map"].get(key, UnknownType()) + ) + if isinstance(type_map[key], InvalidType): + return InvalidType() + return TypedDictType(type_map) + + return InvalidType() + + def assign( + self, py_obj: t.Optional[t.Any] = None + ) -> t.Union["TypedDictType", InvalidType]: + if ( + isinstance(py_obj, dict) + and len(set(py_obj.keys()) - set(self.params["type_map"].keys())) == 0 + ): + type_map = {} + for key in self.params["type_map"]: + type_map[key] = self.params["type_map"][key].assign( + py_obj.get(key, None) + ) + if isinstance(type_map[key], InvalidType): + return InvalidType() + return TypedDictType(type_map) + + return InvalidType() + + def explain(self, other: t.Any, depth=0) -> str: + exp = super().explain(other, depth) + gap = "".join(["\t"] * depth) + if isinstance(other, dict): + extra_keys = set(other.keys()) - set(self.params["type_map"].keys()) + if len(extra_keys) > 0: + exp += "\n{}Found extra keys: {}".format( + gap, ",".join(list(extra_keys)) + ) + + for key in self.params["type_map"]: + val = other.get(key, None) + if isinstance(self.params["type_map"][key].assign(val), InvalidType): + exp += "\n{}Key '{}':\n{}".format( + gap, + key, + self.params["type_map"][key].explain(val, depth=depth + 1), + ) + return exp + + def __repr__(self): + return "{}".format(self.params["type_map"]) + + +# Special Types +TypeRegistry.add(InvalidType) +TypeRegistry.add(AnyType) +TypeRegistry.add(UnknownType) + +# Types with default type mappings +TypeRegistry.add(NoneType) +TypeRegistry.add(StringType) +TypeRegistry.add(TimestampType) +TypeRegistry.add(NumberType) +TypeRegistry.add(BooleanType) +TypeRegistry.add(ListType) +TypeRegistry.add(TypedDictType) + +# Types without default type mappings +TypeRegistry.add(UnionType) +TypeRegistry.add(PythonObjectType) +TypeRegistry.add(ConstType) + +# Common Industry Types +TypeRegistry.add(NDArrayType) + +__all__ = [ + "TypeRegistry", + "InvalidType", + "UnknownType", + "AnyType", + "NoneType", + "StringType", + "NumberType", + "TimestampType", + "BooleanType", + "ListType", + "TypedDictType", + "UnionType", + "PythonObjectType", + "ConstType", + "OptionalType", + "Type", + "NDArrayType", +] diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/_private.py b/lib/python3.12/site-packages/wandb/sdk/data_types/_private.py new file mode 100644 index 0000000000000000000000000000000000000000..887bc8a231014d06dbe65590663064ec792eeb84 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/_private.py @@ -0,0 +1,10 @@ +import atexit +import tempfile + +# Staging directory, so we can encode raw data into files, then hash them before +# we put them into the Run directory to be uploaded. +MEDIA_TMP = tempfile.TemporaryDirectory("wandb-media") + + +def _cleanup_media_tmp_dir() -> None: + atexit.register(MEDIA_TMP.cleanup) diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/audio.py b/lib/python3.12/site-packages/wandb/sdk/data_types/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..c1dc46ed101842b0904917c445ae0ba873d625b0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/audio.py @@ -0,0 +1,180 @@ +import hashlib +import os +import pathlib +from typing import TYPE_CHECKING, Optional, Union + +from wandb import util +from wandb.sdk.lib import filesystem, runid + +from . import _dtypes +from ._private import MEDIA_TMP +from .base_types.media import BatchableMedia + +if TYPE_CHECKING: + import numpy as np + + +class Audio(BatchableMedia): + """Wandb class for audio clips. + + Args: + data_or_path: (string or numpy array) A path to an audio file + or a numpy array of audio data. + sample_rate: (int) Sample rate, required when passing in raw + numpy array of audio data. + caption: (string) Caption to display with audio. + """ + + _log_type = "audio-file" + + def __init__( + self, + data_or_path: Union[ + str, + pathlib.Path, + list, + "np.ndarray", + ], + sample_rate: Optional[int] = None, + caption: Optional[str] = None, + ): + """Accept a path to an audio file or a numpy array of audio data.""" + super().__init__(caption=caption) + self._duration = None + self._sample_rate = sample_rate + + if isinstance(data_or_path, (str, pathlib.Path)): + data_or_path = str(data_or_path) + + if self.path_is_reference(data_or_path): + self._path = data_or_path + self._sha256 = hashlib.sha256(data_or_path.encode("utf-8")).hexdigest() + self._is_tmp = False + else: + self._set_file(data_or_path, is_tmp=False) + else: + if sample_rate is None: + raise ValueError( + 'Argument "sample_rate" is required when instantiating wandb.Audio with raw data.' + ) + + soundfile = util.get_module( + "soundfile", + required='Raw audio requires the soundfile package. To get it, run "pip install soundfile"', + ) + + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".wav") + + soundfile.write(tmp_path, data_or_path, sample_rate) + self._duration = len(data_or_path) / float(sample_rate) + + self._set_file(tmp_path, is_tmp=True) + + @classmethod + def get_media_subdir(cls): + return os.path.join("media", "audio") + + @classmethod + def from_json(cls, json_obj, source_artifact): + return cls( + source_artifact.get_entry(json_obj["path"]).download(), + caption=json_obj["caption"], + ) + + def bind_to_run( + self, run, key, step, id_=None, ignore_copy_err: Optional[bool] = None + ): + if self.path_is_reference(self._path): + raise ValueError( + "Audio media created by a reference to external storage cannot currently be added to a run" + ) + + return super().bind_to_run(run, key, step, id_, ignore_copy_err) + + def to_json(self, run): + json_dict = super().to_json(run) + json_dict.update( + { + "_type": self._log_type, + } + ) + return json_dict + + @classmethod + def seq_to_json(cls, seq, run, key, step): + audio_list = list(seq) + + util.get_module( + "soundfile", + required="wandb.Audio requires the soundfile package. To get it, run: pip install soundfile", + ) + base_path = os.path.join(run.dir, "media", "audio") + filesystem.mkdir_exists_ok(base_path) + meta = { + "_type": "audio", + "count": len(audio_list), + "audio": [a.to_json(run) for a in audio_list], + } + sample_rates = cls.sample_rates(audio_list) + if sample_rates: + meta["sampleRates"] = sample_rates + durations = cls.durations(audio_list) + if durations: + meta["durations"] = durations + captions = cls.captions(audio_list) + if captions: + meta["captions"] = captions + + return meta + + @classmethod + def durations(cls, audio_list): + return [a._duration for a in audio_list] + + @classmethod + def sample_rates(cls, audio_list): + return [a._sample_rate for a in audio_list] + + @classmethod + def captions(cls, audio_list): + captions = [a._caption for a in audio_list] + if all(c is None for c in captions): + return False + else: + return ["" if c is None else c for c in captions] + + def resolve_ref(self): + if self.path_is_reference(self._path): + # this object was already created using a ref: + return self._path + source_artifact = self._artifact_source.artifact + + resolved_name = source_artifact._local_path_to_name(self._path) + if resolved_name is not None: + target_entry = source_artifact.manifest.get_entry_by_path(resolved_name) + if target_entry is not None: + return target_entry.ref + + return None + + def __eq__(self, other): + if self.path_is_reference(self._path) or self.path_is_reference(other._path): + # one or more of these objects is an unresolved reference -- we'll compare + # their reference paths instead of their SHAs: + return ( + self.resolve_ref() == other.resolve_ref() + and self._caption == other._caption + ) + + return super().__eq__(other) and self._caption == other._caption + + def __ne__(self, other): + return not self.__eq__(other) + + +class _AudioFileType(_dtypes.Type): + name = "audio-file" + types = [Audio] + + +_dtypes.TypeRegistry.add(_AudioFileType) diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__init__.py b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40c1533a2cd5df729c18bddef23c4edfb1114be7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/json_metadata.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/json_metadata.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be034c0282147b124d8c59a0c7f1793e2cb87b6a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/json_metadata.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/media.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/media.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..98d432be36b63a0ad7081dc7eeab28cb5ae563e3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/media.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/wb_value.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/wb_value.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63dc76f1a16e1440c24f07eddee16fb97dfade1b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/__pycache__/wb_value.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/json_metadata.py b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/json_metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..b644584abcf1a640c35a1c98e7b7083999f3b446 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/json_metadata.py @@ -0,0 +1,55 @@ +import codecs +import os +from typing import TYPE_CHECKING, Type, Union + +from wandb import util +from wandb.sdk.lib import runid + +from .._private import MEDIA_TMP +from .media import Media + +if TYPE_CHECKING: # pragma: no cover + from wandb.sdk.artifacts.artifact import Artifact + + from ...wandb_run import Run as LocalRun + + +# Allows encoding of arbitrary JSON structures +# as a file +# +# This class should be used as an abstract class +# extended to have validation methods + + +class JSONMetadata(Media): + """JSONMetadata is a type for encoding arbitrary metadata as files.""" + + def __init__(self, val: dict) -> None: + super().__init__() + + self.validate(val) + self._val = val + + ext = "." + self.type_name() + ".json" + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ext) + with codecs.open(tmp_path, "w", encoding="utf-8") as fp: + util.json_dump_uncompressed(self._val, fp) + self._set_file(tmp_path, is_tmp=True, extension=ext) + + @classmethod + def get_media_subdir(cls: Type["JSONMetadata"]) -> str: + return os.path.join("media", "metadata", cls.type_name()) + + def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict: + json_dict = super().to_json(run_or_artifact) + json_dict["_type"] = self.type_name() + + return json_dict + + # These methods should be overridden in the child class + @classmethod + def type_name(cls) -> str: + return "metadata" + + def validate(self, val: dict) -> bool: + return True diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/media.py b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/media.py new file mode 100644 index 0000000000000000000000000000000000000000..6a088c1d129490754b19fafcb2eb10029a03c750 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/media.py @@ -0,0 +1,389 @@ +import hashlib +import os +import pathlib +import platform +import re +import shutil +from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Type, Union, cast + +import wandb +from wandb import util +from wandb.sdk.lib import filesystem +from wandb.sdk.lib.paths import LogicalPath + +from .wb_value import WBValue + +if TYPE_CHECKING: # pragma: no cover + import numpy as np + + from wandb.sdk.artifacts.artifact import Artifact + + from ...wandb_run import Run as LocalRun + + +SYS_PLATFORM = platform.system() + + +def check_windows_valid_filename(path: Union[int, str]) -> bool: + r"""Verify that the given path does not contain any invalid characters for a Windows filename. + + Windows filenames cannot contain the following characters: + < > : " \ / | ? * + + For more details, refer to the official documentation: + https://learn.microsoft.com/en-us/windows/win32/fileio/naming-a-file#naming-conventions + + Args: + path: The file path to check, which can be either an integer or a string. + + Returns: + bool: True if the path does not contain any invalid characters, False otherwise. + """ + return not bool(re.search(r'[<>:"\\?*]', path)) # type: ignore + + +def _wb_filename( + key: Union[str, int], step: Union[str, int], id: Union[str, int], extension: str +) -> str: + r"""Generates a safe filename/path for storing media files, using the provided key, step, and id. + + The filename is made safe by: + 1. Removing any leading slashes to prevent writing to absolute paths + 2. Replacing '.' and '..' with underscores to prevent directory traversal attacks + + If the key contains slashes (e.g. 'images/cats/fluffy.jpg'), subdirectories will be created: + media/ + images/ + cats/ + fluffy.jpg_step_id.ext + + Args: + key: Name/path for the media file + step: Training step number + id: Unique identifier + extension: File extension (e.g. '.jpg', '.mp3') + + Returns: + A sanitized filename string in the format: key_step_id.extension + + Raises: + ValueError: If running on Windows and the key contains invalid filename characters + (\\, :, *, ?, ", <, >, |) + """ + if SYS_PLATFORM == "Windows" and not check_windows_valid_filename(key): + raise ValueError( + f'Media {key} is invalid. Please remove invalid filename characters (\\, :, *, ?, ", <, >, |)' + ) + + # On Windows, convert forward slashes to backslashes. + # This ensures that the key is a valid filename on Windows. + if SYS_PLATFORM == "Windows": + key = str(key).replace("/", os.sep) + + # Avoid writing to absolute paths by striping any leading slashes. + # The key has already been validated for windows operating systems in util.check_windows_valid_filename + # This ensures the key does not contain invalid characters for windows, such as '\' or ':'. + # So we can check only for '/' in the key. + key = str(key).lstrip(os.sep) + + # Avoid directory traversal by replacing dots with underscores. + keys = key.split(os.sep) + keys = [k.replace(".", "_") if k in (os.curdir, os.pardir) else k for k in keys] + + # Recombine the key into a relative path. + key = os.sep.join(keys) + + return f"{str(key)}_{str(step)}_{str(id)}{extension}" + + +class Media(WBValue): + """A WBValue stored as a file outside JSON that can be rendered in a media panel. + + If necessary, we move or copy the file into the Run's media directory so that it + gets uploaded. + """ + + _path: Optional[str] + _run: Optional["LocalRun"] + _caption: Optional[str] + _is_tmp: Optional[bool] + _extension: Optional[str] + _sha256: Optional[str] + _size: Optional[int] + + def __init__(self, caption: Optional[str] = None) -> None: + super().__init__() + self._path = None + # The run under which this object is bound, if any. + self._run = None + self._caption = caption + + def _set_file( + self, + path: str, + is_tmp: bool = False, + extension: Optional[str] = None, + ) -> None: + self._path = path + self._is_tmp = is_tmp + self._extension = extension + assert extension is None or path.endswith(extension), ( + f'Media file extension "{extension}" must occur at the end of path "{path}".' + ) + + with open(self._path, "rb") as f: + self._sha256 = hashlib.sha256(f.read()).hexdigest() + self._size = os.path.getsize(self._path) + + @classmethod + def get_media_subdir(cls: Type["Media"]) -> str: + raise NotImplementedError + + @staticmethod + def captions( + media_items: Sequence["Media"], + ) -> Union[bool, Sequence[Optional[str]]]: + if media_items[0]._caption is not None: + return [m._caption for m in media_items] + else: + return False + + def is_bound(self) -> bool: + return self._run is not None + + def file_is_set(self) -> bool: + return self._path is not None and self._sha256 is not None + + def bind_to_run( + self, + run: "LocalRun", + key: Union[int, str], + step: Union[int, str], + id_: Optional[Union[int, str]] = None, + ignore_copy_err: Optional[bool] = None, + ) -> None: + """Bind this object to a particular Run. + + Calling this function is necessary so that we have somewhere specific to put the + file associated with this object, from which other Runs can refer to it. + """ + assert self.file_is_set(), "bind_to_run called before _set_file" + + # The following two assertions are guaranteed to pass + # by definition file_is_set, but are needed for + # mypy to understand that these are strings below. + assert isinstance(self._path, str) + assert isinstance(self._sha256, str) + + assert run is not None, 'Argument "run" must not be None.' + self._run = run + + if self._extension is None: + _, extension = os.path.splitext(os.path.basename(self._path)) + else: + extension = self._extension + + if id_ is None: + id_ = self._sha256[:20] + + file_path = _wb_filename(key, step, id_, extension) + media_path = os.path.join(self.get_media_subdir(), file_path) + new_path = os.path.join(self._run.dir, media_path) + filesystem.mkdir_exists_ok(os.path.dirname(new_path)) + + if self._is_tmp: + shutil.move(self._path, new_path) + self._path = new_path + self._is_tmp = False + run._publish_file(media_path) + else: + try: + shutil.copy(self._path, new_path) + except shutil.SameFileError: + if not ignore_copy_err: + raise + self._path = new_path + run._publish_file(media_path) + + def to_json(self, run: Union["LocalRun", "Artifact"]) -> dict: + """Serialize the object into a JSON blob. + + Uses run or artifact to store additional data. If `run_or_artifact` is a + wandb.Run then `self.bind_to_run()` must have been previously been called. + + Args: + run_or_artifact (wandb.Run | wandb.Artifact): the Run or Artifact for which + this object should be generating JSON for - this is useful to store + additional data if needed. + + Returns: + dict: JSON representation + """ + # NOTE: uses of Audio in this class are a temporary hack -- when Ref support moves up + # into Media itself we should get rid of them + from wandb import Image + from wandb.data_types import Audio + from wandb.sdk.wandb_run import Run + + json_obj: Dict[str, Any] = {} + + if self._caption is not None: + json_obj["caption"] = self._caption + + if isinstance(run, Run): + json_obj.update( + { + "_type": "file", # TODO(adrian): This isn't (yet) a real media type we support on the frontend. + "sha256": self._sha256, + "size": self._size, + } + ) + + artifact_entry_url = self._get_artifact_entry_ref_url() + if artifact_entry_url is not None: + json_obj["artifact_path"] = artifact_entry_url + artifact_entry_latest_url = self._get_artifact_entry_latest_ref_url() + if artifact_entry_latest_url is not None: + json_obj["_latest_artifact_path"] = artifact_entry_latest_url + + if artifact_entry_url is None or self.is_bound(): + assert self.is_bound(), ( + f"Value of type {type(self).__name__} must be bound to a run with bind_to_run() before being serialized to JSON." + ) + + assert self._run is run, ( + "We don't support referring to media files across runs." + ) + + # The following two assertions are guaranteed to pass + # by definition is_bound, but are needed for + # mypy to understand that these are strings below. + assert isinstance(self._path, str) + json_obj["path"] = LogicalPath( + os.path.relpath(self._path, self._run.dir) + ) + + elif isinstance(run, wandb.Artifact): + if self.file_is_set(): + # The following two assertions are guaranteed to pass + # by definition of the call above, but are needed for + # mypy to understand that these are strings below. + assert isinstance(self._path, str) + assert isinstance(self._sha256, str) + artifact = run # Checks if the concrete image has already been added to this artifact + name = artifact.get_added_local_path_name(self._path) + if name is None: + if self._is_tmp: + name = os.path.join( + self.get_media_subdir(), os.path.basename(self._path) + ) + else: + # If the files is not temporary, include the first 8 characters of the file's SHA256 to + # avoid name collisions. This way, if there are two images `dir1/img.png` and `dir2/img.png` + # we end up with a unique path for each. + name = os.path.join( + self.get_media_subdir(), + self._sha256[:20], + os.path.basename(self._path), + ) + + # if not, check to see if there is a source artifact for this object + if ( + self._artifact_source is not None + # and self._artifact_source.artifact != artifact + ): + default_root = self._artifact_source.artifact._default_root() + # if there is, get the name of the entry (this might make sense to move to a helper off artifact) + if self._path.startswith(default_root): + name = self._path[len(default_root) :] + name = name.lstrip(os.sep) + + # Add this image as a reference + path = self._artifact_source.artifact.get_entry(name) + artifact.add_reference(path.ref_url(), name=name) + elif ( + isinstance(self, Audio) or isinstance(self, Image) + ) and self.path_is_reference(self._path): + artifact.add_reference(self._path, name=name) + else: + entry = artifact.add_file( + self._path, name=name, is_tmp=self._is_tmp + ) + name = entry.path + + json_obj["path"] = name + json_obj["sha256"] = self._sha256 + json_obj["_type"] = self._log_type + return json_obj + + @classmethod + def from_json( + cls: Type["Media"], json_obj: dict, source_artifact: "Artifact" + ) -> "Media": + """Likely will need to override for any more complicated media objects.""" + return cls(source_artifact.get_entry(json_obj["path"]).download()) + + def __eq__(self, other: object) -> bool: + """Likely will need to override for any more complicated media objects.""" + return ( + isinstance(other, self.__class__) + and hasattr(self, "_sha256") + and hasattr(other, "_sha256") + and self._sha256 == other._sha256 + ) + + @staticmethod + def path_is_reference(path: Optional[Union[str, pathlib.Path]]) -> bool: + if path is None or isinstance(path, pathlib.Path): + return False + + return bool(path and re.match(r"^(gs|s3|https?)://", path)) + + +class BatchableMedia(Media): + """Media that is treated in batches. + + E.g. images and thumbnails. Apart from images, we just use these batches to help + organize files by name in the media directory. + """ + + def __init__( + self, + caption: Optional[str] = None, + ) -> None: + super().__init__(caption=caption) + + @classmethod + def seq_to_json( + cls: Type["BatchableMedia"], + seq: Sequence["BatchableMedia"], + run: "LocalRun", + key: str, + step: Union[int, str], + ) -> dict: + raise NotImplementedError + + +def _numpy_arrays_to_lists( + payload: Union[dict, Sequence, "np.ndarray"], +) -> Union[Sequence, dict, str, int, float, bool]: + # Casts all numpy arrays to lists so we don't convert them to histograms, primarily for Plotly + + if isinstance(payload, dict): + res = {} + for key, val in payload.items(): + res[key] = _numpy_arrays_to_lists(val) + return res + elif isinstance(payload, Sequence) and not isinstance(payload, str): + return [_numpy_arrays_to_lists(v) for v in payload] + elif util.is_numpy_array(payload): + if TYPE_CHECKING: + payload = cast("np.ndarray", payload) + return [ + _numpy_arrays_to_lists(v) + for v in (payload.tolist() if payload.ndim > 0 else [payload.tolist()]) + ] + # Protects against logging non serializable objects + elif isinstance(payload, Media): + return str(payload.__class__.__name__) + return payload # type: ignore diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/wb_value.py b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/wb_value.py new file mode 100644 index 0000000000000000000000000000000000000000..0a2c20f435f0319bab431b56cf35199ed065e2f5 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/base_types/wb_value.py @@ -0,0 +1,295 @@ +from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Type, Union + +from wandb import util +from wandb.sdk import wandb_setup + +if TYPE_CHECKING: # pragma: no cover + from wandb.sdk.artifacts.artifact import Artifact + + from ...wandb_run import Run as LocalRun + + TypeMappingType = Dict[str, Type["WBValue"]] + + +def _is_maybe_offline() -> bool: + """Guess whether wandb is configured to be offline. + + This is an anti-pattern because there is no library-level "offline" mode: + only runs can be offline. Online and offline runs can exist in the same + process. This function is a heuristic that works only if there is at most + one run in the process, and could otherwise produce unexpected results. + + Returns: + Whether the user likely configured wandb to be offline. + """ + singleton = wandb_setup.singleton() + + # First check: if there's a run, check if it is offline. + # + # This covers uses like `wandb.init(mode="offline")` which don't modify + # the singleton's settings. + if run := singleton.most_recent_active_run: + return run.offline + + # Second check: default to global defaults derived from environment + # variables or passed explicitly to `wandb.setup()`. + return singleton.settings._offline + + +def _server_accepts_client_ids() -> bool: + from packaging.version import parse + + # There are versions of W&B Server that cannot accept client IDs. Those versions of + # the backend have a max_cli_version of less than "0.11.0." If the backend cannot + # accept client IDs, manifests and artifact data would never be resolvable and lead + # to failed uploads. Our position in 2021/06/29 was to never lose data - and instead take the + # tradeoff in the UI. The results in tables not displaying media correctly, but + # the table can still be accessed via the .artifact op. + # + # The latest SDK version that is < "0.11.0" was released on 2021/06/29. + # AS OF NOW, 2024/11/06, we assume that all customer's server deployments accept + # client IDs. + + if _is_maybe_offline(): + singleton = wandb_setup.singleton() + + if run := singleton.most_recent_active_run: + return run._settings.allow_offline_artifacts + else: + return singleton.settings.allow_offline_artifacts + + # If the script is online, request the max_cli_version and ensure the server + # is of a high enough version. + max_cli_version = util._get_max_cli_version() + if max_cli_version is None: + return False + accepts_client_ids: bool = parse(max_cli_version) >= parse("0.11.0") + return accepts_client_ids + + +class _WBValueArtifactSource: + artifact: "Artifact" + name: Optional[str] + + def __init__(self, artifact: "Artifact", name: Optional[str] = None) -> None: + self.artifact = artifact + self.name = name + + +class _WBValueArtifactTarget: + artifact: "Artifact" + name: Optional[str] + + def __init__(self, artifact: "Artifact", name: Optional[str] = None) -> None: + self.artifact = artifact + self.name = name + + +class WBValue: + """Typed objects that can be logged with `wandb.log()` and visualized by wandb. + + The objects will be serialized as JSON and always have a _type attribute that + indicates how to interpret the other fields. + """ + + # Class Attributes + _type_mapping: ClassVar[Optional["TypeMappingType"]] = None + # override _log_type to indicate the type which the subclass deserializes + _log_type: ClassVar[Optional[str]] = None + + # Instance Attributes + _artifact_source: Optional[_WBValueArtifactSource] + _artifact_target: Optional[_WBValueArtifactTarget] + + def __init__(self) -> None: + self._artifact_source = None + self._artifact_target = None + + def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict: + """Serialize the object into a JSON blob. + + Uses current run or artifact to store additional data. + + Args: + run_or_artifact (wandb.Run | wandb.Artifact): the Run or Artifact for which + this object should be generating JSON for - this is useful to to store + additional data if needed. + + Returns: + dict: JSON representation + """ + raise NotImplementedError + + @classmethod + def from_json(cls, json_obj: dict, source_artifact: "Artifact") -> "WBValue": + """Deserialize a `json_obj` into it's class representation. + + If additional resources were stored in the `run_or_artifact` artifact during the + `to_json` call, then those resources should be in the `source_artifact`. + + Args: + json_obj (dict): A JSON dictionary to deserialize source_artifact + (wandb.Artifact): An artifact which will hold any additional + resources which were stored during the `to_json` function. + """ + raise NotImplementedError + + @classmethod + def with_suffix(cls: Type["WBValue"], name: str, filetype: str = "json") -> str: + """Get the name with the appropriate suffix. + + Args: + name (str): the name of the file + filetype (str, optional): the filetype to use. Defaults to "json". + + Returns: + str: a filename which is suffixed with it's `_log_type` followed by the + filetype. + """ + if cls._log_type is not None: + suffix = cls._log_type + "." + filetype + else: + suffix = filetype + if not name.endswith(suffix): + return name + "." + suffix + return name + + @staticmethod + def init_from_json( + json_obj: dict, source_artifact: "Artifact" + ) -> Optional["WBValue"]: + """Initialize a `WBValue` from a JSON blob based on the class that created it. + + Looks through all subclasses and tries to match the json obj with the class + which created it. It will then call that subclass' `from_json` method. + Importantly, this function will set the return object's `source_artifact` + attribute to the passed in source artifact. This is critical for artifact + bookkeeping. If you choose to create a wandb.Value via it's `from_json` method, + make sure to properly set this `artifact_source` to avoid data duplication. + + Args: + json_obj (dict): A JSON dictionary to deserialize. It must contain a `_type` + key. This is used to lookup the correct subclass to use. + source_artifact (wandb.Artifact): An artifact which will hold any additional + resources which were stored during the `to_json` function. + + Returns: + wandb.Value: a newly created instance of a subclass of wandb.Value + """ + class_option = WBValue.type_mapping().get(json_obj["_type"]) + if class_option is not None: + obj = class_option.from_json(json_obj, source_artifact) + obj._set_artifact_source(source_artifact) + return obj + + return None + + @staticmethod + def type_mapping() -> "TypeMappingType": + """Return a map from `_log_type` to subclass. Used to lookup correct types for deserialization. + + Returns: + dict: dictionary of str:class + """ + if WBValue._type_mapping is None: + WBValue._type_mapping = {} + frontier = [WBValue] + explored = set() + while len(frontier) > 0: + class_option = frontier.pop() + explored.add(class_option) + if class_option._log_type is not None: + WBValue._type_mapping[class_option._log_type] = class_option + for subclass in class_option.__subclasses__(): + if subclass not in explored: + frontier.append(subclass) + return WBValue._type_mapping + + def __eq__(self, other: object) -> bool: + return id(self) == id(other) + + def __ne__(self, other: object) -> bool: + return not self.__eq__(other) + + def to_data_array(self) -> List[Any]: + """Convert the object to a list of primitives representing the underlying data.""" + raise NotImplementedError + + def _set_artifact_source( + self, artifact: "Artifact", name: Optional[str] = None + ) -> None: + assert self._artifact_source is None, ( + f"Cannot update artifact_source. Existing source: {self._artifact_source.artifact}/{self._artifact_source.name}" + ) + self._artifact_source = _WBValueArtifactSource(artifact, name) + + def _set_artifact_target( + self, artifact: "Artifact", name: Optional[str] = None + ) -> None: + assert self._artifact_target is None, ( + f"Cannot update artifact_target. Existing target: {self._artifact_target.artifact}/{self._artifact_target.name}" + ) + self._artifact_target = _WBValueArtifactTarget(artifact, name) + + def _get_artifact_entry_ref_url(self) -> Optional[str]: + # If the object is coming from another artifact + if self._artifact_source and self._artifact_source.name: + ref_entry = self._artifact_source.artifact.get_entry( + type(self).with_suffix(self._artifact_source.name) + ) + return str(ref_entry.ref_url()) + # Else, if the object is destined for another artifact and we support client IDs + elif ( + self._artifact_target + and self._artifact_target.name + and self._artifact_target.artifact._client_id is not None + and self._artifact_target.artifact._final + and _server_accepts_client_ids() + ): + return f"wandb-client-artifact://{self._artifact_target.artifact._client_id}/{type(self).with_suffix(self._artifact_target.name)}" + # Else if we do not support client IDs, but online, then block on upload + # Note: this is old behavior just to stay backwards compatible + # with older server versions. This code path should be removed + # once those versions are no longer supported. This path uses a .wait + # which blocks the user process on artifact upload. + elif ( + self._artifact_target + and self._artifact_target.name + and self._artifact_target.artifact._is_draft_save_started() + and not _is_maybe_offline() + and not _server_accepts_client_ids() + ): + self._artifact_target.artifact.wait() + ref_entry = self._artifact_target.artifact.get_entry( + type(self).with_suffix(self._artifact_target.name) + ) + return str(ref_entry.ref_url()) + return None + + def _get_artifact_entry_latest_ref_url(self) -> Optional[str]: + if ( + self._artifact_target + and self._artifact_target.name + and self._artifact_target.artifact._client_id is not None + and self._artifact_target.artifact._final + and _server_accepts_client_ids() + ): + return f"wandb-client-artifact://{self._artifact_target.artifact._sequence_client_id}:latest/{type(self).with_suffix(self._artifact_target.name)}" + # Else if we do not support client IDs, then block on upload + # Note: this is old behavior just to stay backwards compatible + # with older server versions. This code path should be removed + # once those versions are no longer supported. This path uses a .wait + # which blocks the user process on artifact upload. + elif ( + self._artifact_target + and self._artifact_target.name + and self._artifact_target.artifact._is_draft_save_started() + and not _is_maybe_offline() + and not _server_accepts_client_ids() + ): + self._artifact_target.artifact.wait() + ref_entry = self._artifact_target.artifact.get_entry( + type(self).with_suffix(self._artifact_target.name) + ) + return str(ref_entry.ref_url()) + return None diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/bokeh.py b/lib/python3.12/site-packages/wandb/sdk/data_types/bokeh.py new file mode 100644 index 0000000000000000000000000000000000000000..dfaf6307e1107e3a07a866076731da2edfd94a63 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/bokeh.py @@ -0,0 +1,87 @@ +import codecs +import json +import os +import pathlib +from typing import TYPE_CHECKING, Union + +from wandb import util +from wandb.sdk.lib import runid + +from . import _dtypes +from ._private import MEDIA_TMP +from .base_types.media import Media + +if TYPE_CHECKING: + from bokeh import document, model + + +class Bokeh(Media): + """Wandb class for Bokeh plots. + + Args: + val: Bokeh plot + """ + + _log_type = "bokeh-file" + + def __init__( + self, + data_or_path: Union[ + str, + pathlib.Path, + "document.Document", + "model.Model", + ], + ): + super().__init__() + bokeh = util.get_module("bokeh", required=True) + if isinstance(data_or_path, (str, pathlib.Path)) and os.path.exists( + data_or_path + ): + data_or_path = str(data_or_path) + + with open(data_or_path) as file: + b_json = json.load(file) + self.b_obj = bokeh.document.Document.from_json(b_json) + self._set_file(data_or_path, is_tmp=False, extension=".bokeh.json") + elif isinstance(data_or_path, bokeh.model.Model): + _data = bokeh.document.Document() + _data.add_root(data_or_path) + # serialize/deserialize pairing followed by sorting attributes ensures + # that the file's sha's are equivalent in subsequent calls + self.b_obj = bokeh.document.Document.from_json(_data.to_json()) + b_json = self.b_obj.to_json() + if "references" in b_json["roots"]: + b_json["roots"]["references"].sort(key=lambda x: x["id"]) + + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".bokeh.json") + with codecs.open(tmp_path, "w", encoding="utf-8") as fp: + util.json_dump_safer(b_json, fp) + self._set_file(tmp_path, is_tmp=True, extension=".bokeh.json") + elif not isinstance(data_or_path, bokeh.document.Document): + raise TypeError( + "Bokeh constructor accepts Bokeh document/model or path to Bokeh json file" + ) + + def get_media_subdir(self): + return os.path.join("media", "bokeh") + + def to_json(self, run): + # TODO: (tss) this is getting redundant for all the media objects. We can probably + # pull this into Media#to_json and remove this type override for all the media types. + # There are only a few cases where the type is different between artifacts and runs. + json_dict = super().to_json(run) + json_dict["_type"] = self._log_type + return json_dict + + @classmethod + def from_json(cls, json_obj, source_artifact): + return cls(source_artifact.get_entry(json_obj["path"]).download()) + + +class _BokehFileType(_dtypes.Type): + name = "bokeh-file" + types = [Bokeh] + + +_dtypes.TypeRegistry.add(_BokehFileType) diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/graph.py b/lib/python3.12/site-packages/wandb/sdk/data_types/graph.py new file mode 100644 index 0000000000000000000000000000000000000000..d46c1f7fe86515a7f63f3eea3b5c271c5a24f2aa --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/graph.py @@ -0,0 +1,405 @@ +import codecs +import os +import pprint + +from wandb import util +from wandb.sdk.data_types._private import MEDIA_TMP +from wandb.sdk.data_types.base_types.media import Media, _numpy_arrays_to_lists +from wandb.sdk.data_types.base_types.wb_value import WBValue +from wandb.sdk.lib import runid + + +def _nest(thing): + # Use tensorflows nest function if available, otherwise just wrap object in an array""" + + tfutil = util.get_module("tensorflow.python.util") + if tfutil: + return tfutil.nest.flatten(thing) + else: + return [thing] + + +class Edge(WBValue): + """Edge used in `Graph`.""" + + def __init__(self, from_node, to_node): + self._attributes = {} + self.from_node = from_node + self.to_node = to_node + + def __repr__(self): + temp_attr = dict(self._attributes) + del temp_attr["from_node"] + del temp_attr["to_node"] + temp_attr["from_id"] = self.from_node.id + temp_attr["to_id"] = self.to_node.id + return str(temp_attr) + + def to_json(self, run=None): + return [self.from_node.id, self.to_node.id] + + @property + def name(self): + """Optional, not necessarily unique.""" + return self._attributes.get("name") + + @name.setter + def name(self, val): + self._attributes["name"] = val + return val + + @property + def from_node(self): + return self._attributes.get("from_node") + + @from_node.setter + def from_node(self, val): + self._attributes["from_node"] = val + return val + + @property + def to_node(self): + return self._attributes.get("to_node") + + @to_node.setter + def to_node(self, val): + self._attributes["to_node"] = val + return val + + +class Node(WBValue): + """Node used in `Graph`.""" + + def __init__( + self, + id=None, + name=None, + class_name=None, + size=None, + parameters=None, + output_shape=None, + is_output=None, + num_parameters=None, + node=None, + ): + self._attributes = {"name": None} + self.in_edges = {} # indexed by source node id + self.out_edges = {} # indexed by dest node id + # optional object (e.g. PyTorch Parameter or Module) that this Node represents + self.obj = None + + if node is not None: + self._attributes.update(node._attributes) + del self._attributes["id"] + self.obj = node.obj + + if id is not None: + self.id = id + if name is not None: + self.name = name + if class_name is not None: + self.class_name = class_name + if size is not None: + self.size = size + if parameters is not None: + self.parameters = parameters + if output_shape is not None: + self.output_shape = output_shape + if is_output is not None: + self.is_output = is_output + if num_parameters is not None: + self.num_parameters = num_parameters + + def to_json(self, run=None): + return self._attributes + + def __repr__(self): + return repr(self._attributes) + + @property + def id(self): + """Must be unique in the graph.""" + return self._attributes.get("id") + + @id.setter + def id(self, val): + self._attributes["id"] = val + return val + + @property + def name(self): + """Usually the type of layer or sublayer.""" + return self._attributes.get("name") + + @name.setter + def name(self, val): + self._attributes["name"] = val + return val + + @property + def class_name(self): + """Usually the type of layer or sublayer.""" + return self._attributes.get("class_name") + + @class_name.setter + def class_name(self, val): + self._attributes["class_name"] = val + return val + + @property + def functions(self): + return self._attributes.get("functions", []) + + @functions.setter + def functions(self, val): + self._attributes["functions"] = val + return val + + @property + def parameters(self): + return self._attributes.get("parameters", []) + + @parameters.setter + def parameters(self, val): + self._attributes["parameters"] = val + return val + + @property + def size(self): + return self._attributes.get("size") + + @size.setter + def size(self, val): + """Tensor size.""" + self._attributes["size"] = tuple(val) + return val + + @property + def output_shape(self): + return self._attributes.get("output_shape") + + @output_shape.setter + def output_shape(self, val): + """Tensor output_shape.""" + self._attributes["output_shape"] = val + return val + + @property + def is_output(self): + return self._attributes.get("is_output") + + @is_output.setter + def is_output(self, val): + """Tensor is_output.""" + self._attributes["is_output"] = val + return val + + @property + def num_parameters(self): + return self._attributes.get("num_parameters") + + @num_parameters.setter + def num_parameters(self, val): + """Tensor num_parameters.""" + self._attributes["num_parameters"] = val + return val + + @property + def child_parameters(self): + return self._attributes.get("child_parameters") + + @child_parameters.setter + def child_parameters(self, val): + """Tensor child_parameters.""" + self._attributes["child_parameters"] = val + return val + + @property + def is_constant(self): + return self._attributes.get("is_constant") + + @is_constant.setter + def is_constant(self, val): + """Tensor is_constant.""" + self._attributes["is_constant"] = val + return val + + @classmethod + def from_keras(cls, layer): + node = cls() + + try: + output_shape = layer.output_shape + except AttributeError: + output_shape = ["multiple"] + + node.id = layer.name + node.name = layer.name + node.class_name = layer.__class__.__name__ + node.output_shape = output_shape + node.num_parameters = layer.count_params() + + return node + + +class Graph(Media): + """Wandb class for graphs. + + This class is typically used for saving and displaying neural net models. It + represents the graph as an array of nodes and edges. The nodes can have + labels that can be visualized by wandb. + + Examples: + Import a keras model: + ``` + Graph.from_keras(keras_model) + ``` + + Attributes: + format (string): Format to help wandb display the graph nicely. + nodes ([wandb.Node]): List of wandb.Nodes + nodes_by_id (dict): dict of ids -> nodes + edges ([(wandb.Node, wandb.Node)]): List of pairs of nodes interpreted as edges + loaded (boolean): Flag to tell whether the graph is completely loaded + root (wandb.Node): root node of the graph + """ + + _log_type = "graph-file" + + def __init__(self, format="keras"): + super().__init__() + # LB: TODO: I think we should factor criterion and criterion_passed out + self.format = format + self.nodes = [] + self.nodes_by_id = {} + self.edges = [] + self.loaded = False + self.criterion = None + self.criterion_passed = False + self.root = None # optional root Node if applicable + + def _to_graph_json(self, run=None): + # Needs to be its own function for tests + return { + "format": self.format, + "nodes": [node.to_json() for node in self.nodes], + "edges": [edge.to_json() for edge in self.edges], + } + + def bind_to_run(self, *args, **kwargs): + data = self._to_graph_json() + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".graph.json") + data = _numpy_arrays_to_lists(data) + with codecs.open(tmp_path, "w", encoding="utf-8") as fp: + util.json_dump_safer(data, fp) + self._set_file(tmp_path, is_tmp=True, extension=".graph.json") + if self.is_bound(): + return + super().bind_to_run(*args, **kwargs) + + @classmethod + def get_media_subdir(cls): + return os.path.join("media", "graph") + + def to_json(self, run): + json_dict = super().to_json(run) + json_dict["_type"] = self._log_type + return json_dict + + def __getitem__(self, nid): + return self.nodes_by_id[nid] + + def pprint(self): + for edge in self.edges: + pprint.pprint(edge.attributes) # noqa: T203 + for node in self.nodes: + pprint.pprint(node.attributes) # noqa: T203 + + def add_node(self, node=None, **node_kwargs): + if node is None: + node = Node(**node_kwargs) + elif node_kwargs: + raise ValueError( + f"Only pass one of either node ({node}) or other keyword arguments ({node_kwargs})" + ) + self.nodes.append(node) + self.nodes_by_id[node.id] = node + + return node + + def add_edge(self, from_node, to_node): + edge = Edge(from_node, to_node) + self.edges.append(edge) + + return edge + + @classmethod + def from_keras(cls, model): + # TODO: his method requires a refactor to work with the keras 3. + graph = cls() + # Shamelessly copied (then modified) from keras/keras/utils/layer_utils.py + sequential_like = cls._is_sequential(model) + + relevant_nodes = None + if not sequential_like: + relevant_nodes = [] + for v in model._nodes_by_depth.values(): + relevant_nodes += v + + layers = model.layers + for i in range(len(layers)): + node = Node.from_keras(layers[i]) + if hasattr(layers[i], "_inbound_nodes"): + for in_node in layers[i]._inbound_nodes: + if relevant_nodes and in_node not in relevant_nodes: + # node is not part of the current network + continue + for in_layer in _nest(in_node.inbound_layers): + inbound_keras_node = Node.from_keras(in_layer) + + if inbound_keras_node.id not in graph.nodes_by_id: + graph.add_node(inbound_keras_node) + inbound_node = graph.nodes_by_id[inbound_keras_node.id] + + graph.add_edge(inbound_node, node) + graph.add_node(node) + return graph + + @classmethod + def _is_sequential(cls, model): + sequential_like = True + + if ( + model.__class__.__name__ != "Sequential" + and hasattr(model, "_is_graph_network") + and model._is_graph_network + ): + nodes_by_depth = model._nodes_by_depth.values() + nodes = [] + for v in nodes_by_depth: + # TensorFlow2 doesn't insure inbound is always a list + inbound = v[0].inbound_layers + if not hasattr(inbound, "__len__"): + inbound = [inbound] + if (len(v) > 1) or (len(v) == 1 and len(inbound) > 1): + # if the model has multiple nodes + # or if the nodes have multiple inbound_layers + # the model is no longer sequential + sequential_like = False + break + nodes += v + if sequential_like: + # search for shared layers + for layer in model.layers: + flag = False + if hasattr(layer, "_inbound_nodes"): + for node in layer._inbound_nodes: + if node in nodes: + if flag: + sequential_like = False + break + else: + flag = True + if not sequential_like: + break + return sequential_like diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__init__.py b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a3fe116402fa24b14121a58fc7918ab7e756fd14 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/bounding_boxes_2d.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/bounding_boxes_2d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6e386277211784a668a0868fe76eab9ad07ca2c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/bounding_boxes_2d.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/classes.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/classes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c825a0e45599761665b2df31d946f3609c96057c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/classes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/image_mask.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/image_mask.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1af8d3e0135eb31e825ef96bb12724c67ac8ada3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/__pycache__/image_mask.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/bounding_boxes_2d.py b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/bounding_boxes_2d.py new file mode 100644 index 0000000000000000000000000000000000000000..eb090ead2f3bd4099478a02f6ded00d32cab7a7a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/bounding_boxes_2d.py @@ -0,0 +1,329 @@ +import numbers +from typing import TYPE_CHECKING, Optional, Type, Union + +import wandb +from wandb import util +from wandb.util import has_num + +from ..base_types.json_metadata import JSONMetadata + +if TYPE_CHECKING: # pragma: no cover + from wandb.sdk.artifacts.artifact import Artifact + + from ...wandb_run import Run as LocalRun + + +def _convert_pytorch_tensor_to_list(box_data): + for box in box_data: + if ( + "position" in box + and "middle" in box["position"] + and util.is_pytorch_tensor_typename( + util.get_full_typename(box["position"]["middle"]) + ) + ): + box["position"]["middle"] = box["position"]["middle"].tolist() + + +class BoundingBoxes2D(JSONMetadata): + """Format images with 2D bounding box overlays for logging to W&B. + + Args: + val: (dictionary) A dictionary of the following form: + box_data: (list of dictionaries) One dictionary for each bounding box, containing: + position: (dictionary) the position and size of the bounding box, in one of two formats + Note that boxes need not all use the same format. + {"minX", "minY", "maxX", "maxY"}: (dictionary) A set of coordinates defining + the upper and lower bounds of the box (the bottom left and top right corners) + {"middle", "width", "height"}: (dictionary) A set of coordinates defining the + center and dimensions of the box, with "middle" as a list [x, y] for the + center point and "width" and "height" as numbers + domain: (string) One of two options for the bounding box coordinate domain + null: By default, or if no argument is passed, the coordinate domain + is assumed to be relative to the original image, expressing this box as a fraction + or percentage of the original image. This means all coordinates and dimensions + passed into the "position" argument are floating point numbers between 0 and 1. + "pixel": (string literal) The coordinate domain is set to the pixel space. This means all + coordinates and dimensions passed into "position" are integers within the bounds + of the image dimensions. + class_id: (integer) The class label id for this box + scores: (dictionary of string to number, optional) A mapping of named fields + to numerical values (float or int), can be used for filtering boxes in the UI + based on a range of values for the corresponding field + box_caption: (string, optional) A string to be displayed as the label text above this + box in the UI, often composed of the class label, class name, and/or scores + + class_labels: (dictionary, optional) A map of integer class labels to their readable class names + + key: (string) + The readable name or id for this set of bounding boxes (e.g. predictions, ground_truth) + + Examples: + ### Log bounding boxes for a single image + + ```python + import numpy as np + import wandb + + run = wandb.init() + image = np.random.randint(low=0, high=256, size=(200, 300, 3)) + + class_labels = {0: "person", 1: "car", 2: "road", 3: "building"} + + img = wandb.Image( + image, + boxes={ + "predictions": { + "box_data": [ + { + # one box expressed in the default relative/fractional domain + "position": { + "minX": 0.1, + "maxX": 0.2, + "minY": 0.3, + "maxY": 0.4, + }, + "class_id": 1, + "box_caption": class_labels[1], + "scores": {"acc": 0.2, "loss": 1.2}, + }, + { + # another box expressed in the pixel domain + "position": { + "middle": [150, 20], + "width": 68, + "height": 112, + }, + "domain": "pixel", + "class_id": 3, + "box_caption": "a building", + "scores": {"acc": 0.5, "loss": 0.7}, + }, + # Log as many boxes an as needed + ], + "class_labels": class_labels, + } + }, + ) + + run.log({"driving_scene": img}) + ``` + + ### Log a bounding box overlay to a Table + + ```python + import numpy as np + import wandb + + run = wandb.init() + image = np.random.randint(low=0, high=256, size=(200, 300, 3)) + + class_labels = {0: "person", 1: "car", 2: "road", 3: "building"} + + class_set = wandb.Classes( + [ + {"name": "person", "id": 0}, + {"name": "car", "id": 1}, + {"name": "road", "id": 2}, + {"name": "building", "id": 3}, + ] + ) + + img = wandb.Image( + image, + boxes={ + "predictions": { + "box_data": [ + { + # one box expressed in the default relative/fractional domain + "position": { + "minX": 0.1, + "maxX": 0.2, + "minY": 0.3, + "maxY": 0.4, + }, + "class_id": 1, + "box_caption": class_labels[1], + "scores": {"acc": 0.2, "loss": 1.2}, + }, + { + # another box expressed in the pixel domain + "position": { + "middle": [150, 20], + "width": 68, + "height": 112, + }, + "domain": "pixel", + "class_id": 3, + "box_caption": "a building", + "scores": {"acc": 0.5, "loss": 0.7}, + }, + # Log as many boxes an as needed + ], + "class_labels": class_labels, + } + }, + classes=class_set, + ) + + table = wandb.Table(columns=["image"]) + table.add_data(img) + run.log({"driving_scene": table}) + ``` + """ + + _log_type = "bounding-boxes" + # TODO: when the change is made to have this produce a dict with a _type, define + # it here as _log_type, associate it in to_json + + def __init__(self, val: dict, key: str) -> None: + """Initialize a BoundingBoxes object. + + The input dictionary `val` should contain the keys: + box_data: a list of dictionaries, each of which describes a bounding box. + class_labels: (optional) A map of integer class labels to their readable + class names. + + Each bounding box dictionary should contain the following keys: + position: (dictionary) the position and size of the bounding box. + domain: (string) One of two options for the bounding box coordinate domain. + class_id: (integer) The class label id for this box. + scores: (dictionary of string to number, optional) A mapping of named fields + to numerical values (float or int). + box_caption: (optional) The label text, often composed of the class label, + class name, and/or scores. + + The position dictionary should be in one of two formats: + {"minX", "minY", "maxX", "maxY"}: (dictionary) A set of coordinates defining + the upper and lower bounds of the box (the bottom left and top right + corners). + {"middle", "width", "height"}: (dictionary) A set of coordinates defining + the center and dimensions of the box, with "middle" as a list [x, y] for + the center point and "width" and "height" as numbers. + Note that boxes need not all use the same format. + + Args: + val: (dictionary) A dictionary containing the bounding box data. + key: (string) The readable name or id for this set of bounding boxes (e.g. + predictions, ground_truth) + """ + # Pytorch tensors are not serializable to json, + # so we convert them to lists to avoid errors later on. + _convert_pytorch_tensor_to_list(val.get("box_data", [])) + super().__init__(val) + + self._val = val["box_data"] + self._key = key + # Add default class mapping + if "class_labels" not in val: + np = util.get_module( + "numpy", required="Bounding box support requires numpy" + ) + classes = ( + np.unique(list(box["class_id"] for box in val["box_data"])) + .astype(np.int32) + .tolist() + ) + class_labels = {c: "class_" + str(c) for c in classes} + self._class_labels = class_labels + else: + self._class_labels = val["class_labels"] + + def bind_to_run( + self, + run: "LocalRun", + key: Union[int, str], + step: Union[int, str], + id_: Optional[Union[int, str]] = None, + ignore_copy_err: Optional[bool] = None, + ) -> None: + # bind_to_run key argument is the Image parent key + # the self._key value is the mask's sub key + super().bind_to_run(run, key, step, id_=id_, ignore_copy_err=ignore_copy_err) + run._add_singleton( + "bounding_box/class_labels", + str(key) + "_wandb_delimeter_" + self._key, + self._class_labels, + ) + + @classmethod + def type_name(cls) -> str: + return "boxes2D" + + def validate(self, val: dict) -> bool: + # Optional argument + if "class_labels" in val: + for k, v in list(val["class_labels"].items()): + if (not isinstance(k, numbers.Number)) or (not isinstance(v, str)): + raise TypeError( + "Class labels must be a dictionary of numbers to string" + ) + + boxes = val["box_data"] + if not isinstance(boxes, list): + raise TypeError("Boxes must be a list") + + for box in boxes: + # Required arguments + error_str = ( + "Each box must contain a position with: middle, width, and height or \ + \nminX, maxX, minY, maxY." + ) + if "position" not in box: + raise TypeError(error_str) + else: + valid = False + if ( + "middle" in box["position"] + and len(box["position"]["middle"]) == 2 + and has_num(box["position"], "width") + and has_num(box["position"], "height") + ): + valid = True + elif ( + has_num(box["position"], "minX") + and has_num(box["position"], "maxX") + and has_num(box["position"], "minY") + and has_num(box["position"], "maxY") + ): + valid = True + + if not valid: + raise TypeError(error_str) + + # Optional arguments + if ("scores" in box) and not isinstance(box["scores"], dict): + raise TypeError("Box scores must be a dictionary") + elif "scores" in box: + for k, v in list(box["scores"].items()): + if not isinstance(k, str): + raise TypeError("A score key must be a string") + if not isinstance(v, numbers.Number): + raise TypeError("A score value must be a number") + + if ("class_id" in box) and not isinstance(box["class_id"], int): + raise TypeError("A box's class_id must be an integer") + + # Optional + if ("box_caption" in box) and not isinstance(box["box_caption"], str): + raise TypeError("A box's caption must be a string") + return True + + def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict: + from wandb.sdk.wandb_run import Run + + if isinstance(run_or_artifact, Run): + return super().to_json(run_or_artifact) + elif isinstance(run_or_artifact, wandb.Artifact): + # TODO (tim): I would like to log out a proper dictionary representing this object, but don't + # want to mess with the visualizations that are currently available in the UI. This really should output + # an object with a _type key. Will need to push this change to the UI first to ensure backwards compat + return self._val + else: + raise TypeError("to_json accepts wandb_run.Run or wandb.Artifact") + + @classmethod + def from_json( + cls: Type["BoundingBoxes2D"], json_obj: dict, source_artifact: "Artifact" + ) -> "BoundingBoxes2D": + return cls({"box_data": json_obj}, "") diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/classes.py b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/classes.py new file mode 100644 index 0000000000000000000000000000000000000000..54a22054aa44ad58de5b46a72b1123d73bc8836d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/classes.py @@ -0,0 +1,159 @@ +import os +from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Type, Union + +from .. import _dtypes +from ..base_types.media import Media + +if TYPE_CHECKING: # pragma: no cover + from wandb.sdk.artifacts.artifact import Artifact + + from ...wandb_run import Run as LocalRun + + +class Classes(Media): + _log_type = "classes" + + _class_set: Sequence[dict] + + def __init__(self, class_set: Sequence[dict]) -> None: + """Classes is holds class metadata intended to be used in concert with other objects when visualizing artifacts. + + Args: + class_set (list): list of dicts in the form of {"id":int|str, "name":str} + """ + super().__init__() + for class_obj in class_set: + assert "id" in class_obj and "name" in class_obj + self._class_set = class_set + + @classmethod + def from_json( + cls: Type["Classes"], + json_obj: dict, + source_artifact: Optional["Artifact"], + ) -> "Classes": + return cls(json_obj.get("class_set")) # type: ignore + + def to_json(self, run_or_artifact: Optional[Union["LocalRun", "Artifact"]]) -> dict: + json_obj = {} + # This is a bit of a hack to allow _ClassesIdType to + # be able to operate fully without an artifact in play. + # In all other cases, artifact should be a true artifact. + if run_or_artifact is not None: + json_obj = super().to_json(run_or_artifact) + json_obj["_type"] = Classes._log_type + json_obj["class_set"] = self._class_set + return json_obj + + def get_type(self) -> "_ClassesIdType": + return _ClassesIdType(self) + + def __ne__(self, other: object) -> bool: + return not self.__eq__(other) + + def __eq__(self, other: object) -> bool: + if isinstance(other, Classes): + return self._class_set == other._class_set + else: + return False + + +class _ClassesIdType(_dtypes.Type): + name = "classesId" + legacy_names = ["wandb.Classes_id"] + types = [Classes] + + def __init__( + self, + classes_obj: Optional[Classes] = None, + valid_ids: Optional["_dtypes.UnionType"] = None, + ): + if valid_ids is None: + valid_ids = _dtypes.UnionType() + elif isinstance(valid_ids, list): + valid_ids = _dtypes.UnionType( + [_dtypes.ConstType(item) for item in valid_ids] + ) + elif isinstance(valid_ids, _dtypes.UnionType): + valid_ids = valid_ids + else: + raise TypeError("valid_ids must be None, list, or UnionType") + + if classes_obj is None: + classes_obj = Classes( + [ + {"id": _id.params["val"], "name": str(_id.params["val"])} + for _id in valid_ids.params["allowed_types"] + ] + ) + elif not isinstance(classes_obj, Classes): + raise TypeError("valid_ids must be None, or instance of Classes") + else: + valid_ids = _dtypes.UnionType( + [ + _dtypes.ConstType(class_obj["id"]) + for class_obj in classes_obj._class_set + ] + ) + + self.wb_classes_obj_ref = classes_obj + self.params.update({"valid_ids": valid_ids}) + + def assign(self, py_obj: Optional[Any] = None) -> "_dtypes.Type": + return self.assign_type(_dtypes.ConstType(py_obj)) + + def assign_type(self, wb_type: "_dtypes.Type") -> "_dtypes.Type": + valid_ids = self.params["valid_ids"].assign_type(wb_type) + if not isinstance(valid_ids, _dtypes.InvalidType): + return self + + return _dtypes.InvalidType() + + @classmethod + def from_obj(cls, py_obj: Optional[Any] = None) -> "_dtypes.Type": + return cls(py_obj) + + def to_json(self, artifact: Optional["Artifact"] = None) -> Dict[str, Any]: + cl_dict = super().to_json(artifact) + # TODO (tss): Refactor this block with the similar one in wandb.Image. + # This is a bit of a smell that the classes object does not follow + # the same file-pattern as other media types. + if artifact is not None: + class_name = os.path.join("media", "cls") + classes_entry = artifact.add(self.wb_classes_obj_ref, class_name) + cl_dict["params"]["classes_obj"] = { + "type": "classes-file", + "path": classes_entry.path, + "digest": classes_entry.digest, # is this needed really? + } + else: + cl_dict["params"]["classes_obj"] = self.wb_classes_obj_ref.to_json(artifact) + return cl_dict + + @classmethod + def from_json( + cls, + json_dict: Dict[str, Any], + artifact: Optional["Artifact"] = None, + ) -> "_dtypes.Type": + classes_obj = None + if ( + json_dict.get("params", {}).get("classes_obj", {}).get("type") + == "classes-file" + ): + if artifact is not None: + classes_obj = artifact.get( + json_dict.get("params", {}).get("classes_obj", {}).get("path") + ) + assert classes_obj is None or isinstance(classes_obj, Classes) + else: + raise RuntimeError("Expected artifact to be non-null.") + else: + classes_obj = Classes.from_json( + json_dict["params"]["classes_obj"], artifact + ) + + return cls(classes_obj) + + +_dtypes.TypeRegistry.add(_ClassesIdType) diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/image_mask.py b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/image_mask.py new file mode 100644 index 0000000000000000000000000000000000000000..d26e013f588a7883331af7b1c5f475976c402a45 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/helper_types/image_mask.py @@ -0,0 +1,253 @@ +import numbers +import os +from typing import TYPE_CHECKING, Optional, Type, Union + +import wandb +from wandb import util +from wandb.sdk.lib import runid + +from .._private import MEDIA_TMP +from ..base_types.media import Media + +if TYPE_CHECKING: # pragma: no cover + from wandb.sdk.artifacts.artifact import Artifact + + from ...wandb_run import Run as LocalRun + + +class ImageMask(Media): + """Format image masks or overlays for logging to W&B. + + Args: + val: (dictionary) + One of these two keys to represent the image: + mask_data : (2D numpy array) The mask containing an integer class label + for each pixel in the image + path : (string) The path to a saved image file of the mask + class_labels : (dictionary of integers to strings, optional) A mapping of the + integer class labels in the mask to readable class names. These will default + to class_0, class_1, class_2, etc. + + key: (string) + The readable name or id for this mask type (e.g. predictions, ground_truth) + + Examples: + ### Logging a single masked image + + ```python + import numpy as np + import wandb + + run = wandb.init() + image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8) + predicted_mask = np.empty((100, 100), dtype=np.uint8) + ground_truth_mask = np.empty((100, 100), dtype=np.uint8) + + predicted_mask[:50, :50] = 0 + predicted_mask[50:, :50] = 1 + predicted_mask[:50, 50:] = 2 + predicted_mask[50:, 50:] = 3 + + ground_truth_mask[:25, :25] = 0 + ground_truth_mask[25:, :25] = 1 + ground_truth_mask[:25, 25:] = 2 + ground_truth_mask[25:, 25:] = 3 + + class_labels = {0: "person", 1: "tree", 2: "car", 3: "road"} + + masked_image = wandb.Image( + image, + masks={ + "predictions": { + "mask_data": predicted_mask, + "class_labels": class_labels, + }, + "ground_truth": { + "mask_data": ground_truth_mask, + "class_labels": class_labels, + }, + }, + ) + run.log({"img_with_masks": masked_image}) + ``` + + ### Log a masked image inside a Table + + ```python + import numpy as np + import wandb + + run = wandb.init() + image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8) + predicted_mask = np.empty((100, 100), dtype=np.uint8) + ground_truth_mask = np.empty((100, 100), dtype=np.uint8) + + predicted_mask[:50, :50] = 0 + predicted_mask[50:, :50] = 1 + predicted_mask[:50, 50:] = 2 + predicted_mask[50:, 50:] = 3 + + ground_truth_mask[:25, :25] = 0 + ground_truth_mask[25:, :25] = 1 + ground_truth_mask[:25, 25:] = 2 + ground_truth_mask[25:, 25:] = 3 + + class_labels = {0: "person", 1: "tree", 2: "car", 3: "road"} + + class_set = wandb.Classes( + [ + {"name": "person", "id": 0}, + {"name": "tree", "id": 1}, + {"name": "car", "id": 2}, + {"name": "road", "id": 3}, + ] + ) + + masked_image = wandb.Image( + image, + masks={ + "predictions": { + "mask_data": predicted_mask, + "class_labels": class_labels, + }, + "ground_truth": { + "mask_data": ground_truth_mask, + "class_labels": class_labels, + }, + }, + classes=class_set, + ) + + table = wandb.Table(columns=["image"]) + table.add_data(masked_image) + run.log({"random_field": table}) + ``` + """ + + _log_type = "mask" + + def __init__(self, val: dict, key: str) -> None: + """Initialize an ImageMask object. + + Args: + val: (dictionary) One of these two keys to represent the image: + mask_data : (2D numpy array) The mask containing an integer class label + for each pixel in the image + path : (string) The path to a saved image file of the mask + class_labels : (dictionary of integers to strings, optional) A mapping + of the integer class labels in the mask to readable class names. + These will default to class_0, class_1, class_2, etc. + + key: (string) + The readable name or id for this mask type (e.g. predictions, ground_truth) + """ + super().__init__() + + if "path" in val: + self._set_file(val["path"]) + else: + np = util.get_module("numpy", required="Image mask support requires numpy") + + if util.is_pytorch_tensor_typename( + util.get_full_typename(val["mask_data"]) + ): + val["mask_data"] = val["mask_data"].cpu().numpy() + + # Add default class mapping + if "class_labels" not in val: + classes = np.unique(val["mask_data"]).astype(np.int32).tolist() + class_labels = {c: "class_" + str(c) for c in classes} + val["class_labels"] = class_labels + + self.validate(val) + self._val = val + self._key = key + + ext = "." + self.type_name() + ".png" + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ext) + + pil_image = util.get_module( + "PIL.Image", + required='wandb.Image needs the PIL package. To get it, run "pip install pillow".', + ) + image = pil_image.fromarray(val["mask_data"].astype(np.int8), mode="L") + + image.save(tmp_path, transparency=None) + self._set_file(tmp_path, is_tmp=True, extension=ext) + + def bind_to_run( + self, + run: "LocalRun", + key: Union[int, str], + step: Union[int, str], + id_: Optional[Union[int, str]] = None, + ignore_copy_err: Optional[bool] = None, + ) -> None: + # bind_to_run key argument is the Image parent key + # the self._key value is the mask's sub key + super().bind_to_run(run, key, step, id_=id_, ignore_copy_err=ignore_copy_err) + if hasattr(self, "_val") and "class_labels" in self._val: + class_labels = self._val["class_labels"] + + run._add_singleton( + "mask/class_labels", + str(key) + "_wandb_delimeter_" + self._key, + class_labels, + ) + + @classmethod + def get_media_subdir(cls: Type["ImageMask"]) -> str: + return os.path.join("media", "images", cls.type_name()) + + @classmethod + def from_json( + cls: Type["ImageMask"], json_obj: dict, source_artifact: "Artifact" + ) -> "ImageMask": + return cls( + {"path": source_artifact.get_entry(json_obj["path"]).download()}, + key="", + ) + + def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict: + from wandb.sdk.wandb_run import Run + + json_dict = super().to_json(run_or_artifact) + + if isinstance(run_or_artifact, Run): + json_dict["_type"] = self.type_name() + return json_dict + elif isinstance(run_or_artifact, wandb.Artifact): + # Nothing special to add (used to add "digest", but no longer used.) + return json_dict + else: + raise TypeError("to_json accepts wandb_run.Run or wandb.Artifact") + + @classmethod + def type_name(cls: Type["ImageMask"]) -> str: + return cls._log_type + + def validate(self, val: dict) -> bool: + np = util.get_module("numpy", required="Image mask support requires numpy") + # 2D Make this work with all tensor(like) types + if "mask_data" not in val: + raise TypeError( + 'Missing key "mask_data": An image mask requires mask data: a 2D array representing the predictions' + ) + else: + error_str = "mask_data must be a 2D array" + shape = val["mask_data"].shape + if len(shape) != 2: + raise TypeError(error_str) + if not ( + (val["mask_data"] >= 0).all() and (val["mask_data"] <= 255).all() + ) and issubclass(val["mask_data"].dtype.type, np.integer): + raise TypeError("Mask data must be integers between 0 and 255") + + # Optional argument + if "class_labels" in val: + for k, v in list(val["class_labels"].items()): + if (not isinstance(k, numbers.Number)) or (not isinstance(v, str)): + raise TypeError( + "Class labels must be a dictionary of numbers to strings" + ) + return True diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/histogram.py b/lib/python3.12/site-packages/wandb/sdk/data_types/histogram.py new file mode 100644 index 0000000000000000000000000000000000000000..7dc13a98914d8a32f8f6eba6f4e8e9f389deb97d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/histogram.py @@ -0,0 +1,94 @@ +import sys +from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Union + +from wandb import util + +from .base_types.wb_value import WBValue + +if TYPE_CHECKING: # pragma: no cover + import numpy as np + + from wandb.sdk.artifacts.artifact import Artifact + + from ..wandb_run import Run as LocalRun + + NumpyHistogram = Tuple[np.ndarray, np.ndarray] + + +class Histogram(WBValue): + """wandb class for histograms. + + This object works just like numpy's histogram function + https://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram.html + + Examples: + Generate histogram from a sequence + ```python + wandb.Histogram([1, 2, 3]) + ``` + + Efficiently initialize from np.histogram. + ```python + hist = np.histogram(data) + wandb.Histogram(np_histogram=hist) + ``` + + Args: + sequence: (array_like) input data for histogram + np_histogram: (numpy histogram) alternative input of a precomputed histogram + num_bins: (int) Number of bins for the histogram. The default number of bins + is 64. The maximum number of bins is 512 + + Attributes: + bins: ([float]) edges of bins + histogram: ([int]) number of elements falling in each bin + """ + + MAX_LENGTH: int = 512 + _log_type = "histogram" + + def __init__( + self, + sequence: Optional[Sequence] = None, + np_histogram: Optional["NumpyHistogram"] = None, + num_bins: int = 64, + ) -> None: + if np_histogram: + if len(np_histogram) == 2: + self.histogram = ( + np_histogram[0].tolist() + if hasattr(np_histogram[0], "tolist") + else np_histogram[0] + ) + self.bins = ( + np_histogram[1].tolist() + if hasattr(np_histogram[1], "tolist") + else np_histogram[1] + ) + else: + raise ValueError( + "Expected np_histogram to be a tuple of (values, bin_edges) or sequence to be specified" + ) + else: + np = util.get_module( + "numpy", required="Auto creation of histograms requires numpy" + ) + + histogram, bins = np.histogram(sequence, bins=num_bins) + self.histogram = histogram.tolist() + self.bins = bins.tolist() + if len(self.histogram) > self.MAX_LENGTH: + raise ValueError(f"The maximum length of a histogram is {self.MAX_LENGTH}") + if len(self.histogram) + 1 != len(self.bins): + raise ValueError("len(bins) must be len(histogram) + 1") + + def to_json(self, run: Optional[Union["LocalRun", "Artifact"]] = None) -> dict: + return {"_type": self._log_type, "values": self.histogram, "bins": self.bins} + + def __sizeof__(self) -> int: + """Estimated size in bytes. + + Currently the factor of 1.7 is used to account for the JSON encoding. We use + this in tb_watcher.TBHistory. + """ + return int((sys.getsizeof(self.histogram) + sys.getsizeof(self.bins)) * 1.7) diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/html.py b/lib/python3.12/site-packages/wandb/sdk/data_types/html.py new file mode 100644 index 0000000000000000000000000000000000000000..fa7b482c2df72560d416ce4d0817fd694a99de0d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/html.py @@ -0,0 +1,142 @@ +import os +import pathlib +from typing import TYPE_CHECKING, Sequence, Type, Union + +from wandb.sdk.lib import filesystem, runid + +from . import _dtypes +from ._private import MEDIA_TMP +from .base_types.media import BatchableMedia + +if TYPE_CHECKING: # pragma: no cover + from typing import TextIO + + from wandb.sdk.artifacts.artifact import Artifact + + from ..wandb_run import Run as LocalRun + + +class Html(BatchableMedia): + """A class for logging HTML content to W&B.""" + + _log_type = "html-file" + + def __init__( + self, + data: Union[str, pathlib.Path, "TextIO"], + inject: bool = True, + data_is_not_path: bool = False, + ) -> None: + """Creates a W&B HTML object. + + It can be initialized by providing a path to a file: + ``` + with wandb.init() as run: + run.log({"html": wandb.Html("./index.html")}) + ``` + + Alternatively, it can be initialized by providing literal HTML, + in either a string or IO object: + ``` + with wandb.init() as run: + run.log({"html": wandb.Html("

Hello, world!

")}) + ``` + + Args: + data: + A string that is a path to a file with the extension ".html", + or a string or IO object containing literal HTML. + inject: Add a stylesheet to the HTML object. If set + to False the HTML will pass through unchanged. + data_is_not_path: If set to False, the data will be + treated as a path to a file. + """ + super().__init__() + data_is_path = ( + isinstance(data, (str, pathlib.Path)) + and os.path.isfile(data) + and os.path.splitext(data)[1] == ".html" + ) and not data_is_not_path + data_path = "" + if data_is_path: + data_path = str(data) + with open(data_path, encoding="utf-8") as file: + self.html = file.read() + elif isinstance(data, str): + self.html = data + elif hasattr(data, "read"): + if hasattr(data, "seek"): + data.seek(0) + self.html = data.read() + else: + raise ValueError("data must be a string or an io object") + + if inject: + self.inject_head() + + if inject or not data_is_path: + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".html") + with open(tmp_path, "w", encoding="utf-8") as out: + out.write(self.html) + + self._set_file(tmp_path, is_tmp=True) + else: + self._set_file(data_path, is_tmp=False) + + def inject_head(self) -> None: + join = "" + if "" in self.html: + parts = self.html.split("", 1) + parts[0] = parts[0] + "" + elif "" in self.html: + parts = self.html.split("", 1) + parts[0] = parts[0] + "" + parts[1] = "" + parts[1] + else: + parts = ["", self.html] + parts.insert( + 1, + '', + ) + self.html = join.join(parts).strip() + + @classmethod + def get_media_subdir(cls: Type["Html"]) -> str: + return os.path.join("media", "html") + + def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict: + json_dict = super().to_json(run_or_artifact) + json_dict["_type"] = self._log_type + return json_dict + + @classmethod + def from_json( + cls: Type["Html"], json_obj: dict, source_artifact: "Artifact" + ) -> "Html": + return cls(source_artifact.get_entry(json_obj["path"]).download(), inject=False) + + @classmethod + def seq_to_json( + cls: Type["Html"], + seq: Sequence["BatchableMedia"], + run: "LocalRun", + key: str, + step: Union[int, str], + ) -> dict: + base_path = os.path.join(run.dir, cls.get_media_subdir()) + filesystem.mkdir_exists_ok(base_path) + + meta = { + "_type": "html", + "count": len(seq), + "html": [h.to_json(run) for h in seq], + } + return meta + + +class _HtmlFileType(_dtypes.Type): + name = "html-file" + types = [Html] + + +_dtypes.TypeRegistry.add(_HtmlFileType) diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/image.py b/lib/python3.12/site-packages/wandb/sdk/data_types/image.py new file mode 100644 index 0000000000000000000000000000000000000000..1ce18c0905e47fbeb308ad042906fe92c654d73d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/image.py @@ -0,0 +1,938 @@ +import hashlib +import logging +import os +import pathlib +from io import BytesIO +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Type, Union, cast +from urllib import parse + +from packaging.version import parse as parse_version + +import wandb +from wandb import util +from wandb.sdk.lib import hashutil, runid +from wandb.sdk.lib.paths import LogicalPath + +from . import _dtypes +from ._private import MEDIA_TMP +from .base_types.media import BatchableMedia, Media +from .helper_types.bounding_boxes_2d import BoundingBoxes2D +from .helper_types.classes import Classes +from .helper_types.image_mask import ImageMask + +if TYPE_CHECKING: # pragma: no cover + import matplotlib # type: ignore + import numpy as np + import torch # type: ignore + from PIL.Image import Image as PILImage + + from wandb.sdk.artifacts.artifact import Artifact + + from ..wandb_run import Run as LocalRun + + ImageDataType = Union[ + "matplotlib.artist.Artist", "PILImage", "TorchTensorType", "np.ndarray" + ] + ImageDataOrPathType = Union[str, pathlib.Path, "Image", ImageDataType] + TorchTensorType = Union["torch.Tensor", "torch.Variable"] + + +def _warn_on_invalid_data_range( + data: "np.ndarray", + normalize: bool = True, +) -> None: + if not normalize: + return + + np = util.get_module( + "numpy", + required="wandb.Image requires numpy if not supplying PIL Images: pip install numpy", + ) + + if np.min(data) < 0 or np.max(data) > 255: + wandb.termwarn( + "Data passed to `wandb.Image` should consist of values in the range [0, 255], " + "image data will be normalized to this range, " + "but behavior will be removed in a future version of wandb.", + repeat=False, + ) + + +def _guess_and_rescale_to_0_255(data: "np.ndarray") -> "np.ndarray": + """Guess the image's format and rescale its values to the range [0, 255]. + + This is an unfortunate design flaw carried forward for backward + compatibility. A better design would have been to document the expected + data format and not mangle the data provided by the user. + + If given data in the range [0, 1], we multiply all values by 255 + and round down to get integers. + + If given data in the range [-1, 1], we rescale it by mapping -1 to 0 and + 1 to 255, then round down to get integers. + + We clip and round all other data. + """ + try: + import numpy as np + except ImportError: + raise wandb.Error( + "wandb.Image requires numpy if not supplying PIL images: pip install numpy" + ) from None + + data_min: float = data.min() + data_max: float = data.max() + + if 0 <= data_min and data_max <= 1: + return (data * 255).astype(np.uint8) + + elif -1 <= data_min and data_max <= 1: + return (255 * 0.5 * (data + 1)).astype(np.uint8) + + else: + return data.clip(0, 255).astype(np.uint8) + + +def _convert_to_uint8(data: "np.ndarray") -> "np.ndarray": + np = util.get_module( + "numpy", + required="wandb.Image requires numpy if not supplying PIL Images: pip install numpy", + ) + return data.astype(np.uint8) + + +def _server_accepts_image_filenames(run: "LocalRun") -> bool: + if run.offline: + return True + + # Newer versions of wandb accept large image filenames arrays + # but older versions would have issues with this. + max_cli_version = util._get_max_cli_version() + if max_cli_version is None: + return False + + accepts_image_filenames: bool = parse_version(max_cli_version) >= parse_version( + "0.12.10" + ) + return accepts_image_filenames + + +def _server_accepts_artifact_path(run: "LocalRun") -> bool: + if run.offline: + return False + + max_cli_version = util._get_max_cli_version() + if max_cli_version is None: + return False + + return parse_version(max_cli_version) >= parse_version("0.12.14") + + +class Image(BatchableMedia): + """A class for logging images to W&B.""" + + MAX_ITEMS = 108 + + # PIL limit + MAX_DIMENSION = 65500 + + _log_type = "image-file" + + format: Optional[str] + _grouping: Optional[int] + _caption: Optional[str] + _width: Optional[int] + _height: Optional[int] + _image: Optional["PILImage"] + _classes: Optional["Classes"] + _boxes: Optional[Dict[str, "BoundingBoxes2D"]] + _masks: Optional[Dict[str, "ImageMask"]] + _file_type: Optional[str] + + def __init__( + self, + data_or_path: "ImageDataOrPathType", + mode: Optional[str] = None, + caption: Optional[str] = None, + grouping: Optional[int] = None, + classes: Optional[Union["Classes", Sequence[dict]]] = None, + boxes: Optional[Union[Dict[str, "BoundingBoxes2D"], Dict[str, dict]]] = None, + masks: Optional[Union[Dict[str, "ImageMask"], Dict[str, dict]]] = None, + file_type: Optional[str] = None, + normalize: bool = True, + ) -> None: + """Initialize a wandb.Image object. + + Args: + data_or_path: Accepts numpy array/pytorch tensor of image data, + a PIL image object, or a path to an image file. + + If a numpy array or pytorch tensor is provided, + the image data will be saved to the given file type. + If the values are not in the range [0, 255] or all values are in the range [0, 1], + the image pixel values will be normalized to the range [0, 255] + unless `normalize` is set to False. + - pytorch tensor should be in the format (channel, height, width) + - numpy array should be in the format (height, width, channel) + mode: The PIL mode for an image. Most common are "L", "RGB", + "RGBA". Full explanation at https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes + caption: Label for display of image. + grouping: The grouping number for the image. + classes: A list of class information for the image, + used for labeling bounding boxes, and image masks. + boxes: A dictionary containing bounding box information for the image. + see: https://docs.wandb.ai/ref/python/data-types/boundingboxes2d/ + masks: A dictionary containing mask information for the image. + see: https://docs.wandb.ai/ref/python/data-types/imagemask/ + file_type: The file type to save the image as. + This parameter has no effect if data_or_path is a path to an image file. + normalize: If True, normalize the image pixel values to fall within the range of [0, 255]. + Normalize is only applied if data_or_path is a numpy array or pytorch tensor. + + Examples: + ### Create a wandb.Image from a numpy array + ```python + import numpy as np + import wandb + + with wandb.init() as run: + examples = [] + for i in range(3): + pixels = np.random.randint(low=0, high=256, size=(100, 100, 3)) + image = wandb.Image(pixels, caption=f"random field {i}") + examples.append(image) + run.log({"examples": examples}) + ``` + + ### Create a wandb.Image from a PILImage + ```python + import numpy as np + from PIL import Image as PILImage + import wandb + + with wandb.init() as run: + examples = [] + for i in range(3): + pixels = np.random.randint( + low=0, high=256, size=(100, 100, 3), dtype=np.uint8 + ) + pil_image = PILImage.fromarray(pixels, mode="RGB") + image = wandb.Image(pil_image, caption=f"random field {i}") + examples.append(image) + run.log({"examples": examples}) + ``` + + ### log .jpg rather than .png (default) + ```python + import numpy as np + import wandb + + with wandb.init() as run: + examples = [] + for i in range(3): + pixels = np.random.randint(low=0, high=256, size=(100, 100, 3)) + image = wandb.Image( + pixels, caption=f"random field {i}", file_type="jpg" + ) + examples.append(image) + run.log({"examples": examples}) + ``` + """ + super().__init__(caption=caption) + # TODO: We should remove grouping, it's a terrible name and I don't + # think anyone uses it. + + self._grouping = None + self._width = None + self._height = None + self._image = None + self._classes = None + self._boxes = None + self._masks = None + self._file_type = None + + # Allows the user to pass an Image object as the first parameter and have a perfect copy, + # only overriding additional metadata passed in. If this pattern is compelling, we can generalize. + if isinstance(data_or_path, Image): + self._initialize_from_wbimage(data_or_path) + elif isinstance(data_or_path, (str, pathlib.Path)): + data_or_path = str(data_or_path) + + if self.path_is_reference(data_or_path): + self._initialize_from_reference(data_or_path) + else: + self._initialize_from_path(data_or_path) + else: + self._initialize_from_data(data_or_path, mode, file_type, normalize) + self._set_initialization_meta( + grouping, caption, classes, boxes, masks, file_type + ) + + def _set_initialization_meta( + self, + grouping: Optional[int] = None, + caption: Optional[str] = None, + classes: Optional[Union["Classes", Sequence[dict]]] = None, + boxes: Optional[Union[Dict[str, "BoundingBoxes2D"], Dict[str, dict]]] = None, + masks: Optional[Union[Dict[str, "ImageMask"], Dict[str, dict]]] = None, + file_type: Optional[str] = None, + ) -> None: + if grouping is not None: + self._grouping = grouping + + total_classes = {} + + if boxes: + if not isinstance(boxes, dict): + raise ValueError('Images "boxes" argument must be a dictionary') + boxes_final: Dict[str, BoundingBoxes2D] = {} + for key in boxes: + box_item = boxes[key] + if isinstance(box_item, BoundingBoxes2D): + boxes_final[key] = box_item + elif isinstance(box_item, dict): + # TODO: Consider injecting top-level classes if user-provided is empty + boxes_final[key] = BoundingBoxes2D(box_item, key) + total_classes.update(boxes_final[key]._class_labels) + self._boxes = boxes_final + + if masks: + if not isinstance(masks, dict): + raise ValueError('Images "masks" argument must be a dictionary') + masks_final: Dict[str, ImageMask] = {} + for key in masks: + mask_item = masks[key] + if isinstance(mask_item, ImageMask): + masks_final[key] = mask_item + elif isinstance(mask_item, dict): + # TODO: Consider injecting top-level classes if user-provided is empty + masks_final[key] = ImageMask(mask_item, key) + if hasattr(masks_final[key], "_val"): + total_classes.update(masks_final[key]._val["class_labels"]) + self._masks = masks_final + + if classes is not None: + if isinstance(classes, Classes): + total_classes.update( + {val["id"]: val["name"] for val in classes._class_set} + ) + else: + total_classes.update({val["id"]: val["name"] for val in classes}) + + if len(total_classes.keys()) > 0: + self._classes = Classes( + [ + {"id": key, "name": total_classes[key]} + for key in total_classes.keys() + ] + ) + if self.image is not None: + self._width, self._height = self.image.size + self._free_ram() + + def _initialize_from_wbimage(self, wbimage: "Image") -> None: + self._grouping = wbimage._grouping + self._caption = wbimage._caption + self._width = wbimage._width + self._height = wbimage._height + self._image = wbimage._image + self._classes = wbimage._classes + self._path = wbimage._path + self._is_tmp = wbimage._is_tmp + self._extension = wbimage._extension + self._sha256 = wbimage._sha256 + self._size = wbimage._size + self.format = wbimage.format + self._file_type = wbimage._file_type + self._artifact_source = wbimage._artifact_source + self._artifact_target = wbimage._artifact_target + + # We do not want to implicitly copy boxes or masks, just the image-related data. + # self._boxes = wbimage._boxes + # self._masks = wbimage._masks + + def _initialize_from_path(self, path: str) -> None: + pil_image = util.get_module( + "PIL.Image", + required='wandb.Image needs the PIL package. To get it, run "pip install pillow".', + ) + self._set_file(path, is_tmp=False) + self._image = pil_image.open(path) + assert self._image is not None + self._image.load() + ext = os.path.splitext(path)[1][1:] + self.format = ext + + def _initialize_from_reference(self, path: str) -> None: + self._path = path + self._is_tmp = False + self._sha256 = hashlib.sha256(path.encode("utf-8")).hexdigest() + path = parse.urlparse(path).path + ext = path.split("/")[-1].split(".")[-1] + self.format = ext + + def _initialize_from_data( + self, + data: "ImageDataType", + mode: Optional[str] = None, + file_type: Optional[str] = None, + normalize: bool = True, + ) -> None: + pil_image = util.get_module( + "PIL.Image", + required='wandb.Image needs the PIL package. To get it, run "pip install pillow".', + ) + + accepted_formats = ["png", "jpg", "jpeg", "bmp"] + self.format = file_type or "png" + + if self.format not in accepted_formats: + raise ValueError(f"file_type must be one of {accepted_formats}") + + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + "." + self.format) + + if util.is_matplotlib_typename(util.get_full_typename(data)): + buf = BytesIO() + util.ensure_matplotlib_figure(data).savefig(buf, format=self.format) + self._image = pil_image.open(buf) + elif isinstance(data, pil_image.Image): + self._image = data + elif util.is_pytorch_tensor_typename(util.get_full_typename(data)): + if hasattr(data, "requires_grad") and data.requires_grad: + data = data.detach() # type: ignore + if hasattr(data, "dtype") and str(data.dtype) == "torch.uint8": + data = data.to(float) # type: ignore [union-attr] + mode = mode or self.guess_mode(data, file_type) + data = data.permute(1, 2, 0).cpu().numpy() # type: ignore [union-attr] + + _warn_on_invalid_data_range(data, normalize) + + data = _guess_and_rescale_to_0_255(data) if normalize else data # type: ignore [arg-type] + data = _convert_to_uint8(data) + + if data.ndim > 2: + data = data.squeeze() + + self._image = pil_image.fromarray( + data, + mode=mode, + ) + else: + if hasattr(data, "numpy"): # TF data eager tensors + data = data.numpy() + if data.ndim > 2: # type: ignore [union-attr] + # get rid of trivial dimensions as a convenience + data = data.squeeze() # type: ignore [union-attr] + + _warn_on_invalid_data_range(data, normalize) # type: ignore [arg-type] + + mode = mode or self.guess_mode(data, file_type) + data = _guess_and_rescale_to_0_255(data) if normalize else data # type: ignore [arg-type] + data = _convert_to_uint8(data) # type: ignore [arg-type] + self._image = pil_image.fromarray( + data, + mode=mode, + ) + + assert self._image is not None + self._image.save(tmp_path, transparency=None) + self._set_file(tmp_path, is_tmp=True) + + @classmethod + def from_json( + cls: Type["Image"], json_obj: dict, source_artifact: "Artifact" + ) -> "Image": + classes: Optional[Classes] = None + if json_obj.get("classes") is not None: + value = source_artifact.get(json_obj["classes"]["path"]) + assert isinstance(value, (type(None), Classes)) + classes = value + + masks = json_obj.get("masks") + _masks: Optional[Dict[str, ImageMask]] = None + if masks: + _masks = {} + for key in masks: + _masks[key] = ImageMask.from_json(masks[key], source_artifact) + _masks[key]._set_artifact_source(source_artifact) + _masks[key]._key = key + + boxes = json_obj.get("boxes") + _boxes: Optional[Dict[str, BoundingBoxes2D]] = None + if boxes: + _boxes = {} + for key in boxes: + _boxes[key] = BoundingBoxes2D.from_json(boxes[key], source_artifact) + _boxes[key]._key = key + + return cls( + source_artifact.get_entry(json_obj["path"]).download(), + caption=json_obj.get("caption"), + grouping=json_obj.get("grouping"), + classes=classes, + boxes=_boxes, + masks=_masks, + ) + + @classmethod + def get_media_subdir(cls: Type["Image"]) -> str: + return os.path.join("media", "images") + + def bind_to_run( + self, + run: "LocalRun", + key: Union[int, str], + step: Union[int, str], + id_: Optional[Union[int, str]] = None, + ignore_copy_err: Optional[bool] = None, + ) -> None: + # For Images, we are going to avoid copying the image file to the run. + # We should make this common functionality for all media types, but that + # requires a broader UI refactor. This model can easily be moved to the + # higher level Media class, but that will require every UI surface area + # that depends on the `path` to be able to instead consume + # `artifact_path`. I (Tim) think the media panel makes up most of this + # space, but there are also custom charts, and maybe others. Let's + # commit to getting all that fixed up before moving this to the top + # level Media class. + if self.path_is_reference(self._path): + raise ValueError( + "Image media created by a reference to external storage cannot currently be added to a run" + ) + + if ( + not _server_accepts_artifact_path(run) + or self._get_artifact_entry_ref_url() is None + ): + super().bind_to_run(run, key, step, id_, ignore_copy_err=ignore_copy_err) + if self._boxes is not None: + for i, k in enumerate(self._boxes): + id_ = f"{id_}{i}" if id_ is not None else None + self._boxes[k].bind_to_run( + run, key, step, id_, ignore_copy_err=ignore_copy_err + ) + + if self._masks is not None: + for i, k in enumerate(self._masks): + id_ = f"{id_}{i}" if id_ is not None else None + self._masks[k].bind_to_run( + run, key, step, id_, ignore_copy_err=ignore_copy_err + ) + + def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict: + from wandb.sdk.wandb_run import Run + + json_dict = super().to_json(run_or_artifact) + json_dict["_type"] = Image._log_type + json_dict["format"] = self.format + + if self._width is not None: + json_dict["width"] = self._width + if self._height is not None: + json_dict["height"] = self._height + if self._grouping: + json_dict["grouping"] = self._grouping + + if isinstance(run_or_artifact, wandb.Artifact): + artifact = run_or_artifact + if ( + self._masks is not None or self._boxes is not None + ) and self._classes is None: + raise ValueError( + "classes must be passed to wandb.Image which have masks or bounding boxes when adding to artifacts" + ) + + if self._classes is not None: + class_id = hashutil._md5( + str(self._classes._class_set).encode("utf-8") + ).hexdigest() + class_name = os.path.join( + "media", + "classes", + class_id + "_cls", + ) + classes_entry = artifact.add(self._classes, class_name) + json_dict["classes"] = { + "type": "classes-file", + "path": classes_entry.path, + "digest": classes_entry.digest, + } + + elif not isinstance(run_or_artifact, Run): + raise TypeError("to_json accepts wandb_run.Run or wandb_artifact.Artifact") + + if self._boxes: + json_dict["boxes"] = { + k: box.to_json(run_or_artifact) for (k, box) in self._boxes.items() + } + if self._masks: + json_dict["masks"] = { + k: mask.to_json(run_or_artifact) for (k, mask) in self._masks.items() + } + return json_dict + + def guess_mode( + self, + data: Union["np.ndarray", "torch.Tensor"], + file_type: Optional[str] = None, + ) -> str: + """Guess what type of image the np.array is representing.""" + # TODO: do we want to support dimensions being at the beginning of the array? + ndims = data.ndim + if util.is_pytorch_tensor_typename(util.get_full_typename(data)): + # Torch tenors typically have the channels dimension first + num_channels = data.shape[0] + else: + num_channels = data.shape[-1] + + if ndims == 2 or num_channels == 1: + return "L" + elif num_channels == 3: + return "RGB" + elif num_channels == 4: + if file_type in ["jpg", "jpeg"]: + wandb.termwarn( + "JPEG format does not support transparency. " + "Ignoring alpha channel.", + repeat=False, + ) + return "RGB" + else: + return "RGBA" + else: + raise ValueError( + f"Un-supported shape for image conversion {list(data.shape)}" + ) + + @classmethod + def seq_to_json( + cls: Type["Image"], + seq: Sequence["BatchableMedia"], + run: "LocalRun", + key: str, + step: Union[int, str], + ) -> dict: + """Combine a list of images into a meta dictionary object describing the child images.""" + if TYPE_CHECKING: + seq = cast(Sequence["Image"], seq) + + jsons = [obj.to_json(run) for obj in seq] + + media_dir = cls.get_media_subdir() + + for obj in jsons: + expected = LogicalPath(media_dir) + if "path" in obj and not obj["path"].startswith(expected): + raise ValueError( + "Files in an array of Image's must be in the {} directory, not {}".format( + cls.get_media_subdir(), obj["path"] + ) + ) + + num_images_to_log = len(seq) + width, height = seq[0].image.size # type: ignore + format = jsons[0]["format"] + + def size_equals_image(image: "Image") -> bool: + img_width, img_height = image.image.size # type: ignore + return img_width == width and img_height == height + + sizes_match = all(size_equals_image(img) for img in seq) + if not sizes_match: + logging.warning( + "Images sizes do not match. This will causes images to be display incorrectly in the UI." + ) + + meta = { + "_type": "images/separated", + "width": width, + "height": height, + "format": format, + "count": num_images_to_log, + } + if _server_accepts_image_filenames(run): + meta["filenames"] = [ + obj.get("path", obj.get("artifact_path")) for obj in jsons + ] + else: + wandb.termwarn( + "Unable to log image array filenames. In some cases, this can prevent images from being " + "viewed in the UI. Please upgrade your wandb server", + repeat=False, + ) + + captions = Image.all_captions(seq) + + if captions: + meta["captions"] = captions + + all_masks = Image.all_masks(seq, run, key, step) + + if all_masks: + meta["all_masks"] = all_masks + + all_boxes = Image.all_boxes(seq, run, key, step) + + if all_boxes: + meta["all_boxes"] = all_boxes + + return meta + + @classmethod + def all_masks( + cls: Type["Image"], + images: Sequence["Image"], + run: "LocalRun", + run_key: str, + step: Union[int, str], + ) -> Union[List[Optional[dict]], bool]: + all_mask_groups: List[Optional[dict]] = [] + for image in images: + if image._masks: + mask_group = {} + for k in image._masks: + mask = image._masks[k] + mask_group[k] = mask.to_json(run) + all_mask_groups.append(mask_group) + else: + all_mask_groups.append(None) + if all_mask_groups and not all(x is None for x in all_mask_groups): + return all_mask_groups + else: + return False + + @classmethod + def all_boxes( + cls: Type["Image"], + images: Sequence["Image"], + run: "LocalRun", + run_key: str, + step: Union[int, str], + ) -> Union[List[Optional[dict]], bool]: + all_box_groups: List[Optional[dict]] = [] + for image in images: + if image._boxes: + box_group = {} + for k in image._boxes: + box = image._boxes[k] + box_group[k] = box.to_json(run) + all_box_groups.append(box_group) + else: + all_box_groups.append(None) + if all_box_groups and not all(x is None for x in all_box_groups): + return all_box_groups + else: + return False + + @classmethod + def all_captions( + cls: Type["Image"], images: Sequence["Media"] + ) -> Union[bool, Sequence[Optional[str]]]: + return cls.captions(images) + + def __ne__(self, other: object) -> bool: + return not self.__eq__(other) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Image): + return False + else: + if self.path_is_reference(self._path) and self.path_is_reference( + other._path + ): + return self._path == other._path + self_image = self.image + other_image = other.image + if self_image is not None: + self_image = list(self_image.getdata()) # type: ignore + if other_image is not None: + other_image = list(other_image.getdata()) # type: ignore + + return ( + self._grouping == other._grouping + and self._caption == other._caption + and self._width == other._width + and self._height == other._height + and self_image == other_image + and self._classes == other._classes + ) + + def to_data_array(self) -> List[Any]: + res = [] + if self.image is not None: + data = list(self.image.getdata()) + for i in range(self.image.height): + res.append(data[i * self.image.width : (i + 1) * self.image.width]) + self._free_ram() + return res + + def _free_ram(self) -> None: + if self._path is not None: + self._image = None + + @property + def image(self) -> Optional["PILImage"]: + if self._image is None: + if self._path is not None and not self.path_is_reference(self._path): + pil_image = util.get_module( + "PIL.Image", + required='wandb.Image needs the PIL package. To get it, run "pip install pillow".', + ) + self._image = pil_image.open(self._path) + self._image.load() + return self._image + + +# Custom dtypes for typing system +class _ImageFileType(_dtypes.Type): + name = "image-file" + legacy_names = ["wandb.Image"] + types = [Image] + + def __init__( + self, + box_layers=None, + box_score_keys=None, + mask_layers=None, + class_map=None, + **kwargs, + ): + box_layers = box_layers or {} + box_score_keys = box_score_keys or [] + mask_layers = mask_layers or {} + class_map = class_map or {} + + if isinstance(box_layers, _dtypes.ConstType): + box_layers = box_layers._params["val"] + if not isinstance(box_layers, dict): + raise TypeError("box_layers must be a dict") + else: + box_layers = _dtypes.ConstType( + {layer_key: set(box_layers[layer_key]) for layer_key in box_layers} + ) + + if isinstance(mask_layers, _dtypes.ConstType): + mask_layers = mask_layers._params["val"] + if not isinstance(mask_layers, dict): + raise TypeError("mask_layers must be a dict") + else: + mask_layers = _dtypes.ConstType( + {layer_key: set(mask_layers[layer_key]) for layer_key in mask_layers} + ) + + if isinstance(box_score_keys, _dtypes.ConstType): + box_score_keys = box_score_keys._params["val"] + if not isinstance(box_score_keys, list) and not isinstance(box_score_keys, set): + raise TypeError("box_score_keys must be a list or a set") + else: + box_score_keys = _dtypes.ConstType(set(box_score_keys)) + + if isinstance(class_map, _dtypes.ConstType): + class_map = class_map._params["val"] + if not isinstance(class_map, dict): + raise TypeError("class_map must be a dict") + else: + class_map = _dtypes.ConstType(class_map) + + self.params.update( + { + "box_layers": box_layers, + "box_score_keys": box_score_keys, + "mask_layers": mask_layers, + "class_map": class_map, + } + ) + + def assign_type(self, wb_type=None): + if isinstance(wb_type, _ImageFileType): + box_layers_self = self.params["box_layers"].params["val"] or {} + box_score_keys_self = self.params["box_score_keys"].params["val"] or [] + mask_layers_self = self.params["mask_layers"].params["val"] or {} + class_map_self = self.params["class_map"].params["val"] or {} + + box_layers_other = wb_type.params["box_layers"].params["val"] or {} + box_score_keys_other = wb_type.params["box_score_keys"].params["val"] or [] + mask_layers_other = wb_type.params["mask_layers"].params["val"] or {} + class_map_other = wb_type.params["class_map"].params["val"] or {} + + # Merge the class_ids from each set of box_layers + box_layers = { + str(key): set( + list(box_layers_self.get(key, [])) + + list(box_layers_other.get(key, [])) + ) + for key in set( + list(box_layers_self.keys()) + list(box_layers_other.keys()) + ) + } + + # Merge the class_ids from each set of mask_layers + mask_layers = { + str(key): set( + list(mask_layers_self.get(key, [])) + + list(mask_layers_other.get(key, [])) + ) + for key in set( + list(mask_layers_self.keys()) + list(mask_layers_other.keys()) + ) + } + + # Merge the box score keys + box_score_keys = set(list(box_score_keys_self) + list(box_score_keys_other)) + + # Merge the class_map + class_map = { + str(key): class_map_self.get(key, class_map_other.get(key, None)) + for key in set( + list(class_map_self.keys()) + list(class_map_other.keys()) + ) + } + + return _ImageFileType(box_layers, box_score_keys, mask_layers, class_map) + + return _dtypes.InvalidType() + + @classmethod + def from_obj(cls, py_obj): + if not isinstance(py_obj, Image): + raise TypeError("py_obj must be a wandb.Image") + else: + if hasattr(py_obj, "_boxes") and py_obj._boxes: + box_layers = { + str(key): set(py_obj._boxes[key]._class_labels.keys()) + for key in py_obj._boxes.keys() + } + box_score_keys = { + key + for val in py_obj._boxes.values() + for box in val._val + for key in box.get("scores", {}).keys() + } + + else: + box_layers = {} + box_score_keys = set() + + if hasattr(py_obj, "_masks") and py_obj._masks: + mask_layers = { + str(key): set( + py_obj._masks[key]._val["class_labels"].keys() + if hasattr(py_obj._masks[key], "_val") + else [] + ) + for key in py_obj._masks.keys() + } + else: + mask_layers = {} + + if hasattr(py_obj, "_classes") and py_obj._classes: + class_set = { + str(item["id"]): item["name"] for item in py_obj._classes._class_set + } + else: + class_set = {} + + return cls(box_layers, box_score_keys, mask_layers, class_set) + + +_dtypes.TypeRegistry.add(_ImageFileType) diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/molecule.py b/lib/python3.12/site-packages/wandb/sdk/data_types/molecule.py new file mode 100644 index 0000000000000000000000000000000000000000..de23ecacae8c8af34bbcf833a724f69a14285368 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/molecule.py @@ -0,0 +1,237 @@ +import io +import os +import pathlib +from typing import TYPE_CHECKING, Optional, Sequence, Type, Union + +from wandb import util +from wandb.sdk.lib import runid +from wandb.sdk.lib.paths import LogicalPath + +from ._private import MEDIA_TMP +from .base_types.media import BatchableMedia, Media + +if TYPE_CHECKING: # pragma: no cover + from typing import TextIO + + import rdkit.Chem # type: ignore + + from wandb.sdk.artifacts.artifact import Artifact + + from ..wandb_run import Run as LocalRun + + RDKitDataType = Union[str, "rdkit.Chem.rdchem.Mol"] + + +class Molecule(BatchableMedia): + """Wandb class for 3D Molecular data. + + Args: + data_or_path: (pathlib.Path, string, io) + Molecule can be initialized from a file name or an io object. + caption: (string) + Caption associated with the molecule for display. + """ + + SUPPORTED_TYPES = { + "pdb", + "pqr", + "mmcif", + "mcif", + "cif", + "sdf", + "sd", + "gro", + "mol2", + "mmtf", + } + SUPPORTED_RDKIT_TYPES = {"mol", "sdf"} + _log_type = "molecule-file" + + def __init__( + self, + data_or_path: Union[str, pathlib.Path, "TextIO"], + caption: Optional[str] = None, + **kwargs: str, + ) -> None: + super().__init__(caption=caption) + + if hasattr(data_or_path, "name"): + # if the file has a path, we just detect the type and copy it from there + data_or_path = data_or_path.name + + if hasattr(data_or_path, "read"): + if hasattr(data_or_path, "seek"): + data_or_path.seek(0) + molecule = data_or_path.read() + + extension = kwargs.pop("file_type", None) + if extension is None: + raise ValueError( + "Must pass file_type keyword argument when using io objects." + ) + if extension not in Molecule.SUPPORTED_TYPES: + raise ValueError( + "Molecule 3D only supports files of the type: " + + ", ".join(Molecule.SUPPORTED_TYPES) + ) + + tmp_path = os.path.join( + MEDIA_TMP.name, runid.generate_id() + "." + extension + ) + with open(tmp_path, "w") as f: + f.write(molecule) + + self._set_file(tmp_path, is_tmp=True) + elif isinstance(data_or_path, (str, pathlib.Path)): + data_or_path = str(data_or_path) + + extension = os.path.splitext(data_or_path)[1][1:] + if extension not in Molecule.SUPPORTED_TYPES: + raise ValueError( + "Molecule only supports files of the type: " + + ", ".join(Molecule.SUPPORTED_TYPES) + ) + + self._set_file(data_or_path, is_tmp=False) + else: + raise ValueError("Data must be file name or a file object") + + @classmethod + def from_rdkit( + cls, + data_or_path: "RDKitDataType", + caption: Optional[str] = None, + convert_to_3d_and_optimize: bool = True, + mmff_optimize_molecule_max_iterations: int = 200, + ) -> "Molecule": + """Convert RDKit-supported file/object types to wandb.Molecule. + + Args: + data_or_path: (string, rdkit.Chem.rdchem.Mol) + Molecule can be initialized from a file name or an rdkit.Chem.rdchem.Mol object. + caption: (string) + Caption associated with the molecule for display. + convert_to_3d_and_optimize: (bool) + Convert to rdkit.Chem.rdchem.Mol with 3D coordinates. + This is an expensive operation that may take a long time for complicated molecules. + mmff_optimize_molecule_max_iterations: (int) + Number of iterations to use in rdkit.Chem.AllChem.MMFFOptimizeMolecule + """ + rdkit_chem = util.get_module( + "rdkit.Chem", + required='wandb.Molecule needs the rdkit-pypi package. To get it, run "pip install rdkit-pypi".', + ) + rdkit_chem_all_chem = util.get_module( + "rdkit.Chem.AllChem", + required='wandb.Molecule needs the rdkit-pypi package. To get it, run "pip install rdkit-pypi".', + ) + + if isinstance(data_or_path, str): + # path to a file? + path = pathlib.Path(data_or_path) + extension = path.suffix.split(".")[-1] + if extension not in Molecule.SUPPORTED_RDKIT_TYPES: + raise ValueError( + "Molecule.from_rdkit only supports files of the type: " + + ", ".join(Molecule.SUPPORTED_RDKIT_TYPES) + ) + # use the appropriate method + if extension == "sdf": + with rdkit_chem.SDMolSupplier(data_or_path) as supplier: + molecule = next(supplier) # get only the first molecule + else: + molecule = getattr(rdkit_chem, f"MolFrom{extension.capitalize()}File")( + data_or_path + ) + elif isinstance(data_or_path, rdkit_chem.rdchem.Mol): + molecule = data_or_path + else: + raise TypeError("Data must be file name or an rdkit.Chem.rdchem.Mol object") + + if convert_to_3d_and_optimize: + molecule = rdkit_chem.AddHs(molecule) + rdkit_chem_all_chem.EmbedMolecule(molecule) + rdkit_chem_all_chem.MMFFOptimizeMolecule( + molecule, + maxIters=mmff_optimize_molecule_max_iterations, + ) + # convert to the pdb format supported by Molecule + pdb_block = rdkit_chem.rdmolfiles.MolToPDBBlock(molecule) + + return cls(io.StringIO(pdb_block), caption=caption, file_type="pdb") + + @classmethod + def from_smiles( + cls, + data: str, + caption: Optional[str] = None, + sanitize: bool = True, + convert_to_3d_and_optimize: bool = True, + mmff_optimize_molecule_max_iterations: int = 200, + ) -> "Molecule": + """Convert SMILES string to wandb.Molecule. + + Args: + data: (string) + SMILES string. + caption: (string) + Caption associated with the molecule for display + sanitize: (bool) + Check if the molecule is chemically reasonable by the RDKit's definition. + convert_to_3d_and_optimize: (bool) + Convert to rdkit.Chem.rdchem.Mol with 3D coordinates. + This is an expensive operation that may take a long time for complicated molecules. + mmff_optimize_molecule_max_iterations: (int) + Number of iterations to use in rdkit.Chem.AllChem.MMFFOptimizeMolecule + """ + rdkit_chem = util.get_module( + "rdkit.Chem", + required='wandb.Molecule needs the rdkit-pypi package. To get it, run "pip install rdkit-pypi".', + ) + molecule = rdkit_chem.MolFromSmiles(data, sanitize=sanitize) + if molecule is None: + raise ValueError("Unable to parse the SMILES string.") + + return cls.from_rdkit( + data_or_path=molecule, + caption=caption, + convert_to_3d_and_optimize=convert_to_3d_and_optimize, + mmff_optimize_molecule_max_iterations=mmff_optimize_molecule_max_iterations, + ) + + @classmethod + def get_media_subdir(cls: Type["Molecule"]) -> str: + return os.path.join("media", "molecule") + + def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict: + json_dict = super().to_json(run_or_artifact) + json_dict["_type"] = self._log_type + return json_dict + + @classmethod + def seq_to_json( + cls: Type["Molecule"], + seq: Sequence["BatchableMedia"], + run: "LocalRun", + key: str, + step: Union[int, str], + ) -> dict: + seq = list(seq) + + jsons = [obj.to_json(run) for obj in seq] + + for obj in jsons: + expected = LogicalPath(cls.get_media_subdir()) + if not obj["path"].startswith(expected): + raise ValueError( + "Files in an array of Molecule's must be in the {} directory, not {}".format( + cls.get_media_subdir(), obj["path"] + ) + ) + + return { + "_type": "molecule", + "filenames": [obj["path"] for obj in jsons], + "count": len(jsons), + "captions": Media.captions(seq), + } diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/object_3d.py b/lib/python3.12/site-packages/wandb/sdk/data_types/object_3d.py new file mode 100644 index 0000000000000000000000000000000000000000..e46a149c53db034642b2670315c699febc181de1 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/object_3d.py @@ -0,0 +1,476 @@ +import codecs +import itertools +import json +import os +import pathlib +from typing import ( + TYPE_CHECKING, + ClassVar, + Literal, + Optional, + Sequence, + Set, + TextIO, + Tuple, + Type, + TypedDict, + Union, +) + +import wandb +from wandb import util +from wandb.sdk.lib import runid +from wandb.sdk.lib.paths import LogicalPath + +from . import _dtypes +from ._private import MEDIA_TMP +from .base_types.media import BatchableMedia + +if TYPE_CHECKING: # pragma: no cover + import numpy as np + import numpy.typing as npt + + from wandb.sdk.artifacts.artifact import Artifact + + from ..wandb_run import Run as LocalRun + + numeric = Union[int, float, np.integer, np.float64] + FileFormat3D = Literal[ + "obj", + "gltf", + "glb", + "babylon", + "stl", + "pts.json", + ] + Point3D = Tuple[numeric, numeric, numeric] + Point3DWithCategory = Tuple[numeric, numeric, numeric, numeric] + Point3DWithColors = Tuple[numeric, numeric, numeric, numeric, numeric, numeric] + Point = Union[Point3D, Point3DWithCategory, Point3DWithColors] + PointCloudType = Literal["lidar/beta"] + RGBColor = Tuple[numeric, numeric, numeric] + + Quaternion = Tuple[numeric, numeric, numeric, numeric] + + class Box3D(TypedDict): + corners: Tuple[ + Point3D, + Point3D, + Point3D, + Point3D, + Point3D, + Point3D, + Point3D, + Point3D, + ] + label: Optional[str] + color: RGBColor + score: Optional[numeric] + + class Vector3D(TypedDict): + start: Sequence[Point3D] + end: Sequence[Point3D] + + class Camera(TypedDict): + viewpoint: Sequence[Point3D] + target: Sequence[Point3D] + + +def _install_numpy_error() -> "wandb.Error": + return wandb.Error( + "wandb.Object3D requires NumPy. To get it, run 'pip install numpy'." + ) + + +def _normalize_vec(x: "np.ndarray") -> "np.ndarray": + """Normalizes a non-zero 1-dimensional array.""" + try: + import numpy as np + except ImportError as e: + raise _install_numpy_error() from e + + x = abs(x) + argmax = x.argmax() + + if x[argmax] == 0: + raise ValueError("Unexpected zero quaternion.") + + # Rescale by the largest component first to be more numerically + # stable than dividing by the norm directly. + x = x / x[argmax] + return x / np.linalg.norm(x) + + +def _quaternion_to_rotation(quaternion: "np.ndarray") -> "np.ndarray": + """Returns the rotation matrix corresponding to a non-zero quaternion. + + The corresponding rotation matrix transforms column vectors by + post-multiplication: `x @ R`. This way, it can be used to + transform an Nx3 NumPy array of points as `points @ R`. + """ + try: + import numpy as np + except ImportError as e: + raise _install_numpy_error() from e + + # Precompute a few products to simplify the expression below. + qr, qi, qj, qk = _normalize_vec(quaternion) + qii, qjj, qkk = qi**2, qj**2, qk**2 + qij, qik, qjk = qi * qj, qi * qk, qj * qk + qir, qjr, qkr = qi * qr, qj * qr, qk * qr + + return np.array( + ( + (1 - 2 * (qjj + qkk), 2 * (qij + qkr), 2 * (qik - qjr)), + (2 * (qij - qkr), 1 - 2 * (qii + qkk), 2 * (qjk + qir)), + (2 * (qik + qjr), 2 * (qjk - qir), 1 - 2 * (qii + qjj)), + ), + dtype=np.float64, + ) + + +def box3d( + *, + center: "npt.ArrayLike", + size: "npt.ArrayLike", + orientation: "npt.ArrayLike", + color: "RGBColor", + label: "Optional[str]" = None, + score: "Optional[numeric]" = None, +) -> "Box3D": + """Returns a Box3D. + + Args: + center: The center point of the box as a length-3 ndarray. + size: The box's X, Y and Z dimensions as a length-3 ndarray. + orientation: The rotation transforming global XYZ coordinates + into the box's local XYZ coordinates, given as a length-4 + ndarray [r, x, y, z] corresponding to the non-zero quaternion + r + xi + yj + zk. + color: The box's color as an (r, g, b) tuple with 0 <= r,g,b <= 1. + label: An optional label for the box. + score: An optional score for the box. + """ + try: + import numpy as np + except ImportError as e: + raise _install_numpy_error() from e + + center = np.asarray(center, dtype=np.float64) + size = np.asarray(size, dtype=np.float64) + orientation = np.asarray(orientation, dtype=np.float64) + + assert center.shape == (3,) + assert size.shape == (3,) + assert orientation.shape == (4,) + + # Precompute the rotation matrix. + rot = _quaternion_to_rotation(orientation) + + # Scale, rotate and translate each corner of the unit box. + unit_corners = np.array( + list(itertools.product((-1, 1), (-1, 1), (-1, 1))), + dtype=np.float64, + ) + corners = center + (0.5 * size * unit_corners) @ rot + + return { + # Ignore the type because mypy can't infer that the list has length 8: + # https://github.com/python/mypy/issues/7509 + "corners": tuple(tuple(pt) for pt in corners), # type: ignore + "color": color, + "label": label, + "score": score, + } + + +class Object3D(BatchableMedia): + """Wandb class for 3D point clouds. + + Args: + data_or_path: (numpy array, pathlib.Path, string, io) + Object3D can be initialized from a file or a numpy array. + + You can pass a path to a file or an io object and a file_type + which must be one of SUPPORTED_TYPES + + The shape of the numpy array must be one of either: + ``` + [[x y z], ...] nx3 + [[x y z c], ...] nx4 where c is a category with supported range [1, 14] + [[x y z r g b], ...] nx6 where is rgb is color + ``` + """ + + SUPPORTED_TYPES: ClassVar[Set[str]] = { + "obj", + "gltf", + "glb", + "babylon", + "stl", + "pts.json", + } + SUPPORTED_POINT_CLOUD_TYPES: ClassVar[Set[str]] = {"lidar/beta"} + _log_type: ClassVar[str] = "object3D-file" + + def __init__( + self, + data_or_path: Union["np.ndarray", str, pathlib.Path, "TextIO", dict], + caption: Optional[str] = None, + **kwargs: Optional[Union[str, "FileFormat3D"]], + ) -> None: + super().__init__(caption=caption) + + if hasattr(data_or_path, "name") and not isinstance(data_or_path, pathlib.Path): + # if the file has a path, we just detect the type and copy it from there. + # this does not work for pathlib.Path objects, + # where `.name` returns the last directory in the path. + data_or_path = data_or_path.name + + if hasattr(data_or_path, "read"): + if hasattr(data_or_path, "seek"): + data_or_path.seek(0) + object_3d = data_or_path.read() + + extension = kwargs.pop("file_type", None) + if extension is None: + raise ValueError( + "Must pass file type keyword argument when using io objects." + ) + if extension not in Object3D.SUPPORTED_TYPES: + raise ValueError( + "Object 3D only supports numpy arrays or files of the type: " + + ", ".join(Object3D.SUPPORTED_TYPES) + ) + + extension = "." + extension + + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + extension) + with open(tmp_path, "w") as f: + f.write(object_3d) + + self._set_file(tmp_path, is_tmp=True, extension=extension) + elif isinstance(data_or_path, (str, pathlib.Path)): + data_or_path = str(data_or_path) + + path = data_or_path + extension = None + for supported_type in Object3D.SUPPORTED_TYPES: + if path.endswith(supported_type): + extension = "." + supported_type + break + + if not extension: + raise ValueError( + "File '" + + path + + "' is not compatible with Object3D: supported types are: " + + ", ".join(Object3D.SUPPORTED_TYPES) + ) + + self._set_file(data_or_path, is_tmp=False, extension=extension) + # Supported different types and scene for 3D scenes + elif isinstance(data_or_path, dict) and "type" in data_or_path: + if data_or_path["type"] == "lidar/beta": + data = { + "type": data_or_path["type"], + "vectors": data_or_path["vectors"].tolist() + if "vectors" in data_or_path + else [], + "points": data_or_path["points"].tolist() + if "points" in data_or_path + else [], + "boxes": data_or_path["boxes"].tolist() + if "boxes" in data_or_path + else [], + } + else: + raise ValueError( + "Type not supported, only 'lidar/beta' is currently supported" + ) + + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".pts.json") + with codecs.open(tmp_path, "w", encoding="utf-8") as fp: + json.dump( + data, + fp, + separators=(",", ":"), + sort_keys=True, + indent=4, + ) + self._set_file(tmp_path, is_tmp=True, extension=".pts.json") + elif util.is_numpy_array(data_or_path): + np_data = data_or_path + + # The following assertion is required for numpy to trust that + # np_data is numpy array. The reason it is behind a False + # guard is to ensure that this line does not run at runtime, + # which would cause a runtime error if the user's machine did + # not have numpy installed. + + if TYPE_CHECKING: + assert isinstance(np_data, np.ndarray) + + if len(np_data.shape) != 2 or np_data.shape[1] not in {3, 4, 6}: + raise ValueError( + """ + The shape of the numpy array must be one of either + [[x y z], ...] nx3 + [x y z c], ...] nx4 where c is a category with supported range [1, 14] + [x y z r g b], ...] nx6 where rgb is color + """ + ) + + list_data = np_data.tolist() + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".pts.json") + with codecs.open(tmp_path, "w", encoding="utf-8") as fp: + json.dump( + list_data, + fp, + separators=(",", ":"), + sort_keys=True, + indent=4, + ) + self._set_file(tmp_path, is_tmp=True, extension=".pts.json") + else: + raise ValueError("data must be a numpy array, dict or a file object") + + @classmethod + def from_file( + cls, + data_or_path: Union["TextIO", str], + file_type: Optional["FileFormat3D"] = None, + ) -> "Object3D": + """Initializes Object3D from a file or stream. + + Args: + data_or_path (Union["TextIO", str]): A path to a file or a `TextIO` stream. + file_type (str): Specifies the data format passed to `data_or_path`. Required when `data_or_path` is a + `TextIO` stream. This parameter is ignored if a file path is provided. The type is taken from the file extension. + """ + # if file_type is not None and file_type not in cls.SUPPORTED_TYPES: + # raise ValueError( + # f"Unsupported file type: {file_type}. Supported types are: {cls.SUPPORTED_TYPES}" + # ) + return cls(data_or_path, file_type=file_type) + + @classmethod + def from_numpy(cls, data: "np.ndarray") -> "Object3D": + """Initializes Object3D from a numpy array. + + Args: + data (numpy array): Each entry in the array will + represent one point in the point cloud. + + + The shape of the numpy array must be one of either: + ``` + [[x y z], ...] # nx3. + [[x y z c], ...] # nx4 where c is a category with supported range [1, 14]. + [[x y z r g b], ...] # nx6 where is rgb is color. + ``` + """ + if not util.is_numpy_array(data): + raise ValueError("`data` must be a numpy array") + + if len(data.shape) != 2 or data.shape[1] not in {3, 4, 6}: + raise ValueError( + """ + The shape of the numpy array must be one of either: + [[x y z], ...] nx3 + [x y z c], ...] nx4 where c is a category with supported range [1, 14] + [x y z r g b], ...] nx6 where rgb is color + """ + ) + + return cls(data) + + @classmethod + def from_point_cloud( + cls, + points: Sequence["Point"], + boxes: Sequence["Box3D"], + vectors: Optional[Sequence["Vector3D"]] = None, + point_cloud_type: "PointCloudType" = "lidar/beta", + # camera: Optional[Camera] = None, + ) -> "Object3D": + """Initializes Object3D from a python object. + + Args: + points (Sequence["Point"]): The points in the point cloud. + boxes (Sequence["Box3D"]): 3D bounding boxes for labeling the point cloud. Boxes + are displayed in point cloud visualizations. + vectors (Optional[Sequence["Vector3D"]]): Each vector is displayed in the point cloud + visualization. Can be used to indicate directionality of bounding boxes. Defaults to None. + point_cloud_type ("lidar/beta"): At this time, only the "lidar/beta" type is supported. Defaults to "lidar/beta". + """ + if point_cloud_type not in cls.SUPPORTED_POINT_CLOUD_TYPES: + raise ValueError("Point cloud type not supported") + + numpy = wandb.util.get_module( + "numpy", + required="wandb.Object3D.from_point_cloud requires numpy. Install with `pip install numpy`", + ) + + data = { + "type": point_cloud_type, + "points": numpy.array(points), + "boxes": numpy.array(boxes), + "vectors": numpy.array(vectors) if vectors is not None else numpy.array([]), + } + + return cls(data) + + @classmethod + def get_media_subdir(cls: Type["Object3D"]) -> str: + return os.path.join("media", "object3D") + + def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict: + json_dict = super().to_json(run_or_artifact) + json_dict["_type"] = Object3D._log_type + + if isinstance(run_or_artifact, wandb.Artifact): + if self._path is None or not self._path.endswith(".pts.json"): + raise ValueError( + "Non-point cloud 3D objects are not yet supported with Artifacts" + ) + + return json_dict + + @classmethod + def seq_to_json( + cls: Type["Object3D"], + seq: Sequence["BatchableMedia"], + run: "LocalRun", + key: str, + step: Union[int, str], + ) -> dict: + seq = list(seq) + + jsons = [obj.to_json(run) for obj in seq] + + for obj in jsons: + expected = LogicalPath(cls.get_media_subdir()) + if not obj["path"].startswith(expected): + raise ValueError( + "Files in an array of Object3D's must be in the {} directory, not {}".format( + expected, obj["path"] + ) + ) + + return { + "_type": "object3D", + "filenames": [ + os.path.relpath(j["path"], cls.get_media_subdir()) for j in jsons + ], + "count": len(jsons), + "objects": jsons, + } + + +class _Object3DFileType(_dtypes.Type): + name = "object3D-file" + types = [Object3D] + + +_dtypes.TypeRegistry.add(_Object3DFileType) diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/plotly.py b/lib/python3.12/site-packages/wandb/sdk/data_types/plotly.py new file mode 100644 index 0000000000000000000000000000000000000000..62254699d12b8c81fce61acc5cbdd823700e228f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/plotly.py @@ -0,0 +1,82 @@ +import codecs +import os +from typing import TYPE_CHECKING, Sequence, Type, Union + +from wandb import util +from wandb.sdk.lib import runid + +from ._private import MEDIA_TMP +from .base_types.media import Media, _numpy_arrays_to_lists +from .base_types.wb_value import WBValue +from .image import Image + +if TYPE_CHECKING: # pragma: no cover + import matplotlib # type: ignore + import pandas as pd + import plotly # type: ignore + + from wandb.sdk.artifacts.artifact import Artifact + + from ..wandb_run import Run as LocalRun + + ValToJsonType = Union[ + dict, + "WBValue", + Sequence["WBValue"], + "plotly.Figure", + "matplotlib.artist.Artist", + "pd.DataFrame", + object, + ] + + +class Plotly(Media): + """Wandb class for plotly plots. + + Args: + val: matplotlib or plotly figure + """ + + _log_type = "plotly-file" + + @classmethod + def make_plot_media( + cls: Type["Plotly"], val: Union["plotly.Figure", "matplotlib.artist.Artist"] + ) -> Union[Image, "Plotly"]: + if util.is_matplotlib_typename(util.get_full_typename(val)): + if util.matplotlib_contains_images(val): + return Image(val) + val = util.matplotlib_to_plotly(val) + return cls(val) + + def __init__(self, val: Union["plotly.Figure", "matplotlib.artist.Artist"]): + super().__init__() + # First, check to see if the incoming `val` object is a plotfly figure + if not util.is_plotly_figure_typename(util.get_full_typename(val)): + # If it is not, but it is a matplotlib figure, then attempt to convert it to plotly + if util.is_matplotlib_typename(util.get_full_typename(val)): + if util.matplotlib_contains_images(val): + raise ValueError( + "Plotly does not currently support converting matplotlib figures containing images. \ + You can convert the plot to a static image with `wandb.Image(plt)` " + ) + val = util.matplotlib_to_plotly(val) + else: + raise ValueError( + "Logged plots must be plotly figures, or matplotlib plots convertible to plotly via mpl_to_plotly" + ) + + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".plotly.json") + val = _numpy_arrays_to_lists(val.to_plotly_json()) + with codecs.open(tmp_path, "w", encoding="utf-8") as fp: + util.json_dump_safer(val, fp) + self._set_file(tmp_path, is_tmp=True, extension=".plotly.json") + + @classmethod + def get_media_subdir(cls: Type["Plotly"]) -> str: + return os.path.join("media", "plotly") + + def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict: + json_dict = super().to_json(run_or_artifact) + json_dict["_type"] = self._log_type + return json_dict diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/saved_model.py b/lib/python3.12/site-packages/wandb/sdk/data_types/saved_model.py new file mode 100644 index 0000000000000000000000000000000000000000..88e897f47b9020dc564e62a715d3bb0cd35f3974 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/saved_model.py @@ -0,0 +1,437 @@ +from __future__ import annotations + +import os +import pathlib +import shutil +import sys +from types import ModuleType +from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast + +import wandb +from wandb import util +from wandb.sdk.lib import runid +from wandb.sdk.lib.hashutil import md5_file_hex +from wandb.sdk.lib.paths import LogicalPath + +from ._private import MEDIA_TMP +from .base_types.wb_value import WBValue + +if TYPE_CHECKING: + import sklearn # type: ignore + import tensorflow # type: ignore + import torch # type: ignore + from typing_extensions import Self + + from wandb.sdk.artifacts.artifact import Artifact + from wandb.sdk.wandb_run import Run as LocalRun + + +DEBUG_MODE = False + + +def _add_deterministic_dir_to_artifact( + artifact: Artifact, dir_name: str, target_dir_root: str +) -> str: + file_paths = [] + for dirpath, _, filenames in os.walk(dir_name, topdown=True): + for fn in filenames: + file_paths.append(os.path.join(dirpath, fn)) + dirname = md5_file_hex(*file_paths)[:20] + target_path = LogicalPath(os.path.join(target_dir_root, dirname)) + artifact.add_dir(dir_name, target_path) + return target_path + + +def _load_dir_from_artifact(source_artifact: Artifact, path: str) -> str: + dl_path = None + + # Look through the entire manifest to find all of the files in the directory. + # Construct the directory path by inspecting the target download location. + for p, _ in source_artifact.manifest.entries.items(): + if p.startswith(path): + example_path = source_artifact.get_entry(p).download() + if dl_path is None: + root = example_path[: -len(p)] + dl_path = os.path.join(root, path) + + assert dl_path is not None, f"Could not find directory {path} in artifact" + + return dl_path + + +SavedModelObjType = TypeVar("SavedModelObjType") + + +class _SavedModel(WBValue, Generic[SavedModelObjType]): + """Internal W&B Artifact model storage. + + _model_type_id: (str) The id of the SavedModel subclass used to serialize the model. + """ + + _log_type: ClassVar[str] + _path_extension: ClassVar[str] + + _model_obj: SavedModelObjType | None + _path: str | None + _input_obj_or_path: SavedModelObjType | str | pathlib.Path + + # Public Methods + def __init__( + self, obj_or_path: SavedModelObjType | str | pathlib.Path, **kwargs: Any + ) -> None: + super().__init__() + if self.__class__ == _SavedModel: + raise TypeError( + "Cannot instantiate abstract SavedModel class - please use SavedModel.init(...) instead." + ) + self._model_obj = None + self._path = None + self._input_obj_or_path = obj_or_path + input_is_path = isinstance(obj_or_path, (str, pathlib.Path)) and os.path.exists( + obj_or_path + ) + if input_is_path: + obj_or_path = str(obj_or_path) + self._set_obj(self._deserialize(obj_or_path)) + else: + self._set_obj(obj_or_path) + + self._copy_to_disk() + # At this point, the model will be saved to a temp path, + # and self._path will be set to such temp path. If the model + # provided was a path, then both self._path and self._model_obj + # are copies of the user-provided data. However, if the input + # was a model object, then we want to clear the model object. The first + # accessing of the model object (via .model_obj()) will load the model + # from the temp path. + + if not input_is_path: + self._unset_obj() + + @staticmethod + def init(obj_or_path: Any, **kwargs: Any) -> _SavedModel: + maybe_instance = _SavedModel._maybe_init(obj_or_path, **kwargs) + if maybe_instance is None: + raise ValueError( + f"No suitable SavedModel subclass constructor found for obj_or_path: {obj_or_path}" + ) + return maybe_instance + + @classmethod + def from_json( + cls: type[_SavedModel], json_obj: dict, source_artifact: Artifact + ) -> _SavedModel: + path = json_obj["path"] + + # First, if the entry is a file, the download it. + entry = source_artifact.manifest.entries.get(path) + if entry is not None: + dl_path = str(source_artifact.get_entry(path).download()) + else: + # If not, assume it is directory. + # FUTURE: Add this functionality to the artifact loader + # (would be nice to parallelize) + dl_path = _load_dir_from_artifact(source_artifact, path) + + # Return the SavedModel object instantiated with the downloaded path + # and specified adapter. + return cls(dl_path) + + def to_json(self, run_or_artifact: LocalRun | Artifact) -> dict: + # Unlike other data types, we do not allow adding to a Run directly. There is a + # bit of tech debt in the other data types which requires the input to `to_json` + # to accept a Run or Artifact. However, Run additions should be deprecated in the future. + # This check helps ensure we do not add to the debt. + from wandb.sdk.wandb_run import Run + + if isinstance(run_or_artifact, Run): + raise TypeError("SavedModel cannot be added to run - must use artifact") + artifact = run_or_artifact + json_obj = { + "type": self._log_type, + } + assert self._path is not None, "Cannot add SavedModel to Artifact without path" + if os.path.isfile(self._path): + # If the path is a file, then we can just add it to the artifact, + # First checking to see if the artifact already has the file (use the cache) + # Else, add it directly, allowing the artifact adder to rename the file deterministically. + already_added_path = artifact.get_added_local_path_name(self._path) + if already_added_path is not None: + json_obj["path"] = already_added_path + else: + target_path = os.path.join( + ".wb_data", "saved_models", os.path.basename(self._path) + ) + json_obj["path"] = artifact.add_file(self._path, target_path, True).path + elif os.path.isdir(self._path): + # If the path is a directory, then we need to add all of the files + # The directory must be named deterministically based on the contents of the directory, + # but the files themselves need to have their name preserved. + # FUTURE: Add this functionality to the artifact adder itself + json_obj["path"] = _add_deterministic_dir_to_artifact( + artifact, self._path, os.path.join(".wb_data", "saved_models") + ) + else: + raise ValueError( + f"Expected a path to a file or directory, got {self._path}" + ) + + return json_obj + + def model_obj(self) -> SavedModelObjType: + """Return the model object.""" + if self._model_obj is None: + assert self._path is not None, "Cannot load model object without path" + self._set_obj(self._deserialize(self._path)) + + assert self._model_obj is not None, "Model object is None" + return self._model_obj + + # Methods to be implemented by subclasses + @staticmethod + def _deserialize(path: str) -> SavedModelObjType: + """Return the model object from a path. Allowed to throw errors.""" + raise NotImplementedError + + @staticmethod + def _validate_obj(obj: Any) -> bool: + """Validate the model object. Allowed to throw errors.""" + raise NotImplementedError + + @staticmethod + def _serialize(obj: SavedModelObjType, dir_or_file_path: str) -> None: + """Save the model to disk. + + The method will receive a directory path which all files needed for + deserialization should be saved. A directory will always be passed if + _path_extension is an empty string, else a single file will be passed. Allowed + to throw errors. + """ + raise NotImplementedError + + # Private Class Methods + @classmethod + def _maybe_init( + cls: type[_SavedModel], obj_or_path: Any, **kwargs: Any + ) -> _SavedModel | None: + # _maybe_init is an exception-safe method that will return an instance of this class + # (or any subclass of this class - recursively) OR None if no subclass constructor is found. + # We first try the current class, then recursively call this method on children classes. This pattern + # conforms to the new "Weave-type" pattern developed by Shawn. This way, we can for example have a + # pytorch subclass that can itself have two subclasses: one for a TorchScript model, and one for a PyTorch model. + # The children subclasses will know how to serialize/deserialize their respective payloads, but the pytorch + # parent class can know how to execute inference on the model - regardless of serialization strategy. + try: + return cls(obj_or_path, **kwargs) + except Exception as e: + if DEBUG_MODE: + print(f"{cls}._maybe_init({obj_or_path}) failed: {e}") # noqa: T201 + + for child_cls in cls.__subclasses__(): + maybe_instance = child_cls._maybe_init(obj_or_path, **kwargs) + if maybe_instance is not None: + return maybe_instance + + return None + + @classmethod + def _tmp_path(cls: type[_SavedModel]) -> str: + # Generates a tmp path under our MEDIA_TMP directory which confirms to the file + # or folder preferences of the class. + assert isinstance(cls._path_extension, str), "_path_extension must be a string" + tmp_path = os.path.abspath(os.path.join(MEDIA_TMP.name, runid.generate_id())) + if cls._path_extension != "": + tmp_path += "." + cls._path_extension + return tmp_path + + # Private Instance Methods + def _copy_to_disk(self) -> None: + # Creates a temporary path and writes a fresh copy of the + # model to disk - updating the _path appropriately. + tmp_path = self._tmp_path() + self._dump(tmp_path) + self._path = tmp_path + + def _unset_obj(self) -> None: + assert self._path is not None, "Cannot unset object if path is None" + self._model_obj = None + + def _set_obj(self, model_obj: Any) -> None: + assert model_obj is not None and self._validate_obj(model_obj), ( + f"Invalid model object {model_obj}" + ) + self._model_obj = model_obj + + def _dump(self, target_path: str) -> None: + assert self._model_obj is not None, "Cannot dump if model object is None" + self._serialize(self._model_obj, target_path) + + +def _get_cloudpickle() -> ModuleType: + return cast( + ModuleType, + util.get_module("cloudpickle", "ModelAdapter requires `cloudpickle`"), + ) + + +# TODO: Add pip deps +# TODO: potentially move this up to the saved model class +PicklingSavedModelObjType = TypeVar("PicklingSavedModelObjType") + + +class _PicklingSavedModel(_SavedModel[SavedModelObjType]): + _dep_py_files: list[str] | None = None + _dep_py_files_path: str | None = None + + def __init__( + self, + obj_or_path: SavedModelObjType | str | pathlib.Path, + dep_py_files: list[str] | None = None, + ): + super().__init__(obj_or_path) + if self.__class__ == _PicklingSavedModel: + raise TypeError( + "Cannot instantiate abstract _PicklingSavedModel class - please use SavedModel.init(...) instead." + ) + if dep_py_files is not None and len(dep_py_files) > 0: + self._dep_py_files = dep_py_files + self._dep_py_files_path = os.path.abspath( + os.path.join(MEDIA_TMP.name, runid.generate_id()) + ) + os.makedirs(self._dep_py_files_path, exist_ok=True) + for extra_file in self._dep_py_files: + if os.path.isfile(extra_file): + shutil.copy(extra_file, self._dep_py_files_path) + elif os.path.isdir(extra_file): + shutil.copytree( + extra_file, + os.path.join( + self._dep_py_files_path, os.path.basename(extra_file) + ), + ) + else: + raise ValueError(f"Invalid dependency file: {extra_file}") + + @classmethod + def from_json(cls, json_obj: dict, source_artifact: Artifact) -> Self: + backup_path = [p for p in sys.path] + if ( + "dep_py_files_path" in json_obj + and json_obj["dep_py_files_path"] is not None + ): + dl_path = _load_dir_from_artifact( + source_artifact, json_obj["dep_py_files_path"] + ) + assert dl_path is not None + sys.path.append(dl_path) + inst = super().from_json(json_obj, source_artifact) # type: ignore + sys.path = backup_path + + return inst # type: ignore + + def to_json(self, run_or_artifact: LocalRun | Artifact) -> dict: + json_obj = super().to_json(run_or_artifact) + assert isinstance(run_or_artifact, wandb.Artifact) + if self._dep_py_files_path is not None: + json_obj["dep_py_files_path"] = _add_deterministic_dir_to_artifact( + run_or_artifact, + self._dep_py_files_path, + os.path.join(".wb_data", "extra_files"), + ) + return json_obj + + +def _get_torch() -> ModuleType: + return cast( + ModuleType, + util.get_module("torch", "ModelAdapter requires `torch`"), + ) + + +class _PytorchSavedModel(_PicklingSavedModel["torch.nn.Module"]): + _log_type = "pytorch-model-file" + _path_extension = "pt" + + @staticmethod + def _deserialize(dir_or_file_path: str) -> torch.nn.Module: + return _get_torch().load(dir_or_file_path, weights_only=False) + + @staticmethod + def _validate_obj(obj: Any) -> bool: + return isinstance(obj, _get_torch().nn.Module) + + @staticmethod + def _serialize(model_obj: torch.nn.Module, dir_or_file_path: str) -> None: + _get_torch().save( + model_obj, + dir_or_file_path, + pickle_module=_get_cloudpickle(), + ) + + +def _get_sklearn() -> ModuleType: + return cast( + ModuleType, + util.get_module("sklearn", "ModelAdapter requires `sklearn`"), + ) + + +class _SklearnSavedModel(_PicklingSavedModel["sklearn.base.BaseEstimator"]): + _log_type = "sklearn-model-file" + _path_extension = "pkl" + + @staticmethod + def _deserialize( + dir_or_file_path: str, + ) -> sklearn.base.BaseEstimator: + with open(dir_or_file_path, "rb") as file: + model = _get_cloudpickle().load(file) + return model + + @staticmethod + def _validate_obj(obj: Any) -> bool: + dynamic_sklearn = _get_sklearn() + return cast( + bool, + ( + dynamic_sklearn.base.is_classifier(obj) + or dynamic_sklearn.base.is_outlier_detector(obj) + or dynamic_sklearn.base.is_regressor(obj) + ), + ) + + @staticmethod + def _serialize( + model_obj: sklearn.base.BaseEstimator, dir_or_file_path: str + ) -> None: + dynamic_cloudpickle = _get_cloudpickle() + with open(dir_or_file_path, "wb") as file: + dynamic_cloudpickle.dump(model_obj, file) + + +def _get_tf_keras() -> ModuleType: + return cast( + ModuleType, + util.get_module("tensorflow", "ModelAdapter requires `tensorflow`"), + ).keras + + +class _TensorflowKerasSavedModel(_SavedModel["tensorflow.keras.Model"]): + _log_type = "tfkeras-model-file" + _path_extension = "" + + @staticmethod + def _deserialize( + dir_or_file_path: str, + ) -> tensorflow.keras.Model: + return _get_tf_keras().models.load_model(dir_or_file_path) + + @staticmethod + def _validate_obj(obj: Any) -> bool: + return isinstance(obj, _get_tf_keras().models.Model) + + @staticmethod + def _serialize(model_obj: tensorflow.keras.Model, dir_or_file_path: str) -> None: + _get_tf_keras().models.save_model( + model_obj, dir_or_file_path, include_optimizer=True + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/table.py b/lib/python3.12/site-packages/wandb/sdk/data_types/table.py new file mode 100644 index 0000000000000000000000000000000000000000..448ca0a9e7f891a341bd5d2cb2a0e1770c6231ce --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/table.py @@ -0,0 +1,1439 @@ +import base64 +import binascii +import codecs +import datetime +import json +import logging +import os +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set, Tuple + +import wandb +from wandb import util +from wandb.sdk.lib import runid + +from . import _dtypes +from ._private import MEDIA_TMP +from .base_types.media import Media, _numpy_arrays_to_lists +from .base_types.wb_value import WBValue +from .table_decorators import ( + allow_incremental_logging_after_append, + allow_relogging_after_mutation, + ensure_not_incremental, +) +from .utils import _json_helper + +if TYPE_CHECKING: + from wandb.sdk.artifacts import artifact + + from ...wandb_run import Run as LocalRun + + +class _TableLinkMixin: + def set_table(self, table): + self._table = table + + +class _TableKey(str, _TableLinkMixin): + def set_table(self, table, col_name): + assert col_name in table.columns + self._table = table + self._col_name = col_name + + +class _TableIndex(int, _TableLinkMixin): + def get_row(self): + row = {} + if self._table: + row = { + c: self._table.data[self][i] for i, c in enumerate(self._table.columns) + } + + return row + + +class _PrimaryKeyType(_dtypes.Type): + name = "primaryKey" + legacy_names = ["wandb.TablePrimaryKey"] + + def assign_type(self, wb_type=None): + if isinstance(wb_type, _dtypes.StringType) or isinstance( + wb_type, _PrimaryKeyType + ): + return self + return _dtypes.InvalidType() + + @classmethod + def from_obj(cls, py_obj): + if not isinstance(py_obj, _TableKey): + raise TypeError("py_obj must be a wandb.Table") + else: + return cls() + + +class _ForeignKeyType(_dtypes.Type): + name = "foreignKey" + legacy_names = ["wandb.TableForeignKey"] + types = [_TableKey] + + def __init__(self, table, col_name): + assert isinstance(table, Table) + assert isinstance(col_name, str) + assert col_name in table.columns + self.params.update({"table": table, "col_name": col_name}) + + def assign_type(self, wb_type=None): + if isinstance(wb_type, _dtypes.StringType): + return self + elif ( + isinstance(wb_type, _ForeignKeyType) + and id(self.params["table"]) == id(wb_type.params["table"]) + and self.params["col_name"] == wb_type.params["col_name"] + ): + return self + + return _dtypes.InvalidType() + + @classmethod + def from_obj(cls, py_obj): + if not isinstance(py_obj, _TableKey): + raise TypeError("py_obj must be a _TableKey") + else: + return cls(py_obj._table, py_obj._col_name) + + def to_json(self, artifact=None): + res = super().to_json(artifact) + if artifact is not None: + table_name = f"media/tables/t_{runid.generate_id()}" + entry = artifact.add(self.params["table"], table_name) + res["params"]["table"] = entry.path + else: + raise AssertionError( + "_ForeignKeyType does not support serialization without an artifact" + ) + return res + + @classmethod + def from_json( + cls, + json_dict, + artifact, + ): + table = None + col_name = None + if artifact is None: + raise AssertionError( + "_ForeignKeyType does not support deserialization without an artifact" + ) + else: + table = artifact.get(json_dict["params"]["table"]) + col_name = json_dict["params"]["col_name"] + + if table is None: + raise AssertionError("Unable to deserialize referenced table") + + return cls(table, col_name) + + +class _ForeignIndexType(_dtypes.Type): + name = "foreignIndex" + legacy_names = ["wandb.TableForeignIndex"] + types = [_TableIndex] + + def __init__(self, table): + assert isinstance(table, Table) + self.params.update({"table": table}) + + def assign_type(self, wb_type=None): + if isinstance(wb_type, _dtypes.NumberType): + return self + elif isinstance(wb_type, _ForeignIndexType) and id(self.params["table"]) == id( + wb_type.params["table"] + ): + return self + + return _dtypes.InvalidType() + + @classmethod + def from_obj(cls, py_obj): + if not isinstance(py_obj, _TableIndex): + raise TypeError("py_obj must be a _TableIndex") + else: + return cls(py_obj._table) + + def to_json(self, artifact=None): + res = super().to_json(artifact) + if artifact is not None: + table_name = f"media/tables/t_{runid.generate_id()}" + entry = artifact.add(self.params["table"], table_name) + res["params"]["table"] = entry.path + else: + raise AssertionError( + "_ForeignIndexType does not support serialization without an artifact" + ) + return res + + @classmethod + def from_json( + cls, + json_dict, + artifact, + ): + table = None + if artifact is None: + raise AssertionError( + "_ForeignIndexType does not support deserialization without an artifact" + ) + else: + table = artifact.get(json_dict["params"]["table"]) + + if table is None: + raise AssertionError("Unable to deserialize referenced table") + + return cls(table) + + +_SUPPORTED_LOGGING_MODES = ["IMMUTABLE", "MUTABLE", "INCREMENTAL"] + + +class Table(Media): + """The Table class used to display and analyze tabular data. + + Unlike traditional spreadsheets, Tables support numerous types of data: + scalar values, strings, numpy arrays, and most subclasses of `wandb.data_types.Media`. + This means you can embed `Images`, `Video`, `Audio`, and other sorts of rich, annotated media + directly in Tables, alongside other traditional scalar values. + + This class is the primary class used to generate the W&B Tables + https://docs.wandb.ai/guides/models/tables/. + """ + + MAX_ROWS = 10000 + MAX_ARTIFACT_ROWS = 200000 + _MAX_EMBEDDING_DIMENSIONS = 150 + _log_type = "table" + + def __init__( + self, + columns=None, + data=None, + rows=None, + dataframe=None, + dtype=None, + optional=True, + allow_mixed_types=False, + log_mode: Optional[ + Literal["IMMUTABLE", "MUTABLE", "INCREMENTAL"] + ] = "IMMUTABLE", + ): + """Initializes a Table object. + + The rows is available for legacy reasons and should not be used. + The Table class uses data to mimic the Pandas API. + + Args: + columns: (List[str]) Names of the columns in the table. + Defaults to ["Input", "Output", "Expected"]. + data: (List[List[any]]) 2D row-oriented array of values. + dataframe: (pandas.DataFrame) DataFrame object used to create the table. + When set, `data` and `columns` arguments are ignored. + optional: (Union[bool,List[bool]]) Determines if `None` values are allowed. Default to True + - If a singular bool value, then the optionality is enforced for all + columns specified at construction time + - If a list of bool values, then the optionality is applied to each + column - should be the same length as `columns` + applies to all columns. A list of bool values applies to each respective column. + allow_mixed_types: (bool) Determines if columns are allowed to have mixed types + (disables type validation). Defaults to False + log_mode: Optional[str] Controls how the Table is logged when mutations occur. + Options: + - "IMMUTABLE" (default): Table can only be logged once; subsequent + logging attempts after the table has been mutated will be no-ops. + - "MUTABLE": Table can be re-logged after mutations, creating + a new artifact version each time it's logged. + - "INCREMENTAL": Table data is logged incrementally, with each log creating + a new artifact entry containing the new data since the last log. + """ + super().__init__() + self._validate_log_mode(log_mode) + self.log_mode = log_mode + if self.log_mode == "INCREMENTAL": + self._increment_num: int | None = None + self._last_logged_idx: int | None = None + self._previous_increments_paths: list[str] | None = None + self._run_target_for_increments: LocalRun | None = None + self._pk_col = None + self._fk_cols: set[str] = set() + if allow_mixed_types: + dtype = _dtypes.AnyType + + # This is kept for legacy reasons (tss: personally, I think we should remove this) + if columns is None: + columns = ["Input", "Output", "Expected"] + + # Explicit dataframe option + if dataframe is not None: + self._init_from_dataframe(dataframe, columns, optional, dtype) + else: + # Expected pattern + if data is not None: + if util.is_numpy_array(data): + self._init_from_ndarray(data, columns, optional, dtype) + elif util.is_pandas_data_frame(data): + self._init_from_dataframe(data, columns, optional, dtype) + else: + self._init_from_list(data, columns, optional, dtype) + + # legacy + elif rows is not None: + self._init_from_list(rows, columns, optional, dtype) + + # Default empty case + else: + self._init_from_list([], columns, optional, dtype) + + def _validate_log_mode(self, log_mode): + assert log_mode in _SUPPORTED_LOGGING_MODES, ( + f"Invalid log_mode: {log_mode}. Must be one of {_SUPPORTED_LOGGING_MODES}" + ) + + @staticmethod + def _assert_valid_columns(columns): + valid_col_types = [str, int] + assert isinstance(columns, list), "columns argument expects a `list` object" + assert len(columns) == 0 or all( + [type(col) in valid_col_types for col in columns] + ), "columns argument expects list of strings or ints" + + def _init_from_list(self, data, columns, optional=True, dtype=None): + assert isinstance(data, list), "data argument expects a `list` object" + self.data = [] + self._assert_valid_columns(columns) + self.columns = columns + self._make_column_types(dtype, optional) + for row in data: + self.add_data(*row) + + def _init_from_ndarray(self, ndarray, columns, optional=True, dtype=None): + assert util.is_numpy_array(ndarray), ( + "ndarray argument expects a `numpy.ndarray` object" + ) + self.data = [] + self._assert_valid_columns(columns) + self.columns = columns + self._make_column_types(dtype, optional) + for row in ndarray: + self.add_data(*row) + + def _init_from_dataframe(self, dataframe, columns, optional=True, dtype=None): + assert util.is_pandas_data_frame(dataframe), ( + "dataframe argument expects a `pandas.core.frame.DataFrame` object" + ) + self.data = [] + columns = list(dataframe.columns) + self._assert_valid_columns(columns) + self.columns = columns + self._make_column_types(dtype, optional) + for row in range(len(dataframe)): + self.add_data(*tuple(dataframe[col].values[row] for col in self.columns)) + + def _make_column_types(self, dtype=None, optional=True): + if dtype is None: + dtype = _dtypes.UnknownType() + + if optional.__class__ is not list: + optional = [optional for _ in range(len(self.columns))] + + if dtype.__class__ is not list: + dtype = [dtype for _ in range(len(self.columns))] + + self._column_types = _dtypes.TypedDictType({}) + for col_name, opt, dt in zip(self.columns, optional, dtype): + self.cast(col_name, dt, opt) + + def _load_incremental_table_state_from_resumed_run(self, run: "LocalRun", key: str): + """Handle updating incremental table state for resumed runs. + + This method is called when a run is resumed and there are previous + increments of this table that need to be preserved. It updates the + table's internal state to track previous increments and the current + increment number. + """ + if ( + self._previous_increments_paths is not None + or self._increment_num is not None + ): + raise AssertionError( + "The table has been initialized for a resumed run already" + ) + + self._set_incremental_table_run_target(run) + + summary_from_key = run.summary.get(key) + + if ( + summary_from_key is None + or not isinstance(summary_from_key, dict) + or summary_from_key.get("_type") != "incremental-table-file" + ): + # The key was never logged to the run or its last logged + # value was not an incrementally logged table. + return + + previous_increments_paths = summary_from_key.get( + "previous_increments_paths", [] + ) + + # add the artifact path of the last logged increment + last_artifact_path = summary_from_key.get("artifact_path") + + if last_artifact_path: + previous_increments_paths.append(last_artifact_path) + + # add 1 because a new increment is being logged + last_increment_num = summary_from_key.get("increment_num", 0) + + self._increment_num = last_increment_num + 1 + self._previous_increments_paths = previous_increments_paths + + def _set_incremental_table_run_target(self, run: "LocalRun") -> None: + """Associate a Run object with this incremental Table. + + A Table object in incremental mode can only be logged to a single Run. + Raises an error if the table is already associated to a different run. + """ + if self._run_target_for_increments is None: + self._run_target_for_increments = run + elif self._run_target_for_increments is not run: + raise AssertionError("An incremental Table can only be logged to one Run.") + + @allow_relogging_after_mutation + def cast(self, col_name, dtype, optional=False): + """Casts a column to a specific data type. + + This can be one of the normal python classes, an internal W&B type, or an + example object, like an instance of wandb.Image or wandb.Classes. + + Args: + col_name: (str) - The name of the column to cast. + dtype: (class, wandb.wandb_sdk.interface._dtypes.Type, any) - The target dtype. + optional: (bool) - If the column should allow Nones. + """ + assert col_name in self.columns + + wbtype = _dtypes.TypeRegistry.type_from_dtype(dtype) + + if optional: + wbtype = _dtypes.OptionalType(wbtype) + + # Cast each value in the row, raising an error if there are invalid entries. + col_ndx = self.columns.index(col_name) + for row in self.data: + result_type = wbtype.assign(row[col_ndx]) + if isinstance(result_type, _dtypes.InvalidType): + raise TypeError( + f"Existing data {row[col_ndx]}, of type {_dtypes.TypeRegistry.type_of(row[col_ndx])} cannot be cast to {wbtype}" + ) + wbtype = result_type + + # Assert valid options + is_pk = isinstance(wbtype, _PrimaryKeyType) + is_fk = isinstance(wbtype, _ForeignKeyType) + is_fi = isinstance(wbtype, _ForeignIndexType) + if is_pk or is_fk or is_fi: + assert not optional, ( + "Primary keys, foreign keys, and foreign indexes cannot be optional." + ) + + if (is_fk or is_fk) and id(wbtype.params["table"]) == id(self): + raise AssertionError("Cannot set a foreign table reference to same table.") + + if is_pk: + assert self._pk_col is None, ( + f"Cannot have multiple primary keys - {self._pk_col} is already set as the primary key." + ) + + # Update the column type + self._column_types.params["type_map"][col_name] = wbtype + + # Wrap the data if needed + self._update_keys() + return wbtype + + def __ne__(self, other): + return not self.__eq__(other) + + def _eq_debug(self, other, should_assert=False): + eq = isinstance(other, Table) + assert not should_assert or eq, ( + f"Found type {other.__class__}, expected {Table}" + ) + eq = eq and len(self.data) == len(other.data) + assert not should_assert or eq, ( + f"Found {len(other.data)} rows, expected {len(self.data)}" + ) + eq = eq and self.columns == other.columns + assert not should_assert or eq, ( + f"Found columns {other.columns}, expected {self.columns}" + ) + eq = eq and self._column_types == other._column_types + assert not should_assert or eq, ( + f"Found column type {other._column_types}, expected column type {self._column_types}" + ) + if eq: + for row_ndx in range(len(self.data)): + for col_ndx in range(len(self.data[row_ndx])): + _eq = self.data[row_ndx][col_ndx] == other.data[row_ndx][col_ndx] + # equal if all are equal + if util.is_numpy_array(_eq): + _eq = ((_eq * -1) + 1).sum() == 0 + eq = eq and _eq + assert not should_assert or eq, ( + f"Unequal data at row_ndx {row_ndx} col_ndx {col_ndx}: found {other.data[row_ndx][col_ndx]}, expected {self.data[row_ndx][col_ndx]}" + ) + if not eq: + return eq + return eq + + def __eq__(self, other): + return self._eq_debug(other) + + @allow_relogging_after_mutation + def add_row(self, *row): + """Deprecated; use add_data instead.""" + logging.warning("add_row is deprecated, use add_data") + self.add_data(*row) + + @allow_relogging_after_mutation + @allow_incremental_logging_after_append + def add_data(self, *data): + """Adds a new row of data to the table. The maximum amount of rows in a table is determined by `wandb.Table.MAX_ARTIFACT_ROWS`. + + The length of the data should match the length of the table column. + """ + if len(data) != len(self.columns): + raise ValueError( + f"This table expects {len(self.columns)} columns: {self.columns}, found {len(data)}" + ) + + # Special case to pre-emptively cast a column as a key. + # Needed as String.assign(Key) is invalid + for ndx, item in enumerate(data): + if isinstance(item, _TableLinkMixin): + self.cast( + self.columns[ndx], + _dtypes.TypeRegistry.type_of(item), + optional=False, + ) + + # Update the table's column types + result_type = self._get_updated_result_type(data) + self._column_types = result_type + + # rows need to be mutable + if isinstance(data, tuple): + data = list(data) + # Add the new data + self.data.append(data) + + # Update the wrapper values if needed + self._update_keys(force_last=True) + + def _get_updated_result_type(self, row): + """Returns the updated result type based on the inputted row. + + Raises: + TypeError: if the assignment is invalid. + """ + incoming_row_dict = { + col_key: row[ndx] for ndx, col_key in enumerate(self.columns) + } + current_type = self._column_types + result_type = current_type.assign(incoming_row_dict) + if isinstance(result_type, _dtypes.InvalidType): + raise TypeError( + f"Data row contained incompatible types:\n{current_type.explain(incoming_row_dict)}" + ) + return result_type + + def _to_table_json(self, max_rows=None, warn=True): + # separate this method for easier testing + if max_rows is None: + max_rows = Table.MAX_ROWS + n_rows = len(self.data) + if n_rows > max_rows and warn: + # NOTE: Never raises for reinit="create_new" runs. + # Since this is called by bind_to_run(), this can be fixed by + # propagating the run. It cannot be fixed for to_json() calls + # that are given an artifact, other than by deferring to singleton + # settings. + if wandb.run and ( + wandb.run.settings.table_raise_on_max_row_limit_exceeded + or wandb.run.settings.strict + ): + raise ValueError( + f"Table row limit exceeded: table has {n_rows} rows, limit is {max_rows}. " + f"To increase the maximum number of allowed rows in a wandb.Table, override " + f"the limit with `wandb.Table.MAX_ARTIFACT_ROWS = X` and try again. Note: " + f"this may cause slower queries in the W&B UI." + ) + logging.warning(f"Truncating wandb.Table object to {max_rows} rows.") + + if self.log_mode == "INCREMENTAL" and self._last_logged_idx is not None: + return { + "columns": self.columns, + "data": self.data[ + self._last_logged_idx + 1 : self._last_logged_idx + 1 + max_rows + ], + } + else: + return {"columns": self.columns, "data": self.data[:max_rows]} + + def bind_to_run(self, *args, **kwargs): + # We set `warn=False` since Tables will now always be logged to both + # files and artifacts. The file limit will never practically matter and + # this code path will be ultimately removed. The 10k limit warning confuses + # users given that we publicly say 200k is the limit. + data = self._to_table_json(warn=False) + tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".table.json") + data = _numpy_arrays_to_lists(data) + with codecs.open(tmp_path, "w", encoding="utf-8") as fp: + util.json_dump_safer(data, fp) + self._set_file(tmp_path, is_tmp=True, extension=".table.json") + super().bind_to_run(*args, **kwargs) + + @classmethod + def get_media_subdir(cls): + return os.path.join("media", "table") + + @classmethod + def from_json(cls, json_obj, source_artifact: "artifact.Artifact"): + data = [] + column_types = None + np_deserialized_columns = {} + timestamp_column_indices = set() + log_mode = json_obj.get("log_mode", "IMMUTABLE") + if json_obj.get("column_types") is not None: + column_types = _dtypes.TypeRegistry.type_from_dict( + json_obj["column_types"], source_artifact + ) + for col_name in column_types.params["type_map"]: + col_type = column_types.params["type_map"][col_name] + ndarray_type = None + if isinstance(col_type, _dtypes.NDArrayType): + ndarray_type = col_type + elif isinstance(col_type, _dtypes.UnionType): + for t in col_type.params["allowed_types"]: + if isinstance(t, _dtypes.NDArrayType): + ndarray_type = t + elif isinstance(t, _dtypes.TimestampType): + timestamp_column_indices.add( + json_obj["columns"].index(col_name) + ) + + elif isinstance(col_type, _dtypes.TimestampType): + timestamp_column_indices.add(json_obj["columns"].index(col_name)) + + if ( + ndarray_type is not None + and ndarray_type._get_serialization_path() is not None + ): + serialization_path = ndarray_type._get_serialization_path() + + if serialization_path is None: + continue + + np = util.get_module( + "numpy", + required="Deserializing NumPy columns requires NumPy to be installed.", + ) + deserialized = np.load( + source_artifact.get_entry(serialization_path["path"]).download() + ) + np_deserialized_columns[json_obj["columns"].index(col_name)] = ( + deserialized[serialization_path["key"]] + ) + ndarray_type._clear_serialization_path() + + if log_mode == "INCREMENTAL": + unprocessed_table_data = _get_data_from_increments( + json_obj, source_artifact + ) + else: + unprocessed_table_data = json_obj["data"] + + for r_ndx, row in enumerate(unprocessed_table_data): + data.append( + _process_table_row( + row, + timestamp_column_indices, + np_deserialized_columns, + source_artifact, + r_ndx, + ) + ) + + # construct Table with dtypes for each column if type information exists + dtypes = None + if column_types is not None: + dtypes = [ + column_types.params["type_map"][str(col)] for col in json_obj["columns"] + ] + + new_obj = cls( + columns=json_obj["columns"], data=data, dtype=dtypes, log_mode=log_mode + ) + + if column_types is not None: + new_obj._column_types = column_types + + new_obj._update_keys() + return new_obj + + def to_json(self, run_or_artifact): + json_dict = super().to_json(run_or_artifact) + + if self.log_mode == "INCREMENTAL": + if self._previous_increments_paths is None: + self._previous_increments_paths = [] + if self._increment_num is None: + self._increment_num = 0 + + json_dict.update( + { + "increment_num": self._increment_num, + "previous_increments_paths": self._previous_increments_paths, + } + ) + + if isinstance(run_or_artifact, wandb.wandb_sdk.wandb_run.Run): + if self.log_mode == "INCREMENTAL": + wbvalue_type = "incremental-table-file" + else: + wbvalue_type = "table-file" + + json_dict.update( + { + "_type": wbvalue_type, + "ncols": len(self.columns), + "nrows": len(self.data), + "log_mode": self.log_mode, + } + ) + + elif isinstance(run_or_artifact, wandb.Artifact): + artifact = run_or_artifact + mapped_data = [] + data = self._to_table_json(Table.MAX_ARTIFACT_ROWS)["data"] + + ndarray_col_ndxs = set() + for col_ndx, col_name in enumerate(self.columns): + col_type = self._column_types.params["type_map"][col_name] + ndarray_type = None + if isinstance(col_type, _dtypes.NDArrayType): + ndarray_type = col_type + elif isinstance(col_type, _dtypes.UnionType): + for t in col_type.params["allowed_types"]: + if isinstance(t, _dtypes.NDArrayType): + ndarray_type = t + + # Do not serialize 1d arrays - these are likely embeddings and + # will not have the same cost as higher dimensional arrays + is_1d_array = ( + ndarray_type is not None + and "shape" in ndarray_type._params + and isinstance(ndarray_type._params["shape"], list) + and len(ndarray_type._params["shape"]) == 1 + and ndarray_type._params["shape"][0] + <= self._MAX_EMBEDDING_DIMENSIONS + ) + if is_1d_array: + self._column_types.params["type_map"][col_name] = _dtypes.ListType( + _dtypes.NumberType, ndarray_type._params["shape"][0] + ) + elif ndarray_type is not None: + np = util.get_module( + "numpy", + required="Serializing NumPy requires NumPy to be installed.", + ) + file_name = f"{str(col_name)}_{runid.generate_id()}.npz" + npz_file_name = os.path.join(MEDIA_TMP.name, file_name) + np.savez_compressed( + npz_file_name, + **{ + str(col_name): self.get_column(col_name, convert_to="numpy") + }, + ) + entry = artifact.add_file( + npz_file_name, "media/serialized_data/" + file_name, is_tmp=True + ) + ndarray_type._set_serialization_path(entry.path, str(col_name)) + ndarray_col_ndxs.add(col_ndx) + + for row in data: + mapped_row = [] + for ndx, v in enumerate(row): + if ndx in ndarray_col_ndxs: + mapped_row.append(None) + else: + mapped_row.append(_json_helper(v, artifact)) + mapped_data.append(mapped_row) + + json_dict.update( + { + "_type": Table._log_type, + "columns": self.columns, + "data": mapped_data, + "ncols": len(self.columns), + "nrows": len(mapped_data), + "column_types": self._column_types.to_json(artifact), + "log_mode": self.log_mode, + } + ) + else: + raise TypeError("to_json accepts wandb_run.Run or wandb_artifact.Artifact") + + return json_dict + + def iterrows(self): + """Returns the table data by row, showing the index of the row and the relevant data. + + Yields: + ------ + index : int + The index of the row. Using this value in other W&B tables + will automatically build a relationship between the tables + row : List[any] + The data of the row. + """ + for ndx in range(len(self.data)): + index = _TableIndex(ndx) + index.set_table(self) + yield index, self.data[ndx] + + @allow_relogging_after_mutation + def set_pk(self, col_name): + # TODO: Docs + assert col_name in self.columns + self.cast(col_name, _PrimaryKeyType()) + + @allow_relogging_after_mutation + def set_fk(self, col_name, table, table_col): + # TODO: Docs + assert col_name in self.columns + assert col_name != self._pk_col + self.cast(col_name, _ForeignKeyType(table, table_col)) + + def _update_keys(self, force_last=False): + """Updates the known key-like columns based on current column types. + + If the state has been updated since the last update, wraps the data + appropriately in the Key classes. + + Args: + force_last: (bool) Wraps the last column of data even if there + are no key updates. + """ + _pk_col = None + _fk_cols = set() + + # Buildup the known keys from column types + c_types = self._column_types.params["type_map"] + for t in c_types: + if isinstance(c_types[t], _PrimaryKeyType): + _pk_col = t + elif isinstance(c_types[t], _ForeignKeyType) or isinstance( + c_types[t], _ForeignIndexType + ): + _fk_cols.add(t) + + # If there are updates to perform, safely update them + has_update = _pk_col != self._pk_col or _fk_cols != self._fk_cols + if has_update: + # If we removed the PK + if _pk_col is None and self._pk_col is not None: + raise AssertionError( + f"Cannot unset primary key (column {self._pk_col})" + ) + # If there is a removed FK + if len(self._fk_cols - _fk_cols) > 0: + raise AssertionError( + f"Cannot unset foreign key. Attempted to unset ({self._fk_cols - _fk_cols})" + ) + + self._pk_col = _pk_col + self._fk_cols = _fk_cols + + # Apply updates to data only if there are update or the caller + # requested the final row to be updated + if has_update or force_last: + self._apply_key_updates(not has_update) + + def _apply_key_updates(self, only_last=False): + """Appropriately wraps the underlying data in special Key classes. + + Args: + only_last: only apply the updates to the last row (used for performance when + the caller knows that the only new data is the last row and no updates were + applied to the column types) + """ + c_types = self._column_types.params["type_map"] + + # Define a helper function which will wrap the data of a single row + # in the appropriate class wrapper. + def update_row(row_ndx): + for fk_col in self._fk_cols: + col_ndx = self.columns.index(fk_col) + + # Wrap the Foreign Keys + if isinstance(c_types[fk_col], _ForeignKeyType) and not isinstance( + self.data[row_ndx][col_ndx], _TableKey + ): + self.data[row_ndx][col_ndx] = _TableKey(self.data[row_ndx][col_ndx]) + self.data[row_ndx][col_ndx].set_table( + c_types[fk_col].params["table"], + c_types[fk_col].params["col_name"], + ) + + # Wrap the Foreign Indexes + elif isinstance(c_types[fk_col], _ForeignIndexType) and not isinstance( + self.data[row_ndx][col_ndx], _TableIndex + ): + self.data[row_ndx][col_ndx] = _TableIndex( + self.data[row_ndx][col_ndx] + ) + self.data[row_ndx][col_ndx].set_table( + c_types[fk_col].params["table"] + ) + + # Wrap the Primary Key + if self._pk_col is not None: + col_ndx = self.columns.index(self._pk_col) + self.data[row_ndx][col_ndx] = _TableKey(self.data[row_ndx][col_ndx]) + self.data[row_ndx][col_ndx].set_table(self, self._pk_col) + + if only_last: + update_row(len(self.data) - 1) + else: + for row_ndx in range(len(self.data)): + update_row(row_ndx) + + @ensure_not_incremental + @allow_relogging_after_mutation + def add_column(self, name, data, optional=False): + """Adds a column of data to the table. + + Args: + name: (str) - the unique name of the column + data: (list | np.array) - a column of homogeneous data + optional: (bool) - if null-like values are permitted + """ + assert isinstance(name, str) and name not in self.columns + is_np = util.is_numpy_array(data) + assert isinstance(data, list) or is_np + assert isinstance(optional, bool) + is_first_col = len(self.columns) == 0 + assert is_first_col or len(data) == len(self.data), ( + f"Expected length {len(self.data)}, found {len(data)}" + ) + + # Add the new data + for ndx in range(max(len(data), len(self.data))): + if is_first_col: + self.data.append([]) + if is_np: + self.data[ndx].append(data[ndx]) + else: + self.data[ndx].append(data[ndx]) + # add the column + self.columns.append(name) + + try: + self.cast(name, _dtypes.UnknownType(), optional=optional) + except TypeError: + # Undo the changes + if is_first_col: + self.data = [] + self.columns = [] + else: + for ndx in range(len(self.data)): + self.data[ndx] = self.data[ndx][:-1] + self.columns = self.columns[:-1] + raise + + def get_column(self, name, convert_to=None): + """Retrieves a column from the table and optionally converts it to a NumPy object. + + Args: + name: (str) - the name of the column + convert_to: (str, optional) + - "numpy": will convert the underlying data to numpy object + """ + assert name in self.columns + assert convert_to is None or convert_to == "numpy" + if convert_to == "numpy": + np = util.get_module( + "numpy", required="Converting to NumPy requires installing NumPy" + ) + col = [] + col_ndx = self.columns.index(name) + for row in self.data: + item = row[col_ndx] + if convert_to is not None and isinstance(item, WBValue): + item = item.to_data_array() + col.append(item) + if convert_to == "numpy": + col = np.array(col) + return col + + def get_index(self): + """Returns an array of row indexes for use in other tables to create links.""" + ndxs = [] + for ndx in range(len(self.data)): + index = _TableIndex(ndx) + index.set_table(self) + ndxs.append(index) + return ndxs + + def get_dataframe(self): + """Returns a `pandas.DataFrame` of the table.""" + pd = util.get_module( + "pandas", + required="Converting to pandas.DataFrame requires installing pandas", + ) + return pd.DataFrame.from_records(self.data, columns=self.columns) + + def index_ref(self, index): + """Gets a reference of the index of a row in the table.""" + assert index < len(self.data) + _index = _TableIndex(index) + _index.set_table(self) + return _index + + @ensure_not_incremental + @allow_relogging_after_mutation + def add_computed_columns(self, fn): + """Adds one or more computed columns based on existing data. + + Args: + fn: A function which accepts one or two parameters, ndx (int) and row (dict), + which is expected to return a dict representing new columns for that row, keyed + by the new column names. + + `ndx` is an integer representing the index of the row. Only included if `include_ndx` + is set to `True`. + + `row` is a dictionary keyed by existing columns + """ + new_columns = {} + for ndx, row in self.iterrows(): + row_dict = {self.columns[i]: row[i] for i in range(len(self.columns))} + new_row_dict = fn(ndx, row_dict) + assert isinstance(new_row_dict, dict) + for key in new_row_dict: + new_columns[key] = new_columns.get(key, []) + new_columns[key].append(new_row_dict[key]) + for new_col_name in new_columns: + self.add_column(new_col_name, new_columns[new_col_name]) + + +class _PartitionTablePartEntry: + """Helper class for PartitionTable to track its parts.""" + + def __init__(self, entry, source_artifact): + self.entry = entry + self.source_artifact = source_artifact + self._part = None + + def get_part(self): + if self._part is None: + self._part = self.source_artifact.get(self.entry.path) + return self._part + + def free(self): + self._part = None + + +class PartitionedTable(Media): + """A table which is composed of multiple sub-tables. + + Currently, PartitionedTable is designed to point to a directory within an artifact. + """ + + _log_type = "partitioned-table" + + def __init__(self, parts_path): + """Initialize a PartitionedTable. + + Args: + parts_path (str): path to a directory of tables in the artifact. + """ + super().__init__() + self.parts_path = parts_path + self._loaded_part_entries = {} + + def to_json(self, artifact_or_run): + json_obj = { + "_type": PartitionedTable._log_type, + } + if isinstance(artifact_or_run, wandb.wandb_sdk.wandb_run.Run): + artifact_entry_url = self._get_artifact_entry_ref_url() + if artifact_entry_url is None: + raise ValueError( + "PartitionedTables must first be added to an Artifact before logging to a Run" + ) + json_obj["artifact_path"] = artifact_entry_url + else: + json_obj["parts_path"] = self.parts_path + return json_obj + + @classmethod + def from_json(cls, json_obj, source_artifact): + instance = cls(json_obj["parts_path"]) + entries = source_artifact.manifest.get_entries_in_directory( + json_obj["parts_path"] + ) + for entry in entries: + instance._add_part_entry(entry, source_artifact) + return instance + + def iterrows(self): + """Iterate over rows as (ndx, row). + + Yields: + ------ + index : int + The index of the row. + row : List[any] + The data of the row. + """ + columns = None + ndx = 0 + for entry_path in self._loaded_part_entries: + part = self._loaded_part_entries[entry_path].get_part() + if columns is None: + columns = part.columns + elif columns != part.columns: + raise ValueError( + f"Table parts have non-matching columns. {columns} != {part.columns}" + ) + for _, row in part.iterrows(): + yield ndx, row + ndx += 1 + + self._loaded_part_entries[entry_path].free() + + def _add_part_entry(self, entry, source_artifact): + self._loaded_part_entries[entry.path] = _PartitionTablePartEntry( + entry, source_artifact + ) + + def __ne__(self, other): + return not self.__eq__(other) + + def __eq__(self, other): + return isinstance(other, self.__class__) and self.parts_path == other.parts_path + + def bind_to_run(self, *args, **kwargs): + raise ValueError("PartitionedTables cannot be bound to runs") + + +class JoinedTable(Media): + """Join two tables for visualization in the Artifact UI. + + Args: + table1 (str, wandb.Table, ArtifactManifestEntry): + the path to a wandb.Table in an artifact, the table object, or ArtifactManifestEntry + table2 (str, wandb.Table): + the path to a wandb.Table in an artifact, the table object, or ArtifactManifestEntry + join_key (str, [str, str]): + key or keys to perform the join + """ + + _log_type = "joined-table" + + def __init__(self, table1, table2, join_key): + super().__init__() + + if not isinstance(join_key, str) and ( + not isinstance(join_key, list) or len(join_key) != 2 + ): + raise ValueError( + "JoinedTable join_key should be a string or a list of two strings" + ) + + if not self._validate_table_input(table1): + raise ValueError( + "JoinedTable table1 should be an artifact path to a table or wandb.Table object" + ) + + if not self._validate_table_input(table2): + raise ValueError( + "JoinedTable table2 should be an artifact path to a table or wandb.Table object" + ) + + self._table1 = table1 + self._table2 = table2 + self._join_key = join_key + + @classmethod + def from_json(cls, json_obj, source_artifact): + t1 = source_artifact.get(json_obj["table1"]) + if t1 is None: + t1 = json_obj["table1"] + + t2 = source_artifact.get(json_obj["table2"]) + if t2 is None: + t2 = json_obj["table2"] + + return cls( + t1, + t2, + json_obj["join_key"], + ) + + @staticmethod + def _validate_table_input(table): + """Helper method to validate that the table input is one of the 3 supported types.""" + return ( + (isinstance(table, str) and table.endswith(".table.json")) + or isinstance(table, Table) + or isinstance(table, PartitionedTable) + or (hasattr(table, "ref_url") and table.ref_url().endswith(".table.json")) + ) + + def _ensure_table_in_artifact(self, table, artifact, table_ndx): + """Helper method to add the table to the incoming artifact. Returns the path.""" + if isinstance(table, Table) or isinstance(table, PartitionedTable): + table_name = f"t{table_ndx}_{str(id(self))}" + if ( + table._artifact_source is not None + and table._artifact_source.name is not None + ): + table_name = os.path.basename(table._artifact_source.name) + entry = artifact.add(table, table_name) + table = entry.path + # Check if this is an ArtifactManifestEntry + elif hasattr(table, "ref_url"): + # Give the new object a unique, yet deterministic name + name = binascii.hexlify(base64.standard_b64decode(table.digest)).decode( + "ascii" + )[:20] + entry = artifact.add_reference( + table.ref_url(), "{}.{}.json".format(name, table.name.split(".")[-2]) + )[0] + table = entry.path + + err_str = "JoinedTable table:{} not found in artifact. Add a table to the artifact using Artifact#add(, {}) before adding this JoinedTable" + if table not in artifact._manifest.entries: + raise ValueError(err_str.format(table, table)) + + return table + + def to_json(self, artifact_or_run): + json_obj = { + "_type": JoinedTable._log_type, + } + if isinstance(artifact_or_run, wandb.wandb_sdk.wandb_run.Run): + artifact_entry_url = self._get_artifact_entry_ref_url() + if artifact_entry_url is None: + raise ValueError( + "JoinedTables must first be added to an Artifact before logging to a Run" + ) + json_obj["artifact_path"] = artifact_entry_url + else: + table1 = self._ensure_table_in_artifact(self._table1, artifact_or_run, 1) + table2 = self._ensure_table_in_artifact(self._table2, artifact_or_run, 2) + json_obj.update( + { + "table1": table1, + "table2": table2, + "join_key": self._join_key, + } + ) + return json_obj + + def __ne__(self, other): + return not self.__eq__(other) + + def _eq_debug(self, other, should_assert=False): + eq = isinstance(other, JoinedTable) + assert not should_assert or eq, ( + f"Found type {other.__class__}, expected {JoinedTable}" + ) + eq = eq and self._join_key == other._join_key + assert not should_assert or eq, ( + f"Found {other._join_key} join key, expected {self._join_key}" + ) + eq = eq and self._table1._eq_debug(other._table1, should_assert) + eq = eq and self._table2._eq_debug(other._table2, should_assert) + return eq + + def __eq__(self, other): + return self._eq_debug(other, False) + + def bind_to_run(self, *args, **kwargs): + raise ValueError("JoinedTables cannot be bound to runs") + + +class _TableType(_dtypes.Type): + name = "table" + legacy_names = ["wandb.Table"] + types = [Table] + + def __init__(self, column_types=None): + if column_types is None: + column_types = _dtypes.UnknownType() + if isinstance(column_types, dict): + column_types = _dtypes.TypedDictType(column_types) + elif not ( + isinstance(column_types, _dtypes.TypedDictType) + or isinstance(column_types, _dtypes.UnknownType) + ): + raise TypeError("column_types must be a dict or TypedDictType") + + self.params.update({"column_types": column_types}) + + def assign_type(self, wb_type=None): + if isinstance(wb_type, _TableType): + column_types = self.params["column_types"].assign_type( + wb_type.params["column_types"] + ) + if not isinstance(column_types, _dtypes.InvalidType): + return _TableType(column_types) + + return _dtypes.InvalidType() + + @classmethod + def from_obj(cls, py_obj): + if not isinstance(py_obj, Table): + raise TypeError("py_obj must be a wandb.Table") + else: + return cls(py_obj._column_types) + + +class _JoinedTableType(_dtypes.Type): + name = "joined-table" + types = [JoinedTable] + + +class _PartitionedTableType(_dtypes.Type): + name = "partitioned-table" + types = [PartitionedTable] + + +_dtypes.TypeRegistry.add(_TableType) +_dtypes.TypeRegistry.add(_JoinedTableType) +_dtypes.TypeRegistry.add(_PartitionedTableType) +_dtypes.TypeRegistry.add(_ForeignKeyType) +_dtypes.TypeRegistry.add(_PrimaryKeyType) +_dtypes.TypeRegistry.add(_ForeignIndexType) + + +def _get_data_from_increments( + json_obj: Dict[str, Any], source_artifact: "artifact.Artifact" +) -> List[Any]: + """Get data from incremental table artifacts. + + Args: + json_obj: The JSON object containing table metadata. + source_artifact: The source artifact containing the table data. + + Returns: + List of table rows from all increments. + """ + if "latest" not in source_artifact.aliases: + wandb.termwarn( + ( + "It is recommended to use the latest version of the " + "incremental table artifact for ordering guarantees." + ), + repeat=False, + ) + data: List[Any] = [] + increment_num = json_obj.get("increment_num", None) + if increment_num is None: + return data + + # Sort by increment number first, then by timestamp if present + # Format of name is: "{incr_num}-{timestamp_ms}.{key}.table.json" + def get_sort_key(key: str) -> Tuple[int, int]: + try: + parts = key.split(".") + increment_parts = parts[0].split("-") + increment_num = int(increment_parts[0]) + # If there's a timestamp part, use it for secondary sorting + timestamp = int(increment_parts[1]) if len(increment_parts) > 1 else 0 + except (ValueError, IndexError): + wandb.termwarn( + ( + f"Could not parse artifact entry for increment {key}." + " The entry name does not follow the naming convention" + " -..table.json" + " The data in the table will be out of order." + ), + repeat=False, + ) + return (0, 0) + + return (increment_num, timestamp) + + sorted_increment_keys = [] + for entry_key in source_artifact.manifest.entries: + if entry_key.endswith(".table.json"): + sorted_increment_keys.append(entry_key) + + sorted_increment_keys.sort(key=get_sort_key) + + for entry_key in sorted_increment_keys: + try: + with open(source_artifact.manifest.entries[entry_key].download()) as f: + table_data = json.load(f) + data.extend(table_data["data"]) + except (json.JSONDecodeError, KeyError) as e: + raise wandb.Error(f"Invalid table file {entry_key}") from e + return data + + +def _process_table_row( + row: List[Any], + timestamp_column_indices: Set[_dtypes.TimestampType], + np_deserialized_columns: Dict[int, Any], + source_artifact: "artifact.Artifact", + row_idx: int, +) -> List[Any]: + """Convert special columns in a table row to Python types. + + Processes a single row of table data by converting timestamp values to + datetime objects, replacing np typed cells with numpy array data, + and initializing media objects from their json value. + + + Args: + row: The row data to process. + timestamp_column_indices: Set of column indices containing timestamps. + np_deserialized_columns: Dictionary mapping column indices to numpy arrays. + source_artifact: The source artifact containing the table data. + row_idx: The index of the current row. + + Returns: + Processed row data. + """ + row_data = [] + for c_ndx, item in enumerate(row): + cell: Any + if c_ndx in timestamp_column_indices and isinstance(item, (int, float)): + cell = datetime.datetime.fromtimestamp( + item / 1000, tz=datetime.timezone.utc + ) + elif c_ndx in np_deserialized_columns: + cell = np_deserialized_columns[c_ndx][row_idx] + elif ( + isinstance(item, dict) + and "_type" in item + and (obj := WBValue.init_from_json(item, source_artifact)) + ): + cell = obj + else: + cell = item + row_data.append(cell) + return row_data diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/table_decorators.py b/lib/python3.12/site-packages/wandb/sdk/data_types/table_decorators.py new file mode 100644 index 0000000000000000000000000000000000000000..5c4ead3462d8745ef21c319321c9773667a2ee4d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/table_decorators.py @@ -0,0 +1,108 @@ +"""Decorators for W&B Table operations.""" + +from __future__ import annotations + +from functools import wraps +from typing import Any, Callable, TypeVar + +from typing_extensions import Concatenate, ParamSpec + +import wandb + +_P = ParamSpec("_P") +_T = TypeVar("_T") + + +def allow_relogging_after_mutation( + method: Callable[Concatenate[wandb.Table, _P], _T], +) -> Callable[Concatenate[wandb.Table, _P], _T]: + """Decorator that handles table state after mutations based on log_mode. + + For MUTABLE tables, resets the run and artifact target to allow re-logging. + For IMMUTABLE tables, warns if attempting to mutate after logging. + """ + + @wraps(method) + def wrapper(self: wandb.Table, *args: Any, **kwargs: Any) -> _T: + has_been_logged = self._run is not None or self._artifact_target is not None + + if self.log_mode == "MUTABLE": + self._run = None + self._artifact_target = None + self._path = None + self._sha256 = None + elif self.log_mode == "IMMUTABLE" and has_been_logged: + wandb.termwarn( + "You are mutating a Table with log_mode='IMMUTABLE' that has been " + "logged already. Subsequent log() calls will have no effect. " + "Set log_mode='MUTABLE' to enable re-logging after mutations", + repeat=False, + ) + + return method(self, *args, **kwargs) + + return wrapper + + +def allow_incremental_logging_after_append( + method: Callable[Concatenate[wandb.Table, _P], _T], +) -> Callable[Concatenate[wandb.Table, _P], _T]: + """Decorator that handles incremental logging state after append operations. + + For INCREMENTAL tables, manages artifact references and increments counters + to support partial data logging. + """ + + @wraps(method) + def wrapper(self: wandb.Table, *args: Any, **kwargs: Any) -> _T: + res = method(self, *args, **kwargs) + + if self.log_mode != "INCREMENTAL" or self._artifact_target is None: + return res + + art_entry_url = self._get_artifact_entry_ref_url() + if art_entry_url is not None: + if self._previous_increments_paths is None: + raise ValueError( + "_previous_increments_paths must be set if incremental" + " table has been logged already" + ) + self._previous_increments_paths.append(art_entry_url) + self._run = None + self._artifact_target = None + self._path = None + self._sha256 = None + + if self._increment_num is None: + raise ValueError( + "_increment_num must be set if incremental table has been" + " logged already" + ) + + self._increment_num += 1 + if self._increment_num > 99: + wandb.termwarn( + "You have exceeded 100 increments for this table. " + "Only the latest 100 increments will be visualized in the run workspace.", + repeat=False, + ) + return res + + return wrapper + + +def ensure_not_incremental( + method: Callable[Concatenate[wandb.Table, _P], _T], +) -> Callable[Concatenate[wandb.Table, _P], _T]: + """Decorator that checks if log mode is incremental to disallow methods from being called.""" + + @wraps(method) + def wrapper(self: wandb.Table, *args: Any, **kwargs: Any) -> _T: + if self.log_mode == "INCREMENTAL": + raise wandb.Error( + f"Operation '{method.__name__}' is not supported for tables with " + "log_mode='INCREMENTAL'. Use a different log mode like 'MUTABLE' or 'IMMUTABLE'." + ) + return method(self, *args, **kwargs) + + return wrapper diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/trace_tree.py b/lib/python3.12/site-packages/wandb/sdk/data_types/trace_tree.py new file mode 100644 index 0000000000000000000000000000000000000000..d42d407c8a4ad9397306be9e3ea5a37383496a6e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/trace_tree.py @@ -0,0 +1,440 @@ +"""This module contains the `WBTraceTree` media type, and the supporting dataclasses. + +A `WBTraceTree` is a media object containing a root span and an +arbitrary model dump as a serializable dictionary. Logging such media type will +result in a W&B Trace Debugger panel being created in the workspace UI. +""" + +import dataclasses +import hashlib +import json +import typing +from dataclasses import dataclass, field +from enum import Enum +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import wandb +from wandb.sdk.data_types import _dtypes +from wandb.sdk.data_types.base_types.media import Media +from wandb.sdk.data_types.utils import _json_helper + +if TYPE_CHECKING: # pragma: no cover + from wandb.sdk.artifacts.artifact import Artifact + + from ..wandb_run import Run as LocalRun + + +class StatusCode(str, Enum): + SUCCESS = "SUCCESS" + ERROR = "ERROR" + + def __str__(self) -> str: + return str(self.value) + + +class SpanKind(str, Enum): + LLM = "LLM" + CHAIN = "CHAIN" + AGENT = "AGENT" + TOOL = "TOOL" + + def __str__(self) -> str: + return str(self.value) + + +@dataclass() +class Result: + inputs: Optional[Dict[str, Any]] = field(default=None) + outputs: Optional[Dict[str, Any]] = field(default=None) + + +@dataclass() +class Span: + span_id: Optional[str] = field(default=None) + name: Optional[str] = field(default=None) + start_time_ms: Optional[int] = field(default=None) + end_time_ms: Optional[int] = field(default=None) + status_code: Optional[StatusCode] = field(default=None) + status_message: Optional[str] = field(default=None) + attributes: Optional[Dict[str, Any]] = field(default=None) + results: Optional[List[Result]] = field(default=None) + child_spans: Optional[List["Span"]] = field(default=None) + span_kind: Optional[SpanKind] = field(default=None) + + def add_attribute(self, key: str, value: Any) -> None: + if self.attributes is None: + self.attributes = {} + self.attributes[key] = value + + def add_named_result(self, inputs: Dict[str, Any], outputs: Dict[str, Any]) -> None: + if self.results is None: + self.results = [] + self.results.append(Result(inputs, outputs)) + + def add_child_span(self, span: "Span") -> None: + if self.child_spans is None: + self.child_spans = [] + self.child_spans.append(span) + + +class WBTraceTree(Media): + """Media object for trace tree data. + + Args: + root_span (Span): The root span of the trace tree. + model_dict (dict, optional): A dictionary containing the model dump. + NOTE: model_dict is a completely-user-defined dict. The UI will render + a JSON viewer for this dict, giving special treatment to dictionaries + with a `_kind` key. This is because model vendors have such different + serialization formats that we need to be flexible here. + """ + + _log_type = "wb_trace_tree" + + def __init__( + self, + root_span: Span, + model_dict: typing.Optional[dict] = None, + ): + super().__init__() + self._root_span = root_span + self._model_dict = model_dict + + @classmethod + def get_media_subdir(cls) -> str: + return "media/wb_trace_tree" + + def to_json(self, run: Optional[Union["LocalRun", "Artifact"]]) -> dict: + res = {"_type": self._log_type} + # Here we use `dumps` to put things into string format. This is because + # the complex data structures create problems for gorilla history to parquet. + if self._model_dict is not None: + model_dump_str = _safe_serialize(self._model_dict) + res["model_hash"] = _hash_id(model_dump_str) + res["model_dict_dumps"] = model_dump_str + res["root_span_dumps"] = _safe_serialize(dataclasses.asdict(self._root_span)) + return res + + def is_bound(self) -> bool: + return True + + +class _WBTraceTreeFileType(_dtypes.Type): + name = "wb_trace_tree" + types = [WBTraceTree] + + +_dtypes.TypeRegistry.add(_WBTraceTreeFileType) + + +# generate a deterministic 16 character id based on input string +def _hash_id(s: str) -> str: + return hashlib.md5(s.encode("utf-8")).hexdigest()[:16] + + +def _fallback_serialize(obj: Any) -> str: + try: + return f"<>" + except Exception: + return "<>" + + +def _safe_serialize(obj: dict) -> str: + try: + return json.dumps( + _json_helper(obj, None), + skipkeys=True, + default=_fallback_serialize, + ) + except Exception: + return "{}" + + +class TraceAttribute: + """Descriptor for accessing and setting attributes of the `Trace` class.""" + + def __set_name__(self, owner: type, name: str) -> None: + self.name = name + + def __get__(self, instance: "Trace", owner: type) -> Any: + return getattr(instance._span, self.name) + + def __set__(self, instance: "Trace", value: Any) -> None: + setattr(instance._span, self.name, value) + + +class Trace: + """A simplification of WBTraceTree and Span to manage a trace - a collection of spans, their metadata and hierarchy. + + Args: + name: (str) The name of the root span. + kind: (str, optional) The kind of the root span. + status_code: (str, optional) The status of the root span, either "error" or "success". + status_message: (str, optional) Any status message associated with the root span. + metadata: (dict, optional) Any additional metadata for the root span. + start_time_ms: (int, optional) The start time of the root span in milliseconds. + end_time_ms: (int, optional) The end time of the root span in milliseconds. + inputs: (dict, optional) The named inputs of the root span. + outputs: (dict, optional) The named outputs of the root span. + model_dict: (dict, optional) A json serializable dictionary containing the model architecture details. + + Example: + .. code-block:: python + ``` + trace = Trace( + name="My awesome Model", + kind="LLM", + status_code= "SUCCESS", + metadata={"attr_1": 1, "attr_2": 2,}, + start_time_ms=int(round(time.time() * 1000)), + end_time_ms=int(round(time.time() * 1000))+1000, + inputs={"user": "How old is google?"}, + outputs={"assistant": "25 years old"}, + model_dict={"_kind": "openai", "api_type": "azure"} + ) + run = wandb.init(project=,) + trace.log("my_trace") + wandb.finish() + ``` + """ + + name = TraceAttribute() + status_code = TraceAttribute() + status_message = TraceAttribute() + start_time_ms = TraceAttribute() + end_time_ms = TraceAttribute() + + def __init__( + self, + name: str, + kind: Optional[str] = None, + status_code: Optional[str] = None, + status_message: Optional[str] = None, + metadata: Optional[dict] = None, + start_time_ms: Optional[int] = None, + end_time_ms: Optional[int] = None, + inputs: Optional[dict] = None, + outputs: Optional[dict] = None, + model_dict: Optional[dict] = None, + ): + self._span = self._assert_and_create_span( + name=name, + kind=kind, + status_code=status_code, + status_message=status_message, + metadata=metadata, + start_time_ms=start_time_ms, + end_time_ms=end_time_ms, + inputs=inputs, + outputs=outputs, + ) + if model_dict is not None: + assert isinstance(model_dict, dict), "Model dict must be a dictionary" + self._model_dict = model_dict + + def _assert_and_create_span( + self, + name: str, + kind: Optional[str] = None, + status_code: Optional[str] = None, + status_message: Optional[str] = None, + metadata: Optional[dict] = None, + start_time_ms: Optional[int] = None, + end_time_ms: Optional[int] = None, + inputs: Optional[dict] = None, + outputs: Optional[dict] = None, + ) -> Span: + """Utility to assert the validity of the span parameters and create a span object. + + Args: + name: The name of the span. + kind: The kind of the span. + status_code: The status code of the span. + status_message: The status message of the span. + metadata: Dictionary of metadata to be logged with the span. + start_time_ms: Start time of the span in milliseconds. + end_time_ms: End time of the span in milliseconds. + inputs: Dictionary of inputs to be logged with the span. + outputs: Dictionary of outputs to be logged with the span. + + Returns: + A Span object. + """ + if kind is not None: + assert kind.upper() in SpanKind.__members__, ( + "Invalid span kind, can be one of 'LLM', 'AGENT', 'CHAIN', 'TOOL'" + ) + kind = SpanKind(kind.upper()) + if status_code is not None: + assert status_code.upper() in StatusCode.__members__, ( + "Invalid status code, can be one of 'SUCCESS' or 'ERROR'" + ) + status_code = StatusCode(status_code.upper()) + if inputs is not None: + assert isinstance(inputs, dict), "Inputs must be a dictionary" + if outputs is not None: + assert isinstance(outputs, dict), "Outputs must be a dictionary" + if inputs or outputs: + result = Result(inputs=inputs, outputs=outputs) + else: + result = None + + if metadata is not None: + assert isinstance(metadata, dict), "Metadata must be a dictionary" + + return Span( + name=name, + span_kind=kind, + status_code=status_code, + status_message=status_message, + attributes=metadata, + start_time_ms=start_time_ms, + end_time_ms=end_time_ms, + results=[result] if result else None, + ) + + def add_child( + self, + child: "Trace", + ) -> "Trace": + """Utility to add a child span to the current span of the trace. + + Args: + child: The child span to be added to the current span of the trace. + + Returns: + The current trace object with the child span added to it. + """ + self._span.add_child_span(child._span) + if self._model_dict is not None and child._model_dict is not None: + self._model_dict.update({child._span.name: child._model_dict}) + return self + + def add_inputs_and_outputs(self, inputs: dict, outputs: dict) -> "Trace": + """Add a result to the span of the current trace. + + Args: + inputs: Dictionary of inputs to be logged with the span. + outputs: Dictionary of outputs to be logged with the span. + + Returns: + The current trace object with the result added to it. + """ + if self._span.results is None: + result = Result(inputs=inputs, outputs=outputs) + self._span.results = [result] + else: + result = Result(inputs=inputs, outputs=outputs) + self._span.results.append(result) + return self + + def add_metadata(self, metadata: dict) -> "Trace": + """Add metadata to the span of the current trace.""" + if self._span.attributes is None: + self._span.attributes = metadata + else: + self._span.attributes.update(metadata) + return self + + @property + def metadata(self) -> Optional[Dict[str, str]]: + """Get the metadata of the trace. + + Returns: + Dictionary of metadata. + """ + return self._span.attributes + + @metadata.setter + def metadata(self, value: Dict[str, str]) -> None: + """Set the metadata of the trace. + + Args: + value: Dictionary of metadata to be set. + """ + if self._span.attributes is None: + self._span.attributes = value + else: + self._span.attributes.update(value) + + @property + def inputs(self) -> Optional[Dict[str, str]]: + """Get the inputs of the trace. + + Returns: + Dictionary of inputs. + """ + return self._span.results[-1].inputs if self._span.results else None + + @inputs.setter + def inputs(self, value: Dict[str, str]) -> None: + """Set the inputs of the trace. + + Args: + value: Dictionary of inputs to be set. + """ + if self._span.results is None: + result = Result(inputs=value, outputs={}) + self._span.results = [result] + else: + result = Result(inputs=value, outputs=self._span.results[-1].outputs) + self._span.results.append(result) + + @property + def outputs(self) -> Optional[Dict[str, str]]: + """Get the outputs of the trace. + + Returns: + Dictionary of outputs. + """ + return self._span.results[-1].outputs if self._span.results else None + + @outputs.setter + def outputs(self, value: Dict[str, str]) -> None: + """Set the outputs of the trace. + + Args: + value: Dictionary of outputs to be set. + """ + if self._span.results is None: + result = Result(inputs={}, outputs=value) + self._span.results = [result] + else: + result = Result(inputs=self._span.results[-1].inputs, outputs=value) + self._span.results.append(result) + + @property + def kind(self) -> Optional[str]: + """Get the kind of the trace. + + Returns: + The kind of the trace. + """ + return self._span.span_kind.value if self._span.span_kind else None + + @kind.setter + def kind(self, value: str) -> None: + """Set the kind of the trace. + + Args: + value: The kind of the trace to be set. + """ + assert value.upper() in SpanKind.__members__, ( + "Invalid span kind, can be one of 'LLM', 'AGENT', 'CHAIN', 'TOOL'" + ) + self._span.span_kind = SpanKind(value.upper()) + + def log(self, name: str) -> None: + """Log the trace to a wandb run. + + Args: + name: The name of the trace to be logged + """ + trace_tree = WBTraceTree(self._span, self._model_dict) + # NOTE: Does not work for reinit="create_new" runs. + # This method should be deprecated and users should call run.log(). + assert wandb.run is not None, ( + "You must call wandb.init() before logging a trace" + ) + assert len(name.strip()) > 0, "You must provide a valid name to log the trace" + wandb.run.log({name: trace_tree}) diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/utils.py b/lib/python3.12/site-packages/wandb/sdk/data_types/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0a7515446782ddce590f814f9546aea8e2887c4e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/utils.py @@ -0,0 +1,260 @@ +import datetime +import logging +import os +from decimal import Decimal +from typing import TYPE_CHECKING, Optional, Sequence, Union, cast + +import wandb +from wandb import util + +from ..internal import incremental_table_util +from .base_types.media import BatchableMedia, Media +from .base_types.wb_value import WBValue +from .image import _server_accepts_image_filenames +from .plotly import Plotly + +if TYPE_CHECKING: # pragma: no cover + import matplotlib # type: ignore + import pandas as pd + import plotly # type: ignore + + from ..wandb_run import Run as LocalRun + + ValToJsonType = Union[ + dict, + "WBValue", + Sequence["WBValue"], + "plotly.Figure", + "matplotlib.artist.Artist", + "pd.DataFrame", + object, + ] + + +def history_dict_to_json( + run: Optional["LocalRun"], + payload: dict, + step: Optional[int] = None, + ignore_copy_err: Optional[bool] = None, +) -> dict: + # Converts a History row dict's elements so they're friendly for JSON serialization. + + if step is None: + # We should be at the top level of the History row; assume this key is set. + step = payload["_step"] + + # We use list here because we were still seeing cases of RuntimeError dict changed size + for key in list(payload): + val = payload[key] + if isinstance(val, dict): + payload[key] = history_dict_to_json( + run, val, step=step, ignore_copy_err=ignore_copy_err + ) + else: + payload[key] = val_to_json( + run, key, val, namespace=step, ignore_copy_err=ignore_copy_err + ) + + return payload + + +# TODO: refine this +def val_to_json( + run: Optional["LocalRun"], + key: str, + val: "ValToJsonType", + namespace: Optional[Union[str, int]] = None, + ignore_copy_err: Optional[bool] = None, +) -> Union[Sequence, dict]: + # Converts a wandb datatype to its JSON representation. + if namespace is None: + raise ValueError( + "val_to_json must be called with a namespace(a step number, or 'summary') argument" + ) + + converted = val + + if isinstance(val, (int, float, str, bool)): + # These are already JSON-serializable, + # no need to do the expensive checks below. + return converted # type: ignore[return-value] + + typename = util.get_full_typename(val) + + if util.is_pandas_data_frame(val): + val = wandb.Table(dataframe=val) + + elif util.is_matplotlib_typename(typename) or util.is_plotly_typename(typename): + val = Plotly.make_plot_media(val) + elif isinstance(val, (list, tuple, range)) and all( + isinstance(v, WBValue) for v in val + ): + assert run + # This check will break down if Image/Audio/... have child classes. + if ( + len(val) + and isinstance(val[0], BatchableMedia) + and all(isinstance(v, type(val[0])) for v in val) + ): + if TYPE_CHECKING: + val = cast(Sequence["BatchableMedia"], val) + + items = _prune_max_seq(val) + + if _server_accepts_image_filenames(run): + for item in items: + item.bind_to_run( + run=run, + key=key, + step=namespace, + ignore_copy_err=ignore_copy_err, + ) + else: + for i, item in enumerate(items): + item.bind_to_run( + run=run, + key=key, + step=namespace, + id_=i, + ignore_copy_err=ignore_copy_err, + ) + if run._attach_id and run._init_pid != os.getpid(): + wandb.termwarn( + f"Attempting to log a sequence of {items[0].__class__.__name__} objects from multiple processes might result in data loss. Please upgrade your wandb server", + repeat=False, + ) + + return items[0].seq_to_json(items, run, key, namespace) + else: + # TODO(adrian): Good idea to pass on the same key here? Maybe include + # the array index? + # There is a bug here: if this array contains two arrays of the same type of + # anonymous media objects, their eventual names will collide. + # This used to happen. The frontend doesn't handle heterogeneous arrays + # raise ValueError( + # "Mixed media types in the same list aren't supported") + return [ + val_to_json( + run, key, v, namespace=namespace, ignore_copy_err=ignore_copy_err + ) + for v in val + ] + + if isinstance(val, WBValue): + assert run + if isinstance(val, Media) and not val.is_bound(): + if hasattr(val, "_log_type") and val._log_type in [ + "table", + "partitioned-table", + "joined-table", + ]: + _log_table_artifact(val, key, run) + + # Partitioned tables and joined tables do not support being bound to runs. + if not ( + hasattr(val, "_log_type") + and val._log_type in ["partitioned-table", "joined-table"] + ): + val.bind_to_run(run, key, namespace) + + res = val.to_json(run) + + if isinstance(val, wandb.Table) and val.log_mode == "INCREMENTAL": + # Set the _last_logged_idx AFTER the Table has been logged and + # bound to the run. + val._last_logged_idx = len(val.data) - 1 + return res + + return converted # type: ignore + + +def _log_table_artifact(val: "Media", key: str, run: "LocalRun") -> None: + """Log a table to the run based on the table type and logging mode. + + Creates and logs a `run_table` type for Table, PartitionedTable, and + JoinedTable values. For tables with log_mode="INCREMENTAL", creates and + logs an incremental artifact of type `wandb-run-incremental-table.` + + Args: + val: A wbvalue with log_type "table", "partitioned-table", + or "joined-table." + key: The key used to log val. + run: The LocalRun used to log val. + """ + from wandb.sdk.artifacts._internal_artifact import InternalArtifact + + if isinstance(val, wandb.Table) and val.log_mode == "INCREMENTAL": + if ( + run.resumed + and val._previous_increments_paths is None + and val._increment_num is None + ): + val._load_incremental_table_state_from_resumed_run(run, key) + else: + val._set_incremental_table_run_target(run) + art = incremental_table_util.init_artifact(run, key) + entry_name = incremental_table_util.get_entry_name(val, key) + else: + art = InternalArtifact(f"run-{run.id}-{key}", "run_table") + entry_name = key + + art.add(val, entry_name) + run.log_artifact(art) + + +def _prune_max_seq(seq: Sequence["BatchableMedia"]) -> Sequence["BatchableMedia"]: + # If media type has a max respect it + items = seq + if hasattr(seq[0], "MAX_ITEMS") and seq[0].MAX_ITEMS < len(seq): + logging.warning( + f"Only {seq[0].MAX_ITEMS} {seq[0].__class__.__name__} will be uploaded." + ) + items = seq[: seq[0].MAX_ITEMS] + return items + + +def _json_helper(val, artifact): + if isinstance(val, WBValue): + return val.to_json(artifact) + elif val.__class__ is dict: + res = {} + for key in val: + res[key] = _json_helper(val[key], artifact) + return res + + if hasattr(val, "tolist"): + py_val = val.tolist() + if val.__class__.__name__ == "datetime64" and isinstance(py_val, int): + # when numpy datetime64 .tolist() returns an int, it is nanoseconds. + # need to convert to milliseconds + return _json_helper(py_val / int(1e6), artifact) + return _json_helper(py_val, artifact) + elif hasattr(val, "item"): + return _json_helper(val.item(), artifact) + + if isinstance(val, datetime.datetime): + if val.tzinfo is None: + val = datetime.datetime( + val.year, + val.month, + val.day, + val.hour, + val.minute, + val.second, + val.microsecond, + tzinfo=datetime.timezone.utc, + ) + return int(val.timestamp() * 1000) + elif isinstance(val, datetime.date): + return int( + datetime.datetime( + val.year, val.month, val.day, tzinfo=datetime.timezone.utc + ).timestamp() + * 1000 + ) + elif isinstance(val, (list, tuple)): + return [_json_helper(i, artifact) for i in val] + elif isinstance(val, Decimal): + return float(val) + else: + return util.json_friendly(val)[0] diff --git a/lib/python3.12/site-packages/wandb/sdk/data_types/video.py b/lib/python3.12/site-packages/wandb/sdk/data_types/video.py new file mode 100644 index 0000000000000000000000000000000000000000..9a6bbc499fdb55e0cfd916e9618fb7e3f022ab76 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/data_types/video.py @@ -0,0 +1,276 @@ +import functools +import logging +import os +import pathlib +from io import BytesIO +from typing import TYPE_CHECKING, Any, Literal, Optional, Sequence, Type, Union + +import wandb +from wandb import util +from wandb.sdk.lib import filesystem, printer, printer_asyncio, runid + +from . import _dtypes +from ._private import MEDIA_TMP +from .base_types.media import BatchableMedia + +if TYPE_CHECKING: # pragma: no cover + from typing import TextIO + + import numpy as np + + from wandb.sdk.artifacts.artifact import Artifact + + from ..wandb_run import Run as LocalRun + + +# This helper function is a workaround for the issue discussed here: +# https://github.com/wandb/wandb/issues/3472 +# +# Essentially, the issue is that moviepy's write_gif function fails to close +# the open write / file descriptor returned from `imageio.save`. The following +# function is a simplified copy of the function in the moviepy source code. +# See https://github.com/Zulko/moviepy/blob/7e3e8bb1b739eb6d1c0784b0cb2594b587b93b39/moviepy/video/io/gif_writers.py#L428 +# +# Except, we close the writer! +def write_gif_with_image_io( + clip: Any, filename: str, fps: Optional[int] = None +) -> None: + from packaging.version import parse + + imageio = util.get_module( + "imageio", + required='wandb.Video requires imageio when passing raw data. Install with "pip install wandb[media]"', + ) + + if parse(imageio.__version__) < parse("2.28.1"): + raise ValueError( + "imageio version 2.28.1 or higher is required to encode gifs. " + "Please upgrade imageio with `pip install imageio>=2.28.1`" + ) + + writer = imageio.save( + filename, + quantizer=0, + palettesize=256, + loop=0, + duration=1000 / clip.fps, + ) + + for frame in clip.iter_frames(fps=fps, dtype="uint8"): + writer.append_data(frame) + + writer.close() + + +class Video(BatchableMedia): + """A class for logging videos to W&B.""" + + _log_type = "video-file" + EXTS = ("gif", "mp4", "webm", "ogg") + _width: Optional[int] + _height: Optional[int] + + def __init__( + self, + data_or_path: Union[str, pathlib.Path, "np.ndarray", "TextIO", "BytesIO"], + caption: Optional[str] = None, + fps: Optional[int] = None, + format: Optional[Literal["gif", "mp4", "webm", "ogg"]] = None, + ): + """Initialize a W&B Video object. + + Args: + data_or_path: + Video can be initialized with a path to a file or an io object. + Video can be initialized with a numpy tensor. + The numpy tensor must be either 4 dimensional or 5 dimensional. + The dimensions should be (number of frames, channel, height, width) or + (batch, number of frames, channel, height, width) + The format parameter must be specified with the format argument + when initializing with a numpy array + or io object. + caption: Caption associated with the video for display. + fps: + The frame rate to use when encoding raw video frames. + Default value is 4. + This parameter has no effect when data_or_path is a string, or bytes. + format: + Format of video, necessary if initializing with a numpy array + or io object. This parameter will be used to determine the format + to use when encoding the video data. Accepted values are "gif", + "mp4", "webm", or "ogg". + If no value is provided, the default format will be "gif". + + Examples: + ### Log a numpy array as a video + ```python + import numpy as np + import wandb + + with wandb.init() as run: + # axes are (number of frames, channel, height, width) + frames = np.random.randint( + low=0, high=256, size=(10, 3, 100, 100), dtype=np.uint8 + ) + run.log({"video": wandb.Video(frames, format="mp4", fps=4)}) + ``` + """ + super().__init__(caption=caption) + + if format is None: + wandb.termwarn( + "`format` argument was not provided, defaulting to `gif`. " + "This parameter will be required in v0.20.0, " + "please specify the format explicitly." + ) + self._format = format or "gif" + self._width = None + self._height = None + self._channels = None + if self._format not in Video.EXTS: + raise ValueError( + "wandb.Video accepts {} formats".format(", ".join(Video.EXTS)) + ) + + if isinstance(data_or_path, (BytesIO, str)) and fps: + msg = ( + "`fps` argument does not affect the frame rate of the video " + "when providing a file path or raw bytes." + ) + wandb.termwarn(msg) + + if isinstance(data_or_path, BytesIO): + filename = os.path.join( + MEDIA_TMP.name, runid.generate_id() + "." + self._format + ) + with open(filename, "wb") as f: + f.write(data_or_path.read()) + self._set_file(filename, is_tmp=True) + elif isinstance(data_or_path, (str, pathlib.Path)): + data_or_path = str(data_or_path) + + _, ext = os.path.splitext(data_or_path) + ext = ext[1:].lower() + if ext not in Video.EXTS: + raise ValueError( + "wandb.Video accepts {} formats".format(", ".join(Video.EXTS)) + ) + self._set_file(data_or_path, is_tmp=False) + # ffprobe -v error -select_streams v:0 -show_entries stream=width,height -of csv=p=0 data_or_path + else: + if hasattr(data_or_path, "numpy"): # TF data eager tensors + self.data = data_or_path.numpy() + elif util.is_numpy_array(data_or_path): + self.data = data_or_path + else: + raise ValueError( + "wandb.Video accepts a file path or numpy like data as input" + ) + fps = fps or 4 + printer_asyncio.run_async_with_spinner( + printer.new_printer(), + "Encoding video...", + functools.partial(self.encode, fps=fps), + ) + + def encode(self, fps: int = 4) -> None: + # import ImageSequenceClip from the appropriate MoviePy module + mpy = util.get_module( + "moviepy.video.io.ImageSequenceClip", + required='wandb.Video requires moviepy when passing raw data. Install with "pip install wandb[media]"', + ) + + tensor = self._prepare_video(self.data) + _, self._height, self._width, self._channels = tensor.shape # type: ignore + + # encode sequence of images into gif string + clip = mpy.ImageSequenceClip(list(tensor), fps=fps) + + filename = os.path.join( + MEDIA_TMP.name, runid.generate_id() + "." + self._format + ) + + if self._format == "gif": + write_gif_with_image_io(clip, filename) + else: + clip.write_videofile(filename, logger=None) + + self._set_file(filename, is_tmp=True) + + @classmethod + def get_media_subdir(cls: Type["Video"]) -> str: + return os.path.join("media", "videos") + + def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict: + json_dict = super().to_json(run_or_artifact) + json_dict["_type"] = self._log_type + + if self._width is not None: + json_dict["width"] = self._width + if self._height is not None: + json_dict["height"] = self._height + + return json_dict + + def _prepare_video(self, video: "np.ndarray") -> "np.ndarray": + """This logic was mostly taken from tensorboardX.""" + np = util.get_module( + "numpy", + required='wandb.Video requires numpy when passing raw data. To get it, run "pip install numpy".', + ) + if video.ndim < 4: + raise ValueError( + "Video must be at least 4 dimensions: time, channels, height, width" + ) + if video.ndim == 4: + video = video.reshape(1, *video.shape) + b, t, c, h, w = video.shape + + if video.dtype != np.uint8: + logging.warning("Converting video data to uint8") + video = video.astype(np.uint8) + + def is_power2(num: int) -> bool: + return num != 0 and ((num & (num - 1)) == 0) + + # pad to nearest power of 2, all at once + if not is_power2(video.shape[0]): + len_addition = int(2 ** video.shape[0].bit_length() - video.shape[0]) + video = np.concatenate( + (video, np.zeros(shape=(len_addition, t, c, h, w))), axis=0 + ) + + n_rows = 2 ** ((b.bit_length() - 1) // 2) + n_cols = video.shape[0] // n_rows + + video = video.reshape(n_rows, n_cols, t, c, h, w) + video = np.transpose(video, axes=(2, 0, 4, 1, 5, 3)) + video = video.reshape(t, n_rows * h, n_cols * w, c) + return video + + @classmethod + def seq_to_json( + cls: Type["Video"], + seq: Sequence["BatchableMedia"], + run: "LocalRun", + key: str, + step: Union[int, str], + ) -> dict: + base_path = os.path.join(run.dir, cls.get_media_subdir()) + filesystem.mkdir_exists_ok(base_path) + + meta = { + "_type": "videos", + "count": len(seq), + "videos": [v.to_json(run) for v in seq], + "captions": Video.captions(seq), + } + return meta + + +class _VideoFileType(_dtypes.Type): + name = "video-file" + types = [Video] + + +_dtypes.TypeRegistry.add(_VideoFileType) diff --git a/lib/python3.12/site-packages/wandb/sdk/integration_utils/__init__.py b/lib/python3.12/site-packages/wandb/sdk/integration_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/integration_utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/integration_utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a2bcc16afebfbda6985cfab359df904bbc60ab3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/integration_utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/integration_utils/__pycache__/auto_logging.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/integration_utils/__pycache__/auto_logging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..264c39762ee24039270272d604ddd3dd91c70c1e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/integration_utils/__pycache__/auto_logging.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/integration_utils/__pycache__/data_logging.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/integration_utils/__pycache__/data_logging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c70a973cab9efba982de0c1e287c22812e83e52e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/integration_utils/__pycache__/data_logging.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/integration_utils/auto_logging.py b/lib/python3.12/site-packages/wandb/sdk/integration_utils/auto_logging.py new file mode 100644 index 0000000000000000000000000000000000000000..0aa0a3842944f65f59b989671636af8c59c456e9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/integration_utils/auto_logging.py @@ -0,0 +1,232 @@ +import asyncio +import functools +import inspect +import logging +from typing import Any, Dict, Optional, Protocol, Sequence, TypeVar + +import wandb.sdk +import wandb.util +from wandb.sdk.lib import telemetry as wb_telemetry +from wandb.sdk.lib.timer import Timer + +logger = logging.getLogger(__name__) + + +AutologInitArgs = Optional[Dict[str, Any]] + + +K = TypeVar("K", bound=str) +V = TypeVar("V") + + +class Response(Protocol[K, V]): + def __getitem__(self, key: K) -> V: ... # pragma: no cover + + def get( + self, key: K, default: Optional[V] = None + ) -> Optional[V]: ... # pragma: no cover + + +class ArgumentResponseResolver(Protocol): + def __call__( + self, + args: Sequence[Any], + kwargs: Dict[str, Any], + response: Response, + start_time: float, + time_elapsed: float, + ) -> Optional[Dict[str, Any]]: ... # pragma: no cover + + +class PatchAPI: + def __init__( + self, + name: str, + symbols: Sequence[str], + resolver: ArgumentResponseResolver, + ) -> None: + """Patches the API to log wandb Media or metrics.""" + # name of the LLM provider, e.g. "Cohere" or "OpenAI" or package name like "Transformers" + self.name = name + # api library name, e.g. "cohere" or "openai" or "transformers" + self._api = None + # dictionary of original methods + self.original_methods: Dict[str, Any] = {} + # list of symbols to patch, e.g. ["Client.generate", "Edit.create"] or ["Pipeline.__call__"] + self.symbols = symbols + # resolver callable to convert args/response into a dictionary of wandb media objects or metrics + self.resolver = resolver + + @property + def set_api(self) -> Any: + """Returns the API module.""" + lib_name = self.name.lower() + if self._api is None: + self._api = wandb.util.get_module( + name=lib_name, + required=f"To use the W&B {self.name} Autolog, " + f"you need to have the `{lib_name}` python " + f"package installed. Please install it with `pip install {lib_name}`.", + lazy=False, + ) + return self._api + + def patch(self, run: "wandb.sdk.wandb_run.Run") -> None: + """Patches the API to log media or metrics to W&B.""" + for symbol in self.symbols: + # split on dots, e.g. "Client.generate" -> ["Client", "generate"] + symbol_parts = symbol.split(".") + # and get the attribute from the module + original = functools.reduce(getattr, symbol_parts, self.set_api) + + def method_factory(original_method: Any): + async def async_method(*args, **kwargs): + future = asyncio.Future() + + async def callback(coro): + try: + result = await coro + loggable_dict = self.resolver( + args, kwargs, result, timer.start_time, timer.elapsed + ) + if loggable_dict is not None: + run.log(loggable_dict) + future.set_result(result) + except Exception as e: + logger.warning(e) + + with Timer() as timer: + coro = original_method(*args, **kwargs) + asyncio.ensure_future(callback(coro)) + + return await future + + def sync_method(*args, **kwargs): + with Timer() as timer: + result = original_method(*args, **kwargs) + try: + loggable_dict = self.resolver( + args, kwargs, result, timer.start_time, timer.elapsed + ) + if loggable_dict is not None: + run.log(loggable_dict) + except Exception as e: + logger.warning(e) + return result + + if inspect.iscoroutinefunction(original_method): + return functools.wraps(original_method)(async_method) + else: + return functools.wraps(original_method)(sync_method) + + # save original method + self.original_methods[symbol] = original + # monkey patch the method + if len(symbol_parts) == 1: + setattr(self.set_api, symbol_parts[0], method_factory(original)) + else: + setattr( + functools.reduce(getattr, symbol_parts[:-1], self.set_api), + symbol_parts[-1], + method_factory(original), + ) + + def unpatch(self) -> None: + """Unpatches the API.""" + for symbol, original in self.original_methods.items(): + # split on dots, e.g. "Client.generate" -> ["Client", "generate"] + symbol_parts = symbol.split(".") + # unpatch the method + if len(symbol_parts) == 1: + setattr(self.set_api, symbol_parts[0], original) + else: + setattr( + functools.reduce(getattr, symbol_parts[:-1], self.set_api), + symbol_parts[-1], + original, + ) + + +class AutologAPI: + def __init__( + self, + name: str, + symbols: Sequence[str], + resolver: ArgumentResponseResolver, + telemetry_feature: Optional[str] = None, + ) -> None: + """Autolog API calls to W&B.""" + self._telemetry_feature = telemetry_feature + self._patch_api = PatchAPI( + name=name, + symbols=symbols, + resolver=resolver, + ) + self._name = self._patch_api.name + self._run: Optional[wandb.sdk.wandb_run.Run] = None + self.__run_created_by_autolog: bool = False + + @property + def _is_enabled(self) -> bool: + """Returns whether autologging is enabled.""" + return self._run is not None + + def __call__(self, init: AutologInitArgs = None) -> None: + """Enable autologging.""" + self.enable(init=init) + + def _run_init(self, init: AutologInitArgs = None) -> None: + """Handle wandb run initialization.""" + # - autolog(init: dict = {...}) calls wandb.init(**{...}) + # regardless of whether there is a wandb.run or not, + # we only track if the run was created by autolog + # - todo: autolog(init: dict | run = run) would use the user-provided run + # - autolog() uses the wandb.run if there is one, otherwise it calls wandb.init() + if init: + _wandb_run = wandb.run + # we delegate dealing with the init dict to wandb.init() + self._run = wandb.init(**init) + if _wandb_run != self._run: + self.__run_created_by_autolog = True + elif wandb.run is None: + self._run = wandb.init() + self.__run_created_by_autolog = True + else: + self._run = wandb.run + + def enable(self, init: AutologInitArgs = None) -> None: + """Enable autologging. + + Args: + init: Optional dictionary of arguments to pass to wandb.init(). + + """ + if self._is_enabled: + logger.info( + f"{self._name} autologging is already enabled, disabling and re-enabling." + ) + self.disable() + + logger.info(f"Enabling {self._name} autologging.") + self._run_init(init=init) + + self._patch_api.patch(self._run) + + if self._telemetry_feature: + with wb_telemetry.context(self._run) as tel: + setattr(tel.feature, self._telemetry_feature, True) + + def disable(self) -> None: + """Disable autologging.""" + if self._run is None: + return + + logger.info(f"Disabling {self._name} autologging.") + + if self.__run_created_by_autolog: + self._run.finish() + self.__run_created_by_autolog = False + + self._run = None + + self._patch_api.unpatch() diff --git a/lib/python3.12/site-packages/wandb/sdk/integration_utils/data_logging.py b/lib/python3.12/site-packages/wandb/sdk/integration_utils/data_logging.py new file mode 100644 index 0000000000000000000000000000000000000000..b0b971cc17211228132698f4279b7b7db1a81bed --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/integration_utils/data_logging.py @@ -0,0 +1,475 @@ +# wandb.integrations.data_logging.py +# +# Contains common utility functions that enable +# logging datasets and predictions to wandb. +import sys +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union + +import wandb + +if TYPE_CHECKING: + from wandb.data_types import _TableIndex + +CAN_INFER_IMAGE_AND_VIDEO = sys.version_info.major == 3 and sys.version_info.minor >= 5 + + +class ValidationDataLogger: + """Logs validation data as a wandb.Table. + + ValidationDataLogger is intended to be used inside of library integrations + in order to facilitate the process of optionally building a validation dataset + and logging periodic predictions against such validation data using WandB best + practices. + """ + + validation_inputs: Union[Sequence, Dict[str, Sequence]] + validation_targets: Optional[Union[Sequence, Dict[str, Sequence]]] + validation_indexes: List["_TableIndex"] + prediction_row_processor: Optional[Callable] + class_labels_table: Optional["wandb.Table"] + infer_missing_processors: bool + + def __init__( + self, + inputs: Union[Sequence, Dict[str, Sequence]], + targets: Optional[Union[Sequence, Dict[str, Sequence]]] = None, + indexes: Optional[List["_TableIndex"]] = None, + validation_row_processor: Optional[Callable] = None, + prediction_row_processor: Optional[Callable] = None, + input_col_name: str = "input", + target_col_name: str = "target", + table_name: str = "wb_validation_data", + artifact_type: str = "validation_dataset", + class_labels: Optional[List[str]] = None, + infer_missing_processors: bool = True, + ) -> None: + """Initialize a new ValidationDataLogger. + + Args: + inputs: A list of input vectors or dictionary of lists of input vectors + (used if the model has multiple named inputs) + targets: A list of target vectors or dictionary of lists of target vectors + (used if the model has multiple named targets/putputs). Defaults to `None`. + `targets` and `indexes` cannot both be `None`. + indexes: An ordered list of `wandb.data_types._TableIndex` mapping the + input items to their source table. This is most commonly retrieved by using + `indexes = my_data_table.get_index()`. Defaults to `None`. `targets` + and `indexes` cannot both be `None`. + validation_row_processor: A function to apply to the validation data, + commonly used to visualize the data. The function will receive an `ndx` (`int`) + and a `row` (`dict`). If `inputs` is a list, then `row["input"]` will be the input + data for the row. Else, it will be keyed based on the name of the input slot + (corresponding to `inputs`). If `targets` is a list, then + `row["target"]` will be the target data for the row. Else, it will + be keyed based on `targets`. For example, if your input data is a + single ndarray, but you wish to visualize the data as an image, + then you can provide `lambda ndx, row: {"img": wandb.Image(row["input"])}` + as the processor. If `None`, we will try to guess the appropriate processor. + Ignored if `log_evaluation` is `False` or `val_keys` are present. Defaults to `None`. + prediction_row_processor: Same as validation_row_processor, but applied to the + model's output. `row["output"]` will contain the results of the model output. + Defaults to `None`. + input_col_name: The name to use for the input column. + Defaults to `"input"`. + target_col_name: The name to use for the target column. + Defaults to `"target"`. + table_name: The name to use for the validation table. + Defaults to `"wb_validation_data"`. + artifact_type: The artifact type to use for the validation data. + Defaults to `"validation_dataset"`. + class_labels: Optional list of labels to use in the inferred + processors. If the model's `target` or `output` is inferred to be a class, + we will attempt to map the class to these labels. Defaults to `None`. + infer_missing_processors: Determines if processors are inferred if + they are missing. Defaults to True. + """ + class_labels_table: Optional[wandb.Table] + if isinstance(class_labels, list) and len(class_labels) > 0: + class_labels_table = wandb.Table( + columns=["label"], data=[[label] for label in class_labels] + ) + else: + class_labels_table = None + + if indexes is None: + assert targets is not None + local_validation_table = wandb.Table(columns=[], data=[]) + + if isinstance(targets, dict): + for col_name in targets: + local_validation_table.add_column(col_name, targets[col_name]) + else: + local_validation_table.add_column(target_col_name, targets) + + if isinstance(inputs, dict): + for col_name in inputs: + local_validation_table.add_column(col_name, inputs[col_name]) + else: + local_validation_table.add_column(input_col_name, inputs) + + if validation_row_processor is None and infer_missing_processors: + example_input = _make_example(inputs) + example_target = _make_example(targets) + if example_input is not None and example_target is not None: + validation_row_processor = _infer_validation_row_processor( + example_input, + example_target, + class_labels_table, + input_col_name, + target_col_name, + ) + + if validation_row_processor is not None: + local_validation_table.add_computed_columns(validation_row_processor) + + local_validation_artifact = wandb.Artifact(table_name, artifact_type) + local_validation_artifact.add(local_validation_table, "validation_data") + if wandb.run: + wandb.run.use_artifact(local_validation_artifact) + indexes = local_validation_table.get_index() + else: + local_validation_artifact = None + + self.class_labels_table = class_labels_table + self.validation_inputs = inputs + self.validation_targets = targets + self.validation_indexes = indexes + self.prediction_row_processor = prediction_row_processor + self.infer_missing_processors = infer_missing_processors + self.local_validation_artifact = local_validation_artifact + self.input_col_name = input_col_name + + def make_predictions( + self, predict_fn: Callable + ) -> Union[Sequence, Dict[str, Sequence]]: + """Produce predictions by passing `validation_inputs` to `predict_fn`. + + Args: + predict_fn (Callable): Any function which can accept `validation_inputs` and produce + a list of vectors or dictionary of lists of vectors + + Returns: + (Sequence | Dict[str, Sequence]): The returned value of predict_fn + """ + return predict_fn(self.validation_inputs) + + def log_predictions( + self, + predictions: Union[Sequence, Dict[str, Sequence]], + prediction_col_name: str = "output", + val_ndx_col_name: str = "val_row", + table_name: str = "validation_predictions", + commit: bool = True, + ) -> wandb.data_types.Table: + """Log a set of predictions. + + Intended usage: + + vl.log_predictions(vl.make_predictions(self.model.predict)) + + Args: + predictions (Sequence | Dict[str, Sequence]): A list of prediction vectors or dictionary + of lists of prediction vectors + prediction_col_name (str, optional): the name of the prediction column. Defaults to "output". + val_ndx_col_name (str, optional): The name of the column linking prediction table + to the validation ata table. Defaults to "val_row". + table_name (str, optional): name of the prediction table. Defaults to "validation_predictions". + commit (bool, optional): determines if commit should be called on the logged data. Defaults to False. + """ + pred_table = wandb.Table(columns=[], data=[]) + if isinstance(predictions, dict): + for col_name in predictions: + pred_table.add_column(col_name, predictions[col_name]) + else: + pred_table.add_column(prediction_col_name, predictions) + pred_table.add_column(val_ndx_col_name, self.validation_indexes) + + if self.prediction_row_processor is None and self.infer_missing_processors: + example_prediction = _make_example(predictions) + example_input = _make_example(self.validation_inputs) + if example_prediction is not None and example_input is not None: + self.prediction_row_processor = _infer_prediction_row_processor( + example_prediction, + example_input, + self.class_labels_table, + self.input_col_name, + prediction_col_name, + ) + + if self.prediction_row_processor is not None: + pred_table.add_computed_columns(self.prediction_row_processor) + + wandb.log({table_name: pred_table}, commit=commit) + return pred_table + + +def _make_example(data: Any) -> Optional[Union[Dict, Sequence, Any]]: + """Used to make an example input, target, or output.""" + example: Optional[Union[Dict, Sequence, Any]] + + if isinstance(data, dict): + example = {} + for key in data: + example[key] = data[key][0] + elif hasattr(data, "__len__"): + example = data[0] + else: + example = None + + return example + + +def _get_example_shape(example: Union[Sequence, Any]): + """Get the shape of an object if applicable.""" + shape = [] + if not isinstance(example, str) and hasattr(example, "__len__"): + length = len(example) + shape = [length] + if length > 0: + shape += _get_example_shape(example[0]) + return shape + + +def _bind(lambda_fn: Callable, **closure_kwargs: Any) -> Callable: + """Create a closure around a lambda function by binding `closure_kwargs` to the function.""" + + def closure(*args: Any, **kwargs: Any) -> Any: + _k = {} + _k.update(kwargs) + _k.update(closure_kwargs) + return lambda_fn(*args, **_k) + + return closure + + +def _infer_single_example_keyed_processor( + example: Union[Sequence, Any], + class_labels_table: Optional["wandb.Table"] = None, + possible_base_example: Optional[Union[Sequence, Any]] = None, +) -> Dict[str, Callable]: + """Infers a processor from a single example. + + Infers a processor from a single example with optional class_labels_table + and base_example. Base example is useful for cases such as segmentation masks + """ + shape = _get_example_shape(example) + processors: Dict[str, Callable] = {} + if ( + class_labels_table is not None + and len(shape) == 1 + and shape[0] == len(class_labels_table.data) + ): + np = wandb.util.get_module( + "numpy", + required="Inferring processors require numpy", + ) + # Assume these are logits + class_names = class_labels_table.get_column("label") + + processors["max_class"] = lambda n, d, p: class_labels_table.index_ref( # type: ignore + np.argmax(d) + ) + # TODO: Consider adding back if users ask + # processors["min_class"] = lambda n, d, p: class_labels_table.index_ref( # type: ignore + # np.argmin(d) + # ) + + values = np.unique(example) + is_one_hot = len(values) == 2 and set(values) == {0, 1} + if not is_one_hot: + processors["score"] = lambda n, d, p: { + class_names[i]: d[i] for i in range(shape[0]) + } + elif ( + len(shape) == 1 + and shape[0] == 1 + and ( + isinstance(example[0], int) + or (hasattr(example, "tolist") and isinstance(example.tolist()[0], int)) # type: ignore + ) + ): + # assume this is a class + if class_labels_table is not None: + processors["class"] = ( + lambda n, d, p: class_labels_table.index_ref(d[0]) + if d[0] < len(class_labels_table.data) + else d[0] + ) # type: ignore + else: + processors["val"] = lambda n, d, p: d[0] + elif len(shape) == 1: + np = wandb.util.get_module( + "numpy", + required="Inferring processors require numpy", + ) + # This could be anything + if shape[0] <= 10: + # if less than 10, fan out the results + # processors["node"] = lambda n, d, p: {i: d[i] for i in range(shape[0])} + processors["node"] = lambda n, d, p: [ + d[i].tolist() if hasattr(d[i], "tolist") else d[i] + for i in range(shape[0]) + ] + # just report the argmax and argmin + processors["argmax"] = lambda n, d, p: np.argmax(d) + + values = np.unique(example) + is_one_hot = len(values) == 2 and set(values) == {0, 1} + if not is_one_hot: + processors["argmin"] = lambda n, d, p: np.argmin(d) + elif len(shape) == 2 and CAN_INFER_IMAGE_AND_VIDEO: + if ( + class_labels_table is not None + and possible_base_example is not None + and shape == _get_example_shape(possible_base_example) + ): + # consider this a segmentation mask + processors["image"] = lambda n, d, p: wandb.Image( + p, + masks={ + "masks": { + "mask_data": d, + "class_labels": class_labels_table.get_column("label"), # type: ignore + } + }, + ) + else: + # consider this a 2d image + processors["image"] = lambda n, d, p: wandb.Image(d) + elif len(shape) == 3 and CAN_INFER_IMAGE_AND_VIDEO: + # consider this an image + processors["image"] = lambda n, d, p: wandb.Image(d) + elif len(shape) == 4 and CAN_INFER_IMAGE_AND_VIDEO: + # consider this a video + processors["video"] = lambda n, d, p: wandb.Video(d) + + return processors + + +def _infer_validation_row_processor( + example_input: Union[Dict, Sequence], + example_target: Union[Dict, Sequence, Any], + class_labels_table: Optional["wandb.Table"] = None, + input_col_name: str = "input", + target_col_name: str = "target", +) -> Callable: + """Infers the composite processor for the validation data.""" + single_processors = {} + if isinstance(example_input, dict): + for key in example_input: + key_processors = _infer_single_example_keyed_processor(example_input[key]) + for p_key in key_processors: + single_processors[f"{key}:{p_key}"] = _bind( + lambda ndx, row, key_processor, key: key_processor( + ndx, + row[key], + None, + ), + key_processor=key_processors[p_key], + key=key, + ) + else: + key = input_col_name + key_processors = _infer_single_example_keyed_processor(example_input) + for p_key in key_processors: + single_processors[f"{key}:{p_key}"] = _bind( + lambda ndx, row, key_processor, key: key_processor( + ndx, + row[key], + None, + ), + key_processor=key_processors[p_key], + key=key, + ) + + if isinstance(example_target, dict): + for key in example_target: + key_processors = _infer_single_example_keyed_processor( + example_target[key], class_labels_table + ) + for p_key in key_processors: + single_processors[f"{key}:{p_key}"] = _bind( + lambda ndx, row, key_processor, key: key_processor( + ndx, + row[key], + None, + ), + key_processor=key_processors[p_key], + key=key, + ) + else: + key = target_col_name + key_processors = _infer_single_example_keyed_processor( + example_target, + class_labels_table, + example_input if not isinstance(example_input, dict) else None, + ) + for p_key in key_processors: + single_processors[f"{key}:{p_key}"] = _bind( + lambda ndx, row, key_processor, key: key_processor( + ndx, + row[key], + row[input_col_name] + if not isinstance(example_input, dict) + else None, + ), + key_processor=key_processors[p_key], + key=key, + ) + + def processor(ndx, row): + return {key: single_processors[key](ndx, row) for key in single_processors} + + return processor + + +def _infer_prediction_row_processor( + example_prediction: Union[Dict, Sequence], + example_input: Union[Dict, Sequence], + class_labels_table: Optional["wandb.Table"] = None, + input_col_name: str = "input", + output_col_name: str = "output", +) -> Callable: + """Infers the composite processor for the prediction output data.""" + single_processors = {} + + if isinstance(example_prediction, dict): + for key in example_prediction: + key_processors = _infer_single_example_keyed_processor( + example_prediction[key], class_labels_table + ) + for p_key in key_processors: + single_processors[f"{key}:{p_key}"] = _bind( + lambda ndx, row, key_processor, key: key_processor( + ndx, + row[key], + None, + ), + key_processor=key_processors[p_key], + key=key, + ) + else: + key = output_col_name + key_processors = _infer_single_example_keyed_processor( + example_prediction, + class_labels_table, + example_input if not isinstance(example_input, dict) else None, + ) + for p_key in key_processors: + single_processors[f"{key}:{p_key}"] = _bind( + lambda ndx, row, key_processor, key: key_processor( + ndx, + row[key], + ndx.get_row().get("val_row").get_row().get(input_col_name) + if not isinstance(example_input, dict) + else None, + ), + key_processor=key_processors[p_key], + key=key, + ) + + def processor(ndx, row): + return {key: single_processors[key](ndx, row) for key in single_processors} + + return processor diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__init__.py b/lib/python3.12/site-packages/wandb/sdk/interface/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5491d3a7d7d28a469ec137e02573aa5b4d32db94 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7d9fbd7ef178d7f573c3f393035bd918b43dec6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f01a04576b79adbc23aeaa5fccaa6ed6eeff817 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_queue.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_queue.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f03fc174e74b831749dd7c3e4d4553a5faf4c103 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_queue.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_relay.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_relay.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6bf991d8e07af85bf62aacdba1885430f6f7521f Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_relay.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_shared.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_shared.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08845fce6057fc9e4c8c20cfcca3ed94fe097bd5 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_shared.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_sock.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_sock.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3cb8c9c7ffecacfc2416b7d5f939935d0cd71344 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/interface_sock.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dadeea516d53a681ed9200a2d127503ed4cccbe9 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router_queue.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router_queue.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f687bccc973ed9f2f1ec46ba4d301fa4a6fce88 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router_queue.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router_relay.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router_relay.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ef06d2bd39b7b767d56a9df0e3ce4ce04fa790c9 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router_relay.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router_sock.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router_sock.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f727155574f4a9cda0857b68aa2144c6e0b6c67 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/router_sock.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/summary_record.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/summary_record.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..af335c5c9cfc306d3f60ec6ec8ca3e21886ff29d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/interface/__pycache__/summary_record.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/constants.py b/lib/python3.12/site-packages/wandb/sdk/interface/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..09fe1e26f0ce72e23df8b627ff15e332f6ebc4e6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/constants.py @@ -0,0 +1,4 @@ +# +NOTIFY_PROCESS = 1 +NOTIFY_SHUTDOWN = 2 +NOTIFY_REQUEST = 3 diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/interface.py b/lib/python3.12/site-packages/wandb/sdk/interface/interface.py new file mode 100644 index 0000000000000000000000000000000000000000..89b4a3fee510fd58f4b383097ae6f2d53d5f177a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/interface.py @@ -0,0 +1,1056 @@ +"""Interface base class - Used to send messages to the internal process. + +InterfaceBase: The abstract class +InterfaceShared: Common routines for socket and queue based implementations +InterfaceQueue: Use multiprocessing queues to send and receive messages +InterfaceSock: Use socket to send and receive messages +""" + +import gzip +import logging +import time +from abc import abstractmethod +from pathlib import Path +from secrets import token_hex +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterable, + List, + Literal, + NewType, + Optional, + Tuple, + TypedDict, + Union, +) + +from wandb import termwarn +from wandb.proto import wandb_internal_pb2 as pb +from wandb.proto import wandb_telemetry_pb2 as tpb +from wandb.sdk.artifacts.artifact import Artifact +from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest +from wandb.sdk.artifacts.staging import get_staging_dir +from wandb.sdk.lib import json_util as json +from wandb.sdk.mailbox import HandleAbandonedError, MailboxHandle +from wandb.util import ( + WandBJSONEncoderOld, + get_h5_typename, + json_dumps_safer, + json_dumps_safer_history, + json_friendly, + json_friendly_val, + maybe_compress_summary, +) + +from ..data_types.utils import history_dict_to_json, val_to_json +from . import summary_record as sr + +MANIFEST_FILE_SIZE_THRESHOLD = 100_000 + +GlobStr = NewType("GlobStr", str) + + +PolicyName = Literal["now", "live", "end"] + + +class FilesDict(TypedDict): + files: Iterable[Tuple[GlobStr, PolicyName]] + + +if TYPE_CHECKING: + from ..wandb_run import Run + + +logger = logging.getLogger("wandb") + + +def file_policy_to_enum(policy: "PolicyName") -> "pb.FilesItem.PolicyType.V": + if policy == "now": + enum = pb.FilesItem.PolicyType.NOW + elif policy == "end": + enum = pb.FilesItem.PolicyType.END + elif policy == "live": + enum = pb.FilesItem.PolicyType.LIVE + return enum + + +def file_enum_to_policy(enum: "pb.FilesItem.PolicyType.V") -> "PolicyName": + if enum == pb.FilesItem.PolicyType.NOW: + policy: PolicyName = "now" + elif enum == pb.FilesItem.PolicyType.END: + policy = "end" + elif enum == pb.FilesItem.PolicyType.LIVE: + policy = "live" + return policy + + +class InterfaceBase: + _drop: bool + + def __init__(self) -> None: + self._drop = False + + def publish_header(self) -> None: + header = pb.HeaderRecord() + self._publish_header(header) + + @abstractmethod + def _publish_header(self, header: pb.HeaderRecord) -> None: + raise NotImplementedError + + def deliver_status(self) -> MailboxHandle[pb.Result]: + return self._deliver_status(pb.StatusRequest()) + + @abstractmethod + def _deliver_status( + self, + status: pb.StatusRequest, + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def _make_config( + self, + data: Optional[dict] = None, + key: Optional[Union[Tuple[str, ...], str]] = None, + val: Optional[Any] = None, + obj: Optional[pb.ConfigRecord] = None, + ) -> pb.ConfigRecord: + config = obj or pb.ConfigRecord() + if data: + for k, v in data.items(): + update = config.update.add() + update.key = k + update.value_json = json_dumps_safer(json_friendly(v)[0]) + if key: + update = config.update.add() + if isinstance(key, tuple): + for k in key: + update.nested_key.append(k) + else: + update.key = key + update.value_json = json_dumps_safer(json_friendly(val)[0]) + return config + + def _make_run(self, run: "Run") -> pb.RunRecord: # noqa: C901 + proto_run = pb.RunRecord() + if run._settings.entity is not None: + proto_run.entity = run._settings.entity + if run._settings.project is not None: + proto_run.project = run._settings.project + if run._settings.run_group is not None: + proto_run.run_group = run._settings.run_group + if run._settings.run_job_type is not None: + proto_run.job_type = run._settings.run_job_type + if run._settings.run_id is not None: + proto_run.run_id = run._settings.run_id + if run._settings.run_name is not None: + proto_run.display_name = run._settings.run_name + if run._settings.run_notes is not None: + proto_run.notes = run._settings.run_notes + if run._settings.run_tags is not None: + proto_run.tags.extend(run._settings.run_tags) + if run._start_time is not None: + proto_run.start_time.FromMicroseconds(int(run._start_time * 1e6)) + if run._starting_step is not None: + proto_run.starting_step = run._starting_step + if run._settings.git_remote_url is not None: + proto_run.git.remote_url = run._settings.git_remote_url + if run._settings.git_commit is not None: + proto_run.git.commit = run._settings.git_commit + if run._settings.sweep_id is not None: + proto_run.sweep_id = run._settings.sweep_id + if run._settings.host: + proto_run.host = run._settings.host + if run._settings.resumed: + proto_run.resumed = run._settings.resumed + if run._settings.fork_from: + run_moment = run._settings.fork_from + proto_run.branch_point.run = run_moment.run + proto_run.branch_point.metric = run_moment.metric + proto_run.branch_point.value = run_moment.value + if run._settings.resume_from: + run_moment = run._settings.resume_from + proto_run.branch_point.run = run_moment.run + proto_run.branch_point.metric = run_moment.metric + proto_run.branch_point.value = run_moment.value + if run._forked: + proto_run.forked = run._forked + if run._config is not None: + config_dict = run._config._as_dict() # type: ignore + self._make_config(data=config_dict, obj=proto_run.config) + if run._telemetry_obj: + proto_run.telemetry.MergeFrom(run._telemetry_obj) + if run._start_runtime: + proto_run.runtime = run._start_runtime + return proto_run + + def publish_run(self, run: "Run") -> None: + run_record = self._make_run(run) + self._publish_run(run_record) + + @abstractmethod + def _publish_run(self, run: pb.RunRecord) -> None: + raise NotImplementedError + + def publish_cancel(self, cancel_slot: str) -> None: + cancel = pb.CancelRequest(cancel_slot=cancel_slot) + self._publish_cancel(cancel) + + @abstractmethod + def _publish_cancel(self, cancel: pb.CancelRequest) -> None: + raise NotImplementedError + + def publish_config( + self, + data: Optional[dict] = None, + key: Optional[Union[Tuple[str, ...], str]] = None, + val: Optional[Any] = None, + ) -> None: + cfg = self._make_config(data=data, key=key, val=val) + + self._publish_config(cfg) + + @abstractmethod + def _publish_config(self, cfg: pb.ConfigRecord) -> None: + raise NotImplementedError + + @abstractmethod + def _publish_metric(self, metric: pb.MetricRecord) -> None: + raise NotImplementedError + + def _make_summary_from_dict(self, summary_dict: dict) -> pb.SummaryRecord: + summary = pb.SummaryRecord() + for k, v in summary_dict.items(): + update = summary.update.add() + update.key = k + update.value_json = json.dumps(v) + return summary + + def _summary_encode( + self, + value: Any, + path_from_root: str, + run: "Run", + ) -> dict: + """Normalize, compress, and encode sub-objects for backend storage. + + value: Object to encode. + path_from_root: `str` dot separated string from the top-level summary to the + current `value`. + + Returns: + A new tree of dict's with large objects replaced with dictionaries + with "_type" entries that say which type the original data was. + """ + # Constructs a new `dict` tree in `json_value` that discards and/or + # encodes objects that aren't JSON serializable. + + if isinstance(value, dict): + json_value = {} + for key, value in value.items(): # noqa: B020 + json_value[key] = self._summary_encode( + value, + path_from_root + "." + key, + run=run, + ) + return json_value + else: + friendly_value, converted = json_friendly( + val_to_json(run, path_from_root, value, namespace="summary") + ) + json_value, compressed = maybe_compress_summary( + friendly_value, get_h5_typename(value) + ) + if compressed: + # TODO(jhr): impleement me + pass + # self.write_h5(path_from_root, friendly_value) + + return json_value + + def _make_summary( + self, + summary_record: sr.SummaryRecord, + run: "Run", + ) -> pb.SummaryRecord: + pb_summary_record = pb.SummaryRecord() + + for item in summary_record.update: + pb_summary_item = pb_summary_record.update.add() + key_length = len(item.key) + + assert key_length > 0 + + if key_length > 1: + pb_summary_item.nested_key.extend(item.key) + else: + pb_summary_item.key = item.key[0] + + path_from_root = ".".join(item.key) + json_value = self._summary_encode( + item.value, + path_from_root, + run=run, + ) + json_value, _ = json_friendly(json_value) # type: ignore + + pb_summary_item.value_json = json.dumps( + json_value, + cls=WandBJSONEncoderOld, + ) + + for item in summary_record.remove: + pb_summary_item = pb_summary_record.remove.add() + key_length = len(item.key) + + assert key_length > 0 + + if key_length > 1: + pb_summary_item.nested_key.extend(item.key) + else: + pb_summary_item.key = item.key[0] + + return pb_summary_record + + def publish_summary( + self, + run: "Run", + summary_record: sr.SummaryRecord, + ) -> None: + pb_summary_record = self._make_summary(summary_record, run=run) + self._publish_summary(pb_summary_record) + + @abstractmethod + def _publish_summary(self, summary: pb.SummaryRecord) -> None: + raise NotImplementedError + + def _make_files(self, files_dict: "FilesDict") -> pb.FilesRecord: + files = pb.FilesRecord() + for path, policy in files_dict["files"]: + f = files.files.add() + f.path = path + f.policy = file_policy_to_enum(policy) + return files + + def publish_files(self, files_dict: "FilesDict") -> None: + files = self._make_files(files_dict) + self._publish_files(files) + + @abstractmethod + def _publish_files(self, files: pb.FilesRecord) -> None: + raise NotImplementedError + + def publish_python_packages(self, working_set) -> None: + python_packages = pb.PythonPackagesRequest() + for pkg in working_set: + python_packages.package.add(name=pkg.key, version=pkg.version) + self._publish_python_packages(python_packages) + + @abstractmethod + def _publish_python_packages( + self, python_packages: pb.PythonPackagesRequest + ) -> None: + raise NotImplementedError + + def _make_artifact(self, artifact: "Artifact") -> pb.ArtifactRecord: + proto_artifact = pb.ArtifactRecord() + proto_artifact.type = artifact.type + proto_artifact.name = artifact.name + proto_artifact.client_id = artifact._client_id + proto_artifact.sequence_client_id = artifact._sequence_client_id + proto_artifact.digest = artifact.digest + if artifact.distributed_id: + proto_artifact.distributed_id = artifact.distributed_id + if artifact.description: + proto_artifact.description = artifact.description + if artifact.metadata: + proto_artifact.metadata = json.dumps(json_friendly_val(artifact.metadata)) + if artifact._base_id: + proto_artifact.base_id = artifact._base_id + + ttl_duration_input = artifact._ttl_duration_seconds_to_gql() + if ttl_duration_input: + proto_artifact.ttl_duration_seconds = ttl_duration_input + proto_artifact.incremental_beta1 = artifact.incremental + self._make_artifact_manifest(artifact.manifest, obj=proto_artifact.manifest) + return proto_artifact + + def _make_artifact_manifest( + self, + artifact_manifest: ArtifactManifest, + obj: Optional[pb.ArtifactManifest] = None, + ) -> pb.ArtifactManifest: + proto_manifest = obj or pb.ArtifactManifest() + proto_manifest.version = artifact_manifest.version() + proto_manifest.storage_policy = artifact_manifest.storage_policy.name() + + # Very large manifests need to be written to file to avoid protobuf size limits. + if len(artifact_manifest) > MANIFEST_FILE_SIZE_THRESHOLD: + path = self._write_artifact_manifest_file(artifact_manifest) + proto_manifest.manifest_file_path = path + return proto_manifest + + for k, v in artifact_manifest.storage_policy.config().items() or {}.items(): + cfg = proto_manifest.storage_policy_config.add() + cfg.key = k + cfg.value_json = json.dumps(v) + + for entry in sorted(artifact_manifest.entries.values(), key=lambda k: k.path): + proto_entry = proto_manifest.contents.add() + proto_entry.path = entry.path + proto_entry.digest = entry.digest + if entry.size: + proto_entry.size = entry.size + if entry.birth_artifact_id: + proto_entry.birth_artifact_id = entry.birth_artifact_id + if entry.ref: + proto_entry.ref = entry.ref + if entry.local_path: + proto_entry.local_path = entry.local_path + proto_entry.skip_cache = entry.skip_cache + for k, v in entry.extra.items(): + proto_extra = proto_entry.extra.add() + proto_extra.key = k + proto_extra.value_json = json.dumps(v) + return proto_manifest + + def _write_artifact_manifest_file(self, manifest: ArtifactManifest) -> str: + manifest_dir = Path(get_staging_dir()) / "artifact_manifests" + manifest_dir.mkdir(parents=True, exist_ok=True) + # It would be simpler to use `manifest.to_json()`, but that gets very slow for + # large manifests since it encodes the whole thing as a single JSON object. + filename = f"{time.time()}_{token_hex(8)}.manifest_contents.jl.gz" + manifest_file_path = manifest_dir / filename + with gzip.open(manifest_file_path, mode="wt", compresslevel=1) as f: + for entry in manifest.entries.values(): + f.write(f"{json.dumps(entry.to_json())}\n") + return str(manifest_file_path) + + def deliver_link_artifact( + self, + artifact: "Artifact", + portfolio_name: str, + aliases: Iterable[str], + entity: Optional[str] = None, + project: Optional[str] = None, + organization: Optional[str] = None, + ) -> MailboxHandle[pb.Result]: + link_artifact = pb.LinkArtifactRequest() + if artifact.is_draft(): + link_artifact.client_id = artifact._client_id + else: + link_artifact.server_id = artifact.id if artifact.id else "" + link_artifact.portfolio_name = portfolio_name + link_artifact.portfolio_entity = entity or "" + link_artifact.portfolio_organization = organization or "" + link_artifact.portfolio_project = project or "" + link_artifact.portfolio_aliases.extend(aliases) + + return self._deliver_link_artifact(link_artifact) + + @abstractmethod + def _deliver_link_artifact( + self, link_artifact: pb.LinkArtifactRequest + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + @staticmethod + def _make_partial_source_str( + source: Any, job_info: Dict[str, Any], metadata: Dict[str, Any] + ) -> str: + """Construct use_artifact.partial.source_info.source as str.""" + source_type = job_info.get("source_type", "").strip() + if source_type == "artifact": + info_source = job_info.get("source", {}) + source.artifact.artifact = info_source.get("artifact", "") + source.artifact.entrypoint.extend(info_source.get("entrypoint", [])) + source.artifact.notebook = info_source.get("notebook", False) + build_context = info_source.get("build_context") + if build_context: + source.artifact.build_context = build_context + dockerfile = info_source.get("dockerfile") + if dockerfile: + source.artifact.dockerfile = dockerfile + elif source_type == "repo": + source.git.git_info.remote = metadata.get("git", {}).get("remote", "") + source.git.git_info.commit = metadata.get("git", {}).get("commit", "") + source.git.entrypoint.extend(metadata.get("entrypoint", [])) + source.git.notebook = metadata.get("notebook", False) + build_context = metadata.get("build_context") + if build_context: + source.git.build_context = build_context + dockerfile = metadata.get("dockerfile") + if dockerfile: + source.git.dockerfile = dockerfile + elif source_type == "image": + source.image.image = metadata.get("docker", "") + else: + raise ValueError("Invalid source type") + + source_str: str = source.SerializeToString() + return source_str + + def _make_proto_use_artifact( + self, + use_artifact: pb.UseArtifactRecord, + job_name: str, + job_info: Dict[str, Any], + metadata: Dict[str, Any], + ) -> pb.UseArtifactRecord: + use_artifact.partial.job_name = job_name + use_artifact.partial.source_info._version = job_info.get("_version", "") + use_artifact.partial.source_info.source_type = job_info.get("source_type", "") + use_artifact.partial.source_info.runtime = job_info.get("runtime", "") + + src_str = self._make_partial_source_str( + source=use_artifact.partial.source_info.source, + job_info=job_info, + metadata=metadata, + ) + use_artifact.partial.source_info.source.ParseFromString(src_str) # type: ignore[arg-type] + + return use_artifact + + def publish_use_artifact( + self, + artifact: "Artifact", + ) -> None: + assert artifact.id is not None, "Artifact must have an id" + + use_artifact = pb.UseArtifactRecord( + id=artifact.id, + type=artifact.type, + name=artifact.name, + ) + + # TODO(gst): move to internal process + if "_partial" in artifact.metadata: + # Download source info from logged partial job artifact + job_info = {} + try: + path = artifact.get_entry("wandb-job.json").download() + with open(path) as f: + job_info = json.load(f) + + except Exception as e: + logger.warning( + f"Failed to download partial job info from artifact {artifact}, : {e}" + ) + termwarn( + f"Failed to download partial job info from artifact {artifact}, : {e}" + ) + return + + try: + use_artifact = self._make_proto_use_artifact( + use_artifact=use_artifact, + job_name=artifact.name, + job_info=job_info, + metadata=artifact.metadata, + ) + except Exception as e: + logger.warning(f"Failed to construct use artifact proto: {e}") + termwarn(f"Failed to construct use artifact proto: {e}") + return + + self._publish_use_artifact(use_artifact) + + @abstractmethod + def _publish_use_artifact(self, proto_artifact: pb.UseArtifactRecord) -> None: + raise NotImplementedError + + def deliver_artifact( + self, + run: "Run", + artifact: "Artifact", + aliases: Iterable[str], + tags: Optional[Iterable[str]] = None, + history_step: Optional[int] = None, + is_user_created: bool = False, + use_after_commit: bool = False, + finalize: bool = True, + ) -> MailboxHandle[pb.Result]: + proto_run = self._make_run(run) + proto_artifact = self._make_artifact(artifact) + proto_artifact.run_id = proto_run.run_id + proto_artifact.project = proto_run.project + proto_artifact.entity = proto_run.entity + proto_artifact.user_created = is_user_created + proto_artifact.use_after_commit = use_after_commit + proto_artifact.finalize = finalize + + proto_artifact.aliases.extend(aliases or []) + proto_artifact.tags.extend(tags or []) + + log_artifact = pb.LogArtifactRequest() + log_artifact.artifact.CopyFrom(proto_artifact) + if history_step is not None: + log_artifact.history_step = history_step + log_artifact.staging_dir = get_staging_dir() + resp = self._deliver_artifact(log_artifact) + return resp + + @abstractmethod + def _deliver_artifact( + self, + log_artifact: pb.LogArtifactRequest, + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_download_artifact( + self, + artifact_id: str, + download_root: str, + allow_missing_references: bool, + skip_cache: bool, + path_prefix: Optional[str], + ) -> MailboxHandle[pb.Result]: + download_artifact = pb.DownloadArtifactRequest() + download_artifact.artifact_id = artifact_id + download_artifact.download_root = download_root + download_artifact.allow_missing_references = allow_missing_references + download_artifact.skip_cache = skip_cache + download_artifact.path_prefix = path_prefix or "" + resp = self._deliver_download_artifact(download_artifact) + return resp + + @abstractmethod + def _deliver_download_artifact( + self, download_artifact: pb.DownloadArtifactRequest + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def publish_artifact( + self, + run: "Run", + artifact: "Artifact", + aliases: Iterable[str], + tags: Optional[Iterable[str]] = None, + is_user_created: bool = False, + use_after_commit: bool = False, + finalize: bool = True, + ) -> None: + proto_run = self._make_run(run) + proto_artifact = self._make_artifact(artifact) + proto_artifact.run_id = proto_run.run_id + proto_artifact.project = proto_run.project + proto_artifact.entity = proto_run.entity + proto_artifact.user_created = is_user_created + proto_artifact.use_after_commit = use_after_commit + proto_artifact.finalize = finalize + proto_artifact.aliases.extend(aliases or []) + proto_artifact.tags.extend(tags or []) + self._publish_artifact(proto_artifact) + + @abstractmethod + def _publish_artifact(self, proto_artifact: pb.ArtifactRecord) -> None: + raise NotImplementedError + + def publish_tbdata(self, log_dir: str, save: bool, root_logdir: str = "") -> None: + tbrecord = pb.TBRecord() + tbrecord.log_dir = log_dir + tbrecord.save = save + tbrecord.root_dir = root_logdir + self._publish_tbdata(tbrecord) + + @abstractmethod + def _publish_tbdata(self, tbrecord: pb.TBRecord) -> None: + raise NotImplementedError + + @abstractmethod + def _publish_telemetry(self, telem: tpb.TelemetryRecord) -> None: + raise NotImplementedError + + def publish_environment(self, environment: pb.EnvironmentRecord) -> None: + self._publish_environment(environment) + + @abstractmethod + def _publish_environment(self, environment: pb.EnvironmentRecord) -> None: + raise NotImplementedError + + def publish_partial_history( + self, + run: "Run", + data: dict, + user_step: int, + step: Optional[int] = None, + flush: Optional[bool] = None, + publish_step: bool = True, + ) -> None: + data = history_dict_to_json(run, data, step=user_step, ignore_copy_err=True) + data.pop("_step", None) + + # add timestamp to the history request, if not already present + # the timestamp might come from the tensorboard log logic + if "_timestamp" not in data: + data["_timestamp"] = time.time() + + partial_history = pb.PartialHistoryRequest() + for k, v in data.items(): + item = partial_history.item.add() + item.key = k + item.value_json = json_dumps_safer_history(v) + + if publish_step and step is not None: + partial_history.step.num = step + if flush is not None: + partial_history.action.flush = flush + self._publish_partial_history(partial_history) + + @abstractmethod + def _publish_partial_history(self, history: pb.PartialHistoryRequest) -> None: + raise NotImplementedError + + def publish_history( + self, + run: "Run", + data: dict, + step: Optional[int] = None, + publish_step: bool = True, + ) -> None: + data = history_dict_to_json(run, data, step=step) + history = pb.HistoryRecord() + if publish_step: + assert step is not None + history.step.num = step + data.pop("_step", None) + for k, v in data.items(): + item = history.item.add() + item.key = k + item.value_json = json_dumps_safer_history(v) + self._publish_history(history) + + @abstractmethod + def _publish_history(self, history: pb.HistoryRecord) -> None: + raise NotImplementedError + + def publish_preempting(self) -> None: + preempt_rec = pb.RunPreemptingRecord() + self._publish_preempting(preempt_rec) + + @abstractmethod + def _publish_preempting(self, preempt_rec: pb.RunPreemptingRecord) -> None: + raise NotImplementedError + + def publish_output(self, name: str, data: str) -> None: + # from vendor.protobuf import google3.protobuf.timestamp + # ts = timestamp.Timestamp() + # ts.GetCurrentTime() + # now = datetime.now() + if name == "stdout": + otype = pb.OutputRecord.OutputType.STDOUT + elif name == "stderr": + otype = pb.OutputRecord.OutputType.STDERR + else: + # TODO(jhr): throw error? + termwarn("unknown type") + o = pb.OutputRecord(output_type=otype, line=data) + o.timestamp.GetCurrentTime() + self._publish_output(o) + + @abstractmethod + def _publish_output(self, outdata: pb.OutputRecord) -> None: + raise NotImplementedError + + def publish_output_raw(self, name: str, data: str) -> None: + # from vendor.protobuf import google3.protobuf.timestamp + # ts = timestamp.Timestamp() + # ts.GetCurrentTime() + # now = datetime.now() + if name == "stdout": + otype = pb.OutputRawRecord.OutputType.STDOUT + elif name == "stderr": + otype = pb.OutputRawRecord.OutputType.STDERR + else: + # TODO(jhr): throw error? + termwarn("unknown type") + o = pb.OutputRawRecord(output_type=otype, line=data) + o.timestamp.GetCurrentTime() + self._publish_output_raw(o) + + @abstractmethod + def _publish_output_raw(self, outdata: pb.OutputRawRecord) -> None: + raise NotImplementedError + + def publish_pause(self) -> None: + pause = pb.PauseRequest() + self._publish_pause(pause) + + @abstractmethod + def _publish_pause(self, pause: pb.PauseRequest) -> None: + raise NotImplementedError + + def publish_resume(self) -> None: + resume = pb.ResumeRequest() + self._publish_resume(resume) + + @abstractmethod + def _publish_resume(self, resume: pb.ResumeRequest) -> None: + raise NotImplementedError + + def publish_alert( + self, title: str, text: str, level: str, wait_duration: int + ) -> None: + proto_alert = pb.AlertRecord() + proto_alert.title = title + proto_alert.text = text + proto_alert.level = level + proto_alert.wait_duration = wait_duration + self._publish_alert(proto_alert) + + @abstractmethod + def _publish_alert(self, alert: pb.AlertRecord) -> None: + raise NotImplementedError + + def _make_exit(self, exit_code: Optional[int]) -> pb.RunExitRecord: + exit = pb.RunExitRecord() + if exit_code is not None: + exit.exit_code = exit_code + return exit + + def publish_exit(self, exit_code: Optional[int]) -> None: + exit_data = self._make_exit(exit_code) + self._publish_exit(exit_data) + + @abstractmethod + def _publish_exit(self, exit_data: pb.RunExitRecord) -> None: + raise NotImplementedError + + def publish_keepalive(self) -> None: + keepalive = pb.KeepaliveRequest() + self._publish_keepalive(keepalive) + + @abstractmethod + def _publish_keepalive(self, keepalive: pb.KeepaliveRequest) -> None: + raise NotImplementedError + + def publish_job_input( + self, + include_paths: List[List[str]], + exclude_paths: List[List[str]], + input_schema: Optional[dict], + run_config: bool = False, + file_path: str = "", + ): + """Publishes a request to add inputs to the job. + + If run_config is True, the wandb.config will be added as a job input. + If file_path is provided, the file at file_path will be added as a job + input. + + The paths provided as arguments are sequences of dictionary keys that + specify a path within the wandb.config. If a path is included, the + corresponding field will be treated as a job input. If a path is + excluded, the corresponding field will not be treated as a job input. + + Args: + include_paths: paths within config to include as job inputs. + exclude_paths: paths within config to exclude as job inputs. + input_schema: A JSON Schema describing which attributes will be + editable from the Launch drawer. + run_config: bool indicating whether wandb.config is the input source. + file_path: path to file to include as a job input. + """ + if run_config and file_path: + raise ValueError( + "run_config and file_path are mutually exclusive arguments." + ) + request = pb.JobInputRequest() + include_records = [pb.JobInputPath(path=path) for path in include_paths] + exclude_records = [pb.JobInputPath(path=path) for path in exclude_paths] + request.include_paths.extend(include_records) + request.exclude_paths.extend(exclude_records) + source = pb.JobInputSource( + run_config=pb.JobInputSource.RunConfigSource(), + ) + if run_config: + source.run_config.CopyFrom(pb.JobInputSource.RunConfigSource()) + else: + source.file.CopyFrom( + pb.JobInputSource.ConfigFileSource(path=file_path), + ) + request.input_source.CopyFrom(source) + if input_schema: + request.input_schema = json_dumps_safer(input_schema) + + return self._publish_job_input(request) + + @abstractmethod + def _publish_job_input( + self, request: pb.JobInputRequest + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def join(self) -> None: + # Drop indicates that the internal process has already been shutdown + if self._drop: + return + + handle = self._deliver_shutdown() + + try: + handle.wait_or(timeout=30) + except TimeoutError: + # This can happen if the server fails to respond due to a bug + # or due to being very busy. + logger.warning("timed out communicating shutdown") + except HandleAbandonedError: + # This can happen if the connection to the server is closed + # before a response is read. + logger.warning("handle abandoned while communicating shutdown") + + @abstractmethod + def _deliver_shutdown(self) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_run(self, run: "Run") -> MailboxHandle[pb.Result]: + run_record = self._make_run(run) + return self._deliver_run(run_record) + + def deliver_finish_sync( + self, + ) -> MailboxHandle[pb.Result]: + sync = pb.SyncFinishRequest() + return self._deliver_finish_sync(sync) + + @abstractmethod + def _deliver_finish_sync( + self, sync: pb.SyncFinishRequest + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + @abstractmethod + def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_run_start(self, run: "Run") -> MailboxHandle[pb.Result]: + run_start = pb.RunStartRequest(run=self._make_run(run)) + return self._deliver_run_start(run_start) + + @abstractmethod + def _deliver_run_start( + self, run_start: pb.RunStartRequest + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_attach(self, attach_id: str) -> MailboxHandle[pb.Result]: + attach = pb.AttachRequest(attach_id=attach_id) + return self._deliver_attach(attach) + + @abstractmethod + def _deliver_attach( + self, + status: pb.AttachRequest, + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_stop_status(self) -> MailboxHandle[pb.Result]: + status = pb.StopStatusRequest() + return self._deliver_stop_status(status) + + @abstractmethod + def _deliver_stop_status( + self, + status: pb.StopStatusRequest, + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_network_status(self) -> MailboxHandle[pb.Result]: + status = pb.NetworkStatusRequest() + return self._deliver_network_status(status) + + @abstractmethod + def _deliver_network_status( + self, + status: pb.NetworkStatusRequest, + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_internal_messages(self) -> MailboxHandle[pb.Result]: + internal_message = pb.InternalMessagesRequest() + return self._deliver_internal_messages(internal_message) + + @abstractmethod + def _deliver_internal_messages( + self, internal_message: pb.InternalMessagesRequest + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_get_summary(self) -> MailboxHandle[pb.Result]: + get_summary = pb.GetSummaryRequest() + return self._deliver_get_summary(get_summary) + + @abstractmethod + def _deliver_get_summary( + self, + get_summary: pb.GetSummaryRequest, + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_get_system_metrics(self) -> MailboxHandle[pb.Result]: + get_system_metrics = pb.GetSystemMetricsRequest() + return self._deliver_get_system_metrics(get_system_metrics) + + @abstractmethod + def _deliver_get_system_metrics( + self, get_summary: pb.GetSystemMetricsRequest + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_exit(self, exit_code: Optional[int]) -> MailboxHandle[pb.Result]: + exit_data = self._make_exit(exit_code) + return self._deliver_exit(exit_data) + + @abstractmethod + def _deliver_exit( + self, + exit_data: pb.RunExitRecord, + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + @abstractmethod + def deliver_operation_stats(self) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_poll_exit(self) -> MailboxHandle[pb.Result]: + poll_exit = pb.PollExitRequest() + return self._deliver_poll_exit(poll_exit) + + @abstractmethod + def _deliver_poll_exit( + self, + poll_exit: pb.PollExitRequest, + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_finish_without_exit(self) -> MailboxHandle[pb.Result]: + run_finish_without_exit = pb.RunFinishWithoutExitRequest() + return self._deliver_finish_without_exit(run_finish_without_exit) + + @abstractmethod + def _deliver_finish_without_exit( + self, run_finish_without_exit: pb.RunFinishWithoutExitRequest + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_request_sampled_history(self) -> MailboxHandle[pb.Result]: + sampled_history = pb.SampledHistoryRequest() + return self._deliver_request_sampled_history(sampled_history) + + @abstractmethod + def _deliver_request_sampled_history( + self, sampled_history: pb.SampledHistoryRequest + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError + + def deliver_request_run_status(self) -> MailboxHandle[pb.Result]: + run_status = pb.RunStatusRequest() + return self._deliver_request_run_status(run_status) + + @abstractmethod + def _deliver_request_run_status( + self, run_status: pb.RunStatusRequest + ) -> MailboxHandle[pb.Result]: + raise NotImplementedError diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/interface_queue.py b/lib/python3.12/site-packages/wandb/sdk/interface/interface_queue.py new file mode 100644 index 0000000000000000000000000000000000000000..84b89c66e54e3ba27d58ba997201943e245fb4ff --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/interface_queue.py @@ -0,0 +1,43 @@ +"""InterfaceQueue - Derived from InterfaceShared using queues to send to internal thread. + +See interface.py for how interface classes relate to each other. + +""" + +import logging +from multiprocessing.process import BaseProcess +from typing import TYPE_CHECKING, Optional + +from wandb.sdk.mailbox import Mailbox + +from .interface_shared import InterfaceShared + +if TYPE_CHECKING: + from queue import Queue + + from wandb.proto import wandb_internal_pb2 as pb + + +logger = logging.getLogger("wandb") + + +class InterfaceQueue(InterfaceShared): + def __init__( + self, + record_q: Optional["Queue[pb.Record]"] = None, + result_q: Optional["Queue[pb.Result]"] = None, + process: Optional[BaseProcess] = None, + mailbox: Optional[Mailbox] = None, + ) -> None: + self.record_q = record_q + self.result_q = result_q + self._process = process + super().__init__(mailbox=mailbox) + + def _publish(self, record: "pb.Record", local: Optional[bool] = None) -> None: + if self._process and not self._process.is_alive(): + raise Exception("The wandb backend process has shutdown") + if local: + record.control.local = local + if self.record_q: + self.record_q.put(record) diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/interface_relay.py b/lib/python3.12/site-packages/wandb/sdk/interface/interface_relay.py new file mode 100644 index 0000000000000000000000000000000000000000..ded7989ddba3d7d6842f54ea3ebcfaccdb28aacb --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/interface_relay.py @@ -0,0 +1,38 @@ +"""InterfaceRelay - Derived from InterfaceQueue using RelayRouter to preserve uuid req/resp. + +See interface.py for how interface classes relate to each other. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from wandb.proto import wandb_internal_pb2 as pb +from wandb.sdk.mailbox import Mailbox + +from .interface_queue import InterfaceQueue + +if TYPE_CHECKING: + from queue import Queue + + +logger = logging.getLogger("wandb") + + +class InterfaceRelay(InterfaceQueue): + _mailbox: Mailbox + + def __init__( + self, + mailbox: Mailbox, + record_q: Queue[pb.Record], + result_q: Queue[pb.Result], + relay_q: Queue[pb.Result], + ) -> None: + self.relay_q = relay_q + super().__init__( + record_q=record_q, + result_q=result_q, + mailbox=mailbox, + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/interface_shared.py b/lib/python3.12/site-packages/wandb/sdk/interface/interface_shared.py new file mode 100644 index 0000000000000000000000000000000000000000..8e832743b33bc3225486654c074af69ef35b64b6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/interface_shared.py @@ -0,0 +1,482 @@ +"""InterfaceShared - Derived from InterfaceBase - shared with InterfaceQueue and InterfaceSock. + +See interface.py for how interface classes relate to each other. + +""" + +import logging +from abc import abstractmethod +from typing import Any, Optional, cast + +from wandb.proto import wandb_internal_pb2 as pb +from wandb.proto import wandb_telemetry_pb2 as tpb +from wandb.sdk.mailbox import Mailbox, MailboxHandle +from wandb.util import json_dumps_safer, json_friendly + +from .interface import InterfaceBase + +logger = logging.getLogger("wandb") + + +class InterfaceShared(InterfaceBase): + def __init__(self, mailbox: Optional[Mailbox] = None) -> None: + super().__init__() + self._mailbox = mailbox + + def _publish_output(self, outdata: pb.OutputRecord) -> None: + rec = pb.Record() + rec.output.CopyFrom(outdata) + self._publish(rec) + + def _publish_cancel(self, cancel: pb.CancelRequest) -> None: + rec = self._make_request(cancel=cancel) + self._publish(rec) + + def _publish_output_raw(self, outdata: pb.OutputRawRecord) -> None: + rec = pb.Record() + rec.output_raw.CopyFrom(outdata) + self._publish(rec) + + def _publish_tbdata(self, tbrecord: pb.TBRecord) -> None: + rec = self._make_record(tbrecord=tbrecord) + self._publish(rec) + + def _publish_partial_history( + self, partial_history: pb.PartialHistoryRequest + ) -> None: + rec = self._make_request(partial_history=partial_history) + self._publish(rec) + + def _publish_history(self, history: pb.HistoryRecord) -> None: + rec = self._make_record(history=history) + self._publish(rec) + + def _publish_preempting(self, preempt_rec: pb.RunPreemptingRecord) -> None: + rec = self._make_record(preempting=preempt_rec) + self._publish(rec) + + def _publish_telemetry(self, telem: tpb.TelemetryRecord) -> None: + rec = self._make_record(telemetry=telem) + self._publish(rec) + + def _publish_environment(self, environment: pb.EnvironmentRecord) -> None: + rec = self._make_record(environment=environment) + self._publish(rec) + + def _publish_job_input( + self, job_input: pb.JobInputRequest + ) -> MailboxHandle[pb.Result]: + record = self._make_request(job_input=job_input) + return self._deliver_record(record) + + def _make_stats(self, stats_dict: dict) -> pb.StatsRecord: + stats = pb.StatsRecord() + stats.stats_type = pb.StatsRecord.StatsType.SYSTEM + stats.timestamp.GetCurrentTime() # todo: fix this, this is wrong :) + for k, v in stats_dict.items(): + item = stats.item.add() + item.key = k + item.value_json = json_dumps_safer(json_friendly(v)[0]) + return stats + + def _make_request( # noqa: C901 + self, + get_summary: Optional[pb.GetSummaryRequest] = None, + pause: Optional[pb.PauseRequest] = None, + resume: Optional[pb.ResumeRequest] = None, + status: Optional[pb.StatusRequest] = None, + stop_status: Optional[pb.StopStatusRequest] = None, + internal_messages: Optional[pb.InternalMessagesRequest] = None, + network_status: Optional[pb.NetworkStatusRequest] = None, + operation_stats: Optional[pb.OperationStatsRequest] = None, + poll_exit: Optional[pb.PollExitRequest] = None, + partial_history: Optional[pb.PartialHistoryRequest] = None, + sampled_history: Optional[pb.SampledHistoryRequest] = None, + run_start: Optional[pb.RunStartRequest] = None, + check_version: Optional[pb.CheckVersionRequest] = None, + log_artifact: Optional[pb.LogArtifactRequest] = None, + download_artifact: Optional[pb.DownloadArtifactRequest] = None, + link_artifact: Optional[pb.LinkArtifactRequest] = None, + defer: Optional[pb.DeferRequest] = None, + attach: Optional[pb.AttachRequest] = None, + server_info: Optional[pb.ServerInfoRequest] = None, + keepalive: Optional[pb.KeepaliveRequest] = None, + run_status: Optional[pb.RunStatusRequest] = None, + sender_mark: Optional[pb.SenderMarkRequest] = None, + sender_read: Optional[pb.SenderReadRequest] = None, + sync_finish: Optional[pb.SyncFinishRequest] = None, + status_report: Optional[pb.StatusReportRequest] = None, + cancel: Optional[pb.CancelRequest] = None, + summary_record: Optional[pb.SummaryRecordRequest] = None, + telemetry_record: Optional[pb.TelemetryRecordRequest] = None, + get_system_metrics: Optional[pb.GetSystemMetricsRequest] = None, + python_packages: Optional[pb.PythonPackagesRequest] = None, + job_input: Optional[pb.JobInputRequest] = None, + run_finish_without_exit: Optional[pb.RunFinishWithoutExitRequest] = None, + ) -> pb.Record: + request = pb.Request() + if get_summary: + request.get_summary.CopyFrom(get_summary) + elif pause: + request.pause.CopyFrom(pause) + elif resume: + request.resume.CopyFrom(resume) + elif status: + request.status.CopyFrom(status) + elif stop_status: + request.stop_status.CopyFrom(stop_status) + elif internal_messages: + request.internal_messages.CopyFrom(internal_messages) + elif network_status: + request.network_status.CopyFrom(network_status) + elif operation_stats: + request.operations.CopyFrom(operation_stats) + elif poll_exit: + request.poll_exit.CopyFrom(poll_exit) + elif partial_history: + request.partial_history.CopyFrom(partial_history) + elif sampled_history: + request.sampled_history.CopyFrom(sampled_history) + elif run_start: + request.run_start.CopyFrom(run_start) + elif check_version: + request.check_version.CopyFrom(check_version) + elif log_artifact: + request.log_artifact.CopyFrom(log_artifact) + elif download_artifact: + request.download_artifact.CopyFrom(download_artifact) + elif link_artifact: + request.link_artifact.CopyFrom(link_artifact) + elif defer: + request.defer.CopyFrom(defer) + elif attach: + request.attach.CopyFrom(attach) + elif server_info: + request.server_info.CopyFrom(server_info) + elif keepalive: + request.keepalive.CopyFrom(keepalive) + elif run_status: + request.run_status.CopyFrom(run_status) + elif sender_mark: + request.sender_mark.CopyFrom(sender_mark) + elif sender_read: + request.sender_read.CopyFrom(sender_read) + elif cancel: + request.cancel.CopyFrom(cancel) + elif status_report: + request.status_report.CopyFrom(status_report) + elif summary_record: + request.summary_record.CopyFrom(summary_record) + elif telemetry_record: + request.telemetry_record.CopyFrom(telemetry_record) + elif get_system_metrics: + request.get_system_metrics.CopyFrom(get_system_metrics) + elif sync_finish: + request.sync_finish.CopyFrom(sync_finish) + elif python_packages: + request.python_packages.CopyFrom(python_packages) + elif job_input: + request.job_input.CopyFrom(job_input) + elif run_finish_without_exit: + request.run_finish_without_exit.CopyFrom(run_finish_without_exit) + else: + raise Exception("Invalid request") + record = self._make_record(request=request) + # All requests do not get persisted + record.control.local = True + if status_report: + record.control.flow_control = True + return record + + def _make_record( # noqa: C901 + self, + run: Optional[pb.RunRecord] = None, + config: Optional[pb.ConfigRecord] = None, + files: Optional[pb.FilesRecord] = None, + summary: Optional[pb.SummaryRecord] = None, + history: Optional[pb.HistoryRecord] = None, + stats: Optional[pb.StatsRecord] = None, + exit: Optional[pb.RunExitRecord] = None, + artifact: Optional[pb.ArtifactRecord] = None, + tbrecord: Optional[pb.TBRecord] = None, + alert: Optional[pb.AlertRecord] = None, + final: Optional[pb.FinalRecord] = None, + metric: Optional[pb.MetricRecord] = None, + header: Optional[pb.HeaderRecord] = None, + footer: Optional[pb.FooterRecord] = None, + request: Optional[pb.Request] = None, + telemetry: Optional[tpb.TelemetryRecord] = None, + preempting: Optional[pb.RunPreemptingRecord] = None, + use_artifact: Optional[pb.UseArtifactRecord] = None, + output: Optional[pb.OutputRecord] = None, + output_raw: Optional[pb.OutputRawRecord] = None, + environment: Optional[pb.EnvironmentRecord] = None, + ) -> pb.Record: + record = pb.Record() + if run: + record.run.CopyFrom(run) + elif config: + record.config.CopyFrom(config) + elif summary: + record.summary.CopyFrom(summary) + elif history: + record.history.CopyFrom(history) + elif files: + record.files.CopyFrom(files) + elif stats: + record.stats.CopyFrom(stats) + elif exit: + record.exit.CopyFrom(exit) + elif artifact: + record.artifact.CopyFrom(artifact) + elif tbrecord: + record.tbrecord.CopyFrom(tbrecord) + elif alert: + record.alert.CopyFrom(alert) + elif final: + record.final.CopyFrom(final) + elif header: + record.header.CopyFrom(header) + elif footer: + record.footer.CopyFrom(footer) + elif request: + record.request.CopyFrom(request) + elif telemetry: + record.telemetry.CopyFrom(telemetry) + elif metric: + record.metric.CopyFrom(metric) + elif preempting: + record.preempting.CopyFrom(preempting) + elif use_artifact: + record.use_artifact.CopyFrom(use_artifact) + elif output: + record.output.CopyFrom(output) + elif output_raw: + record.output_raw.CopyFrom(output_raw) + elif environment: + record.environment.CopyFrom(environment) + else: + raise Exception("Invalid record") + return record + + @abstractmethod + def _publish(self, record: pb.Record, local: Optional[bool] = None) -> None: + raise NotImplementedError + + def _publish_defer(self, state: "pb.DeferRequest.DeferState.V") -> None: + defer = pb.DeferRequest(state=state) + rec = self._make_request(defer=defer) + self._publish(rec, local=True) + + def publish_defer(self, state: int = 0) -> None: + self._publish_defer(cast("pb.DeferRequest.DeferState.V", state)) + + def _publish_header(self, header: pb.HeaderRecord) -> None: + rec = self._make_record(header=header) + self._publish(rec) + + def publish_footer(self) -> None: + footer = pb.FooterRecord() + rec = self._make_record(footer=footer) + self._publish(rec) + + def publish_final(self) -> None: + final = pb.FinalRecord() + rec = self._make_record(final=final) + self._publish(rec) + + def _publish_pause(self, pause: pb.PauseRequest) -> None: + rec = self._make_request(pause=pause) + self._publish(rec) + + def _publish_resume(self, resume: pb.ResumeRequest) -> None: + rec = self._make_request(resume=resume) + self._publish(rec) + + def _publish_run(self, run: pb.RunRecord) -> None: + rec = self._make_record(run=run) + self._publish(rec) + + def _publish_config(self, cfg: pb.ConfigRecord) -> None: + rec = self._make_record(config=cfg) + self._publish(rec) + + def _publish_summary(self, summary: pb.SummaryRecord) -> None: + rec = self._make_record(summary=summary) + self._publish(rec) + + def _publish_metric(self, metric: pb.MetricRecord) -> None: + rec = self._make_record(metric=metric) + self._publish(rec) + + def publish_stats(self, stats_dict: dict) -> None: + stats = self._make_stats(stats_dict) + rec = self._make_record(stats=stats) + self._publish(rec) + + def _publish_python_packages( + self, python_packages: pb.PythonPackagesRequest + ) -> None: + rec = self._make_request(python_packages=python_packages) + self._publish(rec) + + def _publish_files(self, files: pb.FilesRecord) -> None: + rec = self._make_record(files=files) + self._publish(rec) + + def _publish_use_artifact(self, use_artifact: pb.UseArtifactRecord) -> Any: + rec = self._make_record(use_artifact=use_artifact) + self._publish(rec) + + def _deliver_artifact( + self, + log_artifact: pb.LogArtifactRequest, + ) -> MailboxHandle[pb.Result]: + rec = self._make_request(log_artifact=log_artifact) + return self._deliver_record(rec) + + def _deliver_download_artifact( + self, download_artifact: pb.DownloadArtifactRequest + ) -> MailboxHandle[pb.Result]: + rec = self._make_request(download_artifact=download_artifact) + return self._deliver_record(rec) + + def _deliver_link_artifact( + self, link_artifact: pb.LinkArtifactRequest + ) -> MailboxHandle[pb.Result]: + rec = self._make_request(link_artifact=link_artifact) + return self._deliver_record(rec) + + def _publish_artifact(self, proto_artifact: pb.ArtifactRecord) -> None: + rec = self._make_record(artifact=proto_artifact) + self._publish(rec) + + def _publish_alert(self, proto_alert: pb.AlertRecord) -> None: + rec = self._make_record(alert=proto_alert) + self._publish(rec) + + def _deliver_status( + self, + status: pb.StatusRequest, + ) -> MailboxHandle[pb.Result]: + req = self._make_request(status=status) + return self._deliver_record(req) + + def _publish_exit(self, exit_data: pb.RunExitRecord) -> None: + rec = self._make_record(exit=exit_data) + self._publish(rec) + + def _publish_keepalive(self, keepalive: pb.KeepaliveRequest) -> None: + record = self._make_request(keepalive=keepalive) + self._publish(record) + + def _deliver_shutdown(self) -> MailboxHandle[pb.Result]: + request = pb.Request(shutdown=pb.ShutdownRequest()) + record = self._make_record(request=request) + return self._deliver_record(record) + + def _get_mailbox(self) -> Mailbox: + mailbox = self._mailbox + assert mailbox + return mailbox + + def _deliver_record(self, record: pb.Record) -> MailboxHandle[pb.Result]: + mailbox = self._get_mailbox() + + handle = mailbox.require_response(record) + self._publish(record) + + return handle.map(lambda resp: resp.result_communicate) + + def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle[pb.Result]: + record = self._make_record(run=run) + return self._deliver_record(record) + + def _deliver_finish_sync( + self, + sync_finish: pb.SyncFinishRequest, + ) -> MailboxHandle[pb.Result]: + record = self._make_request(sync_finish=sync_finish) + return self._deliver_record(record) + + def _deliver_run_start( + self, + run_start: pb.RunStartRequest, + ) -> MailboxHandle[pb.Result]: + record = self._make_request(run_start=run_start) + return self._deliver_record(record) + + def _deliver_get_summary( + self, + get_summary: pb.GetSummaryRequest, + ) -> MailboxHandle[pb.Result]: + record = self._make_request(get_summary=get_summary) + return self._deliver_record(record) + + def _deliver_get_system_metrics( + self, get_system_metrics: pb.GetSystemMetricsRequest + ) -> MailboxHandle[pb.Result]: + record = self._make_request(get_system_metrics=get_system_metrics) + return self._deliver_record(record) + + def _deliver_exit( + self, + exit_data: pb.RunExitRecord, + ) -> MailboxHandle[pb.Result]: + record = self._make_record(exit=exit_data) + return self._deliver_record(record) + + def deliver_operation_stats(self): + record = self._make_request(operation_stats=pb.OperationStatsRequest()) + return self._deliver_record(record) + + def _deliver_poll_exit( + self, + poll_exit: pb.PollExitRequest, + ) -> MailboxHandle[pb.Result]: + record = self._make_request(poll_exit=poll_exit) + return self._deliver_record(record) + + def _deliver_finish_without_exit( + self, run_finish_without_exit: pb.RunFinishWithoutExitRequest + ) -> MailboxHandle[pb.Result]: + record = self._make_request(run_finish_without_exit=run_finish_without_exit) + return self._deliver_record(record) + + def _deliver_stop_status( + self, + stop_status: pb.StopStatusRequest, + ) -> MailboxHandle[pb.Result]: + record = self._make_request(stop_status=stop_status) + return self._deliver_record(record) + + def _deliver_attach( + self, + attach: pb.AttachRequest, + ) -> MailboxHandle[pb.Result]: + record = self._make_request(attach=attach) + return self._deliver_record(record) + + def _deliver_network_status( + self, network_status: pb.NetworkStatusRequest + ) -> MailboxHandle[pb.Result]: + record = self._make_request(network_status=network_status) + return self._deliver_record(record) + + def _deliver_internal_messages( + self, internal_message: pb.InternalMessagesRequest + ) -> MailboxHandle[pb.Result]: + record = self._make_request(internal_messages=internal_message) + return self._deliver_record(record) + + def _deliver_request_sampled_history( + self, sampled_history: pb.SampledHistoryRequest + ) -> MailboxHandle[pb.Result]: + record = self._make_request(sampled_history=sampled_history) + return self._deliver_record(record) + + def _deliver_request_run_status( + self, run_status: pb.RunStatusRequest + ) -> MailboxHandle[pb.Result]: + record = self._make_request(run_status=run_status) + return self._deliver_record(record) diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/interface_sock.py b/lib/python3.12/site-packages/wandb/sdk/interface/interface_sock.py new file mode 100644 index 0000000000000000000000000000000000000000..90aecff3c0f96bf5fe3b4021383b45e93359dca8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/interface_sock.py @@ -0,0 +1,39 @@ +"""InterfaceSock - Derived from InterfaceShared using a socket to send to internal thread. + +See interface.py for how interface classes relate to each other. + +""" + +import logging +from typing import TYPE_CHECKING, Any, Optional + +from wandb.sdk.mailbox import Mailbox + +from ..lib.sock_client import SockClient +from .interface_shared import InterfaceShared + +if TYPE_CHECKING: + from wandb.proto import wandb_internal_pb2 as pb + + +logger = logging.getLogger("wandb") + + +class InterfaceSock(InterfaceShared): + def __init__( + self, + sock_client: SockClient, + mailbox: Mailbox, + stream_id: str, + ) -> None: + super().__init__(mailbox=mailbox) + self._sock_client = sock_client + self._stream_id = stream_id + + def _assign(self, record: Any) -> None: + assert self._stream_id + record._info.stream_id = self._stream_id + + def _publish(self, record: "pb.Record", local: Optional[bool] = None) -> None: + self._assign(record) + self._sock_client.send_record_publish(record) diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/router.py b/lib/python3.12/site-packages/wandb/sdk/interface/router.py new file mode 100644 index 0000000000000000000000000000000000000000..0ebcb75d6d095c354ae200dabcf45599e39480e1 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/router.py @@ -0,0 +1,89 @@ +"""Router - handle message router (base class). + +Router to manage responses. + +""" + +from __future__ import annotations + +import logging +import threading +from abc import abstractmethod +from typing import TYPE_CHECKING + +from wandb.proto import wandb_internal_pb2 as pb +from wandb.proto import wandb_server_pb2 as spb +from wandb.sdk import mailbox + +if TYPE_CHECKING: + from queue import Queue + + +logger = logging.getLogger("wandb") + + +class MessageRouterClosedError(Exception): + """Router has been closed.""" + + +class MessageRouter: + _request_queue: Queue[pb.Record] + _response_queue: Queue[pb.Result] + _mailbox: mailbox.Mailbox | None + + def __init__(self, mailbox: mailbox.Mailbox | None = None) -> None: + self._mailbox = mailbox + self._lock = threading.Lock() + + self._join_event = threading.Event() + self._thread = threading.Thread(target=self.message_loop) + self._thread.name = "MsgRouterThr" + self._thread.daemon = True + self._thread.start() + + @abstractmethod + def _read_message(self) -> pb.Result | spb.ServerResponse | None: + raise NotImplementedError + + @abstractmethod + def _send_message(self, record: pb.Record) -> None: + raise NotImplementedError + + def message_loop(self) -> None: + try: + while not self._join_event.is_set(): + try: + msg = self._read_message() + except EOFError: + # On abnormal shutdown the queue will be destroyed underneath + # resulting in EOFError. message_loop needs to exit.. + logger.warning("EOFError seen in message_loop") + break + except MessageRouterClosedError as e: + logger.warning("message_loop has been closed", exc_info=e) + break + if not msg: + continue + self._handle_msg_rcv(msg) + + finally: + if self._mailbox: + self._mailbox.close() + + def join(self) -> None: + self._join_event.set() + self._thread.join() + + def _handle_msg_rcv(self, msg: pb.Result | spb.ServerResponse) -> None: + if not self._mailbox: + return + + if isinstance(msg, pb.Result) and msg.control.mailbox_slot: + self._mailbox.deliver( + spb.ServerResponse( + request_id=msg.control.mailbox_slot, + result_communicate=msg, + ) + ) + elif isinstance(msg, spb.ServerResponse) and msg.request_id: + self._mailbox.deliver(msg) diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/router_queue.py b/lib/python3.12/site-packages/wandb/sdk/interface/router_queue.py new file mode 100644 index 0000000000000000000000000000000000000000..2fbe509782cc93f34730c5eb299c0f18ee778c09 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/router_queue.py @@ -0,0 +1,43 @@ +"""Router - handle message router (queue). + +Router to manage responses from a queue. + +""" + +from __future__ import annotations + +import queue +from typing import TYPE_CHECKING + +from wandb.proto import wandb_internal_pb2 as pb +from wandb.sdk.mailbox import Mailbox + +from .router import MessageRouter + +if TYPE_CHECKING: + from queue import Queue + + +class MessageQueueRouter(MessageRouter): + _request_queue: Queue[pb.Record] + _response_queue: Queue[pb.Result] + + def __init__( + self, + request_queue: Queue[pb.Record], + response_queue: Queue[pb.Result], + mailbox: Mailbox | None = None, + ) -> None: + self._request_queue = request_queue + self._response_queue = response_queue + super().__init__(mailbox=mailbox) + + def _read_message(self) -> pb.Result | None: + try: + msg = self._response_queue.get(timeout=1) + except queue.Empty: + return None + return msg + + def _send_message(self, record: pb.Record) -> None: + self._request_queue.put(record) diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/router_relay.py b/lib/python3.12/site-packages/wandb/sdk/interface/router_relay.py new file mode 100644 index 0000000000000000000000000000000000000000..f115668399147c17ffdcfbf4f54a7a3f2aa8a802 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/router_relay.py @@ -0,0 +1,50 @@ +"""Router - handle message router (relay). + +Router to manage responses from a queue with relay. + +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from wandb.proto import wandb_internal_pb2 as pb +from wandb.proto import wandb_server_pb2 as spb +from wandb.sdk.mailbox import Mailbox + +from .router_queue import MessageQueueRouter + +if TYPE_CHECKING: + from queue import Queue + + +class MessageRelayRouter(MessageQueueRouter): + _relay_queue: Queue[pb.Result] + + def __init__( + self, + request_queue: Queue[pb.Record], + response_queue: Queue[pb.Result], + relay_queue: Queue[pb.Result], + mailbox: Mailbox, + ) -> None: + self._relay_queue = relay_queue + super().__init__( + request_queue=request_queue, + response_queue=response_queue, + mailbox=mailbox, + ) + + def _handle_msg_rcv(self, msg: pb.Result | spb.ServerResponse) -> None: + if isinstance(msg, pb.Result): + relay_msg = msg + else: + relay_msg = msg.result_communicate + + # This is legacy-service logic for returning responses to the client. + # A different thread reads the "relay queue" and writes responses on + # the socket. + if relay_msg.control.relay_id: + self._relay_queue.put(relay_msg) + else: + super()._handle_msg_rcv(msg) diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/router_sock.py b/lib/python3.12/site-packages/wandb/sdk/interface/router_sock.py new file mode 100644 index 0000000000000000000000000000000000000000..649cc2ff4a7c26bd25031af297dcbbd1f75f3ea4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/router_sock.py @@ -0,0 +1,32 @@ +"""Router - handle message router (sock). + +Router to manage responses from a socket client. + +""" + +from __future__ import annotations + +from wandb.proto import wandb_internal_pb2 as pb +from wandb.proto import wandb_server_pb2 as spb +from wandb.sdk.lib.sock_client import SockClient, SockClientClosedError +from wandb.sdk.mailbox import Mailbox + +from .router import MessageRouter, MessageRouterClosedError + + +class MessageSockRouter(MessageRouter): + _sock_client: SockClient + _mailbox: Mailbox + + def __init__(self, sock_client: SockClient, mailbox: Mailbox) -> None: + self._sock_client = sock_client + super().__init__(mailbox=mailbox) + + def _read_message(self) -> spb.ServerResponse | None: + try: + return self._sock_client.read_server_response(timeout=1) + except SockClientClosedError as e: + raise MessageRouterClosedError from e + + def _send_message(self, record: pb.Record) -> None: + self._sock_client.send_record_communicate(record) diff --git a/lib/python3.12/site-packages/wandb/sdk/interface/summary_record.py b/lib/python3.12/site-packages/wandb/sdk/interface/summary_record.py new file mode 100644 index 0000000000000000000000000000000000000000..2050a39080759ef12c54364aaf76b1dee6ab7e1d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/interface/summary_record.py @@ -0,0 +1,67 @@ +"""Summary Record. + +This module implements a summary record as an intermediate format before being converted +to a protocol buffer. +""" + +import typing as t + + +class SummaryRecord: + """Encodes a diff -- analogous to the SummaryRecord protobuf message.""" + + update: t.List["SummaryItem"] + remove: t.List["SummaryItem"] + + def __init__(self): + self.update = [] + self.remove = [] + + def __str__(self): + s = "SummaryRecord:\n Update:\n " + s += "\n ".join([str(item) for item in self.update]) + s += "\n Remove:\n " + s += "\n ".join([str(item) for item in self.remove]) + s += "\n" + return s + + __repr__ = __str__ + + def _add_next_parent(self, parent_key): + with_next_parent = SummaryRecord() + with_next_parent.update = [ + item._add_next_parent(parent_key) for item in self.update + ] + with_next_parent.remove = [ + item._add_next_parent(parent_key) for item in self.remove + ] + + return with_next_parent + + +class SummaryItem: + """Analogous to the SummaryItem protobuf message.""" + + key: t.Tuple[str] + value: t.Any + + def __init__(self): + self.key = tuple() + self.value = None + + def __str__(self): + return "SummaryItem: key: " + str(self.key) + " value: " + str(self.value) + + __repr__ = __str__ + + def _add_next_parent(self, parent_key): + with_next_parent = SummaryItem() + + key = self.key + if not isinstance(key, tuple): + key = (key,) + + with_next_parent.key = (parent_key,) + self.key + with_next_parent.value = self.value + + return with_next_parent diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb483b5548a7ea33643302105ea4c34589b1bc88 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/_launch.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/_launch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c611e4fffbefefc16e798461a071e80318d3d54d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/_launch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/_launch_add.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/_launch_add.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d9be75493f858de39cd83a22e87d563e807b412 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/_launch_add.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/_project_spec.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/_project_spec.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d001a58b3163f47634ddfa6d9c29421b4ed0a1d0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/_project_spec.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/create_job.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/create_job.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9eedfb8e93149454c46e69e8eefd34148f51930 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/create_job.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/errors.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/errors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb2d57e4b4263810e86746107ee33d76d9d3caf8 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/errors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/git_reference.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/git_reference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..569fc18e308b372928e8fdc8c9d53069908fdf36 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/git_reference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/loader.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/loader.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e981a0b6035483a3b6ffe8c340967ac8ba6676f0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/loader.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25f6af473153e2eb5c7ae5da3d8bfc9b729f40b0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/wandb_reference.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/wandb_reference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21cbadbc6d46d2e23b7a233ff9aa1ed08c0b73a0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/__pycache__/wandb_reference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/__init__.py b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b26cf3280cb4ec183705630aabefec0496849c8b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/abstract.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/abstract.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56f1c0fbfc4c35a2815ed1088cf6ad8e8fca5faa Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/abstract.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/build.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/build.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5bbbfa8aa10af97f113ad413da752e75df73946f Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/build.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/context_manager.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/context_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6072995b5986dccd46bf5da158abb16623b4920c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/context_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/docker_builder.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/docker_builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c7f8d1cfa2f0bf8dc459df5c5e0c92f5c0aa17c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/docker_builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/kaniko_builder.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/kaniko_builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2252f0e2b86e6b22cbe9466663b7fe699b978430 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/kaniko_builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/noop.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/noop.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b6bd8c2e4b6a479adbae9aa53cb3d14b6336327 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/builder/__pycache__/noop.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/abstract.py b/lib/python3.12/site-packages/wandb/sdk/launch/builder/abstract.py new file mode 100644 index 0000000000000000000000000000000000000000..06bd33bc40e1cb398be05367bc326216c7e43a24 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/builder/abstract.py @@ -0,0 +1,156 @@ +"""Abstract plugin class defining the interface needed to build container images for W&B Launch.""" + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, Optional + +from wandb.sdk.launch.environment.abstract import AbstractEnvironment +from wandb.sdk.launch.registry.abstract import AbstractRegistry + +from .._project_spec import EntryPoint, LaunchProject +from ..registry.anon import AnonynmousRegistry +from ..utils import ( + AZURE_CONTAINER_REGISTRY_URI_REGEX, + ELASTIC_CONTAINER_REGISTRY_URI_REGEX, + GCP_ARTIFACT_REGISTRY_URI_REGEX, +) + +if TYPE_CHECKING: + from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker + + +class AbstractBuilder(ABC): + """Abstract plugin class defining the interface needed to build container images for W&B Launch.""" + + builder_type: str + environment: AbstractEnvironment + registry: AbstractRegistry + builder_config: Dict[str, Any] + + @abstractmethod + def __init__( + self, + environment: AbstractEnvironment, + registry: AbstractRegistry, + verify: bool = True, + ) -> None: + """Initialize a builder. + + Arguments: + builder_config: The builder config. + registry: The registry to use. + verify: Whether to verify the functionality of the builder. + + Raises: + LaunchError: If the builder cannot be initialized or verified. + """ + raise NotImplementedError + + @classmethod + @abstractmethod + def from_config( + cls, + config: dict, + environment: AbstractEnvironment, + registry: AbstractRegistry, + ) -> "AbstractBuilder": + """Create a builder from a config dictionary. + + Arguments: + config: The config dictionary. + environment: The environment to use. + registry: The registry to use. + verify: Whether to verify the functionality of the builder. + login: Whether to login to the registry immediately. + + Returns: + The builder. + """ + raise NotImplementedError + + @abstractmethod + async def build_image( + self, + launch_project: LaunchProject, + entrypoint: EntryPoint, + job_tracker: Optional["JobAndRunStatusTracker"] = None, + ) -> str: + """Build the image for the given project. + + Arguments: + launch_project: The project to build. + build_ctx_path: The path to the build context. + + Returns: + The image name. + """ + raise NotImplementedError + + @abstractmethod + async def verify(self) -> None: + """Verify that the builder can be used to build images. + + Raises: + LaunchError: If the builder cannot be used to build images. + """ + raise NotImplementedError + + +def registry_from_uri(uri: str) -> AbstractRegistry: + """Create a registry helper object from a uri. + + This function parses the URI and determines which supported registry it + belongs to. It then creates a registry helper object for that registry. + The supported remote registry types are: + - Azure Container Registry + - Google Container Registry + - AWS Elastic Container Registry + + The format of the URI is as follows: + - Azure Container Registry: .azurecr.io// + - Google Container Registry: -docker.pkg.dev/// + - AWS Elastic Container Registry: .dkr.ecr..amazonaws.com// + + Our classification of the registry is based on the domain name. For example, + if the uri contains `.azurecr.io`, we classify it as an Azure + Container Registry. If the uri contains `.dkr.ecr`, we classify + it as an AWS Elastic Container Registry. If the uri contains + `-docker.pkg.dev`, we classify it as a Google Artifact Registry. + + This function will attempt to load the appropriate cloud helpers for the + + `https://` prefix is optional for all of the above. + + Arguments: + uri: The uri to create a registry from. + + Returns: + The registry. + + Raises: + LaunchError: If the registry helper cannot be loaded for the given URI. + """ + if uri.startswith("https://"): + uri = uri[len("https://") :] + + if AZURE_CONTAINER_REGISTRY_URI_REGEX.match(uri) is not None: + from wandb.sdk.launch.registry.azure_container_registry import ( + AzureContainerRegistry, + ) + + return AzureContainerRegistry(uri=uri) + + elif GCP_ARTIFACT_REGISTRY_URI_REGEX.match(uri) is not None: + from wandb.sdk.launch.registry.google_artifact_registry import ( + GoogleArtifactRegistry, + ) + + return GoogleArtifactRegistry(uri=uri) + + elif ELASTIC_CONTAINER_REGISTRY_URI_REGEX.match(uri) is not None: + from wandb.sdk.launch.registry.elastic_container_registry import ( + ElasticContainerRegistry, + ) + + return ElasticContainerRegistry(uri=uri) + + return AnonynmousRegistry(uri=uri) diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/build.py b/lib/python3.12/site-packages/wandb/sdk/launch/builder/build.py new file mode 100644 index 0000000000000000000000000000000000000000..7e60714b07e8cec6ebdfcadf2eb84a9688328214 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/builder/build.py @@ -0,0 +1,296 @@ +import hashlib +import json +import logging +import os +import pathlib +import shlex +import shutil +from typing import Any, Dict, List, Tuple + +import wandb +import wandb.env +from wandb import docker +from wandb.apis.internal import Api +from wandb.sdk.launch.loader import ( + builder_from_config, + environment_from_config, + registry_from_config, +) +from wandb.util import get_module + +from .._project_spec import EntryPoint, LaunchProject +from ..errors import ExecutionError, LaunchError +from ..utils import LOG_PREFIX, event_loop_thread_exec +from .templates.dockerfile import ( + ACCELERATOR_SETUP_TEMPLATE, + ENTRYPOINT_TEMPLATE, + PIP_TEMPLATE, + PYTHON_SETUP_TEMPLATE, + USER_CREATE_TEMPLATE, +) + +_logger = logging.getLogger(__name__) + + +_WANDB_DOCKERFILE_NAME = "Dockerfile.wandb" + + +async def validate_docker_installation() -> None: + """Verify if Docker is installed on host machine.""" + find_exec = event_loop_thread_exec(shutil.which) + if not await find_exec("docker"): + raise ExecutionError( + "Could not find Docker executable. " + "Ensure Docker is installed as per the instructions " + "at https://docs.docker.com/install/overview/." + ) + + +def join(split_command: List[str]) -> str: + """Return a shell-escaped string from *split_command*. + + Also remove quotes from double quoted strings. Ex: + "'local container queue'" --> "local container queue" + """ + return " ".join(shlex.quote(arg.replace("'", "")) for arg in split_command) + + +async def build_image_from_project( + launch_project: LaunchProject, + api: Api, + launch_config: Dict[str, Any], +) -> str: + """Construct a docker image from a project and returns the URI of the image. + + Arguments: + launch_project: The project to build an image from. + api: The API object to use for fetching the project. + launch_config: The launch config to use for building the image. + + Returns: + The URI of the built image. + """ + assert launch_project.uri, "To build an image on queue a URI must be set." + launch_config = launch_config or {} + env_config = launch_config.get("environment", {}) + if not isinstance(env_config, dict): + wrong_type = type(env_config).__name__ + raise LaunchError( + f"Invalid environment config: {env_config} of type {wrong_type} " + "loaded from launch config. Expected dict." + ) + environment = environment_from_config(env_config) + + registry_config = launch_config.get("registry", {}) + if not isinstance(registry_config, dict): + wrong_type = type(registry_config).__name__ + raise LaunchError( + f"Invalid registry config: {registry_config} of type {wrong_type}" + " loaded from launch config. Expected dict." + ) + registry = registry_from_config(registry_config, environment) + + builder_config = launch_config.get("builder", {}) + if not isinstance(builder_config, dict): + wrong_type = type(builder_config).__name__ + raise LaunchError( + f"Invalid builder config: {builder_config} of type {wrong_type} " + "loaded from launch config. Expected dict." + ) + builder = builder_from_config(builder_config, environment, registry) + + if not builder: + raise LaunchError("Unable to build image. No builder found.") + + launch_project.fetch_and_validate_project() + + entry_point = ( + launch_project.get_job_entry_point() or launch_project.override_entrypoint + ) + assert entry_point is not None + wandb.termlog(f"{LOG_PREFIX}Building docker image from uri source") + image_uri = await builder.build_image(launch_project, entry_point) + if not image_uri: + raise LaunchError("Error building image uri") + else: + return image_uri + + +def image_tag_from_dockerfile_and_source( + launch_project: LaunchProject, dockerfile_contents: str +) -> str: + """Hashes the source and dockerfile contents into a unique tag.""" + image_source_string = launch_project.get_image_source_string() + unique_id_string = image_source_string + dockerfile_contents + image_tag = hashlib.sha256(unique_id_string.encode("utf-8")).hexdigest()[:8] + return image_tag + + +def get_docker_user(launch_project: LaunchProject, runner_type: str) -> Tuple[str, int]: + import getpass + + username = getpass.getuser() + + if runner_type == "sagemaker" and not launch_project.docker_image: + # unless user has provided their own image, sagemaker must run as root but keep the name for workdir etc + return username, 0 + + userid = launch_project.docker_user_id or os.geteuid() + return username, userid + + +def get_base_setup( + launch_project: LaunchProject, py_version: str, py_major: str +) -> str: + """Fill in the Dockerfile templates for stage 2 of build. + + CPU version is built on python, Accelerator version is built on user provided. + """ + minor = int(py_version.split(".")[1]) + if minor < 12: + python_base_image = f"python:{py_version}-buster" + else: + python_base_image = f"python:{py_version}-bookworm" + if launch_project.accelerator_base_image: + _logger.info( + f"Using accelerator base image: {launch_project.accelerator_base_image}" + ) + python_packages = [ + f"python{py_version}", + f"libpython{py_version}", + "python3-pip", + "python3-setuptools", + ] + base_setup = ACCELERATOR_SETUP_TEMPLATE.format( + accelerator_base_image=launch_project.accelerator_base_image, + python_packages=" \\\n".join(python_packages), + py_version=py_version, + ) + else: + python_packages = [ + "python3-dev", + "gcc", + ] # gcc required for python < 3.7 for some reason + base_setup = PYTHON_SETUP_TEMPLATE.format(py_base_image=python_base_image) + return base_setup + + +# Move this into the build context manager. +def get_requirements_section( + launch_project: LaunchProject, build_context_dir: str, builder_type: str +) -> str: + if builder_type == "docker": + buildx_installed = docker.is_buildx_installed() + if not buildx_installed: + wandb.termwarn( + "Docker BuildX is not installed, for faster builds upgrade docker: https://github.com/docker/buildx#installing" + ) + prefix = "RUN WANDB_DISABLE_CACHE=true" + elif builder_type == "kaniko": + prefix = "RUN WANDB_DISABLE_CACHE=true" + buildx_installed = False + + if buildx_installed: + prefix = "RUN --mount=type=cache,mode=0777,target=/root/.cache/pip" + + requirements_files = [] + deps_install_line = None + + base_path = pathlib.Path(build_context_dir) + # If there is a requirements.txt at root of build context, use that. + if (base_path / "src" / "requirements.txt").exists(): + requirements_files += ["src/requirements.txt"] + deps_install_line = "pip install uv && uv pip install -r requirements.txt" + with open(base_path / "src" / "requirements.txt") as f: + requirements = f.readlines() + if not any(["wandb" in r for r in requirements]): + wandb.termwarn(f"{LOG_PREFIX}wandb is not present in requirements.txt.") + return PIP_TEMPLATE.format( + buildx_optional_prefix=prefix, + requirements_files=" ".join(requirements_files), + pip_install=deps_install_line, + ) + + # Elif there is pyproject.toml at build context, convert the dependencies + # section to a requirements.txt and use that. + elif (base_path / "src" / "pyproject.toml").exists(): + tomli = get_module("tomli") + if tomli is None: + wandb.termwarn( + "pyproject.toml found but tomli could not be loaded. To " + "install dependencies from pyproject.toml please run " + "`pip install tomli` and try again." + ) + else: + # First try to read deps from standard pyproject format. + with open(base_path / "src" / "pyproject.toml", "rb") as f: + contents = tomli.load(f) + project_deps = [ + str(d) for d in contents.get("project", {}).get("dependencies", []) + ] + if project_deps: + if not any(["wandb" in d for d in project_deps]): + wandb.termwarn( + f"{LOG_PREFIX}wandb is not present as a dependency in pyproject.toml." + ) + with open(base_path / "src" / "requirements.txt", "w") as f: + f.write("\n".join(project_deps)) + requirements_files += ["src/requirements.txt"] + deps_install_line = ( + "pip install uv && uv pip install -r requirements.txt" + ) + return PIP_TEMPLATE.format( + buildx_optional_prefix=prefix, + requirements_files=" ".join(requirements_files), + pip_install=deps_install_line, + ) + + # Else use frozen requirements from wandb run. + if ( + not deps_install_line + and (base_path / "src" / "requirements.frozen.txt").exists() + ): + requirements_files += [ + "src/requirements.frozen.txt", + "_wandb_bootstrap.py", + ] + deps_install_line = ( + launch_project.parse_existing_requirements() + "python _wandb_bootstrap.py" + ) + + if not deps_install_line: + raise LaunchError(f"No dependency sources found for {launch_project}") + + with open(base_path / "src" / "requirements.frozen.txt") as f: + requirements = f.readlines() + if not any(["wandb" in r for r in requirements]): + wandb.termwarn( + f"{LOG_PREFIX}wandb is not present in requirements.frozen.txt." + ) + + return PIP_TEMPLATE.format( + buildx_optional_prefix=prefix, + requirements_files=" ".join(requirements_files), + pip_install=deps_install_line, + ) + + else: + # this means no deps file was found + requirements_line = "RUN mkdir -p env/" # Docker fails otherwise + wandb.termwarn("No requirements file found. No packages will be installed.") + return requirements_line + + +def get_user_setup(username: str, userid: int, runner_type: str) -> str: + if runner_type == "sagemaker": + # sagemaker must run as root + return "USER root" + user_create = USER_CREATE_TEMPLATE.format(uid=userid, user=username) + user_create += f"\nUSER {username}" + return user_create + + +def get_entrypoint_setup( + entry_point: EntryPoint, +) -> str: + return ENTRYPOINT_TEMPLATE.format(entrypoint=json.dumps(entry_point.command)) diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/context_manager.py b/lib/python3.12/site-packages/wandb/sdk/launch/builder/context_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..c1ceb3a21ce4627e7468339653343364ceac8b1b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/builder/context_manager.py @@ -0,0 +1,235 @@ +import logging +import os +import shutil +import tempfile +from typing import Tuple + +from wandb.sdk.launch._project_spec import LaunchProject +from wandb.sdk.launch.builder.build import image_tag_from_dockerfile_and_source +from wandb.sdk.launch.errors import LaunchError +from wandb.sdk.launch.utils import get_current_python_version + +from .build import ( + _WANDB_DOCKERFILE_NAME, + get_base_setup, + get_docker_user, + get_entrypoint_setup, + get_requirements_section, + get_user_setup, +) +from .templates.dockerfile import DOCKERFILE_TEMPLATE + +_logger = logging.getLogger(__name__) + + +class BuildContextManager: + """Creates a build context for a container image from job source code. + + The dockerfile and build context may be specified by the job itself. If not, + the behavior for creating the build context is as follows: + + - If a Dockerfile.wandb is found adjacent to the entrypoint, the directory + containing the entrypoint is used as the build context and Dockerfile.wandb + is used as the Dockerfile. + + - If `override_dockerfile` is set on the LaunchProject, the directory + containing the Dockerfile is used as the build context and the Dockerfile + is used as the Dockerfile. `override_dockerfile` can be set in a launch + spec via the `-D` flag to `wandb launch` or in the `overrides` section + of the launch drawer. + + - If no dockerfile is set, a Dockerfile is generated from the job's + requirements and entrypoint. + """ + + def __init__(self, launch_project: LaunchProject): + """Initialize a BuildContextManager. + + Arguments: + launch_project: The launch project. + """ + self._launch_project = launch_project + assert self._launch_project.project_dir is not None + self._directory = tempfile.mkdtemp() + + def _generate_dockerfile(self, builder_type: str) -> str: + """Generate a Dockerfile for the container image. + + Arguments: + builder_type: The type of builder to use. One of "docker" or "kaniko". + + Returns: + The contents of the Dockerfile. + """ + launch_project = self._launch_project + entry_point = ( + launch_project.override_entrypoint or launch_project.get_job_entry_point() + ) + + # get python versions truncated to major.minor to ensure image availability + if launch_project.python_version: + spl = launch_project.python_version.split(".")[:2] + py_version, py_major = (".".join(spl), spl[0]) + else: + py_version, py_major = get_current_python_version() + + python_build_image = ( + f"python:{py_version}" # use full python image for package installation + ) + requirements_section = get_requirements_section( + launch_project, self._directory, builder_type + ) + # ----- stage 2: base ----- + python_base_setup = get_base_setup(launch_project, py_version, py_major) + + # set up user info + username, userid = get_docker_user(launch_project, launch_project.resource) + user_setup = get_user_setup(username, userid, launch_project.resource) + workdir = f"/home/{username}" + + assert entry_point is not None + entrypoint_section = get_entrypoint_setup(entry_point) + + dockerfile_contents = DOCKERFILE_TEMPLATE.format( + py_build_image=python_build_image, + requirements_section=requirements_section, + base_setup=python_base_setup, + uid=userid, + user_setup=user_setup, + workdir=workdir, + entrypoint_section=entrypoint_section, + ) + return dockerfile_contents + + def create_build_context(self, builder_type: str) -> Tuple[str, str]: + """Create the build context for the container image. + + Returns: + A pair of str: the path to the build context locally and the image + tag computed from the Dockerfile. + """ + entrypoint = ( + self._launch_project.get_job_entry_point() + or self._launch_project.override_entrypoint + ) + assert entrypoint is not None + assert entrypoint.name is not None + assert self._launch_project.project_dir is not None + + # we use that as the build context. + build_context_root_dir = self._launch_project.project_dir + job_build_context = self._launch_project.job_build_context + if job_build_context: + full_path = os.path.join(build_context_root_dir, job_build_context) + if not os.path.exists(full_path): + raise LaunchError(f"Build context does not exist at {full_path}") + build_context_root_dir = full_path + + # This is the case where the user specifies a Dockerfile to use. + # We use the directory containing the Dockerfile as the build context. + override_dockerfile = self._launch_project.override_dockerfile + if override_dockerfile: + full_path = os.path.join( + build_context_root_dir, + override_dockerfile, + ) + if not os.path.exists(full_path): + raise LaunchError(f"Dockerfile does not exist at {full_path}") + shutil.copytree( + build_context_root_dir, + self._directory, + symlinks=True, + dirs_exist_ok=True, + ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc"), + ) + shutil.copy( + full_path, + os.path.join(self._directory, _WANDB_DOCKERFILE_NAME), + ) + return self._directory, image_tag_from_dockerfile_and_source( + self._launch_project, open(full_path).read() + ) + + # If the job specifies a Dockerfile, we use that as the Dockerfile. + job_dockerfile = self._launch_project.job_dockerfile + if job_dockerfile: + dockerfile_path = os.path.join(build_context_root_dir, job_dockerfile) + if not os.path.exists(dockerfile_path): + raise LaunchError(f"Dockerfile does not exist at {dockerfile_path}") + shutil.copytree( + build_context_root_dir, + self._directory, + symlinks=True, + dirs_exist_ok=True, + ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc"), + ) + shutil.copy( + dockerfile_path, + os.path.join(self._directory, _WANDB_DOCKERFILE_NAME), + ) + return self._directory, image_tag_from_dockerfile_and_source( + self._launch_project, open(dockerfile_path).read() + ) + + # This is the case where we find Dockerfile.wandb adjacent to the + # entrypoint. We use the entrypoint directory as the build context. + entrypoint_dir = os.path.dirname(entrypoint.name) + if entrypoint_dir: + path = os.path.join( + build_context_root_dir, + entrypoint_dir, + _WANDB_DOCKERFILE_NAME, + ) + else: + path = os.path.join(build_context_root_dir, _WANDB_DOCKERFILE_NAME) + if os.path.exists( + path + ): # We found a Dockerfile.wandb adjacent to the entrypoint. + shutil.copytree( + os.path.dirname(path), + self._directory, + symlinks=True, + dirs_exist_ok=True, + ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc"), + ) + # TODO: remove this once we make things more explicit for users + if entrypoint_dir: + new_path = os.path.basename(entrypoint.name) + entrypoint = self._launch_project.get_job_entry_point() + if entrypoint is not None: + entrypoint.update_entrypoint_path(new_path) + with open(path) as f: + docker_file_contents = f.read() + return self._directory, image_tag_from_dockerfile_and_source( + self._launch_project, docker_file_contents + ) + + # This is the case where we use our own Dockerfile template. We move + # the user code into a src directory in the build context. + dst_path = os.path.join(self._directory, "src") + assert self._launch_project.project_dir is not None + shutil.copytree( + src=self._launch_project.project_dir, + dst=dst_path, + symlinks=True, + ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc"), + ) + shutil.copy( + os.path.join(os.path.dirname(__file__), "templates", "_wandb_bootstrap.py"), + os.path.join(self._directory), + ) + if self._launch_project.python_version: + runtime_path = os.path.join(dst_path, "runtime.txt") + with open(runtime_path, "w") as fp: + fp.write(f"python-{self._launch_project.python_version}") + + # TODO: we likely don't need to pass the whole git repo into the container + # with open(os.path.join(directory, ".dockerignore"), "w") as f: + # f.write("**/.git") + with open(os.path.join(self._directory, _WANDB_DOCKERFILE_NAME), "w") as handle: + docker_file_contents = self._generate_dockerfile(builder_type=builder_type) + handle.write(docker_file_contents) + image_tag = image_tag_from_dockerfile_and_source( + self._launch_project, docker_file_contents + ) + return self._directory, image_tag diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/docker_builder.py b/lib/python3.12/site-packages/wandb/sdk/launch/builder/docker_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..d764ba005e38f3346452124816c3b2120ea81fdf --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/builder/docker_builder.py @@ -0,0 +1,177 @@ +"""Implementation of the docker builder.""" + +import logging +import os +from typing import Any, Dict, Optional + +import wandb +import wandb.docker as docker +from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker +from wandb.sdk.launch.builder.abstract import AbstractBuilder, registry_from_uri +from wandb.sdk.launch.environment.abstract import AbstractEnvironment +from wandb.sdk.launch.registry.abstract import AbstractRegistry + +from .._project_spec import EntryPoint, LaunchProject +from ..errors import LaunchDockerError, LaunchError +from ..registry.anon import AnonynmousRegistry +from ..registry.local_registry import LocalRegistry +from ..utils import ( + LOG_PREFIX, + event_loop_thread_exec, + warn_failed_packages_from_build_logs, +) +from .build import _WANDB_DOCKERFILE_NAME, validate_docker_installation +from .context_manager import BuildContextManager + +_logger = logging.getLogger(__name__) + + +class DockerBuilder(AbstractBuilder): + """Builds a docker image for a project. + + Attributes: + builder_config (Dict[str, Any]): The builder config. + + """ + + builder_type = "docker" + target_platform = "linux/amd64" + + def __init__( + self, + environment: AbstractEnvironment, + registry: AbstractRegistry, + config: Dict[str, Any], + ): + """Initialize a DockerBuilder. + + Arguments: + environment (AbstractEnvironment): The environment to use. + registry (AbstractRegistry): The registry to use. + + Raises: + LaunchError: If docker is not installed + """ + self.environment = environment # Docker builder doesn't actually use this. + self.registry = registry + self.config = config + + @classmethod + def from_config( + cls, + config: Dict[str, Any], + environment: AbstractEnvironment, + registry: AbstractRegistry, + ) -> "DockerBuilder": + """Create a DockerBuilder from a config. + + Arguments: + config (Dict[str, Any]): The config. + registry (AbstractRegistry): The registry to use. + verify (bool, optional): Whether to verify the functionality of the builder. + login (bool, optional): Whether to login to the registry. + + Returns: + DockerBuilder: The DockerBuilder. + """ + # If the user provided a destination URI in the builder config + # we use that as the registry. + image_uri = config.get("destination") + if image_uri: + if registry is not None: + wandb.termwarn( + f"{LOG_PREFIX}Overriding registry from registry config" + f" with {image_uri} from builder config." + ) + registry = registry_from_uri(image_uri) + + return cls(environment, registry, config) + + async def verify(self) -> None: + """Verify the builder.""" + await validate_docker_installation() + + async def login(self) -> None: + """Login to the registry.""" + if isinstance(self.registry, LocalRegistry): + _logger.info(f"{LOG_PREFIX}No registry configured, skipping login.") + elif isinstance(self.registry, AnonynmousRegistry): + _logger.info(f"{LOG_PREFIX}Anonymous registry, skipping login.") + else: + username, password = await self.registry.get_username_password() + login = event_loop_thread_exec(docker.login) + await login(username, password, self.registry.uri) + + async def build_image( + self, + launch_project: LaunchProject, + entrypoint: EntryPoint, + job_tracker: Optional[JobAndRunStatusTracker] = None, + ) -> str: + """Build the image for the given project. + + Arguments: + launch_project (LaunchProject): The project to build. + entrypoint (EntryPoint): The entrypoint to use. + """ + await self.verify() + await self.login() + + build_context_manager = BuildContextManager(launch_project=launch_project) + build_ctx_path, image_tag = build_context_manager.create_build_context("docker") + dockerfile = os.path.join(build_ctx_path, _WANDB_DOCKERFILE_NAME) + repository = None if not self.registry else await self.registry.get_repo_uri() + + # if repo is set, use the repo name as the image name + if repository: + image_uri = f"{repository}:{image_tag}" + # otherwise, base the image name off of the source + # which the launch_project checks in image_name + else: + image_uri = f"{launch_project.image_name}:{image_tag}" + + if ( + not launch_project.build_required() + and await self.registry.check_image_exists(image_uri) + ): + return image_uri + + _logger.info( + f"image {image_uri} does not already exist in repository, building." + ) + try: + output = await event_loop_thread_exec(docker.build)( + tags=[image_uri], + file=dockerfile, + context_path=build_ctx_path, + platform=self.config.get("platform"), + ) + + warn_failed_packages_from_build_logs( + output, image_uri, launch_project.api, job_tracker + ) + + except docker.DockerError as e: + if job_tracker: + job_tracker.set_err_stage("build") + raise LaunchDockerError(f"Error communicating with docker client: {e}") + + try: + os.remove(build_ctx_path) + except Exception: + _msg = f"{LOG_PREFIX}Temporary docker context file {build_ctx_path} was not deleted." + _logger.info(_msg) + + if repository: + reg, tag = image_uri.split(":") + wandb.termlog(f"{LOG_PREFIX}Pushing image {image_uri}") + push_resp = await event_loop_thread_exec(docker.push)(reg, tag) + if push_resp is None: + raise LaunchError("Failed to push image to repository") + elif ( + launch_project.resource == "sagemaker" + and f"The push refers to repository [{repository}]" not in push_resp + ): + raise LaunchError(f"Unable to push image to ECR, response: {push_resp}") + + return image_uri diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/kaniko_builder.py b/lib/python3.12/site-packages/wandb/sdk/launch/builder/kaniko_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..3740c05e24eb2b723473a78f03c5c2664a27f8bf --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/builder/kaniko_builder.py @@ -0,0 +1,595 @@ +import asyncio +import base64 +import copy +import json +import logging +import os +import shutil +import tarfile +import tempfile +import time +import traceback +from typing import Any, Dict, Optional + +import wandb +from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker +from wandb.sdk.launch.builder.abstract import AbstractBuilder, registry_from_uri +from wandb.sdk.launch.environment.abstract import AbstractEnvironment +from wandb.sdk.launch.environment.azure_environment import AzureEnvironment +from wandb.sdk.launch.registry.abstract import AbstractRegistry +from wandb.sdk.launch.registry.azure_container_registry import AzureContainerRegistry +from wandb.sdk.launch.registry.elastic_container_registry import ( + ElasticContainerRegistry, +) +from wandb.sdk.launch.registry.google_artifact_registry import GoogleArtifactRegistry +from wandb.util import get_module + +from .._project_spec import EntryPoint, LaunchProject +from ..errors import LaunchError +from ..utils import ( + LOG_PREFIX, + get_kube_context_and_api_client, + warn_failed_packages_from_build_logs, +) +from .build import _WANDB_DOCKERFILE_NAME +from .context_manager import BuildContextManager + +get_module( + "kubernetes_asyncio", + required="Kaniko builder requires the kubernetes_asyncio package. Please install it with `pip install wandb[launch]`.", +) + +import kubernetes_asyncio as kubernetes # type: ignore # noqa: E402 +from kubernetes_asyncio import client # noqa: E402 + +_logger = logging.getLogger(__name__) + +_DEFAULT_BUILD_TIMEOUT_SECS = 1800 # 30 minute build timeout + +SERVICE_ACCOUNT_NAME = os.environ.get("WANDB_LAUNCH_SERVICE_ACCOUNT_NAME", "default") +PVC_NAME = os.environ.get("WANDB_LAUNCH_KANIKO_PVC_NAME") +PVC_MOUNT_PATH = ( + os.environ.get("WANDB_LAUNCH_KANIKO_PVC_MOUNT_PATH", "/kaniko").rstrip("/") + if PVC_NAME + else None +) +DOCKER_CONFIG_SECRET = os.environ.get("WANDB_LAUNCH_KANIKO_AUTH_SECRET") + + +if os.path.exists("/var/run/secrets/kubernetes.io/serviceaccount/namespace"): + with open("/var/run/secrets/kubernetes.io/serviceaccount/namespace") as f: + NAMESPACE = f.read().strip() +else: + NAMESPACE = "wandb" + + +def get_pod_name_safe(job: client.V1Job): + try: + return job.spec.template.metadata.name + except AttributeError: + return None + + +async def _wait_for_completion( + batch_client: client.BatchV1Api, job_name: str, deadline_secs: Optional[int] = None +) -> bool: + start_time = time.time() + while True: + job = await batch_client.read_namespaced_job_status(job_name, NAMESPACE) + if job.status.succeeded is not None and job.status.succeeded >= 1: + return True + elif job.status.failed is not None and job.status.failed >= 1: + wandb.termerror(f"{LOG_PREFIX}Build job {job.status.failed} failed {job}") + return False + wandb.termlog(f"{LOG_PREFIX}Waiting for build job to complete...") + if deadline_secs is not None and time.time() - start_time > deadline_secs: + return False + + await asyncio.sleep(5) + + +class KanikoBuilder(AbstractBuilder): + """Builds a docker image for a project using Kaniko.""" + + type = "kaniko" + + build_job_name: str + build_context_store: str + secret_name: Optional[str] + secret_key: Optional[str] + image: str + + def __init__( + self, + environment: AbstractEnvironment, + registry: AbstractRegistry, + build_job_name: str = "wandb-launch-container-build", + build_context_store: str = "", + secret_name: str = "", + secret_key: str = "", + image: str = "gcr.io/kaniko-project/executor:v1.11.0", + config: Optional[dict] = None, + ): + """Initialize a KanikoBuilder. + + Arguments: + environment (AbstractEnvironment): The environment to use. + registry (AbstractRegistry): The registry to use. + build_job_name (str, optional): The name of the build job. + build_context_store (str, optional): The name of the build context store. + secret_name (str, optional): The name of the secret to use for the registry. + secret_key (str, optional): The key of the secret to use for the registry. + verify (bool, optional): Whether to verify the functionality of the builder. + Defaults to True. + """ + self.environment = environment + self.registry = registry + self.build_job_name = build_job_name + self.build_context_store = build_context_store.rstrip("/") + self.secret_name = secret_name + self.secret_key = secret_key + self.image = image + self.kaniko_config = config or {} + + @classmethod + def from_config( + cls, + config: dict, + environment: AbstractEnvironment, + registry: AbstractRegistry, + verify: bool = True, + login: bool = True, + ) -> "AbstractBuilder": + """Create a KanikoBuilder from a config dict. + + Arguments: + config: A dict containing the builder config. Must contain a "type" key + with value "kaniko". + environment: The environment to use for the build. + registry: The registry to use for the build. + verify: Whether to verify the builder config. + + Returns: + A KanikoBuilder instance. + """ + if config.get("type") != "kaniko": + raise LaunchError( + "Builder config must include 'type':'kaniko' to create a KanikoBuilder." + ) + build_context_store = config.get("build-context-store", "") + if build_context_store is None: + if not PVC_MOUNT_PATH: + raise LaunchError( + "You must specify a build context store for kaniko builds. " + "You can set builder.build-context-store in your agent config " + "to a valid s3, gcs, or azure blog storage URI. Or, configure " + "a persistent volume claim through the agent helm chart: " + "https://github.com/wandb/helm-charts/tree/main/charts/launch-agent" + ) + build_job_name = config.get("build-job-name", "wandb-launch-container-build") + secret_name = config.get("secret-name", "") + secret_key = config.get("secret-key", "") + kaniko_image = config.get( + "kaniko-image", "gcr.io/kaniko-project/executor:v1.11.0" + ) + image_uri = config.get("destination") + if image_uri is not None: + registry = registry_from_uri(image_uri) + kaniko_config = config.get("kaniko-config", {}) + + return cls( + environment, + registry, + build_context_store=build_context_store, + build_job_name=build_job_name, + secret_name=secret_name, + secret_key=secret_key, + image=kaniko_image, + config=kaniko_config, + ) + + async def verify(self) -> None: + """Verify that the builder config is valid. + + Raises: + LaunchError: If the builder config is invalid. + """ + if self.build_context_store: + await self.environment.verify_storage_uri(self.build_context_store) + + def login(self) -> None: + """Login to the registry.""" + + async def _create_docker_ecr_config_map( + self, job_name: str, corev1_client: client.CoreV1Api, repository: str + ) -> None: + username, password = await self.registry.get_username_password() + encoded = base64.b64encode(f"{username}:{password}".encode()).decode("utf-8") + ecr_config_map = client.V1ConfigMap( + api_version="v1", + kind="ConfigMap", + metadata=client.V1ObjectMeta( + name=f"docker-config-{job_name}", + namespace=NAMESPACE, + ), + data={ + "config.json": json.dumps( + { + "auths": { + f"{await self.registry.get_repo_uri()}": {"auth": encoded} + } + } + ) + }, + immutable=True, + ) + await corev1_client.create_namespaced_config_map(NAMESPACE, ecr_config_map) + + async def _delete_docker_ecr_config_map( + self, job_name: str, client: client.CoreV1Api + ) -> None: + if self.secret_name: + await client.delete_namespaced_config_map( + f"docker-config-{job_name}", NAMESPACE + ) + + async def _upload_build_context(self, run_id: str, context_path: str) -> str: + # creat a tar archive of the build context and upload it to s3 + context_file = tempfile.NamedTemporaryFile(delete=False) + with tarfile.TarFile.open(fileobj=context_file, mode="w:gz") as context_tgz: + context_tgz.add(context_path, arcname=".") + context_file.close() + if PVC_MOUNT_PATH is None: + destination = f"{self.build_context_store}/{run_id}.tgz" + if self.environment is None: + raise LaunchError("No environment specified for Kaniko build.") + await self.environment.upload_file(context_file.name, destination) + return destination + else: + destination = f"{PVC_MOUNT_PATH}/{run_id}.tgz" + try: + shutil.copy(context_file.name, destination) + except Exception as e: + raise LaunchError( + f"Error copying build context to PVC mounted at {PVC_MOUNT_PATH}: {e}" + ) from e + return f"tar:///context/{run_id}.tgz" + + async def build_image( + self, + launch_project: LaunchProject, + entrypoint: EntryPoint, + job_tracker: Optional[JobAndRunStatusTracker] = None, + ) -> str: + await self.verify() + + build_contex_manager = BuildContextManager(launch_project=launch_project) + context_path, image_tag = build_contex_manager.create_build_context("kaniko") + run_id = launch_project.run_id + repo_uri = await self.registry.get_repo_uri() + image_uri = repo_uri + ":" + image_tag + + # 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( + kubernetes, launch_project.resource_args + ) + # TODO: use same client as kubernetes_runner.py + batch_v1 = client.BatchV1Api(api_client) + core_v1 = client.CoreV1Api(api_client) + + build_job_name = f"{self.build_job_name}-{run_id}" + + build_context = await self._upload_build_context(run_id, context_path) + build_job = await self._create_kaniko_job( + build_job_name, repo_uri, image_uri, build_context, core_v1, api_client + ) + wandb.termlog(f"{LOG_PREFIX}Created kaniko job {build_job_name}") + + try: + # 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}" + ), + data={ + "config.json": json.dumps( + { + "credHelpers": { + f"{self.registry.registry_name}.azurecr.io": "acr-env" + } + } + ) + }, + ) + await core_v1.create_namespaced_config_map( + "wandb", dockerfile_config_map + ) + if self.secret_name: + await self._create_docker_ecr_config_map( + build_job_name, core_v1, repo_uri + ) + k8s_job = await batch_v1.create_namespaced_job(NAMESPACE, build_job) + # wait for double the job deadline since it might take time to schedule + if not await _wait_for_completion( + batch_v1, build_job_name, 3 * _DEFAULT_BUILD_TIMEOUT_SECS + ): + if job_tracker: + job_tracker.set_err_stage("build") + msg = f"Failed to build image in kaniko for job {run_id}." + pod_name = get_pod_name_safe(k8s_job) + if pod_name: + msg += f" View logs with `kubectl logs -n {NAMESPACE} {pod_name}`." + raise Exception(msg) # noqa: TRY301 + try: + pods_from_job = await core_v1.list_namespaced_pod( + namespace=NAMESPACE, label_selector=f"job-name={build_job_name}" + ) + if len(pods_from_job.items) != 1: + raise Exception( # noqa: TRY301 + f"Expected 1 pod for job {build_job_name}," + f" found {len(pods_from_job.items)}" + ) + pod_name = pods_from_job.items[0].metadata.name + logs = await core_v1.read_namespaced_pod_log(pod_name, NAMESPACE) + warn_failed_packages_from_build_logs( + logs, image_uri, launch_project.api, job_tracker + ) + except Exception as e: + wandb.termwarn( + f"{LOG_PREFIX}Failed to get logs for kaniko job {build_job_name}: {e}" + ) + except Exception as e: + wandb.termerror( + f"{LOG_PREFIX}Exception when creating Kubernetes resources: {e}\n" + ) + raise + finally: + wandb.termlog(f"{LOG_PREFIX}Cleaning up resources") + try: + if ( + isinstance(self.registry, AzureContainerRegistry) + and not DOCKER_CONFIG_SECRET + ): + await core_v1.delete_namespaced_config_map( + f"docker-config-{build_job_name}", "wandb" + ) + if self.secret_name: + await self._delete_docker_ecr_config_map(build_job_name, core_v1) + await batch_v1.delete_namespaced_job(build_job_name, NAMESPACE) + except Exception as e: + traceback.print_exc() + raise LaunchError( + f"Exception during Kubernetes resource clean up {e}" + ) from e + return image_uri + + async def _create_kaniko_job( + self, + job_name: str, + repository: str, + image_tag: str, + build_context_path: str, + core_client: client.CoreV1Api, + api_client, + ) -> Dict[str, Any]: + job = copy.deepcopy(self.kaniko_config) + job_metadata = job.get("metadata", {}) + job_labels = job_metadata.get("labels", {}) + job_spec = job.get("spec", {}) + pod_template = job_spec.get("template", {}) + pod_metadata = pod_template.get("metadata", {}) + pod_labels = pod_metadata.get("labels", {}) + pod_spec = pod_template.get("spec", {}) + volumes = pod_spec.get("volumes", []) + containers = pod_spec.get("containers") or [{}] + if len(containers) > 1: + raise LaunchError( + "Multiple container configs not supported for kaniko builder." + ) + container = containers[0] + volume_mounts = container.get("volumeMounts", []) + env = container.get("env", []) + custom_args = container.get("args", []) + + if PVC_MOUNT_PATH: + volumes.append( + {"name": "kaniko-pvc", "persistentVolumeClaim": {"claimName": PVC_NAME}} + ) + volume_mounts.append({"name": "kaniko-pvc", "mountPath": "/context"}) + + if bool(self.secret_name) != bool(self.secret_key): + raise LaunchError( + "Both secret_name and secret_key or neither must be specified " + "for kaniko build. You provided only one of them." + ) + if isinstance(self.registry, ElasticContainerRegistry): + env.append( + { + "name": "AWS_REGION", + "value": self.registry.region, + } + ) + # TODO(ben): Refactor all of this environment/registry + # specific stuff into methods of those classes. + if isinstance(self.environment, AzureEnvironment): + # Use the core api to check if the secret exists + try: + await core_client.read_namespaced_secret( + "azure-storage-access-key", + "wandb", + ) + except Exception as e: + raise LaunchError( + "Secret azure-storage-access-key does not exist in " + "namespace wandb. Please create it with the key password " + "set to your azure storage access key." + ) from e + env.append( + { + "name": "AZURE_STORAGE_ACCESS_KEY", + "valueFrom": { + "secretKeyRef": { + "name": "azure-storage-access-key", + "key": "password", + } + }, + } + ) + if DOCKER_CONFIG_SECRET: + volumes.append( + { + "name": "kaniko-docker-config", + "secret": { + "secretName": DOCKER_CONFIG_SECRET, + "items": [ + { + "key": ".dockerconfigjson", + "path": "config.json", + } + ], + }, + } + ) + volume_mounts.append( + {"name": "kaniko-docker-config", "mountPath": "/kaniko/.docker"} + ) + elif self.secret_name and self.secret_key: + volumes.append( + { + "name": "docker-config", + "configMap": {"name": f"docker-config-{job_name}"}, + } + ) + volume_mounts.append( + {"name": "docker-config", "mountPath": "/kaniko/.docker"} + ) + # TODO(ben): I don't like conditioning on the registry type here. As a + # future change I want the registry and environment classes to provide + # a list of environment variables and volume mounts that need to be + # added to the job. The environment class provides credentials for + # build context access, and the registry class provides credentials + # for pushing the image. This way we can have separate secrets for + # each and support build contexts and registries that require + # different credentials. + if isinstance(self.registry, ElasticContainerRegistry): + mount_path = "/root/.aws" + key = "credentials" + elif isinstance(self.registry, GoogleArtifactRegistry): + mount_path = "/kaniko/.config/gcloud" + key = "config.json" + env.append( + { + "name": "GOOGLE_APPLICATION_CREDENTIALS", + "value": "/kaniko/.config/gcloud/config.json", + } + ) + else: + wandb.termwarn( + f"{LOG_PREFIX}Automatic credential handling is not supported for registry type {type(self.registry)}. Build job: {self.build_job_name}" + ) + volumes.append( + { + "name": self.secret_name, + "secret": { + "secretName": self.secret_name, + "items": [{"key": self.secret_key, "path": key}], + }, + } + ) + volume_mounts.append( + { + "name": self.secret_name, + "mountPath": mount_path, + "readOnly": True, + } + ) + if ( + isinstance(self.registry, AzureContainerRegistry) + and not DOCKER_CONFIG_SECRET + ): + # Add the docker config map + volumes.append( + { + "name": "docker-config", + "configMap": {"name": f"docker-config-{job_name}"}, + } + ) + volume_mounts.append( + {"name": "docker-config", "mountPath": "/kaniko/.docker/"} + ) + # Kaniko doesn't want https:// at the beginning of the image tag. + destination = image_tag + if destination.startswith("https://"): + destination = destination.replace("https://", "") + args = { + "--context": build_context_path, + "--dockerfile": _WANDB_DOCKERFILE_NAME, + "--destination": destination, + "--cache": "true", + "--cache-repo": repository.replace("https://", ""), + "--snapshot-mode": "redo", + "--compressed-caching": "false", + } + for custom_arg in custom_args: + arg_name, arg_value = custom_arg.split("=", 1) + args[arg_name] = arg_value + parsed_args = [ + f"{arg_name}={arg_value}" for arg_name, arg_value in args.items() + ] + container["args"] = parsed_args + + # Apply the rest of our defaults + pod_labels["wandb"] = "launch" + # This annotation is required to enable azure workload identity. + # 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( + "activeDeadlineSeconds", _DEFAULT_BUILD_TIMEOUT_SECS + ) + pod_spec["serviceAccountName"] = pod_spec.get( + "serviceAccountName", SERVICE_ACCOUNT_NAME + ) + job_spec["backoffLimit"] = job_spec.get("backoffLimit", 0) + job_labels["wandb"] = "launch" + job_metadata["namespace"] = job_metadata.get("namespace", NAMESPACE) + job_metadata["name"] = job_metadata.get("name", job_name) + job["apiVersion"] = "batch/v1" + job["kind"] = "Job" + + # Apply all nested configs from the bottom up + pod_metadata["labels"] = pod_labels + pod_template["metadata"] = pod_metadata + container["name"] = container.get("name", "wandb-container-build") + container["image"] = container.get("image", self.image) + container["volumeMounts"] = volume_mounts + container["env"] = env + pod_spec["containers"] = [container] + pod_spec["volumes"] = volumes + pod_template["spec"] = pod_spec + job_spec["template"] = pod_template + job_metadata["labels"] = job_labels + job["metadata"] = job_metadata + job["spec"] = job_spec + return job diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/noop.py b/lib/python3.12/site-packages/wandb/sdk/launch/builder/noop.py new file mode 100644 index 0000000000000000000000000000000000000000..52a64cf5e17dab77bf1792a540421f99e22b5c55 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/builder/noop.py @@ -0,0 +1,58 @@ +"""NoOp builder implementation.""" + +from typing import Any, Dict, Optional + +from wandb.sdk.launch.builder.abstract import AbstractBuilder +from wandb.sdk.launch.environment.abstract import AbstractEnvironment +from wandb.sdk.launch.errors import LaunchError +from wandb.sdk.launch.registry.abstract import AbstractRegistry + +from .._project_spec import EntryPoint, LaunchProject +from ..agent.job_status_tracker import JobAndRunStatusTracker + + +class NoOpBuilder(AbstractBuilder): + """NoOp builder.""" + + type = "noop" + + def __init__( + self, + builder_config: Dict[str, Any], + environment: AbstractEnvironment, + registry: AbstractRegistry, + ) -> None: + """Initialize a NoOpBuilder.""" + self.environment = environment + self.registry = registry + + @classmethod + def from_config( + cls, + config: dict, + environment: AbstractEnvironment, + registry: AbstractRegistry, + verify: bool = True, + ) -> "AbstractBuilder": + """Create a noop builder from a config.""" + return cls(config, environment, registry) + + async def verify(self) -> None: + """Verify the builder.""" + raise LaunchError("Attempted to verify noop builder.") + + async def build_image( + self, + launch_project: LaunchProject, + entrypoint: EntryPoint, + job_tracker: Optional[JobAndRunStatusTracker] = None, + ) -> str: + """Build the image. + + For this we raise a launch error since it can't build. + """ + raise LaunchError( + "Attempted build with noop builder. Specify a builder in your launch config at ~/.config/wandb/launch-config.yaml.\n" + "Note: Jobs sourced from git repos and code artifacts require a builder, while jobs sourced from Docker images do not.\n" + "See https://docs.wandb.ai/guides/launch/create-job." + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/__pycache__/_wandb_bootstrap.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/__pycache__/_wandb_bootstrap.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f7c61e8ed99fbccef4bac0593add0a7c00cd87e9 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/__pycache__/_wandb_bootstrap.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/__pycache__/dockerfile.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/__pycache__/dockerfile.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aed020afaca1e6f767a85e499c155fa55ae27a28 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/__pycache__/dockerfile.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/_wandb_bootstrap.py b/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/_wandb_bootstrap.py new file mode 100644 index 0000000000000000000000000000000000000000..81e9197775943aa057f74cf2e07a876f1a9452e4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/_wandb_bootstrap.py @@ -0,0 +1,188 @@ +import json +import os +import re +import subprocess +import sys +from typing import List, Optional, Set + +FAILED_PACKAGES_PREFIX = "ERROR: Failed to install: " +FAILED_PACKAGES_POSTFIX = ". During automated build process." +ONLY_INCLUDE = {x for x in os.getenv("WANDB_ONLY_INCLUDE", "").split(",") if x != ""} +OPTS = [] +# If the builder doesn't support buildx no need to use the cache +if os.getenv("WANDB_DISABLE_CACHE"): + OPTS.append("--no-cache-dir") +# When installing all packages from requirements.frozen.txt no need to resolve deps +if len(ONLY_INCLUDE) == 0: + OPTS.append("--no-deps") +# When installing the intersection of requirements.frozen.txt and requirements.txt +# force the frozen versions +else: + OPTS.append("--force") + +TORCH_DEP_REGEX = r"torch(vision|audio)?==\d+\.\d+\.\d+(\+(?:cu[\d]{2,3})|(?:\+cpu))?" + + +def install_deps( + deps: List[str], + failed: Optional[Set[str]] = None, + extra_index: Optional[str] = None, + opts: Optional[List[str]] = None, +) -> Optional[Set[str]]: + """Install pip dependencies. + + Arguments: + deps {List[str]} -- List of dependencies to install + failed (set, None): The libraries that failed to install + + Returns: + deps (str[], None): The dependencies that failed to install + """ + try: + subprocess.check_output(["pip", "install", "uv"], stderr=subprocess.STDOUT) + # Include only uri if @ is present + clean_deps = [d.split("@")[-1].strip() if "@" in d else d for d in deps] + index_args = ["--extra-index-url", extra_index] if extra_index else [] + print("installing {}...".format(", ".join(clean_deps))) + opts = opts or [] + args = ["uv", "pip", "install"] + opts + clean_deps + index_args + sys.stdout.flush() + subprocess.check_output(args, stderr=subprocess.STDOUT) + return failed + except subprocess.CalledProcessError as e: + if failed is None: + failed = set() + num_failed = len(failed) + current_pkg = None + for line in e.output.decode("utf8").splitlines(): + # Since the name of the package might not be on the same line as + # the error msg, keep track of the currently installing package + current_pkg = get_current_package(line, clean_deps, current_pkg) + + if "error: subprocess-exited-with-error" in line: + if current_pkg is not None: + failed.add(current_pkg) + elif line.startswith("ERROR:"): + clean_dep = find_package_in_error_string(clean_deps, line) + if clean_dep is not None: + if clean_dep in deps: + failed.add(clean_dep) + else: + for d in deps: + if clean_dep in d: + failed.add(d.replace(" ", "")) + break + if len(set(clean_deps) - failed) == 0: + return failed + elif len(failed) > num_failed: + return install_deps( + list(set(clean_deps) - failed), + failed, + extra_index=extra_index, + opts=opts, + ) + else: + return failed + + +def main() -> None: + """Install deps in requirements.frozen.txt.""" + extra_index = None + torch_reqs = [] + if os.path.exists("requirements.frozen.txt"): + with open("requirements.frozen.txt") as f: + print("Installing frozen dependencies...") + reqs = [] + for req in f: + if ( + len(ONLY_INCLUDE) == 0 + or req in ONLY_INCLUDE + or req.split("=")[0].lower() in ONLY_INCLUDE + ): + # can't pip install wandb==0.*.*.dev1 through pip. Lets just install wandb for now + if req.startswith("wandb==") and "dev1" in req: + req = "wandb" + match = re.match( + TORCH_DEP_REGEX, + req, + ) + if match: + variant = match.group(2) + if variant: + extra_index = ( + f"https://download.pytorch.org/whl/{variant[1:]}" + ) + torch_reqs.append(req.strip().replace(" ", "")) + else: + reqs.append(req.strip().replace(" ", "")) + else: + print(f"Ignoring requirement: {req} from frozen requirements") + failed = install_deps(reqs, opts=OPTS) or set() + with open("_wandb_bootstrap_errors.json", "w") as f: + f.write(json.dumps({"pip": list(failed)})) + if len(failed) > 0: + sys.stderr.write( + FAILED_PACKAGES_PREFIX + ",".join(failed) + FAILED_PACKAGES_POSTFIX + ) + sys.stderr.flush() + install_deps(torch_reqs, extra_index=extra_index) + else: + print("No frozen requirements found") + + +def add_version_to_package_name(deps: List[str], package: str) -> Optional[str]: + """Add the associated version to a package name. + + For example: `my-package` -> `my-package==1.0.0` + """ + for dep in deps: + if dep.split("==")[0] == package: + return dep + return None + + +def get_current_package( + line: str, deps: List[str], current_pkg: Optional[str] +) -> Optional[str]: + """Tries to pull a package name from the line. + + Used to keep track of what the currently-installing package is, + in case an error message isn't on the same line as the package + """ + # "Collecting my-package==1.0.0" + if line.startswith("Collecting"): + return line.split(" ")[1] + # "Building wheel for my-package (pyproject.toml): finished with status 'error'" + elif line.strip().startswith("Building wheel") and line.strip().endswith( + "finished with status 'error'" + ): + return add_version_to_package_name(deps, line.strip().split(" ")[3]) + # "Running setup.py install for my-package: finished with status 'error'" + elif line.strip().startswith("Running setup.py install") and line.strip().endswith( + "finished with status 'error'" + ): + return add_version_to_package_name(deps, line.strip().split(" ")[4][:-1]) + return current_pkg + + +# hacky way to get the name of the requirement that failed +# attempt last word which is the name of the package often +# fall back to checking all words in the line for the package name +def find_package_in_error_string(deps: List[str], line: str) -> Optional[str]: + # if the last word in the error string is in the list of deps, return it + last_word = line.split(" ")[-1] + if last_word in deps: + return last_word + # if the last word is not in the list of deps, check all words + # TODO: this could report the wrong package if the error string + # contains a reference to another package in the deps + # before the package that failed to install + for word in line.split(" "): + if word.strip(",") in deps: + return word + # if we can't find the package, return None + return None + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/dockerfile.py b/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/dockerfile.py new file mode 100644 index 0000000000000000000000000000000000000000..54ac5bc62201946541541cee6c235391fd9c2d14 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/builder/templates/dockerfile.py @@ -0,0 +1,92 @@ +DOCKERFILE_TEMPLATE = """ +# ----- stage 1: build ----- +FROM {py_build_image} as build + +# requirements section depends on pip vs conda, and presence of buildx +ENV PIP_PROGRESS_BAR off +{requirements_section} + +# ----- stage 2: base ----- +{base_setup} + +COPY --from=build /env /env +ENV PATH="/env/bin:$PATH" + +ENV SHELL /bin/bash + +# some resources (eg sagemaker) must run on root +{user_setup} + +WORKDIR {workdir} +RUN chown -R {uid} {workdir} + +# make artifacts cache dir unrelated to build +RUN mkdir -p {workdir}/.cache && chown -R {uid} {workdir}/.cache + +# copy code/etc +COPY --chown={uid} src/ {workdir} + +ENV PYTHONUNBUFFERED=1 + +{entrypoint_section} +""" + +# this goes into base_setup in TEMPLATE +PYTHON_SETUP_TEMPLATE = """ +FROM {py_base_image} as base +""" + +# this goes into base_setup in TEMPLATE +ACCELERATOR_SETUP_TEMPLATE = """ +FROM {accelerator_base_image} as base + +# make non-interactive so build doesn't block on questions +ENV DEBIAN_FRONTEND=noninteractive + +# install python +RUN apt-get update -qq && apt-get install --no-install-recommends -y \ + {python_packages} \ + && apt-get -qq purge && apt-get -qq clean \ + && rm -rf /var/lib/apt/lists/* + +# make sure `python` points at the right version +RUN update-alternatives --install /usr/bin/python python /usr/bin/python{py_version} 1 \ + && update-alternatives --install /usr/local/bin/python python /usr/bin/python{py_version} 1 +""" + +# this goes into requirements_section in TEMPLATE +PIP_TEMPLATE = """ +RUN python -m venv /env +# make sure we install into the env +ENV PATH="/env/bin:$PATH" + +COPY {requirements_files} ./ +{buildx_optional_prefix} {pip_install} +""" + +# this goes into requirements_section in TEMPLATE +CONDA_TEMPLATE = """ +COPY src/environment.yml . +{buildx_optional_prefix} conda env create -f environment.yml -n env + +# pack the environment so that we can transfer to the base image +RUN conda install -c conda-forge conda-pack +RUN conda pack -n env -o /tmp/env.tar && \ + mkdir /env && cd /env && tar xf /tmp/env.tar && \ + rm /tmp/env.tar +RUN /env/bin/conda-unpack +""" + +USER_CREATE_TEMPLATE = """ +RUN useradd \ + --create-home \ + --no-log-init \ + --shell /bin/bash \ + --gid 0 \ + --uid {uid} \ + {user} || echo "" +""" + +ENTRYPOINT_TEMPLATE = """ +ENTRYPOINT {entrypoint} +""" diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/abstract.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/abstract.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..308b2948cfb30cf42a94e8074918e912c277f265 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/abstract.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/aws_environment.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/aws_environment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..192a2b5ecc8bf0afb350cc388b5b6af3e87ed544 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/aws_environment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/azure_environment.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/azure_environment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54e70d6d8e3896b423e91a9c5d78ac1f4f604e02 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/azure_environment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/gcp_environment.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/gcp_environment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e522537c463b95f4f7af0e67fb39afef0a0befac Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/gcp_environment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/local_environment.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/local_environment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..158c48979316f7b7c0ab30a14be222d8638c3d17 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/environment/__pycache__/local_environment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/environment/abstract.py b/lib/python3.12/site-packages/wandb/sdk/launch/environment/abstract.py new file mode 100644 index 0000000000000000000000000000000000000000..a736ef0404018f0e640a6dad397de281a77895b2 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/environment/abstract.py @@ -0,0 +1,29 @@ +"""Abstract base class for environments.""" + +from abc import ABC, abstractmethod + + +class AbstractEnvironment(ABC): + """Abstract base class for environments.""" + + region: str + + @abstractmethod + async def verify(self) -> None: + """Verify that the environment is configured correctly.""" + raise NotImplementedError + + @abstractmethod + async def upload_file(self, source: str, destination: str) -> None: + """Upload a file from the local filesystem to storage in the environment.""" + raise NotImplementedError + + @abstractmethod + async def upload_dir(self, source: str, destination: str) -> None: + """Upload the contents of a directory from the local filesystem to the environment.""" + raise NotImplementedError + + @abstractmethod + async def verify_storage_uri(self, uri: str) -> None: + """Verify that the storage URI is configured correctly.""" + raise NotImplementedError diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/environment/aws_environment.py b/lib/python3.12/site-packages/wandb/sdk/launch/environment/aws_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..4ca54b6d1ac269916fcd98a7c0a4b6fa25465b45 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/environment/aws_environment.py @@ -0,0 +1,322 @@ +"""Implements the AWS environment.""" + +import logging +import os +from typing import Dict, Optional + +from wandb.sdk.launch.errors import LaunchError +from wandb.util import get_module + +from ..utils import ARN_PARTITION_RE, S3_URI_RE, event_loop_thread_exec +from .abstract import AbstractEnvironment + +boto3 = get_module( + "boto3", + required="AWS environment requires boto3 to be installed. Please install " + "it with `pip install wandb[launch]`.", +) +botocore = get_module( + "botocore", + required="AWS environment requires botocore to be installed. Please install " + "it with `pip install wandb[launch]`.", +) + +_logger = logging.getLogger(__name__) + + +class AwsEnvironment(AbstractEnvironment): + """AWS environment.""" + + def __init__( + self, + region: str, + access_key: str, + secret_key: str, + session_token: str, + ) -> None: + """Initialize the AWS environment. + + Arguments: + region (str): The AWS region. + + Raises: + LaunchError: If the AWS environment is not configured correctly. + """ + super().__init__() + _logger.info(f"Initializing AWS environment in region {region}.") + self._region = region + self._access_key = access_key + self._secret_key = secret_key + self._session_token = session_token + self._account = None + self._partition = None + + @classmethod + def from_default(cls, region: Optional[str] = None) -> "AwsEnvironment": + """Create an AWS environment from the default AWS environment. + + Arguments: + region (str, optional): The AWS region. + verify (bool, optional): Whether to verify the AWS environment. Defaults to True. + + Returns: + AwsEnvironment: The AWS environment. + """ + _logger.info("Creating AWS environment from default credentials.") + try: + session = boto3.Session() + if hasattr(session, "region"): + region = region or session.region + region = region or os.environ.get("AWS_REGION") + credentials = session.get_credentials() + if not credentials: + raise LaunchError( + "Could not create AWS environment from default environment. Please verify that your AWS credentials are configured correctly." + ) + access_key = credentials.access_key + secret_key = credentials.secret_key + session_token = credentials.token + except botocore.client.ClientError as e: + raise LaunchError( + f"Could not create AWS environment from default environment. Please verify that your AWS credentials are configured correctly. {e}" + ) + if not region: + raise LaunchError( + "Could not create AWS environment from default environment. Region not specified." + ) + return cls( + region=region, + access_key=access_key, + secret_key=secret_key, + session_token=session_token, + ) + + @classmethod + def from_config( + cls, + config: Dict[str, str], + ) -> "AwsEnvironment": + """Create an AWS environment from the default AWS environment. + + Arguments: + config (dict): Configuration dictionary. + verify (bool, optional): Whether to verify the AWS environment. Defaults to True. + + Returns: + AwsEnvironment: The AWS environment. + """ + region = str(config.get("region", "")) + if not region: + raise LaunchError( + "Could not create AWS environment from config. Region not specified." + ) + return cls.from_default( + region=region, + ) + + @property + def region(self) -> str: + """The AWS region.""" + return self._region + + @region.setter + def region(self, region: str) -> None: + self._region = region + + async def get_partition(self) -> str: + """Set the partition for the AWS environment.""" + try: + session = await self.get_session() + client = await event_loop_thread_exec(session.client)("sts") + get_caller_identity = event_loop_thread_exec(client.get_caller_identity) + identity = await get_caller_identity() + arn = identity.get("Arn") + if not arn: + raise LaunchError( + "Could not set partition for AWS environment. ARN not found." + ) + matched_partition = ARN_PARTITION_RE.match(arn) + if not matched_partition: + raise LaunchError( + f"Could not set partition for AWS environment. ARN {arn} is not valid." + ) + partition = matched_partition.group(1) + return partition + except botocore.exceptions.ClientError as e: + raise LaunchError( + f"Could not set partition for AWS environment. {e}" + ) from e + + async def verify(self) -> None: + """Verify that the AWS environment is configured correctly. + + Raises: + LaunchError: If the AWS environment is not configured correctly. + """ + _logger.debug("Verifying AWS environment.") + try: + session = await self.get_session() + client = await event_loop_thread_exec(session.client)("sts") + get_caller_identity = event_loop_thread_exec(client.get_caller_identity) + self._account = (await get_caller_identity()).get("Account") + # TODO: log identity details from the response + except botocore.exceptions.ClientError as e: + raise LaunchError( + f"Could not verify AWS environment. Please verify that your AWS credentials are configured correctly. {e}" + ) from e + + async def get_session(self) -> "boto3.Session": # type: ignore + """Get an AWS session. + + Returns: + boto3.Session: The AWS session. + + Raises: + LaunchError: If the AWS session could not be created. + """ + _logger.debug(f"Creating AWS session in region {self._region}") + try: + session = event_loop_thread_exec(boto3.Session) + return await session( + region_name=self._region, + aws_access_key_id=self._access_key, + aws_secret_access_key=self._secret_key, + aws_session_token=self._session_token, + ) + except botocore.exceptions.ClientError as e: + raise LaunchError(f"Could not create AWS session. {e}") + + async def upload_file(self, source: str, destination: str) -> None: + """Upload a file to s3 from local storage. + + The destination is a valid s3 URI, e.g. s3://bucket/key and will + be used as a prefix for the uploaded file. Only the filename of the source + is kept in the upload key. So if the source is "foo/bar" and the + destination is "s3://bucket/key", the file "foo/bar" will be uploaded + to "s3://bucket/key/bar". + + Arguments: + source (str): The path to the file or directory. + destination (str): The uri of the storage destination. This should + be a valid s3 URI, e.g. s3://bucket/key. + + Raises: + LaunchError: If the copy fails, the source path does not exist, or the + destination is not a valid s3 URI, or the upload fails. + """ + _logger.debug(f"Uploading {source} to {destination}") + _err_prefix = f"Error attempting to copy {source} to {destination}." + if not os.path.isfile(source): + raise LaunchError(f"{_err_prefix}: Source {source} does not exist.") + match = S3_URI_RE.match(destination) + if not match: + raise LaunchError( + f"{_err_prefix}: Destination {destination} is not a valid s3 URI." + ) + bucket = match.group(1) + key = match.group(2).lstrip("/") + if not key: + key = "" + session = await self.get_session() + try: + client = await event_loop_thread_exec(session.client)("s3") + client.upload_file(source, bucket, key) + except botocore.exceptions.ClientError as e: + raise LaunchError( + f"{_err_prefix}: botocore error attempting to copy {source} to {destination}. {e}" + ) + + async def upload_dir(self, source: str, destination: str) -> None: + """Upload a directory to s3 from local storage. + + The upload will place the contents of the source directory in the destination + with the same directory structure. So if the source is "foo/bar" and the + destination is "s3://bucket/key", the contents of "foo/bar" will be uploaded + to "s3://bucket/key/bar". + + Arguments: + source (str): The path to the file or directory. + destination (str): The URI of the storage. + recursive (bool, optional): If True, copy the directory recursively. Defaults to False. + + Raises: + LaunchError: If the copy fails, the source path does not exist, or the + destination is not a valid s3 URI. + """ + _logger.debug(f"Uploading {source} to {destination}") + _err_prefix = f"Error attempting to copy {source} to {destination}." + if not os.path.isdir(source): + raise LaunchError(f"{_err_prefix}: Source {source} does not exist.") + match = S3_URI_RE.match(destination) + if not match: + raise LaunchError( + f"{_err_prefix}: Destination {destination} is not a valid s3 URI." + ) + bucket = match.group(1) + key = match.group(2).lstrip("/") + if not key: + key = "" + session = await self.get_session() + try: + client = await event_loop_thread_exec(session.client)("s3") + for path, _, files in os.walk(source): + for file in files: + abs_path = os.path.join(path, file) + key_path = ( + abs_path.replace(source, "").replace("\\", "/").lstrip("/") + ) + client.upload_file( + abs_path, + bucket, + key_path, + ) + except botocore.exceptions.ClientError as e: + raise LaunchError( + f"{_err_prefix}: botocore error attempting to copy {source} to {destination}. {e}" + ) from e + except Exception as e: + raise LaunchError( + f"{_err_prefix}: Unexpected error attempting to copy {source} to {destination}. {e}" + ) from e + + async def verify_storage_uri(self, uri: str) -> None: + """Verify that s3 storage is configured correctly. + + This will check that the bucket exists and that the credentials are + configured correctly. + + Arguments: + uri (str): The URI of the storage. + + Raises: + LaunchError: If the storage is not configured correctly or the URI is + not a valid s3 URI. + + Returns: + None + """ + _logger.debug(f"Verifying storage {uri}") + match = S3_URI_RE.match(uri) + if not match: + raise LaunchError( + f"Failed to validate storage uri: {uri} is not a valid s3 URI." + ) + bucket = match.group(1) + try: + session = await self.get_session() + client = await event_loop_thread_exec(session.client)("s3") + client.head_bucket(Bucket=bucket) + except botocore.exceptions.ClientError as e: + if e.response["Error"]["Code"] == "404": + raise LaunchError( + f"Could not verify AWS storage uri {uri}. Bucket {bucket} does not exist." + ) + if e.response["Error"]["Code"] == "403": + raise LaunchError( + f"Could not verify AWS storage uri {uri}. " + "Bucket {bucket} is not accessible. Please check that this " + "client is authenticated with permission to access the bucket." + ) + raise LaunchError( + f"Failed to verify AWS storage uri {uri}. Response: {e.response} Please verify that your AWS credentials are configured correctly." + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/environment/azure_environment.py b/lib/python3.12/site-packages/wandb/sdk/launch/environment/azure_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..2dbfebbe14d499302b4c22dd2321cdfb4ff8df63 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/environment/azure_environment.py @@ -0,0 +1,105 @@ +"""Implementation of AzureEnvironment class.""" + +from typing import Tuple + +from azure.core.exceptions import HttpResponseError # type: ignore +from azure.identity import DefaultAzureCredential # type: ignore +from azure.storage.blob import BlobClient, BlobServiceClient # type: ignore + +from ..errors import LaunchError +from ..utils import AZURE_BLOB_REGEX +from .abstract import AbstractEnvironment + + +class AzureEnvironment(AbstractEnvironment): + """AzureEnvironment is a helper for accessing Azure resources.""" + + def __init__( + self, + ) -> None: + """Initialize an AzureEnvironment.""" + + @classmethod + def from_config(cls, config: dict, verify: bool = True) -> "AzureEnvironment": + """Create an AzureEnvironment from a config dict.""" + return cls() + + @classmethod + def get_credentials(cls) -> DefaultAzureCredential: + """Get Azure credentials.""" + try: + return DefaultAzureCredential() + except Exception as e: + raise LaunchError( + f"Could not get Azure credentials. Please make sure you have " + f"configured your Azure CLI correctly.\n{e}" + ) from e + + async def upload_file(self, source: str, destination: str) -> None: + """Upload a file to Azure blob storage. + + Arguments: + source (str): The path to the file to upload. + destination (str): The destination path in Azure blob storage. Ex: + https://.blob.core.windows.net// + Raise: + LaunchError: If the file could not be uploaded. + """ + storage_account, storage_container, path = self.parse_uri(destination) + _err_prefix = f"Could not upload file {source} to Azure blob {destination}" + creds = self.get_credentials() + try: + client = BlobClient( + f"https://{storage_account}.blob.core.windows.net", + storage_container, + path, + credential=creds, + ) + with open(source, "rb") as f: + client.upload_blob(f, overwrite=True) + except HttpResponseError as e: + raise LaunchError(f"{_err_prefix}: {e.message}") from e + except Exception as e: + raise LaunchError(f"{_err_prefix}: {e.__class__.__name__}: {e}") from e + + async def upload_dir(self, source: str, destination: str) -> None: + """Upload a directory to Azure blob storage.""" + raise NotImplementedError() + + async def verify_storage_uri(self, uri: str) -> None: + """Verify that the given blob storage prefix exists. + + Args: + uri (str): The URI to verify. + """ + creds = self.get_credentials() + storage_account, storage_container, _ = self.parse_uri(uri) + try: + client = BlobServiceClient( + f"https://{storage_account}.blob.core.windows.net", + credential=creds, + ) + client.get_container_client(storage_container) + except Exception as e: + raise LaunchError( + f"Could not verify storage URI {uri} in container {storage_container}." + ) from e + + async def verify(self) -> None: + """Verify that the AzureEnvironment is valid.""" + self.get_credentials() + + @staticmethod + def parse_uri(uri: str) -> Tuple[str, str, str]: + """Parse an Azure blob storage URI into a storage account and container. + + Args: + uri (str): The URI to parse. + + Returns: + Tuple[str, str, prefix]: The storage account, container, and path. + """ + match = AZURE_BLOB_REGEX.match(uri) + if match is None: + raise LaunchError(f"Could not parse Azure blob URI {uri}.") + return match.group(1), match.group(2), match.group(3) diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/environment/gcp_environment.py b/lib/python3.12/site-packages/wandb/sdk/launch/environment/gcp_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..c9847a9853cb76830d79b4ccf0ca41bdc7ae83db --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/environment/gcp_environment.py @@ -0,0 +1,334 @@ +"""Implementation of the GCP environment for wandb launch.""" + +import logging +import os +import subprocess +from typing import Optional + +from wandb.sdk.launch.errors import LaunchError +from wandb.util import get_module + +from ..utils import GCS_URI_RE, event_loop_thread_exec +from .abstract import AbstractEnvironment + +google = get_module( + "google", + required="Google Cloud Platform support requires the google package. Please" + " install it with `pip install wandb[launch]`.", +) +google.cloud.compute_v1 = get_module( + "google.cloud.compute_v1", + required="Google Cloud Platform support requires the google-cloud-compute package. " + "Please install it with `pip install wandb[launch]`.", +) +google.auth.credentials = get_module( + "google.auth.credentials", + required="Google Cloud Platform support requires google-auth. " + "Please install it with `pip install wandb[launch]`.", +) +google.auth.transport.requests = get_module( + "google.auth.transport.requests", + required="Google Cloud Platform support requires google-auth. " + "Please install it with `pip install wandb[launch]`.", +) +google.api_core.exceptions = get_module( + "google.api_core.exceptions", + required="Google Cloud Platform support requires google-api-core. " + "Please install it with `pip install wandb[launch]`.", +) +google.cloud.storage = get_module( + "google.cloud.storage", + required="Google Cloud Platform support requires google-cloud-storage. " + "Please install it with `pip install wandb[launch].", +) + + +_logger = logging.getLogger(__name__) + +GCP_REGION_ENV_VAR = "GOOGLE_CLOUD_REGION" + + +class GcpEnvironment(AbstractEnvironment): + """GCP Environment. + + Attributes: + region: The GCP region. + """ + + region: str + + def __init__( + self, + region: str, + ) -> None: + """Initialize the GCP environment. + + Arguments: + region: The GCP region. + verify: Whether to verify the credentials, region, and project. + + Raises: + LaunchError: If verify is True and the environment is not properly + configured. + """ + super().__init__() + _logger.info(f"Initializing GcpEnvironment in region {region}") + self.region: str = region + self._project = "" + + @classmethod + def from_config(cls, config: dict) -> "GcpEnvironment": + """Create a GcpEnvironment from a config dictionary. + + Arguments: + config: The config dictionary. + + Returns: + GcpEnvironment: The GcpEnvironment. + """ + if config.get("type") != "gcp": + raise LaunchError( + f"Could not create GcpEnvironment from config. Expected type 'gcp' " + f"but got '{config.get('type')}'." + ) + region = config.get("region", None) + if not region: + raise LaunchError( + "Could not create GcpEnvironment from config. Missing 'region' field." + ) + return cls(region=region) + + @classmethod + def from_default( + cls, + ) -> "GcpEnvironment": + """Create a GcpEnvironment from the default configuration. + + Returns: + GcpEnvironment: The GcpEnvironment. + """ + region = get_default_region() + if region is None: + raise LaunchError( + "Could not create GcpEnvironment from user's gcloud configuration. " + "Please set the default region with `gcloud config set compute/region` " + "or set the environment variable {GCP_REGION_ENV_VAR}. " + "Alternatively, you may specify the region explicitly in your " + "wandb launch configuration at `$HOME/.config/wandb/launch-config.yaml`. " + "See https://docs.wandb.ai/guides/launch/run-agent#environments for more information." + ) + return cls(region=region) + + @property + def project(self) -> str: + """Get the name of the gcp project associated with the credentials. + + Returns: + str: The name of the gcp project. + + Raises: + LaunchError: If the launch environment cannot be verified. + """ + return self._project + + async def get_credentials(self) -> google.auth.credentials.Credentials: # type: ignore + """Get the GCP credentials. + + Uses google.auth.default() to get the credentials. If the credentials + are invalid, this method will refresh them. If the credentials are + still invalid after refreshing, this method will raise an error. + + Returns: + google.auth.credentials.Credentials: The GCP credentials. + + Raises: + LaunchError: If the GCP credentials are invalid. + """ + _logger.debug("Getting GCP credentials") + # TODO: Figure out a minimal set of scopes. + try: + google_auth_default = event_loop_thread_exec(google.auth.default) + creds, project = await google_auth_default() + if not self._project: + self._project = project + _logger.debug("Refreshing GCP credentials") + await event_loop_thread_exec(creds.refresh)( + google.auth.transport.requests.Request() + ) + except google.auth.exceptions.DefaultCredentialsError as e: + raise LaunchError( + "No Google Cloud Platform credentials found. Please run " + "`gcloud auth application-default login` or set the environment " + "variable GOOGLE_APPLICATION_CREDENTIALS to the path of a valid " + "service account key file." + ) from e + except google.auth.exceptions.RefreshError as e: + raise LaunchError( + "Could not refresh Google Cloud Platform credentials. Please run " + "`gcloud auth application-default login` or set the environment " + "variable GOOGLE_APPLICATION_CREDENTIALS to the path of a valid " + "service account key file." + ) from e + if not creds.valid: + raise LaunchError( + "Invalid Google Cloud Platform credentials. Please run " + "`gcloud auth application-default login` or set the environment " + "variable GOOGLE_APPLICATION_CREDENTIALS to the path of a valid " + "service account key file." + ) + return creds + + async def verify(self) -> None: + """Verify the credentials, region, and project. + + Credentials and region are verified by calling get_credentials(). The + region and is verified by calling the compute API. + + Raises: + LaunchError: If the credentials, region, or project are invalid. + + Returns: + None + """ + _logger.debug("Verifying GCP environment") + await self.get_credentials() + + async def verify_storage_uri(self, uri: str) -> None: + """Verify that a storage URI is valid. + + Arguments: + uri: The storage URI. + + Raises: + LaunchError: If the storage URI is invalid. + """ + match = GCS_URI_RE.match(uri) + if not match: + raise LaunchError(f"Invalid GCS URI: {uri}") + bucket = match.group(1) + cloud_storage_client = event_loop_thread_exec(google.cloud.storage.Client) + try: + credentials = await self.get_credentials() + storage_client = await cloud_storage_client(credentials=credentials) + bucket = await event_loop_thread_exec(storage_client.get_bucket)(bucket) + except google.api_core.exceptions.GoogleAPICallError as e: + raise LaunchError( + f"Failed verifying storage uri {uri}: bucket {bucket} does not exist." + ) from e + except google.api_core.exceptions.Forbidden as e: + raise LaunchError( + f"Failed verifying storage uri {uri}: bucket {bucket} is not accessible. Please check your permissions and try again." + ) from e + + async def upload_file(self, source: str, destination: str) -> None: + """Upload a file to GCS. + + Arguments: + source: The path to the local file. + destination: The path to the GCS file. + + Raises: + LaunchError: If the file cannot be uploaded. + """ + _logger.debug(f"Uploading file {source} to {destination}") + _err_prefix = f"Could not upload file {source} to GCS destination {destination}" + if not os.path.isfile(source): + raise LaunchError(f"{_err_prefix}: File {source} does not exist.") + match = GCS_URI_RE.match(destination) + if not match: + raise LaunchError(f"{_err_prefix}: Invalid GCS URI: {destination}") + bucket = match.group(1) + key = match.group(2).lstrip("/") + google_storage_client = event_loop_thread_exec(google.cloud.storage.Client) + credentials = await self.get_credentials() + try: + storage_client = await google_storage_client(credentials=credentials) + bucket = await event_loop_thread_exec(storage_client.bucket)(bucket) + blob = await event_loop_thread_exec(bucket.blob)(key) + await event_loop_thread_exec(blob.upload_from_filename)(source) + except google.api_core.exceptions.GoogleAPICallError as e: + resp = e.response + assert resp is not None + try: + message = resp.json()["error"]["message"] + except Exception: + message = str(resp) + raise LaunchError(f"{_err_prefix}: {message}") from e + + async def upload_dir(self, source: str, destination: str) -> None: + """Upload a directory to GCS. + + Arguments: + source: The path to the local directory. + destination: The path to the GCS directory. + + Raises: + LaunchError: If the directory cannot be uploaded. + """ + _logger.debug(f"Uploading directory {source} to {destination}") + _err_prefix = ( + f"Could not upload directory {source} to GCS destination {destination}" + ) + if not os.path.isdir(source): + raise LaunchError(f"{_err_prefix}: Directory {source} does not exist.") + match = GCS_URI_RE.match(destination) + if not match: + raise LaunchError(f"{_err_prefix}: Invalid GCS URI: {destination}") + bucket = match.group(1) + key = match.group(2).lstrip("/") + google_storage_client = event_loop_thread_exec(google.cloud.storage.Client) + credentials = await self.get_credentials() + try: + storage_client = await google_storage_client(credentials=credentials) + bucket = await event_loop_thread_exec(storage_client.bucket)(bucket) + for root, _, files in os.walk(source): + for file in files: + local_path = os.path.join(root, file) + gcs_path = os.path.join( + key, os.path.relpath(local_path, source) + ).replace("\\", "/") + blob = await event_loop_thread_exec(bucket.blob)(gcs_path) + await event_loop_thread_exec(blob.upload_from_filename)(local_path) + except google.api_core.exceptions.GoogleAPICallError as e: + resp = e.response + assert resp is not None + try: + message = resp.json()["error"]["message"] + except Exception: + message = str(resp) + raise LaunchError(f"{_err_prefix}: {message}") from e + except Exception as e: + raise LaunchError(f"{_err_prefix}: GCS upload failed: {e}") from e + + +def get_gcloud_config_value(config_name: str) -> Optional[str]: + """Get a value from gcloud config. + + Arguments: + config_name: The name of the config value. + + Returns: + str: The config value, or None if the value is not set. + """ + try: + output = subprocess.check_output( + ["gcloud", "config", "get-value", config_name], stderr=subprocess.STDOUT + ) + value = str(output.decode("utf-8").strip()) + if value and "unset" not in value: + return value + return None + except subprocess.CalledProcessError: + return None + + +def get_default_region() -> Optional[str]: + """Get the default region from gcloud config or environment variables. + + Returns: + str: The default region, or None if it cannot be determined. + """ + region = get_gcloud_config_value("compute/region") + if not region: + region = os.environ.get(GCP_REGION_ENV_VAR) + return region diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/environment/local_environment.py b/lib/python3.12/site-packages/wandb/sdk/launch/environment/local_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..fe726354de7cead29d3f663d25b300014aac6380 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/environment/local_environment.py @@ -0,0 +1,65 @@ +"""Dummy local environment implementation. This is the default environment.""" + +from typing import Any, Dict, Union + +from wandb.sdk.launch.errors import LaunchError + +from .abstract import AbstractEnvironment + + +class LocalEnvironment(AbstractEnvironment): + """Local environment class.""" + + def __init__(self) -> None: + """Initialize a local environment by doing nothing.""" + + @classmethod + def from_config( + cls, config: Dict[str, Union[Dict[str, Any], str]] + ) -> "LocalEnvironment": + """Create a local environment from a config. + + Arguments: + config (dict): The config. This is ignored. + + Returns: + LocalEnvironment: The local environment. + """ + return cls() + + async def verify(self) -> None: + """Verify that the local environment is configured correctly.""" + raise LaunchError("Attempted to verify LocalEnvironment.") + + async def verify_storage_uri(self, uri: str) -> None: + """Verify that the storage URI is configured correctly. + + Arguments: + uri (str): The storage URI. This is ignored. + """ + raise LaunchError("Attempted to verify storage uri for LocalEnvironment.") + + async def upload_file(self, source: str, destination: str) -> None: + """Upload a file from the local filesystem to storage in the environment. + + Arguments: + source (str): The source file. This is ignored. + destination (str): The destination file. This is ignored. + """ + raise LaunchError("Attempted to upload file for LocalEnvironment.") + + async def upload_dir(self, source: str, destination: str) -> None: + """Upload the contents of a directory from the local filesystem to the environment. + + Arguments: + source (str): The source directory. This is ignored. + destination (str): The destination directory. This is ignored. + """ + raise LaunchError("Attempted to upload directory for LocalEnvironment.") + + async def get_project(self) -> str: + """Get the project of the local environment. + + Returns: An empty string. + """ + raise LaunchError("Attempted to get project for LocalEnvironment.") diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/inputs/__pycache__/files.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/inputs/__pycache__/files.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e5416b31ce872aaeca5b61f6723768be2b611139 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/inputs/__pycache__/files.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/__init__.py b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3bc81d7f8c3fa6c3509df83c689e69369aa63888 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/abstract.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/abstract.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c71ac8950d6213a8a36b68d710a09dec7cd2d33 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/abstract.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/kubernetes_monitor.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/kubernetes_monitor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb2cae42c3fba34e21eacf0367990be07c15eb99 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/kubernetes_monitor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/kubernetes_runner.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/kubernetes_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62305adf918a79a32bc81a1f893767949cc0591e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/kubernetes_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/local_container.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/local_container.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46225bfedfc4225d30c841a66e29d97010978abb Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/local_container.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/local_process.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/local_process.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..950eaa0f5ff9b0b619c7ae9a4dfbf34b56ad327c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/local_process.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/sagemaker_runner.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/sagemaker_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e206e20742f5d90bcdd3d6d0e015cee99f2a176c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/sagemaker_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/vertex_runner.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/vertex_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e5519e0b5f9d81ce0404e571811610206678913f Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/runner/__pycache__/vertex_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/abstract.py b/lib/python3.12/site-packages/wandb/sdk/launch/runner/abstract.py new file mode 100644 index 0000000000000000000000000000000000000000..d1f42e6b80dd60f933e396f94424b88e8c985786 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/runner/abstract.py @@ -0,0 +1,185 @@ +"""Implementation of the abstract runner class. + +This class defines the interface that the W&B launch runner uses to manage the lifecycle +of runs launched in different environments (e.g. runs launched locally or in a cluster). +""" + +import logging +import os +import shutil +import subprocess +import sys +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Literal, Optional, Union + +import wandb +from wandb.apis.internal import Api +from wandb.sdk.lib import runid + +from .._project_spec import LaunchProject + +_logger = logging.getLogger(__name__) + + +State = Literal[ + "unknown", + "starting", + "running", + "failed", + "finished", + "stopping", + "stopped", + "preempted", +] + + +class Status: + def __init__(self, state: "State" = "unknown", messages: List[str] = None): # type: ignore + self.state = state + self.messages = messages or [] + + def __repr__(self) -> "State": + return self.state + + def __str__(self) -> str: + return self.state + + def __eq__(self, __value: object) -> bool: + if isinstance(__value, Status): + return self.state == __value.state + else: + return self.state == __value + + def __hash__(self) -> int: + return hash(self.state) + + +class AbstractRun(ABC): + """Wrapper around a W&B launch run. + + A launched run is a subprocess running an entry point + command, that exposes methods for waiting on and cancelling the run. + This class defines the interface that the W&B launch runner uses to manage the lifecycle + of runs launched in different environments (e.g. runs launched locally or in a cluster). + ``AbstractRun`` is not thread-safe. That is, concurrent calls to wait() / cancel() + from multiple threads may inadvertently kill resources (e.g. local processes) unrelated to the + run. + """ + + def __init__(self) -> None: + self._status = Status() + + @property + def status(self) -> Status: + return self._status + + @abstractmethod + async def get_logs(self) -> Optional[str]: + """Return the logs associated with the run.""" + + def _run_cmd( + self, cmd: List[str], output_only: Optional[bool] = False + ) -> Optional[Union["subprocess.Popen[bytes]", bytes]]: + """Run the command and returns a popen object or the stdout of the command. + + Arguments: + cmd: The command to run + output_only: If true just return the stdout bytes + """ + try: + env = os.environ + popen = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE) + if output_only: + popen.wait() + if popen.stdout is not None: + return popen.stdout.read() + return popen + except subprocess.CalledProcessError as e: + wandb.termerror(f"Command failed: {e}") + return None + + @abstractmethod + async def wait(self) -> bool: + """Wait for the run to finish, returning True if the run succeeded and false otherwise. + + Note that in some cases, we may wait until the remote job completes rather than until the W&B run completes. + """ + + @abstractmethod + async def get_status(self) -> Status: + """Get status of the run.""" + + @abstractmethod + async def cancel(self) -> None: + """Cancel the run (interrupts the command subprocess, cancels the run, etc). + + Cancels the run and waits for it to terminate. The W&B run status may not be + set correctly upon run cancellation. + """ + + @property + @abstractmethod + def id(self) -> Optional[str]: + pass + + +class AbstractRunner(ABC): + """Abstract plugin class defining the interface needed to execute W&B Launches. + + You can define subclasses of ``AbstractRunner`` and expose them as third-party + plugins to enable running W&B projects against custom execution backends + (e.g. to run projects against your team's in-house cluster or job scheduler). + """ + + _type: str + + def __init__( + self, + api: Api, + backend_config: Dict[str, Any], + ) -> None: + self._api = api + self.backend_config = backend_config + self._cwd = os.getcwd() + self._namespace = runid.generate_id() + + def find_executable( + self, + cmd: str, + ) -> Union[str, None]: + """Cross platform utility for checking if a program is available.""" + return shutil.which(cmd) + + @property + def api_key(self) -> Any: + return self._api.api_key + + def verify(self) -> bool: + """This is called on first boot to verify the needed commands, and permissions are available. + + For now just call `wandb.termerror` and `sys.exit(1)` + """ + if self._api.api_key is None: + wandb.termerror( + "Couldn't find W&B api key, run wandb login or set WANDB_API_KEY" + ) + sys.exit(1) + return True + + @abstractmethod + async def run( + self, + launch_project: LaunchProject, + image_uri: str, + ) -> Optional[AbstractRun]: + """Submit an LaunchProject to be run. + + Returns a SubmittedRun object to track the execution + Arguments: + launch_project: Object of _project_spec.LaunchProject class representing a wandb launch project + + Returns: + A :py:class:`wandb.sdk.launch.runners.SubmittedRun`. This function is expected to run + the project asynchronously, i.e. it should trigger project execution and then + immediately return a `SubmittedRun` to track execution status. + """ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/kubernetes_monitor.py b/lib/python3.12/site-packages/wandb/sdk/launch/runner/kubernetes_monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..fa98dc6b2ed99843725a6ffadd5f3bf82b972d13 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/runner/kubernetes_monitor.py @@ -0,0 +1,472 @@ +"""Monitors kubernetes resources managed by the launch agent.""" + +import asyncio +import logging +import traceback +from typing import Any, Dict, List, Optional, Tuple, Union + +import kubernetes_asyncio # type: ignore +import urllib3 +from kubernetes_asyncio import watch +from kubernetes_asyncio.client import ( # type: ignore + ApiException, + BatchV1Api, + CoreV1Api, + CustomObjectsApi, + V1Pod, + V1PodStatus, +) + +import wandb +from wandb.sdk.launch.agent import LaunchAgent +from wandb.sdk.launch.errors import LaunchError +from wandb.sdk.launch.runner.abstract import State, Status +from wandb.sdk.launch.utils import get_kube_context_and_api_client + +WANDB_K8S_LABEL_NAMESPACE = "wandb.ai" +WANDB_K8S_RUN_ID = f"{WANDB_K8S_LABEL_NAMESPACE}/run-id" +WANDB_K8S_LABEL_AGENT = f"{WANDB_K8S_LABEL_NAMESPACE}/agent" +WANDB_K8S_LABEL_MONITOR = f"{WANDB_K8S_LABEL_NAMESPACE}/monitor" + + +class Resources: + JOBS = "jobs" + PODS = "pods" + + +class CustomResource: + """Class for custom resources.""" + + def __init__(self, group: str, version: str, plural: str) -> None: + """Initialize the CustomResource.""" + self.group = group + self.version = version + self.plural = plural + + def __str__(self) -> str: + """Return a string representation of the CustomResource.""" + return f"{self.group}/{self.version}/{self.plural}" + + def __hash__(self) -> int: + """Return a hash of the CustomResource.""" + return hash(str(self)) + + +# Maps phases and conditions of custom objects to agent's internal run states. +CRD_STATE_DICT: Dict[str, State] = { + "created": "starting", + "pending": "starting", + "running": "running", + "completing": "running", + "succeeded": "finished", + "completed": "finished", + "failed": "failed", + "aborted": "failed", + "timeout": "failed", + "terminated": "failed", + "terminating": "stopping", +} + +_logger = logging.getLogger(__name__) + + +def create_named_task(name: str, coro: Any, *args: Any, **kwargs: Any) -> asyncio.Task: + """Create a named task.""" + task = asyncio.create_task(coro(*args, **kwargs)) + task.set_name(name) + task.add_done_callback(_log_err_task_callback) + return task + + +def _log_err_task_callback(task: asyncio.Task) -> None: + """Callback to log exceptions from tasks.""" + exec = task.exception() + if exec is not None: + if isinstance(exec, asyncio.CancelledError): + wandb.termlog(f"Task {task.get_name()} was cancelled") + return + name = task.get_name() + wandb.termerror(f"Exception in task {name}") + tb = exec.__traceback__ + tb_str = "".join(traceback.format_tb(tb)) + wandb.termerror(tb_str) + + +def _is_preempted(status: "V1PodStatus") -> bool: + """Check if this pod has been preempted.""" + if hasattr(status, "conditions") and status.conditions is not None: + for condition in status.conditions: + if condition.type == "DisruptionTarget" and condition.reason in [ + "EvictionByEvictionAPI", + "PreemptionByScheduler", + "TerminationByKubelet", + ]: + return True + return False + + +def _is_container_creating(status: "V1PodStatus") -> bool: + """Check if this pod has started creating containers.""" + for container_status in status.container_statuses or []: + if ( + container_status.state + and container_status.state.waiting + and container_status.state.waiting.reason == "ContainerCreating" + ): + return True + return False + + +def _is_pod_unschedulable(status: "V1PodStatus") -> Tuple[bool, str]: + """Return whether the pod is unschedulable along with the reason message.""" + if not status.conditions: + return False, "" + for condition in status.conditions: + if ( + condition.type == "PodScheduled" + and condition.status == "False" + and condition.reason == "Unschedulable" + ): + return True, condition.message + return False, "" + + +def _get_crd_job_name(object: "V1Pod") -> Optional[str]: + refs = object.metadata.owner_references + if refs: + return refs[0].name + return None + + +def _state_from_conditions(conditions: List[Dict[str, Any]]) -> Optional[State]: + """Get the status from the pod conditions.""" + true_conditions = [ + c.get("type", "").lower() for c in conditions if c.get("status") == "True" + ] + detected_states = { + CRD_STATE_DICT[c] for c in true_conditions if c in CRD_STATE_DICT + } + # The list below is ordered so that returning the first state detected + # will accurately reflect the state of the job. + states_in_order: List[State] = [ + "finished", + "failed", + "stopping", + "running", + "starting", + ] + for state in states_in_order: + if state in detected_states: + return state + return None + + +def _state_from_replicated_status(status_dict: Dict[str, int]) -> Optional[State]: + """Infer overall job status from replicated job status for jobsets. + + More info on jobset: + https://github.com/kubernetes-sigs/jobset/blob/main/docs/concepts/README.md + + This is useful for detecting when jobsets are starting. + """ + pods_ready = status_dict.get("ready", 0) + pods_active = status_dict.get("active", 0) + if pods_ready >= 1: + return "running" + elif pods_active >= 1: + return "starting" + return None + + +class LaunchKubernetesMonitor: + """Monitors kubernetes resources managed by the launch agent. + + Note: this class is forced to be a singleton in order to prevent multiple + threads from being created that monitor the same kubernetes resources. + """ + + _instance = None # This is used to ensure only one instance is created. + + def __new__(cls, *args: Any, **kwargs: Any) -> "LaunchKubernetesMonitor": + """Create a new instance of the LaunchKubernetesMonitor. + + This method ensures that only one instance of the LaunchKubernetesMonitor + is created. This is done to prevent multiple threads from being created + that monitor the same kubernetes resources. + """ + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__( + self, + core_api: CoreV1Api, + batch_api: BatchV1Api, + custom_api: CustomObjectsApi, + label_selector: str, + ): + """Initialize the LaunchKubernetesMonitor.""" + self._core_api: CoreV1Api = core_api + self._batch_api: BatchV1Api = batch_api + self._custom_api: CustomObjectsApi = custom_api + + self._label_selector: str = label_selector + + # Dict mapping a tuple of (namespace, resource_type) to an + # asyncio.Task that is monitoring that resource type in that namespace. + self._monitor_tasks: Dict[ + Tuple[str, Union[str, CustomResource]], asyncio.Task + ] = dict() + + # Map from job name to job state. + self._job_states: Dict[str, Status] = dict() + + @classmethod + async def ensure_initialized( + cls, + ) -> None: + """Initialize the LaunchKubernetesMonitor.""" + if cls._instance is None: + _, api_client = await get_kube_context_and_api_client( + kubernetes_asyncio, {} + ) + core_api = CoreV1Api(api_client) + batch_api = BatchV1Api(api_client) + custom_api = CustomObjectsApi(api_client) + label_selector = f"{WANDB_K8S_LABEL_MONITOR}=true" + if LaunchAgent.initialized(): + label_selector += f",{WANDB_K8S_LABEL_AGENT}={LaunchAgent.name()}" + cls( + core_api=core_api, + batch_api=batch_api, + custom_api=custom_api, + label_selector=label_selector, + ) + + @classmethod + def monitor_namespace( + cls, namespace: str, custom_resource: Optional[CustomResource] = None + ) -> None: + """Start monitoring a namespaces for resources.""" + if cls._instance is None: + raise LaunchError( + "LaunchKubernetesMonitor not initialized, cannot monitor namespace." + ) + cls._instance.__monitor_namespace(namespace, custom_resource=custom_resource) + + @classmethod + def get_status(cls, job_name: str) -> Status: + """Get the status of a job.""" + if cls._instance is None: + raise LaunchError( + "LaunchKubernetesMonitor not initialized, cannot get status." + ) + return cls._instance.__get_status(job_name) + + @classmethod + def status_count(cls) -> Dict[State, int]: + """Get a dictionary mapping statuses to the # monitored jobs with each status.""" + if cls._instance is None: + raise ValueError( + "LaunchKubernetesMonitor not initialized, cannot get status counts." + ) + return cls._instance.__status_count() + + def __monitor_namespace( + self, namespace: str, custom_resource: Optional[CustomResource] = None + ) -> None: + """Start monitoring a namespaces for resources.""" + if (namespace, Resources.PODS) not in self._monitor_tasks: + self._monitor_tasks[(namespace, Resources.PODS)] = create_named_task( + f"monitor_pods_{namespace}", + self._monitor_pods, + namespace, + ) + # If a custom resource is specified then we will start monitoring + # that resource type in the namespace instead of jobs. + if custom_resource is not None: + if (namespace, custom_resource) not in self._monitor_tasks: + self._monitor_tasks[(namespace, custom_resource)] = create_named_task( + f"monitor_{custom_resource}_{namespace}", + self._monitor_crd, + namespace, + custom_resource=custom_resource, + ) + else: + if (namespace, Resources.JOBS) not in self._monitor_tasks: + self._monitor_tasks[(namespace, Resources.JOBS)] = create_named_task( + f"monitor_jobs_{namespace}", + self._monitor_jobs, + namespace, + ) + + def __get_status(self, job_name: str) -> Status: + """Get the status of a job.""" + if job_name not in self._job_states: + return Status("unknown") + state = self._job_states[job_name] + return state + + def __status_count(self) -> Dict[State, int]: + """Get a dictionary mapping statuses to the # monitored jobs with each status.""" + counts = dict() + for _, status in self._job_states.items(): + state = status.state + if state not in counts: + counts[state] = 1 + else: + counts[state] += 1 + return counts + + def _set_status_state(self, job_name: str, state: State) -> None: + """Set the status of the run.""" + if job_name not in self._job_states: + self._job_states[job_name] = Status(state) + elif self._job_states[job_name].state != state: + self._job_states[job_name].state = state + + def _add_status_message(self, job_name: str, message: str) -> None: + if job_name not in self._job_states: + self._job_states[job_name] = Status("unknown") + wandb.termwarn(f"Warning from Kubernetes for job {job_name}: {message}") + self._job_states[job_name].messages.append(message) + + async def _monitor_pods(self, namespace: str) -> None: + """Monitor a namespace for changes.""" + watcher = SafeWatch(watch.Watch()) + async for event in watcher.stream( + self._core_api.list_namespaced_pod, + namespace=namespace, + label_selector=self._label_selector, + ): + obj = event.get("object") + job_name = obj.metadata.labels.get("job-name") or _get_crd_job_name(obj) + if job_name is None or not hasattr(obj, "status"): + continue + if self.__get_status(job_name) in ["finished", "failed"]: + continue + + is_unschedulable, reason = _is_pod_unschedulable(obj.status) + if is_unschedulable: + self._add_status_message(job_name, reason) + if obj.status.phase == "Running" or _is_container_creating(obj.status): + self._set_status_state(job_name, "running") + elif _is_preempted(obj.status): + self._set_status_state(job_name, "preempted") + + async def _monitor_jobs(self, namespace: str) -> None: + """Monitor a namespace for changes.""" + watcher = SafeWatch(watch.Watch()) + async for event in watcher.stream( + self._batch_api.list_namespaced_job, + namespace=namespace, + label_selector=self._label_selector, + ): + obj = event.get("object") + job_name = obj.metadata.name + + if obj.status.succeeded == 1: + self._set_status_state(job_name, "finished") + elif obj.status.failed is not None and obj.status.failed >= 1: + self._set_status_state(job_name, "failed") + + # If the job is deleted and we haven't seen a terminal state + # then we will consider the job failed. + if event.get("type") == "DELETED": + if self._job_states.get(job_name) != Status("finished"): + self._set_status_state(job_name, "failed") + + async def _monitor_crd( + self, namespace: str, custom_resource: CustomResource + ) -> None: + """Monitor a namespace for changes.""" + watcher = SafeWatch(watch.Watch()) + async for event in watcher.stream( + self._custom_api.list_namespaced_custom_object, + namespace=namespace, + plural=custom_resource.plural, + group=custom_resource.group, + version=custom_resource.version, + label_selector=self._label_selector, + ): + object = event.get("object") + name = object.get("metadata", dict()).get("name") + status = object.get("status") + state = None + if status is None: + continue + replicated_jobs_status = status.get("ReplicatedJobsStatus") + if isinstance(replicated_jobs_status, dict): + state = _state_from_replicated_status(replicated_jobs_status) + state_dict = status.get("state") + if isinstance(state_dict, dict): + phase = state_dict.get("phase") + if phase: + state = CRD_STATE_DICT.get(phase.lower()) + else: + conditions = status.get("conditions") + if isinstance(conditions, list): + state = _state_from_conditions(conditions) + else: + # This should never happen. + _logger.warning( + f"Unexpected conditions type {type(conditions)} " + f"for CRD watcher in {namespace}" + ) + if state is None: + continue + self._set_status_state(name, state) + + +class SafeWatch: + """Wrapper for the kubernetes watch class that can recover in more situations.""" + + def __init__(self, watcher: watch.Watch) -> None: + """Initialize the SafeWatch.""" + self._watcher = watcher + self._last_seen_resource_version: Optional[str] = None + self._stopped = False + + async def stream(self, func: Any, *args: Any, **kwargs: Any) -> Any: + """Stream the watcher. + + This method will automatically resume the stream if it breaks. It will + also save the resource version so that the stream can be resumed from + the last seen resource version. + """ + while True: + try: + async for event in self._watcher.stream( + func, *args, **kwargs, timeout_seconds=30 + ): + if self._stopped: + break + # Save the resource version so that we can resume the stream + # if it breaks. + object = event.get("object") + if isinstance(object, dict): + self._last_seen_resource_version = object.get( + "metadata", dict() + ).get("resourceVersion") + else: + self._last_seen_resource_version = ( + object.metadata.resource_version + ) + kwargs["resource_version"] = self._last_seen_resource_version + yield event + # If stream ends after stop just break + if self._stopped: + break + except urllib3.exceptions.ProtocolError as e: + wandb.termwarn(f"Broken event stream: {e}, attempting to recover") + except ApiException as e: + if e.status == 410: + # If resource version is too old we need to start over. + del kwargs["resource_version"] + self._last_seen_resource_version = None + except Exception as E: + exc_type = type(E).__name__ + stack_trace = traceback.format_exc() + wandb.termerror( + f"Unknown exception in event stream of type {exc_type}: {E}, attempting to recover. Stack trace: {stack_trace}" + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/kubernetes_runner.py b/lib/python3.12/site-packages/wandb/sdk/launch/runner/kubernetes_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..4b29ceaa92d0e1ed725a925a9fd8931531656aa9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/runner/kubernetes_runner.py @@ -0,0 +1,963 @@ +"""Implementation of KubernetesRunner class for wandb launch.""" + +import asyncio +import base64 +import datetime +import json +import logging +import os +from typing import Any, Dict, Iterator, List, Optional, Tuple, Union + +import yaml + +import wandb +from wandb.apis.internal import Api +from wandb.sdk.launch.agent.agent import LaunchAgent +from wandb.sdk.launch.environment.abstract import AbstractEnvironment +from wandb.sdk.launch.registry.abstract import AbstractRegistry +from wandb.sdk.launch.registry.azure_container_registry import AzureContainerRegistry +from wandb.sdk.launch.registry.local_registry import LocalRegistry +from wandb.sdk.launch.runner.abstract import Status +from wandb.sdk.launch.runner.kubernetes_monitor import ( + WANDB_K8S_LABEL_AGENT, + WANDB_K8S_LABEL_MONITOR, + WANDB_K8S_RUN_ID, + CustomResource, + LaunchKubernetesMonitor, +) +from wandb.sdk.lib.retry import ExponentialBackoff, retry_async +from wandb.util import get_module + +from .._project_spec import EntryPoint, LaunchProject +from ..errors import LaunchError +from ..utils import ( + CODE_MOUNT_DIR, + LOG_PREFIX, + MAX_ENV_LENGTHS, + PROJECT_SYNCHRONOUS, + get_kube_context_and_api_client, + make_name_dns_safe, +) +from .abstract import AbstractRun, AbstractRunner + +get_module( + "kubernetes_asyncio", + required="Kubernetes runner requires the kubernetes package. Please install it with `pip install wandb[launch]`.", +) + +import kubernetes_asyncio # type: ignore # noqa: E402 +from kubernetes_asyncio import client # noqa: E402 +from kubernetes_asyncio.client.api.batch_v1_api import ( # type: ignore # noqa: E402 + BatchV1Api, +) +from kubernetes_asyncio.client.api.core_v1_api import ( # type: ignore # noqa: E402 + CoreV1Api, +) +from kubernetes_asyncio.client.api.custom_objects_api import ( # type: ignore # noqa: E402 + CustomObjectsApi, +) +from kubernetes_asyncio.client.models.v1_secret import ( # type: ignore # noqa: E402 + V1Secret, +) +from kubernetes_asyncio.client.rest import ApiException # type: ignore # noqa: E402 + +TIMEOUT = 5 +API_KEY_SECRET_MAX_RETRIES = 5 + +_logger = logging.getLogger(__name__) + + +SOURCE_CODE_PVC_MOUNT_PATH = os.environ.get("WANDB_LAUNCH_CODE_PVC_MOUNT_PATH") +SOURCE_CODE_PVC_NAME = os.environ.get("WANDB_LAUNCH_CODE_PVC_NAME") + + +class KubernetesSubmittedRun(AbstractRun): + """Wrapper for a launched run on Kubernetes.""" + + def __init__( + self, + batch_api: "BatchV1Api", + core_api: "CoreV1Api", + name: str, + namespace: Optional[str] = "default", + secret: Optional["V1Secret"] = None, + ) -> None: + """Initialize a KubernetesSubmittedRun. + + Other implementations of the AbstractRun interface poll on the run + when `get_status` is called, but KubernetesSubmittedRun uses + Kubernetes watch streams to update the run status. One thread handles + events from the job object and another thread handles events from the + rank 0 pod. These threads updated the `_status` attributed of the + KubernetesSubmittedRun object. When `get_status` is called, the + `_status` attribute is returned. + + Arguments: + batch_api: Kubernetes BatchV1Api object. + core_api: Kubernetes CoreV1Api object. + name: Name of the job. + namespace: Kubernetes namespace. + secret: Kubernetes secret. + + Returns: + None. + """ + self.batch_api = batch_api + self.core_api = core_api + self.name = name + self.namespace = namespace + self._fail_count = 0 + self.secret = secret + + @property + def id(self) -> str: + """Return the run id.""" + return self.name + + async def get_logs(self) -> Optional[str]: + try: + pods = await self.core_api.list_namespaced_pod( + label_selector=f"job-name={self.name}", namespace=self.namespace + ) + pod_names = [pi.metadata.name for pi in pods.items] + if not pod_names: + wandb.termwarn(f"Found no pods for kubernetes job: {self.name}") + return None + logs = await self.core_api.read_namespaced_pod_log( + name=pod_names[0], namespace=self.namespace + ) + if logs: + return str(logs) + else: + wandb.termwarn(f"No logs for kubernetes pod(s): {pod_names}") + return None + except Exception as e: + wandb.termerror(f"{LOG_PREFIX}Failed to get pod logs: {e}") + return None + + async def wait(self) -> bool: + """Wait for the run to finish. + + Returns: + True if the run finished successfully, False otherwise. + """ + while True: + status = await self.get_status() + wandb.termlog(f"{LOG_PREFIX}Job {self.name} status: {status.state}") + if status.state in ["finished", "failed", "preempted"]: + break + await asyncio.sleep(5) + + await self._delete_secret() + return ( + status.state == "finished" + ) # todo: not sure if this (copied from aws runner) is the right approach? should we return false on failure + + async def get_status(self) -> Status: + status = LaunchKubernetesMonitor.get_status(self.name) + if status in ["stopped", "failed", "finished", "preempted"]: + await self._delete_secret() + return status + + async def cancel(self) -> None: + """Cancel the run.""" + try: + await self.batch_api.delete_namespaced_job( + namespace=self.namespace, + name=self.name, + ) + await self._delete_secret() + except ApiException as e: + raise LaunchError( + f"Failed to delete Kubernetes Job {self.name} in namespace {self.namespace}: {str(e)}" + ) from e + + async def _delete_secret(self) -> None: + # Cleanup secret if not running in a helm-managed context + if not os.environ.get("WANDB_RELEASE_NAME") and self.secret: + await self.core_api.delete_namespaced_secret( + name=self.secret.metadata.name, + namespace=self.secret.metadata.namespace, + ) + self.secret = None + + +class CrdSubmittedRun(AbstractRun): + """Run submitted to a CRD backend, e.g. Volcano.""" + + def __init__( + self, + group: str, + version: str, + plural: str, + name: str, + namespace: str, + core_api: CoreV1Api, + custom_api: CustomObjectsApi, + ) -> None: + """Create a run object for tracking the progress of a CRD. + + Arguments: + group: The API group of the CRD. + version: The API version of the CRD. + plural: The plural name of the CRD. + name: The name of the CRD instance. + namespace: The namespace of the CRD instance. + core_api: The Kubernetes core API client. + custom_api: The Kubernetes custom object API client. + + Raises: + LaunchError: If the CRD instance does not exist. + """ + self.group = group + self.version = version + self.plural = plural + self.name = name + self.namespace = namespace + self.core_api = core_api + self.custom_api = custom_api + self._fail_count = 0 + + @property + def id(self) -> str: + """Get the name of the custom object.""" + return self.name + + async def get_logs(self) -> Optional[str]: + """Get logs for custom object.""" + # TODO: test more carefully once we release multi-node support + logs: Dict[str, Optional[str]] = {} + try: + pods = await self.core_api.list_namespaced_pod( + label_selector=f"wandb/run-id={self.name}", namespace=self.namespace + ) + pod_names = [pi.metadata.name for pi in pods.items] + for pod_name in pod_names: + logs[pod_name] = await self.core_api.read_namespaced_pod_log( + name=pod_name, namespace=self.namespace + ) + except ApiException as e: + wandb.termwarn(f"Failed to get logs for {self.name}: {str(e)}") + return None + if not logs: + return None + logs_as_array = [f"Pod {pod_name}:\n{log}" for pod_name, log in logs.items()] + return "\n".join(logs_as_array) + + async def get_status(self) -> Status: + """Get status of custom object.""" + return LaunchKubernetesMonitor.get_status(self.name) + + async def cancel(self) -> None: + """Cancel the custom object.""" + try: + await self.custom_api.delete_namespaced_custom_object( + group=self.group, + version=self.version, + namespace=self.namespace, + plural=self.plural, + name=self.name, + ) + except ApiException as e: + raise LaunchError( + f"Failed to delete CRD {self.name} in namespace {self.namespace}: {str(e)}" + ) from e + + async def wait(self) -> bool: + """Wait for this custom object to finish running.""" + while True: + status = await self.get_status() + wandb.termlog(f"{LOG_PREFIX}Job {self.name} status: {status}") + if status.state in ["finished", "failed", "preempted"]: + return status.state == "finished" + await asyncio.sleep(5) + + +class KubernetesRunner(AbstractRunner): + """Launches runs onto kubernetes.""" + + def __init__( + self, + api: Api, + backend_config: Dict[str, Any], + environment: AbstractEnvironment, + registry: AbstractRegistry, + ) -> None: + """Create a Kubernetes runner. + + Arguments: + api: The API client object. + backend_config: The backend configuration. + environment: The environment to launch runs into. + + Raises: + LaunchError: If the Kubernetes configuration is invalid. + """ + super().__init__(api, backend_config) + self.environment = environment + self.registry = registry + + def get_namespace( + self, resource_args: Dict[str, Any], context: Dict[str, Any] + ) -> str: + """Get the namespace to launch into. + + Arguments: + resource_args: The resource args to launch. + context: The k8s config context. + + Returns: + The namespace to launch into. + """ + default_namespace = ( + context["context"].get("namespace", "default") if context else "default" + ) + return ( # type: ignore[no-any-return] + resource_args.get("metadata", {}).get("namespace") + or resource_args.get( + "namespace" + ) # continue support for malformed namespace + or self.backend_config.get("runner", {}).get("namespace") + or default_namespace + ) + + async def _inject_defaults( + self, + resource_args: Dict[str, Any], + launch_project: LaunchProject, + image_uri: str, + namespace: str, + core_api: "CoreV1Api", + ) -> Tuple[Dict[str, Any], Optional["V1Secret"]]: + """Apply our default values, return job dict and api key secret. + + Arguments: + resource_args (Dict[str, Any]): The resource args to launch. + launch_project (LaunchProject): The launch project. + builder (Optional[AbstractBuilder]): The builder. + namespace (str): The namespace. + core_api (CoreV1Api): The core api. + + Returns: + Tuple[Dict[str, Any], Optional["V1Secret"]]: The resource args and api key secret. + """ + job: Dict[str, Any] = { + "apiVersion": "batch/v1", + "kind": "Job", + } + job.update(resource_args) + + job_metadata: Dict[str, Any] = job.get("metadata", {}) + job_spec: Dict[str, Any] = {"backoffLimit": 0, "ttlSecondsAfterFinished": 60} + job_spec.update(job.get("spec", {})) + pod_template: Dict[str, Any] = job_spec.get("template", {}) + pod_spec: Dict[str, Any] = {"restartPolicy": "Never"} + pod_spec.update(pod_template.get("spec", {})) + containers: List[Dict[str, Any]] = pod_spec.get("containers", [{}]) + + # Add labels to job metadata + job_metadata.setdefault("labels", {}) + job_metadata["labels"][WANDB_K8S_RUN_ID] = launch_project.run_id + job_metadata["labels"][WANDB_K8S_LABEL_MONITOR] = "true" + if LaunchAgent.initialized(): + job_metadata["labels"][WANDB_K8S_LABEL_AGENT] = LaunchAgent.name() + # name precedence: name in spec > generated name + if not job_metadata.get("name"): + job_metadata["generateName"] = make_name_dns_safe( + f"launch-{launch_project.target_entity}-{launch_project.target_project}-" + ) + + for i, cont in enumerate(containers): + if "name" not in cont: + cont["name"] = cont.get("name", "launch" + str(i)) + if "securityContext" not in cont: + cont["securityContext"] = { + "allowPrivilegeEscalation": False, + "capabilities": {"drop": ["ALL"]}, + "seccompProfile": {"type": "RuntimeDefault"}, + } + + entry_point = ( + launch_project.override_entrypoint or launch_project.get_job_entry_point() + ) + if launch_project.docker_image: + # dont specify run id if user provided image, could have multiple runs + containers[0]["image"] = image_uri + # TODO: handle secret pulling image from registry + elif not any(["image" in cont for cont in containers]): + assert entry_point is not None + # in the non instance case we need to make an imagePullSecret + # so the new job can pull the image + containers[0]["image"] = image_uri + secret = await maybe_create_imagepull_secret( + core_api, self.registry, launch_project.run_id, namespace + ) + if secret is not None: + pod_spec["imagePullSecrets"] = [ + {"name": f"regcred-{launch_project.run_id}"} + ] + + inject_entrypoint_and_args( + containers, + entry_point, + launch_project.override_args, + launch_project.override_entrypoint is not None, + ) + + env_vars = launch_project.get_env_vars_dict( + self._api, MAX_ENV_LENGTHS[self.__class__.__name__] + ) + api_key_secret = None + for cont in containers: + # Add our env vars to user supplied env vars + env = cont.get("env") or [] + for key, value in env_vars.items(): + if ( + key == "WANDB_API_KEY" + and value + and ( + LaunchAgent.initialized() + or self.backend_config[PROJECT_SYNCHRONOUS] + ) + ): + # Override API key with secret. TODO: Do the same for other runners + release_name = os.environ.get("WANDB_RELEASE_NAME") + secret_name = "wandb-api-key" + if release_name: + secret_name += f"-{release_name}" + else: + secret_name += f"-{launch_project.run_id}" + + 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( + { + "name": key, + "valueFrom": { + "secretKeyRef": { + "name": secret_name, + "key": "password", + } + }, + } + ) + else: + env.append({"name": key, "value": value}) + cont["env"] = env + + pod_spec["containers"] = containers + pod_template["spec"] = pod_spec + job_spec["template"] = pod_template + job["spec"] = job_spec + job["metadata"] = job_metadata + + add_label_to_pods( + job, + WANDB_K8S_LABEL_MONITOR, + "true", + ) + + if launch_project.job_base_image: + apply_code_mount_configuration( + job, + launch_project, + ) + + # Add wandb.ai/agent: current agent label on all pods + if LaunchAgent.initialized(): + add_label_to_pods( + job, + WANDB_K8S_LABEL_AGENT, + LaunchAgent.name(), + ) + + return job, api_key_secret + + async def run( + self, launch_project: LaunchProject, image_uri: str + ) -> Optional[AbstractRun]: + """Execute a launch project on Kubernetes. + + Arguments: + launch_project: The launch project to execute. + builder: The builder to use to build the image. + + Returns: + The run object if the run was successful, otherwise None. + """ + await LaunchKubernetesMonitor.ensure_initialized() + resource_args = launch_project.fill_macros(image_uri).get("kubernetes", {}) + if not resource_args: + wandb.termlog( + f"{LOG_PREFIX}Note: no resource args specified. Add a " + "Kubernetes yaml spec or other options in a json file " + "with --resource-args ." + ) + _logger.info(f"Running Kubernetes job with resource args: {resource_args}") + + context, api_client = await get_kube_context_and_api_client( + kubernetes_asyncio, resource_args + ) + + # If using pvc for code mount, move code there. + if launch_project.job_base_image is not None: + if SOURCE_CODE_PVC_NAME is None or SOURCE_CODE_PVC_MOUNT_PATH is None: + raise LaunchError( + "WANDB_LAUNCH_SOURCE_CODE_PVC_ environment variables not set. " + "Unable to mount source code PVC into base image. " + "Use the `codeMountPvcName` variable in the agent helm chart " + "to enable base image jobs for this agent. See " + "https://github.com/wandb/helm-charts/tree/main/charts/launch-agent " + "for more information." + ) + code_subdir = launch_project.get_image_source_string() + launch_project.change_project_dir( + os.path.join(SOURCE_CODE_PVC_MOUNT_PATH, code_subdir) + ) + + # If the user specified an alternate api, we need will execute this + # run by creating a custom object. + api_version = resource_args.get("apiVersion", "batch/v1") + + if api_version not in ["batch/v1", "batch/v1beta1"]: + env_vars = launch_project.get_env_vars_dict( + self._api, MAX_ENV_LENGTHS[self.__class__.__name__] + ) + # Crawl the resource args and add our env vars to the containers. + add_wandb_env(resource_args, env_vars) + + # Add our labels to the resource args. This is necessary for the + # agent to find the custom object later on. + resource_args["metadata"] = resource_args.get("metadata", {}) + resource_args["metadata"]["labels"] = resource_args["metadata"].get( + "labels", {} + ) + resource_args["metadata"]["labels"][WANDB_K8S_LABEL_MONITOR] = "true" + + # Crawl the resource arsg and add our labels to the pods. This is + # necessary for the agent to find the pods later on. + add_label_to_pods( + resource_args, + WANDB_K8S_LABEL_MONITOR, + "true", + ) + + # Add wandb.ai/agent: current agent label on all pods + if LaunchAgent.initialized(): + add_label_to_pods( + resource_args, + WANDB_K8S_LABEL_AGENT, + LaunchAgent.name(), + ) + resource_args["metadata"]["labels"][WANDB_K8S_LABEL_AGENT] = ( + LaunchAgent.name() + ) + + if launch_project.job_base_image: + apply_code_mount_configuration(resource_args, launch_project) + + overrides = {} + if launch_project.override_args: + overrides["args"] = launch_project.override_args + if launch_project.override_entrypoint: + overrides["command"] = launch_project.override_entrypoint.command + add_entrypoint_args_overrides( + resource_args, + overrides, + ) + api = client.CustomObjectsApi(api_client) + # Infer the attributes of a custom object from the apiVersion and/or + # a kind: attribute in the resource args. + namespace = self.get_namespace(resource_args, context) + group, version, *_ = api_version.split("/") + group = resource_args.get("group", group) + version = resource_args.get("version", version) + kind = resource_args.get("kind", version) + plural = f"{kind.lower()}s" + custom_resource = CustomResource( + group=group, + version=version, + plural=plural, + ) + LaunchKubernetesMonitor.monitor_namespace( + namespace, custom_resource=custom_resource + ) + + try: + response = await api.create_namespaced_custom_object( + group=group, + version=version, + namespace=namespace, + plural=plural, + body=resource_args, + ) + except ApiException as e: + body = json.loads(e.body) + body_yaml = yaml.dump(body) + raise LaunchError( + f"Error creating CRD of kind {kind}: {e.status} {e.reason}\n{body_yaml}" + ) from e + name = response.get("metadata", {}).get("name") + _logger.info(f"Created {kind} {response['metadata']['name']}") + submitted_run = CrdSubmittedRun( + name=name, + group=group, + version=version, + namespace=namespace, + plural=plural, + core_api=client.CoreV1Api(api_client), + custom_api=api, + ) + if self.backend_config[PROJECT_SYNCHRONOUS]: + await submitted_run.wait() + return submitted_run + + batch_api = kubernetes_asyncio.client.BatchV1Api(api_client) + core_api = kubernetes_asyncio.client.CoreV1Api(api_client) + namespace = self.get_namespace(resource_args, context) + job, secret = await self._inject_defaults( + resource_args, launch_project, image_uri, namespace, core_api + ) + msg = "Creating Kubernetes job" + if "name" in resource_args: + msg += f": {resource_args['name']}" + _logger.info(msg) + try: + response = await kubernetes_asyncio.utils.create_from_dict( + api_client, job, namespace=namespace + ) + except kubernetes_asyncio.utils.FailToCreateError as e: + for exc in e.api_exceptions: + resp = json.loads(exc.body) + msg = resp.get("message") + code = resp.get("code") + raise LaunchError( + f"Failed to create Kubernetes job for run {launch_project.run_id} ({code} {exc.reason}): {msg}" + ) + except Exception as e: + raise LaunchError( + f"Unexpected exception when creating Kubernetes job: {str(e)}\n" + ) + job_response = response[0] + job_name = job_response.metadata.name + LaunchKubernetesMonitor.monitor_namespace(namespace) + submitted_job = KubernetesSubmittedRun( + batch_api, core_api, job_name, namespace, secret + ) + if self.backend_config[PROJECT_SYNCHRONOUS]: + await submitted_job.wait() + + return submitted_job + + +def inject_entrypoint_and_args( + containers: List[dict], + entry_point: Optional[EntryPoint], + override_args: List[str], + should_override_entrypoint: bool, +) -> None: + """Inject the entrypoint and args into the containers. + + Arguments: + containers: The containers to inject the entrypoint and args into. + entry_point: The entrypoint to inject. + override_args: The args to inject. + should_override_entrypoint: Whether to override the entrypoint. + + Returns: + None + """ + for i in range(len(containers)): + if override_args: + containers[i]["args"] = override_args + if entry_point and ( + not containers[i].get("command") or should_override_entrypoint + ): + containers[i]["command"] = entry_point.command + + +async def ensure_api_key_secret( + core_api: "CoreV1Api", + secret_name: str, + namespace: str, + api_key: str, +) -> "V1Secret": + """Create a secret containing a user's wandb API key. + + Arguments: + core_api: The Kubernetes CoreV1Api object. + secret_name: The name to use for the secret. + namespace: The namespace to create the secret in. + api_key: The user's wandb API key + + Returns: + The created secret + """ + secret_data = {"password": base64.b64encode(api_key.encode()).decode()} + labels = {"wandb.ai/created-by": "launch-agent"} + secret = client.V1Secret( + data=secret_data, + metadata=client.V1ObjectMeta( + name=secret_name, namespace=namespace, labels=labels + ), + kind="Secret", + type="kubernetes.io/basic-auth", + ) + + try: + try: + return await core_api.create_namespaced_secret(namespace, secret) + except ApiException as e: + # 409 = conflict = secret already exists + if e.status == 409: + existing_secret = await core_api.read_namespaced_secret( + name=secret_name, namespace=namespace + ) + if existing_secret.data != secret_data: + # If it's a previous secret made by launch agent, clean it up + if ( + existing_secret.metadata.labels.get("wandb.ai/created-by") + == "launch-agent" + ): + await core_api.delete_namespaced_secret( + name=secret_name, namespace=namespace + ) + return await core_api.create_namespaced_secret( + namespace, secret + ) + else: + raise LaunchError( + f"Kubernetes secret already exists in namespace {namespace} with incorrect data: {secret_name}" + ) + return existing_secret + raise + except Exception as e: + raise LaunchError( + f"Exception when ensuring Kubernetes API key secret: {str(e)}\n" + ) + + +async def maybe_create_imagepull_secret( + core_api: "CoreV1Api", + registry: AbstractRegistry, + run_id: str, + namespace: str, +) -> Optional["V1Secret"]: + """Create a secret for pulling images from a private registry. + + Arguments: + core_api: The Kubernetes CoreV1Api object. + registry: The registry to pull from. + run_id: The run id. + namespace: The namespace to create the secret in. + + Returns: + A secret if one was created, otherwise None. + """ + secret = None + if isinstance(registry, LocalRegistry) or isinstance( + registry, AzureContainerRegistry + ): + # Secret not required + return None + uname, token = await registry.get_username_password() + creds_info = { + "auths": { + registry.uri: { + "auth": base64.b64encode(f"{uname}:{token}".encode()).decode(), + # need an email but the use is deprecated + "email": "deprecated@wandblaunch.com", + } + } + } + secret_data = { + ".dockerconfigjson": base64.b64encode(json.dumps(creds_info).encode()).decode() + } + secret = client.V1Secret( + data=secret_data, + metadata=client.V1ObjectMeta(name=f"regcred-{run_id}", namespace=namespace), + kind="Secret", + type="kubernetes.io/dockerconfigjson", + ) + try: + try: + return await core_api.create_namespaced_secret(namespace, secret) + except ApiException as e: + # 409 = conflict = secret already exists + if e.status == 409: + return await core_api.read_namespaced_secret( + name=f"regcred-{run_id}", namespace=namespace + ) + raise + except Exception as e: + raise LaunchError(f"Exception when creating Kubernetes secret: {str(e)}\n") + + +def yield_containers(root: Any) -> Iterator[dict]: + """Yield all container specs in a manifest. + + Recursively traverses the manifest and yields all container specs. Container + specs are identified by the presence of a "containers" key in the value. + """ + if isinstance(root, dict): + for k, v in root.items(): + if k == "containers": + if isinstance(v, list): + yield from v + elif isinstance(v, (dict, list)): + yield from yield_containers(v) + elif isinstance(root, list): + for item in root: + yield from yield_containers(item) + + +def add_wandb_env(root: Union[dict, list], env_vars: Dict[str, str]) -> None: + """Injects wandb environment variables into specs. + + Recursively walks the spec and injects the environment variables into + every container spec. Containers are identified by the "containers" key. + + This function treats the WANDB_RUN_ID and WANDB_GROUP_ID environment variables + specially. If they are present in the spec, they will be overwritten. If a setting + for WANDB_RUN_ID is provided in env_vars, then that environment variable will only be + set in the first container modified by this function. + + Arguments: + root: The spec to modify. + env_vars: The environment variables to inject. + + Returns: None. + """ + for cont in yield_containers(root): + env = cont.setdefault("env", []) + env.extend([{"name": key, "value": value} for key, value in env_vars.items()]) + cont["env"] = env + # After we have set WANDB_RUN_ID once, we don't want to set it again + if "WANDB_RUN_ID" in env_vars: + env_vars.pop("WANDB_RUN_ID") + + +def yield_pods(manifest: Any) -> Iterator[dict]: + """Yield all pod specs in a manifest. + + Recursively traverses the manifest and yields all pod specs. Pod specs are + identified by the presence of a "spec" key with a "containers" key in the + value. + """ + if isinstance(manifest, list): + for item in manifest: + yield from yield_pods(item) + elif isinstance(manifest, dict): + if "spec" in manifest and "containers" in manifest["spec"]: + yield manifest + for value in manifest.values(): + if isinstance(value, (dict, list)): + yield from yield_pods(value) + + +def add_label_to_pods( + manifest: Union[dict, list], label_key: str, label_value: str +) -> None: + """Add a label to all pod specs in a manifest. + + Recursively traverses the manifest and adds the label to all pod specs. + Pod specs are identified by the presence of a "spec" key with a "containers" + key in the value. + + Arguments: + manifest: The manifest to modify. + label_key: The label key to add. + label_value: The label value to add. + + Returns: None. + """ + for pod in yield_pods(manifest): + metadata = pod.setdefault("metadata", {}) + labels = metadata.setdefault("labels", {}) + labels[label_key] = label_value + + +def add_entrypoint_args_overrides(manifest: Union[dict, list], overrides: dict) -> None: + """Add entrypoint and args overrides to all containers in a manifest. + + Recursively traverses the manifest and adds the entrypoint and args overrides + to all containers. Containers are identified by the presence of a "spec" key + with a "containers" key in the value. + + Arguments: + manifest: The manifest to modify. + overrides: Dictionary with args and entrypoint keys. + + Returns: None. + """ + if isinstance(manifest, list): + for item in manifest: + add_entrypoint_args_overrides(item, overrides) + elif isinstance(manifest, dict): + if "spec" in manifest and "containers" in manifest["spec"]: + containers = manifest["spec"]["containers"] + for container in containers: + if "command" in overrides: + container["command"] = overrides["command"] + if "args" in overrides: + container["args"] = overrides["args"] + for value in manifest.values(): + add_entrypoint_args_overrides(value, overrides) + + +def apply_code_mount_configuration( + manifest: Union[Dict, list], project: LaunchProject +) -> None: + """Apply code mount configuration to all containers in a manifest. + + Recursively traverses the manifest and adds the code mount configuration to + all containers. Containers are identified by the presence of a "spec" key + with a "containers" key in the value. + + Arguments: + manifest: The manifest to modify. + project: The launch project. + + Returns: None. + """ + assert SOURCE_CODE_PVC_NAME is not None + source_dir = project.get_image_source_string() + for pod in yield_pods(manifest): + for container in yield_containers(pod): + if "volumeMounts" not in container: + container["volumeMounts"] = [] + container["volumeMounts"].append( + { + "name": "wandb-source-code-volume", + "mountPath": CODE_MOUNT_DIR, + "subPath": source_dir, + } + ) + container["workingDir"] = CODE_MOUNT_DIR + spec = pod["spec"] + if "volumes" not in spec: + spec["volumes"] = [] + spec["volumes"].append( + { + "name": "wandb-source-code-volume", + "persistentVolumeClaim": { + "claimName": SOURCE_CODE_PVC_NAME, + }, + } + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/local_container.py b/lib/python3.12/site-packages/wandb/sdk/launch/runner/local_container.py new file mode 100644 index 0000000000000000000000000000000000000000..e831467a5c1967d0258231cfbe305d32936e29e5 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/runner/local_container.py @@ -0,0 +1,301 @@ +import asyncio +import logging +import os +import shlex +import subprocess +import sys +import threading +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import wandb +from wandb.sdk.launch.environment.abstract import AbstractEnvironment +from wandb.sdk.launch.registry.abstract import AbstractRegistry + +from .._project_spec import LaunchProject +from ..errors import LaunchError +from ..utils import ( + CODE_MOUNT_DIR, + LOG_PREFIX, + MAX_ENV_LENGTHS, + PROJECT_SYNCHRONOUS, + _is_wandb_dev_uri, + _is_wandb_local_uri, + docker_image_exists, + event_loop_thread_exec, + pull_docker_image, + sanitize_wandb_api_key, +) +from .abstract import AbstractRun, AbstractRunner, Status + +if TYPE_CHECKING: + from wandb.apis.internal import Api + +_logger = logging.getLogger(__name__) + + +class LocalSubmittedRun(AbstractRun): + """Instance of ``AbstractRun`` corresponding to a subprocess launched to run an entry point command locally.""" + + def __init__(self) -> None: + super().__init__() + self._command_proc: Optional[subprocess.Popen] = None + self._stdout: Optional[str] = None + self._terminate_flag: bool = False + self._thread: Optional[threading.Thread] = None + + def set_command_proc(self, command_proc: subprocess.Popen) -> None: + self._command_proc = command_proc + + def set_thread(self, thread: threading.Thread) -> None: + self._thread = thread + + @property + def id(self) -> Optional[str]: + if self._command_proc is None: + return None + return str(self._command_proc.pid) + + async def wait(self) -> bool: + assert self._thread is not None + # if command proc is not set + # wait for thread to set it + if self._command_proc is None: + while self._thread.is_alive(): + await asyncio.sleep(5) + # command proc can be updated by another thread + if self._command_proc is not None: + break # type: ignore # mypy thinks this is unreachable + else: + return False + wait = event_loop_thread_exec(self._command_proc.wait) + return int(await wait()) == 0 + + async def get_logs(self) -> Optional[str]: + return self._stdout + + async def cancel(self) -> None: + # thread is set immediately after starting, should always exist + assert self._thread is not None + + # cancel called before the thread subprocess has started + # indicates to thread to not start command proc if not already started + self._terminate_flag = True + + async def get_status(self) -> Status: + assert self._thread is not None, "Failed to get status, self._thread = None" + if self._command_proc is None: + if self._thread.is_alive(): + return Status("running") + return Status("stopped") + exit_code = self._command_proc.poll() + if exit_code is None: + return Status("running") + if exit_code == 0: + return Status("finished") + return Status("failed") + + +class LocalContainerRunner(AbstractRunner): + """Runner class, uses a project to create a LocallySubmittedRun.""" + + def __init__( + self, + api: "Api", + backend_config: Dict[str, Any], + environment: AbstractEnvironment, + registry: AbstractRegistry, + ) -> None: + super().__init__(api, backend_config) + self.environment = environment + self.registry = registry + + def _populate_docker_args( + self, launch_project: LaunchProject, image_uri: str + ) -> Dict[str, Any]: + docker_args: Dict[str, Any] = launch_project.fill_macros(image_uri).get( + "local-container", {} + ) + if _is_wandb_local_uri(self._api.settings("base_url")): + if sys.platform == "win32": + docker_args["net"] = "host" + else: + docker_args["network"] = "host" + if sys.platform == "linux" or sys.platform == "linux2": + docker_args["add-host"] = "host.docker.internal:host-gateway" + base_image = launch_project.job_base_image + if base_image is not None: + # Mount code into the container and set the working directory. + if "volume" not in docker_args: + docker_args["volume"] = [] + docker_args["volume"].append( + f"{launch_project.project_dir}:{CODE_MOUNT_DIR}" + ) + docker_args["workdir"] = CODE_MOUNT_DIR + return docker_args + + async def run( + self, + launch_project: LaunchProject, + image_uri: str, + ) -> Optional[AbstractRun]: + docker_args = self._populate_docker_args(launch_project, image_uri) + synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS] + + env_vars = launch_project.get_env_vars_dict( + self._api, MAX_ENV_LENGTHS[self.__class__.__name__] + ) + + # When running against local port, need to swap to local docker host + if ( + _is_wandb_local_uri(self._api.settings("base_url")) + and sys.platform == "darwin" + ): + _, _, port = self._api.settings("base_url").split(":") + env_vars["WANDB_BASE_URL"] = f"http://host.docker.internal:{port}" + elif _is_wandb_dev_uri(self._api.settings("base_url")): + env_vars["WANDB_BASE_URL"] = "http://host.docker.internal:9001" + + if launch_project.docker_image or launch_project.job_base_image: + try: + pull_docker_image(image_uri) + except Exception as e: + wandb.termwarn(f"Error attempting to pull docker image {image_uri}") + if not docker_image_exists(image_uri): + raise LaunchError( + f"Failed to pull docker image {image_uri} with error: {e}" + ) + + entrypoint = launch_project.get_job_entry_point() + entry_cmd = None if entrypoint is None else entrypoint.command + command_str = " ".join( + get_docker_command( + image_uri, + env_vars, + docker_args=docker_args, + entry_cmd=entry_cmd, + additional_args=launch_project.override_args, + ) + ).strip() + sanitized_cmd_str = sanitize_wandb_api_key(command_str) + _msg = f"{LOG_PREFIX}Launching run in docker with command: {sanitized_cmd_str}" + wandb.termlog(_msg) + run = _run_entry_point(command_str, launch_project.project_dir) + if synchronous: + await run.wait() + return run + + +def _run_entry_point(command: str, work_dir: Optional[str]) -> AbstractRun: + """Run an entry point command in a subprocess. + + Arguments: + command: Entry point command to run + work_dir: Working directory in which to run the command + + Returns: + An instance of `LocalSubmittedRun` + """ + if work_dir is None: + work_dir = os.getcwd() + env = os.environ.copy() + run = LocalSubmittedRun() + thread = threading.Thread( + target=_thread_process_runner, + args=(run, ["bash", "-c", command], work_dir, env), + ) + run.set_thread(thread) + thread.start() + return run + + +def _thread_process_runner( + run: LocalSubmittedRun, args: List[str], work_dir: str, env: Dict[str, str] +) -> None: + # cancel was called before we started the subprocess + if run._terminate_flag: + return + # TODO: Make this async + process = subprocess.Popen( + args, + close_fds=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + bufsize=1, + cwd=work_dir, + env=env, + ) + run.set_command_proc(process) + run._stdout = "" + while True: + # the agent thread could set the terminate flag + if run._terminate_flag: + process.terminate() # type: ignore + chunk = os.read(process.stdout.fileno(), 4096) # type: ignore + if not chunk: + break + index = chunk.find(b"\r") + decoded_chunk = None + while not decoded_chunk: + try: + decoded_chunk = chunk.decode() + except UnicodeDecodeError: + # Multi-byte character cut off, try to get the rest of it + chunk += os.read(process.stdout.fileno(), 1) # type: ignore + if index != -1: + run._stdout += decoded_chunk + print(chunk.decode(), end="") + else: + run._stdout += decoded_chunk + "\r" + print(chunk.decode(), end="\r") + + +def get_docker_command( + image: str, + env_vars: Dict[str, str], + entry_cmd: Optional[List[str]] = None, + docker_args: Optional[Dict[str, Any]] = None, + additional_args: Optional[List[str]] = None, +) -> List[str]: + """Construct the docker command using the image and docker args. + + Arguments: + image: a Docker image to be run + env_vars: a dictionary of environment variables for the command + entry_cmd: the entry point command to run + docker_args: a dictionary of additional docker args for the command + """ + docker_path = "docker" + cmd: List[Any] = [docker_path, "run", "--rm"] + + # hacky handling of env vars, needs to be improved + for env_key, env_value in env_vars.items(): + cmd += ["-e", f"{shlex.quote(env_key)}={shlex.quote(env_value)}"] + + if docker_args: + for name, value in docker_args.items(): + if len(name) == 1: + prefix = "-" + shlex.quote(name) + else: + prefix = "--" + shlex.quote(name) + if isinstance(value, list): + for v in value: + cmd += [prefix, shlex.quote(str(v))] + elif isinstance(value, bool) and value: + cmd += [prefix] + else: + cmd += [prefix, shlex.quote(str(value))] + + if entry_cmd: + cmd += ["--entrypoint", entry_cmd[0]] + cmd += [shlex.quote(image)] + if entry_cmd and len(entry_cmd) > 1: + cmd += entry_cmd[1:] + if additional_args: + cmd += additional_args + return cmd + + +def join(split_command: List[str]) -> str: + """Return a shell-escaped string from *split_command*.""" + return " ".join(shlex.quote(arg) for arg in split_command) diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/local_process.py b/lib/python3.12/site-packages/wandb/sdk/launch/runner/local_process.py new file mode 100644 index 0000000000000000000000000000000000000000..5b04361dc2501862f93eea3b8ef7c68807dcbf9f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/runner/local_process.py @@ -0,0 +1,78 @@ +import logging +import shlex +from typing import Any, List, Optional + +import wandb + +from .._project_spec import LaunchProject +from ..errors import LaunchError +from ..utils import ( + LOG_PREFIX, + MAX_ENV_LENGTHS, + PROJECT_SYNCHRONOUS, + sanitize_wandb_api_key, + validate_wandb_python_deps, +) +from .abstract import AbstractRun, AbstractRunner +from .local_container import _run_entry_point + +_logger = logging.getLogger(__name__) + + +class LocalProcessRunner(AbstractRunner): + """Runner class, uses a project to create a LocallySubmittedRun. + + LocalProcessRunner is very similar to a LocalContainerRunner, except it does not + run the command inside a docker container. Instead, it runs the + command specified as a process directly on the bare metal machine. + + """ + + async def run( # type: ignore + self, + launch_project: LaunchProject, + *args, + **kwargs, + ) -> Optional[AbstractRun]: + if args is not None: + _msg = f"{LOG_PREFIX}LocalProcessRunner.run received unused args {args}" + _logger.warning(_msg) + if kwargs is not None: + _msg = f"{LOG_PREFIX}LocalProcessRunner.run received unused kwargs {kwargs}" + _logger.warning(_msg) + + synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS] + entry_point = ( + launch_project.override_entrypoint or launch_project.get_job_entry_point() + ) + + cmd: List[Any] = [] + + if launch_project.project_dir is None: + raise LaunchError("Launch LocalProcessRunner received empty project dir") + + if launch_project.job: + assert launch_project._job_artifact is not None + try: + validate_wandb_python_deps( + "requirements.frozen.txt", + launch_project.project_dir, + ) + except Exception: + wandb.termwarn("Unable to validate python dependencies") + env_vars = launch_project.get_env_vars_dict( + self._api, MAX_ENV_LENGTHS[self.__class__.__name__] + ) + for env_key, env_value in env_vars.items(): + cmd += [f"{shlex.quote(env_key)}={shlex.quote(env_value)}"] + if entry_point is not None: + cmd += entry_point.command + cmd += launch_project.override_args + + command_str = " ".join(cmd).strip() + _msg = f"{LOG_PREFIX}Launching run as a local-process with command {sanitize_wandb_api_key(command_str)}" + wandb.termlog(_msg) + run = _run_entry_point(command_str, launch_project.project_dir) + if synchronous: + await run.wait() + return run diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/sagemaker_runner.py b/lib/python3.12/site-packages/wandb/sdk/launch/runner/sagemaker_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..dfd5947e0d82846773410b2486e3ada46b333558 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/runner/sagemaker_runner.py @@ -0,0 +1,424 @@ +"""Implementation of the SageMakerRunner class.""" + +import asyncio +import logging +from typing import Any, Dict, List, Optional, cast + +if False: + import boto3 # type: ignore + +import wandb +from wandb.apis.internal import Api +from wandb.sdk.launch.environment.aws_environment import AwsEnvironment +from wandb.sdk.launch.errors import LaunchError + +from .._project_spec import EntryPoint, LaunchProject +from ..registry.abstract import AbstractRegistry +from ..utils import ( + LOG_PREFIX, + MAX_ENV_LENGTHS, + PROJECT_SYNCHRONOUS, + event_loop_thread_exec, + to_camel_case, +) +from .abstract import AbstractRun, AbstractRunner, Status + +_logger = logging.getLogger(__name__) + + +class SagemakerSubmittedRun(AbstractRun): + """Instance of ``AbstractRun`` corresponding to a subprocess launched to run an entry point command on aws sagemaker.""" + + def __init__( + self, + training_job_name: str, + client: "boto3.Client", + log_client: Optional["boto3.Client"] = None, + ) -> None: + super().__init__() + self.client = client + self.log_client = log_client + self.training_job_name = training_job_name + self._status = Status("running") + + @property + def id(self) -> str: + return f"sagemaker-{self.training_job_name}" + + async def get_logs(self) -> Optional[str]: + if self.log_client is None: + return None + try: + describe_log_streams = event_loop_thread_exec( + self.log_client.describe_log_streams + ) + describe_res = await describe_log_streams( + logGroupName="/aws/sagemaker/TrainingJobs", + logStreamNamePrefix=self.training_job_name, + ) + if len(describe_res["logStreams"]) == 0: + wandb.termwarn( + f"Failed to get logs for training job: {self.training_job_name}" + ) + return None + log_name = describe_res["logStreams"][0]["logStreamName"] + get_log_events = event_loop_thread_exec(self.log_client.get_log_events) + res = await get_log_events( + logGroupName="/aws/sagemaker/TrainingJobs", + logStreamName=log_name, + ) + assert "events" in res + return "\n".join( + [f"{event['timestamp']}:{event['message']}" for event in res["events"]] + ) + except self.log_client.exceptions.ResourceNotFoundException: + wandb.termwarn( + f"Failed to get logs for training job: {self.training_job_name}" + ) + return None + except Exception as e: + wandb.termwarn( + f"Failed to handle logs for training job: {self.training_job_name} with error {str(e)}" + ) + return None + + async def wait(self) -> bool: + while True: + status_state = (await self.get_status()).state + wandb.termlog( + f"{LOG_PREFIX}Training job {self.training_job_name} status: {status_state}" + ) + if status_state in ["stopped", "failed", "finished"]: + break + await asyncio.sleep(5) + return status_state == "finished" + + async def cancel(self) -> None: + # Interrupt child process if it hasn't already exited + status = await self.get_status() + if status.state == "running": + self.client.stop_training_job(TrainingJobName=self.training_job_name) + await self.wait() + + async def get_status(self) -> Status: + describe_training_job = event_loop_thread_exec( + self.client.describe_training_job + ) + job_status = ( + await describe_training_job(TrainingJobName=self.training_job_name) + )["TrainingJobStatus"] + if job_status == "Completed" or job_status == "Stopped": + self._status = Status("finished") + elif job_status == "Failed": + self._status = Status("failed") + elif job_status == "Stopping": + self._status = Status("stopping") + elif job_status == "InProgress": + self._status = Status("running") + return self._status + + +class SageMakerRunner(AbstractRunner): + """Runner class, uses a project to create a SagemakerSubmittedRun.""" + + def __init__( + self, + api: Api, + backend_config: Dict[str, Any], + environment: AwsEnvironment, + registry: AbstractRegistry, + ) -> None: + """Initialize the SagemakerRunner. + + Arguments: + api (Api): The API instance. + backend_config (Dict[str, Any]): The backend configuration. + environment (AwsEnvironment): The AWS environment. + + Raises: + LaunchError: If the runner cannot be initialized. + """ + super().__init__(api, backend_config) + self.environment = environment + self.registry = registry + + async def run( + self, + launch_project: LaunchProject, + image_uri: str, + ) -> Optional[AbstractRun]: + """Run a project on Amazon Sagemaker. + + Arguments: + launch_project (LaunchProject): The project to run. + + Returns: + Optional[AbstractRun]: The run instance. + + Raises: + LaunchError: If the launch is unsuccessful. + """ + _logger.info("using AWSSagemakerRunner") + + given_sagemaker_args = launch_project.resource_args.get("sagemaker") + if given_sagemaker_args is None: + raise LaunchError( + "No sagemaker args specified. Specify sagemaker args in resource_args" + ) + + default_output_path = self.backend_config.get("runner", {}).get( + "s3_output_path" + ) + if default_output_path is not None and not default_output_path.startswith( + "s3://" + ): + default_output_path = f"s3://{default_output_path}" + + session = await self.environment.get_session() + client = await event_loop_thread_exec(session.client)("sts") + caller_id = client.get_caller_identity() + account_id = caller_id["Account"] + _logger.info(f"Using account ID {account_id}") + partition = await self.environment.get_partition() + role_arn = get_role_arn( + given_sagemaker_args, self.backend_config, account_id, partition + ) + + # Create a sagemaker client to launch the job. + sagemaker_client = session.client("sagemaker") + log_client = None + try: + log_client = session.client("logs") + except Exception as e: + wandb.termwarn( + f"Failed to connect to cloudwatch logs with error {str(e)}, logs will not be available" + ) + + # if the user provided the image they want to use, use that, but warn it won't have swappable artifacts + if ( + given_sagemaker_args.get("AlgorithmSpecification", {}).get("TrainingImage") + is not None + ): + sagemaker_args = build_sagemaker_args( + launch_project, + self._api, + role_arn, + launch_project.override_entrypoint, + launch_project.override_args, + MAX_ENV_LENGTHS[self.__class__.__name__], + given_sagemaker_args.get("AlgorithmSpecification", {}).get( + "TrainingImage" + ), + default_output_path, + ) + _logger.info( + f"Launching sagemaker job on user supplied image with args: {sagemaker_args}" + ) + run = await launch_sagemaker_job( + launch_project, sagemaker_args, sagemaker_client, log_client + ) + if self.backend_config[PROJECT_SYNCHRONOUS]: + await run.wait() + return run + + _logger.info("Connecting to sagemaker client") + entry_point = ( + launch_project.override_entrypoint or launch_project.get_job_entry_point() + ) + command_args = [] + if entry_point is not None: + command_args += entry_point.command + command_args += launch_project.override_args + if command_args: + command_str = " ".join(command_args) + wandb.termlog( + f"{LOG_PREFIX}Launching run on sagemaker with entrypoint: {command_str}" + ) + else: + wandb.termlog( + f"{LOG_PREFIX}Launching run on sagemaker with user-provided entrypoint in image" + ) + sagemaker_args = build_sagemaker_args( + launch_project, + self._api, + role_arn, + entry_point, + launch_project.override_args, + MAX_ENV_LENGTHS[self.__class__.__name__], + image_uri, + default_output_path, + ) + _logger.info(f"Launching sagemaker job with args: {sagemaker_args}") + run = await launch_sagemaker_job( + launch_project, sagemaker_args, sagemaker_client, log_client + ) + if self.backend_config[PROJECT_SYNCHRONOUS]: + await run.wait() + return run + + +def merge_image_uri_with_algorithm_specification( + algorithm_specification: Optional[Dict[str, Any]], + image_uri: Optional[str], + entrypoint_command: List[str], + args: Optional[List[str]], +) -> Dict[str, Any]: + """Create an AWS AlgorithmSpecification. + + AWS Sagemaker algorithms require a training image and an input mode. If the user + does not specify the specification themselves, define the spec minimally using these + two fields. Otherwise, if they specify the AlgorithmSpecification set the training + image if it is not set. + """ + if algorithm_specification is None: + algorithm_specification = { + "TrainingImage": image_uri, + "TrainingInputMode": "File", + } + else: + if image_uri: + algorithm_specification["TrainingImage"] = image_uri + if entrypoint_command: + algorithm_specification["ContainerEntrypoint"] = entrypoint_command + if args: + algorithm_specification["ContainerArguments"] = args + + if algorithm_specification["TrainingImage"] is None: + raise LaunchError("Failed determine tag for training image") + return algorithm_specification + + +def build_sagemaker_args( + launch_project: LaunchProject, + api: Api, + role_arn: str, + entry_point: Optional[EntryPoint], + args: Optional[List[str]], + max_env_length: int, + image_uri: str, + default_output_path: Optional[str] = None, +) -> Dict[str, Any]: + sagemaker_args: Dict[str, Any] = {} + resource_args = launch_project.fill_macros(image_uri) + given_sagemaker_args: Optional[Dict[str, Any]] = resource_args.get("sagemaker") + + if given_sagemaker_args is None: + raise LaunchError( + "No sagemaker args specified. Specify sagemaker args in resource_args" + ) + if ( + given_sagemaker_args.get("OutputDataConfig") is None + and default_output_path is not None + ): + sagemaker_args["OutputDataConfig"] = {"S3OutputPath": default_output_path} + else: + sagemaker_args["OutputDataConfig"] = given_sagemaker_args.get( + "OutputDataConfig" + ) + + if sagemaker_args.get("OutputDataConfig") is None: + raise LaunchError( + "Sagemaker launcher requires an OutputDataConfig Sagemaker resource argument" + ) + training_job_name = cast( + str, (given_sagemaker_args.get("TrainingJobName") or launch_project.run_id) + ) + sagemaker_args["TrainingJobName"] = training_job_name + entry_cmd = entry_point.command if entry_point else [] + + sagemaker_args["AlgorithmSpecification"] = ( + merge_image_uri_with_algorithm_specification( + given_sagemaker_args.get( + "AlgorithmSpecification", + given_sagemaker_args.get("algorithm_specification"), + ), + image_uri, + entry_cmd, + args, + ) + ) + + sagemaker_args["RoleArn"] = role_arn + + camel_case_args = { + to_camel_case(key): item for key, item in given_sagemaker_args.items() + } + sagemaker_args = { + **camel_case_args, + **sagemaker_args, + } + + if sagemaker_args.get("ResourceConfig") is None: + raise LaunchError( + "Sagemaker launcher requires a ResourceConfig resource argument" + ) + + if sagemaker_args.get("StoppingCondition") is None: + raise LaunchError( + "Sagemaker launcher requires a StoppingCondition resource argument" + ) + + given_env = given_sagemaker_args.get( + "Environment", sagemaker_args.get("environment", {}) + ) + calced_env = launch_project.get_env_vars_dict(api, max_env_length) + total_env = {**calced_env, **given_env} + sagemaker_args["Environment"] = total_env + + # Add wandb tag + tags = sagemaker_args.get("Tags", []) + tags.append({"Key": "WandbRunId", "Value": launch_project.run_id}) + sagemaker_args["Tags"] = tags + + # remove args that were passed in for launch but not passed to sagemaker + sagemaker_args.pop("EcrRepoName", None) + sagemaker_args.pop("region", None) + sagemaker_args.pop("profile", None) + + # clear the args that are None so they are not passed + filtered_args = {k: v for k, v in sagemaker_args.items() if v is not None} + + return filtered_args + + +async def launch_sagemaker_job( + launch_project: LaunchProject, + sagemaker_args: Dict[str, Any], + sagemaker_client: "boto3.Client", + log_client: Optional["boto3.Client"] = None, +) -> SagemakerSubmittedRun: + training_job_name = sagemaker_args.get("TrainingJobName") or launch_project.run_id + create_training_job = event_loop_thread_exec(sagemaker_client.create_training_job) + resp = await create_training_job(**sagemaker_args) + + if resp.get("TrainingJobArn") is None: + raise LaunchError("Failed to create training job when submitting to SageMaker") + + run = SagemakerSubmittedRun(training_job_name, sagemaker_client, log_client) + wandb.termlog( + f"{LOG_PREFIX}Run job submitted with arn: {resp.get('TrainingJobArn')}" + ) + url = f"https://{sagemaker_client.meta.region_name}.console.aws.amazon.com/sagemaker/home?region={sagemaker_client.meta.region_name}#/jobs/{training_job_name}" + wandb.termlog(f"{LOG_PREFIX}See training job status at: {url}") + return run + + +def get_role_arn( + sagemaker_args: Dict[str, Any], + backend_config: Dict[str, Any], + account_id: str, + partition: str, +) -> str: + """Get the role arn from the sagemaker args or the backend config.""" + role_arn = sagemaker_args.get("RoleArn") or sagemaker_args.get("role_arn") + if role_arn is None: + role_arn = backend_config.get("runner", {}).get("role_arn") + if role_arn is None or not isinstance(role_arn, str): + raise LaunchError( + "AWS sagemaker require a string RoleArn set this by adding a `RoleArn` key to the sagemaker" + "field of resource_args" + ) + if role_arn.startswith(f"arn:{partition}:iam::"): + return role_arn # type: ignore + + return f"arn:{partition}:iam::{account_id}:role/{role_arn}" diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/runner/vertex_runner.py b/lib/python3.12/site-packages/wandb/sdk/launch/runner/vertex_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..b3254b92b0c9289329c3b0e64d380a1068d44aae --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/runner/vertex_runner.py @@ -0,0 +1,225 @@ +import asyncio +import logging +from typing import Any, Dict, Optional + +if False: + from google.cloud import aiplatform # type: ignore # noqa: F401 + +from wandb.apis.internal import Api +from wandb.util import get_module + +from .._project_spec import LaunchProject +from ..environment.gcp_environment import GcpEnvironment +from ..errors import LaunchError +from ..registry.abstract import AbstractRegistry +from ..utils import MAX_ENV_LENGTHS, PROJECT_SYNCHRONOUS, event_loop_thread_exec +from .abstract import AbstractRun, AbstractRunner, Status + +GCP_CONSOLE_URI = "https://console.cloud.google.com" + +_logger = logging.getLogger(__name__) + + +WANDB_RUN_ID_KEY = "wandb-run-id" + + +class VertexSubmittedRun(AbstractRun): + def __init__(self, job: Any) -> None: + self._job = job + + @property + def id(self) -> str: + # numeric ID of the custom training job + return self._job.name # type: ignore + + async def get_logs(self) -> Optional[str]: + # TODO: implement + return None + + @property + def name(self) -> str: + return self._job.display_name # type: ignore + + @property + def gcp_region(self) -> str: + return self._job.location # type: ignore + + @property + def gcp_project(self) -> str: + return self._job.project # type: ignore + + def get_page_link(self) -> str: + return f"{GCP_CONSOLE_URI}/vertex-ai/locations/{self.gcp_region}/training/{self.id}?project={self.gcp_project}" + + async def wait(self) -> bool: + # TODO: run this in a separate thread. + await self._job.wait() + return (await self.get_status()).state == "finished" + + async def get_status(self) -> Status: + job_state = str(self._job.state) # extract from type PipelineState + if job_state == "JobState.JOB_STATE_SUCCEEDED": + return Status("finished") + if job_state == "JobState.JOB_STATE_FAILED": + return Status("failed") + if job_state == "JobState.JOB_STATE_RUNNING": + return Status("running") + if job_state == "JobState.JOB_STATE_PENDING": + return Status("starting") + return Status("unknown") + + async def cancel(self) -> None: + self._job.cancel() + + +class VertexRunner(AbstractRunner): + """Runner class, uses a project to create a VertexSubmittedRun.""" + + def __init__( + self, + api: Api, + backend_config: Dict[str, Any], + environment: GcpEnvironment, + registry: AbstractRegistry, + ) -> None: + """Initialize a VertexRunner instance.""" + super().__init__(api, backend_config) + self.environment = environment + self.registry = registry + + async def run( + self, launch_project: LaunchProject, image_uri: str + ) -> Optional[AbstractRun]: + """Run a Vertex job.""" + full_resource_args = launch_project.fill_macros(image_uri) + resource_args = full_resource_args.get("vertex") + # We support setting under gcp-vertex for historical reasons. + if not resource_args: + resource_args = full_resource_args.get("gcp-vertex") + if not resource_args: + raise LaunchError( + "No Vertex resource args specified. Specify args via --resource-args with a JSON file or string under top-level key gcp_vertex" + ) + + spec_args = resource_args.get("spec", {}) + run_args = resource_args.get("run", {}) + + synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS] + + entry_point = ( + launch_project.override_entrypoint or launch_project.get_job_entry_point() + ) + + # TODO: Set entrypoint in each container + entry_cmd = [] + if entry_point is not None: + entry_cmd += entry_point.command + entry_cmd += launch_project.override_args + + env_vars = launch_project.get_env_vars_dict( + api=self._api, + max_env_length=MAX_ENV_LENGTHS[self.__class__.__name__], + ) + + worker_specs = spec_args.get("worker_pool_specs", []) + if not worker_specs: + raise LaunchError( + "Vertex requires at least one worker pool spec. Please specify " + "a worker pool spec in resource arguments under the key " + "`vertex.spec.worker_pool_specs`." + ) + + # TODO: Add entrypoint + args to each worker pool spec + for spec in worker_specs: + if not spec.get("container_spec"): + raise LaunchError( + "Vertex requires a container spec for each worker pool spec. " + "Please specify a container spec in resource arguments under " + "the key `vertex.spec.worker_pool_specs[].container_spec`." + ) + spec["container_spec"]["command"] = entry_cmd + + # Add our env vars to user supplied env vars + env = spec["container_spec"].get("env", []) + env.extend( + [{"name": key, "value": value} for key, value in env_vars.items()] + ) + spec["container_spec"]["env"] = env + + if not spec_args.get("staging_bucket"): + raise LaunchError( + "Vertex requires a staging bucket. Please specify a staging bucket " + "in resource arguments under the key `vertex.spec.staging_bucket`." + ) + + _logger.info("Launching Vertex job...") + submitted_run = await launch_vertex_job( + launch_project, + spec_args, + run_args, + self.environment, + synchronous, + ) + return submitted_run + + +async def launch_vertex_job( + launch_project: LaunchProject, + spec_args: Dict[str, Any], + run_args: Dict[str, Any], + environment: GcpEnvironment, + synchronous: bool = False, +) -> VertexSubmittedRun: + try: + await environment.verify() + aiplatform = get_module( + "google.cloud.aiplatform", + "VertexRunner requires google.cloud.aiplatform to be installed", + ) + init = event_loop_thread_exec(aiplatform.init) + await init( + project=environment.project, + location=environment.region, + staging_bucket=spec_args.get("staging_bucket"), + credentials=await environment.get_credentials(), + ) + labels = spec_args.get("labels", {}) + labels[WANDB_RUN_ID_KEY] = launch_project.run_id + job = aiplatform.CustomJob( + display_name=launch_project.name, + worker_pool_specs=spec_args.get("worker_pool_specs"), + base_output_dir=spec_args.get("base_output_dir"), + encryption_spec_key_name=spec_args.get("encryption_spec_key_name"), + labels=labels, + ) + execution_kwargs = dict( + timeout=run_args.get("timeout"), + service_account=run_args.get("service_account"), + network=run_args.get("network"), + enable_web_access=run_args.get("enable_web_access", False), + experiment=run_args.get("experiment"), + experiment_run=run_args.get("experiment_run"), + tensorboard=run_args.get("tensorboard"), + restart_job_on_worker_restart=run_args.get( + "restart_job_on_worker_restart", False + ), + ) + # Unclear if there are exceptions that can be thrown where we should + # retry instead of erroring. For now, just catch all exceptions and they + # go to the UI for the user to interpret. + except Exception as e: + raise LaunchError(f"Failed to create Vertex job: {e}") + + if synchronous: + run = event_loop_thread_exec(job.run) + await run(**execution_kwargs, sync=True) + else: + submit = event_loop_thread_exec(job.submit) + await submit(**execution_kwargs) + submitted_run = VertexSubmittedRun(job) + interval = 1 + while not getattr(job._gca_resource, "name", None): + # give time for the gcp job object to be created and named, this should only loop a couple times max + await asyncio.sleep(interval) + interval = min(30, interval * 2) + return submitted_run diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__init__.py b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..98be169a096015729752916d4ac6763cf7a65024 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__init__.py @@ -0,0 +1,37 @@ +import logging +from typing import Any, Callable, Dict + +log = logging.getLogger(__name__) + + +class SchedulerError(Exception): + """Raised when a known error occurs with wandb sweep scheduler.""" + + +def _import_sweep_scheduler() -> Any: + from .scheduler_sweep import SweepScheduler + + return SweepScheduler + + +_WANDB_SCHEDULERS: Dict[str, Callable] = { + "wandb": _import_sweep_scheduler, +} + + +def load_scheduler(scheduler_type: str) -> Any: + scheduler_type = scheduler_type.lower() + if scheduler_type not in _WANDB_SCHEDULERS: + raise SchedulerError( + f"The `scheduler_name` argument must be one of " + f"{list(_WANDB_SCHEDULERS.keys())}, got: {scheduler_type}" + ) + + log.warning(f"Loading dependencies for Scheduler of type: {scheduler_type}") + import_func = _WANDB_SCHEDULERS[scheduler_type] + return import_func() + + +__all__ = [ + "load_scheduler", +] diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f0499ca9c7df207982ace53ad933cbd9cb180384 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/scheduler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ac114217081942ab5631a5a81309bbabda692da Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/scheduler_sweep.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/scheduler_sweep.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50d7cac3ad573296fbe64bb8d9476a6f763f23cf Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/scheduler_sweep.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a052ca18a183a466b7f0a88e56f28489585ed76 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/scheduler.py b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..bc0e6d20bdbb5422545b40239c1b50c936f6ac17 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/scheduler.py @@ -0,0 +1,740 @@ +"""Abstract Scheduler class.""" + +import asyncio +import base64 +import copy +import logging +import os +import socket +import threading +import time +import traceback +from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Union + +import click +import yaml + +import wandb +from wandb.errors import CommError +from wandb.sdk.launch._launch_add import launch_add +from wandb.sdk.launch.errors import LaunchError +from wandb.sdk.launch.sweeps import SchedulerError +from wandb.sdk.launch.sweeps.utils import ( + create_sweep_command_args, + make_launch_sweep_entrypoint, +) +from wandb.sdk.launch.utils import ( + event_loop_thread_exec, + strip_resource_args_and_template_vars, +) +from wandb.sdk.lib.runid import generate_id + +if TYPE_CHECKING: + import wandb.apis.public as public + from wandb.apis.internal import Api + from wandb.apis.public import QueuedRun, Run + from wandb.sdk.wandb_run import Run as SdkRun + + +_logger = logging.getLogger(__name__) +LOG_PREFIX = f"{click.style('sched:', fg='cyan')} " + +DEFAULT_POLLING_SLEEP = 5.0 + + +class SchedulerState(Enum): + PENDING = 0 + STARTING = 1 + RUNNING = 2 + FLUSH_RUNS = 3 + COMPLETED = 4 + FAILED = 5 + STOPPED = 6 + CANCELLED = 7 + + +class RunState(Enum): + RUNNING = "running", "alive" + PENDING = "pending", "alive" + PREEMPTING = "preempting", "alive" + CRASHED = "crashed", "dead" + FAILED = "failed", "dead" + KILLED = "killed", "dead" + FINISHED = "finished", "dead" + PREEMPTED = "preempted", "dead" + # unknown when api.get_run_state fails or returns unexpected state + # assumed alive, unless we get unknown 2x then move to failed (dead) + UNKNOWN = "unknown", "alive" + + def __new__(cls: Any, *args: List, **kwds: Any) -> "RunState": + obj: RunState = object.__new__(cls) + obj._value_ = args[0] + return obj + + def __init__(self, _: str, life: str = "unknown") -> None: + self._life = life + + @property + def is_alive(self) -> bool: + return self._life == "alive" + + +@dataclass +class _Worker: + agent_config: Dict[str, Any] + agent_id: str + + +@dataclass +class SweepRun: + id: str + worker_id: int + state: RunState = RunState.RUNNING + queued_run: Optional["public.QueuedRun"] = None + args: Optional[Dict[str, Any]] = None + logs: Optional[List[str]] = None + + +class Scheduler(ABC): + """A controller/agent that populates a Launch RunQueue from a hyperparameter sweep.""" + + PLACEHOLDER_URI = "placeholder-uri-scheduler" + SWEEP_JOB_TYPE = "sweep-controller" + ENTRYPOINT = ["wandb", "scheduler", "WANDB_SWEEP_ID"] + + def __init__( + self, + api: "Api", + *args: Optional[Any], + polling_sleep: Optional[float] = None, + sweep_id: Optional[str] = None, + entity: Optional[str] = None, + project: Optional[str] = None, + project_queue: Optional[str] = None, + num_workers: Optional[Union[int, str]] = None, + **kwargs: Optional[Any], + ): + from wandb.apis.public import Api as PublicApi + + self._api = api + self._public_api = PublicApi() + self._entity = ( + entity + or os.environ.get("WANDB_ENTITY") + or api.settings("entity") + or api.default_entity + ) + self._project = ( + project or os.environ.get("WANDB_PROJECT") or api.settings("project") + ) + self._sweep_id: str = sweep_id or "empty-sweep-id" + self._state: SchedulerState = SchedulerState.PENDING + + # Make sure the provided sweep_id corresponds to a valid sweep + try: + resp = self._api.sweep( + sweep_id, "{}", entity=self._entity, project=self._project + ) + if resp.get("state") == SchedulerState.CANCELLED.name: + self._state = SchedulerState.CANCELLED + self._sweep_config = yaml.safe_load(resp["config"]) + self._num_runs_launched: int = self._get_num_runs_launched(resp["runs"]) + if self._num_runs_launched > 0: + wandb.termlog( + f"{LOG_PREFIX}Found {self._num_runs_launched} previous valid runs for sweep {self._sweep_id}" + ) + except Exception as e: + raise SchedulerError( + f"{LOG_PREFIX}Exception when finding sweep ({sweep_id}) {e}" + ) + + # Scheduler may receive additional kwargs which will be piped into the launch command + self._kwargs: Dict[str, Any] = kwargs + + # Dictionary of the runs being managed by the scheduler + self._runs: Dict[str, SweepRun] = {} + # Threading lock to ensure thread-safe access to the runs dictionary + self._threading_lock: threading.Lock = threading.Lock() + self._polling_sleep = ( + polling_sleep if polling_sleep is not None else DEFAULT_POLLING_SLEEP + ) + self._project_queue = project_queue + # Optionally run multiple workers in (pseudo-)parallel. Workers do not + # actually run training workloads, they simply send heartbeat messages + # (emulating a real agent) and add new runs to the launch queue. The + # launch agent is the one that actually runs the training workloads. + self._workers: Dict[int, _Worker] = {} + + # Init wandb scheduler run + self._wandb_run = self._init_wandb_run() + + # Grab params from scheduler wandb run config + num_workers = num_workers or self._wandb_run.config.get("scheduler", {}).get( + "num_workers" + ) + self._num_workers = int(num_workers) if str(num_workers).isdigit() else 8 + self._settings_config: Dict[str, Any] = self._wandb_run.config.get( + "settings", {} + ) + + @abstractmethod + def _get_next_sweep_run(self, worker_id: int) -> Optional[SweepRun]: + """Called when worker available.""" + + @abstractmethod + def _poll(self) -> None: + """Called every polling loop.""" + + @abstractmethod + def _exit(self) -> None: + pass + + @abstractmethod + def _load_state(self) -> None: + pass + + @abstractmethod + def _save_state(self) -> None: + pass + + @property + def state(self) -> SchedulerState: + _logger.debug(f"{LOG_PREFIX}Scheduler state is {self._state.name}") + return self._state + + @state.setter + def state(self, value: SchedulerState) -> None: + _logger.debug(f"{LOG_PREFIX}Scheduler was {self.state.name} is {value.name}") + self._state = value + + @property + def is_alive(self) -> bool: + if self.state in [ + SchedulerState.COMPLETED, + SchedulerState.FAILED, + SchedulerState.STOPPED, + SchedulerState.CANCELLED, + ]: + return False + return True + + @property + def at_runcap(self) -> bool: + """False if under user-specified cap on # of runs.""" + run_cap = self._sweep_config.get("run_cap") + if not run_cap: + return False + at_runcap: bool = self._num_runs_launched >= run_cap + return at_runcap + + @property + def num_active_runs(self) -> int: + return len(self._runs) + + @property + def busy_workers(self) -> Dict[int, _Worker]: + """Returns dict of id:worker already assigned to a launch run. + + runs should always have a worker_id, but are created before + workers are assigned to the run + """ + busy_workers = {} + for _, r in self._yield_runs(): + busy_workers[r.worker_id] = self._workers[r.worker_id] + return busy_workers + + @property + def available_workers(self) -> Dict[int, _Worker]: + """Returns dict of id:worker ready to launch another run.""" + if len(self._workers) == 0: + return {} + return { + _id: w for _id, w in self._workers.items() if _id not in self.busy_workers + } + + 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 + + def stop_sweep(self) -> None: + """Stop the sweep.""" + self._state = SchedulerState.STOPPED + + def fail_sweep(self, err: Optional[str]) -> None: + """Fail the sweep w/ optional exception.""" + self._state = SchedulerState.FAILED + if err: + raise SchedulerError(err) + + def start(self) -> None: + """Start a scheduler, confirms prerequisites, begins execution loop.""" + wandb.termlog(f"{LOG_PREFIX}Scheduler starting.") + if not self.is_alive: + wandb.termerror( + f"{LOG_PREFIX}Sweep already in end state ({self.state.name.lower()}). Exiting..." + ) + self.exit() + return + + self._state = SchedulerState.STARTING + if not self._try_load_executable(): + wandb.termerror( + f"{LOG_PREFIX}No 'job' or 'image_uri' loaded from sweep config." + ) + self.exit() + return + + # For resuming sweeps + self._load_state() + asyncio.run(self._register_agents()) + self.run() + + def run(self) -> None: + """Main run function.""" + wandb.termlog(f"{LOG_PREFIX}Scheduler running") + self.state = SchedulerState.RUNNING + try: + while True: + self._update_scheduler_run_state() + if not self.is_alive: + break + + wandb.termlog(f"{LOG_PREFIX}Polling for new runs to launch") + + self._update_run_states() + self._poll() + if self.state == SchedulerState.FLUSH_RUNS: + if self.num_active_runs == 0: + wandb.termlog(f"{LOG_PREFIX}Done polling on runs, exiting") + break + time.sleep(self._polling_sleep) + continue + + for worker_id in self.available_workers: + if self.at_runcap: + wandb.termlog( + f"{LOG_PREFIX}Sweep at run_cap ({self._num_runs_launched})" + ) + self.state = SchedulerState.FLUSH_RUNS + break + + try: + run: Optional[SweepRun] = self._get_next_sweep_run(worker_id) + if not run: + break + except SchedulerError as e: + raise SchedulerError(e) + except Exception as e: + wandb.termerror( + f"{LOG_PREFIX}Failed to get next sweep run: {e}" + ) + self.state = SchedulerState.FAILED + break + + if self._add_to_launch_queue(run): + self._num_runs_launched += 1 + + time.sleep(self._polling_sleep) + except KeyboardInterrupt: + wandb.termwarn(f"{LOG_PREFIX}Scheduler received KeyboardInterrupt. Exiting") + self.state = SchedulerState.STOPPED + self.exit() + return + except Exception as e: + wandb.termlog(f"{LOG_PREFIX}Scheduler failed with exception {e}") + self.state = SchedulerState.FAILED + self.exit() + raise + else: + # scheduler succeeds if at runcap + if self.state == SchedulerState.FLUSH_RUNS and self.at_runcap: + self.state = SchedulerState.COMPLETED + self.exit() + + def exit(self) -> None: + self._exit() + # _save_state isn't controlled, possibly fails + try: + self._save_state() + except Exception: + wandb.termerror( + f"{LOG_PREFIX}Failed to save state: {traceback.format_exc()}" + ) + + status = "" + if self.state == SchedulerState.FLUSH_RUNS: + self._set_sweep_state("PAUSED") + status = "paused" + elif self.state == SchedulerState.COMPLETED: + self._set_sweep_state("FINISHED") + status = "completed" + elif self.state in [SchedulerState.CANCELLED, SchedulerState.STOPPED]: + self._set_sweep_state("CANCELED") # one L + status = "cancelled" + self._stop_runs() + else: + self.state = SchedulerState.FAILED + self._set_sweep_state("CRASHED") + status = "crashed" + self._stop_runs() + + wandb.termlog(f"{LOG_PREFIX}Scheduler {status}") + self._wandb_run.finish() + + def _get_num_runs_launched(self, runs: List[Dict[str, Any]]) -> int: + """Returns the number of valid runs in the sweep.""" + count = 0 + for run in runs: + # if bad run, shouldn't be counted against run cap + if run.get("state", "") in ["killed", "crashed"] and not run.get( + "summaryMetrics" + ): + _logger.debug( + f"excluding run: {run['name']} with state: {run['state']} from run cap \n{run}" + ) + continue + count += 1 + + return count + + def _try_load_executable(self) -> bool: + """Check existence of valid executable for a run. + + logs and returns False when job is unreachable + """ + if self._kwargs.get("job"): + try: + _job_artifact = self._public_api.job(self._kwargs["job"]) + wandb.termlog( + f"{LOG_PREFIX}Successfully loaded job ({_job_artifact.name}) in scheduler" + ) + except Exception: + wandb.termerror(f"{LOG_PREFIX}{traceback.format_exc()}") + return False + return True + elif self._kwargs.get("image_uri"): + # TODO(gst): check docker existence? Use registry in launch config? + return True + else: + return False + + async def _register_agents(self) -> None: + tasks = [] + register_agent = event_loop_thread_exec(self._api.register_agent) + for worker_id in range(self._num_workers): + _logger.debug(f"{LOG_PREFIX}Starting AgentHeartbeat worker ({worker_id})") + try: + worker = register_agent( + f"{socket.gethostname()}-{worker_id}", # host + sweep_id=self._sweep_id, + project_name=self._project, + entity=self._entity, + ) + tasks.append(worker) + except Exception as e: + _logger.debug(f"failed to register agent: {e}") + self.fail_sweep(f"failed to register agent: {e}") + + finished_tasks = await asyncio.gather(*tasks) + for idx, agent_config in enumerate(finished_tasks): + self._workers[idx] = _Worker( + agent_config=agent_config, + agent_id=agent_config["id"], + ) + + def _yield_runs(self) -> Iterator[Tuple[str, SweepRun]]: + """Thread-safe way to iterate over the runs.""" + with self._threading_lock: + yield from self._runs.items() + + def _cleanup_runs(self, runs_to_remove: List[str]) -> None: + """Helper for removing runs from memory. + + Can be overloaded to prevent deletion of runs, which is useful + for debugging or when polling on completed runs. + """ + with self._threading_lock: + for run_id in runs_to_remove: + wandb.termlog(f"{LOG_PREFIX}Cleaning up finished run ({run_id})") + del self._runs[run_id] + + def _stop_runs(self) -> None: + to_delete = [] + for run_id, _ in self._yield_runs(): + to_delete += [run_id] + + for run_id in to_delete: + wandb.termlog(f"{LOG_PREFIX}Stopping run ({run_id})") + if not self._stop_run(run_id): + wandb.termwarn(f"{LOG_PREFIX}Failed to stop run ({run_id})") + + def _stop_run(self, run_id: str) -> bool: + """Stops a run and removes it from the scheduler.""" + if run_id not in self._runs: + _logger.debug(f"run: {run_id} not in _runs: {self._runs}") + return False + + run = self._runs[run_id] + del self._runs[run_id] + + if not run.queued_run: + _logger.debug( + f"tried to _stop_run but run not queued yet (run_id:{run.id})" + ) + return False + + if not run.state.is_alive: + # run already dead, just delete reference + return True + + # run still alive, send stop signal + encoded_run_id = base64.standard_b64encode( + f"Run:v1:{run_id}:{self._project}:{self._entity}".encode() + ).decode("utf-8") + + try: + success: bool = self._api.stop_run(run_id=encoded_run_id) + if success: + wandb.termlog(f"{LOG_PREFIX}Stopped run {run_id}.") + return True + except Exception as e: + _logger.debug(f"error stopping run ({run_id}): {e}") + + return False + + def _update_scheduler_run_state(self) -> None: + """Update the scheduler state from state of scheduler run and sweep state.""" + state: RunState = self._get_run_state(self._wandb_run.id) + + # map scheduler run-state to scheduler-state + if state == RunState.KILLED: + self.state = SchedulerState.STOPPED + elif state in [RunState.FAILED, RunState.CRASHED]: + self.state = SchedulerState.FAILED + elif state == RunState.FINISHED: + self.state = SchedulerState.COMPLETED + + # check sweep state for completed states, overwrite scheduler state + try: + sweep_state = self._api.get_sweep_state( + self._sweep_id, self._entity, self._project + ) + except Exception as e: + _logger.debug(f"sweep state error: {e}") + return + + if sweep_state == "FINISHED": + self.state = SchedulerState.COMPLETED + elif sweep_state in ["CANCELLED", "STOPPED"]: + self.state = SchedulerState.CANCELLED + elif sweep_state == "PAUSED": + self.state = SchedulerState.FLUSH_RUNS + + def _update_run_states(self) -> None: + """Iterate through runs. + + Get state from backend and deletes runs if not in running state. Threadsafe. + """ + runs_to_remove: List[str] = [] + for run_id, run in self._yield_runs(): + run.state = self._get_run_state(run_id, run.state) + + try: + rqi_state = run.queued_run.state if run.queued_run else None + except (CommError, LaunchError) as e: + _logger.debug(f"Failed to get queued_run.state: {e}") + rqi_state = None + + if not run.state.is_alive or rqi_state == "failed": + _logger.debug(f"({run_id}) states: ({run.state}, {rqi_state})") + runs_to_remove.append(run_id) + self._cleanup_runs(runs_to_remove) + + def _get_metrics_from_run(self, run_id: str) -> List[Any]: + """Use the public api to get metrics from a run. + + Uses the metric name found in the sweep config, any + misspellings will result in an empty list. + """ + try: + queued_run: Optional[QueuedRun] = self._runs[run_id].queued_run + if not queued_run: + return [] + + api_run: Run = self._public_api.run( + f"{queued_run.entity}/{queued_run.project}/{run_id}" + ) + metric_name = self._sweep_config["metric"]["name"] + history = api_run.scan_history(keys=["_step", metric_name]) + metrics = [x[metric_name] for x in history] + + return metrics + except Exception as e: + _logger.debug(f"[_get_metrics_from_run] {e}") + return [] + + def _get_run_info(self, run_id: str) -> Dict[str, Any]: + """Use the public api to get info about a run.""" + try: + info: Dict[str, Any] = self._api.get_run_info( + self._entity, self._project, run_id + ) + if info: + return info + except Exception as e: + _logger.debug(f"[_get_run_info] {e}") + return {} + + def _get_run_state( + self, run_id: str, prev_run_state: RunState = RunState.UNKNOWN + ) -> RunState: + """Use the public api to get state of a run.""" + run_state = None + try: + state = self._api.get_run_state(self._entity, self._project, run_id) + run_state = RunState(state) + except CommError as e: + _logger.debug(f"error getting state for run ({run_id}): {e}") + if prev_run_state == RunState.UNKNOWN: + # triggers when we get an unknown state for the second time + wandb.termwarn( + f"Failed to get runstate for run ({run_id}). Error: {traceback.format_exc()}" + ) + run_state = RunState.FAILED + else: # first time we get unknown state + run_state = RunState.UNKNOWN + except (AttributeError, ValueError): + wandb.termwarn( + f"Bad state ({run_state}) for run ({run_id}). Error: {traceback.format_exc()}" + ) + run_state = RunState.UNKNOWN + return run_state + + def _create_run(self) -> Dict[str, Any]: + """Use the public api to create a blank run.""" + try: + run: List[Dict[str, Any]] = self._api.upsert_run( + project=self._project, + entity=self._entity, + sweep_name=self._sweep_id, + ) + if run: + return run[0] + except Exception as e: + _logger.debug(f"[_create_run] {e}") + raise SchedulerError( + "Error creating run from scheduler, check API connection and CLI version." + ) + return {} + + def _set_sweep_state(self, state: str) -> None: + wandb.termlog(f"{LOG_PREFIX}Updating sweep state to: {state.lower()}") + try: + self._api.set_sweep_state(sweep=self._sweep_id, state=state) + except Exception as e: + _logger.debug(f"[set_sweep_state] {e}") + + def _encode(self, _id: str) -> str: + return ( + base64.b64decode(bytes(_id.encode("utf-8"))).decode("utf-8").split(":")[2] + ) + + def _make_entry_and_launch_config( + self, run: SweepRun + ) -> Tuple[Optional[List[str]], Dict[str, Dict[str, Any]]]: + args = create_sweep_command_args({"args": run.args}) + entry_point, macro_args = make_launch_sweep_entrypoint( + args, self._sweep_config.get("command") + ) + # handle program macro + if entry_point and "${program}" in entry_point: + if not self._sweep_config.get("program"): + raise SchedulerError( + f"{LOG_PREFIX}Program macro in command has no corresponding 'program' in sweep config." + ) + pidx = entry_point.index("${program}") + entry_point[pidx] = self._sweep_config["program"] + + launch_config = copy.deepcopy(self._wandb_run.config.get("launch", {})) + if "overrides" not in launch_config: + launch_config["overrides"] = {"run_config": {}} + if "run_config" not in launch_config["overrides"]: + launch_config["overrides"]["run_config"] = {} + launch_config["overrides"]["run_config"].update(args["args_dict"]) + + if macro_args: # pipe in hyperparam args as params to launch + launch_config["overrides"]["args"] = macro_args + + if entry_point: + unresolved = [x for x in entry_point if str(x).startswith("${")] + if unresolved: + wandb.termwarn( + f"{LOG_PREFIX}Sweep command contains unresolved macros: " + f"{unresolved}, see launch docs for supported macros." + ) + return entry_point, launch_config + + def _add_to_launch_queue(self, run: SweepRun) -> bool: + """Convert a sweeprun into a launch job then push to runqueue.""" + # job and image first from CLI args, then from sweep config + _job = self._kwargs.get("job") or self._sweep_config.get("job") + _sweep_config_uri = self._sweep_config.get("image_uri") + _image_uri = self._kwargs.get("image_uri") or _sweep_config_uri + if _job is None and _image_uri is None: + raise SchedulerError(f"{LOG_PREFIX}No 'job' nor 'image_uri' ({run.id})") + elif _job is not None and _image_uri is not None: + raise SchedulerError(f"{LOG_PREFIX}Sweep has both 'job' and 'image_uri'") + + entry_point, launch_config = self._make_entry_and_launch_config(run) + if entry_point: + wandb.termwarn( + f"{LOG_PREFIX}Sweep command {entry_point} will override" + f" {'job' if _job else 'image_uri'} entrypoint" + ) + + # override resource and args of job + _job_launch_config = copy.deepcopy(self._wandb_run.config.get("launch")) or {} + + # default priority is "medium" + _priority = int(launch_config.get("priority", 2)) # type: ignore + + # strip resource_args and template_variables from launch_config + strip_resource_args_and_template_vars(_job_launch_config) + + run_id = run.id or generate_id() + queued_run = launch_add( + run_id=run_id, + entry_point=entry_point, + config=launch_config, + docker_image=_image_uri, # TODO(gst): make agnostic (github? run uri?) + job=_job, + project=self._project, + entity=self._entity, + queue_name=self._kwargs.get("queue"), + project_queue=self._project_queue, + resource=_job_launch_config.get("resource"), + resource_args=_job_launch_config.get("resource_args"), + template_variables=_job_launch_config.get("template_variables"), + author=self._kwargs.get("author"), + sweep_id=self._sweep_id, + priority=_priority, + ) + run.queued_run = queued_run + # TODO(gst): unify run and queued_run state + run.state = RunState.RUNNING # assume it will get picked up + self._runs[run_id] = run + + wandb.termlog( + f"{LOG_PREFIX}Added run ({run_id}) to queue ({self._kwargs.get('queue')})" + ) + return True diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/scheduler_sweep.py b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/scheduler_sweep.py new file mode 100644 index 0000000000000000000000000000000000000000..40484fdc2ac47aa2856d3c123e465f913e2e290a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/scheduler_sweep.py @@ -0,0 +1,90 @@ +"""Scheduler for classic wandb Sweeps.""" + +import logging +from pprint import pformat as pf +from typing import Any, Dict, List, Optional + +import wandb +from wandb.sdk.launch.sweeps.scheduler import LOG_PREFIX, RunState, Scheduler, SweepRun + +_logger = logging.getLogger(__name__) + + +class SweepScheduler(Scheduler): + """A controller/agent that populates a Launch RunQueue from a sweeps RunQueue.""" + + def __init__( + self, + *args: Any, + **kwargs: Any, + ): + super().__init__(*args, **kwargs) + + def _get_next_sweep_run(self, worker_id: int) -> Optional[SweepRun]: + """Called by the main scheduler execution loop. + + Expected to return a properly formatted SweepRun if the scheduler + is alive, or None and set the appropriate scheduler state: + + FAILED: self.fail_sweep() + STOPPED: self.stop_sweep() + """ + commands: List[Dict[str, Any]] = self._get_sweep_commands(worker_id) + for command in commands: + # The command "type" can be one of "run", "resume", "stop", "exit" + _type = command.get("type") + if _type in ["exit", "stop"]: + self.stop_sweep() + return None + + if _type not in ["run", "resume"]: + self.fail_sweep(f"AgentHeartbeat unknown command: {_type}") + + _run_id: Optional[str] = command.get("run_id") + if not _run_id: + self.fail_sweep(f"No run id in agent heartbeat: {command}") + return None + + if _run_id in self._runs: + wandb.termlog(f"{LOG_PREFIX}Skipping duplicate run: {_run_id}") + continue + + return SweepRun( + id=_run_id, + state=RunState.PENDING, + args=command.get("args", {}), + logs=command.get("logs", []), + worker_id=worker_id, + ) + return None + + def _get_sweep_commands(self, worker_id: int) -> List[Dict[str, Any]]: + """Helper to receive sweep command from backend.""" + # AgentHeartbeat wants a Dict of runs which are running or queued + _run_states: Dict[str, bool] = {} + for run_id, run in self._yield_runs(): + # Filter out runs that are from a different worker thread + if run.worker_id == worker_id and run.state.is_alive: + _run_states[run_id] = True + + _logger.debug(f"Sending states: \n{pf(_run_states)}\n") + commands: List[Dict[str, Any]] = self._api.agent_heartbeat( + agent_id=self._workers[worker_id].agent_id, + metrics={}, + run_states=_run_states, + ) + _logger.debug(f"AgentHeartbeat commands: \n{pf(commands)}\n") + + return commands + + def _exit(self) -> None: + pass + + def _poll(self) -> None: + _logger.debug(f"_poll. _runs: {self._runs}") + + def _load_state(self) -> None: + pass + + def _save_state(self) -> None: + pass diff --git a/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/utils.py b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d09defe9b8d757f04e08e663abc277b171d6bfcf --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/launch/sweeps/utils.py @@ -0,0 +1,324 @@ +import json +import os +import re +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +import yaml + +import wandb +from wandb import util +from wandb.sdk.launch.errors import LaunchError + +if TYPE_CHECKING: + from wandb.apis.public import Api as PublicApi + +DEFAULT_SWEEP_COMMAND: List[str] = [ + "${env}", + "${interpreter}", + "${program}", + "${args}", +] +SWEEP_COMMAND_ENV_VAR_REGEX = re.compile(r"\$\{envvar\:([A-Z0-9_]*)\}") + + +def parse_sweep_id(parts_dict: dict) -> Optional[str]: + """In place parse sweep path from parts dict. + + Arguments: + parts_dict (dict): dict(entity=,project=,name=). Modifies dict inplace. + + Returns: + None or str if there is an error + """ + entity = None + project = None + sweep_id = parts_dict.get("name") + if not isinstance(sweep_id, str): + return "Expected string sweep_id" + + sweep_split = sweep_id.split("/") + if len(sweep_split) == 1: + pass + elif len(sweep_split) == 2: + split_project, sweep_id = sweep_split + project = split_project or project + elif len(sweep_split) == 3: + split_entity, split_project, sweep_id = sweep_split + project = split_project or project + entity = split_entity or entity + else: + return ( + "Expected sweep_id in form of sweep, project/sweep, or entity/project/sweep" + ) + parts_dict.update(dict(name=sweep_id, project=project, entity=entity)) + return None + + +def sweep_config_err_text_from_jsonschema_violations(violations: List[str]) -> str: + """Consolidate schema violation strings from wandb/sweeps into a single string. + + Parameters + ---------- + violations: list of str + The warnings to render. + + Returns: + ------- + violation: str + The consolidated violation text. + + """ + violation_base = ( + "Malformed sweep config detected! This may cause your sweep to behave in unexpected ways.\n" + "To avoid this, please fix the sweep config schema violations below:" + ) + + for i, warning in enumerate(violations): + violations[i] = f" Violation {i + 1}. {warning}" + violation = "\n".join([violation_base] + violations) + + return violation + + +def handle_sweep_config_violations(warnings: List[str]) -> None: + """Echo sweep config schema violation warnings from Gorilla to the terminal. + + Parameters + ---------- + warnings: list of str + The warnings to render. + """ + warning = sweep_config_err_text_from_jsonschema_violations(warnings) + if len(warnings) > 0: + wandb.termwarn(warning) + + +def load_sweep_config(sweep_config_path: str) -> Optional[Dict[str, Any]]: + """Load a sweep yaml from path.""" + try: + yaml_file = open(sweep_config_path) + except OSError: + wandb.termerror(f"Couldn't open sweep file: {sweep_config_path}") + return None + try: + config: Optional[Dict[str, Any]] = yaml.safe_load(yaml_file) + except yaml.YAMLError as err: + wandb.termerror(f"Error in configuration file: {err}") + return None + if not config: + wandb.termerror("Configuration file is empty") + return None + return config + + +def load_launch_sweep_config(config: Optional[str]) -> Any: + if not config: + return {} + + parsed_config = util.load_json_yaml_dict(config) + if parsed_config is None: + raise LaunchError(f"Could not load config from {config}. Check formatting") + return parsed_config + + +def construct_scheduler_args( + sweep_config: Dict[str, Any], + queue: str, + project: str, + author: Optional[str] = None, + return_job: bool = False, +) -> Union[List[str], Dict[str, str], None]: + """Construct sweep scheduler args. + + logs error and returns None if misconfigured, + otherwise returns args as a dict if is_job else a list of strings. + """ + job = sweep_config.get("job") + image_uri = sweep_config.get("image_uri") + if not job and not image_uri: # don't allow empty string + wandb.termerror( + "No 'job' nor 'image_uri' top-level key found in sweep config, exactly one is required for a launch-sweep" + ) + return None + elif job and image_uri: + wandb.termerror( + "Sweep config has both 'job' and 'image_uri' but a launch-sweep can use only one" + ) + return None + + # if scheduler is a job, return args as dict + if return_job: + args_dict: Dict[str, str] = { + "sweep_id": "WANDB_SWEEP_ID", + "queue": queue, + "project": project, + } + if job: + args_dict["job"] = job + elif image_uri: + args_dict["image_uri"] = image_uri + + if author: + args_dict["author"] = author + + return args_dict + + # scheduler uses cli commands, pass args as param list + args = [ + "--queue", + f"{queue!r}", + "--project", + f"{project!r}", + ] + if author: + args += [ + "--author", + f"{author!r}", + ] + if job: + args += [ + "--job", + f"{job!r}", + ] + elif image_uri: + args += ["--image_uri", image_uri] + + return args + + +def create_sweep_command(command: Optional[List] = None) -> List: + """Return sweep command, filling in environment variable macros.""" + # Start from default sweep command + command = command or DEFAULT_SWEEP_COMMAND + for i, chunk in enumerate(command): + # Replace environment variable macros + # Search a str(chunk), but allow matches to be of any (ex: int) type + if SWEEP_COMMAND_ENV_VAR_REGEX.search(str(chunk)): + # Replace from backwards forwards + matches = list(SWEEP_COMMAND_ENV_VAR_REGEX.finditer(chunk)) + for m in matches[::-1]: + # Default to just leaving as is if environment variable does not exist + _var: str = os.environ.get(m.group(1), m.group(1)) + command[i] = f"{command[i][: m.start()]}{_var}{command[i][m.end() :]}" + return command + + +def create_sweep_command_args(command: Dict) -> Dict[str, Any]: + """Create various formats of command arguments for the agent. + + Raises: + ValueError: improperly formatted command dict + + """ + if "args" not in command: + raise ValueError(f'No "args" found in command: {command}') + # four different formats of command args + # (1) standard command line flags (e.g. --foo=bar) + flags: List[str] = [] + # (2) flags without hyphens (e.g. foo=bar) + flags_no_hyphens: List[str] = [] + # (3) flags with false booleans omitted (e.g. --foo) + flags_no_booleans: List[str] = [] + # (4) flags as a dictionary (used for constructing a json) + flags_dict: Dict[str, Any] = {} + # (5) flags without equals (e.g. --foo bar) + args_no_equals: List[str] = [] + # (6) flags for hydra append config value (e.g. +foo=bar) + flags_append_hydra: List[str] = [] + # (7) flags for hydra override config value (e.g. ++foo=bar) + flags_override_hydra: List[str] = [] + for param, config in command["args"].items(): + # allow 'None' as a valid value, but error if no value is found + try: + _value: Any = config["value"] + except KeyError: + raise ValueError(f'No "value" found for command["args"]["{param}"]') + + _flag: str = f"{param}={_value}" + flags.append("--" + _flag) + flags_no_hyphens.append(_flag) + args_no_equals += [f"--{param}", str(_value)] + flags_append_hydra.append("+" + _flag) + flags_override_hydra.append("++" + _flag) + if isinstance(_value, bool): + # omit flags if they are boolean and false + if _value: + flags_no_booleans.append("--" + param) + else: + flags_no_booleans.append("--" + _flag) + flags_dict[param] = _value + return { + "args": flags, + "args_no_equals": args_no_equals, + "args_no_hyphens": flags_no_hyphens, + "args_no_boolean_flags": flags_no_booleans, + "args_json": [json.dumps(flags_dict)], + "args_dict": flags_dict, + "args_append_hydra": flags_append_hydra, + "args_override_hydra": flags_override_hydra, + } + + +def make_launch_sweep_entrypoint( + args: Dict[str, Any], command: Optional[List[str]] +) -> Tuple[Optional[List[str]], Any]: + """Use args dict from create_sweep_command_args to construct entrypoint. + + If replace is True, remove macros from entrypoint, fill them in with args + and then return the args in separate return value. + """ + if not command: + return None, None + + entry_point = create_sweep_command(command) + macro_args = {} + for macro in args: + mstr = "${" + macro + "}" + if mstr in entry_point: + idx = entry_point.index(mstr) + # only supports 1 macro per entrypoint + macro_args = args[macro] + entry_point = entry_point[:idx] + entry_point[idx + 1 :] + + if len(entry_point) == 0: + return None, macro_args + + return entry_point, macro_args + + +def check_job_exists(public_api: "PublicApi", job: Optional[str]) -> bool: + """Check if the job exists using the public api. + + Returns: True if no job is passed, or if the job exists. + Returns: False if the job is misformatted or doesn't exist. + """ + if not job: + return True + + try: + public_api.job(job) + except Exception as e: + wandb.termerror(f"Failed to load job. {e}") + return False + return True + + +def get_previous_args( + run_spec: Dict[str, Any], +) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """Parse through previous scheduler run_spec. + + returns scheduler_args and settings. + """ + scheduler_args = ( + run_spec.get("overrides", {}).get("run_config", {}).get("scheduler", {}) + ) + # also pipe through top level resource setup + if run_spec.get("resource"): + scheduler_args["resource"] = run_spec["resource"] + if run_spec.get("resource_args"): + scheduler_args["resource_args"] = run_spec["resource_args"] + + settings = run_spec.get("overrides", {}).get("run_config", {}).get("settings", {}) + + return scheduler_args, settings diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__init__.py b/lib/python3.12/site-packages/wandb/sdk/lib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4289765c6c7c4fda2dbe42dddc1f5f7d69da0c45 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/__init__.py @@ -0,0 +1,5 @@ +from . import lazyloader +from .disabled import RunDisabled, SummaryDisabled +from .run_moment import RunMoment + +__all__ = ("lazyloader", "RunDisabled", "SummaryDisabled", "RunMoment") diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e8ed91751f32053a10a8f14f487839ec3f90ff60 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/apikey.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/apikey.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65e7cfb4a3451ffe2653a3a264914e6b892e316f Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/apikey.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/asyncio_compat.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/asyncio_compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c4557f7199a4e3bb333b4a78a4224728188117ce Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/asyncio_compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/capped_dict.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/capped_dict.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6783ba740f24bd75876c64dd8f3096e200af6846 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/capped_dict.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/config_util.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/config_util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07f47f0d53cc40c06154349eca1c5d23beaced1d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/config_util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/console_capture.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/console_capture.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cdb470176490cae85b58f80183a22dc1e68709cb Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/console_capture.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/credentials.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/credentials.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fcb7502173be1a04f65335bcbcb3018c5a8695b6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/credentials.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/deprecate.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/deprecate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e55a04707bd0248c561b103843928fbe1f42d24 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/deprecate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/disabled.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/disabled.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4cd521aa2a5cc479918cc77c8a569e7951f1c744 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/disabled.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/exit_hooks.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/exit_hooks.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..39a144b16298a31fbd81ccfb55338c4eed1be3d9 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/exit_hooks.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/file_stream_utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/file_stream_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe568c6fc7c130ba8c60050e015a9fed1afeac98 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/file_stream_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/filenames.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/filenames.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..82b8f4eb085799a7a7e6e96d13832f3343c1f262 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/filenames.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/filesystem.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/filesystem.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3dd5bb67c849e7060afd2ca394b19da0c3aea963 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/filesystem.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/fsm.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/fsm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fcde7fb4e21f1b1cd6eb302116863edbec2c9931 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/fsm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/gitlib.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/gitlib.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6bd21510292b48b062fec7bdc067c28d34855b94 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/gitlib.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/gql_request.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/gql_request.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09973dc5ab1bf03a267e70f074d746b4c469bb6b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/gql_request.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/handler_util.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/handler_util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f2ef640efb0e0f5d01db17e189f3a7f43c7d70a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/handler_util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/hashutil.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/hashutil.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c9beac6a587b2c6bf78cf38ed53767a271f4a45 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/hashutil.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/import_hooks.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/import_hooks.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a62fabd7eccf21e8525c3d7e7442b10cc07ec93b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/import_hooks.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/interrupt.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/interrupt.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..edd5933a8ed924efc17d7032e5c4361130652c42 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/interrupt.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/ipython.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/ipython.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..38333c1d4e504f4a8d4c692984d5174d1f361f93 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/ipython.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/json_util.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/json_util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0b4faf91de52ad39a28e5a2727905b9b66e9e1c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/json_util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/lazyloader.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/lazyloader.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..270edd7c9eb38f92ee21ab6daed2c4c93c14cb44 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/lazyloader.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/module.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9b88c11d4a27bb1bda0fcff79b07dfba1806a48 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/paths.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/paths.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..170c86e3fa268ce05b078254a25cc0839c13534d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/paths.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/preinit.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/preinit.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b180c9f202dc05ee5c307f084ff8af3ecec4397 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/preinit.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/printer.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/printer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9224127d7ad048edb5d346989de2483c9307fe2d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/printer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/printer_asyncio.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/printer_asyncio.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b08ceb64f22f9d4fcda2fcc2f39733e39074278a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/printer_asyncio.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/progress.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/progress.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60a88b4b33b9b37e49568400f6f7381bc9f75092 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/progress.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/proto_util.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/proto_util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e1a3b4cf9d866c64466a4242f593b9b5ffd835eb Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/proto_util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/redirect.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/redirect.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3bd2dcf1e675762bec8e1d4ab8e84aede4d6d7de Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/redirect.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/retry.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/retry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..15727ad296585d97af00b8297cd47bd6942dc2ab Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/retry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/run_moment.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/run_moment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..966ca60d3ecfddaa6d5f5dd066287fd1d2fa3666 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/run_moment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/runid.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/runid.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..058b1cce6d8ba1d916303f58672448fc8e523681 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/runid.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/server.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/server.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..828ce9b4daadafefd99ac382ec70a7621f89636e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/server.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/sock_client.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/sock_client.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b5a247c9e92570200134d54e1a1b45432333a4f Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/sock_client.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/sparkline.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/sparkline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47cc9313f7ff0f5779079120015a5092d37b6bfd Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/sparkline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/telemetry.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/telemetry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..18bf033faa7ac1f5a1536a0bd943160e3f3a68e3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/telemetry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/timed_input.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/timed_input.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c93717044ba447ca885d12a4a676bec884904e44 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/timed_input.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/timer.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/timer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1cd079fa19b9c5deb08800a87dea0bfb1aa33136 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/lib/__pycache__/timer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/apikey.py b/lib/python3.12/site-packages/wandb/sdk/lib/apikey.py new file mode 100644 index 0000000000000000000000000000000000000000..f30e791de800d4d4996b403d163ac9c0a4af9f45 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/apikey.py @@ -0,0 +1,334 @@ +"""apikey util.""" + +from __future__ import annotations + +import dataclasses +import os +import platform +import stat +import sys +import textwrap +from functools import partial + +# import Literal +from typing import TYPE_CHECKING, Callable, Literal +from urllib.parse import urlparse + +import click +from requests.utils import NETRC_FILES, get_netrc_auth + +import wandb +from wandb.apis import InternalApi +from wandb.errors import term +from wandb.errors.links import url_registry +from wandb.sdk import wandb_setup +from wandb.util import _is_databricks, isatty, prompt_choices + +if TYPE_CHECKING: + from wandb.sdk.wandb_settings import Settings + +LOGIN_CHOICE_ANON = "Private W&B dashboard, no account required" +LOGIN_CHOICE_NEW = "Create a W&B account" +LOGIN_CHOICE_EXISTS = "Use an existing W&B account" +LOGIN_CHOICE_DRYRUN = "Don't visualize my results" +LOGIN_CHOICE_NOTTY = "Unconfigured" +LOGIN_CHOICES = [ + LOGIN_CHOICE_ANON, + LOGIN_CHOICE_NEW, + LOGIN_CHOICE_EXISTS, + LOGIN_CHOICE_DRYRUN, +] + + +@dataclasses.dataclass(frozen=True) +class _NetrcPermissions: + exists: bool + read_access: bool + write_access: bool + + +class WriteNetrcError(Exception): + """Raised when we cannot write to the netrc file.""" + + +Mode = Literal["allow", "must", "never", "false", "true"] + + +getpass = partial(click.prompt, hide_input=True, err=True) + + +def _fixup_anon_mode(default: Mode | None) -> Mode | None: + # Convert weird anonymode values from legacy settings files + # into one of our expected values. + anon_mode = default or "never" + mapping: dict[Mode, Mode] = {"true": "allow", "false": "never"} + return mapping.get(anon_mode, anon_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) + + # 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 _api_key_prompt_str(app_url: str, referrer: str | None = None) -> str: + """Generate a prompt string for API key authorization. + + Creates a URL string that directs users to the authorization page where they + can find their API key. + + Args: + app_url: The base URL of the W&B application. + referrer: Optional referrer parameter to include in the URL. + + Returns: + A formatted string with instructions and the authorization URL. + """ + ref = "" + if referrer: + ref = f"?ref={referrer}" + return f"You can find your API key in your browser here: {app_url}/authorize{ref}" + + +def prompt_api_key( # noqa: C901 + settings: Settings, + api: InternalApi | None = None, + input_callback: Callable | None = None, + browser_callback: Callable | None = None, + no_offline: bool = False, + no_create: bool = False, + local: bool = False, + referrer: str | None = None, +) -> str | bool | None: + """Prompt for api key. + + Returns: + str - if key is configured + None - if dryrun is selected + False - if unconfigured (notty) + """ + input_callback = input_callback or getpass + log_string = term.LOG_STRING + api = api or InternalApi(settings) + anon_mode = _fixup_anon_mode(settings.anonymous) # type: ignore + jupyter = settings._jupyter or False + app_url = api.app_url + + choices = [choice for choice in LOGIN_CHOICES] + if anon_mode == "never": + # Omit LOGIN_CHOICE_ANON as a choice if the env var is set to never + choices.remove(LOGIN_CHOICE_ANON) + if (jupyter and not settings.login_timeout) or no_offline: + choices.remove(LOGIN_CHOICE_DRYRUN) + if (jupyter and not settings.login_timeout) or no_create: + choices.remove(LOGIN_CHOICE_NEW) + + if jupyter and "google.colab" in sys.modules: + log_string = term.LOG_STRING_NOCOLOR + key = wandb.jupyter.attempt_colab_login(app_url) # type: ignore + if key is not None: + return key # type: ignore + + if anon_mode == "must": + result = LOGIN_CHOICE_ANON + # If we're not in an interactive environment, default to dry-run. + elif ( + not jupyter and (not isatty(sys.stdout) or not isatty(sys.stdin)) + ) or _is_databricks(): + result = LOGIN_CHOICE_NOTTY + elif local: + result = LOGIN_CHOICE_EXISTS + elif len(choices) == 1: + result = choices[0] + else: + result = prompt_choices( + choices, input_timeout=settings.login_timeout, jupyter=jupyter + ) + + key = None + api_ask = ( + f"{log_string}: Paste an API key from your profile and hit enter" + if jupyter + else f"{log_string}: Paste an API key from your profile and hit enter, or press ctrl+c to quit" + ) + if result == LOGIN_CHOICE_ANON: + key = api.create_anonymous_api_key() + elif result == LOGIN_CHOICE_NEW: + key = browser_callback(signup=True) if browser_callback else None + + if not key: + ref = f"&ref={referrer}" if referrer else "" + wandb.termlog( + f"Create an account here: {app_url}/authorize?signup=true{ref}" + ) + key = input_callback(api_ask).strip() + elif result == LOGIN_CHOICE_EXISTS: + key = browser_callback() if browser_callback else None + + if not key: + if not (settings.is_local or local): + host = app_url + for prefix in ("http://", "https://"): + if app_url.startswith(prefix): + host = app_url[len(prefix) :] + wandb.termlog( + f"Logging into {host}. (Learn how to deploy a W&B server " + f"locally: {url_registry.url('wandb-server')})" + ) + wandb.termlog(_api_key_prompt_str(app_url, referrer)) + key = input_callback(api_ask).strip() + elif result == LOGIN_CHOICE_NOTTY: + # TODO: Needs refactor as this needs to be handled by caller + return False + elif result == LOGIN_CHOICE_DRYRUN: + return None + else: + # Jupyter environments don't have a tty, but we can still try logging in using + # the browser callback if one is supplied. + key, anonymous = ( + browser_callback() if jupyter and browser_callback else (None, False) + ) + + if not key: + raise ValueError("No API key specified.") + return key + + +def check_netrc_access( + netrc_path: str, +) -> _NetrcPermissions: + """Check if we can read and write to the netrc file.""" + file_exists = False + write_access = False + read_access = False + try: + st = os.stat(netrc_path) + file_exists = True + write_access = bool(st.st_mode & stat.S_IWUSR) + read_access = bool(st.st_mode & stat.S_IRUSR) + except FileNotFoundError: + # If the netrc file doesn't exist, we will create it. + write_access = True + read_access = True + except OSError as e: + wandb.termerror(f"Unable to read permissions for {netrc_path}, {e}") + + return _NetrcPermissions( + exists=file_exists, + write_access=write_access, + read_access=read_access, + ) + + +def write_netrc(host: str, entity: str, key: str): + """Add our host and key to .netrc.""" + _, key_suffix = key.split("-", 1) if "-" in key else ("", key) + if len(key_suffix) != 40: + raise ValueError( + f"API-key must be exactly 40 characters long: {key_suffix} ({len(key_suffix)} chars)" + ) + + normalized_host = urlparse(host).netloc + netrc_path = get_netrc_file_path() + netrc_access = check_netrc_access(netrc_path) + + if not netrc_access.write_access or not netrc_access.read_access: + raise WriteNetrcError( + f"Cannot access {netrc_path}. In order to persist your API key, " + "grant read and write permissions for your user to the file " + 'or specify a different file with the environment variable "NETRC=".' + ) + + machine_line = f"machine {normalized_host}" + orig_lines = None + try: + with open(netrc_path) as f: + orig_lines = f.read().strip().split("\n") + except FileNotFoundError: + wandb.termlog("No netrc file found, creating one.") + except OSError as e: + raise WriteNetrcError(f"Unable to read {netrc_path}") from e + + try: + with open(netrc_path, "w") as f: + if orig_lines: + # delete this machine from the file if it's already there. + skip = 0 + for line in orig_lines: + # we fix invalid netrc files with an empty host that we wrote before + # verifying host... + if line == "machine " or machine_line in line: + skip = 2 + elif skip: + skip -= 1 + else: + f.write(f"{line}\n") + + wandb.termlog( + f"Appending key for {normalized_host} to your netrc file: {netrc_path}" + ) + f.write( + textwrap.dedent( + """\ + machine {host} + login {entity} + password {key} + """ + ).format(host=normalized_host, entity=entity, key=key) + ) + os.chmod(netrc_path, stat.S_IRUSR | stat.S_IWUSR) + except OSError as e: + raise WriteNetrcError(f"Unable to write {netrc_path}") from e + + +def write_key( + settings: Settings, + key: str | None, + api: InternalApi | None = None, +) -> None: + if not key: + raise ValueError("No API key specified.") + + # TODO(jhr): api shouldn't be optional or it shouldn't be passed, clean up callers + api = api or InternalApi() + + # Normal API keys are 40-character hex strings. On-prem API keys have a + # variable-length prefix, a dash, then the 40-char string. + _, suffix = key.split("-", 1) if "-" in key else ("", key) + + if len(suffix) != 40: + raise ValueError(f"API key must be 40 characters long, yours was {len(key)}") + + write_netrc(settings.base_url, "user", key) + + +def api_key(settings: Settings | None = None) -> str | None: + if settings is None: + settings = wandb_setup.singleton().settings + if settings.api_key: + return settings.api_key + + netrc_access = check_netrc_access(get_netrc_file_path()) + if netrc_access.exists and not netrc_access.read_access: + wandb.termwarn(f"Cannot access {get_netrc_file_path()}.") + return None + + if netrc_access.exists: + auth = get_netrc_auth(settings.base_url) + if auth: + return auth[-1] + + return None diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/asyncio_compat.py b/lib/python3.12/site-packages/wandb/sdk/lib/asyncio_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..598fd53608b309dab70615e2eef48cf76ca1fe4f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/asyncio_compat.py @@ -0,0 +1,210 @@ +"""Functions for compatibility with asyncio.""" + +from __future__ import annotations + +import asyncio +import concurrent +import concurrent.futures +import contextlib +import threading +from typing import Any, AsyncIterator, Callable, Coroutine, Iterator, TypeVar + +_T = TypeVar("_T") + + +def run(fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T: + """Run `fn` in an asyncio loop in a new thread. + + This must always be used instead of `asyncio.run` which fails if there is + an active `asyncio` event loop in the current thread. Since `wandb` was not + originally designed with `asyncio` in mind, using `asyncio.run` would break + users who were calling `wandb` methods from an `asyncio` loop. + + Note that due to starting a new thread, this is slightly slow. + """ + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + runner = _Runner() + future = executor.submit(runner.run, fn) + + try: + return future.result() + + finally: + runner.cancel() + + +class _RunnerCancelledError(Exception): + """The `_Runner.run()` invocation was cancelled.""" + + +class _Runner: + """Runs an asyncio event loop allowing cancellation. + + This is like `asyncio.run()`, except it provides a `cancel()` method + meant to be called in a `finally` block. + + Without this, it is impossible to make `asyncio.run()` stop if it runs + in a non-main thread. In particular, a KeyboardInterrupt causes the + ThreadPoolExecutor above to block until the asyncio thread completes, + but there is no way to tell the asyncio thread to cancel its work. + A second KeyboardInterrupt makes ThreadPoolExecutor give up while the + asyncio thread still runs in the background, with terrible effects if it + prints to the user's terminal. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + + self._is_cancelled = False + self._started = False + self._done = False + + self._loop: asyncio.AbstractEventLoop | None = None + self._cancel_event: asyncio.Event | None = None + + def run(self, fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T: + """Run a coroutine in asyncio, cancelling it on `cancel()`. + + Returns: + The result of the coroutine returned by `fn`. + + Raises: + _RunnerCancelledError: If `cancel()` is called. + """ + return asyncio.run(self._run_or_cancel(fn)) + + async def _run_or_cancel( + self, + fn: Callable[[], Coroutine[Any, Any, _T]], + ) -> _T: + with self._lock: + if self._is_cancelled: + raise _RunnerCancelledError() + + self._loop = asyncio.get_running_loop() + self._cancel_event = asyncio.Event() + self._started = True + + cancellation_task = asyncio.create_task(self._cancel_event.wait()) + fn_task = asyncio.create_task(fn()) + + try: + await asyncio.wait( + [cancellation_task, fn_task], + return_when=asyncio.FIRST_COMPLETED, + ) + + if fn_task.done(): + return fn_task.result() + else: + raise _RunnerCancelledError() + + finally: + cancellation_task.cancel() + fn_task.cancel() + + with self._lock: + self._done = True + + def cancel(self) -> None: + """Cancel all asyncio work started by `run()`.""" + with self._lock: + if self._is_cancelled: + return + self._is_cancelled = True + + if self._done or not self._started: + # If the runner already finished, no need to cancel it. + # + # If the runner hasn't started the loop yet, then it will not + # as we already set _is_cancelled. + return + + assert self._loop + assert self._cancel_event + self._loop.call_soon_threadsafe(self._cancel_event.set) + + +class TaskGroup: + """Object that `open_task_group()` yields.""" + + def __init__(self) -> None: + self._tasks: list[asyncio.Task] = [] + + def start_soon(self, coro: Coroutine[Any, Any, Any]) -> None: + """Schedule a task in the group. + + Args: + coro: The return value of the `async` function defining the task. + """ + self._tasks.append(asyncio.create_task(coro)) + + async def _wait_all(self) -> None: + """Block until all tasks complete. + + Raises: + Exception: If one or more tasks raises an exception, one of these + is raised arbitrarily. + """ + done, _ = await asyncio.wait( + self._tasks, + # NOTE: Cancelling a task counts as a normal exit, + # not an exception. + return_when=concurrent.futures.FIRST_EXCEPTION, + ) + + for task in done: + try: + if exc := task.exception(): + raise exc + except asyncio.CancelledError: + pass + + def _cancel_all(self) -> None: + """Cancel all tasks.""" + for task in self._tasks: + # NOTE: It is safe to cancel tasks that have already completed. + task.cancel() + + +@contextlib.asynccontextmanager +async def open_task_group() -> AsyncIterator[TaskGroup]: + """Create a task group. + + `asyncio` gained task groups in Python 3.11. + + This is an async context manager, meant to be used with `async with`. + On exit, it blocks until all subtasks complete. If any subtask fails, or if + the current task is cancelled, it cancels all subtasks in the group and + raises the subtask's exception. If multiple subtasks fail simultaneously, + one of their exceptions is chosen arbitrarily. + + NOTE: Subtask exceptions do not propagate until the context manager exits. + This means that the task group cannot cancel code running inside the + `async with` block . + """ + task_group = TaskGroup() + + try: + yield task_group + await task_group._wait_all() + finally: + task_group._cancel_all() + + +@contextlib.contextmanager +def cancel_on_exit(coro: Coroutine[Any, Any, Any]) -> Iterator[None]: + """Schedule a task, cancelling it when exiting the context manager. + + If the given coroutine raises an exception, that exception is raised + when exiting the context manager. + """ + task = asyncio.create_task(coro) + + try: + yield + finally: + if task.done() and (exception := task.exception()): + raise exception + + task.cancel() diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/capped_dict.py b/lib/python3.12/site-packages/wandb/sdk/lib/capped_dict.py new file mode 100644 index 0000000000000000000000000000000000000000..5e162109926cf7c115bf835323de4d292f8725e8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/capped_dict.py @@ -0,0 +1,26 @@ +import collections +from typing import Any, Optional + + +class CappedDict(collections.OrderedDict): + default_max_size = 50 + + def __init__(self, max_size: Optional[int] = None) -> None: + self.max_size = max_size or self.default_max_size + super().__init__() + + def __setitem__(self, key: str, val: Any) -> None: + if key not in self: + max_size = self.max_size - 1 + self._prune_dict(max_size) + super().__setitem__(key, val) + + def update(self, **kwargs: Any) -> None: # type: ignore[override] + super().update(**kwargs) + self._prune_dict(self.max_size) + + def _prune_dict(self, max_size: int) -> None: + if len(self) >= max_size: + diff = len(self) - max_size + for k in list(self.keys())[:diff]: + del self[k] diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/config_util.py b/lib/python3.12/site-packages/wandb/sdk/lib/config_util.py new file mode 100644 index 0000000000000000000000000000000000000000..9c20f844174374d09d1eec1711d4f9eea392c8e0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/config_util.py @@ -0,0 +1,101 @@ +import json +import logging +import os +from typing import Any, Dict, Optional + +import yaml + +import wandb +from wandb.errors import Error +from wandb.util import load_yaml + +from . import filesystem + +logger = logging.getLogger("wandb") + + +class ConfigError(Error): + pass + + +def dict_from_proto_list(obj_list): + d = dict() + for item in obj_list: + d[item.key] = dict(desc=None, value=json.loads(item.value_json)) + return d + + +def dict_strip_value_dict(config_dict): + d = dict() + for k, v in config_dict.items(): + d[k] = v["value"] + return d + + +def dict_no_value_from_proto_list(obj_list): + d = dict() + for item in obj_list: + possible_dict = json.loads(item.value_json) + if not isinstance(possible_dict, dict) or "value" not in possible_dict: + continue + d[item.key] = possible_dict["value"] + + return d + + +# TODO(jhr): these functions should go away once we merge jobspec PR +def save_config_file_from_dict(config_filename, config_dict): + s = b"wandb_version: 1" + if config_dict: # adding an empty dictionary here causes a parse error + s += b"\n\n" + yaml.dump( + config_dict, + Dumper=yaml.SafeDumper, + default_flow_style=False, + allow_unicode=True, + encoding="utf-8", + sort_keys=False, + ) + data = s.decode("utf-8") + filesystem.mkdir_exists_ok(os.path.dirname(config_filename)) + with open(config_filename, "w") as conf_file: + conf_file.write(data) + + +def dict_from_config_file( + filename: str, must_exist: bool = False +) -> Optional[Dict[str, Any]]: + if not os.path.exists(filename): + if must_exist: + raise ConfigError(f"config file {filename} doesn't exist") + logger.debug(f"no default config file found in {filename}") + return None + try: + conf_file = open(filename) + except OSError: + raise ConfigError(f"Couldn't read config file: {filename}") + try: + loaded = load_yaml(conf_file) + except yaml.parser.ParserError: + raise ConfigError("Invalid YAML in config yaml") + if loaded is None: + wandb.termwarn( + "Found an empty default config file (config-defaults.yaml). Proceeding with no defaults." + ) + return None + config_version = loaded.pop("wandb_version", None) + if config_version is not None and config_version != 1: + raise ConfigError("Unknown config version") + data = dict() + for k, v in loaded.items(): + data[k] = v["value"] + return data + + +def merge_dicts(dest: dict, src: dict) -> dict: + """Recursively merge two dictionaries. Similar to Lodash's _.merge().""" + for key, value in src.items(): + if isinstance(value, dict) and key in dest and isinstance(dest[key], dict): + merge_dicts(dest[key], value) + else: + dest[key] = value + return dest diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/console_capture.py b/lib/python3.12/site-packages/wandb/sdk/lib/console_capture.py new file mode 100644 index 0000000000000000000000000000000000000000..f9d56af0d74f518969de5a8c62811b30d49e36b6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/console_capture.py @@ -0,0 +1,219 @@ +"""Module for intercepting stdout/stderr. + +This patches the `write()` method of `stdout` and `stderr` on import. +Once patched, it is not possible to unpatch or repatch, though individual +callbacks can be removed. + +We assume that all other writing methods on the object delegate to `write()`, +like `writelines()`. This is not guaranteed to be true, but it is true for +common implementations. In particular, CPython's implementation of IOBase's +`writelines()` delegates to `write()`. + +It is important to note that this technique interacts poorly with other +code that performs similar patching if it also allows unpatching as this +discards our modification. This is why we patch on import and do not support +unpatching: + + with contextlib.redirect_stderr(...): + from ... import console_capture + # Here, everything works fine. + # Here, callbacks are never called again. + +In particular, it does not work with some combinations of pytest's +`capfd` / `capsys` fixtures and pytest's `--capture` option. +""" + +from __future__ import annotations + +import logging +import sys +import threading +from typing import IO, AnyStr, Callable, Protocol + +from . import wb_logging + +_logger = logging.getLogger(__name__) + + +class CannotCaptureConsoleError(Exception): + """The module failed to patch stdout or stderr.""" + + +class _WriteCallback(Protocol): + """A callback that receives intercepted bytes or string data. + + This may be called from any thread, but is only called from one thread + at a time. + + Note on errors: Any error raised during the callback will clear all + callbacks. This means that if a user presses Ctrl-C at an unlucky time + during a run, we will stop uploading console output---but it's not + likely to be a problem unless something catches the KeyboardInterrupt. + + Regular Exceptions are caught and logged instead of bubbling up to the + user's print() statements; other exceptions like KeyboardInterrupt are + re-raised. + + Callbacks should handle all exceptions---a callback that raises any + Exception is considered buggy. + """ + + def __call__( + self, + data: bytes | str, + written: int, + /, + ) -> None: + """Intercept data passed to `write()`. + + See the protocol docstring for information about exceptions. + + Args: + data: The object passed to stderr's or stdout's `write()`. + written: The number of bytes or characters written. + This is the return value of `write()`. + """ + + +# A reentrant lock is used to catch callbacks that write to stderr/stdout. +_module_rlock = threading.RLock() +_is_writing = False + +_patch_exception: CannotCaptureConsoleError | None = None + +_next_callback_id: int = 1 + +_stdout_callbacks: dict[int, _WriteCallback] = {} +_stderr_callbacks: dict[int, _WriteCallback] = {} + + +def capture_stdout(callback: _WriteCallback) -> Callable[[], None]: + """Install a callback that runs after every write to sys.stdout. + + Args: + callback: A callback to invoke after running `sys.stdout.write`. + + Returns: + A function to uninstall the callback. + + Raises: + CannotCaptureConsoleError: If patching failed on import. + """ + with _module_rlock: + if _patch_exception: + raise _patch_exception + + return _insert_disposably( + _stdout_callbacks, + callback, + ) + + +def capture_stderr(callback: _WriteCallback) -> Callable[[], None]: + """Install a callback that runs after every write to sys.sdterr. + + Args: + callback: A callback to invoke after running `sys.stderr.write`. + + Returns: + A function to uninstall the callback. + + Raises: + CannotCaptureConsoleError: If patching failed on import. + """ + with _module_rlock: + if _patch_exception: + raise _patch_exception + + return _insert_disposably( + _stderr_callbacks, + callback, + ) + + +def _insert_disposably( + callback_dict: dict[int, _WriteCallback], + callback: _WriteCallback, +) -> Callable[[], None]: + global _next_callback_id + id = _next_callback_id + _next_callback_id += 1 + + disposed = False + + def dispose() -> None: + nonlocal disposed + + with _module_rlock: + if disposed: + return + + callback_dict.pop(id, None) + + disposed = True + + callback_dict[id] = callback + return dispose + + +def _patch( + stdout_or_stderr: IO[AnyStr], + callbacks: dict[int, _WriteCallback], +) -> None: + orig_write: Callable[[AnyStr], int] + + @wb_logging.log_to_all_runs() + def write_with_callbacks(s: AnyStr, /) -> int: + global _is_writing + n = orig_write(s) + + # NOTE: Since _module_rlock is reentrant, this is safe. It will not + # deadlock if a callback invokes write() again. + with _module_rlock: + if _is_writing: + return n + + _is_writing = True + try: + for cb in callbacks.values(): + cb(s, n) + + except BaseException as e: + # Clear all callbacks on any exception to avoid infinite loops: + # + # * If we re-raise, an exception handler is likely to print + # the exception to the console and trigger callbacks again + # * If we log, we can't guarantee that this doesn't print + # to console. + # + # This is especially important for KeyboardInterrupt. + _stderr_callbacks.clear() + _stdout_callbacks.clear() + + if isinstance(e, Exception): + # We suppress Exceptions so that bugs in W&B code don't + # cause the user's print() statements to raise errors. + _logger.exception("Error in console callback, clearing all!") + else: + # Re-raise errors like KeyboardInterrupt. + raise + + finally: + _is_writing = False + + return n + + orig_write = stdout_or_stderr.write + + # mypy==1.14.1 fails to type-check this: + # Incompatible types in assignment (expression has type + # "Callable[[bytes], int]", variable has type overloaded function) + stdout_or_stderr.write = write_with_callbacks # type: ignore + + +try: + _patch(sys.stdout, _stdout_callbacks) + _patch(sys.stderr, _stderr_callbacks) +except Exception as _patch_exception_cause: + _patch_exception = CannotCaptureConsoleError() + _patch_exception.__cause__ = _patch_exception_cause diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/credentials.py b/lib/python3.12/site-packages/wandb/sdk/lib/credentials.py new file mode 100644 index 0000000000000000000000000000000000000000..422cf0efe6e219c9edd38f8edd7a8d23b9c89d9d --- /dev/null +++ b/lib/python3.12/site-packages/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 diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/deprecate.py b/lib/python3.12/site-packages/wandb/sdk/lib/deprecate.py new file mode 100644 index 0000000000000000000000000000000000000000..2904ae32f3ca6c2f225885c0df42d1c02e4bd08e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/deprecate.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import wandb +from wandb.proto.wandb_deprecated import DEPRECATED_FEATURES +from wandb.sdk.lib import telemetry + +# Necessary to break import cycle. +if TYPE_CHECKING: + from wandb import wandb_run + + +def deprecate( + field_name: DEPRECATED_FEATURES, + warning_message: str, + run: wandb_run.Run | None = None, +) -> None: + """Warn the user that a feature has been deprecated. + + If a run is provided, the given field on its telemetry is updated. + Otherwise, the global run is used. + + Args: + field_name: The field on the Deprecated proto for this deprecation. + warning_message: The message to display to the user. + run: The run whose telemetry to update. + """ + _run = run or wandb.run + with telemetry.context(run=_run) as tel: + setattr(tel.deprecated, field_name, True) + + wandb.termwarn(warning_message, repeat=False) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/disabled.py b/lib/python3.12/site-packages/wandb/sdk/lib/disabled.py new file mode 100644 index 0000000000000000000000000000000000000000..9b5df0608ea1518697c6883163beeb2a700c6874 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/disabled.py @@ -0,0 +1,30 @@ +from typing import Any + +from wandb.proto.wandb_deprecated import Deprecated +from wandb.sdk.lib import deprecate + + +class SummaryDisabled(dict): + __setattr__ = dict.__setitem__ + __delattr__ = dict.__delitem__ + + def __getattr__(self, key): + return self[key] + + def __getitem__(self, key): + val = dict.__getitem__(self, key) + if isinstance(val, dict) and not isinstance(val, SummaryDisabled): + val = SummaryDisabled(val) + self[key] = val + return val + + +class RunDisabled: + """Compatibility class for integrations that explicitly check for wandb.RunDisabled.""" + + def __getattr__(self, name: str) -> Any: + deprecate.deprecate( + field_name=Deprecated.run_disabled, + warning_message="RunDisabled is deprecated and is a no-op. " + '`wandb.init(mode="disabled")` now returns and instance of `wandb.sdk.wandb_run.Run`.', + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/exit_hooks.py b/lib/python3.12/site-packages/wandb/sdk/lib/exit_hooks.py new file mode 100644 index 0000000000000000000000000000000000000000..aa747495299a646041f7deaeb72c1a9bcc050660 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/exit_hooks.py @@ -0,0 +1,54 @@ +import sys +import traceback +from types import TracebackType +from typing import TYPE_CHECKING, Optional, Type + +import wandb +from wandb.errors import Error + +if TYPE_CHECKING: + from typing import NoReturn + + +class ExitHooks: + exception: Optional[BaseException] = None + + def __init__(self) -> None: + self.exit_code = 0 + self.exception = None + + def hook(self) -> None: + self._orig_exit = sys.exit + sys.exit = self.exit + self._orig_excepthook = ( + sys.excepthook + if sys.excepthook + != sys.__excepthook__ # respect hooks by other libraries like pdb + else None + ) + sys.excepthook = self.exc_handler # type: ignore + + def exit(self, code: object = 0) -> "NoReturn": + orig_code = code + code = code if code is not None else 0 + code = code if isinstance(code, int) else 1 + self.exit_code = code + self._orig_exit(orig_code) # type: ignore + + def was_ctrl_c(self) -> bool: + return isinstance(self.exception, KeyboardInterrupt) + + def exc_handler( + self, exc_type: Type[BaseException], exc: BaseException, tb: TracebackType + ) -> None: + self.exit_code = 1 + self.exception = exc + if issubclass(exc_type, Error): + wandb.termerror(str(exc), repeat=False) + + if self.was_ctrl_c(): + self.exit_code = 255 + + traceback.print_exception(exc_type, exc, tb) + if self._orig_excepthook: + self._orig_excepthook(exc_type, exc, tb) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/file_stream_utils.py b/lib/python3.12/site-packages/wandb/sdk/lib/file_stream_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b7420b99706f8cd413da35b50ab76c1bd1e8fce6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/file_stream_utils.py @@ -0,0 +1,118 @@ +# +from typing import Any, Dict, Iterable + + +def split_files( + files: Dict[str, Any], max_bytes: int = 10 * 1024 * 1024 +) -> Iterable[Dict[str, Dict]]: + """Split a file's dict (see `files` arg) into smaller dicts. + + Each smaller dict will have at most `MAX_BYTES` size. + + This method is used in `FileStreamAPI._send()` to limit the size of post requests + sent to wandb server. + + Args: + files (dict): `dict` of form {file_name: {'content': ".....", 'offset': 0}} + The key `file_name` can also be mapped to a List [{"offset": int, "content": str}] + `max_bytes`: max size for chunk in bytes + """ + current_volume: Dict[str, Dict] = {} + current_size = 0 + + def _str_size(x): + return len(x) if isinstance(x, bytes) else len(x.encode("utf-8")) + + def _file_size(file): + size = file.get("_size") + if size is None: + size = sum(map(_str_size, file["content"])) + file["_size"] = size + return size + + def _split_file(file, num_lines): + offset = file["offset"] + content = file["content"] + name = file["name"] + f1 = {"offset": offset, "content": content[:num_lines], "name": name} + f2 = { + "offset": offset + num_lines, + "content": content[num_lines:], + "name": name, + } + return f1, f2 + + def _num_lines_from_num_bytes(file, num_bytes): + size = 0 + num_lines = 0 + content = file["content"] + while num_lines < len(content): + size += _str_size(content[num_lines]) + if size > num_bytes: + break + num_lines += 1 + return num_lines + + files_stack = [] + for k, v in files.items(): + if isinstance(v, list): + for item in v: + files_stack.append( + {"name": k, "offset": item["offset"], "content": item["content"]} + ) + else: + files_stack.append( + {"name": k, "offset": v["offset"], "content": v["content"]} + ) + + while files_stack: + f = files_stack.pop() + if f["name"] in current_volume: + files_stack.append(f) + yield current_volume + current_volume = {} + current_size = 0 + continue + # For each file, we have to do 1 of 4 things: + # - Add the file as such to the current volume if possible. + # - Split the file and add the first part to the current volume and push the second part back onto the stack. + # - If that's not possible, check if current volume is empty: + # - If empty, add first line of file to current volume and push rest onto stack (This volume will exceed MAX_MB). + # - If not, push file back to stack and yield current volume. + fsize = _file_size(f) + rem = max_bytes - current_size + if fsize <= rem: + current_volume[f["name"]] = { + "offset": f["offset"], + "content": f["content"], + } + current_size += fsize + else: + num_lines = _num_lines_from_num_bytes(f, rem) + if not num_lines and not current_volume: + num_lines = 1 + if num_lines: + f1, f2 = _split_file(f, num_lines) + current_volume[f1["name"]] = { + "offset": f1["offset"], + "content": f1["content"], + } + files_stack.append(f2) + yield current_volume + current_volume = {} + current_size = 0 + continue + else: + files_stack.append(f) + yield current_volume + current_volume = {} + current_size = 0 + continue + if current_size >= max_bytes: + yield current_volume + current_volume = {} + current_size = 0 + continue + + if current_volume: + yield current_volume diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/filenames.py b/lib/python3.12/site-packages/wandb/sdk/lib/filenames.py new file mode 100644 index 0000000000000000000000000000000000000000..c272b7a2c3337ddb6d91ce03ed800e84affcc5ff --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/filenames.py @@ -0,0 +1,64 @@ +import os +from typing import Callable, Generator, Union + +WANDB_DIRS = ("wandb", ".wandb") + +CONFIG_FNAME = "config.yaml" +OUTPUT_FNAME = "output.log" +DIFF_FNAME = "diff.patch" +SUMMARY_FNAME = "wandb-summary.json" +METADATA_FNAME = "wandb-metadata.json" +REQUIREMENTS_FNAME = "requirements.txt" +HISTORY_FNAME = "wandb-history.jsonl" +EVENTS_FNAME = "wandb-events.jsonl" +JOBSPEC_FNAME = "wandb-jobspec.json" +CONDA_ENVIRONMENTS_FNAME = "conda-environment.yaml" + + +def is_wandb_file(name: str) -> bool: + return ( + name.startswith("wandb") + or name == METADATA_FNAME + or name == CONFIG_FNAME + or name == REQUIREMENTS_FNAME + or name == OUTPUT_FNAME + or name == DIFF_FNAME + or name == CONDA_ENVIRONMENTS_FNAME + ) + + +def filtered_dir( + root: str, + include_fn: Union[Callable[[str, str], bool], Callable[[str], bool]], + exclude_fn: Union[Callable[[str, str], bool], Callable[[str], bool]], +) -> Generator[str, None, None]: + """Simple generator to walk a directory.""" + import inspect + + # compatibility with old API, which didn't pass root + def _include_fn(path: str, root: str) -> bool: + return ( + include_fn(path, root) # type: ignore + if len(inspect.signature(include_fn).parameters) == 2 + else include_fn(path) # type: ignore + ) + + def _exclude_fn(path: str, root: str) -> bool: + return ( + exclude_fn(path, root) # type: ignore + if len(inspect.signature(exclude_fn).parameters) == 2 + else exclude_fn(path) # type: ignore + ) + + for dirpath, _, files in os.walk(root): + for fname in files: + file_path = os.path.join(dirpath, fname) + if _include_fn(file_path, root) and not _exclude_fn(file_path, root): + yield file_path + + +def exclude_wandb_fn(path: str, root: str) -> bool: + return any( + os.path.relpath(path, root).startswith(wandb_dir + os.sep) + for wandb_dir in WANDB_DIRS + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/filesystem.py b/lib/python3.12/site-packages/wandb/sdk/lib/filesystem.py new file mode 100644 index 0000000000000000000000000000000000000000..180ebbb2c19aaa522e9d839141c9c0a740666117 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/filesystem.py @@ -0,0 +1,372 @@ +import contextlib +import ctypes +import errno +import logging +import os +import platform +import re +import shutil +import tempfile +import threading +from pathlib import Path +from typing import IO, Any, BinaryIO, Generator, Optional + +from wandb.sdk.lib.paths import StrPath + +logger = logging.getLogger(__name__) + +# https://en.wikipedia.org/wiki/Filename#Comparison_of_filename_limitations +PROBLEMATIC_PATH_CHARS = "".join(chr(i) for i in range(32)) + ':"*<>?|' + + +def mkdir_exists_ok(dir_name: StrPath) -> None: + """Create `dir_name` and any parent directories if they don't exist. + + Raises: + FileExistsError: if `dir_name` exists and is not a directory. + PermissionError: if `dir_name` is not writable. + """ + try: + os.makedirs(dir_name, exist_ok=True) + except FileExistsError as e: + raise FileExistsError(f"{dir_name!s} exists and is not a directory") from e + except PermissionError as e: + raise PermissionError(f"{dir_name!s} is not writable") from e + + +def path_fallbacks(path: StrPath) -> Generator[str, None, None]: + """Yield variations of `path` that may exist on the filesystem. + + Return a sequence of paths that should be checked in order for existence or + create-ability. Essentially, keep replacing "suspect" characters until we run out. + """ + path = str(path) + root, tail = os.path.splitdrive(path) + yield os.path.join(root, tail) + for char in PROBLEMATIC_PATH_CHARS: + if char in tail: + tail = tail.replace(char, "-") + yield os.path.join(root, tail) + + +def mkdir_allow_fallback(dir_name: StrPath) -> StrPath: + """Create `dir_name`, removing invalid path characters if necessary. + + Returns: + The path to the created directory, which may not be the original path. + """ + for new_name in path_fallbacks(dir_name): + try: + os.makedirs(new_name, exist_ok=True) + if Path(new_name) != Path(dir_name): + logger.warning(f"Creating '{new_name}' instead of '{dir_name}'") + return Path(new_name) if isinstance(dir_name, Path) else new_name + except (ValueError, NotADirectoryError): + pass + except OSError as e: + if e.errno != 22: + raise + + raise OSError(f"Unable to create directory '{dir_name}'") + + +def files_in(path: StrPath) -> Generator[os.DirEntry, None, None]: + """Yield a directory entry for each file under a given path (recursive).""" + if not os.path.isdir(path): + return + for entry in os.scandir(path): + if entry.is_dir(): + yield from files_in(entry.path) + else: + yield entry + + +class WriteSerializingFile: + """Wrapper for a file object that serializes writes.""" + + def __init__(self, f: BinaryIO) -> None: + self.lock = threading.Lock() + self.f = f + + def write(self, *args, **kargs) -> None: # type: ignore + self.lock.acquire() + try: + self.f.write(*args, **kargs) + self.f.flush() + finally: + self.lock.release() + + def close(self) -> None: + self.lock.acquire() # wait for pending writes + try: + self.f.close() + finally: + self.lock.release() + + +class CRDedupedFile(WriteSerializingFile): + def __init__(self, f: BinaryIO) -> None: + super().__init__(f=f) + self._buff = b"" + + def write(self, data) -> None: # type: ignore + lines = re.split(b"\r\n|\n", data) + ret = [] # type: ignore + for line in lines: + if line[:1] == b"\r": + if ret: + ret.pop() + elif self._buff: + self._buff = b"" + line = line.split(b"\r")[-1] + if line: + ret.append(line) + if self._buff: + ret.insert(0, self._buff) + if ret: + self._buff = ret.pop() + super().write(b"\n".join(ret) + b"\n") + + def close(self) -> None: + if self._buff: + super().write(self._buff) + super().close() + + +def copy_or_overwrite_changed(source_path: StrPath, target_path: StrPath) -> StrPath: + """Copy source_path to target_path, unless it already exists with the same mtime. + + We liberally add write permissions to deal with the case of multiple users needing + to share the same cache or run directory. + + Args: + source_path: The path to the file to copy. + target_path: The path to copy the file to. + + Returns: + The path to the copied file (which may be different from target_path). + """ + return_type = type(target_path) + + target_path = system_preferred_path(target_path, warn=True) + + need_copy = ( + not os.path.isfile(target_path) + or os.stat(source_path).st_mtime != os.stat(target_path).st_mtime + ) + + permissions_plus_write = os.stat(source_path).st_mode + if need_copy: + dir_name, file_name = os.path.split(target_path) + target_path = os.path.join(mkdir_allow_fallback(dir_name), file_name) + try: + # Use copy2 to preserve file metadata (including modified time). + shutil.copy2(source_path, target_path) + except PermissionError: + # If the file is read-only try to make it writable. + try: + os.chmod(target_path, permissions_plus_write) + shutil.copy2(source_path, target_path) + except PermissionError as e: + raise PermissionError("Unable to overwrite '{target_path!s}'") from e + # Prevent future permissions issues by universal write permissions now. + os.chmod(target_path, permissions_plus_write) + + return return_type(target_path) # type: ignore # 'os.PathLike' is abstract. + + +@contextlib.contextmanager +def safe_open( + path: StrPath, mode: str = "r", *args: Any, **kwargs: Any +) -> Generator[IO, None, None]: + """Open a file, ensuring any changes only apply atomically after close. + + This context manager ensures that even unsuccessful writes will not leave a "dirty" + file or overwrite good data, and that all temp data is cleaned up. + + The semantics and behavior are intended to be nearly identical to the built-in + open() function. Differences: + - It creates any parent directories that don't exist, rather than raising. + - In 'x' mode, it checks at the beginning AND end of the write and fails if the + file exists either time. + """ + path = Path(path).resolve() + path.parent.mkdir(parents=True, exist_ok=True) + + if "x" in mode and path.exists(): + raise FileExistsError(f"{path!s} already exists") + + if "r" in mode and "+" not in mode: + # This is read-only, so we can just open the original file. + # TODO (hugh): create a reflink and read from that. + with path.open(mode, *args, **kwargs) as f: + yield f + return + + with tempfile.TemporaryDirectory(dir=path.parent) as tmp_dir: + tmp_path = Path(tmp_dir) / path.name + + if ("r" in mode or "a" in mode) and path.exists(): + # We need to copy the original file in order to support reads and appends. + # TODO (hugh): use reflinks to avoid the copy on platforms that support it. + shutil.copy2(path, tmp_path) + + with tmp_path.open(mode, *args, **kwargs) as f: + yield f + f.flush() + os.fsync(f.fileno()) + + if "x" in mode: + # Ensure that if another process has beaten us to writing the file we raise + # rather than overwrite. os.link() atomically creates a hard link to the + # target file and will raise FileExistsError if the target already exists. + os.link(tmp_path, path) + os.unlink(tmp_path) + else: + tmp_path.replace(path) + + +def safe_copy(source_path: StrPath, target_path: StrPath) -> StrPath: + """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) + with tempfile.TemporaryDirectory(dir=output_path.parent) as tmp_dir: + tmp_path = (Path(tmp_dir) / Path(source_path).name).with_suffix(".tmp") + shutil.copy2(source_path, tmp_path) + tmp_path.replace(output_path) + return target_path + + +def _reflink_linux(existing_path: Path, new_path: Path) -> None: + """Create a reflink to `existing_path` at `new_path` on Linux.""" + import fcntl + + FICLONE = 0x40049409 # magic number from # noqa: N806 + with open(existing_path, "rb") as t_f, open(new_path, "wb+") as l_f: + fcntl.ioctl(l_f.fileno(), FICLONE, t_f.fileno()) + + +def _reflink_macos(existing_path: Path, new_path: Path) -> None: + try: + clib = ctypes.CDLL("libc.dylib", use_errno=True) + except (FileNotFoundError, OSError) as e: + if ctypes.get_errno() != errno.ENOENT and not isinstance(e, FileNotFoundError): + raise + # Before macOS 11 ( None: + """Create a reflink to `existing_path` at `new_path`. + + A reflink (reflective link) is a copy-on-write reference to a file. Once linked, the + file and link are both "real" files (not symbolic or hard links) and each can be + modified independently without affecting the other; however, they share the same + underlying data blocks on disk so until one is modified they are "zero-cost" copies. + + Reflinks have all the functionality of copies, so we should use them wherever they + are supported if we would otherwise copy a file. (This is not particularly radical-- + GNU `cp` defaults to `reflink=auto`, using it whenever available) However, support + for them is limited to a small number of filesystems. They should work on: + - Linux with a Btrfs or XFS filesystem (NOT ext4) + - macOS 10.13 or later with an APFS filesystem (called clone files) + + Reflinks are also supported on Solaris and Windows with ReFSv2, but we haven't + implemented support for them. + + Like hard links, a reflink can only be created on the same filesystem as the target. + """ + if platform.system() == "Linux": + link_fn = _reflink_linux + elif platform.system() == "Darwin": + link_fn = _reflink_macos + else: + raise OSError( + errno.ENOTSUP, f"reflinks are not supported on {platform.system()}" + ) + + new_path = Path(new_path).resolve() + existing_path = Path(existing_path).resolve() + if new_path.exists(): + if not overwrite: + raise FileExistsError(f"{new_path} already exists") + logger.warning(f"Overwriting existing file {new_path}.") + new_path.unlink() + + # Create any missing parent directories. + new_path.parent.mkdir(parents=True, exist_ok=True) + + try: + link_fn(existing_path, new_path) + except OSError as e: + base_msg = f"failed to create reflink from {existing_path} to {new_path}." + if e.errno in (errno.EPERM, errno.EACCES): + raise PermissionError(f"Insufficient permissions; {base_msg}") from e + if e.errno == errno.ENOENT: + raise FileNotFoundError(f"File not found; {base_msg}") from e + if e.errno == errno.EXDEV: + raise ValueError(f"Cannot link across filesystems; {base_msg}") from e + if e.errno == errno.EISDIR: + raise IsADirectoryError(f"Cannot reflink a directory; {base_msg}") from e + if e.errno in (errno.EOPNOTSUPP, errno.ENOTSUP): + raise OSError( + errno.ENOTSUP, + f"Filesystem does not support reflinks; {base_msg}", + ) from e + if e.errno == errno.EINVAL: + raise ValueError(f"Cannot link file ranges; {base_msg}") from e + raise + + +def check_exists(path: StrPath) -> Optional[StrPath]: + """Look for variations of `path` and return the first found. + + This exists to support former behavior around system-dependent paths; we used to use + ':' in Artifact paths unless we were on Windows, but this has issues when e.g. a + Linux machine is accessing an NTFS filesystem; we might need to look for the + alternate path. This checks all the possible directories we would consider creating. + """ + for dest in path_fallbacks(path): + if os.path.exists(dest): + return Path(dest) if isinstance(path, Path) else dest + return None + + +def system_preferred_path(path: StrPath, warn: bool = False) -> StrPath: + """Replace ':' with '-' in paths on Windows. + + Args: + path: The path to convert. + warn: Whether to warn if ':' is replaced. + """ + if platform.system() != "Windows": + return path + head, tail = os.path.splitdrive(path) + if warn and ":" in tail: + logger.warning(f"Replacing ':' in {tail} with '-'") + new_path = head + tail.replace(":", "-") + return Path(new_path) if isinstance(path, Path) else new_path diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/fsm.py b/lib/python3.12/site-packages/wandb/sdk/lib/fsm.py new file mode 100644 index 0000000000000000000000000000000000000000..d36bb472f7a39668896488115dfcb9cdaee9ae78 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/fsm.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python +"""Finite state machine. + +Simple FSM implementation. + +Usage: + ```python + class A: + def on_output(self, inputs) -> None: + pass + + + class B: + def on_output(self, inputs) -> None: + pass + + + def to_b(inputs) -> bool: + return True + + + def to_a(inputs) -> bool: + return True + + + f = Fsm(states=[A(), B()], table={A: [(to_b, B)], B: [(to_a, A)]}) + f.run({"input1": 1, "input2": 2}) + ``` +""" + +from abc import abstractmethod +from dataclasses import dataclass +from typing import Callable, Dict, Generic, Optional, Sequence, Type, Union + +from typing_extensions import Protocol, TypeAlias, TypeVar, runtime_checkable + +T_FsmInputs = TypeVar("T_FsmInputs", contravariant=True) +T_FsmContext = TypeVar("T_FsmContext") +T_FsmContext_cov = TypeVar("T_FsmContext_cov", covariant=True) +T_FsmContext_contra = TypeVar("T_FsmContext_contra", contravariant=True) + + +@runtime_checkable +class FsmStateCheck(Protocol[T_FsmInputs]): + @abstractmethod + def on_check(self, inputs: T_FsmInputs) -> None: ... # pragma: no cover + + +@runtime_checkable +class FsmStateOutput(Protocol[T_FsmInputs]): + @abstractmethod + def on_state(self, inputs: T_FsmInputs) -> None: ... # pragma: no cover + + +@runtime_checkable +class FsmStateEnter(Protocol[T_FsmInputs]): + @abstractmethod + def on_enter(self, inputs: T_FsmInputs) -> None: ... # pragma: no cover + + +@runtime_checkable +class FsmStateEnterWithContext(Protocol[T_FsmInputs, T_FsmContext_contra]): + @abstractmethod + def on_enter( + self, inputs: T_FsmInputs, context: T_FsmContext_contra + ) -> None: ... # pragma: no cover + + +@runtime_checkable +class FsmStateStay(Protocol[T_FsmInputs]): + @abstractmethod + def on_stay(self, inputs: T_FsmInputs) -> None: ... # pragma: no cover + + +@runtime_checkable +class FsmStateExit(Protocol[T_FsmInputs, T_FsmContext_cov]): + @abstractmethod + def on_exit(self, inputs: T_FsmInputs) -> T_FsmContext_cov: ... # pragma: no cover + + +# It would be nice if python provided optional protocol members, but it does not as described here: +# https://peps.python.org/pep-0544/#support-optional-protocol-members +# Until then, we can only enforce that a state at least supports one protocol interface. This +# unfortunately will not check the signature of other potential protocols. +FsmState: TypeAlias = Union[ + FsmStateCheck[T_FsmInputs], + FsmStateOutput[T_FsmInputs], + FsmStateEnter[T_FsmInputs], + FsmStateEnterWithContext[T_FsmInputs, T_FsmContext], + FsmStateStay[T_FsmInputs], + FsmStateExit[T_FsmInputs, T_FsmContext], +] + + +@dataclass +class FsmEntry(Generic[T_FsmInputs, T_FsmContext]): + condition: Callable[[T_FsmInputs], bool] + target_state: Type[FsmState[T_FsmInputs, T_FsmContext]] + action: Optional[Callable[[T_FsmInputs], None]] = None + + +FsmTableWithContext: TypeAlias = Dict[ + Type[FsmState[T_FsmInputs, T_FsmContext]], + Sequence[FsmEntry[T_FsmInputs, T_FsmContext]], +] + + +FsmTable: TypeAlias = FsmTableWithContext[T_FsmInputs, None] + + +class FsmWithContext(Generic[T_FsmInputs, T_FsmContext]): + _state_dict: Dict[Type[FsmState], FsmState] + _table: FsmTableWithContext[T_FsmInputs, T_FsmContext] + _state: FsmState[T_FsmInputs, T_FsmContext] + _states: Sequence[FsmState] + + def __init__( + self, + states: Sequence[FsmState], + table: FsmTableWithContext[T_FsmInputs, T_FsmContext], + ) -> None: + self._states = states + self._table = table + self._state_dict = {type(s): s for s in states} + self._state = self._state_dict[type(states[0])] + + def _transition( + self, + inputs: T_FsmInputs, + new_state: Type[FsmState[T_FsmInputs, T_FsmContext]], + action: Optional[Callable[[T_FsmInputs], None]], + ) -> None: + if action: + action(inputs) + + context = None + if isinstance(self._state, FsmStateExit): + context = self._state.on_exit(inputs) + + prev_state = type(self._state) + if prev_state == new_state: + if isinstance(self._state, FsmStateStay): + self._state.on_stay(inputs) + else: + self._state = self._state_dict[new_state] + if context and isinstance(self._state, FsmStateEnterWithContext): + self._state.on_enter(inputs, context=context) + elif isinstance(self._state, FsmStateEnter): + self._state.on_enter(inputs) + + def _check_transitions(self, inputs: T_FsmInputs) -> None: + for entry in self._table[type(self._state)]: + if entry.condition(inputs): + self._transition(inputs, entry.target_state, entry.action) + return + + def input(self, inputs: T_FsmInputs) -> None: + if isinstance(self._state, FsmStateCheck): + self._state.on_check(inputs) + self._check_transitions(inputs) + if isinstance(self._state, FsmStateOutput): + self._state.on_state(inputs) + + +Fsm: TypeAlias = FsmWithContext[T_FsmInputs, None] diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/gitlib.py b/lib/python3.12/site-packages/wandb/sdk/lib/gitlib.py new file mode 100644 index 0000000000000000000000000000000000000000..b2f18a42e101ae3a78764282421a2fb0c57ee673 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/gitlib.py @@ -0,0 +1,240 @@ +import configparser +import logging +import os +from typing import TYPE_CHECKING, Any, Optional +from urllib.parse import urlparse, urlunparse + +import wandb + +try: + from git import ( # type: ignore + GitCommandError, + InvalidGitRepositoryError, + NoSuchPathError, + Repo, + ) +except ImportError: + Repo = None # type: ignore + +if TYPE_CHECKING: + from git import Repo + + +logger = logging.getLogger(__name__) + + +class GitRepo: + def __init__( + self, + root: Optional[str] = None, + remote: str = "origin", + lazy: bool = True, + remote_url: Optional[str] = None, + commit: Optional[str] = None, + ) -> None: + self.remote_name = remote if remote_url is None else None + self._root = root + self._remote_url = remote_url + self._commit = commit + self._repo = None + self._repo_initialized = False + if not lazy: + self._repo = self._init_repo() + + def _init_repo(self) -> Optional[Repo]: + self._repo_initialized = True + if Repo is None: + return None + if self.remote_name is None: + return None + try: + return Repo(self._root or os.getcwd(), search_parent_directories=True) + except FileNotFoundError: + wandb.termwarn("current working directory has been invalidated") + logger.warning("current working directory has been invalidated") + except InvalidGitRepositoryError: + logger.debug("git repository is invalid") + except NoSuchPathError: + wandb.termwarn(f"git root {self._root} does not exist") + logger.warning(f"git root {self._root} does not exist") + return None + + @property + def repo(self) -> Optional[Repo]: + if not self._repo_initialized: + self._repo = self._init_repo() + return self._repo + + @property + def auto(self) -> bool: + return self._remote_url is None + + def is_untracked(self, file_name: str) -> Optional[bool]: + if not self.repo: + return True + try: + return file_name in self.repo.untracked_files + except GitCommandError: + return None + + @property + def enabled(self) -> bool: + return bool(self.repo) + + @property + def root(self) -> Any: + if not self.repo: + return None + try: + return self.repo.git.rev_parse("--show-toplevel") + except GitCommandError: + # todo: collect telemetry on this + logger.exception("git root error") + return None + + @property + def dirty(self) -> Any: + if not self.repo: + return False + try: + return self.repo.is_dirty() + except GitCommandError: + return False + + @property + def email(self) -> Optional[str]: + if not self.repo: + return None + try: + return self.repo.config_reader().get_value("user", "email") # type: ignore + except configparser.Error: + return None + + @property + def last_commit(self) -> Any: + if self._commit: + return self._commit + if not self.repo: + return None + if not self.repo.head or not self.repo.head.is_valid(): + return None + # TODO: Saw a user getting a Unicode decode error when parsing refs, + # more details on implementing a real fix in [WB-4064] + try: + if len(self.repo.refs) > 0: # type: ignore[arg-type] + return self.repo.head.commit.hexsha + else: + return self.repo.git.show_ref("--head").split(" ")[0] + except Exception: + logger.exception("Unable to find most recent commit in git") + return None + + @property + def branch(self) -> Any: + if not self.repo: + return None + return self.repo.head.ref.name + + @property + def remote(self) -> Any: + if not self.repo: + return None + try: + return self.repo.remotes[self.remote_name] # type: ignore[index] + except IndexError: + return None + + # the --submodule=diff option doesn't exist in pre-2.11 versions of git (november 2016) + # https://stackoverflow.com/questions/10757091/git-list-of-all-changed-files-including-those-in-submodules + @property + def has_submodule_diff(self) -> bool: + if not self.repo: + return False + return bool(self.repo.git.version_info >= (2, 11, 0)) + + @property + def remote_url(self) -> Any: + if self._remote_url: + return self._remote_url + if not self.remote: + return None + parsed = urlparse(self.remote.url) + hostname = parsed.hostname + if parsed.port is not None: + hostname = f"{hostname}:{parsed.port}" + if parsed.password is not None: + return urlunparse(parsed._replace(netloc=f"{parsed.username}:@{hostname}")) + return urlunparse(parsed._replace(netloc=hostname)) + + @property + def root_dir(self) -> Any: + if not self.repo: + return None + try: + return self.repo.git.rev_parse("--show-toplevel") + except GitCommandError: + return None + + def get_upstream_fork_point(self) -> Any: + """Get the most recent ancestor of HEAD that occurs on an upstream branch. + + First looks at the current branch's tracking branch, if applicable. If + that doesn't work, looks at every other branch to find the most recent + ancestor of HEAD that occurs on a tracking branch. + + Returns: + git.Commit object or None + """ + possible_relatives = [] + try: + if not self.repo: + return None + try: + active_branch = self.repo.active_branch + except (TypeError, ValueError): + logger.debug("git is in a detached head state") + return None # detached head + else: + tracking_branch = active_branch.tracking_branch() + if tracking_branch: + possible_relatives.append(tracking_branch.commit) + + if not possible_relatives: + for branch in self.repo.branches: # type: ignore[attr-defined] + tracking_branch = branch.tracking_branch() + if tracking_branch is not None: + possible_relatives.append(tracking_branch.commit) + + head = self.repo.head + most_recent_ancestor = None + for possible_relative in possible_relatives: + # at most one: + for ancestor in self.repo.merge_base(head, possible_relative): + if most_recent_ancestor is None: + most_recent_ancestor = ancestor + elif self.repo.is_ancestor(most_recent_ancestor, ancestor): # type: ignore + most_recent_ancestor = ancestor + except GitCommandError as e: + logger.debug("git remote upstream fork point could not be found") + logger.debug(str(e)) + return None + + return most_recent_ancestor + + def tag(self, name: str, message: Optional[str]) -> Any: + if not self.repo: + return None + try: + return self.repo.create_tag(f"wandb/{name}", message=message, force=True) + except GitCommandError: + logger.debug("Failed to tag repository.") + return None + + def push(self, name: str) -> Any: + if not self.remote: + return None + try: + return self.remote.push(f"wandb/{name}", force=True) + except GitCommandError: + logger.debug("failed to push git") + return None diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/gql_request.py b/lib/python3.12/site-packages/wandb/sdk/lib/gql_request.py new file mode 100644 index 0000000000000000000000000000000000000000..0ef27e8e8cf235a19859d93ffac9a4daa57c58bf --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/gql_request.py @@ -0,0 +1,65 @@ +"""A simple GraphQL client for sending queries and mutations. + +Note: This was originally wandb/vendor/gql-0.2.0/wandb_gql/transport/requests.py +The only substantial change is to reuse a requests.Session object. +""" + +from typing import Any, Callable, Dict, Optional, Tuple, Union + +import requests +from wandb_gql.transport.http import HTTPTransport +from wandb_graphql.execution import ExecutionResult +from wandb_graphql.language import ast +from wandb_graphql.language.printer import print_ast + + +class GraphQLSession(HTTPTransport): + def __init__( + self, + url: str, + auth: Optional[Union[Tuple[str, str], Callable]] = None, + use_json: bool = False, + timeout: Optional[Union[int, float]] = None, + proxies: Optional[Dict[str, str]] = None, + **kwargs: Any, + ) -> None: + """Setup a session for sending GraphQL queries and mutations. + + Args: + url (str): The GraphQL URL + auth (tuple or callable): Auth tuple or callable for Basic/Digest/Custom HTTP Auth + use_json (bool): Send request body as JSON instead of form-urlencoded + timeout (int, float): Specifies a default timeout for requests (Default: None) + """ + super().__init__(url, **kwargs) + self.session = requests.Session() + if proxies: + self.session.proxies.update(proxies) + self.session.auth = auth + self.default_timeout = timeout + self.use_json = use_json + + def execute( + self, + document: ast.Node, + variable_values: Optional[Dict] = None, + timeout: Optional[Union[int, float]] = None, + ) -> ExecutionResult: + query_str = print_ast(document) + payload = {"query": query_str, "variables": variable_values or {}} + + data_key = "json" if self.use_json else "data" + post_args = { + "headers": self.headers, + "cookies": self.cookies, + "timeout": timeout or self.default_timeout, + data_key: payload, + } + request = self.session.post(self.url, **post_args) + request.raise_for_status() + + result = request.json() + data, errors = result.get("data"), result.get("errors") + if data is None and errors is None: + raise RuntimeError(f"Received non-compatible response: {result}") + return ExecutionResult(data=data, errors=errors) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/handler_util.py b/lib/python3.12/site-packages/wandb/sdk/lib/handler_util.py new file mode 100644 index 0000000000000000000000000000000000000000..b4efd8d571baebab80f9a7702efdd962f0c9028d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/handler_util.py @@ -0,0 +1,21 @@ +import wandb.data_types as data_types + + +def get_types(): + classes = map(data_types.__dict__.get, data_types.__all__) + types = [] + for cls in classes: + if hasattr(cls, "_log_type") and cls._log_type is not None: + types.append(cls._log_type) + # add table-file type because this is a special case + # that does not have a matching _log_type for artifacts + # and files + types.append("table-file") + return types + + +WANDB_TYPES = get_types() + + +def metric_is_wandb_dict(metric): + return "_type" in list(metric.keys()) and metric["_type"] in WANDB_TYPES diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/hashutil.py b/lib/python3.12/site-packages/wandb/sdk/lib/hashutil.py new file mode 100644 index 0000000000000000000000000000000000000000..cd2aa1104927f2c1b023edb9d6261ee16468b847 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/hashutil.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import base64 +import hashlib +import mmap +import sys +from typing import TYPE_CHECKING, NewType + +from wandb.sdk.lib.paths import StrPath + +if TYPE_CHECKING: + import _hashlib # type: ignore[import-not-found] + +ETag = NewType("ETag", str) +HexMD5 = NewType("HexMD5", str) +B64MD5 = NewType("B64MD5", str) + + +def _md5(data: bytes = b"") -> _hashlib.HASH: + """Allow FIPS-compliant md5 hash when supported.""" + if sys.version_info >= (3, 9): + return hashlib.md5(data, usedforsecurity=False) + else: + return hashlib.md5(data) + + +def md5_string(string: str) -> B64MD5: + return _b64_from_hasher(_md5(string.encode("utf-8"))) + + +def _b64_from_hasher(hasher: _hashlib.HASH) -> B64MD5: + return B64MD5(base64.b64encode(hasher.digest()).decode("ascii")) + + +def b64_to_hex_id(string: B64MD5) -> HexMD5: + return HexMD5(base64.standard_b64decode(string).hex()) + + +def hex_to_b64_id(encoded_string: str | bytes) -> B64MD5: + if isinstance(encoded_string, bytes): + encoded_string = encoded_string.decode("utf-8") + as_str = bytes.fromhex(encoded_string) + return B64MD5(base64.standard_b64encode(as_str).decode("utf-8")) + + +def md5_file_b64(*paths: StrPath) -> B64MD5: + return _b64_from_hasher(_md5_file_hasher(*paths)) + + +def md5_file_hex(*paths: StrPath) -> HexMD5: + return HexMD5(_md5_file_hasher(*paths).hexdigest()) + + +_KB: int = 1_024 +_CHUNKSIZE: int = 128 * _KB +"""Chunk size (in bytes) for iteratively reading from file, if needed.""" + + +def _md5_file_hasher(*paths: StrPath) -> _hashlib.HASH: + md5_hash = _md5() + + # Note: We use str paths (instead of pathlib.Path objs) for minor perf improvements. + for path in sorted(map(str, paths)): + with open(path, "rb") as f: + try: + with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_READ) as mview: + md5_hash.update(mview) + except OSError: + # This occurs if the mmap-ed file is on a different/mounted filesystem, + # so we'll fall back on a less performant implementation. + + # Note: At the time of implementation, the walrus operator `:=` + # is avoided to maintain support for users on python 3.7. + # Consider revisiting once 3.7 support is no longer needed. + chunk = f.read(_CHUNKSIZE) + while chunk: + md5_hash.update(chunk) + chunk = f.read(_CHUNKSIZE) + except ValueError: + # This occurs when mmap-ing an empty file, which can be skipped. + # See: https://github.com/python/cpython/blob/986a4e1b6fcae7fe7a1d0a26aea446107dd58dd2/Modules/mmapmodule.c#L1589 + pass + + return md5_hash diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/import_hooks.py b/lib/python3.12/site-packages/wandb/sdk/lib/import_hooks.py new file mode 100644 index 0000000000000000000000000000000000000000..0ef7bb3f4384797fb491650acd0446ef0d1cb35b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/import_hooks.py @@ -0,0 +1,275 @@ +"""Implements a post-import hook mechanism. + +Styled as per PEP-369. Note that it doesn't cope with modules being reloaded. + +Note: This file is based on +https://github.com/GrahamDumpleton/wrapt/blob/1.12.1/src/wrapt/importer.py +and manual backports of later patches up to 1.15.0 in the wrapt repository +(with slight modifications). +""" + +import sys +import threading +from importlib.util import find_spec +from typing import Any, Callable, Dict, Optional, Union + +# The dictionary registering any post import hooks to be triggered once +# the target module has been imported. Once a module has been imported +# and the hooks fired, the list of hooks recorded against the target +# module will be truncated but the list left in the dictionary. This +# acts as a flag to indicate that the module had already been imported. + +_post_import_hooks: Dict = {} +_post_import_hooks_init: bool = False +_post_import_hooks_lock = threading.RLock() + +# Register a new post import hook for the target module name. This +# differs from the PEP-369 implementation in that it also allows the +# hook function to be specified as a string consisting of the name of +# the callback in the form 'module:function'. This will result in a +# proxy callback being registered which will defer loading of the +# specified module containing the callback function until required. + + +def _create_import_hook_from_string(name: str) -> Callable: + def import_hook(module: Any) -> Callable: + module_name, function = name.split(":") + attrs = function.split(".") + __import__(module_name) + callback = sys.modules[module_name] + for attr in attrs: + callback = getattr(callback, attr) + return callback(module) # type: ignore + + return import_hook + + +def register_post_import_hook( + hook: Union[str, Callable], hook_id: str, name: str +) -> None: + # Create a deferred import hook if hook is a string name rather than + # a callable function. + + if isinstance(hook, (str,)): + hook = _create_import_hook_from_string(hook) + + # Automatically install the import hook finder if it has not already + # been installed. + + with _post_import_hooks_lock: + global _post_import_hooks_init + + if not _post_import_hooks_init: + _post_import_hooks_init = True + sys.meta_path.insert(0, ImportHookFinder()) # type: ignore + + # Check if the module is already imported. If not, register the hook + # to be called after import. + + module = sys.modules.get(name, None) + + if module is None: + _post_import_hooks.setdefault(name, {}).update({hook_id: hook}) + + # If the module is already imported, we fire the hook right away. Note that + # the hook is called outside of the lock to avoid deadlocks if code run as a + # consequence of calling the module import hook in turn triggers a separate + # thread which tries to register an import hook. + + if module is not None: + hook(module) + + +def unregister_post_import_hook(name: str, hook_id: Optional[str]) -> None: + # Remove the import hook if it has been registered. + with _post_import_hooks_lock: + hooks = _post_import_hooks.get(name) + + if hooks is not None: + if hook_id is not None: + hooks.pop(hook_id, None) + + if not hooks: + del _post_import_hooks[name] + else: + del _post_import_hooks[name] + + +def unregister_all_post_import_hooks() -> None: + with _post_import_hooks_lock: + _post_import_hooks.clear() + + +# Indicate that a module has been loaded. Any post import hooks which +# were registered against the target module will be invoked. If an +# exception is raised in any of the post import hooks, that will cause +# the import of the target module to fail. + + +def notify_module_loaded(module: Any) -> None: + name = getattr(module, "__name__", None) + + with _post_import_hooks_lock: + hooks = _post_import_hooks.pop(name, {}) + + # Note that the hook is called outside of the lock to avoid deadlocks if + # code run as a consequence of calling the module import hook in turn + # triggers a separate thread which tries to register an import hook. + for hook in hooks.values(): + if hook: + hook(module) + + +# A custom module import finder. This intercepts attempts to import +# modules and watches out for attempts to import target modules of +# interest. When a module of interest is imported, then any post import +# hooks which are registered will be invoked. + + +class _ImportHookChainedLoader: + def __init__(self, loader: Any) -> None: + self.loader = loader + + if hasattr(loader, "load_module"): + self.load_module = self._load_module + if hasattr(loader, "create_module"): + self.create_module = self._create_module + if hasattr(loader, "exec_module"): + self.exec_module = self._exec_module + + def _set_loader(self, module: Any) -> None: + # Set module's loader to self.loader unless it's already set to + # something else. Import machinery will set it to spec.loader if it is + # None, so handle None as well. The module may not support attribute + # assignment, in which case we simply skip it. Note that we also deal + # with __loader__ not existing at all. This is to future proof things + # due to proposal to remove the attribute as described in the GitHub + # issue at https://github.com/python/cpython/issues/77458. Also prior + # to Python 3.3, the __loader__ attribute was only set if a custom + # module loader was used. It isn't clear whether the attribute still + # existed in that case or was set to None. + + class UNDEFINED: + pass + + if getattr(module, "__loader__", UNDEFINED) in (None, self): + try: + module.__loader__ = self.loader + except AttributeError: + pass + + if ( + getattr(module, "__spec__", None) is not None + and getattr(module.__spec__, "loader", None) is self + ): + module.__spec__.loader = self.loader + + def _load_module(self, fullname: str) -> Any: + module = self.loader.load_module(fullname) + self._set_loader(module) + notify_module_loaded(module) + + return module + + # Python 3.4 introduced create_module() and exec_module() instead of + # load_module() alone. Splitting the two steps. + + def _create_module(self, spec: Any) -> Any: + return self.loader.create_module(spec) + + def _exec_module(self, module: Any) -> None: + self._set_loader(module) + self.loader.exec_module(module) + notify_module_loaded(module) + + +class ImportHookFinder: + def __init__(self) -> None: + self.in_progress: Dict = {} + + def find_module( # type: ignore + self, + fullname: str, + path: Optional[str] = None, + ) -> Optional["_ImportHookChainedLoader"]: + # If the module being imported is not one we have registered + # post import hooks for, we can return immediately. We will + # take no further part in the importing of this module. + + with _post_import_hooks_lock: + if fullname not in _post_import_hooks: + return None + + # When we are interested in a specific module, we will call back + # into the import system a second time to defer to the import + # finder that is supposed to handle the importing of the module. + # We set an in progress flag for the target module so that on + # the second time through we don't trigger another call back + # into the import system and cause a infinite loop. + + if fullname in self.in_progress: + return None + + self.in_progress[fullname] = True + + # Now call back into the import system again. + + try: + # For Python 3 we need to use find_spec().loader + # from the importlib.util module. It doesn't actually + # import the target module and only finds the + # loader. If a loader is found, we need to return + # our own loader which will then in turn call the + # real loader to import the module and invoke the + # post import hooks. + loader = getattr(find_spec(fullname), "loader", None) + + if loader and not isinstance(loader, _ImportHookChainedLoader): + return _ImportHookChainedLoader(loader) + + finally: + del self.in_progress[fullname] + + def find_spec( + self, fullname: str, path: Optional[str] = None, target: Any = None + ) -> Any: + # Since Python 3.4, you are meant to implement find_spec() method + # instead of find_module() and since Python 3.10 you get deprecation + # warnings if you don't define find_spec(). + + # If the module being imported is not one we have registered + # post import hooks for, we can return immediately. We will + # take no further part in the importing of this module. + + with _post_import_hooks_lock: + if fullname not in _post_import_hooks: + return None + + # When we are interested in a specific module, we will call back + # into the import system a second time to defer to the import + # finder that is supposed to handle the importing of the module. + # We set an in progress flag for the target module so that on + # the second time through we don't trigger another call back + # into the import system and cause a infinite loop. + + if fullname in self.in_progress: + return None + + self.in_progress[fullname] = True + + # Now call back into the import system again. + + try: + # This should only be Python 3 so find_spec() should always + # exist so don't need to check. + spec = find_spec(fullname) + loader = getattr(spec, "loader", None) + + if loader and not isinstance(loader, _ImportHookChainedLoader): + assert spec is not None + spec.loader = _ImportHookChainedLoader(loader) # type: ignore + + return spec + + finally: + del self.in_progress[fullname] diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/interrupt.py b/lib/python3.12/site-packages/wandb/sdk/lib/interrupt.py new file mode 100644 index 0000000000000000000000000000000000000000..19cc0fe914674d1a328b278024c03efef8b61ad9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/interrupt.py @@ -0,0 +1,37 @@ +"""Utility to send an interrupt (Ctrl+C) signal to the main thread. + +This is necessary because Windows and POSIX use different models for Ctrl+C +interrupts. +""" + +import platform +import signal +import threading + + +def interrupt_main(): + """Interrupt the main Python thread with a SIGINT signal. + + In POSIX, signal.pthread_kill() is the most reliable way to send a signal + to the main thread. + + os.kill() is often recommended, but it isn't guaranteed to deliver the + signal to the main OS thread. Likewise, signal.raise_signal() delivers + the signal to the current thread in POSIX. The issue is that if any other + thread receives the signal, Python will set an internal flag and process it + on the main thread at the next opportunity. If the main thread is executing + C code or is blocked on a syscall (e.g. time.sleep(999999)) the signal + handler won't execute until that's done---i.e. Python won't preempt the OS + thread on its own. + + On Windows, pthread_kill is not available and os.kill() ignores its + second argument and always kills the process. However, + signal.raise_signal() does the right thing. + """ + if platform.system() == "Windows": + signal.raise_signal(signal.SIGINT) + else: + signal.pthread_kill( + threading.main_thread().ident, + signal.SIGINT, + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/ipython.py b/lib/python3.12/site-packages/wandb/sdk/lib/ipython.py new file mode 100644 index 0000000000000000000000000000000000000000..74022e99e318a7b9f85cc933fcaca59b26207ae4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/ipython.py @@ -0,0 +1,126 @@ +import logging +import sys +import warnings +from typing import Literal, Optional + +import wandb + +PythonType = Literal["python", "ipython", "jupyter"] + +logger = logging.getLogger(__name__) + + +def toggle_button(what="run"): + """Returns the HTML for a button used to reveal the element following it. + + The element immediately after the button must have `display: none`. + """ + return ( + "" + ) + + +def _get_python_type() -> PythonType: + if "IPython" not in sys.modules: + return "python" + + try: + from IPython import get_ipython # type: ignore + + # Calling get_ipython can cause an ImportError + if get_ipython() is None: + return "python" + except ImportError: + return "python" + + # jupyter-based environments (e.g. jupyter itself, colab, kaggle, etc) have a connection file + ip_kernel_app_connection_file = ( + (get_ipython().config.get("IPKernelApp", {}) or {}) + .get("connection_file", "") + .lower() + ) or ( + (get_ipython().config.get("ColabKernelApp", {}) or {}) + .get("connection_file", "") + .lower() + ) + + if ( + ("terminal" in get_ipython().__module__) + or ("jupyter" not in ip_kernel_app_connection_file) + or ("spyder" in sys.modules) + ): + return "ipython" + else: + return "jupyter" + + +def in_jupyter() -> bool: + """Returns True if we're in a Jupyter notebook.""" + return _get_python_type() == "jupyter" + + +def in_ipython() -> bool: + """Returns True if we're running in IPython in the terminal.""" + return _get_python_type() == "ipython" + + +def in_notebook() -> bool: + """Returns True if we're running in Jupyter or IPython.""" + return _get_python_type() != "python" + + +class ProgressWidget: + """A simple wrapper to render a nice progress bar with a label.""" + + def __init__(self, widgets, min, max): + from IPython import display + + self._ipython_display = display + + self.widgets = widgets + self._progress = widgets.FloatProgress(min=min, max=max) + self._label = widgets.Label() + self._widget = self.widgets.VBox([self._label, self._progress]) + self._displayed = False + self._disabled = False + + def update(self, value: float, label: str) -> None: + if self._disabled: + return + try: + self._progress.value = value + self._label.value = label + if not self._displayed: + self._displayed = True + self._ipython_display.display(self._widget) + except Exception: + logger.exception("Error in ProgressWidget.update()") + self._disabled = True + wandb.termwarn( + "Unable to render progress bar, see the user log for details" + ) + + def close(self) -> None: + if self._disabled or not self._displayed: + return + self._widget.close() + + +def jupyter_progress_bar(min: float = 0, max: float = 1.0) -> Optional[ProgressWidget]: + """Return an ipywidget progress bar or None if we can't import it.""" + widgets = wandb.util.get_module("ipywidgets") + try: + if widgets is None: + # TODO: this currently works in iPython but it's deprecated since 4.0 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + from IPython.html import widgets # type: ignore + + assert hasattr(widgets, "VBox") + assert hasattr(widgets, "Label") + assert hasattr(widgets, "FloatProgress") + return ProgressWidget(widgets, min=min, max=max) + except (ImportError, AssertionError): + return None diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/json_util.py b/lib/python3.12/site-packages/wandb/sdk/lib/json_util.py new file mode 100644 index 0000000000000000000000000000000000000000..1892e21219bfe062dbd456d05fbd9e8608f9482f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/json_util.py @@ -0,0 +1,75 @@ +import json +import logging +import os +from typing import Any, Union + +logger = logging.getLogger(__name__) + + +try: + import orjson # type: ignore + + # todo: orjson complies with the json standard and does not support + # NaN, Infinity, and -Infinity. Should be fixed in the future. + + # additional safeguard for now + if os.environ.get("_WANDB_ORJSON"): + + def dumps(obj: Any, **kwargs: Any) -> str: + """Wrapper for .dumps.""" + cls = kwargs.pop("cls", None) + try: + _kwargs = kwargs.copy() + if cls: + _kwargs["default"] = cls.default + encoded = orjson.dumps( + obj, option=orjson.OPT_NON_STR_KEYS, **_kwargs + ).decode() + except Exception: + logger.exception("Error using orjson.dumps") + if cls: + kwargs["cls"] = cls + encoded = json.dumps(obj, **kwargs) + + return encoded # type: ignore[no-any-return] + + def dump(obj: Any, fp: Any, **kwargs: Any) -> None: + """Wrapper for .dump.""" + cls = kwargs.pop("cls", None) + try: + _kwargs = kwargs.copy() + if cls: + _kwargs["default"] = cls.default + encoded = orjson.dumps(obj, option=orjson.OPT_NON_STR_KEYS, **_kwargs) + fp.write(encoded) + except Exception: + logger.exception("Error using orjson.dump") + if cls: + kwargs["cls"] = cls + json.dump(obj, fp, **kwargs) + + def loads(obj: Union[str, bytes]) -> Any: + """Wrapper for orjson.loads.""" + try: + decoded = orjson.loads(obj) + except Exception: + logger.exception("Error using orjson.loads") + decoded = json.loads(obj) + + return decoded + + def load(fp: Any) -> Any: + """Wrapper for orjson.load.""" + try: + decoded = orjson.loads(fp.read()) + except Exception: + logger.exception("Error using orjson.load") + decoded = json.load(fp) + + return decoded + + else: + from json import dump, dumps, load, loads # type: ignore[assignment] + +except ImportError: + from json import dump, dumps, load, loads # type: ignore[assignment] # noqa: F401 diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/lazyloader.py b/lib/python3.12/site-packages/wandb/sdk/lib/lazyloader.py new file mode 100644 index 0000000000000000000000000000000000000000..0c7aec3f94ac6bab37154c3fb796ac0ac63ba9a4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/lazyloader.py @@ -0,0 +1,63 @@ +"""module lazyloader.""" + +import importlib +import sys +import types + + +class LazyLoader(types.ModuleType): + """Lazily import a module, mainly to avoid pulling in large dependencies. + + We use this for tensorflow and other optional libraries primarily at the + top module level. + """ + + # The lint error here is incorrect. + def __init__( + self, + local_name, # pylint: disable=super-on-old-class + parent_module_globals, + name, + warning=None, + ): + self._local_name = local_name + self._parent_module_globals = parent_module_globals + self._warning = warning + + super().__init__(str(name)) + + def _load(self): + """Load the module and insert it into the parent's globals.""" + # Import the target module and insert it into the parent's namespace + module = importlib.import_module(self.__name__) + self._parent_module_globals[self._local_name] = module + # print("import", self.__name__) + # print("Set global", self._local_name) + # print("mod", module) + sys.modules[self._local_name] = module + + # Emit a warning if one was specified + if self._warning: + print(self._warning) # noqa: T201 + # Make sure to only warn once. + self._warning = None + + # Update this object's dict so that if someone keeps a reference to the + # LazyLoader, lookups are efficient (__getattr__ is only called on lookups + # that fail). + self.__dict__.update(module.__dict__) + + return module + + # def __getattribute__(self, item): + # print("getattribute", item) + + def __getattr__(self, item): + # print("getattr", item) + module = self._load() + return getattr(module, item) + + def __dir__(self): + # print("dir") + module = self._load() + return dir(module) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/module.py b/lib/python3.12/site-packages/wandb/sdk/lib/module.py new file mode 100644 index 0000000000000000000000000000000000000000..052f39380109b86b70a6a6980ec3409bf4e7ed87 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/module.py @@ -0,0 +1,78 @@ +# +import wandb + +from . import preinit + + +def set_global( + run=None, + config=None, + log=None, + summary=None, + save=None, + use_artifact=None, + log_artifact=None, + define_metric=None, + alert=None, + mark_preempting=None, + log_model=None, + use_model=None, + link_model=None, + watch=None, + unwatch=None, +): + if run: + wandb.run = run + if config is not None: + wandb.config = config + if log: + wandb.log = log + if summary is not None: + wandb.summary = summary + if save: + wandb.save = save + if use_artifact: + wandb.use_artifact = use_artifact + if log_artifact: + wandb.log_artifact = log_artifact + if define_metric: + wandb.define_metric = define_metric + if alert: + wandb.alert = alert + if mark_preempting: + wandb.mark_preempting = mark_preempting + if log_model: + wandb.log_model = log_model + if use_model: + wandb.use_model = use_model + if link_model: + wandb.link_model = link_model + if watch: + wandb.watch = watch + if unwatch: + wandb.unwatch = unwatch + + +def unset_globals(): + wandb.run = None + wandb.config = preinit.PreInitObject("wandb.config") + wandb.summary = preinit.PreInitObject("wandb.summary") + wandb.log = preinit.PreInitCallable("wandb.log", wandb.wandb_sdk.wandb_run.Run.log) + wandb.watch = preinit.PreInitCallable( + "wandb.watch", wandb.wandb_sdk.wandb_run.Run.watch + ) + wandb.unwatch = preinit.PreInitCallable( + "wandb.unwatch", wandb.wandb_sdk.wandb_run.Run.unwatch + ) + wandb.save = preinit.PreInitCallable( + "wandb.save", wandb.wandb_sdk.wandb_run.Run.save + ) + wandb.use_artifact = preinit.PreInitCallable( + "wandb.use_artifact", wandb.wandb_sdk.wandb_run.Run.use_artifact + ) + wandb.log_artifact = preinit.PreInitCallable( + "wandb.log_artifact", wandb.wandb_sdk.wandb_run.Run.log_artifact + ) + wandb.define_metric = preinit.PreInitCallable( + "wandb.define_metric", wandb.wandb_sdk.wandb_run.Run.define_metric + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/paths.py b/lib/python3.12/site-packages/wandb/sdk/lib/paths.py new file mode 100644 index 0000000000000000000000000000000000000000..1fae7110ec5ffe83d46c18c4263b0745a44264c9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/paths.py @@ -0,0 +1,106 @@ +import os +import platform +from functools import wraps +from pathlib import PurePath, PurePosixPath +from typing import Any, NewType, Union + +# Path _inputs_ should generally accept any kind of path. This is named the same and +# modeled after the hint defined in the Python standard library's `typeshed`: +# https://github.com/python/typeshed/blob/0b1cd5989669544866213807afa833a88f649ee7/stdlib/_typeshed/__init__.pyi#L56-L65 +StrPath = Union[str, "os.PathLike[str]"] + +# A native path to a file on a local filesystem. +FilePathStr = NewType("FilePathStr", str) + +URIStr = NewType("URIStr", str) + + +class LogicalPath(str): + """A string that represents a path relative to an artifact or run. + + The format of the string is always as a POSIX path, e.g. "foo/bar.txt". + + A neat trick is that you can use this class as if it were a PurePosixPath. E.g.: + ``` + >>> path = LogicalPath("foo/bar.txt") + >>> path.parts + ('foo', 'bar.txt') + >>> path.parent / "baz.txt" + 'foo/baz.txt' + >>> type(path.relative_to("foo")) + LogicalPath + ``` + """ + + # It should probably always be a relative path, but that would be a behavior change. + # + # These strings used to be the output of `to_forward_slash_path`, which only works + # with strings and whose behavior is pretty simple: + # ``` + # if platform.system() == "Windows": + # path = path.replace("\\", "/") + # ``` + # + # This results in some weird things, such as backslashes being allowed from + # non-Windows platforms (which would probably break if such an artifact was used + # from Windows) and anchors or absolute paths being allowed. E.g., the Windows path + # "C:\foo\bar.txt" becomes "C:/foo/bar.txt", which then would mount as + # "./artifacts/artifact_name:v0/C:/foo/bar.txt" on MacOS and as + # "./artifacts/artifact_name-v0/C-/foo/bar.txt" on Windows. + # + # This implementation preserves behavior for strings but attempts to sanitize other + # formerly unsupported inputs more aggressively. It uses the `.as_posix()` form of + # pathlib objects rather than the `str()` form to reduce how often identical inputs + # will result in different outputs on different platforms; however, it doesn't alter + # absolute paths or check for prohibited characters etc. + + def __new__(cls, path: StrPath) -> "LogicalPath": + if isinstance(path, LogicalPath): + return super().__new__(cls, path) + if hasattr(path, "as_posix"): + path = PurePosixPath(path.as_posix()) + return super().__new__(cls, str(path)) + if hasattr(path, "__fspath__"): + path = path.__fspath__() # Can be str or bytes. + if isinstance(path, bytes): + path = os.fsdecode(path) + # For historical reasons we have to convert backslashes to forward slashes, but + # only on Windows, and need to do it before any pathlib operations. + if platform.system() == "Windows": + path = path.replace("\\", "/") + # This weird contortion and the one above are because in some unusual cases + # PurePosixPath(path.as_posix()).as_posix() != path.as_posix(). + path = PurePath(path).as_posix() + return super().__new__(cls, str(PurePosixPath(path))) + + def to_path(self) -> PurePosixPath: + """Convert this path to a PurePosixPath.""" + return PurePosixPath(self) + + def __getattr__(self, attr: str) -> Any: + """Act like a subclass of PurePosixPath for all methods not defined on str.""" + try: + result = getattr(self.to_path(), attr) + except AttributeError as e: + raise AttributeError(f"LogicalPath has no attribute {attr!r}") from e + + if isinstance(result, PurePosixPath): + return LogicalPath(result) + + # If the result is a callable (a method), wrap it so that it has the same + # behavior: if the call result returns a PurePosixPath, return a LogicalPath. + if callable(result): + + @wraps(result) + def wrapper(*args: Any, **kwargs: Any) -> Any: + inner_result = result(*args, **kwargs) + if isinstance(inner_result, PurePosixPath): + return LogicalPath(inner_result) + return inner_result + + return wrapper + return result + + def __truediv__(self, other: StrPath) -> "LogicalPath": + """Act like a PurePosixPath for the / operator, but return a LogicalPath.""" + return LogicalPath(self.to_path() / LogicalPath(other)) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/preinit.py b/lib/python3.12/site-packages/wandb/sdk/lib/preinit.py new file mode 100644 index 0000000000000000000000000000000000000000..624528198b1cab580f95d29eb91b0531378041e8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/preinit.py @@ -0,0 +1,42 @@ +from typing import Any, Callable, Optional + +import wandb + + +class PreInitObject: + def __init__(self, name: str, destination: Optional[Any] = None) -> None: + self._name = name + + if destination is not None: + self.__doc__ = destination.__doc__ + + def __getitem__(self, key: str) -> None: + raise wandb.Error(f"You must call wandb.init() before {self._name}[{key!r}]") + + def __setitem__(self, key: str, value: Any) -> Any: + raise wandb.Error(f"You must call wandb.init() before {self._name}[{key!r}]") + + def __setattr__(self, key: str, value: Any) -> Any: + if not key.startswith("_"): + raise wandb.Error(f"You must call wandb.init() before {self._name}.{key}") + else: + return object.__setattr__(self, key, value) + + def __getattr__(self, key: str) -> Any: + if not key.startswith("_"): + raise wandb.Error(f"You must call wandb.init() before {self._name}.{key}") + else: + raise AttributeError + + +def PreInitCallable( # noqa: N802 + name: str, destination: Optional[Any] = None +) -> Callable: + def preinit_wrapper(*args: Any, **kwargs: Any) -> Any: + raise wandb.Error(f"You must call wandb.init() before {name}()") + + preinit_wrapper.__name__ = str(name) + if destination: + preinit_wrapper.__wrapped__ = destination # type: ignore + preinit_wrapper.__doc__ = destination.__doc__ + return preinit_wrapper diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/printer.py b/lib/python3.12/site-packages/wandb/sdk/lib/printer.py new file mode 100644 index 0000000000000000000000000000000000000000..c584edf559b175e3484e6007178580f6f6a9e07e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/printer.py @@ -0,0 +1,571 @@ +"""Terminal, Jupyter and file output for W&B.""" + +from __future__ import annotations + +import abc +import contextlib +import itertools +import platform +import sys +from typing import Callable, Iterator + +import click +from typing_extensions import override + +import wandb +import wandb.util +from wandb.errors import term +from wandb.sdk import wandb_setup + +from . import ipython, sparkline + +# Follow the same logic as the python logging module +CRITICAL = 50 +FATAL = CRITICAL +ERROR = 40 +WARNING = 30 +WARN = WARNING +INFO = 20 +DEBUG = 10 +NOTSET = 0 + +_level_to_name = { + CRITICAL: "CRITICAL", + ERROR: "ERROR", + WARNING: "WARNING", + INFO: "INFO", + DEBUG: "DEBUG", + NOTSET: "NOTSET", +} + +_name_to_level = { + "CRITICAL": CRITICAL, + "FATAL": FATAL, + "ERROR": ERROR, + "WARN": WARNING, + "WARNING": WARNING, + "INFO": INFO, + "DEBUG": DEBUG, + "NOTSET": NOTSET, +} + +_PROGRESS_SYMBOL_ANIMATION = "⢿⣻⣽⣾⣷⣯⣟⡿" +"""Sequence of characters for a progress spinner. + +Unicode characters from the Braille Patterns block arranged +to form a subtle clockwise spinning animation. +""" + +_PROGRESS_SYMBOL_COLOR = 0xB2 +"""Color from the 256-color palette for the progress symbol.""" + +_JUPYTER_TABLE_STYLES = """ + +""" + +_JUPYTER_PANEL_STYLES = """ + +""" + + +def new_printer(settings: wandb.Settings | None = None) -> Printer: + """Returns a printer appropriate for the environment we're in. + + Args: + settings: The settings of a run. If not provided and `wandb.setup()` + has been called, then global settings are used. Otherwise, + settings (such as silent mode) are ignored. + """ + if not settings and (singleton := wandb_setup.singleton_if_setup()): + settings = singleton.settings + + if ipython.in_jupyter(): + return _PrinterJupyter(settings=settings) + else: + return _PrinterTerm(settings=settings) + + +class Printer(abc.ABC): + """An object that shows styled text to the user.""" + + @contextlib.contextmanager + @abc.abstractmethod + def dynamic_text(self) -> Iterator[DynamicText | None]: + """A context manager providing a handle to a block of changeable text. + + Since `wandb` may be outputting to a terminal, it's important to only + use this when `wandb` is performing blocking calls, or else text output + by non-`wandb` code may get overwritten. + + Returns None if dynamic text is not supported, such as if stderr is not + a TTY and we're not in a Jupyter notebook. + """ + + @abc.abstractmethod + def display( + self, + text: str | list[str] | tuple[str], + *, + level: str | int | None = None, + ) -> None: + """Display text to the user. + + Args: + text: The text to display. If given an iterable of strings, they're + joined with newlines. + level: The logging level, for controlling verbosity. + """ + + @abc.abstractmethod + def progress_update( + self, + text: str, + percent_done: float | None = None, + ) -> None: + r"""Set the text on the progress indicator. + + Args: + text: The text to set, which must end with \r. + percent_done: The current progress, between 0 and 1. + """ + + @abc.abstractmethod + def progress_close(self, text: str | None = None) -> None: + """Close the progress indicator. + + After this, `progress_update` should not be used. + + Args: + text: The final text to set on the progress indicator. + Ignored in Jupyter notebooks. + """ + + @staticmethod + def _sanitize_level(name_or_level: str | int | None) -> int: + """Returns the number corresponding to the logging level. + + Args: + name_or_level: The logging level passed to `display`. + + Raises: + ValueError: if the input is not a valid logging level. + """ + if isinstance(name_or_level, str): + try: + return _name_to_level[name_or_level.upper()] + except KeyError: + raise ValueError( + f"Unknown level name: {name_or_level}, supported levels: {_name_to_level.keys()}" + ) + + if isinstance(name_or_level, int): + return name_or_level + + if name_or_level is None: + return INFO + + raise ValueError(f"Unknown status level {name_or_level}") + + @property + @abc.abstractmethod + def supports_html(self) -> bool: + """Whether text passed to display may contain HTML styling.""" + + @property + @abc.abstractmethod + def supports_unicode(self) -> bool: + """Whether text passed to display may contain arbitrary Unicode.""" + + def sparklines(self, series: list[int | float]) -> str | None: + """Returns a Unicode art representation of the series of numbers. + + Also known as "ASCII art", except this uses non-ASCII + Unicode characters. + + Returns None if the output doesn't support Unicode. + """ + if self.supports_unicode: + return sparkline.sparkify(series) + else: + return None + + @abc.abstractmethod + def code(self, text: str) -> str: + """Returns the text styled like code.""" + + @abc.abstractmethod + def name(self, text: str) -> str: + """Returns the text styled like a run name.""" + + @abc.abstractmethod + def link(self, link: str, text: str | None = None) -> str: + """Returns the text styled like a link. + + Args: + link: The target link. + text: The text to show for the link. If not set, or if we're not + in an environment that supports clickable links, + this is ignored. + """ + + @abc.abstractmethod + def secondary_text(self, text: str) -> str: + """Returns the text styled to draw less attention.""" + + @abc.abstractmethod + def loading_symbol(self, tick: int) -> str: + """Returns a frame of an animated loading symbol. + + May return an empty string. + + Args: + tick: An index into the animation. + """ + + @abc.abstractmethod + def error(self, text: str) -> str: + """Returns the text colored like an error.""" + + @abc.abstractmethod + def emoji(self, name: str) -> str: + """Returns the string for a named emoji, or an empty string.""" + + @abc.abstractmethod + def files(self, text: str) -> str: + """Returns the text styled like a file path.""" + + @abc.abstractmethod + def grid(self, rows: list[list[str]], title: str | None = None) -> str: + """Returns a grid of strings with an optional title.""" + + @abc.abstractmethod + def panel(self, columns: list[str]) -> str: + """Returns the column text combined in a compact way.""" + + +class DynamicText(abc.ABC): + """A handle to a block of text that's allowed to change.""" + + @abc.abstractmethod + def set_text(self, text: str) -> None: + r"""Change the text. + + Args: + text: The text to put in the block, with lines separated + by \n characters. The text should not end in \n unless + a blank line at the end of the block is desired. + May include styled output from methods on the Printer + that created this. + """ + + +class _PrinterTerm(Printer): + def __init__(self, *, settings: wandb.Settings | None) -> None: + super().__init__() + self._settings = settings + self._progress = itertools.cycle(["-", "\\", "|", "/"]) + + @override + @contextlib.contextmanager + def dynamic_text(self) -> Iterator[DynamicText | None]: + if self._settings and self._settings.silent: + yield None + return + + with term.dynamic_text() as handle: + if not handle: + yield None + else: + yield _DynamicTermText(handle) + + @override + def display( + self, + text: str | list[str] | tuple[str], + *, + level: str | int | None = None, + ) -> None: + if self._settings and self._settings.silent: + return + + text = "\n".join(text) if isinstance(text, (list, tuple)) else text + self._display_fn_mapping(level)(text) + + @staticmethod + def _display_fn_mapping(level: str | int | None = None) -> Callable[[str], None]: + level = Printer._sanitize_level(level) + + if level >= CRITICAL: + return wandb.termerror + elif ERROR <= level < CRITICAL: + return wandb.termerror + elif WARNING <= level < ERROR: + return wandb.termwarn + elif INFO <= level < WARNING: + return wandb.termlog + elif DEBUG <= level < INFO: + return wandb.termlog + else: + return wandb.termlog + + @override + def progress_update(self, text: str, percent_done: float | None = None) -> None: + if self._settings and self._settings.silent: + return + + wandb.termlog(f"{next(self._progress)} {text}", newline=False) + + @override + def progress_close(self, text: str | None = None) -> None: + if self._settings and self._settings.silent: + return + + text = text or " " * 79 + wandb.termlog(text) + + @property + @override + def supports_html(self) -> bool: + return False + + @property + @override + def supports_unicode(self) -> bool: + return wandb.util.is_unicode_safe(sys.stderr) + + @override + def code(self, text: str) -> str: + ret: str = click.style(text, bold=True) + return ret + + @override + def name(self, text: str) -> str: + ret: str = click.style(text, fg="yellow") + return ret + + @override + def link(self, link: str, text: str | None = None) -> str: + ret: str = click.style(link, fg="blue", underline=True) + # ret = f"\x1b[m{text or link}\x1b[0m" + # ret = f"\x1b]8;;{link}\x1b\\{ret}\x1b]8;;\x1b\\" + return ret + + @override + def emoji(self, name: str) -> str: + emojis = dict() + if platform.system() != "Windows" and wandb.util.is_unicode_safe(sys.stdout): + emojis = dict( + star="⭐️", + broom="🧹", + rocket="🚀", + gorilla="🦍", + turtle="🐢", + lightning="️⚡", + ) + + return emojis.get(name, "") + + @override + def secondary_text(self, text: str) -> str: + # NOTE: "white" is really a light gray, and is usually distinct + # from the terminal's foreground color (i.e. default text color) + return click.style(text, fg="white") + + @override + def loading_symbol(self, tick: int) -> str: + if not self.supports_unicode: + return "" + + idx = tick % len(_PROGRESS_SYMBOL_ANIMATION) + return click.style( + _PROGRESS_SYMBOL_ANIMATION[idx], + fg=_PROGRESS_SYMBOL_COLOR, + ) + + @override + def error(self, text: str) -> str: + return click.style(text, fg="red") + + @override + def files(self, text: str) -> str: + ret: str = click.style(text, fg="magenta", bold=True) + return ret + + @override + def grid(self, rows: list[list[str]], title: str | None = None) -> str: + max_len = max(len(row[0]) for row in rows) + format_row = " ".join(["{:>{max_len}}", "{}" * (len(rows[0]) - 1)]) + grid = "\n".join([format_row.format(*row, max_len=max_len) for row in rows]) + if title: + return f"{title}\n{grid}\n" + return f"{grid}\n" + + @override + def panel(self, columns: list[str]) -> str: + return "\n" + "\n".join(columns) + + +class _DynamicTermText(DynamicText): + def __init__(self, handle: term.DynamicBlock) -> None: + self._handle = handle + + @override + def set_text(self, text: str) -> None: + self._handle.set_text(text) + + +class _PrinterJupyter(Printer): + def __init__(self, *, settings: wandb.Settings | None) -> None: + super().__init__() + self._settings = settings + self._progress = ipython.jupyter_progress_bar() + + from IPython import display + + self._ipython_display = display + + @override + @contextlib.contextmanager + def dynamic_text(self) -> Iterator[DynamicText | None]: + if self._settings and self._settings.silent: + yield None + return + + handle = self._ipython_display.display( + self._ipython_display.HTML(""), + display_id=True, + ) + + if handle: + yield _DynamicJupyterText(handle) + handle.update(self._ipython_display.HTML("")) + else: + yield None + + @override + def display( + self, + text: str | list[str] | tuple[str], + *, + level: str | int | None = None, + ) -> None: + if self._settings and self._settings.silent: + return + + text = "
".join(text) if isinstance(text, (list, tuple)) else text + text = "
".join(text.splitlines()) + self._ipython_display.display(self._ipython_display.HTML(text)) + + @property + @override + def supports_html(self) -> bool: + return True + + @property + @override + def supports_unicode(self) -> bool: + return True + + @override + def code(self, text: str) -> str: + return f"{text}" + + @override + def name(self, text: str) -> str: + return f'{text}' + + @override + def link(self, link: str, text: str | None = None) -> str: + return f'{text or link}' + + @override + def emoji(self, name: str) -> str: + return "" + + @override + def secondary_text(self, text: str) -> str: + return text + + @override + def loading_symbol(self, tick: int) -> str: + return "" + + @override + def error(self, text: str) -> str: + return f'{text}' + + @override + def files(self, text: str) -> str: + return f"{text}" + + @override + def progress_update( + self, + text: str, + percent_done: float | None = None, + ) -> None: + if (self._settings and self._settings.silent) or not self._progress: + return + + if percent_done is None: + percent_done = 1.0 + + self._progress.update(percent_done, text) + + @override + def progress_close(self, text: str | None = None) -> None: + if self._progress: + self._progress.close() + + @override + def grid(self, rows: list[list[str]], title: str | None = None) -> str: + format_row = "".join(["", "" * len(rows[0]), ""]) + grid = "".join([format_row.format(*row) for row in rows]) + grid = f'
{}
{grid}
' + if title: + return f"

{title}


{grid}
" + return f"{_JUPYTER_TABLE_STYLES}{grid}
" + + @override + def panel(self, columns: list[str]) -> str: + row = "".join([f'
{col}
' for col in columns]) + return f'{_JUPYTER_PANEL_STYLES}
{row}
' + + +class _DynamicJupyterText(DynamicText): + def __init__(self, handle) -> None: + from IPython import display + + self._ipython_to_html = display.HTML + self._handle: display.DisplayHandle = handle + + @override + def set_text(self, text: str) -> None: + text = "
".join(text.splitlines()) + self._handle.update(self._ipython_to_html(text)) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/printer_asyncio.py b/lib/python3.12/site-packages/wandb/sdk/lib/printer_asyncio.py new file mode 100644 index 0000000000000000000000000000000000000000..da311a16868ef02bc0e72586779f4c153c8bc352 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/printer_asyncio.py @@ -0,0 +1,46 @@ +import asyncio +from typing import Callable, TypeVar + +from wandb.sdk.lib import asyncio_compat, printer + +_T = TypeVar("_T") + + +def run_async_with_spinner( + spinner_printer: printer.Printer, + text: str, + func: Callable[[], _T], +) -> _T: + """Run a slow function while displaying a loading icon. + + Args: + spinner_printer: The printer to use to display text. + text: The text to display next to the spinner while the function runs. + func: The function to run. + + Returns: + The result of func. + """ + + async def _loop_run_with_spinner() -> _T: + func_running = asyncio.Event() + + async def update_spinner() -> None: + tick = 0 + with spinner_printer.dynamic_text() as text_area: + if text_area: + while not func_running.is_set(): + spinner = spinner_printer.loading_symbol(tick) + text_area.set_text(f"{spinner} {text}") + tick += 1 + await asyncio.sleep(0.1) + else: + spinner_printer.display(text) + + async with asyncio_compat.open_task_group() as group: + group.start_soon(update_spinner()) + res = await asyncio.get_running_loop().run_in_executor(None, func) + func_running.set() + return res + + return asyncio_compat.run(_loop_run_with_spinner) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/progress.py b/lib/python3.12/site-packages/wandb/sdk/lib/progress.py new file mode 100644 index 0000000000000000000000000000000000000000..d86272e6cc2d69f73f10480a75e0292f024b5d05 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/progress.py @@ -0,0 +1,339 @@ +"""Defines an object for printing run progress at the end of a script.""" + +from __future__ import annotations + +import asyncio +import contextlib +import time +from typing import Iterable, Iterator, NoReturn + +from wandb.proto import wandb_internal_pb2 as pb +from wandb.sdk.interface import interface +from wandb.sdk.lib import asyncio_compat + +from . import printer as p + + +def print_sync_dedupe_stats( + printer: p.Printer, + final_result: pb.PollExitResponse, +) -> None: + """Print how much W&B sync reduced the amount of uploaded data. + + Args: + final_result: The final PollExit result. + """ + deduped_bytes = final_result.pusher_stats.deduped_bytes + total_bytes = final_result.pusher_stats.total_bytes + + if total_bytes <= 0 or deduped_bytes <= 0: + return + + frac = deduped_bytes / total_bytes + printer.display(f"W&B sync reduced upload amount by {frac:.1%}") + + +async def loop_printing_operation_stats( + progress: ProgressPrinter, + interface: interface.InterfaceBase, +) -> None: + """Poll and display ongoing tasks in the internal service process. + + This never returns and must be cancelled. This is meant to be used with + `mailbox.wait_with_progress()`. + + Args: + progress: The printer to update with operation stats. + interface: The interface to use to poll for updates. + + Raises: + HandleAbandonedError: If the mailbox associated with the interface + becomes closed. + Exception: Any other problem communicating with the service process. + """ + stats: pb.OperationStats | None = None + + async def loop_update_screen() -> NoReturn: + while True: + if stats: + progress.update(stats) + await asyncio.sleep(0.1) + + async def loop_poll_stats() -> NoReturn: + nonlocal stats + while True: + start_time = time.monotonic() + + handle = interface.deliver_operation_stats() + result = await handle.wait_async(timeout=None) + stats = result.response.operations_response.operation_stats + + elapsed_time = time.monotonic() - start_time + if elapsed_time < 0.5: + await asyncio.sleep(0.5 - elapsed_time) + + async with asyncio_compat.open_task_group() as task_group: + task_group.start_soon(loop_update_screen()) + task_group.start_soon(loop_poll_stats()) + + +@contextlib.contextmanager +def progress_printer( + printer: p.Printer, + default_text: str, +) -> Iterator[ProgressPrinter]: + """Context manager providing an object for printing run progress. + + Args: + printer: The printer to use. + default_text: The text to show if no information is available. + """ + with printer.dynamic_text() as text_area: + yield ProgressPrinter( + printer, + text_area, + default_text=default_text, + ) + printer.progress_close() + + +class ProgressPrinter: + """Displays PollExitResponse results to the user.""" + + def __init__( + self, + printer: p.Printer, + progress_text_area: p.DynamicText | None, + default_text: str, + ) -> None: + self._show_operation_stats = True + self._printer = printer + self._progress_text_area = progress_text_area + self._default_text = default_text + self._tick = 0 + self._last_printed_line = "" + + def update( + self, + progress: list[pb.PollExitResponse] | pb.OperationStats, + ) -> None: + """Update the displayed information.""" + if not progress: + return + + if isinstance(progress, pb.OperationStats): + self._update_operation_stats([progress]) + elif self._show_operation_stats: + self._update_operation_stats( + list(response.operation_stats for response in progress) + ) + elif len(progress) == 1: + self._update_single_run(progress[0]) + else: + self._update_multiple_runs(progress) + + self._tick += 1 + + def _update_operation_stats(self, stats_list: list[pb.OperationStats]) -> None: + if self._progress_text_area: + _DynamicOperationStatsPrinter( + self._printer, + self._progress_text_area, + max_lines=6, + loading_symbol=self._printer.loading_symbol(self._tick), + default_text=self._default_text, + ).display(stats_list) + + else: + top_level_operations: list[str] = [] + extra_operations = 0 + for stats in stats_list: + for op in stats.operations: + if len(top_level_operations) < 5: + top_level_operations.append(op.desc) + else: + extra_operations += 1 + + line = "; ".join(top_level_operations) + if extra_operations > 0: + line += f" (+ {extra_operations} more)" + + if line != self._last_printed_line: + self._printer.display(line) + + self._last_printed_line = line + + def _update_single_run( + self, + progress: pb.PollExitResponse, + ) -> None: + stats = progress.pusher_stats + line = ( + f"{_megabytes(stats.uploaded_bytes):.3f} MB" + f" of {_megabytes(stats.total_bytes):.3f} MB uploaded" + ) + + if stats.deduped_bytes > 0: + line += f" ({_megabytes(stats.deduped_bytes):.3f} MB deduped)" + + if stats.total_bytes > 0: + self._update_progress_text( + line, + stats.uploaded_bytes / stats.total_bytes, + ) + else: + self._update_progress_text(line, 1.0) + + def _update_multiple_runs( + self, + progress_list: list[pb.PollExitResponse], + ) -> None: + total_files = 0 + uploaded_bytes = 0 + total_bytes = 0 + + for progress in progress_list: + total_files += progress.file_counts.wandb_count + total_files += progress.file_counts.media_count + total_files += progress.file_counts.artifact_count + total_files += progress.file_counts.other_count + + uploaded_bytes += progress.pusher_stats.uploaded_bytes + total_bytes += progress.pusher_stats.total_bytes + + line = ( + f"Processing {len(progress_list)} runs with {total_files} files" + f" ({_megabytes(uploaded_bytes):.2f} MB" + f" / {_megabytes(total_bytes):.2f} MB)" + ) + + if total_bytes > 0: + self._update_progress_text(line, uploaded_bytes / total_bytes) + else: + self._update_progress_text(line, 1.0) + + def _update_progress_text(self, text: str, progress: float) -> None: + if text == self._last_printed_line: + return + self._last_printed_line = text + + if self._progress_text_area: + self._progress_text_area.set_text(text) + else: + self._printer.progress_update(text + "\r", progress) + + +class _DynamicOperationStatsPrinter: + """Single-use object that writes operation stats into a text area.""" + + def __init__( + self, + printer: p.Printer, + text_area: p.DynamicText, + max_lines: int, + loading_symbol: str, + default_text: str, + ) -> None: + self._printer = printer + self._text_area = text_area + self._max_lines = max_lines + self._loading_symbol = loading_symbol + self._default_text = default_text + + self._lines: list[str] = [] + self._ops_shown = 0 + + def display( + self, + stats_list: Iterable[pb.OperationStats], + ) -> None: + """Show the given stats in the text area.""" + total_operations = 0 + for stats in stats_list: + for op in stats.operations: + self._add_operation(op, is_subtask=False, indent="") + total_operations += stats.total_operations + + if self._ops_shown < total_operations: + if 1 <= self._max_lines <= len(self._lines): + self._lines.pop() + + remaining = total_operations - self._ops_shown + + self._lines.append(f"+ {remaining} more task(s)") + + if len(self._lines) == 0: + if self._loading_symbol: + self._text_area.set_text(f"{self._loading_symbol} {self._default_text}") + else: + self._text_area.set_text(self._default_text) + else: + self._text_area.set_text("\n".join(self._lines)) + + def _add_operation(self, op: pb.Operation, is_subtask: bool, indent: str) -> None: + """Add the operation to `self._lines`.""" + if len(self._lines) >= self._max_lines: + return + + if not is_subtask: + self._ops_shown += 1 + + parts = [] + + # Subtask indicator. + if is_subtask and self._printer.supports_unicode: + parts.append("↳") + + # Loading symbol. + if self._loading_symbol: + parts.append(self._loading_symbol) + + # Task name. + parts.append(op.desc) + + # Progress information. + if op.progress: + parts.append(f"{op.progress}") + + # Task duration. + parts.append(f"({_time_to_string(seconds=op.runtime_seconds)})") + + # Error status. + self._lines.append(indent + " ".join(parts)) + if op.error_status: + error_word = self._printer.error("ERROR") + error_desc = self._printer.secondary_text(op.error_status) + subtask_indent = " " if is_subtask else "" + self._lines.append( + f"{indent}{subtask_indent} {error_word} {error_desc}", + ) + + # Subtasks. + if op.subtasks: + subtask_indent = indent + " " + for task in op.subtasks: + self._add_operation( + task, + is_subtask=True, + indent=subtask_indent, + ) + + +def _time_to_string(seconds: float) -> str: + """Returns a short string representing the duration.""" + if seconds < 10: + return f"{seconds:.1f}s" + if seconds < 60: + return f"{seconds:.0f}s" + if seconds < 60 * 60: + minutes = seconds / 60 + return f"{minutes:.1f}m" + + hours = int(seconds / (60 * 60)) + minutes = int((seconds / 60) % 60) + return f"{hours}h{minutes}m" + + +def _megabytes(bytes: int) -> float: + """Returns the number of megabytes in `bytes`.""" + return bytes / (1 << 20) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/proto_util.py b/lib/python3.12/site-packages/wandb/sdk/lib/proto_util.py new file mode 100644 index 0000000000000000000000000000000000000000..f3a24fc33ed51e15959178692324c5bcd23096d1 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/proto_util.py @@ -0,0 +1,90 @@ +# +import json +from typing import TYPE_CHECKING, Any, Dict, Union + +from wandb.proto import wandb_internal_pb2 as pb + +if TYPE_CHECKING: # pragma: no cover + from google.protobuf.internal.containers import RepeatedCompositeFieldContainer + from google.protobuf.message import Message + + from wandb.proto import wandb_telemetry_pb2 as tpb + + +def dict_from_proto_list(obj_list: "RepeatedCompositeFieldContainer") -> Dict[str, Any]: + result: Dict[str, Any] = {} + + for item in obj_list: + # Start from the root of the result dict + current_level = result + + if len(item.nested_key) > 0: + keys = list(item.nested_key) + else: + keys = [item.key] + + for key in keys[:-1]: + if key not in current_level: + current_level[key] = {} + # Move the reference deeper into the nested dictionary + current_level = current_level[key] + + # Set the value at the final key location, parsing JSON from the value_json field + final_key = keys[-1] + current_level[final_key] = json.loads(item.value_json) + + return result + + +def _result_from_record(record: "pb.Record") -> "pb.Result": + result = pb.Result(uuid=record.uuid, control=record.control) + return result + + +def _assign_record_num(record: "pb.Record", record_num: int) -> None: + record.num = record_num + + +def _assign_end_offset(record: "pb.Record", end_offset: int) -> None: + record.control.end_offset = end_offset + + +def proto_encode_to_dict( + pb_obj: Union["tpb.TelemetryRecord", "pb.MetricRecord"], +) -> Dict[int, Any]: + data: Dict[int, Any] = dict() + fields = pb_obj.ListFields() + for desc, value in fields: + if desc.name.startswith("_"): + continue + if desc.type == desc.TYPE_STRING: + data[desc.number] = value + elif desc.type == desc.TYPE_INT32: + data[desc.number] = value + elif desc.type == desc.TYPE_ENUM: + data[desc.number] = value + elif desc.type == desc.TYPE_MESSAGE: + nested = value.ListFields() + bool_msg = all(d.type == d.TYPE_BOOL for d, _ in nested) + if bool_msg: + items = [d.number for d, v in nested if v] + if items: + data[desc.number] = items + else: + # TODO: for now this code only handles sub-messages with strings + md = {} + for d, v in nested: + if not v or d.type != d.TYPE_STRING: + continue + md[d.number] = v + data[desc.number] = md + return data + + +def message_to_dict( + message: "Message", +) -> Dict[str, Any]: + """Convert a protobuf message into a dictionary.""" + from google.protobuf.json_format import MessageToDict + + return MessageToDict(message, preserving_proto_field_name=True) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/redirect.py b/lib/python3.12/site-packages/wandb/sdk/lib/redirect.py new file mode 100644 index 0000000000000000000000000000000000000000..087c5b0661b3ed8d812f38b7c87a2f268d2db60a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/redirect.py @@ -0,0 +1,876 @@ +from __future__ import annotations + +try: + import fcntl + import pty + import termios + import tty +except ImportError: # windows + pty = tty = termios = fcntl = None # type: ignore + +import itertools +import logging +import os +import queue +import re +import signal +import struct +import sys +import threading +import time +from collections import defaultdict +from typing import Callable, Iterable, Literal + +import wandb +from wandb.sdk.lib import console_capture + + +class _Numpy: # fallback in case numpy is not available + def where(self, x): + return ([i for i in range(len(x)) if x[i]],) + + def diff(self, x): + return [x[i + 1] - x[i] for i in range(len(x) - 1)] + + def arange(self, x): + class Arr(list): + def __getitem__(self, s): + if isinstance(s, slice): + self._start = s.start + return self + return super().__getitem__(s) + + def __getslice__(self, i, j): + self._start = i + return self + + def __iadd__(self, i): # type: ignore + for j in range(self._start, len(self)): + self[j] += i + + return Arr(range(x)) + + +try: + import numpy as np # type: ignore +except ImportError: + np = _Numpy() # type: ignore + + +logger = logging.getLogger("wandb") + + +ANSI_CSI_RE = re.compile("\001?\033\\[((?:\\d|;)*)([a-zA-Z])\002?") +ANSI_OSC_RE = re.compile("\001?\033\\]([^\a]*)(\a)\002?") + +_LAST_WRITE_TOKEN = b"L@stWr!t3T0k3n" + +SEP_RE = re.compile( + "\r|\n|" + # Unprintable ascii characters: + + "|".join([chr(i) for i in range(2**8) if repr(chr(i)).startswith("'\\x")]) +) + +ANSI_FG = list(map(str, itertools.chain(range(30, 40), range(90, 98)))) +ANSI_BG = list(map(str, itertools.chain(range(40, 50), range(100, 108)))) + +ANSI_FG_DEFAULT = "39" +ANSI_BG_DEFAULT = "49" + +ANSI_RESET = "0" + +ANSI_STYLES = { + "1": "bold", + "2": "/bold", + "3": "italics", + "4": "underscore", + "5": "blink", + "7": "reverse", + "9": "strikethrough", + "22": "/bold", + "23": "/italics", + "24": "/underscore", + "25": "/blink", + "27": "/reverse", + "29": "/strikethrough", +} + +ANSI_STYLES_REV = {v: k for k, v in ANSI_STYLES.items()} + + +CSI = "\033[" + + +def _get_char(code): + return "\033[" + str(code) + "m" + + +class Char: + """Class encapsulating a single character, its foreground, background and style attributes.""" + + __slots__ = ( + "data", + "fg", + "bg", + "bold", + "italics", + "underscore", + "blink", + "strikethrough", + "reverse", + ) + + def __init__( + self, + data=" ", + fg=ANSI_FG_DEFAULT, + bg=ANSI_BG_DEFAULT, + bold=False, + italics=False, + underscore=False, + blink=False, + strikethrough=False, + reverse=False, + ): + self.data = data + self.fg = fg + self.bg = bg + self.bold = bold + self.italics = italics + self.underscore = underscore + self.blink = blink + self.strikethrough = strikethrough + self.reverse = reverse + + def reset(self): + # Reset everything other than data to defaults + default = self.__class__() + for k in self.__slots__[1:]: + self[k] = default[k] + + def __getitem__(self, k): + return getattr(self, k) + + def __setitem__(self, k, v): + setattr(self, k, v) + + def copy(self, **kwargs): + attrs = {} + for k in self.__slots__: + if k in kwargs: + attrs[k] = kwargs[k] + else: + attrs[k] = self[k] + return self.__class__(**attrs) + + def __eq__(self, other): + for k in self.__slots__: + if self[k] != other[k]: + return False + return True + + +_defchar = Char() + + +class Cursor: + """A 2D cursor. + + Attributes: + x: x-coordinate. + y: y-coordinate. + char: the character to inherit colors and styles from. + """ + + __slots__ = ("x", "y", "char") + + def __init__(self, x=0, y=0, char=None): + if char is None: + char = Char() + self.x = x + self.y = y + self.char = char + + +class TerminalEmulator: + """An FSM emulating a terminal. + + Characters are stored in a 2D matrix (buffer) indexed by the cursor. + """ + + _MAX_LINES = 100 + + def __init__(self): + self.buffer = defaultdict(lambda: defaultdict(lambda: _defchar)) + self.cursor = Cursor() + self._num_lines = None # Cache + + # For diffing: + self._prev_num_lines = None + self._prev_last_line = None + + def cursor_up(self, n=1): + n = min(n, self.cursor.y) + self.cursor.y -= n + + def cursor_down(self, n=1): + self.cursor.y += n + + def cursor_left(self, n=1): + n = min(n, self.cursor.x) + self.cursor.x -= n + + def cursor_right(self, n=1): + self.cursor.x += n + + def carriage_return(self): + self.cursor.x = 0 + + def cursor_position(self, line, column): + self.cursor.x = min(column, 1) - 1 + self.cursor.y = min(line, 1) - 1 + + def cursor_column(self, column): + self.cursor.x = min(column, 1) - 1 + + def cursor_line(self, line): + self.cursor.y = min(line, 1) - 1 + + def linefeed(self): + self.cursor_down() + self.carriage_return() + + def _get_line_len(self, n): + if n not in self.buffer: + return 0 + line = self.buffer[n] + if not line: + return 0 + n = max(line.keys()) + for i in range(n, -1, -1): + if line[i] != _defchar: + return i + 1 + return 0 + + @property + def num_lines(self): + if self._num_lines is not None: + return self._num_lines + ret = 0 + if self.buffer: + n = max(self.buffer.keys()) + for i in range(n, -1, -1): + if self._get_line_len(i): + ret = i + 1 + break + self._num_lines = ret + return ret + + def display(self): + return [ + [self.buffer[i][j].data for j in range(self._get_line_len(i))] + for i in range(self.num_lines) + ] + + def erase_screen(self, mode=0): + if mode == 0: + for i in range(self.cursor.y + 1, self.num_lines): + if i in self.buffer: + del self.buffer[i] + self.erase_line(mode) + if mode == 1: + for i in range(self.cursor.y): + if i in self.buffer: + del self.buffer[i] + self.erase_line(mode) + elif mode == 2 or mode == 3: + self.buffer.clear() + + def erase_line(self, mode=0): + curr_line = self.buffer[self.cursor.y] + if mode == 0: + for i in range(self.cursor.x, self._get_line_len(self.cursor.y)): + if i in curr_line: + del curr_line[i] + elif mode == 1: + for i in range(self.cursor.x + 1): + if i in curr_line: + del curr_line[i] + else: + curr_line.clear() + + def insert_lines(self, n=1): + for i in range(self.num_lines - 1, self.cursor.y, -1): + self.buffer[i + n] = self.buffer[i] + for i in range(self.cursor.y + 1, self.cursor.y + 1 + n): + if i in self.buffer: + del self.buffer[i] + + def _write_plain_text(self, plain_text): + self.buffer[self.cursor.y].update( + [ + (self.cursor.x + i, self.cursor.char.copy(data=c)) + for i, c in enumerate(plain_text) + ] + ) + self.cursor.x += len(plain_text) + + def _write_text(self, text): + prev_end = 0 + for match in SEP_RE.finditer(text): + start, end = match.span() + self._write_plain_text(text[prev_end:start]) + prev_end = end + c = match.group() + if c == "\n": + self.linefeed() + elif c == "\r": + self.carriage_return() + elif c == "\b": + self.cursor_left() + else: + continue + self._write_plain_text(text[prev_end:]) + + def _remove_osc(self, text): + return re.sub(ANSI_OSC_RE, "", text) + + def write(self, data): + self._num_lines = None # invalidate cache + data = self._remove_osc(data) + prev_end = 0 + for match in ANSI_CSI_RE.finditer(data): + start, end = match.span() + text = data[prev_end:start] + csi = data[start:end] + prev_end = end + self._write_text(text) + self._handle_csi(csi, *match.groups()) + self._write_text(data[prev_end:]) + + def _handle_csi(self, csi, params, command): + try: + if command == "m": + p = params.split(";")[0] + if not p: + p = "0" + if p in ANSI_FG: + self.cursor.char.fg = p + elif p in ANSI_BG: + self.cursor.char.bg = p + elif p == ANSI_RESET: + self.cursor.char.reset() + elif p in ANSI_STYLES: + style = ANSI_STYLES[p] + off = style.startswith("/") + if off: + style = style[1:] + self.cursor.char[style] = not off + else: + abcd = { + "A": "cursor_up", + "B": "cursor_down", + "C": "cursor_right", + "D": "cursor_left", + } + cursor_fn = abcd.get(command) + if cursor_fn: + getattr(self, cursor_fn)(int(params) if params else 1) + elif command == "J": + p = params.split(";")[0] + p = int(p) if p else 0 + self.erase_screen(p) + elif command == "K": + p = params.split(";")[0] + p = int(p) if p else 0 + self.erase_line(p) + elif command == "L": + p = int(params) if params else 1 + self.insert_lines(p) + elif command in "Hf": + p = params.split(";") + if len(p) == 2: + p = (int(p[0]), int(p[1])) + elif len(p) == 1: + p = (int(p[0]), 1) + else: + p = (1, 1) + self.cursor_position(*p) + except Exception: + pass + + def _get_line(self, n): + line = self.buffer[n] + line_len = self._get_line_len(n) + # We have to loop through each character in the line and check if foreground, + # background and other attributes (italics, bold, underline, etc) of the ith + # character are different from those of the (i-1)th character. If different, the + # appropriate ascii character for switching the color/attribute should be + # appended to the output string before appending the actual character. This loop + # and subsequent checks can be expensive, especially because 99% of terminal + # output use default colors and formatting. Even in outputs that do contain + # colors and styles, its unlikely that they will change on a per character + # basis. + + # So instead we create a character list without any ascii codes (`out`), and a + # list of all the foregrounds in the line (`fgs`) on which we call np.diff() and + # np.where() to find the indices where the foreground change, and insert the + # ascii characters in the output list (`out`) on those indices. All of this is + # the done only if there are more than 1 foreground color in the line in the + # first place (`if len(set(fgs)) > 1 else None`). Same logic is repeated for + # background colors and other attributes. + + out = [line[i].data for i in range(line_len)] + + # for dynamic insert using original indices + idxs = np.arange(line_len) + insert = lambda i, c: (out.insert(idxs[i], c), idxs[i:].__iadd__(1)) # noqa + + fgs = [int(_defchar.fg)] + [int(line[i].fg) for i in range(line_len)] + [ + insert(i, _get_char(line[int(i)].fg)) for i in np.where(np.diff(fgs))[0] + ] if len(set(fgs)) > 1 else None + bgs = [int(_defchar.bg)] + [int(line[i].bg) for i in range(line_len)] + [ + insert(i, _get_char(line[int(i)].bg)) for i in np.where(np.diff(bgs))[0] + ] if len(set(bgs)) > 1 else None + attrs = { + k: [False] + [line[i][k] for i in range(line_len)] + for k in Char.__slots__[3:] + } + [ + [ + insert(i, _get_char(ANSI_STYLES_REV[k if line[int(i)][k] else "/" + k])) + for i in np.where(np.diff(v))[0] + ] + for k, v in attrs.items() + if any(v) + ] + return "".join(out) + + def read(self): + num_lines = self.num_lines + if self._prev_num_lines is None: + ret = os.linesep.join(map(self._get_line, range(num_lines))) + if ret: + ret += os.linesep + else: + return ret + else: + curr_line = self._get_line(self._prev_num_lines - 1) + if curr_line == self._prev_last_line: + if num_lines == self._prev_num_lines: + return "" + ret = ( + os.linesep.join( + map(self._get_line, range(self._prev_num_lines, num_lines)) + ) + + os.linesep + ) + else: + ret = ( + "\r" + + os.linesep.join( + map(self._get_line, range(self._prev_num_lines - 1, num_lines)) + ) + + os.linesep + ) + if num_lines > self._MAX_LINES: + shift = num_lines - self._MAX_LINES + for i in range(shift, num_lines): + self.buffer[i - shift] = self.buffer[i] + for i in range(self._MAX_LINES, max(self.buffer.keys())): + if i in self.buffer: + del self.buffer[i] + self.cursor.y -= min(self.cursor.y, shift) + self._num_lines = num_lines = self._MAX_LINES + self._prev_num_lines = num_lines + self._prev_last_line = self._get_line(num_lines - 1) + return ret + + +_MIN_CALLBACK_INTERVAL = 2 # seconds + + +class RedirectBase: + def __init__( + self, + src: Literal["stdout", "stderr"], + cbs: Iterable[Callable[[str], None]] = (), + ) -> None: + """# Arguments. + + `src`: Source stream to be redirected. "stdout" or "stderr". + `cbs`: tuple/list of callbacks. Each callback should take exactly 1 argument (bytes). + + """ + assert hasattr(sys, src) + self.src: Literal["stdout", "stderr"] = src + self.cbs = cbs + + @property + def src_stream(self): + return getattr(sys, f"__{self.src}__") + + @property + def src_fd(self): + return self.src_stream.fileno() + + @property + def src_wrapped_stream(self): + return getattr(sys, self.src) + + def install(self) -> None: + pass + + def uninstall(self) -> None: + pass + + +class StreamWrapper(RedirectBase): + """Patches the write method of current sys.stdout/sys.stderr.""" + + def __init__( + self, + src: Literal["stdout", "stderr"], + cbs: Iterable[Callable[[str], None]] = (), + *, + flush_periodically: bool, + ) -> None: + super().__init__(src=src, cbs=cbs) + self._uninstall: Callable[[], None] | None = None + self._emulator = TerminalEmulator() + self._queue: queue.Queue[str] = queue.Queue() + self._stopped = threading.Event() + self._flush_periodically = flush_periodically + + def _emulator_write(self) -> None: + while True: + if self._queue.empty(): + if self._stopped.is_set(): + return + time.sleep(0.5) + continue + + data: list[str] = [] + while not self._queue.empty(): + data.append(self._queue.get()) + + if self._stopped.is_set() and sum(map(len, data)) > 100000: + wandb.termlog("Terminal output too large. Logging without processing.") + self.flush() + + for line in data: + self.flush(line) + + return + + try: + self._emulator.write("".join(data)) + except Exception: + pass + + def _callback(self) -> None: + while not (self._stopped.is_set() and self._queue.empty()): + self.flush() + time.sleep(_MIN_CALLBACK_INTERVAL) + + def _on_write(self, data: str | bytes, written: int, /) -> None: + if isinstance(data, bytes): + written_data = data[:written].decode("utf-8") + else: + written_data = data[:written] + + self._queue.put(written_data) + + def install(self) -> None: + if self._uninstall: + return + + try: + if self.src == "stdout": + self._uninstall = console_capture.capture_stdout(self._on_write) + else: + self._uninstall = console_capture.capture_stderr(self._on_write) + except console_capture.CannotCaptureConsoleError: + logger.exception("failed to install %s hooks", self.src) + wandb.termwarn( + f"Failed to wrap {self.src}. Console logs will not be captured.", + ) + return + + self._emulator_write_thread = threading.Thread(target=self._emulator_write) + self._emulator_write_thread.daemon = True + self._emulator_write_thread.start() + + if self._flush_periodically: + self._callback_thread = threading.Thread(target=self._callback) + self._callback_thread.daemon = True + self._callback_thread.start() + + def flush(self, data: str | None = None) -> None: + if data is None: + try: + data = self._emulator.read().encode("utf-8") + except Exception: + logger.exception("exception reading TerminalEmulator") + + if data: + for cb in self.cbs: + try: + cb(data) + except Exception: + logger.exception("exception in StreamWrapper callback") + + def uninstall(self) -> None: + if not self._uninstall: + return + + self._uninstall() + + self._stopped.set() + self._emulator_write_thread.join(timeout=5) + if self._emulator_write_thread.is_alive(): + wandb.termlog(f"Processing terminal output ({self.src})...") + self._emulator_write_thread.join() + wandb.termlog("Done.") + self.flush() + + +class StreamRawWrapper(RedirectBase): + """Patches the write method of current sys.stdout/sys.stderr. + + Captures data in a raw form rather than using the emulator + """ + + def __init__( + self, + src: Literal["stdout", "stderr"], + cbs: Iterable[Callable[[str], None]] = (), + ) -> None: + super().__init__(src=src, cbs=cbs) + self._uninstall: Callable[[], None] | None = None + + def _on_write(self, data: str | bytes, written: int, /) -> None: + if isinstance(data, bytes): + written_data = data[:written].decode("utf-8") + else: + written_data = data[:written] + + for cb in self.cbs: + try: + cb(written_data) + except Exception: + logger.exception("error in %s callback", self.src) + + def install(self) -> None: + if self._uninstall: + return + + try: + if self.src == "stdout": + self._uninstall = console_capture.capture_stdout(self._on_write) + else: + self._uninstall = console_capture.capture_stderr(self._on_write) + except console_capture.CannotCaptureConsoleError: + logger.exception("failed to install %s hooks", self.src) + wandb.termwarn( + f"Failed to wrap {self.src}. Console logs will not be captured.", + ) + + def uninstall(self) -> None: + if self._uninstall: + self._uninstall() + + +class _WindowSizeChangeHandler: + def __init__(self): + self._fds = set() + + def _register(self): + old_handler = signal.signal(signal.SIGWINCH, lambda *_: None) + + def handler(signum, frame): + if callable(old_handler): + old_handler(signum, frame) + self.handle_window_size_change() + + signal.signal(signal.SIGWINCH, handler) + self._old_handler = old_handler + + def _unregister(self): + signal.signal(signal.SIGWINCH, self._old_handler) + + def add_fd(self, fd): + if not self._fds: + self._register() + self._fds.add(fd) + self.handle_window_size_change() + + def remove_fd(self, fd): + if fd in self._fds: + self._fds.remove(fd) + if not self._fds: + self._unregister() + + def handle_window_size_change(self): + try: + win_size = fcntl.ioctl(0, termios.TIOCGWINSZ, "\0" * 8) + rows, cols, xpix, ypix = struct.unpack("HHHH", win_size) + # Note: IOError not subclass of OSError in python 2.x + except OSError: # eg. in MPI we can't do this. + return + if cols == 0: + return + win_size = struct.pack("HHHH", rows, cols, xpix, ypix) + for fd in self._fds: + fcntl.ioctl(fd, termios.TIOCSWINSZ, win_size) + + +_WSCH = _WindowSizeChangeHandler() + +_redirects: dict[str, Redirect | None] = {"stdout": None, "stderr": None} + + +class Redirect(RedirectBase): + """Redirect low level file descriptors.""" + + def __init__(self, src, cbs=(), *, flush_periodically: bool): + super().__init__(src=src, cbs=cbs) + self._installed = False + self._emulator = TerminalEmulator() + self._flush_periodically = flush_periodically + + def _pipe(self): + if pty: + r, w = pty.openpty() + else: + r, w = os.pipe() + return r, w + + def install(self): + curr_redirect = _redirects.get(self.src) + if curr_redirect and curr_redirect != self: + curr_redirect.uninstall() + _redirects[self.src] = self + + if self._installed: + return + self._pipe_read_fd, self._pipe_write_fd = self._pipe() + if os.isatty(self._pipe_read_fd): + _WSCH.add_fd(self._pipe_read_fd) + self._orig_src_fd = os.dup(self.src_fd) + self._orig_src = os.fdopen(self._orig_src_fd, "wb", 0) + os.dup2(self._pipe_write_fd, self.src_fd) + self._installed = True + self._queue = queue.Queue() + self._stopped = threading.Event() + self._pipe_relay_thread = threading.Thread(target=self._pipe_relay) + self._pipe_relay_thread.daemon = True + self._pipe_relay_thread.start() + self._emulator_write_thread = threading.Thread(target=self._emulator_write) + self._emulator_write_thread.daemon = True + self._emulator_write_thread.start() + if self._flush_periodically: + self._callback_thread = threading.Thread(target=self._callback) + self._callback_thread.daemon = True + self._callback_thread.start() + + def uninstall(self): + if not self._installed: + return + self._installed = False + # If the user printed a very long string (millions of chars) right + # before run.finish(), it will take a while for it to reach pipe relay. + # 1 second is enough time for ~5 million chars. + time.sleep(1) + self._stopped.set() + os.dup2(self._orig_src_fd, self.src_fd) + os.write(self._pipe_write_fd, _LAST_WRITE_TOKEN) + self._pipe_relay_thread.join() + os.close(self._pipe_read_fd) + os.close(self._pipe_write_fd) + + t = threading.Thread( + target=self.src_wrapped_stream.flush + ) # Calling flush() from the current thread does not flush the buffer instantly. + t.start() + t.join(timeout=10) + + self._emulator_write_thread.join(timeout=5) + if self._emulator_write_thread.is_alive(): + wandb.termlog(f"Processing terminal output ({self.src})...") + self._emulator_write_thread.join() + wandb.termlog("Done.") + self.flush() + + _WSCH.remove_fd(self._pipe_read_fd) + + if _redirects[self.src] == self: + _redirects[self.src] = None + + def flush(self, data=None): + if data is None: + try: + data = self._emulator.read().encode("utf-8") + except Exception: + pass + if data: + for cb in self.cbs: + try: + cb(data) + except Exception: + pass # TODO(frz) + + def _callback(self): + while not self._stopped.is_set(): + self.flush() + time.sleep(_MIN_CALLBACK_INTERVAL) + + def _pipe_relay(self): + while True: + try: + brk = False + data = os.read(self._pipe_read_fd, 4096) + if self._stopped.is_set(): + if _LAST_WRITE_TOKEN not in data: + # _LAST_WRITE_TOKEN could have gotten split up at the 4096 border + n = len(_LAST_WRITE_TOKEN) + while n and data[-n:] != _LAST_WRITE_TOKEN[:n]: + n -= 1 + if n: + data += os.read( + self._pipe_read_fd, len(_LAST_WRITE_TOKEN) - n + ) + if _LAST_WRITE_TOKEN in data: + data = data.replace(_LAST_WRITE_TOKEN, b"") + brk = True + i = self._orig_src.write(data) + if i is not None: # python 3 w/ unbuffered i/o: we need to keep writing + while i < len(data): + i += self._orig_src.write(data[i:]) + self._queue.put(data) + if brk: + return + except OSError: + return + + def _emulator_write(self): + while True: + if self._queue.empty(): + if self._stopped.is_set(): + return + time.sleep(0.5) + continue + data = [] + while not self._queue.empty(): + data.append(self._queue.get()) + if self._stopped.is_set() and sum(map(len, data)) > 100000: + wandb.termlog("Terminal output too large. Logging without processing.") + self.flush() + [self.flush(line) for line in data] + return + try: + self._emulator.write(b"".join(data).decode("utf-8")) + except Exception: + pass diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/retry.py b/lib/python3.12/site-packages/wandb/sdk/lib/retry.py new file mode 100644 index 0000000000000000000000000000000000000000..a920fb6c85311bde47d1510842a286f4169f27a4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/retry.py @@ -0,0 +1,288 @@ +import abc +import asyncio +import datetime +import functools +import logging +import os +import random +import threading +import time +from typing import Any, Awaitable, Callable, Generic, Optional, Tuple, Type, TypeVar + +from requests import HTTPError + +import wandb +import wandb.errors +from wandb.util import CheckRetryFnType + +logger = logging.getLogger(__name__) + + +# To let tests mock out the retry logic's now()/sleep() funcs, this file +# should only use these variables, not call the stdlib funcs directly. +NOW_FN = datetime.datetime.now +SLEEP_FN = time.sleep +SLEEP_ASYNC_FN = asyncio.sleep + + +class RetryCancelledError(wandb.errors.Error): + """A retry did not occur because it was cancelled.""" + + +class TransientError(Exception): + """Exception type designated for errors that may only be temporary. + + Can have its own message and/or wrap another exception. + """ + + def __init__( + self, msg: Optional[str] = None, exc: Optional[BaseException] = None + ) -> None: + super().__init__(msg) + self.message = msg + self.exception = exc + + +_R = TypeVar("_R") + + +class Retry(Generic[_R]): + """Create a retryable version of a function. + + Calling this will call the passed function, retrying if any exceptions in + retryable_exceptions are caught, with exponential backoff. + """ + + MAX_SLEEP_SECONDS = 5 * 60 + + def __init__( + self, + call_fn: Callable[..., _R], + retry_timedelta: Optional[datetime.timedelta] = None, + retry_cancel_event: Optional[threading.Event] = None, + num_retries: Optional[int] = None, + check_retry_fn: CheckRetryFnType = lambda e: True, + retryable_exceptions: Optional[Tuple[Type[Exception], ...]] = None, + error_prefix: str = "Network error", + retry_callback: Optional[Callable[[int, str], Any]] = None, + ) -> None: + self._call_fn = call_fn + self._check_retry_fn = check_retry_fn + self._error_prefix = error_prefix + self._last_print = datetime.datetime.now() - datetime.timedelta(minutes=1) + self._retry_timedelta = retry_timedelta + self._retry_cancel_event = retry_cancel_event + self._num_retries = num_retries + if retryable_exceptions is not None: + self._retryable_exceptions = retryable_exceptions + else: + self._retryable_exceptions = (TransientError,) + self._index = 0 + self.retry_callback = retry_callback + + def _sleep_check_cancelled( + self, wait_seconds: float, cancel_event: Optional[threading.Event] + ) -> bool: + if not cancel_event: + SLEEP_FN(wait_seconds) + return False + cancelled = cancel_event.wait(wait_seconds) + return cancelled + + @property + def num_iters(self) -> int: + """The number of iterations the previous __call__ retried.""" + return self._num_iter + + def __call__(self, *args: Any, **kwargs: Any) -> _R: + """Call the wrapped function, with retries. + + Args: + retry_timedelta (kwarg): amount of time to retry before giving up. + sleep_base (kwarg): amount of time to sleep upon first failure, all other sleeps + are derived from this one. + """ + retry_timedelta = kwargs.pop("retry_timedelta", self._retry_timedelta) + if retry_timedelta is None: + retry_timedelta = datetime.timedelta(days=365) + + retry_cancel_event = kwargs.pop("retry_cancel_event", self._retry_cancel_event) + + num_retries = kwargs.pop("num_retries", self._num_retries) + if num_retries is None: + num_retries = 1000000 + + if os.environ.get("WANDB_TEST"): + num_retries = 0 + + sleep_base: float = kwargs.pop("retry_sleep_base", 1) + + # an extra function to allow performing more logic on the filtered exception + check_retry_fn: CheckRetryFnType = kwargs.pop( + "check_retry_fn", self._check_retry_fn + ) + + sleep = sleep_base + now = NOW_FN() + start_time = now + start_time_triggered = None + + self._num_iter = 0 + + while True: + try: + result = self._call_fn(*args, **kwargs) + # Only print resolved attempts once every minute + if self._num_iter > 2 and now - self._last_print > datetime.timedelta( + minutes=1 + ): + self._last_print = NOW_FN() + if self.retry_callback: + self.retry_callback( + 200, + f"{self._error_prefix} resolved after {NOW_FN() - start_time}, resuming normal operation.", + ) + return result + except self._retryable_exceptions as e: + # if the secondary check fails, re-raise + retry_timedelta_triggered = check_retry_fn(e) + if not retry_timedelta_triggered: + raise + + # always enforce num_retries no matter which type of exception was seen + if self._num_iter >= num_retries: + raise + + now = NOW_FN() + + # handle a triggered secondary check which could have a shortened timeout + if isinstance(retry_timedelta_triggered, datetime.timedelta): + # save the time of the first secondary trigger + if not start_time_triggered: + start_time_triggered = now + + # make sure that we haven't run out of time from secondary trigger + if now - start_time_triggered >= retry_timedelta_triggered: + raise + + # always enforce the default timeout from start of retries + if now - start_time >= retry_timedelta: + raise + + if self._num_iter == 2: + logger.info("Retry attempt failed:", exc_info=e) + if ( + isinstance(e, HTTPError) + and e.response is not None + and self.retry_callback is not None + ): + self.retry_callback(e.response.status_code, e.response.text) + else: + # todo: would like to catch other errors, eg wandb.errors.Error, ConnectionError etc + # but some of these can be raised before the retry handler thread (RunStatusChecker) is + # spawned in wandb_init + wandb.termlog( + f"{self._error_prefix} ({e.__class__.__name__}), entering retry loop." + ) + # if wandb.env.is_debug(): + # traceback.print_exc() + cancelled = self._sleep_check_cancelled( + sleep + random.random() * 0.25 * sleep, cancel_event=retry_cancel_event + ) + if cancelled: + raise RetryCancelledError("retry timeout") + sleep *= 2 + if sleep > self.MAX_SLEEP_SECONDS: + sleep = self.MAX_SLEEP_SECONDS + now = NOW_FN() + + self._num_iter += 1 + + +_F = TypeVar("_F", bound=Callable) + + +def retriable(*args: Any, **kargs: Any) -> Callable[[_F], _F]: + def decorator(fn: _F) -> _F: + retrier: Retry[Any] = Retry(fn, *args, **kargs) + + @functools.wraps(fn) + def wrapped_fn(*args: Any, **kargs: Any) -> Any: + return retrier(*args, **kargs) + + return wrapped_fn # type: ignore + + return decorator + + +class Backoff(abc.ABC): + """A backoff strategy: decides whether to sleep or give up when an exception is raised.""" + + @abc.abstractmethod + def next_sleep_or_reraise(self, exc: Exception) -> datetime.timedelta: + raise NotImplementedError # pragma: no cover + + +class ExponentialBackoff(Backoff): + """Jittered exponential backoff: sleep times increase ~exponentially up to some limit.""" + + def __init__( + self, + initial_sleep: datetime.timedelta, + max_sleep: datetime.timedelta, + max_retries: Optional[int] = None, + timeout_at: Optional[datetime.datetime] = None, + ) -> None: + self._next_sleep = min(max_sleep, initial_sleep) + self._max_sleep = max_sleep + self._remaining_retries = max_retries + self._timeout_at = timeout_at + + def next_sleep_or_reraise(self, exc: Exception) -> datetime.timedelta: + if self._remaining_retries is not None: + if self._remaining_retries <= 0: + raise exc + self._remaining_retries -= 1 + + if self._timeout_at is not None and NOW_FN() > self._timeout_at: + raise exc + + result, self._next_sleep = ( + self._next_sleep, + min(self._max_sleep, self._next_sleep * (1 + random.random())), + ) + + return result + + +class FilteredBackoff(Backoff): + """Re-raise any exceptions that fail a predicate; delegate others to another Backoff.""" + + def __init__(self, filter: Callable[[Exception], bool], wrapped: Backoff) -> None: + self._filter = filter + self._wrapped = wrapped + + def next_sleep_or_reraise(self, exc: Exception) -> datetime.timedelta: + if not self._filter(exc): + raise exc + return self._wrapped.next_sleep_or_reraise(exc) + + +async def retry_async( + backoff: Backoff, + fn: Callable[..., Awaitable[_R]], + *args: Any, + on_exc: Optional[Callable[[Exception], None]] = None, + **kwargs: Any, +) -> _R: + """Call `fn` repeatedly until either it succeeds, or `backoff` decides we should give up. + + Each time `fn` fails, `on_exc` is called with the exception. + """ + while True: + try: + return await fn(*args, **kwargs) + except Exception as e: + if on_exc is not None: + on_exc(e) + await SLEEP_ASYNC_FN(backoff.next_sleep_or_reraise(e).total_seconds()) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/run_moment.py b/lib/python3.12/site-packages/wandb/sdk/lib/run_moment.py new file mode 100644 index 0000000000000000000000000000000000000000..aa293b5a480382f5b2bb51528be2833c84ce3cdb --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/run_moment.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Union +from urllib import parse + + +@dataclass +class RunMoment: + """A moment in a run. + + Defines a branching point in a finished run to fork or resume from. + A run moment is identified by a run ID and a metric value. + Currently, only the metric '_step' is supported. + """ + + run: str + """run ID""" + + value: Union[int, float] + """Value of the metric.""" + + metric: str = "_step" + """Metric to use to determine the moment in the run. + + Currently, only the metric '_step' is supported. + In future, this will be relaxed to be any metric. + """ + + 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 TypeError( + f"Only int or float values are supported, got '{self.value}'." + ) + if not isinstance(self.run, str): + raise TypeError(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 '?='. " + 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=metric, value=num_value) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/runid.py b/lib/python3.12/site-packages/wandb/sdk/lib/runid.py new file mode 100644 index 0000000000000000000000000000000000000000..355d24df01a2fc9b5a6174a343e84ed083696277 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/runid.py @@ -0,0 +1,12 @@ +"""runid util.""" + +import secrets +import string + + +def generate_id(length: int = 8) -> str: + """Generate a random base-36 string of `length` digits.""" + # There are ~2.8T base-36 8-digit strings. If we generate 210k ids, + # we'll have a ~1% chance of collision. + alphabet = string.ascii_lowercase + string.digits + return "".join(secrets.choice(alphabet) for _ in range(length)) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/server.py b/lib/python3.12/site-packages/wandb/sdk/lib/server.py new file mode 100644 index 0000000000000000000000000000000000000000..8b12362d51d5780d03820f2b9c43b3d2b837d279 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/server.py @@ -0,0 +1,58 @@ +"""module server.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any + +from wandb import util +from wandb.apis import InternalApi + +if TYPE_CHECKING: + from wandb.sdk.wandb_settings import Settings + + +class Server: + def __init__( + self, + settings: Settings, + ) -> None: + self._api = InternalApi(default_settings=settings) + self._viewer: dict[str, Any] = {} + self._flags: dict[str, Any] = {} + self._settings: Settings = settings + + def query_with_timeout(self, timeout: int | float = 5) -> None: + if self._settings.x_disable_viewer: + return + + async_viewer = util.async_call(self._api.viewer_server_info, timeout=timeout) + try: + viewer_tuple, viewer_thread = async_viewer() + except Exception: + return + if viewer_thread.is_alive(): + # this is likely a DNS hang + return + # TODO(jhr): should we kill the thread? + self._viewer, self._serverinfo = viewer_tuple + self._flags = json.loads(self._viewer.get("flags", "{}")) + + @property + def viewer(self) -> dict[str, Any]: + """Returns information about the currently authenticated user. + + If the API key is valid, the following is returned: + - id + - entity + - username + - flags + - teams + + If the API key is not valid or the server is not reachable, + an empty dict is returned. + """ + if not self._viewer and not self._settings._offline: + self.query_with_timeout() + + return self._viewer diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/sock_client.py b/lib/python3.12/site-packages/wandb/sdk/lib/sock_client.py new file mode 100644 index 0000000000000000000000000000000000000000..f8488bc29933aa97386180e4cf02dc6d6d1bf37f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/sock_client.py @@ -0,0 +1,236 @@ +import socket +import struct +import threading +import time +import uuid +from typing import TYPE_CHECKING, Any, List, Optional + +from wandb.proto import wandb_server_pb2 as spb + +if TYPE_CHECKING: + from wandb.proto import wandb_internal_pb2 as pb + + +class SockClientClosedError(Exception): + """Raised on operations on a closed socket.""" + + +class SockClientTimeoutError(Exception): + """Raised if the server didn't respond before the timeout.""" + + +class SockBuffer: + _buf_list: List[bytes] + _buf_lengths: List[int] + _buf_total: int + + def __init__(self) -> None: + self._buf_list = [] + self._buf_lengths = [] + self._buf_total = 0 + + @property + def length(self) -> int: + return self._buf_total + + def _get(self, start: int, end: int, peek: bool = False) -> bytes: + index: Optional[int] = None + buffers = [] + need = end + + # compute buffers needed + for i, (buf_len, buf_data) in enumerate(zip(self._buf_lengths, self._buf_list)): + buffers.append(buf_data[:need] if need < buf_len else buf_data) + if need <= buf_len: + index = i + break + need -= buf_len + + # buffer not large enough, caller should have made sure there was enough data + if index is None: + raise IndexError("SockBuffer index out of range") + + # advance buffer internals if we are not peeking into the data + if not peek: + self._buf_total -= end + if need < buf_len: + # update partially used buffer list + self._buf_list = self._buf_list[index:] + self._buf_lengths = self._buf_lengths[index:] + self._buf_list[0] = self._buf_list[0][need:] + self._buf_lengths[0] -= need + else: + # update fully used buffer list + self._buf_list = self._buf_list[index + 1 :] + self._buf_lengths = self._buf_lengths[index + 1 :] + + return b"".join(buffers)[start:end] + + def get(self, start: int, end: int) -> bytes: + return self._get(start, end) + + def peek(self, start: int, end: int) -> bytes: + return self._get(start, end, peek=True) + + def put(self, data: bytes, data_len: int) -> None: + self._buf_list.append(data) + self._buf_lengths.append(data_len) + self._buf_total += data_len + + +class SockClient: + # current header is magic byte "W" followed by 4 byte length of the message + HEADLEN = 1 + 4 + + def __init__(self, sock: socket.socket) -> None: + """Create a SockClient. + + Args: + sock: A connected socket. + """ + self._sock = sock + + # TODO: use safe uuid's (python3.7+) or emulate this + self._sockid = uuid.uuid4().hex + self._retry_delay = 0.1 + self._lock = threading.Lock() + self._bufsize = 4096 + self._buffer = SockBuffer() + + self._detect_bufsize() + + def _detect_bufsize(self) -> None: + sndbuf_size = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF) + rcvbuf_size = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF) + self._bufsize = min(sndbuf_size, rcvbuf_size, 65536) + + def close(self) -> None: + self._sock.close() + + def shutdown(self, val: int) -> None: + self._sock.shutdown(val) + + def set_socket(self, sock: socket.socket) -> None: + self._sock = sock + self._detect_bufsize() + + def _sendall_with_error_handle(self, data: bytes) -> None: + # This is a helper function for sending data in a retry fashion. + # Similar to the sendall() function in the socket module, but with + # an error handling in case of timeout. + total_sent = 0 + total_data = len(data) + while total_sent < total_data: + start_time = time.monotonic() + try: + sent = self._sock.send(data) + # sent equal to 0 indicates a closed socket + if sent == 0: + raise SockClientClosedError("socket connection broken") + total_sent += sent + # truncate our data to save memory + data = data[sent:] + # we handle the timeout case for the cases when timeout is set + # on a system level by another application + except socket.timeout: + # adding sleep to avoid tight loop + delta_time = time.monotonic() - start_time + if delta_time < self._retry_delay: + time.sleep(self._retry_delay - delta_time) + + def _send_message(self, msg: Any) -> None: + raw_size = msg.ByteSize() + data = msg.SerializeToString() + assert len(data) == raw_size, "invalid serialization" + header = struct.pack(" None: + self._send_message(msg) + + def send_server_response(self, msg: spb.ServerResponse) -> None: + try: + self._send_message(msg) + except BrokenPipeError: + # TODO(jhr): user thread might no longer be around to receive responses to + # things like network status poll loop, there might be a better way to quiesce + pass + + def send_record_communicate(self, record: "pb.Record") -> None: + server_req = spb.ServerRequest() + server_req.request_id = record.control.mailbox_slot + server_req.record_communicate.CopyFrom(record) + self.send_server_request(server_req) + + def send_record_publish(self, record: "pb.Record") -> None: + server_req = spb.ServerRequest() + server_req.request_id = record.control.mailbox_slot + server_req.record_publish.CopyFrom(record) + self.send_server_request(server_req) + + def _extract_packet_bytes(self) -> Optional[bytes]: + # Do we have enough data to read the header? + start_offset = self.HEADLEN + if self._buffer.length >= start_offset: + header = self._buffer.peek(0, start_offset) + fields = struct.unpack("= end_offset: + rec_data = self._buffer.get(start_offset, end_offset) + return rec_data + return None + + def _read_packet_bytes(self, timeout: Optional[int] = None) -> Optional[bytes]: + """Read full message from socket. + + Args: + timeout: number of seconds to wait on socket data. + + Raises: + SockClientClosedError: socket has been closed. + """ + while True: + rec = self._extract_packet_bytes() + if rec: + return rec + + if timeout: + self._sock.settimeout(timeout) + try: + data = self._sock.recv(self._bufsize) + except socket.timeout: + break + except OSError as e: + raise SockClientClosedError from e + finally: + if timeout: + self._sock.settimeout(None) + data_len = len(data) + if data_len == 0: + # socket.recv() will return 0 bytes if socket was shutdown + # caller will handle this condition like other connection problems + raise SockClientClosedError + self._buffer.put(data, data_len) + return None + + def read_server_request(self) -> Optional[spb.ServerRequest]: + data = self._read_packet_bytes() + if not data: + return None + rec = spb.ServerRequest() + rec.ParseFromString(data) + return rec + + def read_server_response( + self, timeout: Optional[int] = None + ) -> Optional[spb.ServerResponse]: + data = self._read_packet_bytes(timeout=timeout) + if not data: + return None + rec = spb.ServerResponse() + rec.ParseFromString(data) + return rec diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/sparkline.py b/lib/python3.12/site-packages/wandb/sdk/lib/sparkline.py new file mode 100644 index 0000000000000000000000000000000000000000..e7328cc328cb0df804ca9a8cfc68f6ce9a5ea9ec --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/sparkline.py @@ -0,0 +1,44 @@ +# +# From pysparklines (BSD License): https://pypi.python.org/pypi/pysparklines + +import math +from typing import List, Union + +spark_chars = "▁▂▃▄▅▆▇█" + + +# math.isfinite doesn't exist in python2, so provider our own +def isfinite(f): + return not (math.isinf(f) or math.isnan(f)) + + +def sparkify(series: List[Union[float, int]]) -> str: + """Convert to a sparkline string. + + Example: + >>> sparkify([0.5, 1.2, 3.5, 7.3, 8.0, 12.5, 13.2, 15.0, 14.2, 11.8, 6.1, 1.9]) + u'▁▁▂▄▅▇▇██▆▄▂' + + >>> sparkify([1, 1, -2, 3, -5, 8, -13]) + u'▆▆▅▆▄█▁' + + Raises ValueError if input data cannot be converted to float. + Raises TypeError if series is not an iterable. + """ + series = [float(i) for i in series] + finite_series = [x for x in series if isfinite(x)] + if not finite_series: + return "" + minimum = min(finite_series) + maximum = max(finite_series) + data_range = maximum - minimum + if data_range == 0.0: + # Graph a baseline if every input value is equal. + return "".join([spark_chars[0] if isfinite(x) else " " for x in series]) + coefficient = (len(spark_chars) - 1.0) / data_range + return "".join( + [ + spark_chars[int(round((x - minimum) * coefficient))] if isfinite(x) else " " + for x in series + ] + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/telemetry.py b/lib/python3.12/site-packages/wandb/sdk/lib/telemetry.py new file mode 100644 index 0000000000000000000000000000000000000000..b1fc67efe9cb6a89f7906909519b8314f7f2bb81 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/telemetry.py @@ -0,0 +1,100 @@ +import re +import sys +from types import TracebackType +from typing import TYPE_CHECKING, ContextManager, Dict, List, Optional, Set, Type + +import wandb +from wandb.proto.wandb_telemetry_pb2 import Imports as TelemetryImports +from wandb.proto.wandb_telemetry_pb2 import TelemetryRecord + +# avoid cycle, use string type reference + +if TYPE_CHECKING: + from .. import wandb_run + + +_LABEL_TOKEN: str = "@wandbcode{" + + +class _TelemetryObject: + _run: Optional["wandb_run.Run"] + _obj: TelemetryRecord + + def __init__( + self, + run: Optional["wandb_run.Run"] = None, + obj: Optional[TelemetryRecord] = None, + ) -> None: + self._run = run or wandb.run + self._obj = obj or TelemetryRecord() + + def __enter__(self) -> TelemetryRecord: + return self._obj + + def __exit__( + self, + exctype: Optional[Type[BaseException]], + excinst: Optional[BaseException], + exctb: Optional[TracebackType], + ) -> None: + if not self._run: + return + self._run._telemetry_callback(self._obj) + + +def context( + run: Optional["wandb_run.Run"] = None, obj: Optional[TelemetryRecord] = None +) -> ContextManager[TelemetryRecord]: + return _TelemetryObject(run=run, obj=obj) + + +MATCH_RE = re.compile(r"(?P[a-zA-Z0-9_-]+)[,}](?P.*)") + + +def _parse_label_lines(lines: List[str]) -> Dict[str, str]: + seen = False + ret = {} + for line in lines: + idx = line.find(_LABEL_TOKEN) + if idx < 0: + # Stop parsing on first non token line after match + if seen: + break + continue + seen = True + label_str = line[idx + len(_LABEL_TOKEN) :] + + # match identifier (first token without key=value syntax (optional) + # Note: Parse is fairly permissive as it does not enforce strict syntax + r = MATCH_RE.match(label_str) + if r: + ret["code"] = r.group("code").replace("-", "_") + label_str = r.group("rest") + + # match rest of tokens on one line + tokens = re.findall( + r'([a-zA-Z0-9_]+)\s*=\s*("[a-zA-Z0-9_-]*"|[a-zA-Z0-9_-]*)[,}]', label_str + ) + for k, v in tokens: + ret[k] = v.strip('"').replace("-", "_") + return ret + + +def list_telemetry_imports(only_imported: bool = False) -> Set[str]: + import_telemetry_set = { + desc.name + for desc in TelemetryImports.DESCRIPTOR.fields + if desc.type == desc.TYPE_BOOL + } + if only_imported: + imported_modules_set = set(sys.modules) + return imported_modules_set.intersection(import_telemetry_set) + return import_telemetry_set + + +__all__ = [ + "TelemetryImports", + "TelemetryRecord", + "context", + "list_telemetry_imports", +] diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/timed_input.py b/lib/python3.12/site-packages/wandb/sdk/lib/timed_input.py new file mode 100644 index 0000000000000000000000000000000000000000..6d963ae0521263453ac7f68ed2906934dd7c0f95 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/timed_input.py @@ -0,0 +1,133 @@ +"""timed_input: add a timeout to standard input. + +Approach was inspired by: https://github.com/johejo/inputimeout +""" + +import sys +import threading + +import wandb + +SP = " " +CR = "\r" +LF = "\n" +CRLF = CR + LF + + +def _echo(prompt: str) -> None: + sys.stdout.write(prompt) + sys.stdout.flush() + + +def _posix_timed_input(prompt: str, timeout: float) -> str: + _echo(prompt) + sel = selectors.DefaultSelector() + sel.register(sys.stdin, selectors.EVENT_READ, data=sys.stdin.readline) + events = sel.select(timeout=timeout) + + for key, _ in events: + input_callback = key.data + input_data: str = input_callback() + if not input_data: # end-of-file - treat as timeout + raise TimeoutError + return input_data.rstrip(LF) + + _echo(LF) + termios.tcflush(sys.stdin, termios.TCIFLUSH) + raise TimeoutError + + +def _windows_timed_input(prompt: str, timeout: float) -> str: + interval = 0.1 + + _echo(prompt) + begin = time.monotonic() + end = begin + timeout + line = "" + + while time.monotonic() < end: + if msvcrt.kbhit(): # type: ignore[attr-defined] + c = msvcrt.getwche() # type: ignore[attr-defined] + if c in (CR, LF): + _echo(CRLF) + return line + if c == "\003": + raise KeyboardInterrupt + if c == "\b": + line = line[:-1] + cover = SP * len(prompt + line + SP) + _echo("".join([CR, cover, CR, prompt, line])) + else: + line += c + time.sleep(interval) + + _echo(CRLF) + raise TimeoutError + + +def _jupyter_timed_input(prompt: str, timeout: float) -> str: + clear = True + try: + from IPython.core.display import clear_output # type: ignore + except ImportError: + clear = False + wandb.termwarn( + "Unable to clear output, can't import clear_output from ipython.core" + ) + + _echo(prompt) + + user_inp = None + event = threading.Event() + + def get_input() -> None: + nonlocal user_inp + raw = input() + if event.is_set(): + return + user_inp = raw + + t = threading.Thread(target=get_input) + t.start() + t.join(timeout) + event.set() + if user_inp: + return user_inp + if clear: + clear_output() + raise TimeoutError + + +def timed_input( + prompt: str, timeout: float, show_timeout: bool = True, jupyter: bool = False +) -> str: + """Behaves like builtin `input()` but adds timeout. + + Args: + prompt (str): Prompt to output to stdout. + timeout (float): Timeout to wait for input. + show_timeout (bool): Show timeout in prompt + jupyter (bool): If True, use jupyter specific code. + + Raises: + TimeoutError: exception raised if timeout occurred. + """ + if show_timeout: + prompt = f"{prompt}({timeout:.0f} second timeout) " + if jupyter: + return _jupyter_timed_input(prompt=prompt, timeout=timeout) + + return _timed_input(prompt=prompt, timeout=timeout) + + +try: + import msvcrt +except ImportError: + import selectors + import termios + + _timed_input = _posix_timed_input +else: + import time + + _timed_input = _windows_timed_input diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/timer.py b/lib/python3.12/site-packages/wandb/sdk/lib/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..0ff2be95dc9dcfd5a307558ad936d1b22242f882 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/timer.py @@ -0,0 +1,19 @@ +import time +from typing import Any + + +class Timer: + def __init__(self) -> None: + self.start_time: float = time.time() + self.start: float = time.perf_counter() + self.stop: float = self.start + + def __enter__(self) -> "Timer": + return self + + def __exit__(self, *args: Any) -> None: + self.stop = time.perf_counter() + + @property + def elapsed(self) -> float: + return self.stop - self.start diff --git a/lib/python3.12/site-packages/wandb/sdk/lib/wb_logging.py b/lib/python3.12/site-packages/wandb/sdk/lib/wb_logging.py new file mode 100644 index 0000000000000000000000000000000000000000..af744b9374bfae99403b5120c56095555230b3d2 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/lib/wb_logging.py @@ -0,0 +1,161 @@ +"""Logging configuration for the "wandb" logger. + +Most log statements in wandb are made in the context of a run and should be +redirected to that run's log file (usually named 'debug.log'). This module +provides a context manager to temporarily set the current run ID and registers +a global handler for the 'wandb' logger that sends log statements to the right +place. + +All functions in this module are threadsafe. + +NOTE: The pytest caplog fixture will fail to capture logs from the wandb logger +because they are not propagated to the root logger. +""" + +from __future__ import annotations + +import contextlib +import contextvars +import logging +import pathlib +from typing import Iterator + + +class _NotRunSpecific: + """Sentinel for `not_run_specific()`.""" + + +_NOT_RUN_SPECIFIC = _NotRunSpecific() + + +_run_id: contextvars.ContextVar[str | _NotRunSpecific | None] = contextvars.ContextVar( + "_run_id", + default=None, +) + +_logger = logging.getLogger("wandb") + + +def configure_wandb_logger() -> None: + """Configures the global 'wandb' logger. + + The wandb logger is not intended to be customized by users. Instead, it is + used as a mechanism to redirect log messages into wandb run-specific log + files. + + This function is idempotent: calling it multiple times has the same effect. + """ + # Send all DEBUG and above messages to registered handlers. + # + # Per-run handlers can set different levels. + _logger.setLevel(logging.DEBUG) + + # Do not propagate wandb logs to the root logger, which the user may have + # configured to point elsewhere. All wandb log messages should go to a run's + # log file. + _logger.propagate = False + + # If no handlers are configured for the 'wandb' logger, don't activate the + # "lastResort" handler which sends messages to stderr with a level of + # WARNING by default. + # + # This occurs in wandb code that runs outside the context of a Run and + # not as part of the CLI. + # + # Most such code uses the `termlog` / `termwarn` / `termerror` methods + # to communicate with the user. When that code executes while a run is + # active, its logger messages go to that run's log file. + if not _logger.handlers: + _logger.addHandler(logging.NullHandler()) + + +@contextlib.contextmanager +def log_to_run(run_id: str | None) -> Iterator[None]: + """Direct all wandb log messages to the given run. + + Args: + id: The current run ID, or None if actions in the context manager are + not associated to a specific run. In the latter case, log messages + will go to all runs. + + Usage: + + with wb_logging.run_id(...): + ... # Log messages here go to the specified run's logger. + """ + token = _run_id.set(run_id) + try: + yield + finally: + _run_id.reset(token) + + +@contextlib.contextmanager +def log_to_all_runs() -> Iterator[None]: + """Direct wandb log messages to all runs. + + Unlike `log_to_run(None)`, this indicates an intentional choice. + This is often convenient to use as a decorator: + + @wb_logging.log_to_all_runs() + def my_func(): + ... # Log messages here go to the specified run's logger. + """ + token = _run_id.set(_NOT_RUN_SPECIFIC) + try: + yield + finally: + _run_id.reset(token) + + +def add_file_handler(run_id: str, filepath: pathlib.Path) -> logging.Handler: + """Direct log messages for a run to a file. + + Args: + run_id: The run for which to create a log file. + filepath: The file to write log messages to. + + Returns: + The added handler which can then be configured further or removed + from the 'wandb' logger directly. + + The default logging level is INFO. + """ + handler = logging.FileHandler(filepath) + handler.setLevel(logging.INFO) + handler.addFilter(_RunIDFilter(run_id)) + handler.setFormatter( + logging.Formatter( + "%(asctime)s %(levelname)-7s %(threadName)-10s:%(process)d" + " [%(filename)s:%(funcName)s():%(lineno)s]%(run_id_tag)s" + " %(message)s" + ) + ) + + _logger.addHandler(handler) + return handler + + +class _RunIDFilter(logging.Filter): + """Filters out messages logged for a different run.""" + + def __init__(self, run_id: str) -> None: + """Create a _RunIDFilter. + + Args: + run_id: Allows messages when the run ID is this or None. + """ + self._run_id = run_id + + def filter(self, record: logging.LogRecord) -> bool: + run_id = _run_id.get() + + if run_id is None: + record.run_id_tag = " [no run ID]" + return True + elif isinstance(run_id, _NotRunSpecific): + record.run_id_tag = " [all runs]" + return True + else: + record.run_id_tag = "" + return run_id == self._run_id diff --git a/lib/python3.12/site-packages/wandb/sdk/mailbox/__init__.py b/lib/python3.12/site-packages/wandb/sdk/mailbox/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e89b3d0e3b8d7124ac97264cfa66c5330a2a39a0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/mailbox/__init__.py @@ -0,0 +1,23 @@ +"""A message protocol for the internal service process. + +The core of W&B is implemented by a side process that asynchronously uploads +data. The client process (such as this Python code) sends requests to the +service, and for some requests, the service eventually sends a response. + +The client can send multiple requests before the service provides a response. +The Mailbox handles matching responses to requests. An internal thread +continuously reads data from the service and passes it to the mailbox. +""" + +from .mailbox import Mailbox, MailboxClosedError +from .mailbox_handle import HandleAbandonedError, MailboxHandle +from .wait_with_progress import wait_all_with_progress, wait_with_progress + +__all__ = [ + "Mailbox", + "MailboxClosedError", + "HandleAbandonedError", + "MailboxHandle", + "wait_all_with_progress", + "wait_with_progress", +] diff --git a/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5362aa652098b5eda5c14bfa293336441b01be4e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/mailbox.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/mailbox.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..66d6919ddfc58bfa25033255c766fd1af4197847 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/mailbox.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/mailbox_handle.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/mailbox_handle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9fc28d7486a2da8966e65d2f1c8d157d5c830a5c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/mailbox_handle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/response_handle.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/response_handle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2994f6f32ad42639cd5f340a8b99e61f373e6c4d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/response_handle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/wait_with_progress.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/wait_with_progress.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7db46663cd9b7ec615fddc5734bcd51c99443cc3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/mailbox/__pycache__/wait_with_progress.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/mailbox/mailbox.py b/lib/python3.12/site-packages/wandb/sdk/mailbox/mailbox.py new file mode 100644 index 0000000000000000000000000000000000000000..1f1851ed61d927b35dbd984437c0335bc9f35d39 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/mailbox/mailbox.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +import logging +import secrets +import string +import threading + +from wandb.proto import wandb_internal_pb2 as pb +from wandb.proto import wandb_server_pb2 as spb + +from .mailbox_handle import MailboxHandle +from .response_handle import MailboxResponseHandle + +_logger = logging.getLogger(__name__) + + +class MailboxClosedError(Exception): + """The mailbox has been closed and cannot be used.""" + + +class Mailbox: + """Matches service responses to requests. + + The mailbox can set an address on a server request and create a handle for + waiting for a response to that record. Responses are delivered by calling + `deliver()`. The `close()` method abandons all handles in case the + service process becomes unreachable. + """ + + def __init__(self) -> None: + self._handles: dict[str, MailboxResponseHandle] = {} + self._handles_lock = threading.Lock() + self._closed = False + + def require_response( + self, + request: spb.ServerRequest | pb.Record, + ) -> MailboxHandle[spb.ServerResponse]: + """Set a response address on a request. + + Args: + request: The request on which to set a request ID or mailbox slot. + This is mutated. An address must not already be set. + + Returns: + A handle for waiting for the response to the request. + + Raises: + MailboxClosedError: If the mailbox has been closed, in which case + no new responses are expected to be delivered and new handles + cannot be created. + """ + if isinstance(request, spb.ServerRequest): + if address := request.request_id: + raise ValueError(f"Request already has an address ({address})") + + address = self._new_address() + request.request_id = address + else: + if address := request.control.mailbox_slot: + raise ValueError(f"Request already has an address ({address})") + + address = self._new_address() + request.control.mailbox_slot = address + + with self._handles_lock: + if self._closed: + raise MailboxClosedError() + + handle = MailboxResponseHandle(address) + self._handles[address] = handle + + return handle + + def _new_address(self) -> str: + """Returns an unused address for a request. + + Assumes `_handles_lock` is held. + """ + + def generate(): + return "".join( + secrets.choice(string.ascii_lowercase + string.digits) + for i in range(12) + ) + + address = generate() + + # Being extra cautious. This loop will almost never be entered. + while address in self._handles: + address = generate() + + return address + + def deliver(self, response: spb.ServerResponse) -> None: + """Deliver a response from the service. + + If the response address is invalid, this does nothing. + It is a no-op if the mailbox has been closed. + """ + address = response.request_id + if not address: + kind: str | None = response.WhichOneof("server_response_type") + if kind == "result_communicate": + result_type = response.result_communicate.WhichOneof("result_type") + kind = f"result_communicate.{result_type}" + + _logger.error(f"Received response with no mailbox slot: {kind}") + return + + with self._handles_lock: + # NOTE: If the mailbox is closed, this returns None because + # we clear the dict. + handle = self._handles.pop(address, None) + + # It is not an error if there is no handle for the address: + # handles can be abandoned if the result is no longer needed. + if handle: + handle.deliver(response) + + def close(self) -> None: + """Indicate no further responses will be delivered. + + Abandons all handles. + """ + with self._handles_lock: + self._closed = True + + _logger.info( + f"Closing mailbox, abandoning {len(self._handles)} handles.", + ) + + for handle in self._handles.values(): + handle.abandon() + self._handles.clear() diff --git a/lib/python3.12/site-packages/wandb/sdk/mailbox/mailbox_handle.py b/lib/python3.12/site-packages/wandb/sdk/mailbox/mailbox_handle.py new file mode 100644 index 0000000000000000000000000000000000000000..eda3df5b557f1656dce625c9ccb6382aa11f2de8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/mailbox/mailbox_handle.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import abc +from typing import TYPE_CHECKING, Callable, Generic + +from typing_extensions import TypeVar, override + +# Necessary to break an import loop. +if TYPE_CHECKING: + from wandb.sdk.interface import interface + + +_T = TypeVar("_T") +_S = TypeVar("_S") + + +class HandleAbandonedError(Exception): + """The handle has no response and has been abandoned.""" + + +class MailboxHandle(abc.ABC, Generic[_T]): + """A thread-safe handle that allows waiting for a response to a request.""" + + def map(self, fn: Callable[[_T], _S]) -> MailboxHandle[_S]: + """Returns a transformed handle. + + Methods on the returned handle call methods on this handle, but the + response type is derived using the given function. + + Args: + fn: A function to apply to this handle's result to get the new + handle's result. The function should be pure and fast. + """ + return _MailboxMappedHandle(self, fn) + + @abc.abstractmethod + def abandon(self) -> None: + """Abandon the handle, indicating it will not receive a response.""" + + @abc.abstractmethod + def cancel(self, iface: interface.InterfaceBase) -> None: + """Cancel the handle, requesting any associated work to not complete. + + This automatically abandons the handle, as a response is no longer + guaranteed. + + Args: + iface: The interface on which to publish the cancel request. + """ + + @abc.abstractmethod + def check(self) -> _T | None: + """Returns the response if it's ready.""" + + @abc.abstractmethod + def wait_or(self, *, timeout: float | None) -> _T: + """Wait for a response or a timeout. + + Args: + timeout: A finite number of seconds or None to never time out. + If less than or equal to zero, times out immediately unless + the response is available. + + Returns: + The response if it arrives before the timeout or has already arrived. + + Raises: + TimeoutError: If the timeout is reached. + HandleAbandonedError: If the handle becomes abandoned. + """ + + @abc.abstractmethod + async def wait_async(self, *, timeout: float | None) -> _T: + """Wait for a response or timeout. + + This must run in an `asyncio` event loop. + + Args: + timeout: A finite number of seconds or None to never time out. + + Returns: + The response if it arrives before the timeout or has already arrived. + + Raises: + TimeoutError: If the timeout is reached. + HandleAbandonedError: If the handle becomes abandoned. + """ + + +class _MailboxMappedHandle(Generic[_S], MailboxHandle[_S]): + """A mailbox handle whose result is derived from another handle.""" + + def __init__( + self, + handle: MailboxHandle[_T], + fn: Callable[[_T], _S], + ) -> None: + self._handle = handle + self._fn = fn + + @override + def abandon(self) -> None: + self._handle.abandon() + + @override + def cancel(self, iface: interface.InterfaceBase) -> None: + self._handle.cancel(iface) + + @override + def check(self) -> _S | None: + if response := self._handle.check(): + return self._fn(response) + else: + return None + + @override + def wait_or(self, *, timeout: float | None) -> _S: + return self._fn(self._handle.wait_or(timeout=timeout)) + + @override + async def wait_async(self, *, timeout: float | None) -> _S: + response = await self._handle.wait_async(timeout=timeout) + return self._fn(response) diff --git a/lib/python3.12/site-packages/wandb/sdk/mailbox/response_handle.py b/lib/python3.12/site-packages/wandb/sdk/mailbox/response_handle.py new file mode 100644 index 0000000000000000000000000000000000000000..5680b2e01b747c35f4def7fbb81ccab5315e1ac2 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/mailbox/response_handle.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +import asyncio +import math +import threading +from typing import TYPE_CHECKING + +from typing_extensions import override + +from wandb.proto import wandb_server_pb2 as spb + +from .mailbox_handle import HandleAbandonedError, MailboxHandle + +# Necessary to break an import loop. +if TYPE_CHECKING: + from wandb.sdk.interface import interface + + +class MailboxResponseHandle(MailboxHandle[spb.ServerResponse]): + """A general handle for any ServerResponse.""" + + def __init__(self, address: str) -> None: + self._address = address + self._lock = threading.Lock() + self._event = threading.Event() + + self._abandoned = False + self._response: spb.ServerResponse | None = None + + self._asyncio_events: dict[asyncio.Event, _AsyncioEvent] = dict() + + def deliver(self, response: spb.ServerResponse) -> None: + """Deliver the response. + + This may only be called once. It is an error to respond to the same + request more than once. It is a no-op if the handle has been abandoned. + """ + with self._lock: + if self._abandoned: + return + + if self._response: + raise ValueError( + f"A response has already been delivered to {self._address}." + ) + + self._response = response + self._signal_done() + + @override + def cancel(self, iface: interface.InterfaceBase) -> None: + iface.publish_cancel(self._address) + self.abandon() + + @override + def abandon(self) -> None: + with self._lock: + self._abandoned = True + self._signal_done() + + def _signal_done(self) -> None: + """Indicate that the handle either got a response or became abandoned. + + The lock must be held. + """ + # Unblock threads blocked on `wait_or`. + self._event.set() + + # Unblock asyncio loops blocked on `wait_async`. + for asyncio_event in self._asyncio_events.values(): + asyncio_event.set_threadsafe() + self._asyncio_events.clear() + + @override + def check(self) -> spb.ServerResponse | None: + with self._lock: + return self._response + + @override + def wait_or(self, *, timeout: float | None) -> spb.ServerResponse: + if timeout is not None and not math.isfinite(timeout): + raise ValueError("Timeout must be finite or None.") + + if not self._event.wait(timeout=timeout): + raise TimeoutError( + f"Timed out waiting for response on {self._address}", + ) + + with self._lock: + if self._response: + return self._response + + assert self._abandoned + raise HandleAbandonedError() + + @override + async def wait_async(self, *, timeout: float | None) -> spb.ServerResponse: + if timeout is not None and not math.isfinite(timeout): + raise ValueError("Timeout must be finite or None.") + + evt = asyncio.Event() + self._add_asyncio_event(asyncio.get_event_loop(), evt) + + try: + await asyncio.wait_for(evt.wait(), timeout=timeout) + + except (asyncio.TimeoutError, TimeoutError) as e: + with self._lock: + if self._response: + return self._response + elif self._abandoned: + raise HandleAbandonedError() + else: + raise TimeoutError( + f"Timed out waiting for response on {self._address}" + ) from e + + else: + with self._lock: + if self._response: + return self._response + + assert self._abandoned + raise HandleAbandonedError() + + finally: + self._forget_asyncio_event(evt) + + def _add_asyncio_event( + self, + loop: asyncio.AbstractEventLoop, + event: asyncio.Event, + ) -> None: + """Add an event to signal when a response is received. + + If a response already exists, this notifies the event loop immediately. + """ + asyncio_event = _AsyncioEvent(loop, event) + + with self._lock: + if self._response or self._abandoned: + asyncio_event.set_threadsafe() + else: + self._asyncio_events[event] = asyncio_event + + def _forget_asyncio_event(self, event: asyncio.Event) -> None: + """Cancel signalling an event when a response is received.""" + with self._lock: + self._asyncio_events.pop(event, None) + + +class _AsyncioEvent: + def __init__( + self, + loop: asyncio.AbstractEventLoop, + event: asyncio.Event, + ): + self._loop = loop + self._event = event + + def set_threadsafe(self) -> None: + """Set the asyncio event in its own loop.""" + self._loop.call_soon_threadsafe(self._event.set) diff --git a/lib/python3.12/site-packages/wandb/sdk/mailbox/wait_with_progress.py b/lib/python3.12/site-packages/wandb/sdk/mailbox/wait_with_progress.py new file mode 100644 index 0000000000000000000000000000000000000000..dda74b030143b2514ead85c36e37edfe5359bdb9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/mailbox/wait_with_progress.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +import time +from typing import Any, Callable, Coroutine, List, TypeVar, cast + +from wandb.sdk.lib import asyncio_compat + +from .mailbox_handle import MailboxHandle + +_T = TypeVar("_T") + + +def wait_with_progress( + handle: MailboxHandle[_T], + *, + timeout: float | None, + progress_after: float, + display_progress: Callable[[], Coroutine[Any, Any, None]], +) -> _T: + """Wait for a handle, possibly displaying progress to the user. + + Equivalent to passing a single handle to `wait_all_with_progress`. + """ + return wait_all_with_progress( + [handle], + timeout=timeout, + progress_after=progress_after, + display_progress=display_progress, + )[0] + + +def wait_all_with_progress( + handle_list: list[MailboxHandle[_T]], + *, + timeout: float | None, + progress_after: float, + display_progress: Callable[[], Coroutine[Any, Any, None]], +) -> list[_T]: + """Wait for multiple handles, possibly displaying progress to the user. + + Args: + handle_list: The handles to wait for. + timeout: A number of seconds after which to raise a TimeoutError, + or None if this should never timeout. + progress_after: A number of seconds after which to start the + display_progress callback. Starting the callback creates a thread + and starts an asyncio loop, so we want to avoid doing it if + the handle is resolved quickly. + display_progress: An asyncio function that displays progress to + the user. This function is executed on a new thread and cancelled + if the timeout is exceeded. + + Returns: + A list where the Nth item is the Nth handle's result. + + Raises: + TimeoutError: If the overall timeout expires. + HandleAbandonedError: If any handle becomes abandoned. + Exception: Any exception from the display function is propagated. + """ + if not handle_list: + return [] + + if timeout is not None and timeout <= progress_after: + return _wait_handles(handle_list, timeout=timeout) + + start_time = time.monotonic() + + try: + return _wait_handles(handle_list, timeout=progress_after) + except TimeoutError: + pass + + async def progress_loop_with_timeout() -> list[_T]: + with asyncio_compat.cancel_on_exit(display_progress()): + if timeout is not None: + elapsed_time = time.monotonic() - start_time + remaining_timeout = timeout - elapsed_time + else: + remaining_timeout = None + + return await _wait_handles_async( + handle_list, + timeout=remaining_timeout, + ) + + return asyncio_compat.run(progress_loop_with_timeout) + + +def _wait_handles( + handle_list: list[MailboxHandle[_T]], + *, + timeout: float, +) -> list[_T]: + """Wait for multiple mailbox handles. + + Returns: + Each handle's result, in the same order as the given handles. + + Raises: + TimeoutError: If the overall timeout expires. + HandleAbandonedError: If any handle becomes abandoned. + """ + results: list[_T] = [] + + start_time = time.monotonic() + for handle in handle_list: + elapsed_time = time.monotonic() - start_time + remaining_timeout = timeout - elapsed_time + results.append(handle.wait_or(timeout=remaining_timeout)) + + return results + + +async def _wait_handles_async( + handle_list: list[MailboxHandle[_T]], + *, + timeout: float | None, +) -> list[_T]: + """Asynchronously wait for multiple mailbox handles. + + Just like _wait_handles. + """ + results: list[_T | None] = [None for _ in handle_list] + + async def wait_single(index: int) -> None: + handle = handle_list[index] + results[index] = await handle.wait_async(timeout=timeout) + + async with asyncio_compat.open_task_group() as task_group: + for index in range(len(handle_list)): + task_group.start_soon(wait_single(index)) + + # NOTE: `list` is not subscriptable until Python 3.10, so we use List. + return cast(List[_T], results) diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__init__.py b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b8b57cecbaa97634b0f6beb8ccd1a778201c19d3 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__init__.py @@ -0,0 +1,47 @@ +# Generated by ariadne-codegen + +from .delete_project import DeleteProject, DeleteProjectDeleteModel +from .fetch_registry import FetchRegistry, FetchRegistryEntity +from .fragments import ( + RegistryFragment, + RegistryFragmentArtifactTypes, + RegistryFragmentArtifactTypesEdges, + RegistryFragmentArtifactTypesEdgesNode, +) +from .input_types import ArtifactTypeInput +from .operations import ( + DELETE_PROJECT_GQL, + FETCH_REGISTRY_GQL, + RENAME_PROJECT_GQL, + UPSERT_REGISTRY_PROJECT_GQL, +) +from .rename_project import ( + RenameProject, + RenameProjectRenameProject, + RenameProjectRenameProjectProject, +) +from .upsert_registry_project import ( + UpsertRegistryProject, + UpsertRegistryProjectUpsertModel, +) + +__all__ = [ + "DELETE_PROJECT_GQL", + "FETCH_REGISTRY_GQL", + "RENAME_PROJECT_GQL", + "UPSERT_REGISTRY_PROJECT_GQL", + "FetchRegistry", + "FetchRegistryEntity", + "RenameProject", + "RenameProjectRenameProject", + "RenameProjectRenameProjectProject", + "UpsertRegistryProject", + "UpsertRegistryProjectUpsertModel", + "DeleteProject", + "DeleteProjectDeleteModel", + "ArtifactTypeInput", + "RegistryFragment", + "RegistryFragmentArtifactTypes", + "RegistryFragmentArtifactTypesEdges", + "RegistryFragmentArtifactTypesEdgesNode", +] diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6cfe91aae7af71fe316914bcb86679194328235c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/delete_project.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/delete_project.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ab548335fff01b4022744dc0b1185de61e42f41 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/delete_project.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/enums.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/enums.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..88f28b59c909f2eac6b9727dd581cc8c91fc0d75 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/enums.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/fetch_registry.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/fetch_registry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..01ce1bc4412e70379b595ffd65b59c9cf13e433d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/fetch_registry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/fragments.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/fragments.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f0dbab35b9d5d9267198a5418fd23181af1703e7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/fragments.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/input_types.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/input_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f0964eedcf81177209bc9eb26933a87246e97a6c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/input_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/operations.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/operations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2076fff889020d3a9e9f4b2f677ff14467a8633c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/operations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/rename_project.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/rename_project.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d29da072c8506b22c6009b6a424a6d4ab79a7db0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/rename_project.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/upsert_registry_project.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/upsert_registry_project.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c951b308382f06b41550d4e0676037bd35074af Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/__pycache__/upsert_registry_project.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/delete_project.py b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/delete_project.py new file mode 100644 index 0000000000000000000000000000000000000000..d1d26b3ba7561ff8426c6ddab20527d2295f0e7f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/delete_project.py @@ -0,0 +1,22 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/projects/ + +from __future__ import annotations + +from typing import Literal, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, Typename + + +class DeleteProject(GQLBase): + delete_model: Optional[DeleteProjectDeleteModel] = Field(alias="deleteModel") + + +class DeleteProjectDeleteModel(GQLBase): + success: Optional[bool] + typename__: Typename[Literal["DeleteModelPayload"]] + + +DeleteProject.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/enums.py b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/enums.py new file mode 100644 index 0000000000000000000000000000000000000000..cc7d61d95669bab416e88c3acc378913704e64e8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/enums.py @@ -0,0 +1,4 @@ +# Generated by ariadne-codegen +# Source: core/api/graphql/schemas/schema-latest.graphql + +from __future__ import annotations diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/fetch_registry.py b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/fetch_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..5b5f4e8ccec68df79519b5a610760c2af340a0b9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/fetch_registry.py @@ -0,0 +1,22 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/projects/ + +from __future__ import annotations + +from typing import Optional + +from wandb._pydantic import GQLBase + +from .fragments import RegistryFragment + + +class FetchRegistry(GQLBase): + entity: Optional[FetchRegistryEntity] + + +class FetchRegistryEntity(GQLBase): + project: Optional[RegistryFragment] + + +FetchRegistry.model_rebuild() +FetchRegistryEntity.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/fragments.py b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/fragments.py new file mode 100644 index 0000000000000000000000000000000000000000..69f780b4e2c5b1333b35b217dee1767f64ac0c1d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/fragments.py @@ -0,0 +1,41 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/projects/ + +from __future__ import annotations + +from typing import List, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, GQLId + + +class RegistryFragment(GQLBase): + id: GQLId + allow_all_artifact_types_in_registry: bool = Field( + alias="allowAllArtifactTypesInRegistry" + ) + artifact_types: RegistryFragmentArtifactTypes = Field(alias="artifactTypes") + name: str + description: Optional[str] + created_at: str = Field(alias="createdAt") + updated_at: Optional[str] = Field(alias="updatedAt") + access: Optional[str] + + +class RegistryFragmentArtifactTypes(GQLBase): + edges: List[RegistryFragmentArtifactTypesEdges] + + +class RegistryFragmentArtifactTypesEdges(GQLBase): + node: Optional[RegistryFragmentArtifactTypesEdgesNode] + + +class RegistryFragmentArtifactTypesEdgesNode(GQLBase): + name: str + + +RegistryFragment.model_rebuild() +RegistryFragmentArtifactTypes.model_rebuild() +RegistryFragmentArtifactTypesEdges.model_rebuild() +RegistryFragmentArtifactTypesEdgesNode.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/input_types.py b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/input_types.py new file mode 100644 index 0000000000000000000000000000000000000000..ce360f5e155482771a8f7b0a98c1e015510df73d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/input_types.py @@ -0,0 +1,13 @@ +# Generated by ariadne-codegen +# Source: core/api/graphql/schemas/schema-latest.graphql + +from __future__ import annotations + +from typing import Optional + +from wandb._pydantic import GQLBase + + +class ArtifactTypeInput(GQLBase): + name: str + description: Optional[str] = None diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/operations.py b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/operations.py new file mode 100644 index 0000000000000000000000000000000000000000..3598c6d664da52611c490c10cd64b326814f610c --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/operations.py @@ -0,0 +1,88 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/projects/ + +__all__ = [ + "DELETE_PROJECT_GQL", + "FETCH_REGISTRY_GQL", + "RENAME_PROJECT_GQL", + "UPSERT_REGISTRY_PROJECT_GQL", +] + +FETCH_REGISTRY_GQL = """ +query FetchRegistry($name: String, $entityName: String) { + entity(name: $entityName) { + project(name: $name) { + ...RegistryFragment + } + } +} + +fragment RegistryFragment on Project { + id + allowAllArtifactTypesInRegistry + artifactTypes(includeAll: true) { + edges { + node { + name + } + } + } + name + description + createdAt + updatedAt + access +} +""" + +RENAME_PROJECT_GQL = """ +mutation renameProject($entityName: String!, $oldProjectName: String!, $newProjectName: String!) { + renameProject( + input: {entityName: $entityName, oldProjectName: $oldProjectName, newProjectName: $newProjectName} + ) { + project { + name + } + inserted + } +} +""" + +UPSERT_REGISTRY_PROJECT_GQL = """ +mutation UpsertRegistryProject($description: String, $entityName: String, $name: String, $access: String, $allowAllArtifactTypesInRegistry: Boolean, $artifactTypes: [ArtifactTypeInput!]) { + upsertModel( + input: {description: $description, entityName: $entityName, name: $name, access: $access, allowAllArtifactTypesInRegistry: $allowAllArtifactTypesInRegistry, artifactTypes: $artifactTypes} + ) { + project { + ...RegistryFragment + } + inserted + } +} + +fragment RegistryFragment on Project { + id + allowAllArtifactTypesInRegistry + artifactTypes(includeAll: true) { + edges { + node { + name + } + } + } + name + description + createdAt + updatedAt + access +} +""" + +DELETE_PROJECT_GQL = """ +mutation deleteProject($id: String!) { + deleteModel(input: {id: $id}) { + success + __typename + } +} +""" diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/rename_project.py b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/rename_project.py new file mode 100644 index 0000000000000000000000000000000000000000..34648ebda7698c62f4f1865b5cb578b04505bf32 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/rename_project.py @@ -0,0 +1,27 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/projects/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + + +class RenameProject(GQLBase): + rename_project: Optional[RenameProjectRenameProject] = Field(alias="renameProject") + + +class RenameProjectRenameProject(GQLBase): + project: Optional[RenameProjectRenameProjectProject] + inserted: Optional[bool] + + +class RenameProjectRenameProjectProject(GQLBase): + name: str + + +RenameProject.model_rebuild() +RenameProjectRenameProject.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/verify/__init__.py b/lib/python3.12/site-packages/wandb/sdk/verify/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/verify/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/verify/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a6c0e44eb2b7a7fdb80d47c240f6db73a093af15 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/verify/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/verify/__pycache__/verify.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/verify/__pycache__/verify.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4ccfe2eaf9c570515724d0b6f7c9357ef0f49b8c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/verify/__pycache__/verify.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/verify/verify.py b/lib/python3.12/site-packages/wandb/sdk/verify/verify.py new file mode 100644 index 0000000000000000000000000000000000000000..8cadbf310a9b5c273011fee2c66c78f5c18ed6ea --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/verify/verify.py @@ -0,0 +1,555 @@ +"""Utilities for wandb verify.""" + +import contextlib +import getpass +import io +import os +import time +from functools import partial +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import click +import requests +from wandb_gql import gql + +import wandb +from wandb.sdk.artifacts.artifact import Artifact +from wandb.sdk.lib import runid + +from ...apis.internal import Api + +PROJECT_NAME = "verify" +GET_RUN_MAX_TIME = 10 +MIN_RETRYS = 3 +CHECKMARK = "\u2705" +RED_X = "\u274c" +ID_PREFIX = runid.generate_id() + + +def nice_id(name): + return ID_PREFIX + "-" + name + + +def print_results( + failed_test_or_tests: Optional[Union[str, List[str]]], warning: bool +) -> None: + if warning: + color = "yellow" + else: + color = "red" + if isinstance(failed_test_or_tests, str): + print(RED_X) # noqa: T201 + print(click.style(failed_test_or_tests, fg=color, bold=True)) # noqa: T201 + elif isinstance(failed_test_or_tests, list) and len(failed_test_or_tests) > 0: + print(RED_X) # noqa: T201 + print( # noqa: T201 + "\n".join( + [click.style(f, fg=color, bold=True) for f in failed_test_or_tests] + ) + ) + else: + print(CHECKMARK) # noqa: T201 + + +def check_host(host: str) -> bool: + if host in ("api.wandb.ai", "http://api.wandb.ai", "https://api.wandb.ai"): + print_results("Cannot run wandb verify against api.wandb.ai", False) + return False + return True + + +def check_logged_in(api: Api, host: str) -> bool: + print("Checking if logged in".ljust(72, "."), end="") # noqa: T201 + login_doc_url = "https://docs.wandb.ai/ref/cli/wandb-login" + fail_string = None + if api.api_key is None: + fail_string = ( + "Not logged in. Please log in using `wandb login`. See the docs: {}".format( + click.style(login_doc_url, underline=True, fg="blue") + ) + ) + # check that api key is correct + # TODO: Better check for api key is correct + else: + res = api.api.viewer() + if not res: + fail_string = ( + "Could not get viewer with default API key. " + f"Please relogin using `WANDB_BASE_URL={host} wandb login --relogin` and try again" + ) + + print_results(fail_string, False) + return fail_string is None + + +def check_secure_requests(url: str, test_url_string: str, failure_output: str) -> None: + # check if request is over https + print(test_url_string.ljust(72, "."), end="") # noqa: T201 + fail_string = None + if not url.startswith("https"): + fail_string = failure_output + print_results(fail_string, True) + + +def check_cors_configuration(url: str, origin: str) -> None: + print("Checking CORs configuration of the bucket".ljust(72, "."), end="") # noqa: T201 + fail_string = None + res_get = requests.options( + url, headers={"Origin": origin, "Access-Control-Request-Method": "GET"} + ) + + if res_get.headers.get("Access-Control-Allow-Origin") is None: + fail_string = ( + "Your object store does not have a valid CORs configuration, " + f"you must allow GET and PUT to Origin: {origin}" + ) + + print_results(fail_string, True) + + +def check_run(api: Api) -> bool: + print( # noqa: T201 + "Checking logged metrics, saving and downloading a file".ljust(72, "."), end="" + ) + failed_test_strings = [] + + # set up config + n_epochs = 4 + string_test = "A test config" + dict_test = {"config_val": 2, "config_string": "config string"} + list_test = [0, "one", "2"] + config = { + "epochs": n_epochs, + "stringTest": string_test, + "dictTest": dict_test, + "listTest": list_test, + } + # create a file to save + filepath = "./test with_special-characters.txt" + f = open(filepath, "w") + f.write("test") + f.close() + + with wandb.init( + id=nice_id("check_run"), + reinit=True, + config=config, + project=PROJECT_NAME, + ) as run: + run_id = run.id + entity = run.entity + logged = True + try: + for i in range(1, 11): + run.log({"loss": 1.0 / i}, step=i) + log_dict = {"val1": 1.0, "val2": 2} + run.log({"dict": log_dict}, step=i + 1) + except Exception: + logged = False + failed_test_strings.append( + "Failed to log values to run. Contact W&B for support." + ) + + try: + run.log({"HT%3ML ": wandb.Html('Link')}) + except Exception: + failed_test_strings.append( + "Failed to log to media. Contact W&B for support." + ) + + run.save(filepath) + public_api = wandb.Api() + prev_run = public_api.run(f"{entity}/{PROJECT_NAME}/{run_id}") + # raise Exception(prev_run.__dict__) + if prev_run is None: + failed_test_strings.append( + "Failed to access run through API. Contact W&B for support." + ) + print_results(failed_test_strings, False) + return False + for key, value in config.items(): + if prev_run.config.get(key) != value: + failed_test_strings.append( + "Read config values don't match run config. Contact W&B for support." + ) + break + if logged and ( + prev_run.history_keys["keys"]["loss"]["previousValue"] != 0.1 + or prev_run.history_keys["lastStep"] != 11 + or prev_run.history_keys["keys"]["dict.val1"]["previousValue"] != 1.0 + or prev_run.history_keys["keys"]["dict.val2"]["previousValue"] != 2 + ): + failed_test_strings.append( + "History metrics don't match logged values. Check database encoding." + ) + + if logged and prev_run.summary["loss"] != 1.0 / 10: + failed_test_strings.append( + "Read summary values don't match expected value. Check database encoding, or contact W&B for support." + ) + # TODO: (kdg) refactor this so it doesn't rely on an exception handler + try: + read_file = retry_fn(partial(prev_run.file, filepath)) + # There's a race where the file hasn't been processed in the queue, + # we just retry until we get a download + read_file = retry_fn(partial(read_file.download, replace=True)) + except Exception: + failed_test_strings.append( + "Unable to download file. Check SQS configuration, topic configuration and bucket permissions." + ) + + print_results(failed_test_strings, False) + return False + contents = read_file.read() + if contents != "test": + failed_test_strings.append( + "Contents of downloaded file do not match uploaded contents. Contact W&B for support." + ) + print_results(failed_test_strings, False) + return len(failed_test_strings) == 0 + + +def verify_manifest( + downloaded_manifest: Dict[str, Any], + computed_manifest: Dict[str, Any], + fails_list: List[str], +) -> None: + try: + for key in computed_manifest.keys(): + assert ( + computed_manifest[key]["digest"] == downloaded_manifest[key]["digest"] + ) + assert computed_manifest[key]["size"] == downloaded_manifest[key]["size"] + except AssertionError: + fails_list.append( + "Artifact manifest does not appear as expected. Contact W&B for support." + ) + + +def verify_digest( + downloaded: "Artifact", computed: "Artifact", fails_list: List[str] +) -> None: + if downloaded.digest != computed.digest: + fails_list.append( + "Artifact digest does not appear as expected. Contact W&B for support." + ) + + +def artifact_with_path_or_paths( + name: str, verify_dir: Optional[str] = None, singular: bool = False +) -> "Artifact": + art = wandb.Artifact(type="artsy", name=name) + # internal file + with open("verify_int_test.txt", "w") as f: + f.write("test 1") + f.close() + art.add_file(f.name) + if singular: + return art + if verify_dir is None: + verify_dir = "./" + with art.new_file("verify_a.txt") as f: + f.write("test 2") + if not os.path.exists(verify_dir): + os.makedirs(verify_dir) + with open(f"{verify_dir}/verify_1.txt", "w") as f: + f.write("1") + art.add_dir(verify_dir) + file3 = Path(verify_dir) / "verify_3.txt" + file3.write_text("3") + + # reference to local file + art.add_reference(file3.resolve().as_uri()) + + return art + + +def log_use_download_artifact( + artifact: "Artifact", + alias: str, + name: str, + download_dir: str, + failed_test_strings: List[str], + add_extra_file: bool, +) -> Tuple[bool, Optional["Artifact"], List[str]]: + with wandb.init( + id=nice_id("log_artifact"), + reinit=True, + project=PROJECT_NAME, + config={"test": "artifact log"}, + ) as log_art_run: + if add_extra_file: + with open("verify_2.txt", "w") as f: + f.write("2") + f.close() + artifact.add_file(f.name) + + try: + log_art_run.log_artifact(artifact, aliases=alias) + except Exception as e: + failed_test_strings.append(f"Unable to log artifact. {e}") + return False, None, failed_test_strings + + with wandb.init( + id=nice_id("use_artifact"), + project=PROJECT_NAME, + config={"test": "artifact use"}, + ) as use_art_run: + try: + used_art = use_art_run.use_artifact(f"{name}:{alias}") + except Exception as e: + failed_test_strings.append(f"Unable to use artifact. {e}") + return False, None, failed_test_strings + try: + used_art.download(root=download_dir) + except Exception: + failed_test_strings.append( + "Unable to download artifact. Check bucket permissions." + ) + return False, None, failed_test_strings + + return True, used_art, failed_test_strings + + +def check_artifacts() -> bool: + print("Checking artifact save and download workflows".ljust(72, "."), end="") # noqa: T201 + failed_test_strings: List[str] = [] + + # test checksum + sing_art_dir = "./verify_sing_art" + alias = "sing_art1" + name = nice_id("sing-artys") + singular_art = artifact_with_path_or_paths(name, singular=True) + cont_test, download_artifact, failed_test_strings = log_use_download_artifact( + singular_art, alias, name, sing_art_dir, failed_test_strings, False + ) + if not cont_test or download_artifact is None: + print_results(failed_test_strings, False) + return False + try: + download_artifact.verify(root=sing_art_dir) + except ValueError: + failed_test_strings.append( + "Artifact does not contain expected checksum. Contact W&B for support." + ) + + # test manifest and digest + multi_art_dir = "./verify_art" + alias = "art1" + name = nice_id("my-artys") + art1 = artifact_with_path_or_paths(name, "./verify_art_dir", singular=False) + cont_test, download_artifact, failed_test_strings = log_use_download_artifact( + art1, alias, name, multi_art_dir, failed_test_strings, True + ) + if not cont_test or download_artifact is None: + print_results(failed_test_strings, False) + return False + if set(os.listdir(multi_art_dir)) != { + "verify_a.txt", + "verify_2.txt", + "verify_1.txt", + "verify_3.txt", + "verify_int_test.txt", + }: + failed_test_strings.append( + "Artifact directory is missing files. Contact W&B for support." + ) + + computed = wandb.Artifact("computed", type="dataset") + computed.add_dir(multi_art_dir) + verify_digest(download_artifact, computed, failed_test_strings) + + computed_manifest = computed.manifest.to_manifest_json()["contents"] + downloaded_manifest = download_artifact.manifest.to_manifest_json()["contents"] + verify_manifest(downloaded_manifest, computed_manifest, failed_test_strings) + + print_results(failed_test_strings, False) + return len(failed_test_strings) == 0 + + +def check_graphql_put(api: Api, host: str) -> Tuple[bool, Optional[str]]: + # check graphql endpoint using an upload + print("Checking signed URL upload".ljust(72, "."), end="") # noqa: T201 + failed_test_strings = [] + gql_fp = "gql_test_file.txt" + f = open(gql_fp, "w") + f.write("test2") + f.close() + with wandb.init( + id=nice_id("graphql_put"), + reinit=True, + project=PROJECT_NAME, + config={"test": "put to graphql"}, + ) as run: + run.save(gql_fp) + public_api = wandb.Api() + prev_run = public_api.run(f"{run.entity}/{PROJECT_NAME}/{run.id}") + if prev_run is None: + failed_test_strings.append( + "Unable to access previous run through public API. Contact W&B for support." + ) + print_results(failed_test_strings, False) + return False, None + # TODO: (kdg) refactor this so it doesn't rely on an exception handler + try: + read_file = retry_fn(partial(prev_run.file, gql_fp)) + url = read_file.url + read_file = retry_fn(partial(read_file.download, replace=True)) + except Exception: + failed_test_strings.append( + "Unable to read file successfully saved through a put request. Check SQS configurations, bucket permissions and topic configs." + ) + print_results(failed_test_strings, False) + return False, None + contents = read_file.read() + try: + assert contents == "test2" + except AssertionError: + failed_test_strings.append( + "Read file contents do not match saved file contents. Contact W&B for support." + ) + + print_results(failed_test_strings, False) + return len(failed_test_strings) == 0, url + + +def check_large_post() -> bool: + print( # noqa: T201 + "Checking ability to send large payloads through proxy".ljust(72, "."), end="" + ) + descy = "a" * int(10**7) + + username = getpass.getuser() + failed_test_strings = [] + query = gql( + """ + query Project($entity: String!, $name: String!, $runName: String!, $desc: String!){ + project(entityName: $entity, name: $name) { + run(name: $runName, desc: $desc) { + name + summaryMetrics + } + } + } + """ + ) + public_api = wandb.Api() + client = public_api._base_client + + try: + client._get_result( + query, + variable_values={ + "entity": username, + "name": PROJECT_NAME, + "runName": "", + "desc": descy, + }, + timeout=60, + ) + except Exception as e: + if ( + isinstance(e, requests.HTTPError) + and e.response is not None + and e.response.status_code == 413 + ): + failed_test_strings.append( + 'Failed to send a large payload. Check nginx.ingress.kubernetes.io/proxy-body-size is "0".' + ) + else: + failed_test_strings.append( + f"Failed to send a large payload with error: {e}." + ) + print_results(failed_test_strings, False) + return len(failed_test_strings) == 0 + + +def check_wandb_version(api: Api) -> None: + print("Checking wandb package version is up to date".ljust(72, "."), end="") # noqa: T201 + _, server_info = api.viewer_server_info() + fail_string = None + warning = False + max_cli_version = server_info.get("cliVersionInfo", {}).get("max_cli_version", None) + min_cli_version = server_info.get("cliVersionInfo", {}).get( + "min_cli_version", "0.0.1" + ) + + from packaging.version import parse + + if parse(wandb.__version__) < parse(min_cli_version): + fail_string = f"wandb version out of date, please run pip install --upgrade wandb=={max_cli_version}" + elif parse(wandb.__version__) > parse(max_cli_version): + fail_string = ( + "wandb version is not supported by your local installation. This could " + "cause some issues. If you're having problems try: please run `pip " + f"install --upgrade wandb=={max_cli_version}`" + ) + warning = True + + print_results(fail_string, warning) + + +def check_sweeps(api: Api) -> bool: + print("Checking sweep creation and agent execution".ljust(72, "."), end="") # noqa: T201 + failed_test_strings: List[str] = [] + + sweep_config = { + "method": "random", + "metric": {"goal": "minimize", "name": "score"}, + "parameters": { + "x": {"values": [0.01, 0.05, 0.1]}, + "y": {"values": [1, 2, 3]}, + }, + "name": "verify_sweep", + } + + try: + with contextlib.redirect_stdout(io.StringIO()): + sweep_id = wandb.sweep( + sweep=sweep_config, project=PROJECT_NAME, entity=api.default_entity + ) + except Exception as e: + failed_test_strings.append(f"Failed to create sweep: {e}") + print_results(failed_test_strings, False) + return False + + if not sweep_id: + failed_test_strings.append("Sweep creation returned an invalid ID.") + print_results(failed_test_strings, False) + return False + + try: + + def objective(config): + score = config.x**3 + config.y + return score + + def main(): + with wandb.init(project=PROJECT_NAME) as run: + score = objective(run.config) + run.log({"score": score}) + + wandb.agent(sweep_id, function=main, count=10) + except Exception as e: + failed_test_strings.append(f"Failed to run sweep agent: {e}") + print_results(failed_test_strings, False) + return False + + print_results(failed_test_strings, False) + return len(failed_test_strings) == 0 + + +def retry_fn(fn: Callable) -> Any: + ini_time = time.time() + res = None + i = 0 + while i < MIN_RETRYS or time.time() - ini_time < GET_RUN_MAX_TIME: + i += 1 + try: + res = fn() + break + except Exception: + time.sleep(1) + continue + return res