awx-zipline-ai 0.2.0__py3-none-any.whl

ai/chronon/repo/hub_runner.py

@@ -0,0 +1,171 @@
+import json
+import os
+from datetime import datetime
+from urllib.parse import quote_plus
+
+import click
+from attr import dataclass
+
+from ai.chronon.cli.git_utils import get_current_branch
+from ai.chronon.repo import hub_uploader, utils
+from ai.chronon.repo.constants import RunMode
+from ai.chronon.repo.zipline_hub import ZiplineHub
+
+
+@click.group()
+def hub():
+    pass
+
+
+#### Common click options
+def common_options(func):
+    func = click.option("--repo", help="Path to chronon repo", default=".")(func)
+    func = click.option("--conf", required=True, help="Conf param - required for every mode")(func)
+    return func
+
+def ds_option(func):
+    return click.option("--ds", help="the end partition to backfill the data")(func)
+
+def start_ds_option(func):
+    return click.option(
+    "--start-ds",
+    help="override the original start partition for a range backfill. "
+    "It only supports staging query, group by backfill and join jobs. "
+    "It could leave holes in your final output table due to the override date range.",)(func)
+
+
+def end_ds_option(func):
+    return click.option("--end-ds", help="the end ds for a range backfill")(func)
+
+
+def submit_workflow(repo,
+                    conf,
+                    mode,
+                    start_ds,
+                    end_ds):
+
+    hub_conf = get_hub_conf(conf)
+    zipline_hub = ZiplineHub(base_url=hub_conf.hub_url)
+    conf_name_to_hash_dict = hub_uploader.build_local_repo_hashmap(root_dir= repo)
+    branch = get_current_branch()
+
+    hub_uploader.compute_and_upload_diffs(branch, zipline_hub=zipline_hub, local_repo_confs=conf_name_to_hash_dict)
+
+    # get conf name
+    conf_name = utils.get_metadata_name_from_conf(repo, conf)
+
+
+    response_json = zipline_hub.call_workflow_start_api(
+        conf_name=conf_name,
+        mode=mode,
+        branch=branch,  # Get the current branch
+        user=os.environ.get('USER'),
+        start=start_ds,
+        end=end_ds,
+        conf_hash=conf_name_to_hash_dict[conf_name].hash,
+    )
+
+    print(" 🆔 Workflow Id:", response_json.get("workflowId", "N/A"))
+    print_wf_url(
+        conf=conf,
+        conf_name=conf_name,
+        mode=RunMode.BACKFILL.value,
+        start_ds=start_ds,
+        end_ds=end_ds,
+        branch=branch
+    )
+
+
+# zipline hub backfill --conf=compiled/joins/join
+# adhoc backfills
+@hub.command()
+@common_options
+@start_ds_option
+@end_ds_option
+def backfill(repo,
+             conf,
+             start_ds,
+             end_ds):
+    """
+    - Submit a backfill job to Zipline.
+    Response should contain a list of confs that are different from what's on remote.
+    - Call upload API to upload the conf contents for the list of confs that were different.
+    - Call the actual run API with mode set to backfill.
+    """
+    submit_workflow(repo, conf, RunMode.BACKFILL.value, start_ds, end_ds)
+
+
+
+# zipline hub deploy --conf=compiled/joins/join
+# currently only supports one-off deploy node submission
+@hub.command()
+@common_options
+@end_ds_option
+def deploy(repo,
+           conf,
+           end_ds):
+    """
+    - Submit a one-off deploy job to Zipline.
+    Response should contain a list of confs that are different from what's on remote.
+    - Call upload API to upload the conf contents for the list of confs that were different.
+    - Call the actual run API with mode set to deploy
+    """
+    submit_workflow(repo, conf, RunMode.DEPLOY.value, end_ds, end_ds)
+
+
+def get_common_env_map(file_path):
+    with open(file_path, 'r') as f:
+        data = json.load(f)
+    common_env_map = data['metaData']['executionInfo']['env']['common']
+    return common_env_map
+
+
+@dataclass
+class HubConfig:
+    hub_url: str
+    frontend_url: str
+
+
+def get_hub_conf(conf_path):
+    common_env_map = get_common_env_map(conf_path)
+    hub_url = common_env_map.get("HUB_URL", os.environ.get("HUB_URL"))
+    frontend_url = common_env_map.get("FRONTEND_URL", os.environ.get("FRONTEND_URL"))
+    return HubConfig(hub_url=hub_url, frontend_url=frontend_url)
+
+
+def print_wf_url(conf, conf_name, mode, start_ds, end_ds, branch):
+
+    hub_conf = get_hub_conf(conf)
+    frontend_url = hub_conf.frontend_url
+
+    if "compiled/joins" in conf:
+        hub_conf_type = "joins"
+    elif "compiled/staging_queries" in conf:
+        hub_conf_type = "stagingqueries"
+    elif "compiled/group_by" in conf:
+        hub_conf_type = "groupby"
+    elif "compiled/models" in conf:
+        hub_conf_type = "models"
+    else:
+        raise ValueError(f"Unsupported conf type: {conf}")
+
+    # TODO: frontend uses localtime to create the millis, we should make it use UTC and make this align
+    def _millis(date_str):
+        return int(datetime.strptime(date_str, "%Y-%m-%d").timestamp() * 1000)
+
+    def _mode_string(mode):
+        if mode == "backfill":
+            return "offline"
+        elif mode == "deploy":
+            return "online"
+        else:
+            raise ValueError(f"Unsupported mode: {mode}")
+
+    workflow_url = f"{frontend_url.rstrip('/')}/{hub_conf_type}/{conf_name}/{_mode_string(mode)}?start={_millis(start_ds)}&end={_millis(end_ds)}&branch={quote_plus(branch)}"
+
+    print(" 🔗 Workflow : " + workflow_url + "\n")
+
+if __name__ == "__main__":
+    hub()
+
+

