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

ai/chronon/query.py

@@ -0,0 +1,126 @@
+#     Copyright (C) 2023 The Chronon Authors.
+#
+#     Licensed 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 collections import OrderedDict
+from typing import Dict, List
+
+import ai.chronon.api.ttypes as api
+
+
+def Query(
+    selects: Dict[str, str] = None,
+    wheres: List[str] = None,
+    start_partition: str = None,
+    end_partition: str = None,
+    time_column: str = None,
+    setups: List[str] = None,
+    mutation_time_column: str = None,
+    reversal_column: str = None,
+    partition_column: str = None,
+    partition_format: str = None,
+    sub_partitions_to_wait_for: List[str] = None,
+) -> api.Query:
+    """
+    Create a query object that is used to scan data from various data sources.
+    This contains partition ranges, row level transformations and filtering logic.
+    Additionally we also require a time_column for TEMPORAL events, mutation_time_column & reversal
+    for TEMPORAL entities.
+
+    :param selects: Spark sql expressions with only arithmetic, function application & inline lambdas.
+        You can also apply udfs see setups param below.::
+
+            Example: {
+                "alias": "built_in_function(col1) * my_udf(col2)",
+                "alias1": "aggregate(array_col, 0, (acc, x) -> acc + x)"
+            }
+
+        See: https://spark.apache.org/docs/latest/api/sql/#built-in-functions
+        When none, we will assume that no transformations are needed and will pick columns necessary for aggregations.
+    :type selects: List[str], optional
+    :param wheres: Used for filtering. Same as above, but each expression must return boolean.
+        Expressions are joined using AND.
+    :type wheres: List[str], optional
+    :param start_partition: From which partition of the source is the data valid from - inclusive.
+        When absent we will consider all available data is usable.
+    :type start_partition: str, optional
+    :param end_partition: Till what partition of the source is the data valid till - inclusive.
+        Not specified unless you know for a fact that a particular source has expired after a partition and you
+        should instead use another source after this partition.
+    :type end_partition: str, optional
+    :param time_column: a single expression to produce time as ** milliseconds since epoch**.
+    :type time_column: str, optional
+    :param setups: you can register UDFs using setups
+        ["ADD JAR YOUR_JAR", "create temporary function YOU_UDF_NAME as YOUR_CLASS"]
+    :type setups: List[str], optional
+    :param mutation_time_column: For entities, with real time accuracy, you need to specify an expression that
+        represents mutation time. Time should be milliseconds since epoch.
+        This is not necessary for event sources, defaults to "mutation_ts"
+    :type mutation_time_column: str, optional
+    :param reversal_column: (defaults to "is_before")
+        For entities with realtime accuracy, we divide updates into two additions & reversal.
+        updates have two rows - one with is_before = True (the old value) & is_before = False (the new value)
+        inserts only have is_before = false (just the new value).
+        deletes only have is_before = true (just the old value).
+        This is not necessary for event sources.
+    :type reversal_column: str, optional
+    :param partition_column:
+        Specify this to override spark.chronon.partition.column set in teams.py for this particular query.
+    :type partition_column: str, optional
+    :param sub_partitions_to_wait_for:
+        Additional partitions to be used in sensing that the source data has landed. Should be a full partition string, such as `hr=23:00'
+    :type sub_partitions_to_wait_for: List[str], optional
+    :param partition_format:
+        Date format string to expect the partition values to be in.
+    :type partition_format: str, optional
+    :return: A Query object that Chronon can use to scan just the necessary data efficiently.
+    """
+    return api.Query(
+        selects=selects,
+        wheres=wheres,
+        startPartition=start_partition,
+        endPartition=end_partition,
+        timeColumn=time_column,
+        setups=setups,
+        mutationTimeColumn=mutation_time_column,
+        reversalColumn=reversal_column,
+        partitionColumn=partition_column,
+        subPartitionsToWaitFor=sub_partitions_to_wait_for,
+        partitionFormat=partition_format
+    )
+
+
+def selects(*args, **kwargs):
+    """
+    Create a dictionary required for the selects parameter of Query.
+
+    .. code-block:: python
+        selects(
+            "event_id",
+            user_id="user_id",
+        )
+
+    creates the following dictionary:
+
+    .. code-block:: python
+        {
+            "event_id": "event_id",
+            "user_id": "user_id"
+        }
+    """
+    result = OrderedDict()
+    for x in args:
+        result[x] = x
+    for k, v in kwargs.items():
+        result[k] = v
+    return result

