From a96b25cd8f98f0c5dcb60541eaff344d10ec27a5 Mon Sep 17 00:00:00 2001 From: Kyle Montemayor Date: Wed, 15 Jul 2026 18:04:41 +0000 Subject: [PATCH] Promote generic utils/constants from gigl.src.common to gigl.common Promote the reusable, application-agnostic helpers out of the pipeline source tree into the top-level gigl.common package: - gigl/common/utils/{timeout,file_loader,model,time,bq}.py - datetime format constants merged into gigl/common/constants.py Each old gigl.src.common path is left as a re-export shim that logs a deprecation warning, so external imports keep working. All internal call sites (gigl, examples, scripts, tests) are repointed to the new gigl.common locations; the dedicated bq tests move into the common test tree. Application/task logic (pb_wrappers, task types, launchers, pipeline path constants, translators, graph_builder, modeling_task_specs) stays in gigl.src by design. Co-Authored-By: Claude Opus 4.8 --- examples/MAG240M/preprocessor_config.py | 2 +- .../preprocessor_config_heterogeneous.py | 2 +- .../gowalla_data_preprocessor_config.py | 2 +- .../graph_store/heterogeneous_inference.py | 4 +- .../graph_store/heterogeneous_training.py | 2 +- .../graph_store/homogeneous_inference.py | 4 +- .../graph_store/homogeneous_training.py | 2 +- .../heterogeneous_inference.py | 4 +- .../link_prediction/heterogeneous_training.py | 2 +- .../link_prediction/homogeneous_inference.py | 4 +- .../link_prediction/homogeneous_training.py | 2 +- .../homogeneous_inference.py | 4 +- .../homogeneous_training.py | 2 +- .../KDD_2025/heterogeneous_inference.py | 2 +- .../KDD_2025/heterogeneous_training.py | 2 +- .../graph_validation/bq_graph_validator.py | 2 +- gigl/common/constants.py | 8 + gigl/common/data/dataloaders.py | 2 +- gigl/common/data/export.py | 2 +- gigl/common/utils/bq.py | 574 +++++++++++++++++ gigl/common/utils/file_loader.py | 322 ++++++++++ gigl/common/utils/jupyter_magics.py | 4 +- gigl/common/utils/model.py | 70 +++ gigl/common/utils/proto_utils.py | 2 +- gigl/common/utils/retry.py | 2 +- gigl/common/utils/time.py | 52 ++ gigl/common/utils/timeout.py | 64 ++ gigl/common/utils/yaml_loader.py | 2 +- gigl/distributed/utils/networking.py | 2 +- .../common/dist_checkpoint.py | 2 +- .../common/graph_dataset.py | 2 +- .../lib/config/parsing.py | 2 +- .../lib/constants/bq.py | 2 +- .../lib/data/edge_dataset.py | 2 +- .../kubeflow/kfp_orchestrator.py | 4 +- gigl/orchestration/kubeflow/runner.py | 2 +- .../kubeflow/utils/log_metrics.py | 2 +- gigl/scripts/load_gowalla_to_bq.py | 2 +- gigl/src/common/constants/bq.py | 2 +- gigl/src/common/constants/time.py | 25 +- .../model_eval_metrics_translator.py | 2 +- .../common/types/pb_wrappers/gbml_config.py | 2 +- gigl/src/common/utils/bq.py | 580 +----------------- gigl/src/common/utils/file_loader.py | 328 +--------- gigl/src/common/utils/model.py | 89 +-- gigl/src/common/utils/time.py | 73 +-- gigl/src/common/utils/timeout.py | 73 +-- .../data_preprocessor/data_preprocessor.py | 2 +- .../data_preprocessor/lib/enumerate/utils.py | 2 +- .../data_preprocessor/lib/transform/utils.py | 2 +- gigl/src/inference/lib/assets.py | 2 +- gigl/src/inference/v1/gnn_inferencer.py | 4 +- gigl/src/mocking/dataset_asset_mocker.py | 2 +- .../mocking/dataset_asset_mocking_suite.py | 2 +- gigl/src/mocking/lib/constants.py | 2 +- .../lib/mock_input_for_data_preprocessor.py | 2 +- .../mocking/lib/mock_input_for_inference.py | 2 +- .../lib/mock_input_for_subgraph_sampler.py | 2 +- .../mocking/lib/mock_output_for_inference.py | 4 +- gigl/src/mocking/lib/pyg_datasets_forks.py | 2 +- gigl/src/post_process/post_processor.py | 2 +- .../post_process/utils/cosine_similarity.py | 4 +- gigl/src/post_process/utils/unenumeration.py | 2 +- gigl/src/subgraph_sampler/subgraph_sampler.py | 2 +- gigl/src/training/trainer.py | 2 +- gigl/src/training/v1/lib/training_process.py | 8 +- .../libs/frozen_config_path_checks.py | 2 +- scripts/bootstrap_resource_config.py | 2 +- scripts/launch_graph_store_local.py | 2 +- tests/api_test/api_test_inference.py | 4 +- tests/api_test/api_test_training.py | 2 +- tests/config_tests/notebooks_test.py | 2 +- tests/integration/common/data/export_test.py | 2 +- tests/integration/common/file_loader_test.py | 2 +- tests/integration/common/utils/__init__.py | 0 .../{src => }/common/utils/bq_test.py | 2 +- .../distributed/utils/networking_test.py | 2 +- .../config_populator_pipeline_test.py | 4 +- .../data_preprocessor_pipeline_test.py | 6 +- .../data_preprocessor/enumerator_test.py | 6 +- .../pipeline/inferencer/inferencer_test.py | 4 +- .../split_generator_pipeline_test.py | 4 +- .../subgraph_sampler/subgraph_sampler_test.py | 2 +- .../pipeline/subgraph_sampler/utils.py | 2 +- tests/unit/common/data/export_test.py | 16 +- tests/unit/{src => }/common/utils/bq_test.py | 4 +- tests/unit/common/utils/retry_test.py | 2 +- tests/unit/common/utils/timeout_test.py | 2 +- .../unit/src/common/utils/gbml_config_test.py | 2 +- .../config_populator_functionality_test.py | 2 +- 90 files changed, 1291 insertions(+), 1179 deletions(-) create mode 100644 gigl/common/utils/bq.py create mode 100644 gigl/common/utils/file_loader.py create mode 100644 gigl/common/utils/model.py create mode 100644 gigl/common/utils/time.py create mode 100644 gigl/common/utils/timeout.py create mode 100644 tests/integration/common/utils/__init__.py rename tests/integration/{src => }/common/utils/bq_test.py (99%) rename tests/unit/{src => }/common/utils/bq_test.py (98%) diff --git a/examples/MAG240M/preprocessor_config.py b/examples/MAG240M/preprocessor_config.py index a3bc3c5c3..a794eaade 100644 --- a/examples/MAG240M/preprocessor_config.py +++ b/examples/MAG240M/preprocessor_config.py @@ -15,13 +15,13 @@ query_template_reindex_author_writes_paper_table, ) from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.env.pipelines_config import get_resource_config from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, EdgeUsageType, NodeType, Relation from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( GiglResourceConfigWrapper, ) -from gigl.src.common.utils.bq import BqUtils from gigl.src.data_preprocessor.data_preprocessor import DataPreprocessorConfig from gigl.src.data_preprocessor.lib.data_preprocessor_config import ( DataPreprocessorConfig, diff --git a/examples/MAG240M/preprocessor_config_heterogeneous.py b/examples/MAG240M/preprocessor_config_heterogeneous.py index 0bc88d0b8..e0cfc038c 100644 --- a/examples/MAG240M/preprocessor_config_heterogeneous.py +++ b/examples/MAG240M/preprocessor_config_heterogeneous.py @@ -5,13 +5,13 @@ from examples.MAG240M.common import NUM_PAPER_FEATURES from examples.MAG240M.queries import query_template_compute_average_features from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.env.pipelines_config import get_resource_config from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, EdgeUsageType, NodeType, Relation from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( GiglResourceConfigWrapper, ) -from gigl.src.common.utils.bq import BqUtils from gigl.src.data_preprocessor.lib.data_preprocessor_config import ( DataPreprocessorConfig, build_ingestion_feature_spec_fn, diff --git a/examples/id_embeddings/gowalla_data_preprocessor_config.py b/examples/id_embeddings/gowalla_data_preprocessor_config.py index 6467eb4df..826d601f4 100644 --- a/examples/id_embeddings/gowalla_data_preprocessor_config.py +++ b/examples/id_embeddings/gowalla_data_preprocessor_config.py @@ -17,13 +17,13 @@ from google.cloud.bigquery.job import WriteDisposition from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.env.pipelines_config import get_resource_config from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, EdgeUsageType, NodeType, Relation from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( GiglResourceConfigWrapper, ) -from gigl.src.common.utils.bq import BqUtils from gigl.src.data_preprocessor.lib.data_preprocessor_config import ( DataPreprocessorConfig, build_ingestion_feature_spec_fn, diff --git a/examples/link_prediction/graph_store/heterogeneous_inference.py b/examples/link_prediction/graph_store/heterogeneous_inference.py index 398d0772f..b24b4cc17 100644 --- a/examples/link_prediction/graph_store/heterogeneous_inference.py +++ b/examples/link_prediction/graph_store/heterogeneous_inference.py @@ -98,7 +98,9 @@ from gigl.common import GcsUri, Uri, UriFactory from gigl.common.data.export import EmbeddingExporter, load_embeddings_to_bigquery from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils +from gigl.common.utils.model import load_state_dict_from_uri from gigl.distributed.graph_store import ( GraphStoreInfo, RemoteDistDataset, @@ -110,8 +112,6 @@ from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.bq import BqUtils -from gigl.src.common.utils.model import load_state_dict_from_uri from gigl.src.inference.lib.assets import InferenceAssets from gigl.utils.sampling import parse_fanout diff --git a/examples/link_prediction/graph_store/heterogeneous_training.py b/examples/link_prediction/graph_store/heterogeneous_training.py index 749cd1537..5375eb767 100644 --- a/examples/link_prediction/graph_store/heterogeneous_training.py +++ b/examples/link_prediction/graph_store/heterogeneous_training.py @@ -91,6 +91,7 @@ from examples.link_prediction.models import init_example_gigl_heterogeneous_model from gigl.common import Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.common.utils.torch_training import is_distributed_available_and_initialized from gigl.distributed import DistABLPLoader from gigl.distributed.distributed_neighborloader import DistNeighborLoader @@ -113,7 +114,6 @@ EvalMetricType, ) from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.utils.iterator import InfiniteIterator from gigl.utils.sampling import parse_fanout diff --git a/examples/link_prediction/graph_store/homogeneous_inference.py b/examples/link_prediction/graph_store/homogeneous_inference.py index 157bc0045..56c0ec7e0 100644 --- a/examples/link_prediction/graph_store/homogeneous_inference.py +++ b/examples/link_prediction/graph_store/homogeneous_inference.py @@ -98,7 +98,9 @@ from gigl.common import GcsUri, Uri, UriFactory from gigl.common.data.export import EmbeddingExporter, load_embeddings_to_bigquery from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils +from gigl.common.utils.model import load_state_dict_from_uri from gigl.distributed.graph_store import ( GraphStoreInfo, RemoteDistDataset, @@ -110,8 +112,6 @@ from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.bq import BqUtils -from gigl.src.common.utils.model import load_state_dict_from_uri from gigl.src.inference.lib.assets import InferenceAssets from gigl.utils.sampling import parse_fanout diff --git a/examples/link_prediction/graph_store/homogeneous_training.py b/examples/link_prediction/graph_store/homogeneous_training.py index e3ce65a8d..6047040ab 100644 --- a/examples/link_prediction/graph_store/homogeneous_training.py +++ b/examples/link_prediction/graph_store/homogeneous_training.py @@ -135,6 +135,7 @@ from examples.link_prediction.models import init_example_gigl_homogeneous_model from gigl.common import Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.common.utils.torch_training import is_distributed_available_and_initialized from gigl.distributed import DistABLPLoader from gigl.distributed.distributed_neighborloader import DistNeighborLoader @@ -157,7 +158,6 @@ EvalMetricType, ) from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.utils.iterator import InfiniteIterator from gigl.utils.sampling import parse_fanout diff --git a/examples/link_prediction/heterogeneous_inference.py b/examples/link_prediction/heterogeneous_inference.py index 0c54b8f04..42c1add92 100644 --- a/examples/link_prediction/heterogeneous_inference.py +++ b/examples/link_prediction/heterogeneous_inference.py @@ -35,14 +35,14 @@ from gigl.common import GcsUri, Uri, UriFactory from gigl.common.data.export import EmbeddingExporter, load_embeddings_to_bigquery from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils +from gigl.common.utils.model import load_state_dict_from_uri from gigl.distributed import DistDataset, build_dataset_from_task_config_uri from gigl.nn import LinkPredictionGNN from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.bq import BqUtils -from gigl.src.common.utils.model import load_state_dict_from_uri from gigl.src.inference.lib.assets import InferenceAssets from gigl.utils.sampling import parse_fanout diff --git a/examples/link_prediction/heterogeneous_training.py b/examples/link_prediction/heterogeneous_training.py index 107554073..c9031733d 100644 --- a/examples/link_prediction/heterogeneous_training.py +++ b/examples/link_prediction/heterogeneous_training.py @@ -43,6 +43,7 @@ from examples.link_prediction.models import init_example_gigl_heterogeneous_model from gigl.common import Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.common.utils.torch_training import is_distributed_available_and_initialized from gigl.distributed import ( DistABLPLoader, @@ -62,7 +63,6 @@ EvalMetricType, ) from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.utils.iterator import InfiniteIterator from gigl.utils.sampling import parse_fanout diff --git a/examples/link_prediction/homogeneous_inference.py b/examples/link_prediction/homogeneous_inference.py index 084d2375e..6bbd53ffd 100644 --- a/examples/link_prediction/homogeneous_inference.py +++ b/examples/link_prediction/homogeneous_inference.py @@ -34,14 +34,14 @@ from gigl.common import GcsUri, Uri, UriFactory from gigl.common.data.export import EmbeddingExporter, load_embeddings_to_bigquery from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils +from gigl.common.utils.model import load_state_dict_from_uri from gigl.distributed import DistDataset, build_dataset_from_task_config_uri from gigl.nn import LinkPredictionGNN from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.bq import BqUtils -from gigl.src.common.utils.model import load_state_dict_from_uri from gigl.src.inference.lib.assets import InferenceAssets from gigl.utils.sampling import parse_fanout diff --git a/examples/link_prediction/homogeneous_training.py b/examples/link_prediction/homogeneous_training.py index 124879101..cfdbcb8fd 100644 --- a/examples/link_prediction/homogeneous_training.py +++ b/examples/link_prediction/homogeneous_training.py @@ -38,6 +38,7 @@ from examples.link_prediction.models import init_example_gigl_homogeneous_model from gigl.common import Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.common.utils.torch_training import is_distributed_available_and_initialized from gigl.distributed import ( DistABLPLoader, @@ -57,7 +58,6 @@ EvalMetricType, ) from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.types.graph import to_homogeneous from gigl.utils.iterator import InfiniteIterator from gigl.utils.sampling import parse_fanout diff --git a/examples/node_classification/homogeneous_inference.py b/examples/node_classification/homogeneous_inference.py index 0b0c3a125..539e45210 100644 --- a/examples/node_classification/homogeneous_inference.py +++ b/examples/node_classification/homogeneous_inference.py @@ -44,13 +44,13 @@ from gigl.common import GcsUri, Uri, UriFactory from gigl.common.data.export import PredictionExporter, load_predictions_to_bigquery from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils +from gigl.common.utils.model import load_state_dict_from_uri from gigl.distributed import DistDataset, build_dataset_from_task_config_uri from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.bq import BqUtils -from gigl.src.common.utils.model import load_state_dict_from_uri from gigl.src.inference.lib.assets import InferenceAssets from gigl.utils.sampling import parse_fanout diff --git a/examples/node_classification/homogeneous_training.py b/examples/node_classification/homogeneous_training.py index 5308ce1f3..62cd89094 100644 --- a/examples/node_classification/homogeneous_training.py +++ b/examples/node_classification/homogeneous_training.py @@ -48,6 +48,7 @@ ) from gigl.common import Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.common.utils.torch_training import is_distributed_available_and_initialized from gigl.distributed import DistDataset, build_dataset_from_task_config_uri from gigl.distributed.distributed_neighborloader import DistNeighborLoader @@ -62,7 +63,6 @@ EvalMetricType, ) from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.types.graph import to_homogeneous from gigl.utils.iterator import InfiniteIterator from gigl.utils.sampling import parse_fanout diff --git a/examples/tutorial/KDD_2025/heterogeneous_inference.py b/examples/tutorial/KDD_2025/heterogeneous_inference.py index 772d695e7..1aedb780d 100644 --- a/examples/tutorial/KDD_2025/heterogeneous_inference.py +++ b/examples/tutorial/KDD_2025/heterogeneous_inference.py @@ -38,6 +38,7 @@ from gigl.common import Uri, UriFactory from gigl.common.data.export import EmbeddingExporter from gigl.common.logger import Logger +from gigl.common.utils.model import load_state_dict_from_uri from gigl.distributed import ( DistDataset, DistNeighborLoader, @@ -45,7 +46,6 @@ ) from gigl.distributed.utils import get_free_port from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.model import load_state_dict_from_uri logger = Logger() diff --git a/examples/tutorial/KDD_2025/heterogeneous_training.py b/examples/tutorial/KDD_2025/heterogeneous_training.py index d8a09646c..a7be429a7 100644 --- a/examples/tutorial/KDD_2025/heterogeneous_training.py +++ b/examples/tutorial/KDD_2025/heterogeneous_training.py @@ -47,6 +47,7 @@ from examples.tutorial.KDD_2025.utils import LOCAL_SAVED_MODEL_URI, init_model from gigl.common import UriFactory from gigl.common.logger import Logger +from gigl.common.utils.model import save_state_dict from gigl.distributed import ( DistABLPLoader, DistDataset, @@ -55,7 +56,6 @@ from gigl.distributed.utils import get_free_port from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.model import save_state_dict from gigl.utils.iterator import InfiniteIterator logger = Logger() diff --git a/gigl/analytics/graph_validation/bq_graph_validator.py b/gigl/analytics/graph_validation/bq_graph_validator.py index 615daab69..0b293c4ab 100644 --- a/gigl/analytics/graph_validation/bq_graph_validator.py +++ b/gigl/analytics/graph_validation/bq_graph_validator.py @@ -3,7 +3,7 @@ from google.cloud.bigquery.table import RowIterator from gigl.common.logger import Logger -from gigl.src.common.utils.bq import BqUtils +from gigl.common.utils.bq import BqUtils logger = Logger() diff --git a/gigl/common/constants.py b/gigl/common/constants.py index f8c5afd33..a63b1f9bc 100644 --- a/gigl/common/constants.py +++ b/gigl/common/constants.py @@ -79,3 +79,11 @@ def parse_makefile_vars(makefile_path: Path) -> dict[str, str]: "DEFAULT_GIGL_RELEASE_KFP_PIPELINE_PATH" ] # =============================================================================================== + + +# Datetime format constants +DEFAULT_DATETIME_FORMAT: Final[str] = "%Y-%m-%d--%H-%M-%S-%Z" +DEFAULT_DATE_FORMAT: Final[str] = "%Y%m%d" +NODASH_DATE_FMT: Final[str] = "%Y%m%d" +NODASH_DATETIME_FORMAT: Final[str] = "%Y%m%d%H%M%S" +DASH_DATE_FMT: Final[str] = "%Y-%m-%d" diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 0ac83f0e7..c182faa3d 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -11,8 +11,8 @@ from gigl.common import Uri from gigl.common.logger import Logger from gigl.common.utils.decorator import tf_on_cpu +from gigl.common.utils.file_loader import FileLoader from gigl.src.common.types.features import FeatureTypes -from gigl.src.common.utils.file_loader import FileLoader from gigl.src.data_preprocessor.lib.types import FeatureSpecDict logger = Logger() diff --git a/gigl/common/data/export.py b/gigl/common/data/export.py index e4ebcca2a..406f5e316 100644 --- a/gigl/common/data/export.py +++ b/gigl/common/data/export.py @@ -24,8 +24,8 @@ from gigl.common import GcsUri, LocalUri, Uri from gigl.common.logger import Logger +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.retry import retry -from gigl.src.common.utils.file_loader import FileLoader logger = Logger() diff --git a/gigl/common/utils/bq.py b/gigl/common/utils/bq.py new file mode 100644 index 000000000..7850e5aee --- /dev/null +++ b/gigl/common/utils/bq.py @@ -0,0 +1,574 @@ +import datetime +import itertools +import re +from typing import Iterable, Optional, Tuple, Union + +import google.api_core.retry +import google.cloud.bigquery as bigquery +from google.api_core.exceptions import NotFound +from google.cloud.bigquery._helpers import _record_field_to_json +from google.cloud.bigquery.job import _AsyncJob +from google.cloud.bigquery.table import RowIterator + +from gigl.common import GcsUri, LocalUri, Uri +from gigl.common.constants import DEFAULT_DATE_FORMAT +from gigl.common.logger import Logger +from gigl.common.utils.retry import retry +from gigl.common.utils.time import convert_days_to_ms, current_datetime + +logger = Logger() + + +def _load_file_to_bq( + source: Uri, + bq_path: str, + client: bigquery.Client, + job_config: bigquery.LoadJobConfig, +) -> _AsyncJob: + try: + if isinstance(source, GcsUri): + # if gcs, need to use load_table_from_uri + load_job = client.load_table_from_uri( + source.uri, bq_path, job_config=job_config + ) + elif isinstance(source, LocalUri): + # if local, need to use load_table_from_file + with open(source.uri, "rb") as source_path_obj: + # API request -- starts the job + load_job = client.load_table_from_file( + source_path_obj, bq_path, job_config=job_config + ) + else: + raise ValueError(f"Unsupported source type: {type(source)}") + logger.info(f"Loading {source} to {bq_path}") + return load_job.result() # Waits for job to complete. + except Exception as e: + logger.exception(f"Could not load file to BQ. {repr(e)}") + raise e + + +@retry() +def _load_file_to_bq_with_retry( + source_path: Uri, + bq_path: str, + client: bigquery.Client, + job_config: bigquery.LoadJobConfig, +) -> _AsyncJob: + return _load_file_to_bq(source_path, bq_path, client, job_config) + + +class BqUtils: + def __init__(self, project: Optional[str] = None) -> None: + logger.info(f"BqUtils initialized with project: {project}") + self.__bq_client = bigquery.Client(project=project) + + def create_bq_dataset(self, dataset_id, exists_ok=True) -> None: + dataset = bigquery.Dataset(dataset_id) + dataset.location = "US" + try: + self.__bq_client.create_dataset(dataset, exists_ok=exists_ok) # API request + logger.info(f"Created dataset {dataset_id}") + except Exception as e: + logger.exception(f"Could not create dataset. {repr(e)}") + + def get_dataset_name_from_table(self, bq_path: str) -> str: + dataset_id = ".".join(bq_path.split(".")[:-1]) + return dataset_id + + def create_or_empty_bq_table( + self, bq_path: str, schema: Optional[list[bigquery.SchemaField]] = None + ) -> None: + bq_path = self.format_bq_path(bq_path) + split_bq_path = bq_path.split(".") + if len(split_bq_path) == 2: + dataset_name = split_bq_path[0] + elif len(split_bq_path) == 3: + dataset_name = split_bq_path[1] + else: + raise Exception(f"Could not parse BQ table path: {bq_path}") + self.__bq_client.create_dataset( + dataset=dataset_name, exists_ok=True + ) # No-Op if dataset exists + self.__bq_client.delete_table( + table=bq_path, not_found_ok=True + ) # Deletes if table exists + table: Union[str, bigquery.Table] = ( + bigquery.Table(table_ref=bq_path, schema=schema) if schema else bq_path + ) + self.__bq_client.create_table( + table=table, exists_ok=False + ) # Recreate the table + + def count_number_of_rows_in_bq_table( + self, + bq_table: str, + labels: dict[str, str] = {}, + ) -> int: + bq_table = bq_table.replace(":", ".") + ROW_COUNTING_QUERY = f""" + SELECT count(1) AS ct FROM `{bq_table}` + """ + result = self.run_query(query=ROW_COUNTING_QUERY, labels=labels) + for row in result: + n_rows = row["ct"] + return n_rows + + def count_number_of_columns_in_bq_table( + self, + bq_table: str, + ) -> int: + schema = self.fetch_bq_table_schema(bq_table=bq_table) + return len(schema.keys()) + + def run_query( + self, + query, + labels: dict[str, str], + **job_config_args, + ) -> RowIterator: + logger.info(f"Running query: {query}") + job_config = bigquery.QueryJobConfig(**job_config_args) + job_config.labels = labels + # Start the query, passing in the extra configuration. + try: + query_job = self.__bq_client.query( + query, location="US", job_config=job_config + ) # API request - starts the query + + # Waits for the query to finish and returns the result row iterator object. + result = query_job.result() + return result + except Exception as e: + logger.exception(f"Could not run query: {e}") + raise e + + @staticmethod + def format_bq_path(bq_path: str, format_for_table_reference: bool = False) -> str: + """Formats BQ paths. + + Args: + bq_path (str): expected to be one of: + ".." or ":.
" + "." or ":" + ".
" + format_for_table_reference (bool, optional): If project, dataset, and + table are all specified; add the `:` seperator between project and dataset. + Useful for when "table_reference" is required instead of path i.e. for + using BigQuery IO operator for beam pipelines. + Defaults to False. + + Returns: + str: Formatted bq path + """ + + bq_path = bq_path.replace(":", ".") + count_bq_path_parts = bq_path.count(".") + assert count_bq_path_parts > 0 and count_bq_path_parts < 3, ( + "BQ path expected to contain project + dataset and/or table." + ) + split_path = bq_path.split(".") + project = split_path[0] + dataset_and_table_name = ".".join(split_path[1:]) + return_path = ".".join([project, dataset_and_table_name]) + if format_for_table_reference and count_bq_path_parts == 2: + return_path = return_path.replace(".", ":", 1) + + return return_path + + @staticmethod + def join_path(path: str, *paths) -> str: + joined_path = ".".join([path, *paths]) + assert joined_path.count(".") <= 2, f"Invalid BQ path: {joined_path}" + return BqUtils.format_bq_path(joined_path) + + @staticmethod + def parse_bq_table_path(bq_table_path: str) -> Tuple[str, str, str]: + """ + Parses a joined bq table path into its project, dataset, and table names + Args: + bq_table_path (str): Joined bq table path of format `project.dataset.table` + Returns: + bq_project_id (str): Parsed BQ Project ID + bq_dataset_id (str): Parsed Dataset ID + bq_table_name (str): Parsed Table Name + """ + split_bq_table_path = BqUtils.format_bq_path(bq_table_path).split(".") + assert len(split_bq_table_path) == 3, ( + "bqtable_path should be in the format project.dataset.table" + ) + bq_project_id, bq_dataset_id, bq_table_name = split_bq_table_path + + return bq_project_id, bq_dataset_id, bq_table_name + + def update_bq_dataset_retention( + self, + bq_dataset_path: str, + retention_in_days: int, + apply_retroactively: Optional[bool] = False, + ) -> None: + """ + Update default retention for a whole BQ dataset. + This applies only to new tables unless apply_retroactively=True. + + :param bq_dataset_path: The BigQuery dataset path in the format `project_id.dataset_id`. + :param retention_in_days: The number of days to retain data in BigQuery tables. + :param apply_retroactively: If True, applies this retention policy retroactively to all existing tables in the dataset. + """ + bq_dataset_path = BqUtils.format_bq_path(bq_dataset_path) + dataset = self.__bq_client.get_dataset(bq_dataset_path) + retention_in_ms = convert_days_to_ms(retention_in_days) + + dataset.default_table_expiration_ms = retention_in_ms + try: + self.__bq_client.update_dataset(dataset, ["default_table_expiration_ms"]) + logger.info( + f"Updated dataset {bq_dataset_path} with default expiration in {retention_in_days} days." + ) + except Exception as e: + logger.exception(e) + + if apply_retroactively: + for table_item in self.__bq_client.list_tables(dataset): + table_id = table_item.full_table_id + self.update_bq_table_retention(table_id, retention_in_days) + + def update_bq_table_retention( + self, bq_table_path: str, retention_in_days: int + ) -> None: + """ + Update retention of a single BQ table. + :param bq_table_path: + :param retention_in_days: + :param client: + :return: + """ + + bq_table_path = BqUtils.format_bq_path(bq_table_path) + table = bigquery.Table(bq_table_path) + expiration_dt = current_datetime() + datetime.timedelta(days=retention_in_days) + table.expires = expiration_dt + try: + self.__bq_client.update_table(table, ["expires"]) + logger.info( + f"Updated table {bq_table_path} to expire in {retention_in_days} days." + ) + except Exception as e: + logger.exception(e) + + def does_bq_table_exist(self, bq_table_path: str) -> bool: + exists = False + try: + bq_table_path = BqUtils.format_bq_path(bq_table_path) + self.__bq_client.get_table(bq_table_path) # Make an API request. + exists = True + logger.info(f"Table {bq_table_path} exists.") + except NotFound: + logger.info(f"Table {bq_table_path} not found.") + except Exception as e: + logger.info(f"Could not evaluate table existence. {repr(e)}") + return exists + + def list_matching_tables( + self, bq_dataset_path: str, table_match_string: str + ) -> list[str]: + bq_dataset_path = BqUtils.format_bq_path(bq_dataset_path) + tables = self.__bq_client.list_tables(bq_dataset_path) + matching_tables = list() + for table_list_item in tables: + if table_match_string in table_list_item.table_id: + formatted_table_path = BqUtils.format_bq_path( + table_list_item.full_table_id + ) + matching_tables.append(formatted_table_path) + return matching_tables + + def get_latest_table( + self, + bq_table_path_prefix: str, + table_partition_suffix: str = "YYYYMMDD", + cap_date: Optional[str] = None, + ) -> str: + """Return the latest table matching *bq_table_path_prefix*. + + Tables are expected to have a partition suffix whose length matches + *table_partition_suffix*. All supported GCP partition suffixes (``YYYY``, + ``YYYYMM``, ``YYYYMMDD``, ``YYYYMMDDHH``, integer ranges) are + lexicographically sortable, so the latest table is the + lexicographic maximum. + See: https://docs.cloud.google.com/bigquery/docs/partitioned-tables#date_timestamp_partitioned_tables + + + If *cap_date* is given, only tables with suffix <= cap_date are + considered. + + Args: + bq_table_path_prefix: Fully qualified prefix in + ``"project.dataset.table_prefix"`` format. + table_partition_suffix: A representative suffix string whose length + determines how many trailing characters form the sortable + partition key. Defaults to ``"YYYYMMDD"`` (8 characters). + cap_date: Optional upper-bound date string that is in the format of + *table_partition_suffix*. Only tables with a suffix <= this value are + considered. + + Returns: + The fully qualified table path of the latest matching table. + """ + project_id, dataset_name, table_prefix = BqUtils.parse_bq_table_path( + bq_table_path_prefix + ) + + bq_dataset_path = BqUtils.join_path(project_id, dataset_name) + tables = self.__bq_client.list_tables(bq_dataset_path) + + suffix_len = len(table_partition_suffix) + expected_table_path_len = len(bq_table_path_prefix) + suffix_len + latest_table_path: str + curr_latest_candidate_suffix: Optional[str] = None + for table_list_item in tables: + if not table_list_item.table_id.startswith(table_prefix): + continue + table_path = BqUtils.format_bq_path(table_list_item.full_table_id) + if len(table_path) != expected_table_path_len: + continue + suffix = table_path[-suffix_len:] + if cap_date is None or suffix <= cap_date: + if ( + curr_latest_candidate_suffix is None + or suffix > curr_latest_candidate_suffix + ): + curr_latest_candidate_suffix = suffix + latest_table_path = table_path + if curr_latest_candidate_suffix is None: + raise ValueError( + f"No tables found with prefix {bq_table_path_prefix} and cap date {cap_date}" + ) + return latest_table_path + + def delete_matching_tables( + self, bq_dataset_path: str, table_match_string: str + ) -> None: + try: + bq_dataset_path = BqUtils.format_bq_path(bq_dataset_path) + tables = self.__bq_client.list_tables(bq_dataset_path) + for table_list_item in tables: + if table_match_string in table_list_item.table_id: + formatted_table_path = BqUtils.format_bq_path( + table_list_item.full_table_id + ) + self.__bq_client.delete_table(formatted_table_path) + logger.info(f"Deleted table {formatted_table_path}.") + except Exception as e: + logger.exception("Error in deleting tables." + repr(e)) + + def get_table_names_within_date_range( + self, + bq_dataset_path: str, + table_match_string: str, + start_date: str, + end_date: str, + ) -> list[str]: + """ + start_date and end_date are in the format of 'YYYYMMDD' + table_match_string is a regex string to match table names + """ + _start_date = datetime.datetime.strptime(start_date, DEFAULT_DATE_FORMAT).date() + _end_date = datetime.datetime.strptime(end_date, DEFAULT_DATE_FORMAT).date() + filtered_tables_by_name = list() + filtered_tables_by_date = list() + bq_dataset_path = BqUtils.format_bq_path(bq_dataset_path) + all_tables = self.__bq_client.list_tables(bq_dataset_path) + for table in all_tables: + if re.search(table_match_string, table.table_id): + filtered_tables_by_name.append(table) + sorted_tables_by_date = sorted( + filtered_tables_by_name, key=lambda x: x.created, reverse=True + ) + for table in sorted_tables_by_date: + if _start_date <= table.created.date() <= _end_date: + filtered_tables_by_date.append( + ".".join([bq_dataset_path, table.table_id]) + ) + return filtered_tables_by_date + + def delete_bq_table_if_exist( + self, bq_table_path: str, not_found_ok: bool = True + ) -> None: + """bq_table_path = 'your-project.your_dataset.your_table'""" + bq_table_path = BqUtils.format_bq_path(bq_table_path) + try: + self.__bq_client.delete_table(bq_table_path, not_found_ok=not_found_ok) + logger.info(f"Table deleted '{bq_table_path}'") + except Exception as e: + logger.exception(f"Failed to delete table '{bq_table_path}' due to \n {e}") + + def fetch_bq_table_schema(self, bq_table: str) -> dict[str, bigquery.SchemaField]: + """ + Create a dictionary representation for SchemaFields from BigQuery table. + """ + + bq_table = bq_table.replace(":", ".") + bq_schema = self.__bq_client.get_table(bq_table).schema + schema_dict = {field.name: field for field in bq_schema} + return schema_dict + + def load_file_to_bq( + self, + source_path: Uri, + bq_path: str, + job_config: bigquery.LoadJobConfig, + retry: bool = False, + ) -> _AsyncJob: + """ + Uploads a single file to biqquery. + + Args: + source_path (Uri): The source file to upload. + bq_path (str): The BigQuery table path to upload to. + job_config (bigquery.LoadJobConfig): The job configuration for the upload. + retry (bool, optional): Whether to retry the upload if it fails. Defaults to False. + Returns: The job object for the upload. + """ + if retry: + result = _load_file_to_bq_with_retry( + source_path, bq_path, self.__bq_client, job_config + ) + else: + result = _load_file_to_bq( + source_path, bq_path, self.__bq_client, job_config + ) + return result + + def load_rows_to_bq( + self, + bq_path: str, + schema: list[bigquery.SchemaField], + rows: Iterable[Tuple], + ) -> None: + first_item = next(iter(rows), None) + if first_item is None: + logger.warning(f"No rows to insert into {bq_path}.") + return + + _BQ_INSERT_REQUEST_LIMIT_BYTES = ( + 10_000_000 - 1_000_000 + ) # 10MB is the limit; we leave 1MB of buffer + _RETRY_BACKOFF = google.api_core.retry.Retry( + # retries if and only if the 'reason' is 'backendError' or 'rateLimitExceeded' + predicate=bigquery.DEFAULT_RETRY._predicate, + initial=1.0, # 1 second + maximum=60.0 * 5, # 5 minutes + multiplier=1.5, + timeout=60 * 30, # 30 mins + ) + + table = bigquery.Table(table_ref=bq_path, schema=schema) + batch_rows = [] + estimated_batch_rows_size_bytes = 0 + + for row in itertools.chain([first_item], rows): + json_row: dict = _record_field_to_json(schema, row) + batch_rows.append(json_row) + estimated_batch_rows_size_bytes += len(str(json_row)) + if estimated_batch_rows_size_bytes > _BQ_INSERT_REQUEST_LIMIT_BYTES: + self.__bq_client.insert_rows_json( + table=table, + json_rows=batch_rows, + retry=_RETRY_BACKOFF, + ) + batch_rows = [] + estimated_batch_rows_size_bytes = 0 + + if len(batch_rows) > 0: + self.__bq_client.insert_rows_json( + table=table, + json_rows=batch_rows, + retry=_RETRY_BACKOFF, + ) + + def check_columns_exist_in_table( + self, bq_table: str, columns: Iterable[str] + ) -> None: + schema = self.fetch_bq_table_schema(bq_table=bq_table) + all_fields = set(schema.keys()) + missing_fields = [field for field in columns if field not in all_fields] + if missing_fields: + raise ValueError(f"Fields {missing_fields} missing from table {bq_table}.") + else: + logger.info(f"All requisite fields found in table {bq_table}") + + def copy_table( + self, + source_table: str, + destination_table: str, + overwrite: bool = True, + ) -> None: + """Copy a BigQuery table using the copy_table API (more efficient than SELECT). + + This method copies the table at the storage level without processing data through + the query engine, making it faster and more cost-effective than using CREATE TABLE AS SELECT. + + See https://cloud.google.com/bigquery/docs/managing-tables#copy-table for more info. + Args: + source_table (str): Fully qualified source table path (project.dataset.table) + destination_table (str): Fully qualified destination table path (project.dataset.table) + overwrite (bool, optional): Whether to overwrite if destination exists. Defaults to True. + """ + source_table = self.format_bq_path(source_table) + destination_table = self.format_bq_path(destination_table) + + job_config = bigquery.CopyJobConfig( + write_disposition="WRITE_TRUNCATE" if overwrite else "WRITE_EMPTY" + ) + + try: + logger.info(f"Copying table from {source_table} to {destination_table}") + copy_job = self.__bq_client.copy_table( + source_table, destination_table, job_config=job_config + ) + copy_job.result() # Wait for completion + logger.info( + f"Successfully copied table from {source_table} to {destination_table}" + ) + except Exception as e: + logger.exception( + f"Failed to copy table from {source_table} to {destination_table}" + ) + raise e + + def export_to_gcs( + self, + bq_table_path: str, + destination_gcs_uri: GcsUri, + destination_format: str = "NEWLINE_DELIMITED_JSON", + ) -> None: + """ + Export a BigQuery table to Google Cloud Storage. + + Args: + bq_table_path (str): The full BigQuery table path to export. + destination_gcs_uri (str): The destination GCS URI where the table will be exported. + If the gcs uri has * in it, the table will be exported to multiple shards. + destination_format (str, optional): The format of the exported data. Defaults to 'NEWLINE_DELIMITED_JSON'. + 'CSV', 'AVRO', 'PARQUET' also available. + """ + try: + job_config = bigquery.job.ExtractJobConfig() + job_config.destination_format = destination_format + + extract_job = self.__bq_client.extract_table( + source=bigquery.TableReference.from_string(bq_table_path), + destination_uris=destination_gcs_uri.uri, + job_config=job_config, + ) + + logger.info( + f"Exporting `{bq_table_path}` to {destination_gcs_uri} with format '{destination_format}'..." + ) + extract_job.result() # Waits for job to complete. + logger.info( + f"Exported `{bq_table_path}` to {destination_gcs_uri} successfully." + ) + except Exception as e: + logger.exception(f"Failed to export table to GCS.") + raise e diff --git a/gigl/common/utils/file_loader.py b/gigl/common/utils/file_loader.py new file mode 100644 index 000000000..15cc5478d --- /dev/null +++ b/gigl/common/utils/file_loader.py @@ -0,0 +1,322 @@ +import shutil +import tempfile +from collections.abc import Mapping +from tempfile import _TemporaryFileWrapper as TemporaryFileWrapper +from typing import IO, AnyStr, Optional, Sequence, Tuple, Type, Union, cast + +from gigl.common import GcsUri, HttpUri, LocalUri, Uri, UriFactory +from gigl.common.logger import Logger +from gigl.common.utils.gcs import GcsUtils +from gigl.common.utils.http import HttpUtils +from gigl.common.utils.local_fs import ( + FileSystemEntity, + copy_files, + count_files_with_uri_prefix, + create_file_symlinks, + does_path_exist, + list_at_path, + remove_file_or_folder_if_exist, +) + +logger = Logger() + + +class FileLoader: + def __init__(self, project: Optional[str] = None): + gcs_utils = GcsUtils(project) + self.__gcs_utils = gcs_utils + self.__unsupported_uri_message = ( + f"{self.__class__.__name__} does not support Uris of this type." + ) + + @staticmethod + def __get_uri_map_schema( + uri_map: Mapping[Uri, Uri], + ) -> Tuple[Optional[Type[Uri]], Optional[Type[Uri]]]: + uniform_src_type: Optional[Type[Uri]] = None + uniform_dst_type: Optional[Type[Uri]] = None + src_types: list[Type[Uri]] = [uri.__class__ for uri in uri_map.keys()] + dst_types: list[Type[Uri]] = [uri.__class__ for uri in uri_map.values()] + if all([src_types[0] == x for x in src_types]): + uniform_src_type = src_types[0] + if all([dst_types[0] == x for x in dst_types]): + uniform_dst_type = dst_types[0] + return uniform_src_type, uniform_dst_type + + def load_directories( + self, + source_to_dest_directory_map: dict[Uri, Uri], + ): + for dir_uri_src, dir_uri_dst in source_to_dest_directory_map.items(): + self.load_directory(dir_uri_src=dir_uri_src, dir_uri_dst=dir_uri_dst) + + def load_directory(self, dir_uri_src: Uri, dir_uri_dst: Uri): + uri_map_schema = self.__get_uri_map_schema(uri_map={dir_uri_src: dir_uri_dst}) + + if uri_map_schema == (GcsUri, LocalUri): + dir_uri_src = cast(GcsUri, dir_uri_src) + dir_uri_dst = cast(LocalUri, dir_uri_dst) + self.__gcs_utils.download_files_from_gcs_paths_to_local_dir( + gcs_paths=[dir_uri_src], local_path_dir=dir_uri_dst + ) + elif uri_map_schema == (LocalUri, GcsUri): + dir_uri_src = cast(LocalUri, dir_uri_src) + dir_uri_dst = cast(GcsUri, dir_uri_dst) + local_paths: list[LocalUri] = list_at_path( + local_path=dir_uri_src, file_system_entity=FileSystemEntity.FILE + ) + gcs_paths: list[GcsUri] = [ + GcsUri.join(dir_uri_dst, local_fn.uri) + for local_fn in list_at_path(dir_uri_src, names_only=True) + ] + local_file_path_to_gcs_path_map: dict[LocalUri, GcsUri] = { + src: dst for src, dst in zip(local_paths, gcs_paths) + } + self.load_files( + source_to_dest_file_uri_map=cast( + dict[Uri, Uri], local_file_path_to_gcs_path_map + ) + ) + elif uri_map_schema == (LocalUri, LocalUri): + dir_uri_src = cast(LocalUri, dir_uri_src) + dir_uri_dst = cast(LocalUri, dir_uri_dst) + + local_src_paths: list[LocalUri] = list_at_path( + local_path=dir_uri_src, file_system_entity=FileSystemEntity.FILE + ) + local_dst_paths: list[LocalUri] = [ + LocalUri.join(dir_uri_dst, local_src_fn) + for local_src_fn in list_at_path( + local_path=dir_uri_src, + names_only=True, + file_system_entity=FileSystemEntity.FILE, + ) + ] + source_to_dest_file_uri_map = { + src: dst for src, dst in zip(local_src_paths, local_dst_paths) + } + self.load_files( + source_to_dest_file_uri_map=cast( + dict[Uri, Uri], source_to_dest_file_uri_map + ) + ) + else: + raise TypeError(self.__unsupported_uri_message) + + def load_files( + self, + source_to_dest_file_uri_map: Mapping[Uri, Uri], + should_create_symlinks_if_possible: bool = True, + ) -> None: + uri_map_schema = self.__get_uri_map_schema(uri_map=source_to_dest_file_uri_map) + + if uri_map_schema == (GcsUri, LocalUri): + logger.info("Downloading from GCS to Local") + self.__gcs_utils.download_files_from_gcs_paths_to_local_paths( + file_map=cast(dict[GcsUri, LocalUri], source_to_dest_file_uri_map) + ) + elif uri_map_schema == (LocalUri, GcsUri): + logger.info("Uploading from Local to GCS") + self.__gcs_utils.upload_files_to_gcs( + local_file_path_to_gcs_path_map=cast( + dict[LocalUri, GcsUri], source_to_dest_file_uri_map + ), + parallel=True, + ) + elif uri_map_schema == (LocalUri, LocalUri): + logger.info("Copying from Local to Local") + local_source_to_link_path_map = source_to_dest_file_uri_map + if should_create_symlinks_if_possible: + logger.info("Will create symlinks") + create_file_symlinks( + local_source_to_link_path_map=cast( + dict[LocalUri, LocalUri], local_source_to_link_path_map + ), + should_overwrite=True, + ) + else: + logger.info("Will copy files") + copy_files( + local_source_to_local_dst_path_map=cast( + dict[LocalUri, LocalUri], local_source_to_link_path_map + ), + should_overwrite=True, + ) + elif uri_map_schema == (HttpUri, LocalUri): + logger.info("Downloading from HTTP to Local") + HttpUtils.download_files_from_http( + http_to_local_path_map=cast( + dict[HttpUri, LocalUri], source_to_dest_file_uri_map + ), + ) + else: + for file_uri_src, file_uri_dst in source_to_dest_file_uri_map.items(): + self.load_file( + file_uri_src=file_uri_src, + file_uri_dst=file_uri_dst, + should_create_symlinks_if_possible=should_create_symlinks_if_possible, + ) + + def load_file( + self, + file_uri_src: Uri, + file_uri_dst: Uri, + should_create_symlinks_if_possible: bool = True, + ) -> None: + if isinstance(file_uri_src, GcsUri) and isinstance(file_uri_dst, GcsUri): + self.__gcs_utils.copy_gcs_blob(file_uri_src, file_uri_dst) + elif isinstance(file_uri_src, GcsUri) and isinstance(file_uri_dst, LocalUri): + self.__gcs_utils.download_file_from_gcs( + gcs_path=file_uri_src, + dest_file_path=file_uri_dst, + ) + elif isinstance(file_uri_src, HttpUri) and isinstance(file_uri_dst, LocalUri): + HttpUtils.download_files_from_http( + http_to_local_path_map={file_uri_src: file_uri_dst} + ) + elif isinstance(file_uri_src, LocalUri) and isinstance(file_uri_dst, GcsUri): + self.__gcs_utils.upload_files_to_gcs( + local_file_path_to_gcs_path_map={file_uri_src: file_uri_dst}, + parallel=False, + ) + elif isinstance(file_uri_src, LocalUri) and isinstance(file_uri_dst, LocalUri): + local_source_to_link_path_map = {file_uri_src: file_uri_dst} + if should_create_symlinks_if_possible: + create_file_symlinks( + local_source_to_link_path_map=local_source_to_link_path_map, + should_overwrite=True, + ) + else: + copy_files( + local_source_to_local_dst_path_map=local_source_to_link_path_map, + should_overwrite=True, + ) + else: + logger.error( + f"Unsupported src/dst combo: {type(file_uri_src)} -> {type(file_uri_dst)}" + ) + raise TypeError(self.__unsupported_uri_message) + + def load_from_filelike(self, uri: Uri, filelike: IO[AnyStr]) -> None: + """ + Load a file from a file-like object into the specified URI. + + Args: + uri (Uri): The URI to load the file into. + filelike (IO[AnyStr]): A file-like object containing the data to be loaded. + """ + if isinstance(uri, GcsUri): + self.__gcs_utils.upload_from_filelike(gcs_path=uri, filelike=filelike) + elif isinstance(uri, LocalUri): + ptr = filelike.tell() + first = filelike.read(1) + filelike.seek(ptr) # Reset the file pointer after reading + if isinstance(first, bytes): + with open(uri.uri, "wb") as dest: + shutil.copyfileobj( + filelike, # ty: ignore[invalid-argument-type] + dest, + ) + else: + with open(uri.uri, "w", encoding="utf-8") as dest: + shutil.copyfileobj( + filelike, # ty: ignore[invalid-argument-type] + dest, + ) + + else: + raise NotImplementedError( + f"Cannot load file from buffer to URI {uri.uri} of type {type(uri)}; {self.__unsupported_uri_message}" + ) + + def load_to_temp_file( + self, + file_uri_src: Uri, + delete: bool = False, + should_create_symlinks_if_possible: bool = True, + ) -> TemporaryFileWrapper: + temp_file_handle = tempfile.NamedTemporaryFile(delete=delete) + temp_file_path = LocalUri(str(temp_file_handle.name)) + self.load_file( + file_uri_src=file_uri_src, + file_uri_dst=temp_file_path, + should_create_symlinks_if_possible=should_create_symlinks_if_possible, + ) + return temp_file_handle + + def count_assets(self, uri_prefix: Uri, suffix: Optional[str] = None) -> int: + if isinstance(uri_prefix, GcsUri): + return self.__gcs_utils.count_blobs_in_gcs_path( + gcs_path=uri_prefix, suffix=suffix + ) + elif isinstance(uri_prefix, LocalUri): + return count_files_with_uri_prefix(uri_prefix=uri_prefix, suffix=suffix) + else: + raise TypeError( + f"Uri type not supported, got {uri_prefix} in type {type(uri_prefix)}" + ) + + def does_uri_exist(self, uri: Union[str, Uri]) -> bool: + ( + """"" + Check if a URI exists + + Args: + uri (Union[str, Uri]): uri to check + Returns: + bool: True if URI exists, False otherwise + """ + "" + ) + + _uri = UriFactory.create_uri(uri=uri) if isinstance(uri, str) else uri + + exists: bool + if isinstance(_uri, GcsUri): + exists = self.__gcs_utils.does_gcs_file_exist(gcs_path=_uri) + elif isinstance(_uri, LocalUri): + exists = does_path_exist(_uri) + elif isinstance(_uri, HttpUri): + exists = HttpUtils.does_http_path_resolve(http_path=_uri) + else: + raise NotImplementedError(f"{self.__unsupported_uri_message} : {_uri}") + return exists + + def delete_files(self, uris: list[Uri]) -> None: + """ + Recursively delete files in the specified URIs. + + Args: + uris (list[Uri]): URIs to delete + Returns + None + """ + for uri in uris: + if isinstance(uri, LocalUri): + remove_file_or_folder_if_exist(local_path=uri) + elif isinstance(uri, GcsUri): + self.__gcs_utils.delete_files_in_bucket_dir(gcs_path=uri) + else: + raise NotImplementedError( + f"Cannot delete URI {uri.uri} of type {type(uri)}; {self.__unsupported_uri_message}" + ) + + def list_children(self, uri: Uri, pattern: Optional[str] = None) -> Sequence[Uri]: + """ + List all children of the given URI. + Args: + uri (Uri): The URI to list children of. + pattern (Optional[str]): Optional regex to match. If not provided then all children will be returned. + Returns: + list[Uri]: A list of URIs for the children of the given URI. + """ + if isinstance(uri, GcsUri): + return self.__gcs_utils.list_uris_with_gcs_path_pattern( + gcs_path=uri, pattern=pattern + ) + elif isinstance(uri, LocalUri): + return list_at_path(local_path=uri, regex=pattern) + else: + raise NotImplementedError( + f"Cannot list children of URI {uri.uri} of type {type(uri)}; {self.__unsupported_uri_message}" + ) diff --git a/gigl/common/utils/jupyter_magics.py b/gigl/common/utils/jupyter_magics.py index a010ca437..4f4f755b4 100644 --- a/gigl/common/utils/jupyter_magics.py +++ b/gigl/common/utils/jupyter_magics.py @@ -17,11 +17,11 @@ from gigl.common import Uri from gigl.common.collections.frozen_dict import FrozenDict +from gigl.common.utils.file_loader import FileLoader from gigl.src.common.graph_builder.pyg_graph_builder import PygGraphBuilder from gigl.src.common.translators.gbml_protos_translator import GbmlProtosTranslator from gigl.src.common.types.graph_data import CondensedNodeType, EdgeType, Node, NodeType from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.file_loader import FileLoader from snapchat.research.gbml import training_samples_schema_pb2 gigl_root_dir = pathlib.Path(__file__).parent.parent.parent.parent.parent @@ -60,7 +60,7 @@ def __init__(self, frozen_task_config: GbmlConfigPbWrapper): ) graph_metadata_pb_wrapper = frozen_task_config.graph_metadata_pb_wrapper - from gigl.src.common.utils.bq import BqUtils + from gigl.common.utils.bq import BqUtils bq_utils = BqUtils() diff --git a/gigl/common/utils/model.py b/gigl/common/utils/model.py new file mode 100644 index 000000000..5ab48ab3a --- /dev/null +++ b/gigl/common/utils/model.py @@ -0,0 +1,70 @@ +import tempfile +from collections import OrderedDict + +import torch + +from gigl.common import LocalUri, Uri +from gigl.common.utils.file_loader import FileLoader + + +def save_scripted_model(model: torch.nn.Module, save_to_path_uri: Uri) -> None: + assert isinstance(model, torch.nn.Module), ( + "Can only save model of type torch.nn.Module" + ) + file_loader = FileLoader() + tmp_save_model_file = tempfile.NamedTemporaryFile(delete=False) + + model_scripted = torch.jit.script(model) # Export to TorchScript + model_scripted.save(tmp_save_model_file.name) # Save + + file_loader.load_file( + file_uri_src=LocalUri(tmp_save_model_file.name), file_uri_dst=save_to_path_uri + ) + tmp_save_model_file.close() + + +def save_state_dict(model: torch.nn.Module, save_to_path_uri: Uri) -> None: + if isinstance(model, torch.nn.parallel.DistributedDataParallel): + model = model.module + + assert isinstance(model, torch.nn.Module), ( + "Can only save model of type torch.nn.Module" + ) + file_loader = FileLoader() + + tmp_save_model_file = tempfile.NamedTemporaryFile(delete=False) + torch.save(model.state_dict(), tmp_save_model_file.name) + + file_loader.load_file( + file_uri_src=LocalUri(tmp_save_model_file.name), file_uri_dst=save_to_path_uri + ) + tmp_save_model_file.close() + + +def load_state_dict_from_uri( + load_from_uri: Uri, + device: torch.device = torch.device("cpu"), +) -> OrderedDict[str, torch.Tensor]: + state_dict: OrderedDict[str, torch.Tensor] + + file_loader = FileLoader() + tmp_file = file_loader.load_to_temp_file(load_from_uri) + + state_dict = torch.load(tmp_file.name, map_location=device) + tmp_file.close() + + return state_dict + + +def load_scripted_model_from_uri( + load_from_uri: Uri, +) -> torch.nn.Module: + scripted_model: torch.jit.ScriptModule + + file_loader = FileLoader() + tmp_file = file_loader.load_to_temp_file(load_from_uri) + + scripted_model = torch.jit.load(tmp_file.name) + tmp_file.close() + + return scripted_model diff --git a/gigl/common/utils/proto_utils.py b/gigl/common/utils/proto_utils.py index e61258756..74afeb076 100644 --- a/gigl/common/utils/proto_utils.py +++ b/gigl/common/utils/proto_utils.py @@ -9,7 +9,7 @@ from gigl.common import LocalUri, Uri from gigl.common.logger import Logger from gigl.common.omegaconf_resolvers import register_resolvers -from gigl.src.common.utils.file_loader import FileLoader +from gigl.common.utils.file_loader import FileLoader logger = Logger() diff --git a/gigl/common/utils/retry.py b/gigl/common/utils/retry.py index 6be63cb4e..6b6209715 100644 --- a/gigl/common/utils/retry.py +++ b/gigl/common/utils/retry.py @@ -4,7 +4,7 @@ from xmlrpc.client import Boolean from gigl.common.logger import Logger -from gigl.src.common.utils.timeout import TimedOutException, timeout +from gigl.common.utils.timeout import TimedOutException, timeout T = TypeVar("T") # return type diff --git a/gigl/common/utils/time.py b/gigl/common/utils/time.py new file mode 100644 index 000000000..f2d1f4031 --- /dev/null +++ b/gigl/common/utils/time.py @@ -0,0 +1,52 @@ +import datetime + +from gigl.common.constants import DEFAULT_DATETIME_FORMAT + + +def parse_formatted_datetime( + stringified_date: str, fmt: str = DEFAULT_DATETIME_FORMAT +) -> datetime.datetime: + return datetime.datetime.strptime(stringified_date, fmt) + + +def format_datetime(dt: datetime.datetime, fmt: str = DEFAULT_DATETIME_FORMAT) -> str: + return dt.strftime(fmt) + + +def current_formatted_datetime(fmt: str = DEFAULT_DATETIME_FORMAT) -> str: + """ + Returns: + str: current timezone aware utc datetime formatted as a string + as per provided date format string. + """ + return current_datetime().strftime(fmt) + + +def current_datetime() -> datetime.datetime: + """ + Returns: + str: current timezone aware utc datetime. + """ + return datetime.datetime.now(datetime.timezone.utc) + + +def convert_days_to_ms(days: int) -> int: + return days * 24 * 60 * 60 * 1000 + + +def is_datetime_str_format_valid(datetime_str: str, datetime_format: str) -> bool: + """ + Validates a datetime string with specified format. ex "%Y%m%d" for "YYYYmmdd" + Return True/False so users can decide to raise error or reformat string + + :param datetime_str: + :param datetime_format: + :return: + """ + try: + datetime.datetime.strptime(datetime_str, datetime_format) + return True + except ValueError: + return False + except Exception as e: + return False diff --git a/gigl/common/utils/timeout.py b/gigl/common/utils/timeout.py new file mode 100644 index 000000000..b57893c2c --- /dev/null +++ b/gigl/common/utils/timeout.py @@ -0,0 +1,64 @@ +import errno +import os +import signal +from functools import wraps + +from gigl.common.logger import Logger + +logger = Logger() + + +class TimedOutException(Exception): + pass + + +def timeout( + seconds=10, + error_message=os.strerror(errno.ETIME), + timeout_action_func=None, + exception_thrown_on_timeout=TimedOutException, + **timeout_action_func_params, +): + """ + Decorator to exit a program when a function execution exceeds a specified timeout. + + This decorator exits the program when the decorated function timed out, and executes + timeout_action_func before exiting. The timeout_action_func can be useful for cases + like environment clean up (e.g., ray.shutdown()). Another way to handle timeouts, + especially when there are multiple threads or child threads, is to run `func` with + multiprocessing. However, multiprocessing does not work with Ray. + + References to the current timeout implementation: + - https://docs.python.org/3/library/signal.html + - https://www.saltycrane.com/blog/2010/04/using-python-timeout-decorator-uploading-s3/ + + Parameters: + - seconds (int): Timeout duration in seconds. Defaults to 10. + - error_message (str): Error message to log on timeout. Defaults to os.strerror(errno.ETIME). + - timeout_action_func (callable, optional): Function to execute before exiting on timeout. + - exception_thrown_on_timeout (Exception): Exception to raise on timeout. Defaults to TimedOutException. + - **timeout_action_func_params: Arbitrary keyword arguments passed to timeout_action_func. + """ + + def decorator(func): + def _handler(signum, frame): + logger.info(error_message) + if timeout_action_func: + timeout_action_func(**timeout_action_func_params) + raise exception_thrown_on_timeout() + + def wrapper(*args, **kwargs): + old = signal.signal(signal.SIGALRM, _handler) + signal.alarm(seconds) + try: + result = func(*args, **kwargs) + finally: + # cancel the alarm + signal.alarm(0) + # reinstall the old signal handler + signal.signal(signal.SIGALRM, old) + return result + + return wraps(func)(wrapper) + + return decorator diff --git a/gigl/common/utils/yaml_loader.py b/gigl/common/utils/yaml_loader.py index 0dda4f208..bb9b0d9fc 100644 --- a/gigl/common/utils/yaml_loader.py +++ b/gigl/common/utils/yaml_loader.py @@ -5,7 +5,7 @@ from gigl.common import Uri from gigl.common.logger import Logger from gigl.common.omegaconf_resolvers import register_resolvers -from gigl.src.common.utils.file_loader import FileLoader +from gigl.common.utils.file_loader import FileLoader logger = Logger() diff --git a/gigl/distributed/utils/networking.py b/gigl/distributed/utils/networking.py index ea204f570..e042cbf48 100644 --- a/gigl/distributed/utils/networking.py +++ b/gigl/distributed/utils/networking.py @@ -8,6 +8,7 @@ from gigl.common import Uri from gigl.common.logger import Logger +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.vertex_ai_context import ( ClusterSpec, get_cluster_spec, @@ -20,7 +21,6 @@ GraphStoreInfo, ) from gigl.env.pipelines_config import get_resource_config -from gigl.src.common.utils.file_loader import FileLoader logger = Logger() diff --git a/gigl/experimental/knowledge_graph_embedding/common/dist_checkpoint.py b/gigl/experimental/knowledge_graph_embedding/common/dist_checkpoint.py index b1e4b170d..f4abe78ea 100644 --- a/gigl/experimental/knowledge_graph_embedding/common/dist_checkpoint.py +++ b/gigl/experimental/knowledge_graph_embedding/common/dist_checkpoint.py @@ -15,8 +15,8 @@ from gigl.common import GcsUri, LocalUri, Uri from gigl.common.logger import Logger +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.local_fs import delete_local_directory -from gigl.src.common.utils.file_loader import FileLoader logger = Logger() diff --git a/gigl/experimental/knowledge_graph_embedding/common/graph_dataset.py b/gigl/experimental/knowledge_graph_embedding/common/graph_dataset.py index 71666a1a9..88e72fead 100644 --- a/gigl/experimental/knowledge_graph_embedding/common/graph_dataset.py +++ b/gigl/experimental/knowledge_graph_embedding/common/graph_dataset.py @@ -13,8 +13,8 @@ from gigl.common.types.uri.gcs_uri import GcsUri from gigl.common.types.uri.uri_factory import UriFactory +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.torch_training import get_rank, get_world_size -from gigl.src.common.utils.file_loader import FileLoader from gigl.src.training.v1.lib.data_loaders.utils import ( get_data_split_for_current_worker, ) diff --git a/gigl/experimental/knowledge_graph_embedding/lib/config/parsing.py b/gigl/experimental/knowledge_graph_embedding/lib/config/parsing.py index d43bf8ddf..716c7a99e 100644 --- a/gigl/experimental/knowledge_graph_embedding/lib/config/parsing.py +++ b/gigl/experimental/knowledge_graph_embedding/lib/config/parsing.py @@ -10,12 +10,12 @@ import gigl.experimental.knowledge_graph_embedding.lib.constants.local as local_constants from gigl.common import GcsUri, LocalUri, UriFactory from gigl.common.types.uri.uri import Uri +from gigl.common.utils.file_loader import FileLoader from gigl.env.pipelines_config import get_resource_config from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( GiglResourceConfigWrapper, ) -from gigl.src.common.utils.file_loader import FileLoader def build_modeling_and_resource_config_from_args( diff --git a/gigl/experimental/knowledge_graph_embedding/lib/constants/bq.py b/gigl/experimental/knowledge_graph_embedding/lib/constants/bq.py index e16f690b3..7e4246a80 100644 --- a/gigl/experimental/knowledge_graph_embedding/lib/constants/bq.py +++ b/gigl/experimental/knowledge_graph_embedding/lib/constants/bq.py @@ -1,7 +1,7 @@ import gigl.src.common.constants.bq as bq_constants +from gigl.common.utils.bq import BqUtils from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType -from gigl.src.common.utils.bq import BqUtils def get_src_node_embedding_table_for_edge_type( diff --git a/gigl/experimental/knowledge_graph_embedding/lib/data/edge_dataset.py b/gigl/experimental/knowledge_graph_embedding/lib/data/edge_dataset.py index 9480813bc..61514e108 100644 --- a/gigl/experimental/knowledge_graph_embedding/lib/data/edge_dataset.py +++ b/gigl/experimental/knowledge_graph_embedding/lib/data/edge_dataset.py @@ -8,6 +8,7 @@ import gigl.experimental.knowledge_graph_embedding.lib.constants.gcs as gcs_constants from gigl.common.logger import Logger from gigl.common.types.uri.gcs_uri import GcsUri +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils from gigl.common.utils.torch_training import is_distributed_available_and_initialized from gigl.distributed.dist_context import DistributedContext @@ -24,7 +25,6 @@ from gigl.src.common.types.dataset_split import DatasetSplit from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation from gigl.src.common.types.pb_wrappers.graph_metadata import GraphMetadataPbWrapper -from gigl.src.common.utils.bq import BqUtils from gigl.src.data_preprocessor.lib.enumerate.utils import EnumeratorEdgeTypeMetadata from gigl.src.data_preprocessor.lib.ingest.bigquery import BigqueryEdgeDataReference diff --git a/gigl/orchestration/kubeflow/kfp_orchestrator.py b/gigl/orchestration/kubeflow/kfp_orchestrator.py index 374347ac1..08544f16c 100644 --- a/gigl/orchestration/kubeflow/kfp_orchestrator.py +++ b/gigl/orchestration/kubeflow/kfp_orchestrator.py @@ -12,12 +12,12 @@ from gigl.common.logger import Logger from gigl.common.services.vertex_ai import VertexAIService from gigl.common.types.resource_config import CommonPipelineComponentConfigs +from gigl.common.utils.file_loader import FileLoader +from gigl.common.utils.time import current_formatted_datetime from gigl.env.pipelines_config import get_resource_config from gigl.orchestration.kubeflow.kfp_pipeline import generate_pipeline from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.types import AppliedTaskIdentifier -from gigl.src.common.utils.file_loader import FileLoader -from gigl.src.common.utils.time import current_formatted_datetime from gigl.src.validation_check.libs.name_checks import ( check_if_kfp_pipeline_job_name_valid, ) diff --git a/gigl/orchestration/kubeflow/runner.py b/gigl/orchestration/kubeflow/runner.py index 790b000dc..41c377e1c 100644 --- a/gigl/orchestration/kubeflow/runner.py +++ b/gigl/orchestration/kubeflow/runner.py @@ -84,6 +84,7 @@ DEFAULT_GIGL_RELEASE_SRC_IMAGE_DATAFLOW_CPU, ) from gigl.common.logger import Logger +from gigl.common.utils.time import current_formatted_datetime from gigl.orchestration.img_builder import build_and_push_customer_src_images from gigl.orchestration.kubeflow.kfp_orchestrator import ( DEFAULT_KFP_COMPILED_PIPELINE_DEST_PATH, @@ -92,7 +93,6 @@ from gigl.orchestration.kubeflow.kfp_pipeline import SPECED_COMPONENTS from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.types import AppliedTaskIdentifier -from gigl.src.common.utils.time import current_formatted_datetime DEFAULT_JOB_NAME = f"gigl_run_at_{current_formatted_datetime()}" DEFAULT_START_AT = GiGLComponents.ConfigPopulator.value diff --git a/gigl/orchestration/kubeflow/utils/log_metrics.py b/gigl/orchestration/kubeflow/utils/log_metrics.py index 124803dec..4d3e051a0 100644 --- a/gigl/orchestration/kubeflow/utils/log_metrics.py +++ b/gigl/orchestration/kubeflow/utils/log_metrics.py @@ -42,9 +42,9 @@ def _log_eval_metrics_to_ui( from gigl.common import UriFactory from gigl.common.logger import Logger + from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.proto_utils import ProtoUtils from gigl.src.common.constants.components import GiGLComponents - from gigl.src.common.utils.file_loader import FileLoader from snapchat.research.gbml import gbml_config_pb2 logger = Logger() diff --git a/gigl/scripts/load_gowalla_to_bq.py b/gigl/scripts/load_gowalla_to_bq.py index 3df7386e4..1c188024b 100644 --- a/gigl/scripts/load_gowalla_to_bq.py +++ b/gigl/scripts/load_gowalla_to_bq.py @@ -28,7 +28,7 @@ from gigl.common import UriFactory from gigl.common.logger import Logger -from gigl.src.common.utils.bq import BqUtils +from gigl.common.utils.bq import BqUtils logger = Logger() diff --git a/gigl/src/common/constants/bq.py b/gigl/src/common/constants/bq.py index ab39e6886..c2afc45c1 100644 --- a/gigl/src/common/constants/bq.py +++ b/gigl/src/common/constants/bq.py @@ -1,7 +1,7 @@ +from gigl.common.utils.bq import BqUtils from gigl.env.pipelines_config import get_resource_config from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import NodeType -from gigl.src.common.utils.bq import BqUtils def get_embeddings_dataset_bq_path() -> str: diff --git a/gigl/src/common/constants/time.py b/gigl/src/common/constants/time.py index 9eef7436e..b540f2129 100644 --- a/gigl/src/common/constants/time.py +++ b/gigl/src/common/constants/time.py @@ -1,5 +1,20 @@ -DEFAULT_DATETIME_FORMAT = "%Y-%m-%d--%H-%M-%S-%Z" -DEFAULT_DATE_FORMAT = "%Y%m%d" -NODASH_DATE_FMT = "%Y%m%d" -NODASH_DATETIME_FORMAT = "%Y%m%d%H%M%S" -DASH_DATE_FMT = "%Y-%m-%d" +"""Deprecated shim - datetime format constants moved to :mod:`gigl.common.constants`. + +Importing from here continues to work but logs a deprecation warning. This shim +will be removed in a future release. +""" + +from gigl.common.constants import ( # noqa: F401 (re-export for backwards compat) + DASH_DATE_FMT, + DEFAULT_DATE_FORMAT, + DEFAULT_DATETIME_FORMAT, + NODASH_DATE_FMT, + NODASH_DATETIME_FORMAT, +) +from gigl.common.logger import Logger + +logger = Logger() +logger.warning( + "gigl.common.constants is deprecated and will be removed in a future " + "release. Please import from `gigl.common.constants` instead." +) diff --git a/gigl/src/common/translators/model_eval_metrics_translator.py b/gigl/src/common/translators/model_eval_metrics_translator.py index 7214128b3..17002a321 100644 --- a/gigl/src/common/translators/model_eval_metrics_translator.py +++ b/gigl/src/common/translators/model_eval_metrics_translator.py @@ -5,8 +5,8 @@ from gigl.common import LocalUri, Uri from gigl.common.logger import Logger +from gigl.common.utils.file_loader import FileLoader from gigl.src.common.types.model_eval_metrics import EvalMetricsCollection -from gigl.src.common.utils.file_loader import FileLoader logger = Logger() diff --git a/gigl/src/common/types/pb_wrappers/gbml_config.py b/gigl/src/common/types/pb_wrappers/gbml_config.py index f153a6805..c728ea6a5 100644 --- a/gigl/src/common/types/pb_wrappers/gbml_config.py +++ b/gigl/src/common/types/pb_wrappers/gbml_config.py @@ -6,6 +6,7 @@ from gigl.common import Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.proto_utils import ProtoUtils from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.common.types.pb_wrappers.dataset_metadata import DatasetMetadataPbWrapper @@ -23,7 +24,6 @@ from gigl.src.common.types.pb_wrappers.trained_model_metadata import ( TrainedModelMetadataPbWrapper, ) -from gigl.src.common.utils.file_loader import FileLoader from gigl.src.data_preprocessor.lib.types import FeatureSchema from snapchat.research.gbml import ( dataset_metadata_pb2, diff --git a/gigl/src/common/utils/bq.py b/gigl/src/common/utils/bq.py index dfdec91b9..b29876062 100644 --- a/gigl/src/common/utils/bq.py +++ b/gigl/src/common/utils/bq.py @@ -1,574 +1,16 @@ -import datetime -import itertools -import re -from typing import Iterable, Optional, Tuple, Union +"""Deprecated shim - implementation moved to :mod:`gigl.common.utils.bq`. -import google.api_core.retry -import google.cloud.bigquery as bigquery -from google.api_core.exceptions import NotFound -from google.cloud.bigquery._helpers import _record_field_to_json -from google.cloud.bigquery.job import _AsyncJob -from google.cloud.bigquery.table import RowIterator +Importing from here continues to work but logs a deprecation warning. This shim +will be removed in a future release. +""" -from gigl.common import GcsUri, LocalUri, Uri from gigl.common.logger import Logger -from gigl.common.utils.retry import retry -from gigl.src.common.constants.time import DEFAULT_DATE_FORMAT -from gigl.src.common.utils.time import convert_days_to_ms, current_datetime +from gigl.common.utils.bq import ( # noqa: F401 (re-export for backwards compat) + BqUtils, +) logger = Logger() - - -def _load_file_to_bq( - source: Uri, - bq_path: str, - client: bigquery.Client, - job_config: bigquery.LoadJobConfig, -) -> _AsyncJob: - try: - if isinstance(source, GcsUri): - # if gcs, need to use load_table_from_uri - load_job = client.load_table_from_uri( - source.uri, bq_path, job_config=job_config - ) - elif isinstance(source, LocalUri): - # if local, need to use load_table_from_file - with open(source.uri, "rb") as source_path_obj: - # API request -- starts the job - load_job = client.load_table_from_file( - source_path_obj, bq_path, job_config=job_config - ) - else: - raise ValueError(f"Unsupported source type: {type(source)}") - logger.info(f"Loading {source} to {bq_path}") - return load_job.result() # Waits for job to complete. - except Exception as e: - logger.exception(f"Could not load file to BQ. {repr(e)}") - raise e - - -@retry() -def _load_file_to_bq_with_retry( - source_path: Uri, - bq_path: str, - client: bigquery.Client, - job_config: bigquery.LoadJobConfig, -) -> _AsyncJob: - return _load_file_to_bq(source_path, bq_path, client, job_config) - - -class BqUtils: - def __init__(self, project: Optional[str] = None) -> None: - logger.info(f"BqUtils initialized with project: {project}") - self.__bq_client = bigquery.Client(project=project) - - def create_bq_dataset(self, dataset_id, exists_ok=True) -> None: - dataset = bigquery.Dataset(dataset_id) - dataset.location = "US" - try: - self.__bq_client.create_dataset(dataset, exists_ok=exists_ok) # API request - logger.info(f"Created dataset {dataset_id}") - except Exception as e: - logger.exception(f"Could not create dataset. {repr(e)}") - - def get_dataset_name_from_table(self, bq_path: str) -> str: - dataset_id = ".".join(bq_path.split(".")[:-1]) - return dataset_id - - def create_or_empty_bq_table( - self, bq_path: str, schema: Optional[list[bigquery.SchemaField]] = None - ) -> None: - bq_path = self.format_bq_path(bq_path) - split_bq_path = bq_path.split(".") - if len(split_bq_path) == 2: - dataset_name = split_bq_path[0] - elif len(split_bq_path) == 3: - dataset_name = split_bq_path[1] - else: - raise Exception(f"Could not parse BQ table path: {bq_path}") - self.__bq_client.create_dataset( - dataset=dataset_name, exists_ok=True - ) # No-Op if dataset exists - self.__bq_client.delete_table( - table=bq_path, not_found_ok=True - ) # Deletes if table exists - table: Union[str, bigquery.Table] = ( - bigquery.Table(table_ref=bq_path, schema=schema) if schema else bq_path - ) - self.__bq_client.create_table( - table=table, exists_ok=False - ) # Recreate the table - - def count_number_of_rows_in_bq_table( - self, - bq_table: str, - labels: dict[str, str] = {}, - ) -> int: - bq_table = bq_table.replace(":", ".") - ROW_COUNTING_QUERY = f""" - SELECT count(1) AS ct FROM `{bq_table}` - """ - result = self.run_query(query=ROW_COUNTING_QUERY, labels=labels) - for row in result: - n_rows = row["ct"] - return n_rows - - def count_number_of_columns_in_bq_table( - self, - bq_table: str, - ) -> int: - schema = self.fetch_bq_table_schema(bq_table=bq_table) - return len(schema.keys()) - - def run_query( - self, - query, - labels: dict[str, str], - **job_config_args, - ) -> RowIterator: - logger.info(f"Running query: {query}") - job_config = bigquery.QueryJobConfig(**job_config_args) - job_config.labels = labels - # Start the query, passing in the extra configuration. - try: - query_job = self.__bq_client.query( - query, location="US", job_config=job_config - ) # API request - starts the query - - # Waits for the query to finish and returns the result row iterator object. - result = query_job.result() - return result - except Exception as e: - logger.exception(f"Could not run query: {e}") - raise e - - @staticmethod - def format_bq_path(bq_path: str, format_for_table_reference: bool = False) -> str: - """Formats BQ paths. - - Args: - bq_path (str): expected to be one of: - "..
" or ":.
" - "." or ":" - ".
" - format_for_table_reference (bool, optional): If project, dataset, and - table are all specified; add the `:` seperator between project and dataset. - Useful for when "table_reference" is required instead of path i.e. for - using BigQuery IO operator for beam pipelines. - Defaults to False. - - Returns: - str: Formatted bq path - """ - - bq_path = bq_path.replace(":", ".") - count_bq_path_parts = bq_path.count(".") - assert count_bq_path_parts > 0 and count_bq_path_parts < 3, ( - "BQ path expected to contain project + dataset and/or table." - ) - split_path = bq_path.split(".") - project = split_path[0] - dataset_and_table_name = ".".join(split_path[1:]) - return_path = ".".join([project, dataset_and_table_name]) - if format_for_table_reference and count_bq_path_parts == 2: - return_path = return_path.replace(".", ":", 1) - - return return_path - - @staticmethod - def join_path(path: str, *paths) -> str: - joined_path = ".".join([path, *paths]) - assert joined_path.count(".") <= 2, f"Invalid BQ path: {joined_path}" - return BqUtils.format_bq_path(joined_path) - - @staticmethod - def parse_bq_table_path(bq_table_path: str) -> Tuple[str, str, str]: - """ - Parses a joined bq table path into its project, dataset, and table names - Args: - bq_table_path (str): Joined bq table path of format `project.dataset.table` - Returns: - bq_project_id (str): Parsed BQ Project ID - bq_dataset_id (str): Parsed Dataset ID - bq_table_name (str): Parsed Table Name - """ - split_bq_table_path = BqUtils.format_bq_path(bq_table_path).split(".") - assert len(split_bq_table_path) == 3, ( - "bqtable_path should be in the format project.dataset.table" - ) - bq_project_id, bq_dataset_id, bq_table_name = split_bq_table_path - - return bq_project_id, bq_dataset_id, bq_table_name - - def update_bq_dataset_retention( - self, - bq_dataset_path: str, - retention_in_days: int, - apply_retroactively: Optional[bool] = False, - ) -> None: - """ - Update default retention for a whole BQ dataset. - This applies only to new tables unless apply_retroactively=True. - - :param bq_dataset_path: The BigQuery dataset path in the format `project_id.dataset_id`. - :param retention_in_days: The number of days to retain data in BigQuery tables. - :param apply_retroactively: If True, applies this retention policy retroactively to all existing tables in the dataset. - """ - bq_dataset_path = BqUtils.format_bq_path(bq_dataset_path) - dataset = self.__bq_client.get_dataset(bq_dataset_path) - retention_in_ms = convert_days_to_ms(retention_in_days) - - dataset.default_table_expiration_ms = retention_in_ms - try: - self.__bq_client.update_dataset(dataset, ["default_table_expiration_ms"]) - logger.info( - f"Updated dataset {bq_dataset_path} with default expiration in {retention_in_days} days." - ) - except Exception as e: - logger.exception(e) - - if apply_retroactively: - for table_item in self.__bq_client.list_tables(dataset): - table_id = table_item.full_table_id - self.update_bq_table_retention(table_id, retention_in_days) - - def update_bq_table_retention( - self, bq_table_path: str, retention_in_days: int - ) -> None: - """ - Update retention of a single BQ table. - :param bq_table_path: - :param retention_in_days: - :param client: - :return: - """ - - bq_table_path = BqUtils.format_bq_path(bq_table_path) - table = bigquery.Table(bq_table_path) - expiration_dt = current_datetime() + datetime.timedelta(days=retention_in_days) - table.expires = expiration_dt - try: - self.__bq_client.update_table(table, ["expires"]) - logger.info( - f"Updated table {bq_table_path} to expire in {retention_in_days} days." - ) - except Exception as e: - logger.exception(e) - - def does_bq_table_exist(self, bq_table_path: str) -> bool: - exists = False - try: - bq_table_path = BqUtils.format_bq_path(bq_table_path) - self.__bq_client.get_table(bq_table_path) # Make an API request. - exists = True - logger.info(f"Table {bq_table_path} exists.") - except NotFound: - logger.info(f"Table {bq_table_path} not found.") - except Exception as e: - logger.info(f"Could not evaluate table existence. {repr(e)}") - return exists - - def list_matching_tables( - self, bq_dataset_path: str, table_match_string: str - ) -> list[str]: - bq_dataset_path = BqUtils.format_bq_path(bq_dataset_path) - tables = self.__bq_client.list_tables(bq_dataset_path) - matching_tables = list() - for table_list_item in tables: - if table_match_string in table_list_item.table_id: - formatted_table_path = BqUtils.format_bq_path( - table_list_item.full_table_id - ) - matching_tables.append(formatted_table_path) - return matching_tables - - def get_latest_table( - self, - bq_table_path_prefix: str, - table_partition_suffix: str = "YYYYMMDD", - cap_date: Optional[str] = None, - ) -> str: - """Return the latest table matching *bq_table_path_prefix*. - - Tables are expected to have a partition suffix whose length matches - *table_partition_suffix*. All supported GCP partition suffixes (``YYYY``, - ``YYYYMM``, ``YYYYMMDD``, ``YYYYMMDDHH``, integer ranges) are - lexicographically sortable, so the latest table is the - lexicographic maximum. - See: https://docs.cloud.google.com/bigquery/docs/partitioned-tables#date_timestamp_partitioned_tables - - - If *cap_date* is given, only tables with suffix <= cap_date are - considered. - - Args: - bq_table_path_prefix: Fully qualified prefix in - ``"project.dataset.table_prefix"`` format. - table_partition_suffix: A representative suffix string whose length - determines how many trailing characters form the sortable - partition key. Defaults to ``"YYYYMMDD"`` (8 characters). - cap_date: Optional upper-bound date string that is in the format of - *table_partition_suffix*. Only tables with a suffix <= this value are - considered. - - Returns: - The fully qualified table path of the latest matching table. - """ - project_id, dataset_name, table_prefix = BqUtils.parse_bq_table_path( - bq_table_path_prefix - ) - - bq_dataset_path = BqUtils.join_path(project_id, dataset_name) - tables = self.__bq_client.list_tables(bq_dataset_path) - - suffix_len = len(table_partition_suffix) - expected_table_path_len = len(bq_table_path_prefix) + suffix_len - latest_table_path: str - curr_latest_candidate_suffix: Optional[str] = None - for table_list_item in tables: - if not table_list_item.table_id.startswith(table_prefix): - continue - table_path = BqUtils.format_bq_path(table_list_item.full_table_id) - if len(table_path) != expected_table_path_len: - continue - suffix = table_path[-suffix_len:] - if cap_date is None or suffix <= cap_date: - if ( - curr_latest_candidate_suffix is None - or suffix > curr_latest_candidate_suffix - ): - curr_latest_candidate_suffix = suffix - latest_table_path = table_path - if curr_latest_candidate_suffix is None: - raise ValueError( - f"No tables found with prefix {bq_table_path_prefix} and cap date {cap_date}" - ) - return latest_table_path - - def delete_matching_tables( - self, bq_dataset_path: str, table_match_string: str - ) -> None: - try: - bq_dataset_path = BqUtils.format_bq_path(bq_dataset_path) - tables = self.__bq_client.list_tables(bq_dataset_path) - for table_list_item in tables: - if table_match_string in table_list_item.table_id: - formatted_table_path = BqUtils.format_bq_path( - table_list_item.full_table_id - ) - self.__bq_client.delete_table(formatted_table_path) - logger.info(f"Deleted table {formatted_table_path}.") - except Exception as e: - logger.exception("Error in deleting tables." + repr(e)) - - def get_table_names_within_date_range( - self, - bq_dataset_path: str, - table_match_string: str, - start_date: str, - end_date: str, - ) -> list[str]: - """ - start_date and end_date are in the format of 'YYYYMMDD' - table_match_string is a regex string to match table names - """ - _start_date = datetime.datetime.strptime(start_date, DEFAULT_DATE_FORMAT).date() - _end_date = datetime.datetime.strptime(end_date, DEFAULT_DATE_FORMAT).date() - filtered_tables_by_name = list() - filtered_tables_by_date = list() - bq_dataset_path = BqUtils.format_bq_path(bq_dataset_path) - all_tables = self.__bq_client.list_tables(bq_dataset_path) - for table in all_tables: - if re.search(table_match_string, table.table_id): - filtered_tables_by_name.append(table) - sorted_tables_by_date = sorted( - filtered_tables_by_name, key=lambda x: x.created, reverse=True - ) - for table in sorted_tables_by_date: - if _start_date <= table.created.date() <= _end_date: - filtered_tables_by_date.append( - ".".join([bq_dataset_path, table.table_id]) - ) - return filtered_tables_by_date - - def delete_bq_table_if_exist( - self, bq_table_path: str, not_found_ok: bool = True - ) -> None: - """bq_table_path = 'your-project.your_dataset.your_table'""" - bq_table_path = BqUtils.format_bq_path(bq_table_path) - try: - self.__bq_client.delete_table(bq_table_path, not_found_ok=not_found_ok) - logger.info(f"Table deleted '{bq_table_path}'") - except Exception as e: - logger.exception(f"Failed to delete table '{bq_table_path}' due to \n {e}") - - def fetch_bq_table_schema(self, bq_table: str) -> dict[str, bigquery.SchemaField]: - """ - Create a dictionary representation for SchemaFields from BigQuery table. - """ - - bq_table = bq_table.replace(":", ".") - bq_schema = self.__bq_client.get_table(bq_table).schema - schema_dict = {field.name: field for field in bq_schema} - return schema_dict - - def load_file_to_bq( - self, - source_path: Uri, - bq_path: str, - job_config: bigquery.LoadJobConfig, - retry: bool = False, - ) -> _AsyncJob: - """ - Uploads a single file to biqquery. - - Args: - source_path (Uri): The source file to upload. - bq_path (str): The BigQuery table path to upload to. - job_config (bigquery.LoadJobConfig): The job configuration for the upload. - retry (bool, optional): Whether to retry the upload if it fails. Defaults to False. - Returns: The job object for the upload. - """ - if retry: - result = _load_file_to_bq_with_retry( - source_path, bq_path, self.__bq_client, job_config - ) - else: - result = _load_file_to_bq( - source_path, bq_path, self.__bq_client, job_config - ) - return result - - def load_rows_to_bq( - self, - bq_path: str, - schema: list[bigquery.SchemaField], - rows: Iterable[Tuple], - ) -> None: - first_item = next(iter(rows), None) - if first_item is None: - logger.warning(f"No rows to insert into {bq_path}.") - return - - _BQ_INSERT_REQUEST_LIMIT_BYTES = ( - 10_000_000 - 1_000_000 - ) # 10MB is the limit; we leave 1MB of buffer - _RETRY_BACKOFF = google.api_core.retry.Retry( - # retries if and only if the 'reason' is 'backendError' or 'rateLimitExceeded' - predicate=bigquery.DEFAULT_RETRY._predicate, - initial=1.0, # 1 second - maximum=60.0 * 5, # 5 minutes - multiplier=1.5, - timeout=60 * 30, # 30 mins - ) - - table = bigquery.Table(table_ref=bq_path, schema=schema) - batch_rows = [] - estimated_batch_rows_size_bytes = 0 - - for row in itertools.chain([first_item], rows): - json_row: dict = _record_field_to_json(schema, row) - batch_rows.append(json_row) - estimated_batch_rows_size_bytes += len(str(json_row)) - if estimated_batch_rows_size_bytes > _BQ_INSERT_REQUEST_LIMIT_BYTES: - self.__bq_client.insert_rows_json( - table=table, - json_rows=batch_rows, - retry=_RETRY_BACKOFF, - ) - batch_rows = [] - estimated_batch_rows_size_bytes = 0 - - if len(batch_rows) > 0: - self.__bq_client.insert_rows_json( - table=table, - json_rows=batch_rows, - retry=_RETRY_BACKOFF, - ) - - def check_columns_exist_in_table( - self, bq_table: str, columns: Iterable[str] - ) -> None: - schema = self.fetch_bq_table_schema(bq_table=bq_table) - all_fields = set(schema.keys()) - missing_fields = [field for field in columns if field not in all_fields] - if missing_fields: - raise ValueError(f"Fields {missing_fields} missing from table {bq_table}.") - else: - logger.info(f"All requisite fields found in table {bq_table}") - - def copy_table( - self, - source_table: str, - destination_table: str, - overwrite: bool = True, - ) -> None: - """Copy a BigQuery table using the copy_table API (more efficient than SELECT). - - This method copies the table at the storage level without processing data through - the query engine, making it faster and more cost-effective than using CREATE TABLE AS SELECT. - - See https://cloud.google.com/bigquery/docs/managing-tables#copy-table for more info. - Args: - source_table (str): Fully qualified source table path (project.dataset.table) - destination_table (str): Fully qualified destination table path (project.dataset.table) - overwrite (bool, optional): Whether to overwrite if destination exists. Defaults to True. - """ - source_table = self.format_bq_path(source_table) - destination_table = self.format_bq_path(destination_table) - - job_config = bigquery.CopyJobConfig( - write_disposition="WRITE_TRUNCATE" if overwrite else "WRITE_EMPTY" - ) - - try: - logger.info(f"Copying table from {source_table} to {destination_table}") - copy_job = self.__bq_client.copy_table( - source_table, destination_table, job_config=job_config - ) - copy_job.result() # Wait for completion - logger.info( - f"Successfully copied table from {source_table} to {destination_table}" - ) - except Exception as e: - logger.exception( - f"Failed to copy table from {source_table} to {destination_table}" - ) - raise e - - def export_to_gcs( - self, - bq_table_path: str, - destination_gcs_uri: GcsUri, - destination_format: str = "NEWLINE_DELIMITED_JSON", - ) -> None: - """ - Export a BigQuery table to Google Cloud Storage. - - Args: - bq_table_path (str): The full BigQuery table path to export. - destination_gcs_uri (str): The destination GCS URI where the table will be exported. - If the gcs uri has * in it, the table will be exported to multiple shards. - destination_format (str, optional): The format of the exported data. Defaults to 'NEWLINE_DELIMITED_JSON'. - 'CSV', 'AVRO', 'PARQUET' also available. - """ - try: - job_config = bigquery.job.ExtractJobConfig() - job_config.destination_format = destination_format - - extract_job = self.__bq_client.extract_table( - source=bigquery.TableReference.from_string(bq_table_path), - destination_uris=destination_gcs_uri.uri, - job_config=job_config, - ) - - logger.info( - f"Exporting `{bq_table_path}` to {destination_gcs_uri} with format '{destination_format}'..." - ) - extract_job.result() # Waits for job to complete. - logger.info( - f"Exported `{bq_table_path}` to {destination_gcs_uri} successfully." - ) - except Exception as e: - logger.exception(f"Failed to export table to GCS.") - raise e +logger.warning( + "gigl.common.utils.bq is deprecated and will be removed in a future release. " + "Please import from `gigl.common.utils.bq` instead." +) diff --git a/gigl/src/common/utils/file_loader.py b/gigl/src/common/utils/file_loader.py index 15cc5478d..35e61969d 100644 --- a/gigl/src/common/utils/file_loader.py +++ b/gigl/src/common/utils/file_loader.py @@ -1,322 +1,16 @@ -import shutil -import tempfile -from collections.abc import Mapping -from tempfile import _TemporaryFileWrapper as TemporaryFileWrapper -from typing import IO, AnyStr, Optional, Sequence, Tuple, Type, Union, cast +"""Deprecated shim - implementation moved to :mod:`gigl.common.utils.file_loader`. + +Importing from here continues to work but logs a deprecation warning. This shim +will be removed in a future release. +""" -from gigl.common import GcsUri, HttpUri, LocalUri, Uri, UriFactory from gigl.common.logger import Logger -from gigl.common.utils.gcs import GcsUtils -from gigl.common.utils.http import HttpUtils -from gigl.common.utils.local_fs import ( - FileSystemEntity, - copy_files, - count_files_with_uri_prefix, - create_file_symlinks, - does_path_exist, - list_at_path, - remove_file_or_folder_if_exist, +from gigl.common.utils.file_loader import ( # noqa: F401 (re-export for backwards compat) + FileLoader, ) logger = Logger() - - -class FileLoader: - def __init__(self, project: Optional[str] = None): - gcs_utils = GcsUtils(project) - self.__gcs_utils = gcs_utils - self.__unsupported_uri_message = ( - f"{self.__class__.__name__} does not support Uris of this type." - ) - - @staticmethod - def __get_uri_map_schema( - uri_map: Mapping[Uri, Uri], - ) -> Tuple[Optional[Type[Uri]], Optional[Type[Uri]]]: - uniform_src_type: Optional[Type[Uri]] = None - uniform_dst_type: Optional[Type[Uri]] = None - src_types: list[Type[Uri]] = [uri.__class__ for uri in uri_map.keys()] - dst_types: list[Type[Uri]] = [uri.__class__ for uri in uri_map.values()] - if all([src_types[0] == x for x in src_types]): - uniform_src_type = src_types[0] - if all([dst_types[0] == x for x in dst_types]): - uniform_dst_type = dst_types[0] - return uniform_src_type, uniform_dst_type - - def load_directories( - self, - source_to_dest_directory_map: dict[Uri, Uri], - ): - for dir_uri_src, dir_uri_dst in source_to_dest_directory_map.items(): - self.load_directory(dir_uri_src=dir_uri_src, dir_uri_dst=dir_uri_dst) - - def load_directory(self, dir_uri_src: Uri, dir_uri_dst: Uri): - uri_map_schema = self.__get_uri_map_schema(uri_map={dir_uri_src: dir_uri_dst}) - - if uri_map_schema == (GcsUri, LocalUri): - dir_uri_src = cast(GcsUri, dir_uri_src) - dir_uri_dst = cast(LocalUri, dir_uri_dst) - self.__gcs_utils.download_files_from_gcs_paths_to_local_dir( - gcs_paths=[dir_uri_src], local_path_dir=dir_uri_dst - ) - elif uri_map_schema == (LocalUri, GcsUri): - dir_uri_src = cast(LocalUri, dir_uri_src) - dir_uri_dst = cast(GcsUri, dir_uri_dst) - local_paths: list[LocalUri] = list_at_path( - local_path=dir_uri_src, file_system_entity=FileSystemEntity.FILE - ) - gcs_paths: list[GcsUri] = [ - GcsUri.join(dir_uri_dst, local_fn.uri) - for local_fn in list_at_path(dir_uri_src, names_only=True) - ] - local_file_path_to_gcs_path_map: dict[LocalUri, GcsUri] = { - src: dst for src, dst in zip(local_paths, gcs_paths) - } - self.load_files( - source_to_dest_file_uri_map=cast( - dict[Uri, Uri], local_file_path_to_gcs_path_map - ) - ) - elif uri_map_schema == (LocalUri, LocalUri): - dir_uri_src = cast(LocalUri, dir_uri_src) - dir_uri_dst = cast(LocalUri, dir_uri_dst) - - local_src_paths: list[LocalUri] = list_at_path( - local_path=dir_uri_src, file_system_entity=FileSystemEntity.FILE - ) - local_dst_paths: list[LocalUri] = [ - LocalUri.join(dir_uri_dst, local_src_fn) - for local_src_fn in list_at_path( - local_path=dir_uri_src, - names_only=True, - file_system_entity=FileSystemEntity.FILE, - ) - ] - source_to_dest_file_uri_map = { - src: dst for src, dst in zip(local_src_paths, local_dst_paths) - } - self.load_files( - source_to_dest_file_uri_map=cast( - dict[Uri, Uri], source_to_dest_file_uri_map - ) - ) - else: - raise TypeError(self.__unsupported_uri_message) - - def load_files( - self, - source_to_dest_file_uri_map: Mapping[Uri, Uri], - should_create_symlinks_if_possible: bool = True, - ) -> None: - uri_map_schema = self.__get_uri_map_schema(uri_map=source_to_dest_file_uri_map) - - if uri_map_schema == (GcsUri, LocalUri): - logger.info("Downloading from GCS to Local") - self.__gcs_utils.download_files_from_gcs_paths_to_local_paths( - file_map=cast(dict[GcsUri, LocalUri], source_to_dest_file_uri_map) - ) - elif uri_map_schema == (LocalUri, GcsUri): - logger.info("Uploading from Local to GCS") - self.__gcs_utils.upload_files_to_gcs( - local_file_path_to_gcs_path_map=cast( - dict[LocalUri, GcsUri], source_to_dest_file_uri_map - ), - parallel=True, - ) - elif uri_map_schema == (LocalUri, LocalUri): - logger.info("Copying from Local to Local") - local_source_to_link_path_map = source_to_dest_file_uri_map - if should_create_symlinks_if_possible: - logger.info("Will create symlinks") - create_file_symlinks( - local_source_to_link_path_map=cast( - dict[LocalUri, LocalUri], local_source_to_link_path_map - ), - should_overwrite=True, - ) - else: - logger.info("Will copy files") - copy_files( - local_source_to_local_dst_path_map=cast( - dict[LocalUri, LocalUri], local_source_to_link_path_map - ), - should_overwrite=True, - ) - elif uri_map_schema == (HttpUri, LocalUri): - logger.info("Downloading from HTTP to Local") - HttpUtils.download_files_from_http( - http_to_local_path_map=cast( - dict[HttpUri, LocalUri], source_to_dest_file_uri_map - ), - ) - else: - for file_uri_src, file_uri_dst in source_to_dest_file_uri_map.items(): - self.load_file( - file_uri_src=file_uri_src, - file_uri_dst=file_uri_dst, - should_create_symlinks_if_possible=should_create_symlinks_if_possible, - ) - - def load_file( - self, - file_uri_src: Uri, - file_uri_dst: Uri, - should_create_symlinks_if_possible: bool = True, - ) -> None: - if isinstance(file_uri_src, GcsUri) and isinstance(file_uri_dst, GcsUri): - self.__gcs_utils.copy_gcs_blob(file_uri_src, file_uri_dst) - elif isinstance(file_uri_src, GcsUri) and isinstance(file_uri_dst, LocalUri): - self.__gcs_utils.download_file_from_gcs( - gcs_path=file_uri_src, - dest_file_path=file_uri_dst, - ) - elif isinstance(file_uri_src, HttpUri) and isinstance(file_uri_dst, LocalUri): - HttpUtils.download_files_from_http( - http_to_local_path_map={file_uri_src: file_uri_dst} - ) - elif isinstance(file_uri_src, LocalUri) and isinstance(file_uri_dst, GcsUri): - self.__gcs_utils.upload_files_to_gcs( - local_file_path_to_gcs_path_map={file_uri_src: file_uri_dst}, - parallel=False, - ) - elif isinstance(file_uri_src, LocalUri) and isinstance(file_uri_dst, LocalUri): - local_source_to_link_path_map = {file_uri_src: file_uri_dst} - if should_create_symlinks_if_possible: - create_file_symlinks( - local_source_to_link_path_map=local_source_to_link_path_map, - should_overwrite=True, - ) - else: - copy_files( - local_source_to_local_dst_path_map=local_source_to_link_path_map, - should_overwrite=True, - ) - else: - logger.error( - f"Unsupported src/dst combo: {type(file_uri_src)} -> {type(file_uri_dst)}" - ) - raise TypeError(self.__unsupported_uri_message) - - def load_from_filelike(self, uri: Uri, filelike: IO[AnyStr]) -> None: - """ - Load a file from a file-like object into the specified URI. - - Args: - uri (Uri): The URI to load the file into. - filelike (IO[AnyStr]): A file-like object containing the data to be loaded. - """ - if isinstance(uri, GcsUri): - self.__gcs_utils.upload_from_filelike(gcs_path=uri, filelike=filelike) - elif isinstance(uri, LocalUri): - ptr = filelike.tell() - first = filelike.read(1) - filelike.seek(ptr) # Reset the file pointer after reading - if isinstance(first, bytes): - with open(uri.uri, "wb") as dest: - shutil.copyfileobj( - filelike, # ty: ignore[invalid-argument-type] - dest, - ) - else: - with open(uri.uri, "w", encoding="utf-8") as dest: - shutil.copyfileobj( - filelike, # ty: ignore[invalid-argument-type] - dest, - ) - - else: - raise NotImplementedError( - f"Cannot load file from buffer to URI {uri.uri} of type {type(uri)}; {self.__unsupported_uri_message}" - ) - - def load_to_temp_file( - self, - file_uri_src: Uri, - delete: bool = False, - should_create_symlinks_if_possible: bool = True, - ) -> TemporaryFileWrapper: - temp_file_handle = tempfile.NamedTemporaryFile(delete=delete) - temp_file_path = LocalUri(str(temp_file_handle.name)) - self.load_file( - file_uri_src=file_uri_src, - file_uri_dst=temp_file_path, - should_create_symlinks_if_possible=should_create_symlinks_if_possible, - ) - return temp_file_handle - - def count_assets(self, uri_prefix: Uri, suffix: Optional[str] = None) -> int: - if isinstance(uri_prefix, GcsUri): - return self.__gcs_utils.count_blobs_in_gcs_path( - gcs_path=uri_prefix, suffix=suffix - ) - elif isinstance(uri_prefix, LocalUri): - return count_files_with_uri_prefix(uri_prefix=uri_prefix, suffix=suffix) - else: - raise TypeError( - f"Uri type not supported, got {uri_prefix} in type {type(uri_prefix)}" - ) - - def does_uri_exist(self, uri: Union[str, Uri]) -> bool: - ( - """"" - Check if a URI exists - - Args: - uri (Union[str, Uri]): uri to check - Returns: - bool: True if URI exists, False otherwise - """ - "" - ) - - _uri = UriFactory.create_uri(uri=uri) if isinstance(uri, str) else uri - - exists: bool - if isinstance(_uri, GcsUri): - exists = self.__gcs_utils.does_gcs_file_exist(gcs_path=_uri) - elif isinstance(_uri, LocalUri): - exists = does_path_exist(_uri) - elif isinstance(_uri, HttpUri): - exists = HttpUtils.does_http_path_resolve(http_path=_uri) - else: - raise NotImplementedError(f"{self.__unsupported_uri_message} : {_uri}") - return exists - - def delete_files(self, uris: list[Uri]) -> None: - """ - Recursively delete files in the specified URIs. - - Args: - uris (list[Uri]): URIs to delete - Returns - None - """ - for uri in uris: - if isinstance(uri, LocalUri): - remove_file_or_folder_if_exist(local_path=uri) - elif isinstance(uri, GcsUri): - self.__gcs_utils.delete_files_in_bucket_dir(gcs_path=uri) - else: - raise NotImplementedError( - f"Cannot delete URI {uri.uri} of type {type(uri)}; {self.__unsupported_uri_message}" - ) - - def list_children(self, uri: Uri, pattern: Optional[str] = None) -> Sequence[Uri]: - """ - List all children of the given URI. - Args: - uri (Uri): The URI to list children of. - pattern (Optional[str]): Optional regex to match. If not provided then all children will be returned. - Returns: - list[Uri]: A list of URIs for the children of the given URI. - """ - if isinstance(uri, GcsUri): - return self.__gcs_utils.list_uris_with_gcs_path_pattern( - gcs_path=uri, pattern=pattern - ) - elif isinstance(uri, LocalUri): - return list_at_path(local_path=uri, regex=pattern) - else: - raise NotImplementedError( - f"Cannot list children of URI {uri.uri} of type {type(uri)}; {self.__unsupported_uri_message}" - ) +logger.warning( + "gigl.common.utils.file_loader is deprecated and will be removed in a future release. " + "Please import from `gigl.common.utils.file_loader` instead." +) diff --git a/gigl/src/common/utils/model.py b/gigl/src/common/utils/model.py index 41c04a1b7..28a814bc5 100644 --- a/gigl/src/common/utils/model.py +++ b/gigl/src/common/utils/model.py @@ -1,70 +1,19 @@ -import tempfile -from collections import OrderedDict - -import torch - -from gigl.common import LocalUri, Uri -from gigl.src.common.utils.file_loader import FileLoader - - -def save_scripted_model(model: torch.nn.Module, save_to_path_uri: Uri) -> None: - assert isinstance(model, torch.nn.Module), ( - "Can only save model of type torch.nn.Module" - ) - file_loader = FileLoader() - tmp_save_model_file = tempfile.NamedTemporaryFile(delete=False) - - model_scripted = torch.jit.script(model) # Export to TorchScript - model_scripted.save(tmp_save_model_file.name) # Save - - file_loader.load_file( - file_uri_src=LocalUri(tmp_save_model_file.name), file_uri_dst=save_to_path_uri - ) - tmp_save_model_file.close() - - -def save_state_dict(model: torch.nn.Module, save_to_path_uri: Uri) -> None: - if isinstance(model, torch.nn.parallel.DistributedDataParallel): - model = model.module - - assert isinstance(model, torch.nn.Module), ( - "Can only save model of type torch.nn.Module" - ) - file_loader = FileLoader() - - tmp_save_model_file = tempfile.NamedTemporaryFile(delete=False) - torch.save(model.state_dict(), tmp_save_model_file.name) - - file_loader.load_file( - file_uri_src=LocalUri(tmp_save_model_file.name), file_uri_dst=save_to_path_uri - ) - tmp_save_model_file.close() - - -def load_state_dict_from_uri( - load_from_uri: Uri, - device: torch.device = torch.device("cpu"), -) -> OrderedDict[str, torch.Tensor]: - state_dict: OrderedDict[str, torch.Tensor] - - file_loader = FileLoader() - tmp_file = file_loader.load_to_temp_file(load_from_uri) - - state_dict = torch.load(tmp_file.name, map_location=device) - tmp_file.close() - - return state_dict - - -def load_scripted_model_from_uri( - load_from_uri: Uri, -) -> torch.nn.Module: - scripted_model: torch.jit.ScriptModule - - file_loader = FileLoader() - tmp_file = file_loader.load_to_temp_file(load_from_uri) - - scripted_model = torch.jit.load(tmp_file.name) - tmp_file.close() - - return scripted_model +"""Deprecated shim - implementation moved to :mod:`gigl.common.utils.model`. + +Importing from here continues to work but logs a deprecation warning. This shim +will be removed in a future release. +""" + +from gigl.common.logger import Logger +from gigl.common.utils.model import ( # noqa: F401 (re-export for backwards compat) + load_scripted_model_from_uri, + load_state_dict_from_uri, + save_scripted_model, + save_state_dict, +) + +logger = Logger() +logger.warning( + "gigl.common.utils.model is deprecated and will be removed in a future release. " + "Please import from `gigl.common.utils.model` instead." +) diff --git a/gigl/src/common/utils/time.py b/gigl/src/common/utils/time.py index 37dc0cf8a..7d3ec4639 100644 --- a/gigl/src/common/utils/time.py +++ b/gigl/src/common/utils/time.py @@ -1,52 +1,21 @@ -import datetime - -from gigl.src.common.constants.time import DEFAULT_DATETIME_FORMAT - - -def parse_formatted_datetime( - stringified_date: str, fmt: str = DEFAULT_DATETIME_FORMAT -) -> datetime.datetime: - return datetime.datetime.strptime(stringified_date, fmt) - - -def format_datetime(dt: datetime.datetime, fmt: str = DEFAULT_DATETIME_FORMAT) -> str: - return dt.strftime(fmt) - - -def current_formatted_datetime(fmt: str = DEFAULT_DATETIME_FORMAT) -> str: - """ - Returns: - str: current timezone aware utc datetime formatted as a string - as per provided date format string. - """ - return current_datetime().strftime(fmt) - - -def current_datetime() -> datetime.datetime: - """ - Returns: - str: current timezone aware utc datetime. - """ - return datetime.datetime.now(datetime.timezone.utc) - - -def convert_days_to_ms(days: int) -> int: - return days * 24 * 60 * 60 * 1000 - - -def is_datetime_str_format_valid(datetime_str: str, datetime_format: str) -> bool: - """ - Validates a datetime string with specified format. ex "%Y%m%d" for "YYYYmmdd" - Return True/False so users can decide to raise error or reformat string - - :param datetime_str: - :param datetime_format: - :return: - """ - try: - datetime.datetime.strptime(datetime_str, datetime_format) - return True - except ValueError: - return False - except Exception as e: - return False +"""Deprecated shim - implementation moved to :mod:`gigl.common.utils.time`. + +Importing from here continues to work but logs a deprecation warning. This shim +will be removed in a future release. +""" + +from gigl.common.logger import Logger +from gigl.common.utils.time import ( # noqa: F401 (re-export for backwards compat) + convert_days_to_ms, + current_datetime, + current_formatted_datetime, + format_datetime, + is_datetime_str_format_valid, + parse_formatted_datetime, +) + +logger = Logger() +logger.warning( + "gigl.common.utils.time is deprecated and will be removed in a future release. " + "Please import from `gigl.common.utils.time` instead." +) diff --git a/gigl/src/common/utils/timeout.py b/gigl/src/common/utils/timeout.py index b57893c2c..b1a7177a9 100644 --- a/gigl/src/common/utils/timeout.py +++ b/gigl/src/common/utils/timeout.py @@ -1,64 +1,17 @@ -import errno -import os -import signal -from functools import wraps +"""Deprecated shim - implementation moved to :mod:`gigl.common.utils.timeout`. + +Importing from here continues to work but logs a deprecation warning. This shim +will be removed in a future release. +""" from gigl.common.logger import Logger +from gigl.common.utils.timeout import ( # noqa: F401 (re-export for backwards compat) + TimedOutException, + timeout, +) logger = Logger() - - -class TimedOutException(Exception): - pass - - -def timeout( - seconds=10, - error_message=os.strerror(errno.ETIME), - timeout_action_func=None, - exception_thrown_on_timeout=TimedOutException, - **timeout_action_func_params, -): - """ - Decorator to exit a program when a function execution exceeds a specified timeout. - - This decorator exits the program when the decorated function timed out, and executes - timeout_action_func before exiting. The timeout_action_func can be useful for cases - like environment clean up (e.g., ray.shutdown()). Another way to handle timeouts, - especially when there are multiple threads or child threads, is to run `func` with - multiprocessing. However, multiprocessing does not work with Ray. - - References to the current timeout implementation: - - https://docs.python.org/3/library/signal.html - - https://www.saltycrane.com/blog/2010/04/using-python-timeout-decorator-uploading-s3/ - - Parameters: - - seconds (int): Timeout duration in seconds. Defaults to 10. - - error_message (str): Error message to log on timeout. Defaults to os.strerror(errno.ETIME). - - timeout_action_func (callable, optional): Function to execute before exiting on timeout. - - exception_thrown_on_timeout (Exception): Exception to raise on timeout. Defaults to TimedOutException. - - **timeout_action_func_params: Arbitrary keyword arguments passed to timeout_action_func. - """ - - def decorator(func): - def _handler(signum, frame): - logger.info(error_message) - if timeout_action_func: - timeout_action_func(**timeout_action_func_params) - raise exception_thrown_on_timeout() - - def wrapper(*args, **kwargs): - old = signal.signal(signal.SIGALRM, _handler) - signal.alarm(seconds) - try: - result = func(*args, **kwargs) - finally: - # cancel the alarm - signal.alarm(0) - # reinstall the old signal handler - signal.signal(signal.SIGALRM, old) - return result - - return wraps(func)(wrapper) - - return decorator +logger.warning( + "gigl.common.utils.timeout is deprecated and will be removed in a future release. " + "Please import from `gigl.common.utils.timeout` instead." +) diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index 543954753..a76d8b820 100644 --- a/gigl/src/data_preprocessor/data_preprocessor.py +++ b/gigl/src/data_preprocessor/data_preprocessor.py @@ -21,6 +21,7 @@ from gigl.common.logger import Logger from gigl.common.metrics.decorators import flushes_metrics, profileit from gigl.common.utils import os_utils +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.proto_utils import ProtoUtils from gigl.env.pipelines_config import get_resource_config from gigl.src.common.constants.components import GiGLComponents @@ -35,7 +36,6 @@ NodeType, ) from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.file_loader import FileLoader from gigl.src.common.utils.gigl_runtime import initialize_gigl_runtime from gigl.src.common.utils.metrics_service_provider import ( get_metrics_service_instance, diff --git a/gigl/src/data_preprocessor/lib/enumerate/utils.py b/gigl/src/data_preprocessor/lib/enumerate/utils.py index b606d7edb..ea1f283e3 100644 --- a/gigl/src/data_preprocessor/lib/enumerate/utils.py +++ b/gigl/src/data_preprocessor/lib/enumerate/utils.py @@ -7,12 +7,12 @@ import google.cloud.bigquery as bigquery from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.env.pipelines_config import get_resource_config from gigl.src.common.constants import bq as bq_constants from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, EdgeUsageType, NodeType -from gigl.src.common.utils.bq import BqUtils from gigl.src.data_preprocessor.lib.enumerate import queries as enumeration_queries from gigl.src.data_preprocessor.lib.ingest.bigquery import ( BigqueryEdgeDataReference, diff --git a/gigl/src/data_preprocessor/lib/transform/utils.py b/gigl/src/data_preprocessor/lib/transform/utils.py index 7e9be7081..bbfef15a0 100644 --- a/gigl/src/data_preprocessor/lib/transform/utils.py +++ b/gigl/src/data_preprocessor/lib/transform/utils.py @@ -17,11 +17,11 @@ from gigl.common import GcsUri, LocalUri, Uri from gigl.common.beam.better_tfrecordio import BetterWriteToTFRecord from gigl.common.logger import Logger +from gigl.common.utils.time import current_formatted_datetime from gigl.env.pipelines_config import get_resource_config from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.utils.dataflow import init_beam_pipeline_options -from gigl.src.common.utils.time import current_formatted_datetime from gigl.src.data_preprocessor.lib.ingest.reference import ( DataReference, EdgeDataReference, diff --git a/gigl/src/inference/lib/assets.py b/gigl/src/inference/lib/assets.py index a1fe82188..9b551c89a 100644 --- a/gigl/src/inference/lib/assets.py +++ b/gigl/src/inference/lib/assets.py @@ -1,10 +1,10 @@ from gigl.common import GcsUri from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils from gigl.src.common.constants import gcs as gcs_constants from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.bq import BqUtils logger = Logger() diff --git a/gigl/src/inference/v1/gnn_inferencer.py b/gigl/src/inference/v1/gnn_inferencer.py index 7f17a0439..ad1083169 100644 --- a/gigl/src/inference/v1/gnn_inferencer.py +++ b/gigl/src/inference/v1/gnn_inferencer.py @@ -18,6 +18,8 @@ from gigl.common.logger import Logger from gigl.common.metrics.decorators import flushes_metrics, profileit from gigl.common.utils import os_utils +from gigl.common.utils.bq import BqUtils +from gigl.common.utils.model import load_state_dict_from_uri from gigl.env.pipelines_config import get_resource_config from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.constants.metrics import TIMER_INFERENCER_S @@ -25,12 +27,10 @@ from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import NodeType from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.bq import BqUtils from gigl.src.common.utils.gigl_runtime import initialize_gigl_runtime from gigl.src.common.utils.metrics_service_provider import ( get_metrics_service_instance, ) -from gigl.src.common.utils.model import load_state_dict_from_uri from gigl.src.inference.lib.assets import InferenceAssets from gigl.src.inference.v1.lib.base_inference_blueprint import BaseInferenceBlueprint from gigl.src.inference.v1.lib.base_inferencer import BaseInferencer diff --git a/gigl/src/mocking/dataset_asset_mocker.py b/gigl/src/mocking/dataset_asset_mocker.py index 4230e979e..2eaf5c286 100644 --- a/gigl/src/mocking/dataset_asset_mocker.py +++ b/gigl/src/mocking/dataset_asset_mocker.py @@ -3,13 +3,13 @@ import gigl.src.mocking.lib.constants as mocking_constants from gigl.common import GcsUri from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils from gigl.common.utils.proto_utils import ProtoUtils from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.common.types.pb_wrappers.graph_metadata import GraphMetadataPbWrapper from gigl.src.common.types.task_metadata import TaskMetadataType -from gigl.src.common.utils.bq import BqUtils from gigl.src.config_populator.config_populator import ConfigPopulator from gigl.src.mocking.lib import ( mock_input_for_data_preprocessor, diff --git a/gigl/src/mocking/dataset_asset_mocking_suite.py b/gigl/src/mocking/dataset_asset_mocking_suite.py index 6b924aadf..e8fc9ec59 100644 --- a/gigl/src/mocking/dataset_asset_mocking_suite.py +++ b/gigl/src/mocking/dataset_asset_mocking_suite.py @@ -11,9 +11,9 @@ from gigl.common.constants import GIGL_ROOT_DIR from gigl.common.logger import Logger +from gigl.common.utils.time import current_formatted_datetime from gigl.src.common.types.graph_data import EdgeType, EdgeUsageType, NodeType, Relation from gigl.src.common.types.task_metadata import TaskMetadataType -from gigl.src.common.utils.time import current_formatted_datetime from gigl.src.mocking.dataset_asset_mocker import DatasetAssetMocker from gigl.src.mocking.lib.mocked_dataset_resources import MockedDatasetInfo from gigl.src.mocking.lib.pyg_datasets_forks import CoraFromGCS, DBLPFromGCS diff --git a/gigl/src/mocking/lib/constants.py b/gigl/src/mocking/lib/constants.py index cb188623d..2d43ceef1 100644 --- a/gigl/src/mocking/lib/constants.py +++ b/gigl/src/mocking/lib/constants.py @@ -1,8 +1,8 @@ import gigl.env.dep_constants as dep_constants from gigl.common import GcsUri, LocalUri +from gigl.common.utils.bq import BqUtils from gigl.src.common.constants.local_fs import get_gigl_root_directory from gigl.src.common.types.graph_data import EdgeType, EdgeUsageType, NodeType -from gigl.src.common.utils.bq import BqUtils MOCK_DATA_GCS_BUCKET = GcsUri(f"gs://{dep_constants.GIGL_PUBLIC_BUCKET_NAME}/") MOCK_DATA_BQ_DATASET_NAME = dep_constants.GIGL_PUBLIC_DATASET_NAME diff --git a/gigl/src/mocking/lib/mock_input_for_data_preprocessor.py b/gigl/src/mocking/lib/mock_input_for_data_preprocessor.py index c1a01bbb5..0c056e57c 100644 --- a/gigl/src/mocking/lib/mock_input_for_data_preprocessor.py +++ b/gigl/src/mocking/lib/mock_input_for_data_preprocessor.py @@ -7,8 +7,8 @@ from gigl.common import UriFactory from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.src.common.types.graph_data import EdgeType, EdgeUsageType, NodeType -from gigl.src.common.utils.bq import BqUtils from gigl.src.mocking.lib.constants import ( get_example_task_edges_bq_table_path, get_example_task_nodes_bq_table_path, diff --git a/gigl/src/mocking/lib/mock_input_for_inference.py b/gigl/src/mocking/lib/mock_input_for_inference.py index cff5f8f0d..93f801a24 100644 --- a/gigl/src/mocking/lib/mock_input_for_inference.py +++ b/gigl/src/mocking/lib/mock_input_for_inference.py @@ -1,6 +1,6 @@ import torch.nn -import gigl.src.common.utils.model as model_utils +import gigl.common.utils.model as model_utils from gigl.common import UriFactory from gigl.common.logger import Logger from gigl.common.utils.os_utils import import_obj diff --git a/gigl/src/mocking/lib/mock_input_for_subgraph_sampler.py b/gigl/src/mocking/lib/mock_input_for_subgraph_sampler.py index 006dd99f9..b22edfe5d 100644 --- a/gigl/src/mocking/lib/mock_input_for_subgraph_sampler.py +++ b/gigl/src/mocking/lib/mock_input_for_subgraph_sampler.py @@ -8,6 +8,7 @@ from gigl.common import GcsUri, LocalUri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.proto_utils import ProtoUtils from gigl.src.common.types.graph_data import ( CondensedEdgeType, @@ -16,7 +17,6 @@ EdgeUsageType, NodeType, ) -from gigl.src.common.utils.file_loader import FileLoader from gigl.src.data_preprocessor.lib.transform.tf_value_encoder import TFValueEncoder from gigl.src.data_preprocessor.lib.types import FeatureSpecDict, InstanceDict from gigl.src.mocking.lib.constants import ( diff --git a/gigl/src/mocking/lib/mock_output_for_inference.py b/gigl/src/mocking/lib/mock_output_for_inference.py index efaed1627..47ca21857 100644 --- a/gigl/src/mocking/lib/mock_output_for_inference.py +++ b/gigl/src/mocking/lib/mock_output_for_inference.py @@ -6,9 +6,10 @@ import tensorflow as tf from google.cloud import bigquery -import gigl.src.common.utils.model as model_utils +import gigl.common.utils.model as model_utils from gigl.common import UriFactory from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.os_utils import import_obj from gigl.src.common.graph_builder.abstract_graph_builder import GraphBuilder from gigl.src.common.graph_builder.pyg_graph_builder import PygGraphBuilder @@ -23,7 +24,6 @@ PreprocessedMetadataPbWrapper, ) from gigl.src.common.types.task_metadata import TaskMetadataType -from gigl.src.common.utils.bq import BqUtils from gigl.src.inference.v1.lib.base_inferencer import BaseInferencer, InferBatchResults from gigl.src.inference.v1.lib.inference_output_schema import ( DEFAULT_EMBEDDING_FIELD, diff --git a/gigl/src/mocking/lib/pyg_datasets_forks.py b/gigl/src/mocking/lib/pyg_datasets_forks.py index e83abfe0c..09717081b 100644 --- a/gigl/src/mocking/lib/pyg_datasets_forks.py +++ b/gigl/src/mocking/lib/pyg_datasets_forks.py @@ -9,8 +9,8 @@ import gigl.env.dep_constants as dep_constants from gigl.common import GcsUri, LocalUri +from gigl.common.utils.file_loader import FileLoader from gigl.env.pipelines_config import get_resource_config -from gigl.src.common.utils.file_loader import FileLoader unprocessed_datasets_gcs_uri = GcsUri( f"gs://{dep_constants.GIGL_PUBLIC_BUCKET_NAME}/unprocessed_datasets/" diff --git a/gigl/src/post_process/post_processor.py b/gigl/src/post_process/post_processor.py index 6642dac89..48abdd5af 100644 --- a/gigl/src/post_process/post_processor.py +++ b/gigl/src/post_process/post_processor.py @@ -8,6 +8,7 @@ from gigl.common.logger import Logger from gigl.common.metrics.decorators import flushes_metrics, profileit from gigl.common.utils import os_utils +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.gcs import GcsUtils from gigl.src.common.constants import gcs as gcs_constants from gigl.src.common.constants.components import GiGLComponents @@ -18,7 +19,6 @@ from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.model_eval_metrics import EvalMetricsCollection from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.file_loader import FileLoader from gigl.src.common.utils.gigl_runtime import initialize_gigl_runtime from gigl.src.common.utils.metrics_service_provider import ( get_metrics_service_instance, diff --git a/gigl/src/post_process/utils/cosine_similarity.py b/gigl/src/post_process/utils/cosine_similarity.py index adbd6a57c..f5656b9f1 100644 --- a/gigl/src/post_process/utils/cosine_similarity.py +++ b/gigl/src/post_process/utils/cosine_similarity.py @@ -4,9 +4,9 @@ import pandas as pd +from gigl.common.constants import DEFAULT_DATE_FORMAT +from gigl.common.utils.bq import BqUtils from gigl.env.pipelines_config import get_resource_config -from gigl.src.common.constants.time import DEFAULT_DATE_FORMAT -from gigl.src.common.utils.bq import BqUtils from gigl.src.inference.v1.lib.inference_output_schema import ( DEFAULT_EMBEDDING_FIELD, DEFAULT_NODE_ID_FIELD, diff --git a/gigl/src/post_process/utils/unenumeration.py b/gigl/src/post_process/utils/unenumeration.py index bc7972dbe..5c00246b7 100644 --- a/gigl/src/post_process/utils/unenumeration.py +++ b/gigl/src/post_process/utils/unenumeration.py @@ -8,11 +8,11 @@ import gigl.src.inference.v1.lib.queries as inference_queries from gigl.common.env_config import get_available_cpus from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.env.pipelines_config import get_resource_config from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.types.graph_data import NodeType from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.bq import BqUtils from gigl.src.inference.lib.assets import InferenceAssets from gigl.src.inference.v1.lib.inference_output_schema import ( DEFAULT_EMBEDDINGS_TABLE_SCHEMA, diff --git a/gigl/src/subgraph_sampler/subgraph_sampler.py b/gigl/src/subgraph_sampler/subgraph_sampler.py index 8b7c140d1..0af340cee 100644 --- a/gigl/src/subgraph_sampler/subgraph_sampler.py +++ b/gigl/src/subgraph_sampler/subgraph_sampler.py @@ -15,13 +15,13 @@ from gigl.common.logger import Logger from gigl.common.metrics.decorators import flushes_metrics, profileit from gigl.common.utils import os_utils +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.gcs import GcsUtils from gigl.env.pipelines_config import get_resource_config from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.constants.metrics import TIMER_SUBGRAPH_SAMPLER_S from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.file_loader import FileLoader from gigl.src.common.utils.metrics_service_provider import ( get_metrics_service_instance, initialize_metrics, diff --git a/gigl/src/training/trainer.py b/gigl/src/training/trainer.py index 197aab0f3..f84f66cf2 100644 --- a/gigl/src/training/trainer.py +++ b/gigl/src/training/trainer.py @@ -4,10 +4,10 @@ import gigl.src.common.constants.gcs as gcs_constants from gigl.common import Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.file_loader import FileLoader from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.file_loader import FileLoader from gigl.src.common.utils.gigl_runtime import initialize_gigl_runtime # TODO: (svij) Rename Trainer to TrainerV1 diff --git a/gigl/src/training/v1/lib/training_process.py b/gigl/src/training/v1/lib/training_process.py index c79bd6983..d09940166 100644 --- a/gigl/src/training/v1/lib/training_process.py +++ b/gigl/src/training/v1/lib/training_process.py @@ -13,12 +13,15 @@ import torch.distributed import torch.nn.parallel -import gigl.src.common.utils.model as model_utils +import gigl.common.utils.model as model_utils from gigl.common import GcsUri, LocalUri, Uri, UriFactory from gigl.common.logger import Logger from gigl.common.metrics.decorators import flushes_metrics, profileit from gigl.common.utils import os_utils, torch_training +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.local_fs import does_path_exist +from gigl.common.utils.model import load_state_dict_from_uri +from gigl.common.utils.time import current_formatted_datetime from gigl.common.utils.torch_training import ( get_distributed_backend, get_rank, @@ -42,13 +45,10 @@ from gigl.src.common.types.model_eval_metrics import EvalMetricsCollection from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.common.types.task_metadata import TaskMetadataType -from gigl.src.common.utils.file_loader import FileLoader from gigl.src.common.utils.gigl_runtime import initialize_gigl_runtime from gigl.src.common.utils.metrics_service_provider import ( get_metrics_service_instance, ) -from gigl.src.common.utils.model import load_state_dict_from_uri -from gigl.src.common.utils.time import current_formatted_datetime from gigl.src.training.v1.lib.base_trainer import BaseTrainer logger = Logger() diff --git a/gigl/src/validation_check/libs/frozen_config_path_checks.py b/gigl/src/validation_check/libs/frozen_config_path_checks.py index 8769dc4c9..77e668f27 100644 --- a/gigl/src/validation_check/libs/frozen_config_path_checks.py +++ b/gigl/src/validation_check/libs/frozen_config_path_checks.py @@ -2,10 +2,10 @@ from gigl.common import Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.file_loader import FileLoader from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.common.types.pb_wrappers.task_metadata import TaskMetadataPbWrapper from gigl.src.common.types.task_metadata import TaskMetadataType -from gigl.src.common.utils.file_loader import FileLoader from snapchat.research.gbml import gbml_config_pb2 logger = Logger() diff --git a/scripts/bootstrap_resource_config.py b/scripts/bootstrap_resource_config.py index f3517138e..9f544cfcd 100644 --- a/scripts/bootstrap_resource_config.py +++ b/scripts/bootstrap_resource_config.py @@ -12,7 +12,7 @@ import yaml from gigl.common import GcsUri, HttpUri, LocalUri, UriFactory -from gigl.src.common.utils.file_loader import FileLoader +from gigl.common.utils.file_loader import FileLoader GIGL_ROOT_DIR = pathlib.Path(__file__).resolve().parent.parent LOCAL_DEV_TEMPLATE_RES_CONF = LocalUri( diff --git a/scripts/launch_graph_store_local.py b/scripts/launch_graph_store_local.py index 368a57dcd..610e728c3 100644 --- a/scripts/launch_graph_store_local.py +++ b/scripts/launch_graph_store_local.py @@ -204,7 +204,7 @@ def _load_yaml_payload(uri: str) -> dict[str, Any]: return payload or {} from gigl.common import UriFactory - from gigl.src.common.utils.file_loader import FileLoader + from gigl.common.utils.file_loader import FileLoader temp_file = FileLoader().load_to_temp_file( file_uri_src=UriFactory.create_uri(uri), diff --git a/tests/api_test/api_test_inference.py b/tests/api_test/api_test_inference.py index 4381e131e..7709484ec 100644 --- a/tests/api_test/api_test_inference.py +++ b/tests/api_test/api_test_inference.py @@ -34,14 +34,14 @@ from gigl.common import GcsUri, UriFactory from gigl.common.data.export import EmbeddingExporter, load_embeddings_to_bigquery from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils +from gigl.common.utils.model import load_state_dict_from_uri from gigl.distributed import DistDataset, build_dataset_from_task_config_uri from gigl.nn.models import LinkPredictionGNN from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import NodeType from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.bq import BqUtils -from gigl.src.common.utils.model import load_state_dict_from_uri from gigl.src.inference.lib.assets import InferenceAssets from gigl.utils.sampling import parse_fanout diff --git a/tests/api_test/api_test_training.py b/tests/api_test/api_test_training.py index 5b8ff162a..c3c751d4b 100644 --- a/tests/api_test/api_test_training.py +++ b/tests/api_test/api_test_training.py @@ -38,6 +38,7 @@ from examples.link_prediction.models import init_example_gigl_homogeneous_model from gigl.common import Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.common.utils.torch_training import is_distributed_available_and_initialized from gigl.distributed import ( DistABLPLoader, @@ -49,7 +50,6 @@ from gigl.nn.loss import RetrievalLoss from gigl.nn.models import LinkPredictionGNN from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.model import load_state_dict_from_uri, save_state_dict from gigl.types.graph import to_homogeneous from gigl.utils.iterator import InfiniteIterator from gigl.utils.sampling import parse_fanout diff --git a/tests/config_tests/notebooks_test.py b/tests/config_tests/notebooks_test.py index d3118b10e..0b0270007 100644 --- a/tests/config_tests/notebooks_test.py +++ b/tests/config_tests/notebooks_test.py @@ -12,8 +12,8 @@ from gigl.common.logger import Logger from gigl.common.types.uri.gcs_uri import GcsUri from gigl.common.types.uri.uri_factory import UriFactory +from gigl.common.utils.file_loader import FileLoader from gigl.env.pipelines_config import get_resource_config -from gigl.src.common.utils.file_loader import FileLoader logger = Logger() diff --git a/tests/integration/common/data/export_test.py b/tests/integration/common/data/export_test.py index cb54344b5..2b0be4def 100644 --- a/tests/integration/common/data/export_test.py +++ b/tests/integration/common/data/export_test.py @@ -7,9 +7,9 @@ from gigl.common import GcsUri from gigl.common.data.export import EmbeddingExporter, load_embeddings_to_bigquery from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils from gigl.env.pipelines_config import get_resource_config -from gigl.src.common.utils.bq import BqUtils from tests.test_assets.test_case import TestCase logger = Logger() diff --git a/tests/integration/common/file_loader_test.py b/tests/integration/common/file_loader_test.py index efa05aee0..c8c3a30b9 100644 --- a/tests/integration/common/file_loader_test.py +++ b/tests/integration/common/file_loader_test.py @@ -7,9 +7,9 @@ import gigl.common.utils.local_fs as local_fs from gigl.common import GcsUri, HttpUri, LocalUri, Uri +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.gcs import GcsUtils from gigl.env.pipelines_config import get_resource_config -from gigl.src.common.utils.file_loader import FileLoader from tests.test_assets.test_case import TestCase diff --git a/tests/integration/common/utils/__init__.py b/tests/integration/common/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integration/src/common/utils/bq_test.py b/tests/integration/common/utils/bq_test.py similarity index 99% rename from tests/integration/src/common/utils/bq_test.py rename to tests/integration/common/utils/bq_test.py index bf4922355..3fdf2cf6b 100644 --- a/tests/integration/src/common/utils/bq_test.py +++ b/tests/integration/common/utils/bq_test.py @@ -4,8 +4,8 @@ from absl.testing import absltest from parameterized import param, parameterized +from gigl.common.utils.bq import BqUtils from gigl.env.pipelines_config import get_resource_config -from gigl.src.common.utils.bq import BqUtils from tests.test_assets.test_case import TestCase diff --git a/tests/integration/distributed/utils/networking_test.py b/tests/integration/distributed/utils/networking_test.py index 18f3c1e29..a0d47b003 100644 --- a/tests/integration/distributed/utils/networking_test.py +++ b/tests/integration/distributed/utils/networking_test.py @@ -6,6 +6,7 @@ from parameterized import param, parameterized from gigl.common.services.vertex_ai import VertexAiJobConfig, VertexAIService +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.proto_utils import ProtoUtils from gigl.distributed.utils import get_graph_store_info from gigl.env.constants import ( @@ -13,7 +14,6 @@ GIGL_RESOURCE_CONFIG_URI_ENV_KEY, ) from gigl.env.pipelines_config import get_resource_config -from gigl.src.common.utils.file_loader import FileLoader from tests.test_assets.test_case import TestCase # Short timeout so a broken fresh image fails fast instead of hanging CI until the diff --git a/tests/integration/pipeline/config_populator/config_populator_pipeline_test.py b/tests/integration/pipeline/config_populator/config_populator_pipeline_test.py index 54c06f82a..e586f858b 100644 --- a/tests/integration/pipeline/config_populator/config_populator_pipeline_test.py +++ b/tests/integration/pipeline/config_populator/config_populator_pipeline_test.py @@ -5,9 +5,9 @@ from gigl.common.logger import Logger from gigl.common.utils.gcs import GcsUtils from gigl.common.utils.proto_utils import ProtoUtils +from gigl.common.utils.time import current_formatted_datetime +from gigl.common.utils.timeout import timeout from gigl.src.common.types import AppliedTaskIdentifier -from gigl.src.common.utils.time import current_formatted_datetime -from gigl.src.common.utils.timeout import timeout from gigl.src.config_populator.config_populator import ConfigPopulator from snapchat.research.gbml import gbml_config_pb2 from tests.test_assets.test_case import TestCase diff --git a/tests/integration/pipeline/data_preprocessor/data_preprocessor_pipeline_test.py b/tests/integration/pipeline/data_preprocessor/data_preprocessor_pipeline_test.py index 9a651ede6..5ed8b147a 100644 --- a/tests/integration/pipeline/data_preprocessor/data_preprocessor_pipeline_test.py +++ b/tests/integration/pipeline/data_preprocessor/data_preprocessor_pipeline_test.py @@ -14,8 +14,11 @@ import gigl.src.common.constants.local_fs as local_fs_constants from gigl.common import GcsUri, LocalUri, Uri from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils from gigl.common.utils.proto_utils import ProtoUtils +from gigl.common.utils.time import current_formatted_datetime +from gigl.common.utils.timeout import timeout from gigl.src.common.constants.graph_metadata import ( DEFAULT_CONDENSED_EDGE_TYPE, DEFAULT_CONDENSED_NODE_TYPE, @@ -33,9 +36,6 @@ from gigl.src.common.types.pb_wrappers.preprocessed_metadata import ( PreprocessedMetadataPbWrapper, ) -from gigl.src.common.utils.bq import BqUtils -from gigl.src.common.utils.time import current_formatted_datetime -from gigl.src.common.utils.timeout import timeout from gigl.src.data_preprocessor.data_preprocessor import DataPreprocessor from gigl.src.mocking.lib.feature_handling import get_feature_field_name from gigl.src.mocking.lib.mocked_dataset_resources import MockedDatasetInfo diff --git a/tests/integration/pipeline/data_preprocessor/enumerator_test.py b/tests/integration/pipeline/data_preprocessor/enumerator_test.py index 0d06a9b1a..a89195932 100644 --- a/tests/integration/pipeline/data_preprocessor/enumerator_test.py +++ b/tests/integration/pipeline/data_preprocessor/enumerator_test.py @@ -5,13 +5,13 @@ import gigl.src.data_preprocessor.lib.enumerate.queries as enumeration_queries from gigl.common.beam.sharded_read import BigQueryShardedReadConfig +from gigl.common.constants import NODASH_DATETIME_FORMAT from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils +from gigl.common.utils.time import current_formatted_datetime from gigl.env.pipelines_config import get_resource_config -from gigl.src.common.constants.time import NODASH_DATETIME_FORMAT from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import EdgeType, EdgeUsageType, NodeType, Relation -from gigl.src.common.utils.bq import BqUtils -from gigl.src.common.utils.time import current_formatted_datetime from gigl.src.data_preprocessor.lib.enumerate.utils import ( Enumerator, EnumeratorEdgeTypeMetadata, diff --git a/tests/integration/pipeline/inferencer/inferencer_test.py b/tests/integration/pipeline/inferencer/inferencer_test.py index 905d996cc..91be45a23 100644 --- a/tests/integration/pipeline/inferencer/inferencer_test.py +++ b/tests/integration/pipeline/inferencer/inferencer_test.py @@ -3,16 +3,16 @@ import gigl.src.common.constants.gcs as gcs_constants from gigl.common import LocalUri from gigl.common.logger import Logger +from gigl.common.utils.bq import BqUtils from gigl.common.utils.gcs import GcsUtils from gigl.common.utils.proto_utils import ProtoUtils +from gigl.common.utils.time import current_formatted_datetime from gigl.env.pipelines_config import get_resource_config from gigl.src.common.constants.bq import get_embeddings_table, get_predictions_table from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.graph_data import NodeType from gigl.src.common.types.pb_wrappers.task_metadata import TaskMetadataPbWrapper from gigl.src.common.types.task_metadata import TaskMetadataType -from gigl.src.common.utils.bq import BqUtils -from gigl.src.common.utils.time import current_formatted_datetime from gigl.src.inference.v1.gnn_inferencer import InferencerV1 from gigl.src.mocking.lib.mocked_dataset_resources import MockedDatasetInfo from gigl.src.mocking.lib.versioning import get_mocked_dataset_artifact_metadata diff --git a/tests/integration/pipeline/split_generator/split_generator_pipeline_test.py b/tests/integration/pipeline/split_generator/split_generator_pipeline_test.py index 2a187992a..997fcf860 100644 --- a/tests/integration/pipeline/split_generator/split_generator_pipeline_test.py +++ b/tests/integration/pipeline/split_generator/split_generator_pipeline_test.py @@ -10,7 +10,9 @@ from gigl.common import GcsUri, LocalUri, Uri, UriFactory from gigl.common.constants import SPARK_35_TFRECORD_JAR_LOCAL_PATH from gigl.common.logger import Logger +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.proto_utils import ProtoUtils +from gigl.common.utils.time import current_formatted_datetime from gigl.env.pipelines_config import get_resource_config from gigl.src.common.constants.local_fs import get_project_root_directory from gigl.src.common.graph_builder.graph_builder_factory import GraphBuilderFactory @@ -22,8 +24,6 @@ from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.common.types.pb_wrappers.task_metadata import TaskMetadataPbWrapper from gigl.src.common.types.task_metadata import TaskMetadataType -from gigl.src.common.utils.file_loader import FileLoader -from gigl.src.common.utils.time import current_formatted_datetime from gigl.src.mocking.lib.mocked_dataset_resources import MockedDatasetInfo from gigl.src.mocking.lib.versioning import get_mocked_dataset_artifact_metadata from gigl.src.mocking.mocking_assets.mocked_datasets_for_pipeline_tests import ( diff --git a/tests/integration/pipeline/subgraph_sampler/subgraph_sampler_test.py b/tests/integration/pipeline/subgraph_sampler/subgraph_sampler_test.py index 63ae84560..7b8ef63fc 100644 --- a/tests/integration/pipeline/subgraph_sampler/subgraph_sampler_test.py +++ b/tests/integration/pipeline/subgraph_sampler/subgraph_sampler_test.py @@ -5,6 +5,7 @@ from gigl.common import LocalUri, Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.file_loader import FileLoader from gigl.env.pipelines_config import get_resource_config from gigl.src.common.types.graph_data import ( CondensedEdgeType, @@ -22,7 +23,6 @@ from gigl.src.common.types.pb_wrappers.graph_data_types_utils import ( get_dehydrated_node_pb_wrappers_from_edge_wrapper, ) -from gigl.src.common.utils.file_loader import FileLoader from gigl.src.mocking.lib.mocked_dataset_resources import MockedDatasetInfo from gigl.src.mocking.lib.versioning import get_mocked_dataset_artifact_metadata from gigl.src.mocking.mocking_assets.mocked_datasets_for_pipeline_tests import ( diff --git a/tests/integration/pipeline/subgraph_sampler/utils.py b/tests/integration/pipeline/subgraph_sampler/utils.py index 3a2dfcd72..59ffb61c0 100644 --- a/tests/integration/pipeline/subgraph_sampler/utils.py +++ b/tests/integration/pipeline/subgraph_sampler/utils.py @@ -14,6 +14,7 @@ ) from gigl.common.logger import Logger from gigl.common.utils.proto_utils import ProtoUtils +from gigl.common.utils.time import current_formatted_datetime from gigl.src.common.constants.local_fs import get_project_root_directory from gigl.src.common.translators.gbml_protos_translator import GbmlProtosTranslator from gigl.src.common.types.graph_data import ( @@ -35,7 +36,6 @@ get_dehydrated_node_pb_wrappers_from_edge_wrapper, ) from gigl.src.common.types.task_metadata import TaskMetadataType -from gigl.src.common.utils.time import current_formatted_datetime from gigl.src.training.v1.lib.data_loaders.tf_records_iterable_dataset import ( get_np_iterator_from_tfrecords, ) diff --git a/tests/unit/common/data/export_test.py b/tests/unit/common/data/export_test.py index 8f1f300b9..a0cccc2b5 100644 --- a/tests/unit/common/data/export_test.py +++ b/tests/unit/common/data/export_test.py @@ -271,7 +271,7 @@ def test_flushes_after_maximum_buffer_size(self): records = list(reader) self.assertEqual(records, expected_records_by_batch[i]) - @patch("gigl.src.common.utils.file_loader.GcsUtils") + @patch("gigl.common.utils.file_loader.GcsUtils") def test_flush_resets_buffer(self, mock_gcs_utils_class): # Mock inputs gcs_base_uri = GcsUri("gs://test-bucket/test-folder") @@ -324,7 +324,7 @@ def mock_upload(gcs_path: Uri, filelike: io.BytesIO): self.assertEqual(records, expected_records) @patch("time.sleep") - @patch("gigl.src.common.utils.file_loader.GcsUtils") + @patch("gigl.common.utils.file_loader.GcsUtils") def test_write_embeddings_to_gcs_upload_retries_on_google_cloud_error_and_fails( self, mock_gcs_utils_class, mock_sleep ): @@ -347,7 +347,7 @@ def test_write_embeddings_to_gcs_upload_retries_on_google_cloud_error_and_fails( self.assertEqual(mock_gcs_utils.upload_from_filelike.call_count, 6) @patch("time.sleep") - @patch("gigl.src.common.utils.file_loader.GcsUtils") + @patch("gigl.common.utils.file_loader.GcsUtils") def test_write_embeddings_to_gcs_upload_retries_on_request_exception_and_fails( self, mock_gcs_utils_class, mock_sleep ): @@ -372,7 +372,7 @@ def test_write_embeddings_to_gcs_upload_retries_on_request_exception_and_fails( exporter.flush_records() self.assertEqual(mock_gcs_utils.upload_from_filelike.call_count, 6) - @patch("gigl.src.common.utils.file_loader.GcsUtils") + @patch("gigl.common.utils.file_loader.GcsUtils") def test_skips_flush_if_empty(self, mock_gcs_utils_class): gcs_base_uri = GcsUri("gs://test-bucket/test-folder") @@ -632,7 +632,7 @@ def test_flushes_after_maximum_buffer_size(self): records = list(reader) self.assertRecordsAlmostEqual(records, expected_records_by_batch[i]) - @patch("gigl.src.common.utils.file_loader.GcsUtils") + @patch("gigl.common.utils.file_loader.GcsUtils") def test_flush_resets_buffer(self, mock_gcs_utils_class): # Mock inputs gcs_base_uri = GcsUri("gs://test-bucket/test-folder") @@ -676,7 +676,7 @@ def mock_upload(gcs_path: Uri, filelike: io.BytesIO): self.assertRecordsAlmostEqual(records, expected_records) @patch("time.sleep") - @patch("gigl.src.common.utils.file_loader.GcsUtils") + @patch("gigl.common.utils.file_loader.GcsUtils") def test_write_predictions_to_gcs_upload_retries_on_google_cloud_error_and_fails( self, mock_gcs_utils_class, mock_sleep ): @@ -699,7 +699,7 @@ def test_write_predictions_to_gcs_upload_retries_on_google_cloud_error_and_fails self.assertEqual(mock_gcs_utils.upload_from_filelike.call_count, 6) @patch("time.sleep") - @patch("gigl.src.common.utils.file_loader.GcsUtils") + @patch("gigl.common.utils.file_loader.GcsUtils") def test_write_predictions_to_gcs_upload_retries_on_request_exception_and_fails( self, mock_gcs_utils_class, mock_sleep ): @@ -724,7 +724,7 @@ def test_write_predictions_to_gcs_upload_retries_on_request_exception_and_fails( exporter.flush_records() self.assertEqual(mock_gcs_utils.upload_from_filelike.call_count, 6) - @patch("gigl.src.common.utils.file_loader.GcsUtils") + @patch("gigl.common.utils.file_loader.GcsUtils") def test_skips_flush_if_empty(self, mock_gcs_utils_class): gcs_base_uri = GcsUri("gs://test-bucket/test-folder") diff --git a/tests/unit/src/common/utils/bq_test.py b/tests/unit/common/utils/bq_test.py similarity index 98% rename from tests/unit/src/common/utils/bq_test.py rename to tests/unit/common/utils/bq_test.py index efa02e342..62259ecc7 100644 --- a/tests/unit/src/common/utils/bq_test.py +++ b/tests/unit/common/utils/bq_test.py @@ -3,7 +3,7 @@ from parameterized import param, parameterized -from gigl.src.common.utils.bq import BqUtils +from gigl.common.utils.bq import BqUtils from tests.test_assets.test_case import TestCase @@ -60,7 +60,7 @@ def _make_mock_table_list_items( return mock_tables -@patch("gigl.src.common.utils.bq.bigquery.Client") +@patch("gigl.common.utils.bq.bigquery.Client") class GetLatestTableTest(TestCase): PREFIX = "myproject.mydataset.events_" diff --git a/tests/unit/common/utils/retry_test.py b/tests/unit/common/utils/retry_test.py index 7b4735065..2e75c79aa 100644 --- a/tests/unit/common/utils/retry_test.py +++ b/tests/unit/common/utils/retry_test.py @@ -2,7 +2,7 @@ from gigl.common.logger import Logger from gigl.common.utils.retry import retry -from gigl.src.common.utils.timeout import TimedOutException +from gigl.common.utils.timeout import TimedOutException from tests.test_assets.test_case import TestCase logger = Logger() diff --git a/tests/unit/common/utils/timeout_test.py b/tests/unit/common/utils/timeout_test.py index 0cc6afa6d..d58aa94d5 100644 --- a/tests/unit/common/utils/timeout_test.py +++ b/tests/unit/common/utils/timeout_test.py @@ -1,6 +1,6 @@ from time import sleep -from gigl.src.common.utils.timeout import TimedOutException, timeout +from gigl.common.utils.timeout import TimedOutException, timeout from tests.test_assets.test_case import TestCase diff --git a/tests/unit/src/common/utils/gbml_config_test.py b/tests/unit/src/common/utils/gbml_config_test.py index 1c1df4662..4bcd692e0 100644 --- a/tests/unit/src/common/utils/gbml_config_test.py +++ b/tests/unit/src/common/utils/gbml_config_test.py @@ -2,9 +2,9 @@ from uuid import uuid4 from gigl.common import LocalUri +from gigl.common.utils.file_loader import FileLoader from gigl.common.utils.proto_utils import ProtoUtils from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper -from gigl.src.common.utils.file_loader import FileLoader from snapchat.research.gbml import gbml_config_pb2 from tests.test_assets.graph_metadata_constants import ( EXAMPLE_HETEROGENEOUS_CONDENSED_EDGE_TYPES, diff --git a/tests/unit/src/config_populator/config_populator_functionality_test.py b/tests/unit/src/config_populator/config_populator_functionality_test.py index a4c177f95..9eb8d3268 100644 --- a/tests/unit/src/config_populator/config_populator_functionality_test.py +++ b/tests/unit/src/config_populator/config_populator_functionality_test.py @@ -5,13 +5,13 @@ from gigl.common import GcsUri, LocalUri from gigl.common.logger import Logger from gigl.common.metrics.base_metrics import NopMetricsPublisher +from gigl.common.utils.time import current_formatted_datetime from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.types.pb_wrappers.dataset_metadata import DatasetMetadataPbWrapper from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.common.types.task_metadata import TaskMetadataType from gigl.src.common.utils import metrics_service_provider -from gigl.src.common.utils.time import current_formatted_datetime from gigl.src.config_populator.config_populator import ConfigPopulator from snapchat.research.gbml import ( dataset_metadata_pb2,