ai/chronon/repo/hub_uploader.py

@@ -0,0 +1,108 @@
+import glob
+import hashlib
+import json
+import os
+
+from ai.chronon.orchestration.ttypes import Conf
+from ai.chronon.repo import (
+    FOLDER_NAME_TO_CLASS,
+    FOLDER_NAME_TO_CONF_TYPE,
+)
+from ai.chronon.repo.zipline_hub import ZiplineHub
+
+
+def build_local_repo_hashmap(root_dir: str):
+    compiled_dir = os.path.join(root_dir, "compiled")
+    # Returns a map of name -> (tbinary, file_hash)
+    results = {}
+
+    # Iterate through each object type folder (staging_queries, group_bys, joins etc)
+    for folder_name, _ in FOLDER_NAME_TO_CLASS.items():
+        folder_path = os.path.join(compiled_dir, folder_name)
+        if not os.path.exists(folder_path):
+            continue
+
+        # Find all json files recursively in this folder
+        json_files = [
+            f
+            for f in glob.glob(os.path.join(folder_path, "**/*"), recursive=True)
+            if os.path.isfile(f)
+        ]
+
+        exceptions = []
+
+        for json_file in json_files:
+            try:
+                # Read the json file
+                with open(json_file, "r") as f:
+                    thrift_json = f.read()
+
+                # Extract name from metadata in json
+                json_obj = json.loads(thrift_json)
+                name = json_obj["metaData"]["name"]
+
+                # Load the json into the appropriate object type based on folder
+                # binary = json2binary(thrift_json, obj_class)
+
+                md5_hash = hashlib.md5(thrift_json.encode()).hexdigest()
+                # md5_hash = hashlib.md5(thrift_json.encode()).hexdigest() + "123"
+                # results[name] = (binary, md5_hash, FOLDER_NAME_TO_CONF_TYPE[folder_name])
+                results[name] = Conf(
+                    name=name,
+                    hash=md5_hash,
+                    # contents=binary,
+                    contents=thrift_json,
+                    confType=FOLDER_NAME_TO_CONF_TYPE[folder_name],
+                )
+
+            except Exception as e:
+                exceptions.append(f"{json_file} - {e}")
+
+        if exceptions:
+            error_msg = (
+                "The following files had exceptions during upload: \n"
+                + "\n".join(exceptions)
+                + "\n\n Consider deleting the files (safe operation) and checking "
+                + "your thrift version before rerunning your command."
+            )
+            raise RuntimeError(error_msg)
+
+    return results
+
+
+def compute_and_upload_diffs(
+    branch: str, zipline_hub: ZiplineHub, local_repo_confs: dict[str, Conf]
+):
+    # Determine which confs are different from the ZiplineHub
+    # Call Zipline hub with `names_and_hashes` as the argument to get back
+    names_to_hashes = {name: local_conf.hash for name, local_conf in local_repo_confs.items()}
+    print(f"\n 🧮 Computed hashes for {len(names_to_hashes)} local files.")
+
+    changed_conf_names = zipline_hub.call_diff_api(names_to_hashes)
+
+    if not changed_conf_names:
+        print(f" ✅ Remote contains all local files. No need to upload '{branch}'.")
+        diffed_confs = {}
+    else:
+        unchanged = len(names_to_hashes) - len(changed_conf_names)
+        print(
+            f" 🔍 Detected {len(changed_conf_names)} changes on local branch '{branch}'. {unchanged} unchanged."
+        )
+
+        # a list of names for diffed hashes on branch
+        diffed_confs = {k: local_repo_confs[k] for k in changed_conf_names}
+        conf_names_str = "\n    - ".join(diffed_confs.keys())
+        print(f"    - {conf_names_str}")
+
+        diff_confs = []
+        for _, conf in diffed_confs.items():
+            diff_confs.append(conf.__dict__)
+
+        # Make PUT request to ZiplineHub
+        zipline_hub.call_upload_api(branch=branch, diff_confs=diff_confs)
+        print(f" ⬆️ Uploaded {len(diffed_confs)} changed confs to branch '{branch}'.")
+
+    zipline_hub.call_sync_api(branch=branch, names_to_hashes=names_to_hashes)
+
+    print(f" ✅ {len(names_to_hashes)} hashes updated on branch '{branch}'.\n")
+    return diffed_confs