ai/chronon/repo/__init__.py

@@ -0,0 +1,40 @@
+#     Copyright (C) 2023 The Chronon Authors.
+#
+#     Licensed 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 ai.chronon.api.ttypes import GroupBy, Join, Model, StagingQuery
+from ai.chronon.orchestration.ttypes import ConfType
+
+JOIN_FOLDER_NAME = "joins"
+GROUP_BY_FOLDER_NAME = "group_bys"
+STAGING_QUERY_FOLDER_NAME = "staging_queries"
+MODEL_FOLDER_NAME = "models"
+# TODO - make team part of thrift API?
+TEAMS_FILE_PATH = "teams.json"
+OUTPUT_ROOT = "production"
+
+# This is set in the main function -
+# from command line or from env variable during invocation
+FOLDER_NAME_TO_CLASS = {
+    GROUP_BY_FOLDER_NAME: GroupBy,
+    JOIN_FOLDER_NAME: Join,
+    STAGING_QUERY_FOLDER_NAME: StagingQuery,
+    MODEL_FOLDER_NAME: Model,
+}
+
+FOLDER_NAME_TO_CONF_TYPE = {
+    GROUP_BY_FOLDER_NAME: ConfType.GROUP_BY,
+    JOIN_FOLDER_NAME: ConfType.JOIN,
+    STAGING_QUERY_FOLDER_NAME: ConfType.STAGING_QUERY,
+    MODEL_FOLDER_NAME: ConfType.MODEL,
+}

ai/chronon/repo/aws.py

