apache-hamilton 1.90.0.dev0__py3-none-any.whl

hamilton/plugins/h_experiments/__init__.py

@@ -0,0 +1,22 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from hamilton.plugins.h_experiments.hook import ExperimentTracker
+
+__all__ = [
+    "ExperimentTracker",
+]

hamilton/plugins/h_experiments/__main__.py

@@ -0,0 +1,62 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import argparse
+import os
+from pathlib import Path
+
+from hamilton.plugins.h_experiments.cache import JsonCache
+
+
+def main():
+    try:
+        import fastapi  # noqa: F401
+        import fastui  # noqa: F401
+        import uvicorn
+    except ModuleNotFoundError as e:
+        raise ModuleNotFoundError(
+            "Some dependencies are missing. Make sure to `pip install sf-hamilton[experiments]`"
+        ) from e
+    parser = argparse.ArgumentParser(prog="hamilton-experiments")
+    parser.description = "Hamilton Experiment Server launcher"
+
+    parser.add_argument(
+        "path",
+        metavar="path",
+        type=str,
+        default="./experiments",
+        nargs="?",
+        help="Set HAMILTON_EXPERIMENTS_PATH environment variable",
+    )
+    parser.add_argument("--host", default="127.0.0.1", type=str, help="Bind to this address")
+    parser.add_argument("--port", default=8123, type=int, help="Bind to this port")
+
+    args = parser.parse_args()
+
+    try:
+        JsonCache(args.path)
+    except Exception as e:
+        raise ValueError(f"Server failed to launch. No metadata cache found at {args.path}") from e
+
+    # set environment variable that FastAPI will use
+    os.environ["HAMILTON_EXPERIMENTS_PATH"] = str(Path(args.path).resolve())
+
+    uvicorn.run("hamilton.plugins.h_experiments.server:app", host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

hamilton/plugins/h_experiments/cache.py

@@ -0,0 +1,39 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import dbm
+
+
+class JsonCache:
+    def __init__(self, cache_path: str):
+        self.cache_path = f"{cache_path}/json_cache"
+
+    def keys(self) -> list:
+        with dbm.open(self.cache_path, "r") as db:
+            return list(db.keys())
+
+    def write(self, data: str, id_: str) -> None:
+        with dbm.open(self.cache_path, "c") as db:
+            db[id_] = data
+
+    def read(self, id_: str) -> str:
+        with dbm.open(self.cache_path, "r") as db:
+            return db[id_].decode()
+
+    def delete(self, id_: str) -> None:
+        with dbm.open(self.cache_path, "r") as db:
+            del db[id_]

hamilton/plugins/h_experiments/data_model.py

@@ -0,0 +1,68 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Any
+
+from pydantic import BaseModel, create_model, model_validator
+
+
+def model_from_values(name: str, specs: dict[str, Any]) -> dict[str, tuple[type, ...]]:
+    return create_model(name, **{str(k): (type(v), ...) for k, v in specs.items()})
+
+
+class NodeMaterializer(BaseModel):
+    source_nodes: list[str]
+    path: str
+    sink: str
+    data_saver: str
+
+
+def denormalize_node_list(nodes: list) -> dict:
+    denorm = dict()
+    for node in nodes:
+        name = node["name"]
+        del node["name"]
+        denorm[name] = node
+    return denorm
+
+
+class RunMetadata(BaseModel):
+    experiment: str
+    run_id: str
+    run_dir: str
+    success: bool
+    date_completed: datetime.datetime
+    graph_hash: str
+    modules: list[str]
+    config: dict
+    inputs: dict
+    overrides: dict
+    materialized: list[NodeMaterializer]
+    graph_version: int | None = None
+
+    @model_validator(mode="before")
+    def pre_root(cls, v: dict[str, Any]):
+        dt = datetime.datetime.fromisoformat(v["date_completed"])
+        res = datetime.timedelta(seconds=1)
+        nsecs = dt.hour * 3600 + dt.minute * 60 + dt.second + dt.microsecond * 1e-6
+        delta = nsecs % res.seconds
+        v["date_completed"] = dt - datetime.timedelta(seconds=delta)
+
+        for field in ["inputs", "overrides"]:
+            v[field] = denormalize_node_list(v[field])
+        return v

hamilton/plugins/h_experiments/hook.py

@@ -0,0 +1,219 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import hashlib
+import inspect
+import json
+import logging
+import os
+import string
+import uuid
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+from hamilton import graph_types, lifecycle
+from hamilton.plugins.h_experiments.cache import JsonCache
+
+logger = logging.getLogger(__name__)
+
+
+def validate_string_input(user_input):
+    """Validate the experiment name will make a valid directory name"""
+    allowed = set(string.ascii_letters + string.digits + "_" + "-")
+    for char in user_input:
+        if char not in allowed:
+            raise ValueError(f"`{char}` from `{user_input}` is an invalid character.")
+
+
+def _get_default_input(node) -> Any:
+    """Get default input node value from originating function signature"""
+    param_name = node.name
+    origin_function = node.originating_functions[0]
+    param = inspect.signature(origin_function).parameters[param_name]
+    return None if param.default is inspect._empty else param.default
+
+
+def json_encoder(obj: Any):
+    """convert non JSON-serializable objects to a serializable format
+
+    set[T] -> list[T]
+    else -> dict[type: str, byte_hash: str]
+    """
+    if isinstance(obj, set):
+        serialized = list(obj)
+    else:
+        obj_hash = hashlib.sha256()
+        obj_hash.update(obj)
+        serialized = dict(
+            dtype=type(obj).__name__,
+            obj_hash=obj_hash.hexdigest(),
+        )
+    return serialized
+
+
+@dataclass
+class NodeImplementation:
+    name: str
+    source_code: str
+
+
+@dataclass
+class NodeInput:
+    name: str
+    value: Any
+    default_value: Any | None
+
+
+@dataclass
+class NodeOverride:
+    name: str
+    value: Any
+
+
+@dataclass
+class NodeMaterializer:
+    source_nodes: list[str]
+    path: str
+    sink: str
+    data_saver: str
+
+
+@dataclass
+class RunMetadata:
+    """Metadata about an Hamilton to store in cache"""
+
+    experiment: str
+    run_id: str
+    run_dir: str
+    success: bool
+    date_completed: datetime.datetime
+    graph_hash: str
+    modules: list[str]
+    config: dict
+    inputs: list[NodeInput]
+    overrides: list[NodeOverride]
+    materialized: list[NodeMaterializer]
+
+
+class ExperimentTracker(
+    lifecycle.NodeExecutionHook,
+    lifecycle.GraphExecutionHook,
+    lifecycle.GraphConstructionHook,
+):
+    def __init__(self, experiment_name: str, base_directory: str = "./experiments"):
+        validate_string_input(experiment_name)
+
+        self.experiment_name = experiment_name
+        self.cache = JsonCache(cache_path=base_directory)
+        self.run_id = str(uuid.uuid4())
+
+        self.init_directory = Path.cwd()
+        self.run_directory = (
+            Path(base_directory).resolve().joinpath(self.experiment_name, self.run_id)
+        )
+        self.run_directory.mkdir(exist_ok=True, parents=True)
+
+        self.graph_hash: str = ""
+        self.modules: set[str] = set()
+        self.config = dict()
+        self.inputs: list[NodeInput] = list()
+        self.overrides: list[NodeOverride] = list()
+        self.materializers: list[NodeMaterializer] = list()
+
+    def run_after_graph_construction(self, *, config: dict[str, Any], **kwargs):
+        """Store the Driver config before creating the graph"""
+        self.config = config
+
+    def run_before_graph_execution(
+        self,
+        *,
+        graph: graph_types.HamiltonGraph,
+        inputs: dict[str, Any],
+        overrides: dict[str, Any],
+        **kwargs,
+    ):
+        """Store execution metadata: graph hash, inputs, overrides"""
+        self.graph_hash = graph.version
+
+        for node in graph.nodes:
+            if node.tags.get("module"):
+                self.modules.add(node.tags["module"])
+
+            # filter out config nodes
+            elif node.is_external_input and node.originating_functions:
+                self.inputs.append(
+                    NodeInput(
+                        name=node.name,
+                        value=inputs.get(node.name),
+                        default_value=_get_default_input(node),
+                    )
+                )
+
+        if overrides:
+            self.overrides = [NodeOverride(name=k, value=v) for k, v in overrides.items()]
+
+    def run_before_node_execution(self, *args, node_tags: dict, **kwargs):
+        """Move to run directory before executing materializer"""
+        if node_tags.get("hamilton.data_saver") is True:
+            os.chdir(self.run_directory)  # before materialization
+
+    def run_after_node_execution(
+        self, *, node_name: str, node_tags: dict, node_kwargs: dict, result: Any, **kwargs
+    ):
+        """Move back to init directory after executing materializer.
+        Then, save materialization metadata
+        """
+        if node_tags.get("hamilton.data_saver") is True:
+            if "path" in result:
+                path = result["path"]
+            elif "file_metadata" in result:
+                path = result["file_metadata"]["path"]
+            else:
+                logger.warning(
+                    f"Materialization result from node={node_name} has no recordable path: {result}. Materializer must have either "
+                    f"'path' or 'file_metadata' keys."
+                )
+            self.materializers.append(
+                NodeMaterializer(
+                    source_nodes=list(node_kwargs.keys()),
+                    path=str(Path(path).resolve()),
+                    sink=node_tags["hamilton.data_saver.sink"],
+                    data_saver=node_tags["hamilton.data_saver.classname"],
+                )
+            )
+            os.chdir(self.init_directory)  # after materialization
+
+    def run_after_graph_execution(self, *, success: bool, **kwargs):
+        """Encode run metadata as JSON and store in cache"""
+        run_data = dict(
+            experiment=self.experiment_name,
+            run_id=self.run_id,
+            run_dir=str(self.run_directory),
+            date_completed=datetime.datetime.now().isoformat(),
+            success=success,
+            graph_hash=self.graph_hash,
+            modules=list(self.modules),
+            config=self.config,
+            inputs=[] if len(self.inputs) == 0 else [asdict(i) for i in self.inputs],
+            overrides=[] if len(self.overrides) == 0 else [asdict(o) for o in self.overrides],
+            materialized=[asdict(m) for m in self.materializers],
+        )
+
+        run_json_string = json.dumps(run_data, default=str, sort_keys=True)
+        self.cache.write(run_json_string, self.run_id)