ai/chronon/repo/init.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python
+
+import os
+import shutil
+
+import click
+from importlib_resources import files
+from rich.prompt import Prompt
+from rich.syntax import Syntax
+
+from ai.chronon.cli.compile.display.console import console
+
+
+@click.command(name="init")
+@click.option(
+    "--cloud-provider",
+    envvar="CLOUD_PROVIDER",
+    help="Cloud provider to use.",
+    required=True,
+    type=click.Choice(['aws', 'gcp'], case_sensitive=False)
+)
+@click.option(
+    "--chronon-root",
+    help="Path to the root chronon folder.",
+    default=os.path.join(os.getcwd(), "zipline"),
+    type=click.Path(file_okay=False, writable=True),
+)
+@click.pass_context
+def main(ctx, chronon_root, cloud_provider):
+    template_path = files("ai.chronon").joinpath("resources", cloud_provider.lower())
+    target_path = os.path.abspath(chronon_root)
+
+    if os.path.exists(target_path) and os.listdir(target_path):
+        choice = Prompt.ask(f"[bold yellow] Warning: [/]{target_path} is not empty. Proceed?",
+                   choices=["y", "n"],
+                   default="y")
+        if choice == "n":
+            return
+
+    console.print(f"Generating scaffolding at {target_path} ...")
+
+    try:
+        shutil.copytree(template_path, target_path, dirs_exist_ok=True)
+        console.print("[bold green] Project scaffolding created successfully! 🎉\n")
+        export_cmd = Syntax(f"`export PYTHONPATH={target_path}:$PYTHONPATH`", "bash", theme="github-dark", line_numbers=False)
+        console.print("Please copy the following command to your shell config:")
+        console.print(export_cmd)
+    except Exception:
+        console.print_exception()
+
+
+if __name__ == "__main__":
+    main()