@@ -0,0 +1,298 @@
+import json
+import multiprocessing
+import os
+from typing import List
+
+import boto3
+
+from ai.chronon.logger import get_logger
+from ai.chronon.repo.constants import ROUTES, ZIPLINE_DIRECTORY
+from ai.chronon.repo.default_runner import Runner
+from ai.chronon.repo.utils import (
+    JobType,
+    check_call,
+    extract_filename_from_path,
+    get_customer_id,
+    split_date_range,
+)
+
+LOG = get_logger()
+
+# AWS SPECIFIC CONSTANTS
+EMR_ENTRY = "ai.chronon.integrations.aws.EmrSubmitter"
+ZIPLINE_AWS_JAR_DEFAULT = "cloud_aws_lib_deploy.jar"
+ZIPLINE_AWS_ONLINE_CLASS_DEFAULT = "ai.chronon.integrations.aws.AwsApiImpl"
+ZIPLINE_AWS_FLINK_JAR_DEFAULT = "flink_assembly_deploy.jar"
+ZIPLINE_AWS_SERVICE_JAR = "service_assembly_deploy.jar"
+
+LOCAL_FILE_TO_ETAG_JSON = f"{ZIPLINE_DIRECTORY}/local_file_to_etag.json"
+
+EMR_MOUNT_FILE_PREFIX = "/mnt/zipline/"
+
+
+class AwsRunner(Runner):
+    def __init__(self, args):
+        aws_jar_path = AwsRunner.download_zipline_aws_jar(
+            ZIPLINE_DIRECTORY, get_customer_id(), args["version"], ZIPLINE_AWS_JAR_DEFAULT
+        )
+        service_jar_path = AwsRunner.download_zipline_aws_jar(
+            ZIPLINE_DIRECTORY, get_customer_id(), args["version"], ZIPLINE_AWS_SERVICE_JAR
+        )
+        jar_path = (
+            f"{service_jar_path}:{aws_jar_path}" if args['mode'] == "fetch" else aws_jar_path
+        )
+        self.version = args.get("version", "latest")
+
+        super().__init__(args, os.path.expanduser(jar_path))
+
+    @staticmethod
+    def upload_s3_file(
+        bucket_name: str, source_file_name: str, destination_blob_name: str
+    ):
+        """Uploads a file to the bucket."""
+        obj = boto3.client("s3")
+        try:
+            obj.upload_file(source_file_name, bucket_name, destination_blob_name)
+            print(
+                f"File {source_file_name} uploaded to {destination_blob_name} in bucket {bucket_name}."
+            )
+            return f"s3://{bucket_name}/{destination_blob_name}"
+        except Exception as e:
+            raise RuntimeError(f"Failed to upload {source_file_name}: {str(e)}") from e
+
+    @staticmethod
+    def download_zipline_aws_jar(destination_dir: str, customer_id: str, version: str, jar_name: str):
+        s3_client = boto3.client("s3")
+        destination_path = f"{destination_dir}/{jar_name}"
+        source_key_name = f"release/{version}/jars/{jar_name}"
+        bucket_name = f"zipline-artifacts-{customer_id}"
+
+        are_identical = (
+            AwsRunner.compare_s3_and_local_file_hashes(
+                bucket_name, source_key_name, destination_path
+            )
+            if os.path.exists(destination_path)
+            else False
+        )
+
+        if are_identical:
+            print(f"{destination_path} matches S3 {bucket_name}/{source_key_name}")
+        else:
+            print(
+                f"{destination_path} does NOT match S3 {bucket_name}/{source_key_name}"
+            )
+            print(f"Downloading {jar_name} from S3...")
+
+            s3_client.download_file(
+                Filename=destination_path, Bucket=bucket_name, Key=source_key_name
+            )
+            # Persist ETag to prevent downloading the same file next time
+            etag = AwsRunner.get_s3_file_hash(bucket_name, source_key_name)
+            if os.path.exists(LOCAL_FILE_TO_ETAG_JSON):
+                with open(LOCAL_FILE_TO_ETAG_JSON, "r") as file:
+                    data = json.load(file)
+
+                # Add the new entry
+                data[destination_path] = etag
+
+                # Write the updated dictionary back to the file
+                with open(LOCAL_FILE_TO_ETAG_JSON, "w") as file:
+                    json.dump(data, file)
+            else:
+                with open(LOCAL_FILE_TO_ETAG_JSON, "w") as file:
+                    data = {destination_path: etag}
+                    json.dump(data, file)
+
+        return destination_path
+
+    @staticmethod
+    def get_s3_file_hash(bucket_name: str, file_name: str):
+        s3_client = boto3.client("s3")
+        response = s3_client.head_object(Bucket=bucket_name, Key=file_name)
+        return response["ETag"].strip('"')
+
+    @staticmethod
+    def get_local_file_hash(file_name: str):
+        # read in the json file
+        if os.path.exists(LOCAL_FILE_TO_ETAG_JSON):
+            with open(LOCAL_FILE_TO_ETAG_JSON, "r") as f:
+                data = json.load(f)
+                if file_name in data:
+                    return data[file_name]
+        return None
+
+    @staticmethod
+    def compare_s3_and_local_file_hashes(
+        bucket_name: str, s3_file_path: str, local_file_path: str
+    ):
+        try:
+            s3_hash = AwsRunner.get_s3_file_hash(bucket_name, s3_file_path)
+            local_hash = AwsRunner.get_local_file_hash(local_file_path)
+            print(f"Local hash: {local_hash}, S3 hash: {s3_hash}")
+            return s3_hash == local_hash
+        except Exception as e:
+            print(f"Error comparing files: {str(e)}")
+            return False
+
+    def generate_emr_submitter_args(
+        self,
+        user_args: str,
+        job_type: JobType = JobType.SPARK,
+        local_files_to_upload: List[str] = None,
+    ):
+        customer_warehouse_bucket_name = f"zipline-warehouse-{get_customer_id()}"
+        s3_files = []
+        for source_file in local_files_to_upload:
+            # upload to `metadata` folder
+            destination_file_path = (
+                f"metadata/{extract_filename_from_path(source_file)}"
+            )
+            s3_files.append(
+                AwsRunner.upload_s3_file(
+                    customer_warehouse_bucket_name, source_file, destination_file_path
+                )
+            )
+
+        # we also want the additional-confs included here. it should already be in the bucket
+
+        zipline_artifacts_bucket_prefix = "s3://zipline-artifacts"
+
+        s3_files.append(
+            f"{zipline_artifacts_bucket_prefix}-{get_customer_id()}/confs/additional-confs.yaml"
+        )
+
+        s3_file_args = ",".join(s3_files)
+
+        # include jar uri. should also already be in the bucket
+        jar_uri = (
+            f"{zipline_artifacts_bucket_prefix}-{get_customer_id()}"
+            + f"/release/{self.version}/jars/{ZIPLINE_AWS_JAR_DEFAULT}"
+        )
+
+        final_args = "{user_args} --jar-uri={jar_uri} --job-type={job_type} --main-class={main_class}"
+
+        if job_type == JobType.FLINK:
+            main_class = "ai.chronon.flink.FlinkJob"
+            flink_jar_uri = (
+                f"{zipline_artifacts_bucket_prefix}-{get_customer_id()}"
+                + f"/jars/{ZIPLINE_AWS_FLINK_JAR_DEFAULT}"
+            )
+            return (
+                final_args.format(
+                    user_args=user_args,
+                    jar_uri=jar_uri,
+                    job_type=job_type.value,
+                    main_class=main_class,
+                )
+                + f" --flink-main-jar-uri={flink_jar_uri}"
+            )
+
+        elif job_type == JobType.SPARK:
+            main_class = "ai.chronon.spark.Driver"
+            return (
+                final_args.format(
+                    user_args=user_args,
+                    jar_uri=jar_uri,
+                    job_type=job_type.value,
+                    main_class=main_class,
+                )
+                + f" --additional-conf-path={EMR_MOUNT_FILE_PREFIX}additional-confs.yaml"
+                  f" --files={s3_file_args}"
+            )
+        else:
+            raise ValueError(f"Invalid job type: {job_type}")
+
+    def run(self):
+        command_list = []
+        if self.mode == "info":
+            command_list.append(
+                "python3 {script} --conf {conf} --ds {ds} --repo {repo}".format(
+                    script=self.render_info, conf=self.conf, ds=self.ds, repo=self.repo
+                )
+            )
+        elif self.sub_help or self.mode == "fetch":
+            entrypoint = "ai.chronon.online.fetcher.FetcherMain"
+            command_list.append(
+                "java -cp {jar} {entrypoint} {subcommand} {args}".format(
+                    jar=self.jar_path,
+                    entrypoint=entrypoint,
+                    args="--help" if self.sub_help else self._gen_final_args(),
+                    subcommand=ROUTES[self.conf_type][self.mode],
+                )
+            )
+        elif self.mode in ["streaming", "streaming-client"]:
+            raise ValueError("Streaming is not supported for AWS yet.")
+        else:
+            local_files_to_upload_to_aws = []
+            if self.conf:
+                local_files_to_upload_to_aws.append(os.path.join(self.repo, self.conf))
+            if self.parallelism > 1:
+                assert self.start_ds is not None and self.ds is not None, (
+                    "To use parallelism, please specify --start-ds and --end-ds to "
+                    "break down into multiple backfill jobs"
+                )
+                date_ranges = split_date_range(self.start_ds, self.ds, self.parallelism)
+                for start_ds, end_ds in date_ranges:
+                    user_args = "{subcommand} {args} {additional_args}".format(
+                        subcommand=ROUTES[self.conf_type][self.mode],
+                        args=self._gen_final_args(
+                            start_ds=start_ds,
+                            end_ds=end_ds,
+                            # when we download files from s3 to emr, they'll be mounted at /mnt/zipline
+                            override_conf_path=(
+                                EMR_MOUNT_FILE_PREFIX
+                                + extract_filename_from_path(self.conf)
+                                if self.conf
+                                else None
+                            ),
+                        ),
+                        additional_args=os.environ.get(
+                            "CHRONON_CONFIG_ADDITIONAL_ARGS", ""
+                        ),
+                    )
+
+                    emr_args = self.generate_emr_submitter_args(
+                        local_files_to_upload=local_files_to_upload_to_aws,
+                        # for now, self.conf is the only local file that requires uploading to gcs
+                        user_args=user_args,
+                    )
+                    command = f"java -cp {self.jar_path} {EMR_ENTRY} {emr_args}"
+                    command_list.append(command)
+            else:
+                user_args = ("{subcommand} {args} {additional_args}").format(
+                    subcommand=ROUTES[self.conf_type][self.mode],
+                    args=self._gen_final_args(
+                        start_ds=self.start_ds,
+                        # when we download files from s3 to emr, they'll be mounted at /mnt/zipline
+                        override_conf_path=(
+                            EMR_MOUNT_FILE_PREFIX
+                            + extract_filename_from_path(self.conf)
+                            if self.conf
+                            else None
+                        ),
+                    ),
+                    additional_args=os.environ.get(
+                        "CHRONON_CONFIG_ADDITIONAL_ARGS", ""
+                    ),
+                )
+
+                emr_args = self.generate_emr_submitter_args(
+                    # for now, self.conf is the only local file that requires uploading
+                    local_files_to_upload=local_files_to_upload_to_aws,
+                    user_args=user_args,
+                )
+                command = f"java -cp {self.jar_path} {EMR_ENTRY} {emr_args}"
+                command_list.append(command)
+
+        if len(command_list) > 1:
+            # parallel backfill mode
+            with multiprocessing.Pool(processes=int(self.parallelism)) as pool:
+                LOG.info(
+                    "Running args list {} with pool size {}".format(
+                        command_list, self.parallelism
+                    )
+                )
+                pool.map(check_call, command_list)
+        elif len(command_list) == 1:
+            # TODO: add log tailing
+            check_call(command_list[0])

