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

ai/chronon/repo/hub_runner.py

@@ -0,0 +1,261 @@
+import json
+import os
+from dataclasses import dataclass
+from datetime import date, timedelta
+from typing import Optional
+
+import click
+from gen_thrift.planner.ttypes import Mode
+
+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.utils import handle_conf_not_found, print_possible_confs
+from ai.chronon.repo.zipline_hub import ZiplineHub
+
+ALLOWED_DATE_FORMATS = ["%Y-%m-%d"]
+
+
+@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)
+    func = click.option(
+        "--hub_url", help="Zipline Hub address, e.g. http://localhost:3903", default=None
+    )(func)
+    return func
+
+
+def ds_option(func):
+    return click.option(
+        "--ds",
+        help="the end partition to backfill the data",
+        type=click.DateTime(formats=ALLOWED_DATE_FORMATS),
+    )(func)
+
+
+def start_ds_option(func):
+    return click.option(
+        "--start-ds",
+        type=click.DateTime(formats=ALLOWED_DATE_FORMATS),
+        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",
+        type=click.DateTime(formats=ALLOWED_DATE_FORMATS),
+        default=str(date.today() - timedelta(days=2)),
+    )(func)
+
+
+def submit_workflow(repo, conf, mode, start_ds, end_ds, hub_url=None):
+    hub_conf = get_hub_conf(conf, root_dir=repo)
+    if hub_url is not None:
+        zipline_hub = ZiplineHub(base_url=hub_url, sa_name=hub_conf.sa_name)
+    else:
+        zipline_hub = ZiplineHub(base_url=hub_conf.hub_url, sa_name=hub_conf.sa_name)
+    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,
+        skip_long_running=False,
+    )
+
+    workflow_id = response_json.get("workflowId", "N/A")
+    print(" 🆔 Workflow Id:", workflow_id)
+    print_wf_url(
+        conf=conf,
+        conf_name=conf_name,
+        mode=mode,
+        workflow_id=workflow_id,
+        repo=repo
+    )
+
+
+def submit_schedule(repo, conf, hub_url=None):
+    hub_conf = get_hub_conf(conf, root_dir=repo)
+    if hub_url is not None:
+        zipline_hub = ZiplineHub(base_url=hub_url, sa_name=hub_conf.sa_name)
+    else:
+        zipline_hub = ZiplineHub(base_url=hub_conf.hub_url, sa_name=hub_conf.sa_name)
+    conf_name_to_obj_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_obj_dict
+    )
+
+    # get conf name
+    conf_name = utils.get_metadata_name_from_conf(repo, conf)
+    schedule_modes = get_schedule_modes(os.path.join(repo, conf))
+    # create a dict for RunMode.BACKFILL.value and RunMode.DEPLOY.value to schedule_modes.offline_schedule and schedule_modes.online
+    modes = {
+        RunMode.BACKFILL.value.upper(): schedule_modes.offline_schedule,
+        RunMode.DEPLOY.value.upper(): schedule_modes.online,
+    }
+
+    response_json = zipline_hub.call_schedule_api(
+        modes=modes,
+        branch=branch,
+        conf_name=conf_name,
+        conf_hash=conf_name_to_obj_dict[conf_name].hash,
+    )
+
+    schedules = response_json.get("schedules", "N/A")
+    readable_schedules = {Mode._VALUES_TO_NAMES[int(k)]: v for k, v in schedules.items()}
+    print(" 🗓️ Schedules Deployed:", readable_schedules)
+
+
+# zipline hub backfill --conf=compiled/joins/join
+# adhoc backfills
+@hub.command()
+@common_options
+@start_ds_option
+@end_ds_option
+@handle_conf_not_found(log_error=True, callback=print_possible_confs)
+def backfill(repo, conf, hub_url, 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, hub_url=hub_url
+    )
+
+
+# zipline hub run-adhoc --conf=compiled/joins/join
+# currently only supports one-off deploy node submission
+@hub.command()
+@common_options
+@end_ds_option
+@handle_conf_not_found(log_error=True, callback=print_possible_confs)
+def run_adhoc(repo, conf, hub_url, end_ds):
+    """
+    - Submit a one-off deploy job to Zipline. This submits the various jobs to allow your conf to be tested online.
+    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, hub_url=hub_url)
+
+
+# zipline hub schedule --conf=compiled/joins/join
+@hub.command()
+@common_options
+@handle_conf_not_found(log_error=True, callback=print_possible_confs)
+def schedule(repo, conf, hub_url):
+    """
+    - Deploys a schedule for the specified conf to Zipline. This allows your conf to have various associated jobs run on a schedule.
+    This verb will introspect your conf to determine which of its jobs need to be scheduled (or paused if turned off) based on the
+    'offline_schedule' and 'online' fields.
+    """
+    submit_schedule(repo, conf, hub_url=hub_url)
+
+
+def get_metadata_map(file_path):
+    with open(file_path, "r") as f:
+        data = json.load(f)
+    metadata_map = data["metaData"]
+    return metadata_map
+
+
+def get_common_env_map(file_path):
+    metadata_map = get_metadata_map(file_path)
+    common_env_map = metadata_map["executionInfo"]["env"]["common"]
+    return common_env_map
+
+
+@dataclass
+class HubConfig:
+    hub_url: str
+    frontend_url: str
+    sa_name: Optional[str] = None
+
+
+@dataclass
+class ScheduleModes:
+    online: str
+    offline_schedule: str
+
+
+def get_hub_conf(conf_path, root_dir="."):
+    file_path = os.path.join(root_dir, conf_path)
+    common_env_map = get_common_env_map(file_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"))
+    sa_name = common_env_map.get("SA_NAME", os.environ.get("SA_NAME"))
+    return HubConfig(hub_url=hub_url, frontend_url=frontend_url, sa_name=sa_name)
+
+
+def get_schedule_modes(conf_path):
+    metadata_map = get_metadata_map(conf_path)
+    online_value = metadata_map.get("online", False)
+    online = "true" if bool(online_value) else "false"
+    offline_schedule = metadata_map["executionInfo"].get("scheduleCron", None)
+
+    # check if offline_schedule is null or 'None' or '@daily' else throw an error
+    valid_schedules = {None, "None", "@daily"}
+    if offline_schedule not in valid_schedules:
+        raise ValueError(
+            f"Unsupported offline_schedule: {offline_schedule}. Only null, 'None', or '@daily' are supported."
+        )
+    offline_schedule = offline_schedule or "None"
+    return ScheduleModes(online=online, offline_schedule=offline_schedule)
+
+
+def print_wf_url(conf, conf_name, mode, workflow_id, repo="."):
+    hub_conf = get_hub_conf(conf, root_dir=repo)
+    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 = "groupbys"
+    elif "compiled/models" in conf:
+        hub_conf_type = "models"
+    else:
+        raise ValueError(f"Unsupported conf type: {conf}")
+
+    def _mode_string():
+        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()}?workflowId={workflow_id}"
+
+    print(" 🔗 Workflow : " + workflow_url + "\n")
+
+
+if __name__ == "__main__":
+    hub()

ai/chronon/repo/hub_uploader.py

@@ -0,0 +1,109 @@
+import glob
+import hashlib
+import json
+import os
+
+from gen_thrift.api.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,60 @@
+#!/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,119 @@
+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,
+    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(self.join.table)}",
+            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()