ai/chronon/repo/join_backfill.py

@@ -0,0 +1,105 @@
+import json
+import os
+from typing import Optional
+
+from ai.chronon.scheduler.interfaces.flow import Flow
+from ai.chronon.scheduler.interfaces.node import Node
+from ai.chronon.utils import (
+    convert_json_to_obj,
+    dict_to_bash_commands,
+    dict_to_exports,
+    get_join_output_table_name,
+    join_part_name,
+    sanitize,
+)
+
+TASK_PREFIX = "compute_join"
+DEFAULT_SPARK_SETTINGS = {
+    "default": {
+        "spark_version": "3.1.1",
+        "executor_memory": "4G",
+        "driver_memory": "4G",
+        "executor_cores": 2,
+    }
+}
+
+
+class JoinBackfill:
+    def __init__(
+        self,
+        start_date: str,
+        end_date: str,
+        config_path: str,
+        settings: dict = DEFAULT_SPARK_SETTINGS,
+    ):
+        self.dag_id = "_".join(
+            map(
+                sanitize, ["chronon_joins_backfill", os.path.basename(config_path).split("/")[-1], start_date, end_date]
+            )
+        )
+        self.start_date = start_date
+        self.end_date = end_date
+        self.config_path = config_path
+        self.settings = settings
+        with open(self.config_path, "r") as file:
+            config = file.read()
+        self.join = convert_json_to_obj(json.loads(config))
+
+    def build_flow(self) -> Flow:
+        """
+        Build a flow from a Join object. Each join part is a node and will run in parallel.
+        The next step is final join, which is a node that depends on all join parts.
+        The final join will run after all join parts are done.
+
+        :param join: The Join object to build a flow from
+        :return: A Flow object that represents the flow of the Join
+        """
+        flow = Flow(self.join.metaData.name)
+        final_node = Node(
+            f"{TASK_PREFIX}__{sanitize(get_join_output_table_name(self.join, full_name=True))}", self.run_final_join()
+        )
+        left_node = Node(f"{TASK_PREFIX}__left_table", self.run_left_table())
+        flow.add_node(final_node)
+        flow.add_node(left_node)
+        for join_part in self.join.joinParts:
+            jp_full_name = join_part_name(join_part)
+            jp_node = Node(f"{TASK_PREFIX}__{jp_full_name}", self.run_join_part(jp_full_name))
+            flow.add_node(jp_node)
+            jp_node.add_dependency(left_node)
+            final_node.add_dependency(jp_node)
+        return flow
+
+    def export_template(self, settings: dict):
+        return f"{dict_to_exports(settings)}"
+
+    def command_template(self, extra_args: dict):
+        if self.start_date:
+            extra_args.update({"start_ds": self.start_date})
+        return f"""python3 /tmp/run.py --conf=/tmp/{self.config_path} --ds={self.end_date} \
+{dict_to_bash_commands(extra_args)}"""
+
+    def run_join_part(self, join_part: str):
+        args = {
+            "mode": "backfill",
+            "selected_join_parts": join_part,
+            "use_cached_left": None,
+        }
+        settings = self.settings.get(join_part, self.settings["default"])
+        return self.export_template(settings) + " && " + self.command_template(extra_args=args)
+
+    def run_left_table(self):
+        settings = self.settings.get("left_table", self.settings["default"])
+        return self.export_template(settings) + " && " + self.command_template(extra_args={"mode": "backfill-left"})
+
+    def run_final_join(self):
+        settings = self.settings.get("final_join", self.settings["default"])
+        return self.export_template(settings) + " && " + self.command_template(extra_args={"mode": "backfill-final"})
+
+    def run(self, orchestrator: str, overrides: Optional[dict] = None):
+        from ai.chronon.constants import ADAPTERS
+
+        ADAPTERS.update(overrides)
+        orchestrator = ADAPTERS[orchestrator](dag_id=self.dag_id, start_date=self.start_date)
+        orchestrator.setup()
+        orchestrator.build_dag_from_flow(self.build_flow())
+        orchestrator.trigger_run()