ai/chronon/repo/cluster.py

@@ -0,0 +1,65 @@
+import json
+
+
+def generate_dataproc_cluster_config(num_workers, project_id, artifact_prefix, master_host_type="n2-highmem-64",
+                                     worker_host_type="n2-highmem-16",
+                                     subnetwork="default", idle_timeout="7200s", initialization_actions=None, tags=None):
+    """
+    Create a configuration for a Dataproc cluster.
+    :return: A json string representing the configuration.
+    """
+    if initialization_actions is None:
+        initialization_actions = []
+    return json.dumps({
+        "gceClusterConfig": {
+            "subnetworkUri": subnetwork,
+            "serviceAccount": "dataproc@" + project_id + ".iam.gserviceaccount.com",
+            "serviceAccountScopes": [
+                "https://www.googleapis.com/auth/cloud-platform",
+                "https://www.googleapis.com/auth/cloud.useraccounts.readonly",
+                "https://www.googleapis.com/auth/devstorage.read_write",
+                "https://www.googleapis.com/auth/logging.write"
+            ],
+            "metadata": {
+                "hive-version": "3.1.2",
+                "SPARK_BQ_CONNECTOR_URL": "gs://spark-lib/bigquery/spark-3.5-bigquery-0.42.1.jar",
+                "artifact_prefix": artifact_prefix.rstrip("/"),
+            },
+            "tags": tags or []
+        },
+        "masterConfig": {
+            "numInstances": 1,
+            "machineTypeUri": master_host_type,
+            "diskConfig": {
+                "bootDiskType": "pd-standard",
+                "bootDiskSizeGb": 1024
+            }
+        },
+        "workerConfig": {
+            "numInstances": num_workers,
+            "machineTypeUri": worker_host_type,
+            "diskConfig": {
+                "bootDiskType": "pd-standard",
+                "bootDiskSizeGb": 64,
+                "numLocalSsds": 2
+            }
+        },
+        "softwareConfig": {
+            "imageVersion": "2.2.50-debian12",
+            "optionalComponents": [
+                "FLINK",
+                "JUPYTER",
+            ],
+            "properties": {
+
+            }
+        },
+        "initializationActions": [{"executable_file": initialization_action} for initialization_action in (
+                    (initialization_actions or []) + [artifact_prefix.rstrip("/")+"/scripts/copy_java_security.sh"])],
+        "endpointConfig": {
+            "enableHttpPortAccess": True,
+        },
+        "lifecycleConfig": {
+            "idleDeleteTtl": idle_timeout,
+        }
+    })

