Compare commits

..
Author SHA1 Message Date
Andrew Ferlitsch c0c2583f2d debug: timeout fix 2022-12-17 17:56:04 +00:00
Andrew Ferlitsch e83a341a00 fix: timeout issue for notebook test 2022-12-17 02:09:38 +00:00
1060 changed files with 64210 additions and 357087 deletions
+6 -33
View File
@@ -1,14 +1,5 @@
from typing import List
from ratemate import RateLimit
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--dry_run",
type=bool,
default=False)
args = parser.parse_args()
from resource_cleanup_manager import (
DatasetResourceCleanupManager,
ModelResourceCleanupManager,
@@ -16,15 +7,6 @@ from resource_cleanup_manager import (
ResourceCleanupManager,
MatchingEngineIndexEndpointResourceCleanupManager,
MatchingEngineIndexResourceCleanupManager,
FeatureStoreLegacyCleanupManager,
FeatureStoreCleanupManager,
PipelineJobCleanupManager,
TrainingJobCleanupManager,
HyperparameterTuningCleanupManager,
BatchPredictionJobCleanupManager,
ExperimentCleanupManager,
BucketCleanupManager,
ArtifactRegistryCleanupManager
)
rate_limit = RateLimit(max_count=25, per=60, greedy=False)
@@ -36,14 +18,12 @@ def run_cleanup_managers(managers: List[ResourceCleanupManager], is_dry_run: boo
print(f"Fetching {type_name}'s...")
resources = manager.list()
try:
print(f"Found {len(resources)} {type_name}'s")
except Exception as e:
print(f"{type_name} {e}")
print(f"Found {len(resources)} {type_name}'s")
for resource in resources:
try:
if not manager.is_deletable(resource):
continue
if is_dry_run:
resource_name = manager.resource_name(resource)
print(f"Will delete '{type_name}': {resource_name}")
@@ -56,7 +36,9 @@ def run_cleanup_managers(managers: List[ResourceCleanupManager], is_dry_run: boo
print("")
if args.dry_run:
is_dry_run = False
if is_dry_run:
print("Starting cleanup in dry run mode...")
# List of all cleanup managers
@@ -66,15 +48,6 @@ managers: List[ResourceCleanupManager] = [
ModelResourceCleanupManager(), # ModelResourceCleanupManager must follow EndpointResourceCleanupManager due to deployed models blocking model deletion.
MatchingEngineIndexEndpointResourceCleanupManager(),
MatchingEngineIndexResourceCleanupManager(),
FeatureStoreLegacyCleanupManager(),
FeatureStoreCleanupManager(),
PipelineJobCleanupManager(),
TrainingJobCleanupManager(),
HyperparameterTuningCleanupManager(),
BatchPredictionJobCleanupManager(),
ExperimentCleanupManager(), # Experiment missing _resource_noun
BucketCleanupManager(),
ArtifactRegistryCleanupManager()
]
run_cleanup_managers(managers=managers, is_dry_run=args.dry_run)
run_cleanup_managers(managers=managers, is_dry_run=is_dry_run)
@@ -1,26 +1,10 @@
'''
READ FIRST BEFORE MAKING CHANGES
- Create a convention for resources created from vertex-ai-samples GH. We already have one IIRC
- Only delete those objects as part of our clean-up script.
- Don't run any tests on python-docs-samples-tests project, especially ones that affect resources created outside of our purview
- Add --dry-run option to the clean-up script. This option will just output the list of resources the script will delete instead of actually deleting the resources.
- Have a larger conversation in DEE before touching any resources that were not created as part of vertex-ai-samples
'''
import os
import abc
from typing import Any, Type
from google.cloud import aiplatform
from google.cloud.aiplatform import base
from google.cloud.aiplatform_v1beta1 import (FeatureOnlineStoreAdminServiceClient,
FeatureOnlineStore)
from google.cloud import storage
from proto.datetime_helpers import DatetimeWithNanoseconds
PROJECT_ID = "python-docs-samples-tests"
REGION = "us-central1"
API_ENDPOINT = f"{REGION}-aiplatform.googleapis.com"
# If a resource was updated within this number of seconds, do not delete.
RESOURCE_UPDATE_BUFFER_IN_SECONDS = 60 * 60 * 8
@@ -85,7 +69,7 @@ class VertexAIResourceCleanupManager(ResourceCleanupManager):
def delete(self, resource):
resource.delete()
def get_seconds_since_modification(self, resource: Any) -> float:
def get_seconds_since_modification(self, resource: Any) -> bool:
update_time = resource.update_time
current_time = DatetimeWithNanoseconds.now(tz=update_time.tzinfo)
return (current_time - update_time).total_seconds()
@@ -113,10 +97,13 @@ class EndpointResourceCleanupManager(VertexAIResourceCleanupManager):
vertex_ai_resource = aiplatform.Endpoint
def delete(self, resource):
# TODO: Remove this once https://github.com/googleapis/python-aiplatform/issues/1441 is fixed
resource._sync_gca_resource()
for deployed_model_id in [
models.id for models in resource._gca_resource.deployed_models
]:
resource._undeploy(deployed_model_id=deployed_model_id)
resource.delete(force=True)
@@ -130,176 +117,3 @@ class MatchingEngineIndexResourceCleanupManager(VertexAIResourceCleanupManager):
class MatchingEngineIndexEndpointResourceCleanupManager(VertexAIResourceCleanupManager):
vertex_ai_resource = aiplatform.MatchingEngineIndexEndpoint
def delete(self, resource):
resource.undeploy_all()
resource.delete(force=True)
class FeatureStoreLegacyCleanupManager(VertexAIResourceCleanupManager):
# TODO: only deleting legacy
# not deleting ingestions jobs
# ingest_from_xxx methods do not return a job ID, there is no list command, aka no python way to delete
# not deleting batch serving jobs
# batch_serve_to_xxx methods do not return a job ID, there is no list command, aka no python way to delete
vertex_ai_resource = aiplatform.Featurestore
def resource_name(self, resource: Any) -> str:
return resource.name
def delete(self, resource):
resource.delete(force=True)
class FeatureStoreCleanupManager(VertexAIResourceCleanupManager):
# for FS 2.0
# TODO: use _v1beta1, and gapic clients
# delete features, feature groups, feature views, feature online stores
vertex_ai_resource = FeatureOnlineStore
admin_client = FeatureOnlineStoreAdminServiceClient(
client_options={"api_endpoint": API_ENDPOINT}
)
def resource_name(self, resource: Any) -> str:
return resource.name
def type_name(self) -> str:
return "FeatureOnlineStore"
def list(self) -> Any:
try:
return self.admin_client.list_feature_online_stores(parent=f"projects/{PROJECT_ID}/locations/{REGION}")
except Exception as e:
print(e)
return []
def delete(self, resource):
try:
self.admin_client.delete_feature_online_store(name=resource.name, force=True)
except Exception as e:
print(e)
class PipelineJobCleanupManager(VertexAIResourceCleanupManager):
vertex_ai_resource = aiplatform.PipelineJob
class TrainingJobCleanupManager(VertexAIResourceCleanupManager):
vertex_ai_resource = aiplatform.training_jobs._CustomTrainingJob
job_types = [
aiplatform.AutoMLImageTrainingJob,
aiplatform.AutoMLTextTrainingJob,
aiplatform.AutoMLTabularTrainingJob,
aiplatform.AutoMLVideoTrainingJob,
aiplatform.AutoMLForecastingTrainingJob,
aiplatform.CustomJob,
aiplatform.CustomTrainingJob,
aiplatform.CustomContainerTrainingJob,
aiplatform.CustomPythonPackageTrainingJob
]
def list(self) -> Any:
return [
job
for job_type in self.job_types
for job in job_type.list()
]
class HyperparameterTuningCleanupManager(VertexAIResourceCleanupManager):
vertex_ai_resource = aiplatform.HyperparameterTuningJob
class BatchPredictionJobCleanupManager(VertexAIResourceCleanupManager):
vertex_ai_resource = aiplatform.BatchPredictionJob
class ExperimentCleanupManager(VertexAIResourceCleanupManager):
vertex_ai_resource = aiplatform.Experiment
@property
def type_name(self) -> str:
return "Experiment"
def resource_name(self, resource: Any) -> str:
return resource.name
def get_seconds_since_modification(self, resource: Any) -> float:
update_time = resource._metadata_context.update_time
current_time = DatetimeWithNanoseconds.now()
return float(current_time.timestamp() - update_time.timestamp())
class BucketCleanupManager(ResourceCleanupManager):
vertex_ai_resource = storage.bucket.Bucket
def list(self) -> Any:
storage_client = storage.Client()
return list(storage_client.list_buckets())
def delete(self, resource):
try:
resource.delete(force=True)
except Exception as e:
print(e)
@property
def type_name(self) -> str:
return "Bucket"
def get_seconds_since_modification(self, resource: Any) -> float:
# Bucket has no last_update property, only time created
created_time = resource.time_created
current_time = DatetimeWithNanoseconds.now()
return float(current_time.timestamp() - created_time.timestamp())
def resource_name(self, resource: Any) -> str:
return resource.name
def is_deletable(self, resource: Any) -> bool:
time_difference = self.get_seconds_since_modification(resource)
if not self.resource_name(resource).startswith('your-bucket-name'):
print(f"Skipping '{resource}' not a Vertex AI notebook bucket")
return False
# Check that it wasn't created too recently, to prevent race conditions
if time_difference <= RESOURCE_UPDATE_BUFFER_IN_SECONDS:
print(
f"Skipping '{resource}' due to update_time being '{time_difference}', which is less than '{RESOURCE_UPDATE_BUFFER_IN_SECONDS}'."
)
return False
return True
class ArtifactRegistryCleanupManager(ResourceCleanupManager):
vertex_ai_resource = "Artifact Registry"
def list(self) -> Any:
import subprocess
result = subprocess.run(["gcloud artifacts repositories list --location=us-central1"],
shell=True, capture_output=True, text=True)
ret = []
lines = result.stdout.split('\n')[2:]
for line in lines:
repo = line.split(' ')[0]
if repo.startswith("my-docker-repo"):
ret.append(repo)
return ret
def delete(self, resource):
os.system(f"! gcloud artifacts repositories delete {resource} --location=us-central1")
@property
def type_name(self) -> str:
return "ArtifactRepository"
def resource_name(self, resource: Any) -> str:
return resource
# delete repository regardless of age
def get_seconds_since_modification(self, resource: Any) -> float:
return RESOURCE_UPDATE_BUFFER_IN_SECONDS + 1
def is_deleteable(self, resource: Any) -> bool:
return True
+14 -103
View File
@@ -17,8 +17,6 @@
import argparse
import pathlib
import os
import csv
import execute_changed_notebooks_helper
@@ -38,22 +36,9 @@ parser = argparse.ArgumentParser(description="Run changed notebooks.")
parser.add_argument(
"--test_paths_file",
type=pathlib.Path,
help="The path to the file that has newline-delimited folders of notebooks that should be tested.",
help="The path to the file that has newline-limited folders of notebooks that should be tested.",
required=True,
)
parser.add_argument(
"--test_percent",
type=int,
help="The percent of notebooks to be tested (between 1 and 100).",
required=False,
default=100,
)
parser.add_argument(
"--build_id",
type=str,
help="The build id (which may be a Cloud Build job specific or user explicit.",
required=True
)
parser.add_argument(
"--base_branch",
help="The base git branch to diff against to find changed files.",
@@ -122,98 +107,24 @@ parser.add_argument(
default=True,
help="Should run notebooks in parallel.",
)
parser.add_argument(
"--concurrent_notebooks",
type=int,
help="Maximum number of parallel notebook executions per minute",
default=10,
required=False,
)
parser.add_argument(
"--run_first_file",
type=pathlib.Path,
help="The path to the file that has newline-delimited of notebooks to run in the first batch",
default=None,
required=False,
)
parser.add_argument(
"--aiplatform_whl",
type=str,
help="The GCS path to a whl version google-cloud-aiplatform",
default=None,
required=False,
)
parser.add_argument(
"--dry_run",
type=str2bool,
default=False,
help="Dry run for testing - no execution",
)
args = parser.parse_args()
changed_notebooks = execute_changed_notebooks_helper.get_changed_notebooks(
notebooks = execute_changed_notebooks_helper.get_changed_notebooks(
test_paths_file=args.test_paths_file,
base_branch=args.base_branch,
)
results_bucket = f"{args.artifacts_bucket}"
# artifacts_bucket may get set by trigger to a full gs:// folder path
if results_bucket.startswith("gs://"):
results_bucket = results_bucket[5:]
results_bucket = results_bucket.split('/')[0]
results_file = f"build_results/{args.build_id}.json"
if args.test_percent == 100:
notebooks = changed_notebooks
accumulative_results = {}
else:
accumulative_results = execute_changed_notebooks_helper.load_results(results_bucket, results_file)
notebooks = [changed_notebook for changed_notebook in changed_notebooks if execute_changed_notebooks_helper.select_notebook(changed_notebook, accumulative_results, args.test_percent)]
# cap the number of notebooks to the specified percentage
max_notebooks = int((len(changed_notebooks) * (args.test_percent/100)))
if (len(notebooks) > max_notebooks):
notebooks = notebooks[:max_notebooks]
run_first = []
if args.run_first_file:
if not os.path.isfile(args.run_first_file):
print("Error: file does not exist", args.run_first_file)
else:
with open(args.run_first_file, 'r') as csvfile:
reader = csv.reader(csvfile)
for row in reader:
notebook = row[0]
run_first.append(notebook)
for notebook in run_first:
if notebook in notebooks:
# remove from existing list
notebooks.remove(notebook)
# add back to the front of the list
notebooks.insert(0, notebook)
print(f"Run first: {notebook}")
if args.dry_run:
print("Dry run ...\n")
for notebook in notebooks:
print(f"Would execute: {notebook}")
else:
execute_changed_notebooks_helper.process_and_execute_notebooks(
notebooks=notebooks,
container_uri=args.container_uri,
staging_bucket=args.staging_bucket,
artifacts_bucket=args.artifacts_bucket,
results_file=results_file,
should_parallelize=args.should_parallelize,
timeout=args.timeout,
variable_project_id=args.variable_project_id,
variable_region=args.variable_region,
variable_service_account=args.variable_service_account,
variable_vpc_network=args.variable_vpc_network,
private_pool_id=args.private_pool_id,
concurrent_notebooks=args.concurrent_notebooks,
aiplatform_whl=args.aiplatform_whl
execute_changed_notebooks_helper.process_and_execute_notebooks(
notebooks=notebooks,
container_uri=args.container_uri,
staging_bucket=args.staging_bucket,
artifacts_bucket=args.artifacts_bucket,
should_parallelize=args.should_parallelize,
timeout=args.timeout,
variable_project_id=args.variable_project_id,
variable_region=args.variable_region,
variable_service_account=args.variable_service_account,
variable_vpc_network=args.variable_vpc_network,
private_pool_id=args.private_pool_id,
)
+13 -183
View File
@@ -21,34 +21,25 @@ import json
import git
import operator
import os
import io
import json
import pathlib
import re
import subprocess
import random
from google.cloud import storage
import utils
from typing import List, Optional, Dict, Any
from typing import List, Optional
from utils import util
import execute_notebook_helper
import execute_notebook_remote
import nbformat
from google.cloud.devtools.cloudbuild_v1.types import BuildOperationMetadata
from ratemate import RateLimit
from tabulate import tabulate
from utils import NotebookProcessors, util
# A buffer so that workers finish before the orchestrating job
WORKER_TIMEOUT_BUFFER_IN_SECONDS: int = 60 * 60
PYTHON_VERSION = "3.9" # Set default python version
# rolling time window for accumulating build results for selecting notebooks
MAX_RESULTS_AGE_SECONDS: int = (60 * 60) * 24 * 60 # 60 days
# maximum time since last run to force a run on the current build
MAX_AGE_BEFORE_FORCE_RUN: int = (60 * 60) * 24 * 30
def format_timedelta(delta: datetime.timedelta) -> str:
"""Formats a timedelta duration to [N days] %H:%M:%S format"""
@@ -74,9 +65,7 @@ def format_timedelta(delta: datetime.timedelta) -> str:
@dataclasses.dataclass
class NotebookExecutionResult:
name: str
path: str
duration: datetime.timedelta
start_time: datetime.datetime
is_pass: bool
log_url: str
output_uri: str
@@ -92,97 +81,6 @@ class NotebookExecutionResult:
return None
def load_results(results_bucket: str,
results_file: str) -> Dict[str, Any]:
'''
Load accumulated notebook test results
'''
print("Loading existing accumulative results ...")
accumulative_results = {}
try:
client = storage.Client()
bucket = client.bucket(results_bucket)
build_results_dir = os.path.dirname(results_file)
blobs = client.list_blobs(results_bucket, prefix=build_results_dir)
for blob in blobs:
time_created = blob.time_created.replace(tzinfo=None)
if (datetime.datetime.now().replace(tzinfo=None) - time_created).total_seconds() > MAX_RESULTS_AGE_SECONDS:
continue
content = util.download_blob_into_memory(results_bucket, blob.name, download_as_text=True)
try:
build_results = json.loads(content)
except:
continue # skip corrupted build results files
for notebook in build_results:
if notebook in accumulative_results:
accumulative_results[notebook]['passed'] += build_results[notebook]['passed']
accumulative_results[notebook]['failed'] += build_results[notebook]['failed']
if accumulative_results[notebook]['last_time_ran'] < time_created:
accumulative_results[notebook]['last_time_ran'] = time_created
else:
accumulative_results[notebook] = build_results[notebook]
accumulative_results[notebook]['failed_on_latest_run'] = build_results[notebook]['failed']
accumulative_results[notebook]['last_time_ran'] = time_created
print(accumulative_results)
except Exception as e:
print(e)
# If there are no accumulative results, an empty dict is returned
return accumulative_results
def select_notebook(changed_notebook: str,
accumulative_results: Dict[str, Any],
test_percent: int) -> bool:
'''
Algorithm to randomly select a notebook, but weight the propbability of selected based on past failures
'''
if changed_notebook in accumulative_results:
pass_count = accumulative_results[changed_notebook]['passed']
fail_count = accumulative_results[changed_notebook]['failed']
failed_on_latest_run = accumulative_results[changed_notebook]['failed_on_latest_run']
last_time_ran = accumulative_results[changed_notebook]['last_time_ran']
else:
pass_count = 1
fail_count = 0
failed_on_latest_run = 0
last_time_ran = datetime.datetime.now().replace(tzinfo=None)
# If notebook has not been ran in a long time, force running it
if (datetime.datetime.now().replace(tzinfo=None) - last_time_ran).total_seconds() > MAX_AGE_BEFORE_FORCE_RUN:
should_test_do_to_age = True
else:
should_test_do_to_age = False
# if failed on the last time it was ran, select the notebook
if failed_on_latest_run:
inferred_failure_rate = 1
# otherwise, calculate the frequency of failure
else:
inferred_failure_rate = fail_count / (pass_count + fail_count)
# If failure rate is high, the chance of testing should be higher
should_test_due_to_failure = random.uniform(0, 1) <= inferred_failure_rate
#if accumulative_resultsi[changed_notebook]['latest_date_ran']
# Additionally, only test a percentage of these
should_test_due_to_random_subset = random.uniform(0, 1) <= (test_percent / 100)
if should_test_due_to_failure or should_test_due_to_random_subset or should_test_do_to_age:
print(f"Selected: {changed_notebook}, {should_test_due_to_failure}, {should_test_due_to_random_subset}")
return True
else:
print(f"Not Selected: {changed_notebook}, pass {pass_count}, fail {fail_count}")
return False
def _process_notebook(
notebook_path: str,
variable_project_id: str,
@@ -238,7 +136,7 @@ def _get_notebook_python_version(notebook_path: str) -> str:
# Look for the python version specification pattern
re_match = re.search(
r"python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
"python version = (\d\.\d)", markdown, flags=re.IGNORECASE
)
if re_match:
# get the version number
@@ -258,6 +156,7 @@ def _create_tag(filepath: str) -> str:
return tag
rate_limit = RateLimit(max_count=50, per=60, greedy=True)
def process_and_execute_notebook(
@@ -271,8 +170,9 @@ def process_and_execute_notebook(
private_pool_id: Optional[str],
deadline: datetime.datetime,
notebook: str,
should_get_tail_logs: bool = True,
should_get_tail_logs: bool = False,
) -> NotebookExecutionResult:
rate_limit.wait() # wait before creating the task
print(f"Running notebook: {notebook}")
@@ -291,9 +191,7 @@ def process_and_execute_notebook(
result = NotebookExecutionResult(
name=tag,
path=notebook,
duration=datetime.timedelta(seconds=0),
start_time=datetime.datetime.now(),
is_pass=False,
output_uri=notebook_output_uri,
log_url="",
@@ -303,6 +201,7 @@ def process_and_execute_notebook(
)
# TODO: Handle cases where multiple notebooks have the same name
time_start = datetime.datetime.now()
operation = None
try:
# Get the python version for running the notebook if specified
@@ -346,12 +245,11 @@ def process_and_execute_notebook(
result.logs_bucket = operation_metadata.build.logs_bucket
# Block and wait for the result
operation_result = operation.result(timeout=timeout_in_seconds)
operation_result = operation.result(timeout=86400)
result.duration = datetime.datetime.now() - result.start_time
result.duration = datetime.datetime.now() - time_start
result.is_pass = True
print(f"{notebook} PASSED in {format_timedelta(result.duration)}.")
except Exception as error:
result.error_message = str(error)
@@ -365,12 +263,12 @@ def process_and_execute_notebook(
# Use gcloud to get tail
try:
result.error_message = subprocess.check_output(
["gcloud", "storage", "cat", "--range", "-1000", log_file_uri], encoding="UTF-8"
["gsutil", "cat", "-r", "-1000", log_file_uri], encoding="UTF-8"
)
except Exception as error:
result.error_message = str(error)
result.duration = datetime.datetime.now() - result.start_time
result.duration = datetime.datetime.now() - time_start
result.is_pass = False
print(
@@ -438,68 +336,12 @@ def get_changed_notebooks(
return notebooks
def _save_results(results: List[NotebookExecutionResult],
artifacts_bucket: str,
results_file: str):
artifacts_bucket = artifacts_bucket.replace("gs://", "").split('/')[0]
print("Updating build results ...")
build_results = {}
for result in results:
if result.is_pass:
pass_count = 1
fail_count = 0
else:
pass_count = 0
fail_count = 1
if result.error_message is None:
error_type = ''
elif '500 Internal' in result.error_message or 'INTERNAL' in result.error_message or 'internal error' in result.error_message:
error_type = 'INTERNAL'
elif 'context deadline exceeded' in result.error_message or 'TIMEOUT' in result.error_message:
error_type = 'TIMEOUT'
elif 'Quota' in result.error_message or 'quotas are exceeded' in result.error_message:
error_type = 'QUOTA'
elif 'ServiceUnavailable' in result.error_message:
error_type = 'SERVICEUNAVAILABLE'
elif 'ModuleNotFoundError' in result.error_message:
error_type = 'IMPORT'
elif result.is_pass:
error_type = ''
else:
error_type = 'undetermined'
if error_type != '':
log_url = result.log_url
else:
log_url = ''
build_results[result.path] = {
'duration': result.duration.total_seconds(),
'start_time': str(result.start_time),
'passed': pass_count,
'failed': fail_count,
'error_type': error_type,
'log_url': log_url
}
print(f"adding {result.path}")
print(f"Saving accumulative results to {results_file}, nentries {len(build_results)}")
content = json.dumps(build_results)
client = storage.Client()
bucket = client.get_bucket(artifacts_bucket)
bucket.blob(str(results_file)).upload_from_string(content, 'text/json')
def process_and_execute_notebooks(
notebooks: List[str],
container_uri: str,
staging_bucket: str,
artifacts_bucket: str,
results_file: str,
should_parallelize: bool,
timeout: int,
variable_project_id: str,
@@ -507,8 +349,6 @@ def process_and_execute_notebooks(
variable_service_account: str,
variable_vpc_network: Optional[str] = None,
private_pool_id: Optional[str] = None,
concurrent_notebooks: Optional[int] = 10,
aiplatform_whl: Optional[str] = None,
):
"""
Run the notebooks that exist under the folders defined in the test_paths_file.
@@ -529,8 +369,6 @@ def process_and_execute_notebooks(
Required. The GCS staging bucket to write source code to.
artifacts_bucket (str):
Required. The GCS staging bucket to write executed notebooks to.
results_file (str):
Required: The path to the artifacts bucket to save results
variable_project_id (str):
Required. The value for PROJECT_ID to inject into notebooks.
variable_region (str):
@@ -539,8 +377,6 @@ def process_and_execute_notebooks(
Required. Should run notebooks in parallel using a thread pool as opposed to in sequence.
timeout (str):
Required. Timeout string according to https://cloud.google.com/build/docs/build-config-file-schema#timeout.
concurrent_notebooks (int): Max number of notebooks per minute to run in parallel.
aiplatform_whl: alternate whl version of Vertex AI SDK to install
"""
# Calculate deadline
@@ -557,9 +393,7 @@ def process_and_execute_notebooks(
print(
"Running notebooks in parallel, so no logs will be displayed. Please wait..."
)
with concurrent.futures.ThreadPoolExecutor(max_workers=concurrent_notebooks) as executor:
with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor:
print(f"Max workers: {executor._max_workers}")
notebook_execution_results = list(
@@ -637,7 +471,7 @@ def process_and_execute_notebooks(
print("=" * 100)
build_id = results_sorted[0].build_id
logs_bucket_name = (results_sorted[0].logs_bucket).replace("gs://", "")
logs_bucket_name = (results_sorted[0].logs_bucket).removeprefix("gs://")
log_file_name = f"log-{build_id}.txt"
log_contents = util.download_blob_into_memory(
@@ -655,10 +489,6 @@ def process_and_execute_notebooks(
else:
print(log_contents)
_save_results(results_sorted,
artifacts_bucket,
results_file)
print("\n=== END RESULTS===\n")
total_notebook_duration = functools.reduce(
+2 -2
View File
@@ -56,8 +56,8 @@ def execute_notebook(
print("\n=== DOWNLOAD EXECUTED NOTEBOOK ===\n")
print(f"Please debug the executed notebook by downloading the executed notebook:")
print("Option 1. Using gcloud storage. Run the following command in your terminal.")
print(f'\tgcloud storage cp "{output_file_or_uri}" .')
print("Option 1. Using gsutil. Run the following command in your terminal.")
print(f'\tgsutil cp "{output_file_or_uri}" .')
print("Option 2. Using this link.")
print(f"\thttps://storage.googleapis.com/{output_file_or_uri[5:]}")
@@ -36,7 +36,7 @@ steps:
- -c
- |
. workspace/env/bin/activate &&
python3 .cloud-build/execute_changed_notebooks_cli.py --test_paths_file "${_TEST_PATHS_FILE}" --base_branch "${_FORCED_BASE_BRANCH}" --container_uri ${_PYTHON_IMAGE} --staging_bucket ${_GCS_STAGING_BUCKET} --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION} --variable_service_account ${_GCP_SERVICE_ACCOUNT} --variable_vpc_network "${_GCP_VPC_NETWORK_NAME}" `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi` --build_id ${BUILD_ID} --test_percent=${_TEST_PERCENT} --concurrent_notebooks=${_CONCURRENT_NOTEBOOKS} --run_first_file=${_RUN_FIRST_FILE}
python3 .cloud-build/execute_changed_notebooks_cli.py --test_paths_file "${_TEST_PATHS_FILE}" --base_branch "${_FORCED_BASE_BRANCH}" --container_uri ${_PYTHON_IMAGE} --staging_bucket ${_GCS_STAGING_BUCKET} --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION} --variable_service_account ${_GCP_SERVICE_ACCOUNT} --variable_vpc_network "${_GPC_VPC_NETWORK_NAME}" `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi`
env:
- 'IS_TESTING=1'
timeout: 86400s
+2 -6
View File
@@ -3,15 +3,11 @@ numpy
jupyter
nbconvert
papermill
pandas
matplotlib
tabulate
google-cloud-aiplatform
google-cloud-storage
google-cloud-build
google-cloud-storage
ratemate
GitPython
tqdm
fsspec
pandas
GitPython
-8
View File
@@ -1,8 +0,0 @@
notebooks/official/pipelines/google_cloud_pipeline_components_automl_text.ipynb
notebooks/official/generative_ai/rlhf_tune_llm.ipynb
notebooks/official/generative_ai/tune_peft.ipynb
notebooks/official/prediction/llm_streaming_prediction.ipynb
notebooks/official/migration/sdk-automl-text-classification-batch-prediction.ipynb
notebooks/official/vizier/get_started_vertex_vizier.ipynb
notebooks/official/workbench/sentiment_analysis/Sentiment_Analysis.ipynb
notebooks/official/model_monitoring/get_started_with_model_monitoring_automl.ipynb
1 notebooks/official/pipelines/google_cloud_pipeline_components_automl_text.ipynb
2 notebooks/official/generative_ai/rlhf_tune_llm.ipynb
3 notebooks/official/generative_ai/tune_peft.ipynb
4 notebooks/official/prediction/llm_streaming_prediction.ipynb
5 notebooks/official/migration/sdk-automl-text-classification-batch-prediction.ipynb
6 notebooks/official/vizier/get_started_vertex_vizier.ipynb
7 notebooks/official/workbench/sentiment_analysis/Sentiment_Analysis.ipynb
8 notebooks/official/model_monitoring/get_started_with_model_monitoring_automl.ipynb
+2
View File
@@ -1,3 +1,5 @@
notebooks/official/vizier/gapic-vizier-multi-objective-optimization.ipynb
notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb
notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb
notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb
.cloud-build/tests/python_version_test.ipynb
-46
View File
@@ -1,46 +0,0 @@
# grep PASSED tests.txt | cut -c 10-100 >passed.txt
import os
repo_dir = '/home/jupyter/vertex-ai-samples/'
repo_dir_len = len(repo_dir)
official_dir = repo_dir + 'notebooks/official'
entries = os.scandir(official_dir)
folders = []
for entry in entries:
if entry.is_dir():
folders.append(entry.path)
# Passing
with open('passed.txt', 'r') as pass_file:
notebook_names = pass_file.readlines()
notebooks = []
for folder in folders:
entries = os.scandir(folder)
for entry in entries:
for notebook in notebook_names:
if entry.name == notebook.rstrip():
notebooks.append(entry.path[repo_dir_len:])
with open('passing_tests.txt', 'w') as f:
for notebook in notebooks:
f.write(notebook + '\n')
# Failing
with open('failed.txt', 'r') as fail_file:
notebook_names = fail_file.readlines()
notebooks = []
for folder in folders:
entries = os.scandir(folder)
for entry in entries:
for notebook in notebook_names:
if entry.name == notebook.rstrip():
notebooks.append(entry.path[repo_dir_len:])
with open('failing_tests.txt', 'w') as f:
for notebook in notebooks:
f.write(notebook + '\n')
@@ -1,33 +0,0 @@
import sys
from execute_changed_notebooks_helper import (load_results, select_notebook)
def test_load_results():
bucket: str = "cloud-build-notebooks-presubmit"
bucket_file: str = "build_results"
accum = load_results(bucket, bucket_file)
print(accum)
assert len(accum) > 0
def test_select_notebook():
bucket: str = "cloud-build-notebooks-presubmit"
bucket_file: str = "build_results"
accum = load_results(bucket, bucket_file)
n_select = 0
n_notselect = 0
for notebook in accum:
if select_notebook(notebook, accum, 50):
n_select += 1
else:
n_notselect += 1
print(f"SELECTED {n_select}, NOT SELECTED {n_notselect}")
assert n_select > 0
assert n_notselect > 0
+1 -26
View File
@@ -35,7 +35,7 @@ class RemoveNoExecuteCells(Preprocessor):
class UpdateVariablesPreprocessor(Preprocessor):
def __init__(self, replacement_map: Dict[str, str]):
def __init__(self, replacement_map: Dict):
self._replacement_map = replacement_map
@staticmethod
@@ -98,28 +98,3 @@ class UniqueStringsPreprocessor(Preprocessor):
executable_cells.append(cell)
notebook.cells = executable_cells
return notebook, resources
class VertexAIInstallProprocessor(Preprocessor):
def __init__(self, vertex_ai_wheel):
self.vertex_ai_wheel = vertex_ai_wheel
@staticmethod
def update_vertex_ai_install(content: str):
if "google-cloud-aiplatform" not in content:
return content
return (
f"gcloud storage cp {self.vertex_ai_wheel} google-cloud-aiplatform.whl\n" +
content.replace("google-cloud-aiplatform\n", "google-cloud-aiplatform.whl\n")
.replace("google-cloud-aiplatform ", "google-cloud-aiplatform.whl ")
)
def preprocess(self, notebook, resources=None):
executable_cells = []
for cell in notebook.cells:
if cell.cell_type == "code":
cell.source = self.update_vertex_ai_install(
content=cell.source,
)
executable_cells.append(cell)
notebook.cells = executable_cells
@@ -1,73 +0,0 @@
'''
Viewer for the weekly regression testing of the official notebooks
Cloud Storage location: gs://cloud-build-notebooks-presubmit/build_results/
'''
import argparse
import json
from util import download_file
import csv
import datetime
from google.cloud import storage
BUILD_BUCKET = "cloud-build-notebooks-presubmit"
BUILD_FOLDER = "build_results"
parser = argparse.ArgumentParser()
parser.add_argument('--file', dest='file',
default=None, type=str, help='build results filei (local or GCS)')
args = parser.parse_args()
investigate = {}
with open('investigate.csv', 'r') as csvfile:
reader = csv.reader(csvfile)
for row in reader:
investigate[row[0][:-6]] = row[1]
if not args.file:
client = storage.Client()
blobs = client.list_blobs(BUILD_BUCKET, prefix=BUILD_FOLDER)
newest_time = datetime.datetime(2000, 1, 1)
for blob in blobs:
# individual PR
if blob.size < 2000:
continue
time_created = blob.time_created.replace(tzinfo=None)
if time_created > newest_time:
newest_time = time_created
args.file = f"gs://{BUILD_BUCKET}/{blob.name}"
if args.file.startswith("gs://"):
path = args.file[5:]
bucket = path.split('/')[0]
file = path[len(bucket)+1:]
download_file(bucket, file, "build.json")
args.file = "build.json"
with open(args.file, 'r') as f:
results = json.load(f)
for item in results.items():
notebook = item[0][len("/notebooks/official/")-1:-6]
if item[1]['passed']:
passed = "PASS"
else:
if notebook in investigate:
passed = "INVG"
else:
passed = "FAIL"
error = item[1]['error_type']
if passed == "FAIL":
if error == '':
error = "undetermined"
if 'log_url' in item[1]:
log_url = item[1]['log_url']
else:
log_url = ''
else:
log_url = ''
print(f"{notebook:75} {passed} {error:10} {log_url}")
-19
View File
@@ -1,19 +0,0 @@
notebook,status
prediction/llm_streaming_prediction.ipynb,wait_for_fix
custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb,issue 2527
feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store.ipynb,wait_for_reaper
feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb,wait_for_reaper
pipelines/google_cloud_pipeline_components_TPU_model_train_upload_deploy.ipynb,wait_for_fix
explainable_ai/sdk_custom_image_classification_batch_explain.ipynb,issue 2528
explainable_ai/sdk_custom_image_classification_online_explain.ipynb,issue 2528
explainable_ai/sdk_custom_tabular_regression_online_explain.ipynb,issue 2528
explainable_ai/sdk_automl_tabular_binary_classification_batch_explain.ipynb,issue 2528
explainable_ai/xai_image_classification_feature_attributions.ipynb,issue 2528
matching_engine,sdk_matching_engine_create_stack_overflow_embeddings.ipynb,issue 2530
automl/automl_forecasting_bqml_arima_plus_comparison.ipynb,flaky
model_evaluation/custom_tabular_regression_model_evaluation.ipynb,regr
experiments/get_started_with_vertex_experiments.ipynb,regr
experiments/comparing_local_trained_models.ipynb,regr
generative_ai/tune_peft.ipynb,internal
pipelines/custom_model_training_and_batch_prediction.ipynb,regr
feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store_optimized.ipynb,wait_for_reaper
1 notebook status
2 prediction/llm_streaming_prediction.ipynb wait_for_fix
3 custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb issue 2527
4 feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store.ipynb wait_for_reaper
5 feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb wait_for_reaper
6 pipelines/google_cloud_pipeline_components_TPU_model_train_upload_deploy.ipynb wait_for_fix
7 explainable_ai/sdk_custom_image_classification_batch_explain.ipynb issue 2528
8 explainable_ai/sdk_custom_image_classification_online_explain.ipynb issue 2528
9 explainable_ai/sdk_custom_tabular_regression_online_explain.ipynb issue 2528
10 explainable_ai/sdk_automl_tabular_binary_classification_batch_explain.ipynb issue 2528
11 explainable_ai/xai_image_classification_feature_attributions.ipynb issue 2528
12 matching_engine sdk_matching_engine_create_stack_overflow_embeddings.ipynb issue 2530
13 automl/automl_forecasting_bqml_arima_plus_comparison.ipynb flaky
14 model_evaluation/custom_tabular_regression_model_evaluation.ipynb regr
15 experiments/get_started_with_vertex_experiments.ipynb regr
16 experiments/comparing_local_trained_models.ipynb regr
17 generative_ai/tune_peft.ipynb internal
18 pipelines/custom_model_training_and_batch_prediction.ipynb regr
19 feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store_optimized.ipynb wait_for_reaper
+2 -2
View File
@@ -15,7 +15,7 @@ def download_file(bucket_name: str, blob_name: str, destination_file: str) -> st
remote_file_path = "".join(["gs://", "/".join([bucket_name, blob_name])])
subprocess.check_output(
["gcloud", "storage", "cp", remote_file_path, destination_file], encoding="UTF-8"
["gsutil", "cp", remote_file_path, destination_file], encoding="UTF-8"
)
return destination_file
@@ -27,7 +27,7 @@ def upload_file(
) -> str:
"""Copies a local file to a GCS path"""
subprocess.check_output(
["gcloud", "storage", "cp", local_file_path, remote_file_path], encoding="UTF-8"
["gsutil", "cp", local_file_path, remote_file_path], encoding="UTF-8"
)
return remote_file_path
-11
View File
@@ -1,11 +0,0 @@
sdk2_remote_tabnet_training.ipynb
remote_hyperparameter_tuning.ipynb
remote_prediction.ipynb
remote_training_bigframes_pytorch.ipynb
remote_training_bigframes_sklearn.ipynb
remote_training_bigframes_tensorflow.ipynb
remote_training_lightning.ipynb
remote_training_pytorch.ipynb
remote_training_sklearn.ipynb
remote_training_tensorflow_with_autologging.ipynb
@@ -1,23 +0,0 @@
steps:
# Fetch full repo for diff purposes
- name: gcr.io/cloud-builders/git
args: [fetch, --unshallow, --quiet]
# Create a virtual environment
- name: ${_PYTHON_IMAGE}
entrypoint: /bin/sh
args:
- -c
- python3 -m venv workspace/env
# Install Python dependencies and run testing script
- name: ${_PYTHON_IMAGE}
entrypoint: /bin/sh
args:
- -c
- |
. workspace/env/bin/activate &&
python3 notebooks/notebook_template_review.py --web --title --steps --desc --linkback --notebook-dir=notebooks/official --skip-file=${_DO_NOT_INDEX_FILE} >web.html
artifacts:
objects:
location: gs://${_GCS_ARTIFACTS_BUCKET}/webdoc
paths: ['web.html']
timeout: 86400s
-10
View File
@@ -1,10 +0,0 @@
version: 2
updates:
# Ignore model garden dockerfiles:
- package-ecosystem: "npm"
directory: "/community-content/vertex_model_garden"
schedule:
interval: "monthly"
ignore:
- dependency-name: "*"
+3 -3
View File
@@ -7,11 +7,11 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Set up Python
uses: actions/setup-python@v6
uses: actions/setup-python@v4
with:
python-version: '3.12'
python-version: '3.x'
- name: Fetch pull request branch
uses: actions/checkout@v6
uses: actions/checkout@v3
with:
fetch-depth: 0
- name: Fetch base main branch
+1 -1
View File
@@ -4,7 +4,7 @@
# 2. To lint specific notebooks:
# docker run -v ${PWD}:/setup/app gcr.io/python-docs-samples-tests/notebook_linter:latest notebooks/1.ipynb notebooks/2.ipynb
FROM python:3.14
FROM python:3.10
WORKDIR setup
+5 -5
View File
@@ -2,9 +2,9 @@ git+https://github.com/tensorflow/docs
ipython
jupyter
nbconvert
black==26.1.0
pyupgrade==3.21.2
isort==7.0.0
flake8==7.3.0
nbqa==1.9.1
black==22.10.0
pyupgrade==2.38.4
isort==5.10.1
flake8==4.0.1
nbqa==1.5.3
+3 -3
View File
@@ -58,7 +58,7 @@ done
# Only check notebooks in test folders modified in this pull request.
# Note: Use process substitution to persist the data in the array
if [ ${#notebooks[@]} -eq 0 ]; then
echo "Checking for changed notebooks using git"
echo "Checking for changed notebooked using git"
while read -r file || [ -n "$line" ]; do
notebooks+=("$file")
done < <(git diff --name-only main... | grep '\.ipynb$')
@@ -84,7 +84,7 @@ if [ ${#notebooks[@]} -gt 0 ]; then
# python3 -m nbqa black "$notebook" --check
# BLACK_RTN=$?
echo "Running pyupgrade..."
python3 -m nbqa pyupgrade --exit-zero-even-if-changed "$notebook"
python3 -m nbqa pyupgrade "$notebook"
PYUPGRADE_RTN=$?
echo "Running isort..."
python3 -m nbqa isort "$notebook" --check
@@ -97,7 +97,7 @@ if [ ${#notebooks[@]} -gt 0 ]; then
python3 -m nbqa black "$notebook"
BLACK_RTN=$?
echo "Running pyupgrade..."
python3 -m nbqa pyupgrade --exit-zero-even-if-changed "$notebook"
python3 -m nbqa pyupgrade "$notebook"
PYUPGRADE_RTN=$?
echo "Running isort..."
python3 -m nbqa isort "$notebook"
+5 -3
View File
@@ -44,10 +44,12 @@ Finally, run this code block to check for errors. Each step will attempt to
automatically fix any issues. If the fixes can't be performed automatically,
then you will need to manually address them before submitting your PR.
Note: For official, only submit one notebook per PR.
```shell
docker run -v ${PWD}:/setup/app gcr.io/cloud-devrel-public-resources/notebook_linter:latest your_notebook
nbqa black "$notebook"
nbqa pyupgrade "$notebook"
nbqa isort "$notebook"
nbqa flake8 "$notebook" --extend-ignore=W391,E501,F821,E402,F404,W503,E203,E722,W293,W291
python3 -m tensorflow_docs.tools.nbfmt --remove_outputs "$notebook"
```
## Code Reviews
+13 -152
View File
@@ -1,176 +1,37 @@
# ![Google Cloud](https://avatars.githubusercontent.com/u/2810941?s=60&v=4) Google Cloud Vertex AI Samples
# Google Cloud Vertex AI Samples
This repository contains notebooks, code samples, sample apps, and other resources that demonstrate how to use, develop and manage machine learning and generative AI workflows using Google Cloud Vertex AI.
[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](LICENSE)
Welcome to the Google Cloud [Vertex AI](https://cloud.google.com/vertex-ai/docs/) sample repository.
## Overview
[Vertex AI](https://cloud.google.com/vertex-ai) is a fully-managed, unified AI development platform for building and using generative AI. This repository is designed to help you get started with Vertex AI. Whether you're new to Vertex AI or an experienced ML practitioner, you'll find valuable resources here.
For more Vertex AI Generative AI notebook samples, please visit the Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository.
## Explore, learn and contribute
You can explore, learn, and contribute to this repository to unleash the full potential of machine learning on Vertex AI!
### Explore and learn
Explore this repository, follow the links in the header section of each of the notebooks to -
![Colab](https://cloud.google.com/ml-engine/images/colab-logo-32px.png) Open and run the notebook in [Colab](https://colab.google/)\
![Colab Enterprise](https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png) Open and run the notebook in [Colab Enterprise](https://cloud.google.com/colab/docs/introduction)\
![Workbench](https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32) Open and run the notebook in [Vertex AI Workbench](https://cloud.google.com/vertex-ai/docs/workbench/introduction)\
![Github](https://cloud.google.com/ml-engine/images/github-logo-32px.png) View the notebook on Github
### Contribute
See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
## Get started
To get started using Vertex AI, you must have a Google Cloud project.
- If you don't have a Google Cloud project, you can learn and build on GCP for free using [Free Trail](https://cloud.google.com/free).
- Once you have a Google Cloud project, you can learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).
The repository contains [notebooks](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/notebooks) and [community content](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/community-content) that demonstrate how to develop and manage ML workflows using Google Cloud Vertex AI.
## Repository structure
```bash
├── community-content - Sample code and tutorials contributed by the community
├── notebooks
│ ├── community - Notebooks contributed by the community
│ ├── official - Notebooks demonstrating use of each Vertex AI service
│ │ ├── automl
│ │ ├── custom
│ │ ├── ...
│ ├── community - Notebooks contributed by the community
│ │ ├── model_garden
│ │ ├── ...
├── community-content - Sample code and tutorials contributed by the community
```
## Examples
<!-- markdownlint-disable MD033 -->
<table>
## Contributing
<tr>
<th style="text-align: center;">Category</th>
<th style="text-align: center;">Product</th>
<th style="text-align: center;">Description</th>
</tr>
<tr>
<td>Model</td>
<td>
<a href="notebooks/community/model_garden"><code>Model Garden/</code></a>
</td>
<td>
Curated collection of first-party, open-source, and third-party models available on Vertex AI including Gemini, Gemma, Llama 3, Claude 3 and many more.
</td>
</tr>
<tr>
<td>Data</td>
<td>
<a href="notebooks/official/feature_store"><code>Feature Store/</code></a>
</td>
<td>
Set up and manage online serving using Vertex AI Feature Store.
</td>
</tr>
<tr>
<td></td>
<td>
<a href="notebooks/official/datasets"><code>datasets/</code></a>
</td>
<td>
Use BigQuery and Data Labeling service with Vertex AI.
</td>
</tr>
<tr>
<td>Model development</td>
<td>
<a href="notebooks/official/automl"><code>automl/</code></a>
</td>
<td>
Train and make predictions on AutoML models
</td>
</tr>
<tr>
<td></td>
<td>
<a href="notebooks/official/custom"><code>custom/</code></a>
</td>
<td>
Create, deploy and serve custom models on Vertex AI
</td>
</tr>
<tr>
<td></td>
<td>
<a href="notebooks/official/ray_on_vertex_ai"><code>ray_on_vertex_ai/</code></a>
</td>
<td>
Use Colab Enterprise and Vertex AI SDK for Python to connect to the Ray Cluster.
</td>
</tr>
<tr>
<td>Deploy and use</td>
<td>
<a href="notebooks/official/prediction"><code>prediction/</code></a>
</td>
<td>
Build, train and deploy models using prebuilt containers for custom training and prediction.
</td>
</tr>
<tr>
<td></td>
<td>
<a href="notebooks/official/model_registry"><code>model_registry/</code></a>
</td>
<td>
Use Model Registry to create and register a model.
</td>
</tr>
<tr>
<td></td>
<td>
<a href="notebooks/official/explainable_ai"><code>Explainable AI/</code></a>
</td>
<td>
Use Vertex Explainable AI's feature-based and example-based explanations to explain how or why a model produced a specific prediction.
</td>
</tr>
<tr>
<td></td>
<td>
<a href="notebooks/official/ml_metadata"><code>ml_metadata/</code></a>
</td>
<td>
Record the metadata and artifacts and query that metadata to help analyze, debug, and audit the performance of your ML system.
</td>
</tr>
<tr>
<td>Tools</td>
<td>
<a href="notebooks/official/pipelines"><code>Pipelines/</code></a>
</td>
<td>
Use `Vertex AI Pipelines` and `Google Cloud Pipeline Components` to build, tune, or deploy a custom model.
</td>
</tr>
</table>
<!-- markdownlint-enable MD033 -->
Contributions welcome! See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
## Getting help
## Get help
Please use the [Issues page](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues) to provide feedback or submit a bug report.
Please use the [issues page](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues) to provide feedback or submit a bug report.
## Disclaimer
This is not an officially supported Google product. The code in this repository is for demonstrative purposes only.
## Feedback
## References
- [Vertex AI Jupyter Notebook tutorials](https://cloud.google.com/vertex-ai/docs/tutorials/jupyter-notebooks)
- Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository
- [Vertex AI documentaton](https://cloud.google.com/vertex-ai/docs)
Please feel free to fill out our [survey](https://bit.ly/vertex-ai-samples-survey) to give us feedback on the repo and its content.
-23
View File
@@ -8,26 +8,3 @@
/cpr-examples @samthrasher
/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines @Ark-kun
/pipeline_components @Ark-kun
/pipeline_components/image_ml_model_training @lakeyk
/prediction_featurestore_integration @googleapis/vertex-prediction-team
/vertex_model_garden/model_oss/notebook_util @minwoo33park
/vertex_model_garden/model_oss/util @weigary
/vertex_model_garden/model_oss/diffusers @weigary
/vertex_model_garden/model_oss/keras @dstnluong-google
/vertex_model_garden/model_oss/transformers @dstnluong-google
/vertex_model_garden/model_oss/pic2word @jismailyan-google
/vertex_model_garden/model_oss/open_clip @lydhr
/vertex_model_garden/model_oss/movinet @KCFindstr
/vertex_model_garden/model_oss/data_converter @KCFindstr
/vertex_model_garden/model_oss/peft @weigary
/vertex_model_garden/model_oss/peft/templates @rayandasoriya
/vertex_model_garden/model_oss/lm-evaluation-harness @kathyyu-google
/vertex_model_garden/model_oss/tfvision @dstnluong-google
/vertex_model_garden/model_oss/fvlm @minwoo33park
/vertex_model_garden/model_oss/imagebind @kathyyu-google
/vertex_model_garden/model_oss/llava @py4
/vertex_model_garden/model_oss/vllm @kathyyu-google
/vertex_model_garden/benchmarking_reports @lavraicse
/vertex_model_garden/model_oss/autogluon @lavraicse
/vertex_distributed_training/a3mega/llama-3-8b-nemo-pretraining @mstyer-google @erwinh85 @mchrestkha
@@ -6,8 +6,8 @@ download_from_gcs_op = components.load_component_from_url("https://raw.githubuse
select_columns_using_Pandas_on_CSV_data_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/pandas/Select_columns/in_CSV_format/component.yaml")
fill_all_missing_values_using_Pandas_on_CSV_data_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/pandas/Fill_all_missing_values/in_CSV_format/component.yaml")
binarize_column_using_Pandas_on_CSV_data_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/pandas/Binarize_column/in_CSV_format/component.yaml")
train_logistic_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_logistic_regression_model/from_CSV/component.yaml")
upload_Scikit_learn_pickle_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_Scikit-learn_pickle_model/component.yaml")
train_logistic_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_logistic_regression_model/from_CSV/component.yaml")
upload_Scikit_learn_pickle_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_Scikit-learn_pickle_model/component.yaml")
deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
# %% Pipeline definition
@@ -9,7 +9,7 @@ binarize_column_using_Pandas_on_CSV_data_op = components.load_component_from_url
split_rows_into_subsets_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/dataset_manipulation/Split_rows_into_subsets/in_CSV/component.yaml")
train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
# %% Pipeline definition
@@ -23,7 +23,7 @@ upload_PyTorch_model_archive_to_Google_Cloud_Vertex_AI_op = components.load_comp
# XGBoost
train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
# Scikit-learn
#train_linear_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_linear_regression_model/from_CSV/component.yaml")
@@ -8,7 +8,7 @@ fill_all_missing_values_using_Pandas_on_CSV_data_op = components.load_component_
split_rows_into_subsets_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/dataset_manipulation/Split_rows_into_subsets/in_CSV/component.yaml")
train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
# %% Pipeline definition
@@ -22,7 +22,7 @@ upload_PyTorch_model_archive_to_Google_Cloud_Vertex_AI_op = components.load_comp
# XGBoost
train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
# Scikit-learn
train_linear_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_linear_regression_model/from_CSV/component.yaml")
@@ -1,5 +1,5 @@
absl-py==1.1.0
fastapi==0.109.1
fastapi==0.75.2
uvicorn==0.18.2
timm==0.5.4
smart_open==6.0.0
@@ -148,7 +148,7 @@ implementation:
# Downloading the model archive from GCS
# TODO: Fix gsutil bugs (requires project ID, has auth issues) and use gsutil instead.
# gcloud storage cp "$model_archive_uri" "$model_archive_local_path"
# gsutil cp "$model_archive_uri" "$model_archive_local_path"
pip install google-cloud-storage
python -c '
import sys
@@ -64,8 +64,8 @@ implementation:
labels["component-source"] = "github-com-ark-kun-pipeline-components"
# The serving container decides the model type based on the model file extension.
# So we need to rename the mode file (e.g. /tmp/inputs/model/data) to *.bst
_, renamed_model_path = tempfile.mkstemp(suffix=".bst")
# So we need to rename the mode file (e.g. /tmp/inputs/model/data) to *.pkl
_, renamed_model_path = tempfile.mkstemp(suffix=".pkl")
shutil.copyfile(src=model_path, dst=renamed_model_path)
model = aiplatform.Model.upload_xgboost_model_file(
@@ -24,12 +24,12 @@ implementation:
# Checking whether the URI points to a single blob, a directory or a URI pattern
# URI points to a blob when that URI does not end with slash and listing that URI only yields the same URI
if [[ "$uri" != */ ]] && (gcloud storage ls "$uri" | grep --fixed-strings --line-regexp "$uri"); then
if [[ "$uri" != */ ]] && (gsutil ls "$uri" | grep --fixed-strings --line-regexp "$uri"); then
mkdir -p "$(dirname "$output_path")"
gcloud storage cp --recursive "$uri" "$output_path"
gsutil -m cp -r "$uri" "$output_path"
else
mkdir -p "$output_path" # When source path is a directory, gsutil requires the destination to also be a directory
gcloud storage rsync --recursive "$uri" "$output_path" # gsutil cp has different path handling than Linux cp. It always puts the source directory (name) inside the destination directory. gsutil rsync does not have that problem.
gsutil -m rsync -r "$uri" "$output_path" # gsutil cp has different path handling than Linux cp. It always puts the source directory (name) inside the destination directory. gsutil rsync does not have that problem.
fi
- inputValue: GCS path
- outputPath: Data
@@ -1,112 +0,0 @@
name: Load image classification model from tfhub
description: |
Loads specified model from TFHub, creates layer to receive additional (3 channel) imagery data.
Args:
class_names (Sequence[str]):
Sequence of strings of categories for classification corresponding to input data.
loaded_model_path (str):
Output path for the loaded model.
image_size_path (str):
Output path for the model expected image size.
model_name (Optional[str]):
Name of the pre-trained image classification model to load from TFHub.
Eligible model_name:
- efficientnetv2-s
- efficientnetv2-m
- efficientnetv2-l
- efficientnetv2-s-21k
- efficientnetv2-m-21k
- efficientnetv2-l-21k
- efficientnetv2-xl-21k
- efficientnetv2-b0-21k
- efficientnetv2-b1-21k
- efficientnetv2-b2-21k
- efficientnetv2-b3-21k
- efficientnetv2-s-21k-ft1k
- efficientnetv2-m-21k-ft1k
- efficientnetv2-l-21k-ft1k
- efficientnetv2-xl-21k-ft1k
- efficientnetv2-b0-21k-ft1k
- efficientnetv2-b1-21k-ft1k
- efficientnetv2-b2-21k-ft1k
- efficientnetv2-b3-21k-ft1k
- efficientnetv2-b0
- efficientnetv2-b1
- efficientnetv2-b2
- efficientnetv2-b3
- efficientnet_b0
- efficientnet_b1
- efficientnet_b2
- efficientnet_b3
- efficientnet_b4
- efficientnet_b5
- efficientnet_b6
- efficientnet_b7
- bit_s-r50x1
- inception_v3
- inception_resnet_v2
- resnet_v1_50
- resnet_v1_101
- resnet_v1_152
- resnet_v2_50
- resnet_v2_101
- resnet_v2_152
- nasnet_large
- nasnet_mobile
- pnasnet_large
- mobilenet_v2_100_224
- mobilenet_v2_130_224
- mobilenet_v2_140_224
- mobilenet_v3_small_100_224
- mobilenet_v3_small_075_224
- mobilenet_v3_large_100_224
- mobilenet_v3_large_075_224
dropout_rate (Optional[float]):
Fraction of input units to drop in the last layer. Value should be between 0.0 and 1.0.
trainable (Optional[bool]):
If true fine tuning will be performed on entire Hub model. If false only additional
layers will be trained.
l2_regularization_penalty (Optional[float]):
l2 regularization penalty.
inputs:
- {name: class_names, type: 'typing.List[str]', description: List of class names corresponding
to the input image data}
- {name: model_name, type: String, description: Name of the TFHub model to load, default: efficientnetv2-xl-21k,
optional: true}
- {name: dropout_rate, type: Float, description: Dropout rate, default: '0.2', optional: true}
- name: trainable
type: Boolean
description: True if fine tuning should be performed
default: "True"
optional: true
- {name: l2_regularization_penalty, type: Float, description: Regularization penalty,
default: '0.0001', optional: true}
outputs:
- {name: loaded_model_path, type: TensorflowSavedModel, description: Output path for
the loaded model}
- {name: image_size_path, type: HeightWidth}
implementation:
container:
image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
# command is a list of strings (command-line arguments).
# The YAML language has two syntaxes for lists and you can use either of them.
# Here we use the "flow syntax" - comma-separated strings inside square brackets.
command: [
python3,
# Path of the program inside the container
/pipelines/component/src/loading_component.py,
--loaded-model-path,
{outputPath: loaded_model_path},
--class-names,
{inputValue: class_names},
--model-name,
{inputValue: model_name},
--dropout-rate,
{inputValue: dropout_rate},
--trainable,
{inputValue: trainable},
--l2-regularization-penalty,
{inputValue: l2_regularization_penalty},
--image-size-path,
{outputPath: image_size_path},
]
@@ -1,62 +0,0 @@
# python3 -m pip install "kfp<2.0.0" "google-cloud-aiplatform>=1.16.0" --upgrade --quiet
from kfp import components
from kfp.v2 import dsl
# %% Loading components
upload_Tensorflow_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_Tensorflow_model/component.yaml')
deploy_model_to_endpoint_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml')
transcode_imagedataset_tfrecord_from_csv_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/community-content/pipeline_components/image_ml_model_training/transcode_tfrecord_image_dataset_from_csv/component.yaml')
load_image_classification_model_from_tfhub_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/b5b65198a6c2ffe8c0fa2aa70127e3325752df68/community-content/pipeline_components/image_ml_model_training/load_image_classification_model/component.yaml')
preprocess_image_data_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/community-content/pipeline_components/image_ml_model_training/preprocess_image_data/component.yaml')
train_tensorflow_image_classification_model_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/community-content/pipeline_components/image_ml_model_training/train_image_classification_model/component.yaml')
# %% Pipeline definition
def image_classification_pipeline():
class_names = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
csv_image_data_path = 'gs://cloud-samples-data/ai-platform/flowers/flowers.csv'
deploy_model = False
image_data = dsl.importer(
artifact_uri=csv_image_data_path, artifact_class=dsl.Dataset).output
image_tfrecord_data = transcode_imagedataset_tfrecord_from_csv_op(
csv_image_data_path=image_data,
class_names=class_names
).outputs['tfrecord_image_data_path']
loaded_model_outputs = load_image_classification_model_from_tfhub_op(
class_names=class_names,
).outputs
preprocessed_data = preprocess_image_data_op(
image_tfrecord_data,
height_width_path=loaded_model_outputs['image_size_path'],
).outputs
trained_model = (train_tensorflow_image_classification_model_op(
preprocessed_training_data_path = preprocessed_data['preprocessed_training_data_path'],
preprocessed_validation_data_path = preprocessed_data['preprocessed_validation_data_path'],
model_path=loaded_model_outputs['loaded_model_path']).
set_cpu_limit('96').
set_memory_limit('128G').
add_node_selector_constraint('cloud.google.com/gke-accelerator', 'NVIDIA_TESLA_A100').
set_gpu_limit('8').
outputs['trained_model_path'])
vertex_model_name = upload_Tensorflow_model_to_Google_Cloud_Vertex_AI_op(
model=trained_model,
).outputs['model_name']
# Deploying the model might incur additional costs over time
if deploy_model:
vertex_endpoint_name = deploy_model_to_endpoint_op(
model_name=vertex_model_name,
).outputs['endpoint_name']
pipeline_func = image_classification_pipeline
# %% Pipeline submission
if __name__ == '__main__':
from google.cloud import aiplatform
aiplatform.PipelineJob.from_pipeline_func(pipeline_func=pipeline_func).submit()
@@ -1,57 +0,0 @@
name: Preprocess image data
description: |
Preprocess the image data and split between train and validation.
Args:
input_data_path (str):
Input path for the TFRecord image data. Data will be formatted as 'label' (encoded image
label), and 'image_raw' (the binary string of the image data).
height_width_path (str):
Path to square height and width to resize images to. File should contain single float value.
Value is dependent on training model.
preprocessed_training_data_path (str):
Output path for the TFRecord training data. Data will be formatted as 'label' (encoded image
label), and 'image_raw' (the binary string of the image data).
preprocessed_validation_data_path (str):
Output path for the TFRecord validation data. Data will be formatted as 'label' (encoded
image label), and 'image_raw' (the binary string of the image data).
validation_split (Optional[float]):
Fraction of data that will make up validation dataset. Value should be between 0.0 and 1.0.
seed (Optional[int]):
The global random seed to ensure the system gets a unique random sequence
that is deterministic (https://www.tensorflow.org/api_docs/python/tf/random/set_seed).
inputs:
- {name: input_data_path, type: ImageDatasetTFRecord, description: 'Input path for
the TFRecord image data,'}
- {name: height_width_path, type: HeightWidth, description: 'Path to square height and width to
resize images to,'}
- {name: validation_split, type: Float, description: 'Fraction of data that will make
up validation dataset,', default: '0.2', optional: true}
- {name: seed, type: Integer, description: Random seed, default: '0', optional: true}
outputs:
- {name: preprocessed_training_data_path, type: ImageDatasetTFRecord, description: 'Output
path for the training data,'}
- {name: preprocessed_validation_data_path, type: ImageDatasetTFRecord, description: 'Output
path for the validation data,'}
implementation:
container:
image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
# command is a list of strings (command-line arguments).
# The YAML language has two syntaxes for lists and you can use either of them.
# Here we use the "flow syntax" - comma-separated strings inside square brackets.
command: [
python3,
# Path of the program inside the container
/pipelines/component/src/preprocessing_component.py,
--input-data-path,
{inputPath: input_data_path},
--height-width-path,
{inputPath: height_width_path},
--validation-split,
{inputValue: validation_split},
--seed,
{inputValue: seed},
--preprocessed-training-data-path,
{outputPath: preprocessed_training_data_path},
--preprocessed-validation-data-path,
{outputPath: preprocessed_validation_data_path},
]
@@ -1,90 +0,0 @@
name: Train tensorflow image classification model
description: |
Creates a trained image classification TensorFlow model.
Args:
preprocessed_training_data_path (str):
Input path to the TFRecord training data. Data will be formatted as 'label' (encoded image
label), and 'image_raw' (the binary string of the image data).
preprocessed_validation_data_path (str):
Input path to the TFRecord validation data. Data will be formatted as 'label' (encoded
image label), and 'image_raw' (the binary string of the image data).
model_path (str):
Input path to the loaded pre-trained model.
trained_model_path (str):
Output path to save the trained model to.
optimizer_name (Optional[str]):
Name of the tf.keras optimizer. Available optimizers are listed at
https://keras.io/api/optimizers/
optimizer_parameters (Optional[Dict[str, str]]):
Optimizer parameters.
loss_function_name (Optional[str]):
Name of the loss function.
loss_function_parameters (Optional[Dict[str, str]]):
Loss function parameters.
number_of_epochs (Optional[int]):
Number of training iterations over data.
metric_names (Optional[Sequence[str]]):
List of tf.keras.metrics to be evaluated by the model during training and testing. Available
metrics are listed at https://keras.io/api/metrics/.
seed Optional(int):
The global random seed to ensure the system gets a unique random sequence
that is deterministic (https://www.tensorflow.org/api_docs/python/tf/random/set_seed).
inputs:
- {name: preprocessed_training_data_path, type: ImageDatasetTFRecord, description: 'Input
path for the training data,'}
- {name: preprocessed_validation_data_path, type: ImageDatasetTFRecord, description: 'Input
path for the validation data,'}
- {name: model_path, type: TensorflowSavedModel, description: 'Input path for the
model,'}
- {name: optimizer_name, type: String, description: 'Name of the optimizer,', default: SGD,
optional: true}
- {name: optimizer_parameters, type: 'typing.Dict[str, str]', description: 'Optimizer
parameters,', default: '{}', optional: true}
- {name: loss_function_name, type: String, description: 'Name of the loss function,',
default: CategoricalCrossentropy, optional: true}
- {name: loss_function_parameters, type: 'typing.Dict[str, str]', description: 'Loss
function parameters,', default: '{}', optional: true}
- {name: number_of_epochs, type: Integer, description: 'Number of epochs,', default: '10',
optional: true}
- {name: metric_names, type: 'typing.List[str]', description: 'List of metrics to
use,', default: '["accuracy"]', optional: true}
- {name: seed, type: Integer, description: 'Random seed,', default: '0', optional: true}
- {name: batch_size, type: Integer, description: Batch size, default: '16', optional: true}
outputs:
- {name: trained_model_path, type: TensorflowSavedModel, description: 'Output path
for the saved model,'}
implementation:
container:
image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
# command is a list of strings (command-line arguments).
# The YAML language has two syntaxes for lists and you can use either of them.
# Here we use the "flow syntax" - comma-separated strings inside square brackets.
command: [
python3,
# Path of the program inside the container
/pipelines/component/src/training_component.py,
--preprocessed-training-data-path,
{inputPath: preprocessed_training_data_path},
--preprocessed-validation-data-path,
{inputPath: preprocessed_validation_data_path},
--model-path,
{inputPath: model_path},
--trained-model-path,
{outputPath: trained_model_path},
--optimizer-name,
{inputValue: optimizer_name},
--loss-function-name,
{inputValue: loss_function_name},
--number-of-epochs,
{inputValue: number_of_epochs},
--seed,
{inputValue: seed},
--batch-size,
{inputValue: batch_size},
--metric-names,
{inputValue: metric_names},
--optimizer-parameters,
{inputValue: optimizer_parameters},
--loss-function-parameters,
{inputValue: loss_function_parameters},
]
@@ -1,37 +0,0 @@
name: Transcode imagedataset tfrecord from csv
description: |
Transcodes CSV Data into TFRecord file of TFExamples.
Args:
csv_image_data_path (str):
Path to the CSV image data. Data must include 'image_filepath' (Path to image file) and
'image_label' (output for a prediction) fields.
class_names (Sequence[str]):
Sequence of strings of categories for classification corresponding to input data.
tfrecord_image_data_path (str):
Output path for the TFRecord image data. Data will be formatted as 'label' (encoded image
label), and 'image_raw' (the binary string of the image data).
inputs:
- {name: csv_image_data_path, type: ImageDatasetCSV, description: Input path for the
CSV image data}
- {name: class_names, type: 'typing.List[str]', description: List of class names corresponding
to the input image data}
outputs:
- {name: tfrecord_image_data_path, type: ImageDatasetTFRecord, description: Output
path for the TFRecord image data}
implementation:
container:
image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
# command is a list of strings (command-line arguments).
# The YAML language has two syntaxes for lists and you can use either of them.
# Here we use the "flow syntax" - comma-separated strings inside square brackets.
command: [
python3,
# Path of the program inside the container
/pipelines/component/src/transcoding_csv_component.py,
--csv-image-data-path,
{inputPath: csv_image_data_path},
--tfrecord-image-data-path,
{outputPath: tfrecord_image_data_path},
--class-names,
{inputValue: class_names},
]
@@ -1,39 +0,0 @@
name: Transcode imagedataset tfrecord from jsonlines
description: |
Transcodes JSONL Data into TFRecord file of TFExamples.
Args:
jsonl_image_data_path (str):
Input path for the JSONL image data
Path to the JSONL image data. Each line corresponds to a JSON input describing an image.
Schema follows AutoML image classification JSONL format
https://cloud.google.com/vertex-ai/docs/image-data/classification/prepare-data#json-lines.
class_names (Sequence[str]):
Sequence of strings of categories for classification corresponding to input data.
tfrecord_image_data_path (str):
Output path for the TFRecord image data. Data will be formatted as 'label' (encoded image
label), and 'image_raw' (the binary string of the image data).
inputs:
- {name: jsonl_image_data_path, type: ImageDatasetJsonLines, description: Input path
for the JSONL image data}
- {name: class_names, type: 'typing.List[str]', description: List of class names corresponding
to the input image data}
outputs:
- {name: tfrecord_image_data_path, type: ImageDatasetTFRecord, description: Output
path for the TFRecord image data}
implementation:
container:
image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
# command is a list of strings (command-line arguments).
# The YAML language has two syntaxes for lists and you can use either of them.
# Here we use the "flow syntax" - comma-separated strings inside square brackets.
command: [
python3,
# Path of the program inside the container
/pipelines/component/src/transcoding_jsonl_component.py,
--jsonl-image-data-path,
{inputPath: jsonl_image_data_path},
--tfrecord-image-data-path,
{outputPath: tfrecord_image_data_path},
--class-names,
{inputValue: class_names},
]
@@ -1,40 +1,16 @@
# Stage 1: Build Environment
FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime AS builder
# Install necessary tools and dependencies
RUN apt-get update && \
apt-get install -y curl gnupg && \
echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
apt-get update -y && \
apt-get install -y google-cloud-sdk
# Copy application code
COPY . /trainer
# Set working directory
WORKDIR /trainer
# Install Python dependencies
RUN pip install --no-cache-dir -r requirements.txt
# Stage 2: Runtime Environment
FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime
# Install Google Cloud SDK
RUN apt-get update && \
apt-get install -y curl gnupg && \
echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
apt-get update -y && \
apt-get install -y google-cloud-sdk && \
apt-get clean && rm -rf /var/lib/apt/lists/*
apt-get install google-cloud-sdk -y
# Copy from the builder stage
COPY --from=builder /trainer /trainer
COPY . /trainer
# Set working directory
WORKDIR /trainer
# Set the entry point
ENTRYPOINT ["python", "-m", "task"]
RUN pip install -r requirements.txt
ENTRYPOINT ["python", "-m", "task"]
@@ -1,3 +1,3 @@
torch==2.8.0
torch==1.8.1
torchvision==0.9.1
tensorboard==2.5.0
@@ -1,3 +1,3 @@
torch==2.7.0
torch==1.8.1
torchvision==0.9.1
tensorboard==2.5.0
@@ -110,7 +110,7 @@
},
"outputs": [],
"source": [
"! gcloud storage ls $gcs_output_uri_prefix"
"! gsutil ls $gcs_output_uri_prefix"
]
},
{
@@ -192,7 +192,7 @@
},
"outputs": [],
"source": [
"! gcloud storage cp --recursive $gcs_output_uri_prefix/model ./model_server/"
"! gsutil cp -r $gcs_output_uri_prefix/model ./model_server/"
]
},
{
@@ -556,7 +556,7 @@
},
"outputs": [],
"source": [
"! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
"! gsutil rm -rf $gcs_output_uri_prefix"
]
},
{
@@ -412,7 +412,7 @@
},
"outputs": [],
"source": [
"! gcloud storage ls $gcs_output_uri_prefix"
"! gsutil ls $gcs_output_uri_prefix"
]
}
],
@@ -31,7 +31,17 @@
"source": [
"# Deploying a PyTorch Text Classification Model on [Vertex AI](https://cloud.google.com/vertex-ai)\n",
"\n",
"**Kindly reach out to Vertex AI before you run any scale tests or you have any questions.**\n"
"**This is an Experimental release**, covered by the Pre-GA Offerings Terms of your Google Cloud Platform [Terms of Service](https://cloud.google.com/terms).\n",
"\n",
"Experiments are focused on validating a prototype and are not guaranteed to be released. They are not intended for production use or covered by any SLA, support obligation, or deprecation policy and might be subject to backward-incompatible changes.\n",
"\n",
"**Kindly drop us a note before you run any scale tests.**\n",
"\n",
"**Do not hesitate to contact vertexai-prediction-preview-feedback@google.com if you have any questions or run into any issues.**\n",
"\n",
"The usage of the product is free during the Experimental release period: you will still incur charges for other GCP products usage, such as storage.\n",
"\n",
"The projects need to be added to the allowlist in order to deploy PyTorch models using Vertex AI Prediction pre-built PyTorch images. If you are interested in the feature, please send an email to vertexai-prediction-preview-feedback@google.com to provide your project numbers OR project ids."
]
},
{
@@ -77,4 +77,4 @@ echo "After the job is completed successfully, model files will be saved at $JOB
# # Verify the model was exported
# echo "Verify the model was exported:"
# gcloud storage ls ${JOB_DIR}/
# gsutil ls ${JOB_DIR}/
@@ -34,4 +34,4 @@ RUN echo "service_envelope=json\n" "inference_address=http://0.0.0.0:${AIP_H
USER model-server
# run Torchserve HTTP serve to respond to prediction requests
CMD ["echo", "AIP_STORAGE_URI=${AIP_STORAGE_URI}", ";", "gcloud", "storage", "cp", "--recursive", "${AIP_STORAGE_URI}/${MODEL_NAME}.mar", "/home/model-server/model-store/", ";", "ls", "-ltr", "/home/model-server/model-store/", ";", "torchserve", "--start", "--ts-config=/home/model-server/config.properties", "--models", "${MODEL_NAME}=${MODEL_NAME}.mar", "--model-store", "/home/model-server/model-store"]
CMD ["echo", "AIP_STORAGE_URI=${AIP_STORAGE_URI}", ";", "gsutil", "cp", "-r", "${AIP_STORAGE_URI}/${MODEL_NAME}.mar", "/home/model-server/model-store/", ";", "ls", "-ltr", "/home/model-server/model-store/", ";", "torchserve", "--start", "--ts-config=/home/model-server/config.properties", "--models", "${MODEL_NAME}=${MODEL_NAME}.mar", "--model-store", "/home/model-server/model-store"]
@@ -67,4 +67,4 @@ echo "After the job is completed successfully, model files will be saved at $JOB
# # Verify the model was exported
# echo "Verify the model was exported:"
# gcloud storage ls ${JOB_DIR}/
# gsutil ls ${JOB_DIR}/
@@ -478,7 +478,8 @@
},
"outputs": [],
"source": [
"! gcloud storage buckets create --location $REGION $BUCKET_NAME" ]
"! gsutil mb -l $REGION $BUCKET_NAME"
]
},
{
"cell_type": "markdown",
@@ -497,7 +498,8 @@
},
"outputs": [],
"source": [
"! gcloud storage ls --all-versions --long $BUCKET_NAME" ]
"! gsutil ls -al $BUCKET_NAME"
]
},
{
"cell_type": "markdown",
@@ -580,7 +582,8 @@
"outputs": [],
"source": [
"# Download the sample data into your RAW_DATA_PATH\n",
"! gcloud storage cp \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\" $RAW_DATA_PATH" ]
"! gsutil cp \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\" $RAW_DATA_PATH"
]
},
{
"cell_type": "code",
@@ -1618,7 +1621,9 @@
"! gcloud scheduler jobs delete $SIMULATOR_SCHEDULER_JOB --quiet\n",
"\n",
"# Delete Cloud Storage objects that were created.\n",
"! gcloud storage rm --recursive $PIPELINE_ROOT\n", "! gcloud storage rm --recursive $TRAINING_ARTIFACTS_DIR" ]
"! gsutil -m rm -r $PIPELINE_ROOT\n",
"! gsutil -m rm -r $TRAINING_ARTIFACTS_DIR"
]
}
],
"metadata": {
@@ -1,4 +1,4 @@
google-cloud-bigquery==2.20.0
tensorflow==2.12.1
pillow==10.3.0
tensorflow==2.7.2
pillow==9.0.1
tf-agents==0.8.0
@@ -1,4 +1,4 @@
google-cloud-pubsub==2.5.0
pillow==10.3.0
pillow==9.0.1
tf-agents==0.8.0
tensorflow==2.12.1
tensorflow==2.7.2
@@ -1,5 +1,5 @@
dataclasses==0.6
google-cloud-aiplatform==1.8.1
tensorflow==2.12.1
pillow==10.3.0
tensorflow==2.7.2
pillow==9.0.1
tf-agents==0.8.0
@@ -398,7 +398,6 @@
"if not IS_GOOGLE_CLOUD_NOTEBOOK:\n",
" if \"google.colab\" in sys.modules:\n",
" from google.colab import auth as google_auth\n",
"\n",
" google_auth.authenticate_user()\n",
"\n",
" # If you are running this notebook locally, replace the string below with the\n",
@@ -473,7 +472,7 @@
},
"outputs": [],
"source": [
"! gcloud storage buckets create --location $REGION $BUCKET_NAME"
"! gsutil mb -l $REGION $BUCKET_NAME"
]
},
{
@@ -493,7 +492,7 @@
},
"outputs": [],
"source": [
"! gcloud storage ls --all-versions --long $BUCKET_NAME"
"! gsutil ls -al $BUCKET_NAME"
]
},
{
@@ -566,7 +565,7 @@
"outputs": [],
"source": [
"# Copy the sample data into your DATA_PATH\n",
"! gcloud storage cp \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\" $DATA_PATH"
"! gsutil cp \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\" $DATA_PATH"
]
},
{
@@ -580,15 +579,11 @@
"# Set hyperparameters.\n",
"BATCH_SIZE = 8 # @param {type:\"integer\"} Training and prediction batch size.\n",
"TRAINING_LOOPS = 5 # @param {type:\"integer\"} Number of training iterations.\n",
"STEPS_PER_LOOP = (\n",
" 2 # @param {type:\"integer\"} Number of driver steps per training iteration.\n",
")\n",
"STEPS_PER_LOOP = 2 # @param {type:\"integer\"} Number of driver steps per training iteration.\n",
"\n",
"# Set MovieLens simulation environment parameters.\n",
"RANK_K = 20 # @param {type:\"integer\"} Rank for matrix factorization in the MovieLens environment; also the observation dimension.\n",
"NUM_ACTIONS = (\n",
" 20 # @param {type:\"integer\"} Number of actions (movie items) to choose from.\n",
")\n",
"NUM_ACTIONS = 20 # @param {type:\"integer\"} Number of actions (movie items) to choose from.\n",
"PER_ARM = False # Use the non-per-arm version of the MovieLens environment.\n",
"\n",
"# Set agent parameters.\n",
@@ -626,8 +621,7 @@
"source": [
"# Define RL environment.\n",
"env = movielens_py_environment.MovieLensPyEnvironment(\n",
" DATA_PATH, RANK_K, BATCH_SIZE, num_movies=NUM_ACTIONS, csv_delimiter=\"\\t\"\n",
")\n",
" DATA_PATH, RANK_K, BATCH_SIZE, num_movies=NUM_ACTIONS, csv_delimiter=\"\\t\")\n",
"environment = tf_py_environment.TFPyEnvironment(env)\n",
"\n",
"# Define RL agent/algorithm.\n",
@@ -637,8 +631,7 @@
" tikhonov_weight=TIKHONOV_WEIGHT,\n",
" alpha=AGENT_ALPHA,\n",
" dtype=tf.float32,\n",
" accepts_per_arm_features=PER_ARM,\n",
")\n",
" accepts_per_arm_features=PER_ARM)\n",
"print(\"TimeStep Spec (for each batch):\\n\", agent.time_step_spec, \"\\n\")\n",
"print(\"Action Spec (for each batch):\\n\", agent.action_spec, \"\\n\")\n",
"print(\"Reward Spec (for each batch):\\n\", environment.reward_spec(), \"\\n\")\n",
@@ -646,8 +639,7 @@
"# Define RL metric.\n",
"optimal_reward_fn = functools.partial(\n",
" environment_utilities.compute_optimal_reward_with_movielens_environment,\n",
" environment=environment,\n",
")\n",
" environment=environment)\n",
"regret_metric = tf_bandit_metrics.RegretMetric(optimal_reward_fn)\n",
"metrics = [regret_metric]"
]
@@ -712,38 +704,35 @@
" if training_data_spec_transformation_fn is None:\n",
" data_spec = agent.policy.trajectory_spec\n",
" else:\n",
" data_spec = training_data_spec_transformation_fn(agent.policy.trajectory_spec)\n",
" replay_buffer = trainer.get_replay_buffer(\n",
" data_spec, environment.batch_size, steps_per_loop\n",
" )\n",
" data_spec = training_data_spec_transformation_fn(\n",
" agent.policy.trajectory_spec)\n",
" replay_buffer = trainer.get_replay_buffer(data_spec, environment.batch_size,\n",
" steps_per_loop)\n",
"\n",
" # `step_metric` records the number of individual rounds of bandit interaction;\n",
" # that is, (number of trajectories) * batch_size.\n",
" step_metric = tf_metrics.EnvironmentSteps()\n",
" metrics = [\n",
" tf_metrics.NumberOfEpisodes(),\n",
" tf_metrics.AverageEpisodeLengthMetric(batch_size=environment.batch_size),\n",
" tf_metrics.AverageEpisodeLengthMetric(batch_size=environment.batch_size)\n",
" ]\n",
" if additional_metrics:\n",
" metrics += additional_metrics\n",
"\n",
" if isinstance(environment.reward_spec(), dict):\n",
" metrics += [\n",
" tf_metrics.AverageReturnMultiMetric(\n",
" reward_spec=environment.reward_spec(), batch_size=environment.batch_size\n",
" )\n",
" ]\n",
" metrics += [tf_metrics.AverageReturnMultiMetric(\n",
" reward_spec=environment.reward_spec(),\n",
" batch_size=environment.batch_size)]\n",
" else:\n",
" metrics += [tf_metrics.AverageReturnMetric(batch_size=environment.batch_size)]\n",
" metrics += [\n",
" tf_metrics.AverageReturnMetric(batch_size=environment.batch_size)]\n",
"\n",
" # Store intermediate metric results, indexed by metric names.\n",
" metric_results = defaultdict(list)\n",
"\n",
" if training_data_spec_transformation_fn is not None:\n",
"\n",
" def add_batch_fn(data):\n",
" return replay_buffer.add_batch(training_data_spec_transformation_fn(data))\n",
"\n",
" def add_batch_fn(data): return replay_buffer.add_batch(training_data_spec_transformation_fn(data)) \n",
" \n",
" else:\n",
" add_batch_fn = replay_buffer.add_batch\n",
"\n",
@@ -753,12 +742,10 @@
" env=environment,\n",
" policy=agent.collect_policy,\n",
" num_steps=steps_per_loop * environment.batch_size,\n",
" observers=observers,\n",
" )\n",
" observers=observers)\n",
"\n",
" training_loop = trainer.get_training_loop_fn(\n",
" driver, replay_buffer, agent, steps_per_loop\n",
" )\n",
" driver, replay_buffer, agent, steps_per_loop)\n",
" saver = policy_saver.PolicySaver(agent.policy)\n",
"\n",
" for _ in range(training_loops):\n",
@@ -796,8 +783,7 @@
" environment=environment,\n",
" training_loops=TRAINING_LOOPS,\n",
" steps_per_loop=STEPS_PER_LOOP,\n",
" additional_metrics=metrics,\n",
")\n",
" additional_metrics=metrics)\n",
"\n",
"tf.profiler.experimental.stop()"
]
@@ -1106,15 +1092,11 @@
},
"outputs": [],
"source": [
"RUN_HYPERPARAMETER_TUNING = (\n",
" True # Execute hyperparameter tuning instead of regular training.\n",
")\n",
"RUN_HYPERPARAMETER_TUNING = True # Execute hyperparameter tuning instead of regular training.\n",
"TRAIN_WITH_BEST_HYPERPARAMETERS = False # Do not train.\n",
"\n",
"HPTUNING_RESULT_DIR = \"hptuning/\" # @param {type: \"string\"} Directory to store the best hyperparameter(s) in `BUCKET_NAME` and locally (temporarily).\n",
"HPTUNING_RESULT_PATH = os.path.join(\n",
" HPTUNING_RESULT_DIR, \"result.json\"\n",
") # @param {type: \"string\"} Path to the file containing the best hyperparameter(s)."
"HPTUNING_RESULT_PATH = os.path.join(HPTUNING_RESULT_DIR, \"result.json\") # @param {type: \"string\"} Path to the file containing the best hyperparameter(s)."
]
},
{
@@ -1142,7 +1124,7 @@
" image_uri: str,\n",
" args: List[str],\n",
" location: str = \"us-central1\",\n",
" api_endpoint: str = \"us-central1-aiplatform.googleapis.com\",\n",
" api_endpoint: str = \"us-central1-aiplatform.googleapis.com\"\n",
") -> None:\n",
" \"\"\"Creates a hyperparameter tuning job using a custom container.\n",
"\n",
@@ -1215,8 +1197,8 @@
"\n",
" # Create job\n",
" response = client.create_hyperparameter_tuning_job(\n",
" parent=parent, hyperparameter_tuning_job=hyperparameter_tuning_job\n",
" )\n",
" parent=parent,\n",
" hyperparameter_tuning_job=hyperparameter_tuning_job)\n",
" job_id = response.name.split(\"/\")[-1]\n",
" print(\"Job ID:\", job_id)\n",
" print(\"Job config:\", response)\n",
@@ -1260,8 +1242,7 @@
" image_uri=f\"gcr.io/{PROJECT_ID}/{HPTUNING_TRAINING_CONTAINER}:latest\",\n",
" args=args,\n",
" location=REGION,\n",
" api_endpoint=f\"{REGION}-aiplatform.googleapis.com\",\n",
")"
" api_endpoint=f\"{REGION}-aiplatform.googleapis.com\")"
]
},
{
@@ -1311,8 +1292,7 @@
" name = client.hyperparameter_tuning_job_path(\n",
" project=project,\n",
" location=location,\n",
" hyperparameter_tuning_job=hyperparameter_tuning_job_id,\n",
" )\n",
" hyperparameter_tuning_job=hyperparameter_tuning_job_id)\n",
" response = client.get_hyperparameter_tuning_job(name=name)\n",
" return response"
]
@@ -1333,8 +1313,7 @@
" location=REGION,\n",
" api_endpoint=f\"{REGION}-aiplatform.googleapis.com\")\n",
" if response.state.name == 'JOB_STATE_SUCCEEDED':\n",
" print(\"Job succeeded.\n",
"Job Time:\", response.update_time - response.create_time)\n",
" print(\"Job succeeded.\\nJob Time:\", response.update_time - response.create_time)\n",
" trials = response.trials\n",
" print(\"Trials:\", trials)\n",
" break\n",
@@ -1369,8 +1348,8 @@
"if trials:\n",
" # Dict mapping from metric names to the best metric values seen so far\n",
" best_objective_values = dict.fromkeys(\n",
" [metric.metric_id for metric in trials[0].final_measurement.metrics], -np.inf\n",
" )\n",
" [metric.metric_id for metric in trials[0].final_measurement.metrics],\n",
" -np.inf)\n",
" # Dict mapping from metric names to a list of the best combination(s) of\n",
" # hyperparameter(s). Each combination is a dict mapping from hyperparameter\n",
" # names to their values.\n",
@@ -1379,13 +1358,12 @@
" # `final_measurement` and `parameters` are `RepeatedComposite` objects.\n",
" # Reference the structure above to extract the value of your interest.\n",
" for metric in trial.final_measurement.metrics:\n",
" params = {param.parameter_id: param.value for param in trial.parameters}\n",
" params = {\n",
" param.parameter_id: param.value for param in trial.parameters}\n",
" if metric.value > best_objective_values[metric.metric_id]:\n",
" best_params[metric.metric_id] = [params]\n",
" elif metric.value == best_objective_values[metric.metric_id]:\n",
" best_params[param.parameter_id].append(\n",
" params\n",
" ) # Handle cases where multiple hyperparameter values lead to the same performance.\n",
" best_params[param.parameter_id].append(params) # Handle cases where multiple hyperparameter values lead to the same performance.\n",
" print(\"Best hyperparameter value(s):\")\n",
" for metric, params in best_params.items():\n",
" print(f\"Metric={metric}: {sorted(params)}\")\n",
@@ -1465,9 +1443,7 @@
},
"outputs": [],
"source": [
"PREDICTION_CONTAINER = (\n",
" \"prediction-custom-container\" # @param {type:\"string\"} Name of the container image.\n",
")"
"PREDICTION_CONTAINER = \"prediction-custom-container\" # @param {type:\"string\"} Name of the container image."
]
},
{
@@ -1499,7 +1475,7 @@
" machineType: 'E2_HIGHCPU_8'\"\"\".format(\n",
" PROJECT_ID=PROJECT_ID,\n",
" PREDICTION_CONTAINER=PREDICTION_CONTAINER,\n",
" ARTIFACTS_DIR=ARTIFACTS_DIR,\n",
" ARTIFACTS_DIR=ARTIFACTS_DIR\n",
")\n",
"\n",
"with open(\"cloudbuild.yaml\", \"w\") as fp:\n",
@@ -1616,12 +1592,8 @@
},
"outputs": [],
"source": [
"RUN_HYPERPARAMETER_TUNING = (\n",
" False # Execute regular training instead of hyperparameter tuning.\n",
")\n",
"TRAIN_WITH_BEST_HYPERPARAMETERS = (\n",
" True # @param {type:\"bool\"} Whether to use learned hyperparameters in training.\n",
")"
"RUN_HYPERPARAMETER_TUNING = False # Execute regular training instead of hyperparameter tuning.\n",
"TRAIN_WITH_BEST_HYPERPARAMETERS = True # @param {type:\"bool\"} Whether to use learned hyperparameters in training."
]
},
{
@@ -1661,12 +1633,10 @@
"job = aiplatform.CustomContainerTrainingJob(\n",
" display_name=\"train-movielens\",\n",
" container_uri=f\"gcr.io/{PROJECT_ID}/{HPTUNING_TRAINING_CONTAINER}:latest\",\n",
" command=[\"python3\", \"-m\", \"src.training.task\"]\n",
" + args, # Pass in training arguments, including hyperparameters.\n",
" command=[\"python3\", \"-m\", \"src.training.task\"] + args, # Pass in training arguments, including hyperparameters.\n",
" model_serving_container_image_uri=f\"gcr.io/{PROJECT_ID}/{PREDICTION_CONTAINER}:latest\",\n",
" model_serving_container_predict_route=\"/predict\",\n",
" model_serving_container_health_route=\"/health\",\n",
")\n",
" model_serving_container_health_route=\"/health\")\n",
"\n",
"print(\"Training Spec:\", job._managed_model)\n",
"\n",
@@ -1675,8 +1645,7 @@
" replica_count=1,\n",
" machine_type=\"n1-standard-4\",\n",
" accelerator_type=\"ACCELERATOR_TYPE_UNSPECIFIED\",\n",
" accelerator_count=0,\n",
")"
" accelerator_count=0)"
]
},
{
@@ -1815,7 +1784,7 @@
"! gcloud ai models delete $model.name --quiet\n",
"\n",
"# Delete Cloud Storage objects that were created\n",
"! gcloud storage rm --recursive $ARTIFACTS_DIR"
"! gsutil -m rm -r $ARTIFACTS_DIR"
]
}
],
@@ -324,7 +324,7 @@
},
"outputs": [],
"source": [
"! gcloud storage ls $gcs_output_uri_prefix"
"! gsutil ls $gcs_output_uri_prefix"
]
},
{
@@ -344,7 +344,7 @@
},
"outputs": [],
"source": [
"! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
"! gsutil rm -rf $gcs_output_uri_prefix"
]
}
],
@@ -328,7 +328,7 @@
},
"outputs": [],
"source": [
"! gcloud storage ls $gcs_output_uri_prefix"
"! gsutil ls $gcs_output_uri_prefix"
]
},
{
@@ -348,7 +348,7 @@
},
"outputs": [],
"source": [
"! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
"! gsutil rm -rf $gcs_output_uri_prefix"
]
}
],
@@ -1 +1 @@
tensorflow==2.12.1
tensorflow==2.7.2
@@ -341,7 +341,7 @@
},
"outputs": [],
"source": [
"! gcloud storage ls $gcs_output_uri_prefix"
"! gsutil ls $gcs_output_uri_prefix"
]
},
{
@@ -361,7 +361,7 @@
},
"outputs": [],
"source": [
"! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
"! gsutil rm -rf $gcs_output_uri_prefix"
]
}
],
@@ -1,126 +0,0 @@
# Vertex AI Training: Llama 3.1 8B pre-training using Nvidia A3 Mega VMs (H100)
This document provides a step-by-step guide for pre-training a Llama 3.1 8B model on the `en-wiki` dataset using multiple [Vertex AI Custom Training](https://cloud.google.com/vertex-ai/docs/training/overview) `a3-megagpu-8g` nodes.
We will use a custom container based on NVIDIA's [NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/24.07/overview.html) to demonstrate a scalable, multi-node training workflow. All required artifacts and commands are included.
## 1. Prerequisites
### 1.1. Google Cloud Project setup
- **Enable APIs:** Ensure the Vertex AI API is [enabled for your project](http://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).
- **H100 Mega Quota:** A3 Mega VMs are powered by H100 GPUs. Request quota for `custom_model_training_nvidia_h100_mega_gpus` in one of the [supported regions](https://cloud.google.com/vertex-ai/docs/general/locations#accelerator_support). If using Spot VMs, request `custom_model_training_preemptible_nvidia_h100_mega_gpus` quota instead.
- **Reservations (Optional but recommended):** For guaranteed capacity, [create a reservation](https://cloud.google.com/compute/docs/instances/reservations-shared) and ensure the reservation is shared with the Vertex AI service account. This guide requires a minimum of **16 H100 GPUs** (2 full A3 Mega nodes).
### 1.2. GCS bucket
Create a [Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) in the same region where you have quota. If you're using Hierarchical Namespace for your bucket, you may need to update permissions of the Vertex AI Custom Code Service Agent .
This bucket is used for:
- Staging the training application.
- Storing model checkpoints and logs.
- Storing data if you use your own data.
## 2. Setup & configuration
### 2.1. Clone the repo
First clone the repo into your development environment.
```bash
git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
```
Navigate to the root folder for this sample.
### 2.2. Environment Setup
First, configure your local environment. These variables are used in subsequent commands.
```bash
# Required: Update with your values
export PROJECT_ID="<your-project-id>"
export REPOSITORY="<your-artifact-registry-repo-name>" # e.g., "my-containers"
export BUCKET="<your-gcs-bucket-name>"
# Optional: Change if needed
export REGION="us-central1"
# --- Do not change the lines below ---
export ARTIFACT_REGISTRY="${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPOSITORY}"
export REPO_ROOT=$(git rev-parse --show-toplevel)
```
## 3. Build and push a docker container image to Artifact Registry
Normally, you can use any custom training container on Vertex AI Training. In this example you build a NeMo Docker image that is based on the [Nvidia’s NeMo 24.09](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags) image. Use Cloud Build to build and push the container image.
This document picked NeMo as the demonstrating container since it’s a widely adopted GPU LLM training framework providing high performance and versatile training functionalities.
In addition to the base image, some customizations are included to form the final prebuilt image:
- Some dependencies are installed to integrate with Vertex AI Training.
- An entrypoint script that sets up required environments and calls the training job.
- Some patches are applied to the NeMo code to let it load the dataset from a GCS bucket.
Run this command to build the container and push the container into the Google Artifact Registry.
```bash
cd "${REPO_ROOT}/community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining"
export IMAGE_NAME="vertex-nemo-llama"
gcloud builds submit . \
--project="${PROJECT_ID}" \
--region="${REGION}" \
--config=docker/cloudbuild.yml \
--substitutions="_ARTIFACT_REGISTRY=${ARTIFACT_REGISTRY},_IMAGE_NAME=${IMAGE_NAME}" \
--timeout="2h" \
--machine-type="e2-highcpu-32"
```
## 4. Launch the Training Job
### 4.1. Job Configuration File
Once the container is built, update the job_config.json to set up the training job.
File: job_config.json
```json
{
"project_id": "<project-id>",
"region": "<region>",
"zone": "<zone if using reservation>",
"bucket": "<bucket>",
"dataset_bucket": "github-repo/data/third-party/enwiki-latest-pages-articles",
"image_uri": "<docker image uri from artifact registry>",
"strategy": "spot",
"nodes": "2",
"machine_type": "a3-megagpu-8g",
"gpu_type": "NVIDIA_H100_MEGA_80GB",
"gpus_per_node": "8",
"recipe_name": "llama3_1_8b_pretrain_a3mega",
"job_prefix": "vertex-spot-",
"reservation_name": ""
}
```
### 4.2 Launch the Training Job
First, create a Python virtual environment using your tool of choice, then install
the requirements specified in `requirements.txt`. Using `pip`, the command would be:
```bash
pip install -r requirements.txt
```
Now launch the Vertex AI training job using the provided Python script.
```bash
python3 scripts/launch.py --config_file=job_config.json
```
This script reads job_config.json, defines the cluster specification (2 nodes, 8 GPUs each), and submits the custom training job to Vertex AI.
## 5. Monitor and Clean Up
### 5.1. Monitoring
Vertex AI Console: Track the job's status in the Google Cloud Console under Vertex AI > Training > Custom Jobs.
Logs: View detailed logs in Cloud Logging by filtering for your job name.
Checkpoints: Model checkpoints are saved to your GCS bucket at the path specified in your training script's configuration.
### 5.2. Cleaning Up
To avoid ongoing charges, delete the resources you created:
- The Artifact Registry image.
- The contents of the GCS bucket (checkpoints, logs).
- The Vertex AI Custom Job will eventually complete or fail, incurring no further cost.
@@ -1,265 +0,0 @@
# Reference:
# https://github.com/NVIDIA/NeMo-Framework-Launcher/blob/24.07/launcher_scripts/conf/training/llama/llama3_1_8b.yaml
name: llama3_1_8b_pretrain_a3mega
restore_from_path: null # used when starting from a .nemo file
trainer:
devices: 8
num_nodes: 1
accelerator: gpu
precision: bf16
logger: false # logger provided by exp_manager
enable_checkpointing: false
use_distributed_sampler: false
max_epochs: -1 # PTL default. In practice, max_steps will be reached first.
max_steps: 30 # consumed_samples = global_step * micro_batch_size * data_parallel_size * accumulate_grad_batches
log_every_n_steps: 1
val_check_interval: null
limit_val_batches: 1
limit_test_batches: 1
accumulate_grad_batches: 1 # do not modify, grad acc is automatic for training megatron models
gradient_clip_val: 1.0
benchmark: false
enable_model_summary: false # default PTL callback for this does not support model parallelism, instead we log manually
exp_manager:
explicit_log_dir: null
exp_dir: /data
name: ${name}
create_dllogger_logger: true
dllogger_logger_kwargs:
verbose: true
stdout: true
json_file: "/data/dllogger.json"
create_wandb_logger: false
wandb_logger_kwargs:
project: null
name: null
resume_if_exists: true
resume_ignore_no_checkpoint: true
create_checkpoint_callback: false
checkpoint_callback_params:
monitor: val_loss
save_top_k: 3
mode: min
always_save_nemo: false # saves nemo file during validation, not implemented for model parallel
save_nemo_on_train_end: false # not recommended when training large models on clusters with short time limits
filename: 'megatron_gpt--{val_loss:.2f}-{step}-{consumed_samples}'
model_parallel_size: ${multiply:${model.tensor_model_parallel_size}, ${model.pipeline_model_parallel_size}}
seconds_to_sleep: 5 # Allows node_rank!=0 to sleep and let node0 to init, like preparing data
model:
mcore_gpt: true
# specify micro_batch_size, global_batch_size, and model parallelism
# gradient accumulation will be done automatically based on data_parallel_size
micro_batch_size: 1 # limited by GPU memory
global_batch_size: 1024 # will use more micro batches to reach global batch size
tensor_model_parallel_size: 1 # intra-layer model parallelism
pipeline_model_parallel_size: 2 # inter-layer model parallelism
context_parallel_size: 1
virtual_pipeline_model_parallel_size: null # interleaved pipeline
## Sequence Parallelism
# Makes tensor parallelism more memory efficient for LLMs (20B+) by parallelizing layer norms and dropout sequentially
# See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
sequence_parallel: false
fsdp: false
fsdp_cpu_offload: true
fsdp_sharding_strategy: "full" # Method to shard model states. Available options are 'full', 'hybrid', and 'grad'.
fsdp_grad_reduce_dtype: "16" # Gradient reduction data type.
fsdp_sharded_checkpoint: false # Store and load FSDP shared checkpoint.
fsdp_use_orig_params: false # Set to True to use FSDP for specific peft scheme.
# Distributed checkpoint setup
dist_ckpt_format: "torch_dist" # Set to 'torch_dist' to use PyTorch distributed checkpoint format.
dist_ckpt_load_on_device: true # whether to load checkpoint weights directly on GPU or to CPU
dist_ckpt_parallel_save: true # if true, each worker will write its own part of the dist checkpoint
dist_ckpt_parallel_save_within_dp: false # if true, save will be parallelized only within a DP group (whole world otherwise), which might slightly reduce the save overhead
dist_ckpt_parallel_load: false # if true, each worker will load part of the dist checkpoint and exchange with NCCL. Might use some extra GPU memory
dist_ckpt_torch_dist_multiproc: 2 # number of extra processes per rank used during ckpt save with PyTorch distributed format
dist_ckpt_assume_constant_structure: false # set to True only if the state dict structure doesn't change within a single job. Allows caching some computation across checkpoint saves.
dist_ckpt_parallel_dist_opt: true # parallel save/load of a DistributedOptimizer. 'True' allows performant save and reshardable checkpoints. Set to 'False' only in order to minimize the number of checkpoint files.
dist_ckpt_load_strictness: null # defines checkpoint keys mismatch behavior (only during dist-ckpt load). Choices: assume_ok_unexpected (default - try loading without any check), log_all (log mismatches), raise_all (raise mismatches)
# model architecture
encoder_seq_length: 8192
max_position_embeddings: ${.encoder_seq_length}
num_layers: 32 # 8b: 32 | 70b: 80 | 405b: 126
hidden_size: 4096 # 8b: 4096 | 70b: 8192 | 405b: 16384
ffn_hidden_size: 14336 # 8b: 14336 | 70b: 28672 | 405b: 53248
num_attention_heads: 32 # 8b: 32 | 70b: 64 | 405b: 128
num_query_groups: 8 # Number of query groups for group query attention. If None, normal attention is used. 8b: 8 | 70b: 8 | 405b: 16
init_method_std: 0.01 # Standard deviation of the zero mean normal distribution used for weight initialization. 8b: 0.01 | 70b: 0.008944 | 405b: 0.02
use_scaled_init_method: true # use scaled residuals initialization
hidden_dropout: 0.0 # Dropout probability for hidden state transformer.
attention_dropout: 0.0 # Dropout probability for attention
ffn_dropout: 0.0 # Dropout probability in the feed-forward layer.
kv_channels: null # Projection weights dimension in multi-head attention. Set to hidden_size // num_attention_heads if null
apply_query_key_layer_scaling: true # scale Q * K^T by 1 / layer-number.
normalization: 'rmsnorm' # Normalization layer to use. Options are 'layernorm', 'rmsnorm'
layernorm_epsilon: 1e-5
do_layer_norm_weight_decay: false # True means weight decay on all params
make_vocab_size_divisible_by: 128 # Pad the vocab size to be divisible by this value for computation efficiency.
pre_process: true # add embedding
post_process: true # add pooler
persist_layer_norm: true # Use of persistent fused layer norm kernel.
bias: false # Whether to use bias terms in all weight matrices.
activation: 'fast-swiglu' # Options ['gelu', 'geglu', 'swiglu', 'reglu', 'squared-relu', 'fast-geglu', 'fast-swiglu', 'fast-reglu']
headscale: false # Whether to learn extra parameters that scale the output of the each self-attention head.
transformer_block_type: 'pre_ln' # Options ['pre_ln', 'post_ln', 'normformer']
openai_gelu: false # Use OpenAI's GELU instead of the default GeLU
normalize_attention_scores: true # Whether to scale the output Q * K^T by 1 / sqrt(hidden_size_per_head). This arg is provided as a configuration option mostly for compatibility with models that have been weight-converted from HF. You almost always want to se this to True.
position_embedding_type: 'rope' # Position embedding type. Options ['learned_absolute', 'rope']
rotary_percentage: 1.0 # If using position_embedding_type=rope, then the per head dim is multiplied by this.
attention_type: 'multihead' # Attention type. Options ['multihead']
share_embeddings_and_output_weights: false # Share embedding and output layer weights.
scale_positional_embedding: true # This is false for llama3 models. Only used for >= llama3.1.
# Use GPT2BPETokenizer for test, because the testing dataset is tokenized by this tokenizer.
# https://docs.nvidia.com/nemo-framework/user-guide/24.07/playbooks/singlenodepretrain.html#data-download-and-pre-processing
tokenizer:
library: megatron
type: GPT2BPETokenizer
model: null # /path/to/tokenizer.model
vocab_file: null
merge_file: null
delimiter: null # only used for tabular tokenizer
sentencepiece_legacy: false # Legacy=True allows you to add special tokens to sentencepiece tokenizers.
# Mixed precision
native_amp_init_scale: 4294967296 # 2 ** 32
native_amp_growth_interval: 1000
hysteresis: 2 # Gradient scale hysteresis
fp32_residual_connection: false # Move residual connections to fp32
fp16_lm_cross_entropy: false # Move the cross entropy unreduced loss calculation for lm head to fp16
# Megatron O2-style half-precision
megatron_amp_O2: true # Enable O2-level automatic mixed precision using main parameters
grad_allreduce_chunk_size_mb: 125
# Fusion
grad_div_ar_fusion: true # Fuse grad division into torch.distributed.all_reduce. Only used with O2 and no pipeline parallelism..
gradient_accumulation_fusion: true # Fuse weight gradient accumulation to GEMMs. Only used with pipeline parallelism and O2.
bias_activation_fusion: true # Use a kernel that fuses the bias addition from weight matrices with the subsequent activation function.
bias_dropout_add_fusion: true # Use a kernel that fuses the bias addition, dropout and residual connection addition.
masked_softmax_fusion: true # Use a kernel that fuses the attention softmax with it's mask.
apply_rope_fusion: true # Use a kernel to add rotary positional embeddings. Only used if position_embedding_type=rope
cross_entropy_loss_fusion: true
# Miscellaneous
seed: 1234
resume_from_checkpoint: null # manually set the checkpoint file to load from
use_cpu_initialization: false # Init weights on the CPU (slow for large models)
onnx_safe: false # Use work-arounds for known problems with Torch ONNX exporter.
apex_transformer_log_level: 30 # Python logging level displays logs with severity greater than or equal to this
gradient_as_bucket_view: true # PyTorch DDP argument. Allocate gradients in a contiguous bucket to save memory (less fragmentation and buffer memory)
sync_batch_comm: false # Enable stream synchronization after each p2p communication between pipeline stages
## Activation Checkpointing
# NeMo Megatron supports 'selective' activation checkpointing where only the memory intensive part of attention is checkpointed.
# These memory intensive activations are also less compute intensive which makes activation checkpointing more efficient for LLMs (20B+).
# See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
# 'full' will checkpoint the entire transformer layer.
activations_checkpoint_granularity: null # 'selective' or 'full'
activations_checkpoint_method: null # 'uniform', 'block'
# 'uniform' divides the total number of transformer layers and checkpoints the input activation
# of each chunk at the specified granularity. When used with 'selective', 'uniform' checkpoints all attention blocks in the model.
# 'block' checkpoints the specified number of layers per pipeline stage at the specified granularity
activations_checkpoint_num_layers: null
# when using 'uniform' this creates groups of transformer layers to checkpoint. Usually set to 1. Increase to save more memory.
# when using 'block' this this will checkpoint the first activations_checkpoint_num_layers per pipeline stage.
num_micro_batches_with_partial_activation_checkpoints: null
# This feature is valid only when used with pipeline-model-parallelism.
# When an integer value is provided, it sets the number of micro-batches where only a partial number of Transformer layers get checkpointed
# and recomputed within a window of micro-batches. The rest of micro-batches in the window checkpoint all Transformer layers. The size of window is
# set by the maximum outstanding micro-batch backpropagations, which varies at different pipeline stages. The number of partial layers to checkpoint
# per micro-batch is set by 'activations_checkpoint_num_layers' with 'activations_checkpoint_method' of 'block'.
# This feature enables using activation checkpoint at a fraction of micro-batches up to the point of full GPU memory usage.
activations_checkpoint_layers_per_pipeline: null
# This feature is valid only when used with pipeline-model-parallelism.
# When an integer value (rounded down when float is given) is provided, it sets the number of Transformer layers to skip checkpointing at later
# pipeline stages. For example, 'activations_checkpoint_layers_per_pipeline' of 3 makes pipeline stage 1 to checkpoint 3 layers less than
# stage 0 and stage 2 to checkpoint 6 layers less stage 0, and so on. This is possible because later pipeline stage
# uses less GPU memory with fewer outstanding micro-batch backpropagations. Used with 'num_micro_batches_with_partial_activation_checkpoints',
# this feature removes most of activation checkpoints at the last pipeline stage, which is the critical execution path.
## Transformer Engine
transformer_engine: true
fp8: false # enables fp8 in TransformerLayer forward
fp8_e4m3: false # sets fp8_format = recipe.Format.E4M3
fp8_hybrid: false # sets fp8_format = recipe.Format.HYBRID
fp8_margin: 0 # scaling margin
fp8_interval: 1 # scaling update interval
fp8_amax_history_len: 1024 # Number of steps for which amax history is recorded per tensor
fp8_amax_compute_algo: 'max' # 'most_recent' or 'max'. Algorithm for computing amax from history
ub_tp_comm_overlap: false # do not turn on because of b/397797926
use_flash_attention: true
gc_interval: 100
## Offloading Activations/Weights to CPU
cpu_offloading: false
cpu_offloading_num_layers: ${sum:${.num_layers},-1} # This value should be between [1,num_layers-1] as we don't want to offload the final layer's activations and expose any offloading duration for the final layer
cpu_offloading_activations: true
cpu_offloading_weights: true
data:
# Path to data must be specified by the user.
# Supports List, String and Dictionary
# List : can override from the CLI: "model.data.data_prefix=[.5,/raid/data/pile/my-gpt3_00_text_document,.5,/raid/data/pile/my-gpt3_01_text_document]",
# Or see example below:
# data_prefix:
# - .5
# - /raid/data/pile/my-gpt3_00_text_document
# - .5
# - /raid/data/pile/my-gpt3_01_text_document
# Dictionary: can override from CLI "model.data.data_prefix"={"train":[1.0, /path/to/data], "validation":/path/to/data, "test":/path/to/test}
# Or see example below:
# "model.data.data_prefix: {train:[1.0,/path/to/data], validation:[/path/to/data], test:[/path/to/test]}"
data_prefix: [1.0, /data/hfbpe_gpt_training_data_text_document]
index_mapping_dir: null # path to save index mapping .npy files, by default will save in the same location as data_prefix
data_impl: mmap
splits_string: 900,50,50
seq_length: ${model.encoder_seq_length}
skip_warmup: true
num_workers: 2
dataloader_type: single # cyclic
reset_position_ids: false # Reset position ids after end-of-document token
reset_attention_mask: false # Reset attention mask after end-of-document token
eod_mask_loss: false # Mask loss for the end of document tokens
validation_drop_last: true # Set to false if the last partial validation samples is to be consumed
no_seqlen_plus_one_input_tokens: false # Set to True to disable fetching (sequence length + 1) input tokens, instead get (sequence length) input tokens and mask the last token
pad_samples_to_global_batch_size: false # Set to True if you want to pad the last partial batch with -1's to equal global batch size
shuffle_documents: true # Set to False to disable documents shuffling. Sample index will still be shuffled
# Nsys profiling options
nsys_profile:
enabled: false
start_step: 0 # Global batch to start profiling
end_step: 1 # Global batch to end profiling
ranks: [0] # Global rank IDs to profile
gen_shape: false # Generate model and kernel details including input shapes
memory_profile:
enabled: false
start_step: 0
end_step: 1
ranks: [0]
output_path: /data # Must be a dir
optim:
name: distributed_fused_adam # E.g., fused_adam or set _target_: torch.optim.AdamW field
lr: 2e-5
weight_decay: 0.01
betas:
- 0.9
- 0.98
bucket_cap_mb: 125
overlap_grad_sync: true
overlap_param_sync: true
contiguous_grad_buffer: true
contiguous_param_buffer: true
sched:
name: CosineAnnealing
warmup_steps: 400
constant_steps: 0
min_lr: 2e-6
@@ -1,26 +0,0 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
steps:
- name: 'gcr.io/cloud-builders/docker'
args:
- 'build'
- '--tag=${_ARTIFACT_REGISTRY}/${_IMAGE_NAME}'
- '--file=docker/vertex-dist-recipes.Dockerfile'
- '.'
automapSubstitutions: true
env:
- 'DOCKER_BUILDKIT=1'
images:
- '${_ARTIFACT_REGISTRY}/${_IMAGE_NAME}'
@@ -1,41 +0,0 @@
diff --git a/nemo/collections/nlp/parts/megatron_trainer_builder.py b/nemo/collections/nlp/parts/megatron_trainer_builder.py
index b2c85cde4..a3a9670c3 100644
--- a/nemo/collections/nlp/parts/megatron_trainer_builder.py
+++ b/nemo/collections/nlp/parts/megatron_trainer_builder.py
@@ -19,6 +19,7 @@ from lightning_fabric.utilities.exceptions import MisconfigurationException
from omegaconf import DictConfig
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelSummary
+from pytorch_lightning.callbacks import Callback
from pytorch_lightning.plugins.environments import TorchElasticEnvironment
from nemo.collections.common.metrics.perf_metrics import FLOPsMeasurementCallback
@@ -38,6 +39,23 @@ from nemo.utils.callbacks.dist_ckpt_io import (
AsyncFinalizerCallback,
DistributedCheckpointIO,
)
+from vmg.util.device_stats import gpu_stats_str
+
+class GpuStatsMon(Callback):
+ def on_train_start(self, trainer, pl_module) -> None:
+ rank=pl_module.global_rank
+ print(f'train_start: {rank=} {gpu_stats_str()}', flush=True)
+
+ def on_train_batch_start(self, trainer, pl_module, batch, batch_idx) -> None:
+ rank=pl_module.global_rank
+ print(f'batch_start: {rank=} {gpu_stats_str()}', flush=True)
+
+ def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx) -> None:
+ rank=pl_module.global_rank
+ print(f'batch_end: {rank=} {gpu_stats_str()}', flush=True)
class MegatronTrainerBuilder:
@@ -178,6 +196,7 @@ class MegatronTrainerBuilder:
if self.cfg.get('exp_manager', {}).get('log_tflops_per_sec_per_gpu', True):
callbacks.append(FLOPsMeasurementCallback(self.cfg))
+ callbacks.append(GpuStatsMon())
return callbacks
def create_trainer(self, callbacks=None) -> Trainer:
@@ -1,41 +0,0 @@
diff -ruN old-datasets/blended_megatron_dataset_builder.py datasets/blended_megatron_dataset_builder.py
--- old-datasets/blended_megatron_dataset_builder.py 2025-05-02 04:08:45.369199665 +0000
+++ datasets/blended_megatron_dataset_builder.py 2025-05-02 04:10:47.369119891 +0000
@@ -2,6 +2,7 @@
import logging
import math
+import os
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Callable, Iterable, List, Optional, Type, Union
@@ -353,7 +354,7 @@
num_dataset_builder_threads = self.config.num_dataset_builder_threads
if torch.distributed.is_initialized():
- rank = torch.distributed.get_rank()
+ rank = int(os.getenv("LOCAL_RANK", "0"))
# First, build on rank 0
if rank == 0:
num_workers = num_dataset_builder_threads
@@ -475,7 +476,7 @@
Optional[Union[DistributedDataset, Iterable]]: The DistributedDataset instantion, the Iterable instantiation, or None
"""
if torch.distributed.is_initialized():
- rank = torch.distributed.get_rank()
+ rank = int(os.getenv("LOCAL_RANK", "0"))
dataset = None
diff -ruN old-datasets/gpt_dataset.py datasets/gpt_dataset.py
--- old-datasets/gpt_dataset.py 2025-05-02 04:08:45.369199665 +0000
+++ datasets/gpt_dataset.py 2025-05-02 04:09:30.309170278 +0000
@@ -351,7 +351,7 @@
if not path_to_cache or (
not cache_hit
- and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0)
+ and (not torch.distributed.is_initialized() or int(os.getenv("LOCAL_RANK", "0")) == 0)
):
log_single_rank(
@@ -1,13 +0,0 @@
diff --git a/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py b/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py
index 8da15148d..005cae6c9 100644
--- a/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py
+++ b/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py
@@ -104,6 +104,8 @@ def convert(input_nemo_file, output_hf_file, precision=None, cpu_only=False) ->
dummy_trainer = Trainer(devices=1, accelerator='cpu', strategy=NLPDDPStrategy())
model_config = MegatronGPTModel.restore_from(input_nemo_file, trainer=dummy_trainer, return_config=True)
model_config.tensor_model_parallel_size = 1
+ model_config.virtual_pipeline_model_parallel_size = None
+ model_config.sequence_parallel = False
model_config.pipeline_model_parallel_size = 1
if cpu_only:
map_location = torch.device('cpu')
@@ -1,24 +0,0 @@
diff --git a/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py b/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py
index bfe8ea359..dfeaf93b5 100644
--- a/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py
+++ b/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py
@@ -13,6 +13,8 @@
# limitations under the License.
import torch.multiprocessing as mp
+import torch.distributed as dist
+
from omegaconf.omegaconf import OmegaConf
from nemo.collections.nlp.models.language_modeling.megatron_gpt_sft_model import MegatronGPTSFTModel
@@ -76,6 +78,10 @@ def main(cfg) -> None:
trainer.fit(model)
+ if dist.is_available() and dist.is_initialized():
+ dist.barrier()
+ dist.destroy_process_group()
+
if __name__ == '__main__':
main()
@@ -1,13 +0,0 @@
diff --git a/src/utils/training_metrics/process_training_results.py b/src/utils/training_metrics/process_training_results.py
index 3e82a66..e61e1d8 100644
--- a/src/utils/training_metrics/process_training_results.py
+++ b/src/utils/training_metrics/process_training_results.py
@@ -134,7 +134,7 @@ def get_average_step_time(file: str, start_step: int, end_step: int) -> float:
for line in datajson:
if line.get("step") != "PARAMETER":
step = line.get("step")
- if step >= start_step and step <= end_step:
+ if step >= start_step and step <= end_step and "train_step_timing in s" in line["data"]:
time_step_accumulator += line["data"].get("train_step_timing in s")
num_steps += 1
if num_steps == 0:
@@ -1,10 +0,0 @@
dllogger@git+https://github.com/NVIDIA/dllogger@v1.0.0
# Fixing these libraries versions to avoid conflicting or broken packages.
immutabledict==4.2.1
protobuf==5.29.6
opencv-python-headless==4.11.0.86
docutils==0.16
urllib3==2.6.3
google-cloud-storage==3.0.0
retrying
@@ -1,18 +0,0 @@
# cuml-cu12==24.8.0 was installed in nemo:24.09
# Removing cuml=24.4.0 to avoid conflicting packages.
cudf==24.4.0
cugraph==24.4.0
cugraph-service-server==24.4.0
cuml==24.4.0
dask-cudf==24.4.0
raft-dask==24.4.0
cugraph-dgl==24.4.0
cugraph-pyg==24.4.0
# The following packages are removed temporarily to avoid conflicting packages
# and can be brought back if needed.
tensorrt-llm==0.12.0
img2dataset==1.45.0
Sphinx==8.1.3
sphinxcontrib-bibtex==2.6.3
torchx==0.7.0
nemo-run
@@ -1,66 +0,0 @@
# Dockerfile wrapping NeMo.
#
# To workaround base nemo docker image using too many layers, we use Multi-stage
# build to first collect the additional files we'll need.
FROM alpine:latest AS prep_files
WORKDIR /workspace
RUN mkdir -p configs vdt vdt/util
COPY scripts/*.py vdt/
COPY scripts/util/*.py vdt/util/
COPY configs/* configs/
COPY docker/patches/24.09/* vdt/patches/
RUN chmod a+rwX -R vdt
# Copy license.
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
# Available tags
# https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags
# It installs NeMo source code in /opt/NeMo folder, with tag=r2.0.0
FROM nvcr.io/nvidia/nemo:24.09
RUN apt-get update && apt-get install -y sudo zsh tmux && \
rm -rf /var/lib/apt/lists*
RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | \
tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | \
apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
apt-get update -y && apt-get install google-cloud-sdk -y && \
rm -rf /var/lib/apt/lists*
# Install libraries with pip
ENV PIP_ROOT_USER_ACTION=ignore
# We expect this will be run in the root directory of the vertex-dist-recipes repo
ARG HOST_SRC_DIR="."
# The pre-installed NeMo introduces a lot of deps conflicts.
# We uninstall the confilicting libs and reinstall some of them as needed.
COPY ${HOST_SRC_DIR}/docker/uninstall.txt /tmp/uninstall.txt
RUN cat /tmp/uninstall.txt | grep -v '#' | xargs pip uninstall -y
COPY ${HOST_SRC_DIR}/docker/requirements.txt /tmp/requirements.txt
RUN pip install -r /tmp/requirements.txt
# Make sure there's no inconsistent pip libraries.
RUN pip check
WORKDIR /workspace
# Copy configs
COPY ${HOST_SRC_DIR}/configs/* /opt/NeMo/examples/nlp/language_modeling/conf/
# Copy all additional files we need from `prep_files` image.
COPY --from=prep_files /workspace/ .
# Install for `src/utils/training_metrics/process_training_results.py` to report
# throughput and MFU numbers.
RUN git clone https://github.com/AI-Hypercomputer/gpu-recipes.git
# This hack is needed for multi-node training while not using a sharing file system.
RUN patch --verbose -l -d /opt/megatron-lm/megatron/core/datasets -p1 -i /workspace/vdt/patches/local_rank.patch; \
git -C /workspace/gpu-recipes apply /workspace/vdt/patches/throughput_calc.patch; \
git -C /opt/NeMo apply /workspace/vdt/patches/nemo2hf.patch; \
git -C /opt/NeMo apply /workspace/vdt/patches/sigabort.patch;
# git -C /opt/NeMo apply /workspace/vdt/patches/gpu_stats.patch;
# Do not put an entrypoint here. Specify the entrypoint in the docker run script.
@@ -1,16 +0,0 @@
{
"project_id": "<your_project_id>",
"region": "us-central1",
"zone": "us-central1-c",
"bucket": "<your_bucket",
"dataset_bucket": "github-repo/data/third-party/enwiki-latest-pages-articles",
"image_uri": "<your_image_uri>",
"strategy": "spot",
"nodes": "2",
"machine_type": "a3-megagpu-8g",
"gpu_type": "NVIDIA_H100_MEGA_80GB",
"gpus_per_node": "8",
"recipe_name": "llama3_1_8b_pretrain_a3mega",
"job_prefix": "vertex-ai",
"reservation_name": ""
}
@@ -1,49 +0,0 @@
absl-py==2.2.2
annotated-types==0.7.0
anyio==4.9.0
black==25.1.0
cachetools==5.5.2
certifi==2025.4.26
charset-normalizer==3.4.2
click==8.1.8
docstring_parser==0.16
google-api-core==2.24.2
google-auth==2.40.1
google-cloud-aiplatform==1.92.0
google-cloud-bigquery==3.31.0
google-cloud-core==2.4.3
google-cloud-resource-manager==1.14.2
google-cloud-storage==2.19.0
google-crc32c==1.7.1
google-genai==1.14.0
google-resumable-media==2.7.2
googleapis-common-protos==1.70.0
grpc-google-iam-v1==0.14.2
grpcio==1.71.0
grpcio-status==1.71.0
h11==0.16.0
httpcore==1.0.9
httpx==0.28.1
idna==3.10
mypy_extensions==1.1.0
numpy==2.2.5
packaging==25.0
pathspec==0.12.1
platformdirs==4.3.8
proto-plus==1.26.1
protobuf==5.29.6
pyasn1==0.6.2
pyasn1_modules==0.4.2
pydantic==2.11.4
pydantic_core==2.33.2
python-dateutil==2.9.0.post0
pytz==2025.2
requests==2.32.4
rsa==4.9.1
shapely==2.1.0
six==1.17.0
sniffio==1.3.1
typing-inspection==0.4.0
typing_extensions==4.13.2
urllib3==2.6.3
websockets==15.0.1
@@ -1,173 +0,0 @@
"""Launch script for Vertex distributed training"""
# Copy the sample_job_config.json file to job_config.json
# to define the job parameters.
#
# Run like this:
#
# python3 vertex_dist_train/launch.py --config_file=job_config.json
#
import datetime
import json
import os
import pprint
from collections.abc import Sequence
from typing import Any, List
from absl import app, flags
from google.cloud import aiplatform
from google.cloud.aiplatform_v1.types.custom_job import Scheduling
from pytz import timezone
FLAGS = flags.FLAGS
flags.DEFINE_string("config_file", None, "Path to JSON config file")
flags.DEFINE_boolean(
"debug", False, "Debug mode: just print the command, don't run it."
)
def launch_job(
job_name: str,
project: str,
region: str,
gcs_bucket: str,
image_uri: str,
entrypoint_cmd: List[str],
trainer_args: List[Any],
num_nodes: int,
machine_type: str,
num_gpus_per_node: int,
gpu_type: str,
strategy: str,
reservation_name: str = "",
):
assert strategy in ("dws", "spot", "reservation")
aiplatform.init(
project=project, location=region, staging_bucket=gcs_bucket
)
train_job = aiplatform.CustomContainerTrainingJob(
display_name=job_name,
container_uri=image_uri,
command=entrypoint_cmd,
)
job_args = dict(
args=trainer_args,
enable_web_access=True,
replica_count=num_nodes,
machine_type=machine_type,
accelerator_type=gpu_type,
accelerator_count=num_gpus_per_node,
boot_disk_size_gb=1000,
restart_job_on_worker_restart=True,
#restart_job_on_worker_restart=False,
)
if strategy == "spot":
job_args.update({"scheduling_strategy": Scheduling.Strategy.SPOT.name})
elif strategy == "dws":
job_args.update(
{"scheduling_strategy": Scheduling.Strategy.FLEX_START.name}
)
elif strategy == "reservation":
assert reservation_name != "", (
"If using a reservation, provide the reservation_name in the "
"format `projects/{project_id_or_number}/zones/{zone}/"
"reservations/{reservation_name}`"
)
job_args.update(
{
"reservation_affinity_type": "SPECIFIC_RESERVATION",
"reservation_affinity_key": "compute.googleapis.com/reservation-name",
"reservation_affinity_values": [reservation_name],
}
)
pprint.pprint(job_args)
if not FLAGS.debug:
train_job.submit(**job_args)
def main(argv: Sequence[str]) -> None:
config_file_path = FLAGS.config_file
print(f"Reading job config from {config_file_path}")
with open(config_file_path, encoding="utf-8") as config_file:
config = json.load(config_file)
project_id = config["project_id"]
region = config["region"]
zone = config["zone"]
bucket = config["bucket"]
dataset_bucket = config["dataset_bucket"]
n_nodes = int(config["nodes"])
machine_type = config["machine_type"]
num_gpus_per_node = int(config["gpus_per_node"])
gpu_type = config["gpu_type"]
reservation_name = config.get("reservation_name")
reservation_full_name = (
f"projects/{project_id}/zones/{zone}/reservations/{reservation_name}"
if "reservation_name" in config
else ""
)
strategy = config["strategy"]
recipe_name = config["recipe_name"]
job_prefix = config["job_prefix"]
image_uri = config["image_uri"]
# Job name
timestamp = (
datetime.datetime.now()
.astimezone(timezone("US/Pacific"))
.strftime("%Y%m%d_%H%M%S")
)
job_name = f"{recipe_name}-{timestamp}"
if job_prefix:
job_name = f"{job_prefix}-{job_name}"
base_output_dir = os.path.join("/gcs", bucket, job_name)
# Training command and args
entrypoint_cmd = ["python3", "vdt/run.py"]
dataset_bucket = f"gs://{config['dataset_bucket']}"
trainer_args = [
f"--train_data_gcs={dataset_bucket}",
"/opt/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py",
"--config-path=conf/",
f"--config-name={recipe_name}.yaml",
f"exp_manager.explicit_log_dir={base_output_dir}",
f"exp_manager.dllogger_logger_kwargs.json_file={base_output_dir}/dllogger.json",
"+exp_manager.create_tensorboard_logger=true",
"exp_manager.create_checkpoint_callback=false",
f"trainer.num_nodes={n_nodes}",
f"trainer.devices={num_gpus_per_node}",
"trainer.max_steps=10",
"trainer.log_every_n_steps=1",
"model.tokenizer.vocab_file=/data/gpt2-vocab.json",
"model.tokenizer.merge_file=/data/gpt2-merges.txt",
"model.data.data_prefix=[1.0,/data/hfbpe_gpt_training_data_text_document]",
]
launch_job(
job_name=job_name,
project=project_id,
region=region,
gcs_bucket=bucket,
image_uri=image_uri,
entrypoint_cmd=entrypoint_cmd,
trainer_args=trainer_args,
num_nodes=n_nodes,
machine_type=machine_type,
num_gpus_per_node=num_gpus_per_node,
gpu_type=gpu_type,
strategy=strategy,
reservation_name=reservation_full_name,
)
if __name__ == "__main__":
app.run(main)
@@ -1,85 +0,0 @@
"""Entrypoint for Vertex Distributed Training container."""
import argparse
import os
import sys
from collections.abc import Sequence
from subprocess import STDOUT, check_output, run
from absl import app, flags, logging
from util import cluster_spec
from retrying import retry
# PyTorch barrier call which synchronizes all of the nodes before launching the training process.
# This makes sure that processes will block until all processes are ready.
# Improves the reliability of spot VM usage for multi-node training jobs
@retry(stop_max_attempt_number=100, wait_exponential_multiplier=1000)
def barrier_with_retry() -> None:
import torch
logging.info("Starting barrier on RANK {}".format(os.environ["RANK"]))
torch.distributed.init_process_group()
torch.distributed.barrier()
torch.distributed.destroy_process_group()
logging.info("Finished barrier on RANK {}".format(os.environ["RANK"]))
def main(unused_argv: Sequence[str]) -> None:
parser = argparse.ArgumentParser()
parser.add_argument(
"--train_data_gcs",
type=str,
help="Download training data from gcs path",
)
args, unknown = parser.parse_known_args()
for key, val in os.environ.items():
logging.info("ENV %s=%s", key, val)
if args.train_data_gcs:
local_dir = "/data"
if not os.path.exists(local_dir):
os.mkdir(local_dir)
logging.info("downloading %s to %s...", args.train_data_gcs, local_dir)
check_output(
[
"gcloud",
"storage",
"cp",
"-r",
f"{args.train_data_gcs}/*",
local_dir,
],
stderr=STDOUT,
)
logging.info("%s downloaded.", args.train_data_gcs)
primary_node_addr, primary_node_port, node_rank, num_nodes = (
cluster_spec.get_cluster_spec()
)
cmd = [
"torchrun",
"--nproc-per-node=8",
f"--nnodes={num_nodes}",
f"--node_rank={node_rank}",
]
if num_nodes > 1:
cmd += [
"--max-restarts=3",
"--rdzv-backend=static",
f'--rdzv_id={os.getenv("CLOUD_ML_JOB_ID", primary_node_port)}',
f"--rdzv-endpoint={primary_node_addr}:{primary_node_port}",
]
cmd += unknown
logging.info("launching with cmd: \n%s", " \\\n".join(cmd))
barrier_with_retry()
run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
if __name__ == "__main__":
logging.get_absl_handler().python_handler.stream = sys.stdout
app.run(
main, flags_parser=lambda _args: flags.FLAGS(_args, known_only=True)
)
@@ -1,81 +0,0 @@
"""Get cluster info from environment variables."""
import dataclasses
import json
import os
from absl import logging
@dataclasses.dataclass
class ClusterInfo:
"""Contains information about the cluster.
Attributes:
primary_node_addr: The address of the primary node.
primary_node_port: The port of the primary node.
node_rank: The rank of the node.
num_nodes: The number of nodes in the cluster.
"""
primary_node_addr: str | None = None
primary_node_port: str | None = None
node_rank: int = 0
num_nodes: int = 1
# Allows unpacking operation like
# primary_node_addr, primary_node_port, _, _ = ClusterInfo()
# See https://stackoverflow.com/a/70753113
def __iter__(self):
return iter(dataclasses.astuple(self))
def get_cluster_spec() -> ClusterInfo:
"""Parses CLUSTER_SPEC environment variable and returns the cluster info.
Returns:
A ClusterInfo object.
"""
cluster_spec = os.getenv("CLUSTER_SPEC", None)
# If CLUSTER_SPEC is not set, use individual vars to construct cluster info.
if not cluster_spec:
cluster_info = ClusterInfo(
primary_node_addr=os.getenv("MASTER_ADDR", None),
primary_node_port=os.getenv("MASTER_PORT", None),
node_rank=int(os.getenv("RANK", "0")),
num_nodes=int(os.getenv("NNODES", "1")),
)
return cluster_info
cluster_data = json.loads(cluster_spec)
# Get primary node info
primary_node = cluster_data["cluster"]["workerpool0"][0]
logging.info("primary node: %s", primary_node)
primary_node_addr, primary_node_port = primary_node.split(":")
logging.info("primary node address: %s", primary_node_addr)
logging.info("primary node port: %s", primary_node_port)
# Determine node rank of this machine
workerpool = cluster_data["task"]["type"]
if workerpool == "workerpool0":
node_rank = 0
elif workerpool == "workerpool1":
# Add 1 for the primary node, since `index` is the index of workerpool1.
node_rank = cluster_data["task"]["index"] + 1
else:
raise ValueError(
"Only workerpool0 and workerpool1 are supported. Unknown workerpool:"
f" {workerpool}"
)
logging.info("node rank: %s", node_rank)
# Calculate total nodes.
num_nodes = 1 # For the primary node.
if "workerpool1" in cluster_data["cluster"]:
num_nodes += len(cluster_data["cluster"]["workerpool1"])
logging.info("num nodes: %s", num_nodes)
return ClusterInfo(
primary_node_addr, primary_node_port, node_rank, num_nodes
)
@@ -1,59 +0,0 @@
"""Add tests for cluster_spec.py."""
import os
from . import cluster_spec
# TODO(styer): Use pytest instead
class ClusterSpecTest(googletest.TestCase):
def setUp(self):
super().setUp()
self.curr_env_var = os.environ.copy()
def tearDown(self):
super().tearDown()
os.environ = self.curr_env_var
def test_get_cluster_spec_from_env_vars(self):
os.environ["CLUSTER_SPEC"] = ""
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = "8080"
os.environ["RANK"] = "0"
os.environ["NNODES"] = "2"
cluster_info = cluster_spec.get_cluster_spec()
self.assertEqual(cluster_info.primary_node_addr, "127.0.0.1")
self.assertEqual(cluster_info.primary_node_port, "8080")
self.assertEqual(cluster_info.node_rank, 0)
self.assertEqual(cluster_info.num_nodes, 2)
def test_get_cluster_spec_from_cluster_spec(self):
os.environ[
"CLUSTER_SPEC"
] = """
{
"cluster": {
"workerpool0": [
"127.0.0.1:8080"
],
"workerpool1": [
"127.0.0.2:8080",
"127.0.0.3:8080"
]
},
"task": {
"type": "workerpool1",
"index": 0
}
}
"""
cluster_info = cluster_spec.get_cluster_spec()
self.assertEqual(cluster_info.primary_node_addr, "127.0.0.1")
self.assertEqual(cluster_info.primary_node_port, "8080")
self.assertEqual(cluster_info.node_rank, 1)
self.assertEqual(cluster_info.num_nodes, 3)
if __name__ == "__main__":
googletest.main()
@@ -1,15 +0,0 @@
# Vertex AI custom prediction routines samples
## Overview
Vertex Custom Prediction Routines(CPR) simplify the process of building custom containers
and make local model testing easy. Here are the sameple codes for different libraries.
### Objectives
The objective is to provide various samples for Vertex Custom Prediction Routine(CPR).
### Supporting libraries
* torch
* sklearn
* xgboost
@@ -1,33 +0,0 @@
import numpy as np
import os
import pickle
from google.cloud.aiplatform.constants import prediction
from google.cloud.aiplatform.utils import prediction_utils
from google.cloud.aiplatform.prediction.predictor import Predictor
from sklearn.datasets import load_breast_cancer
from sklearn.linear_model import RidgeClassifier
class LinearRegressionPredictor(Predictor):
def __init__(self):
return
def load(self, artifacts_uri: str) -> None:
prediction_utils.download_model_artifacts(artifacts_uri)
if os.path.exists(prediction.MODEL_FILENAME_PKL):
self._model = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
else:
self._model = RidgeClassifier()
X, y = load_breast_cancer(return_X_y=True)
self._model.fit(X, y)
def preprocess(self, prediction_input: dict) -> np.ndarray:
instances = prediction_input["instances"]
return np.asarray(instances)
def predict(self, instances: np.ndarray) -> np.ndarray:
return self._model.predict(instances)
def postprocess(self, prediction_results: np.ndarray) -> dict:
return {"predictions": prediction_results.tolist()}
@@ -1,33 +0,0 @@
import numpy as np
import os
import pickle
from google.cloud.aiplatform.constants import prediction
from google.cloud.aiplatform.utils import prediction_utils
from google.cloud.aiplatform.prediction.predictor import Predictor
from sklearn.datasets import make_blobs
from sklearn.linear_model import LinearRegression
class LinearRegressionPredictor(Predictor):
def __init__(self):
return
def load(self, artifacts_uri: str) -> None:
prediction_utils.download_model_artifacts(artifacts_uri)
if os.path.exists(prediction.MODEL_FILENAME_PKL):
self._model = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
else:
self._model = LogisticRegression()
X, y = make_blobs(n_samples=100, centers=2, n_features=2, random_state=1)
self._model.fit(X, y)
def preprocess(self, prediction_input: dict) -> np.ndarray:
instances = prediction_input["instances"]
return np.asarray(instances)
def predict(self, instances: np.ndarray) -> np.ndarray:
return self._model.predict_proba(instances)
def postprocess(self, prediction_results: np.ndarray) -> dict:
return {"predictions": prediction_results.tolist()}
@@ -1,33 +0,0 @@
import numpy as np
import os
import pickle
from google.cloud.aiplatform.constants import prediction
from google.cloud.aiplatform.utils import prediction_utils
from google.cloud.aiplatform.prediction.predictor import Predictor
from sklearn.linear_model import SGDClassifier
class SGDClassifierPredictor(Predictor):
def __init__(self):
return
def load(self, artifacts_uri: str) -> None:
prediction_utils.download_model_artifacts(artifacts_uri)
if os.path.exists(prediction.MODEL_FILENAME_PKL):
self._model = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
else:
self._model = SGDClassifier(max_iter=5)
X = [[0., 0.], [1., 1.]]
y = [0, 1]
self._model.fit(X, y)
def preprocess(self, prediction_input: dict) -> np.ndarray:
instances = prediction_input["instances"]
return np.asarray(instances)
def predict(self, instances: np.ndarray) -> np.ndarray:
return self._model.predict(instances)
def postprocess(self, prediction_results: np.ndarray) -> dict:
return {"predictions": prediction_results.tolist()}
@@ -1,34 +0,0 @@
import os
import torch
from google.cloud.aiplatform.utils import prediction_utils
from google.cloud.aiplatform.prediction.predictor import Predictor
from torchvision.models import detection, resnet50, ResNet50_Weights
from typing import Dict, List
class ResNetPredictor(Predictor):
def __init__(self):
return
def load(self, artifacts_uri: str) -> None:
prediction_utils.download_model_artifacts(artifacts_uri)
if os.path.exists("model.pth.tar"):
self.model = detection.fasterrcnn_resnet50_fpn(pretrained=True)
stat_dic = torch.load("model.pth.tar")
self.model.load_state_dict(stat_dic['state_dict'])
else:
weights = ResNet50_Weights.DEFAULT
self.model = resnet50(weights=weights)
self.model.eval()
def preprocess(self, prediction_input: dict) -> torch.Tensor:
instances = prediction_input["instances"]
return torch.Tensor(instances)
@torch.inference_mode()
def predict(self, instances: torch.Tensor) -> List[str]:
return self._model(instances)
def postprocess(self, prediction_results: List[str]) -> Dict:
return {"predictions": prediction_results}
@@ -1,73 +0,0 @@
import ast
import json
import os
import pickle
import torch
from google.cloud.aiplatform.utils import prediction_utils
from google.cloud.aiplatform.prediction.predictor import Predictor
from transformers import AutoModelForQuestionAnswering
from typing import Dict, List
class TorchTransformersPredictor(Predictor):
def __init__(self):
return
def load(self, artifacts_uri: str) -> None:
prediction_utils.download_model_artifacts(artifacts_uri)
if os.path.isfile("setup_config.json"):
with open("setup_config.json") as setup_config_file:
self.setup_config = json.load(setup_config_file)
if os.path.exists("model.pt"):
self.model = AutoModelForQuestionAnswering.from_pretrained("model.pt")
self.model.eval()
else:
raise ValueError("One of the following model files must be provided: model.pt.")
def preprocess(self, prediction_input: dict) -> torch.Tensor:
max_length = self.setup_config["max_length"]
instances = prediction_input["instances"]
question_context = ast.literal_eval(instances)
question = question_context["question"]
context = question_context["context"]
inputs = self.tokenizer.encode_plus(
question,
context,
max_length=int(max_length),
pad_to_max_length=True,
add_special_tokens=True,
return_tensors="pt",
)
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]
return torch.Tensor(input_ids, attention_mask)
@torch.inference_mode()
def predict(self, instances: torch.Tensor) -> List[str]:
input_ids, attention_mask = instances
outputs = self._model(input_ids, attention_mask)
answer_start_scores = outputs.start_logits
answer_end_scores = outputs.end_logits
num_rows, num_cols = answer_start_scores.shape
inferences = []
for i in range(num_rows):
answer_start_scores_one_seq = answer_start_scores[i].unsqueeze(0)
answer_start = torch.argmax(answer_start_scores_one_seq)
answer_end_scores_one_seq = answer_end_scores[i].unsqueeze(0)
answer_end = torch.argmax(answer_end_scores_one_seq) + 1
prediction = self.tokenizer.convert_tokens_to_string(
self.tokenizer.convert_ids_to_tokens(
input_ids[i].tolist()[answer_start:answer_end]
)
)
inferences.append(prediction)
return inferences
def postprocess(self, prediction_results: List[str]) -> Dict:
return {"predictions": prediction_results}
@@ -1,37 +0,0 @@
import os
import numpy as np
import pickle
import xgboost as xgb
from google.cloud.aiplatform.constants import prediction
from google.cloud.aiplatform.utils import prediction_utils
from google.cloud.aiplatform.prediction.predictor import Predictor
from sklearn.datasets import make_blobs
from xgboost import XGBClassifier
class ClassifierPredictor(Predictor):
def __init__(self):
return
def load(self, artifacts_uri: str) -> None:
prediction_utils.download_model_artifacts(artifacts_uri)
if os.path.exists(prediction.MODEL_FILENAME_PKL):
booster = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
else:
X, y = make_blobs(n_samples=100, centers=2, n_features=2, random_state=1)
model = XGBClassifier()
model.fit(X, y)
booster = model.get_booster()
self._booster = booster
def preprocess(self, prediction_input: dict) -> xgb.DMatrix:
instances = prediction_input["instances"]
return xgb.DMatrix(instances)
def predict(self, instances: xgb.DMatrix) -> np.ndarray:
return self._booster.predict(instances)
def postprocess(self, prediction_results: np.ndarray) -> dict:
return {"predictions": prediction_results.tolist()}
@@ -1,41 +0,0 @@
import os
import numpy as np
import pandas as pd
import pickle
import xgboost as xgb
from google.cloud.aiplatform.constants import prediction
from google.cloud.aiplatform.utils import prediction_utils
from google.cloud.aiplatform.prediction.predictor import Predictor
class XGBRankerPredictor(Predictor):
def __init__(self):
return
def load(self, artifacts_uri: str) -> None:
prediction_utils.download_model_artifacts(artifacts_uri)
if os.path.exists(prediction.MODEL_FILENAME_PKL):
booster = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
self._booster = booster
else:
N = 500
dates = pd.date_range(start='2023-01-01', end='2023-01-12', periods=N)
X = pd.DataFrame(np.random.randn(N, 5), columns=list('ABCDE'), index=dates)
y = pd.Series(np.random.randint(0, 10, size=N), index=dates, name='label')
group = X.groupby(dates + pd.offsets.MonthEnd(0)).size()
sample_weight = pd.Series(np.arange(len(group)), index=group.index)
model = xgb.XGBRanker(objective='rank:pairwise', max_depth=3, learning_rate=0.1, booster='gbtree', tree_method='hist', n_jobs=4, n_estimators=50, enable_categorical=False, random_state=42)
model.fit(X=X, y=y, group=group, sample_weight=sample_weight, verbose=True)
booster = model.get_booster()
self._booster = booster
def preprocess(self, prediction_input: dict) -> xgb.DMatrix:
instances = prediction_input["instances"]
return xgb.DMatrix(instances)
def predict(self, instances: xgb.DMatrix) -> np.ndarray:
return self._booster.predict(instances, output_margin=False, ntree_limit=0)
def postprocess(self, prediction_results: np.ndarray) -> dict:
return {"predictions": prediction_results.tolist()}
@@ -1,4 +0,0 @@
[MASTER]
generated-members=get_concrete_function,cv2.*
ignored-modules=tensorflow,google.cloud

Some files were not shown because too many files have changed in this diff Show More