ai/chronon/repo/compile.py

@@ -0,0 +1,56 @@
+import os
+import sys
+
+import click
+
+from ai.chronon.cli.compile.compile_context import CompileContext
+from ai.chronon.cli.compile.compiler import Compiler
+from ai.chronon.cli.compile.display.console import console
+
+
+@click.command(name="compile")
+@click.option(
+    "--chronon-root",
+    envvar="CHRONON_ROOT",
+    help="Path to the root chronon folder",
+    default=os.getcwd(),
+)
+def compile(chronon_root):
+
+    print()
+
+    if chronon_root not in sys.path:
+        console.print(
+            f"Adding [cyan italic]{chronon_root}[/cyan italic] to python path, during compile."
+        )
+        sys.path.append(chronon_root)
+    else:
+        console.print(
+            f"[cyan italic]{chronon_root}[/cyan italic] already on python path."
+        )
+
+    return __compile(chronon_root)
+
+
+def __compile(chronon_root):
+    if chronon_root:
+        chronon_root_path = os.path.expanduser(chronon_root)
+        os.chdir(chronon_root_path)
+
+    # check that a "teams.py" file exists in the current directory
+    if not (os.path.exists("teams.py") or os.path.exists("teams.json")):
+        raise click.ClickException(
+            (
+                "teams.py or teams.json file not found in current directory."
+                " Please run from the top level of conf directory."
+            )
+        )
+
+    compile_context = CompileContext()
+    compiler = Compiler(compile_context)
+    results = compiler.compile()
+    return results
+
+
+if __name__ == "__main__":
+    compile()