Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[GPartition/SageMaker] Improvements to gpartition and SageMaker partition. #880

Merged
merged 1 commit into from
Jun 27, 2024
Merged

[GPartition/SageMaker] Improvements to gpartition and SageMaker partition. #880

merged 1 commit into from
Jun 27, 2024

Conversation

thvasilo
Copy link
Contributor

We use this commit as a base for range partitioning.

Issue #, if available:

Fixes #802

Description of changes:

  • Download files in parallel from S3 for SageMaker partition.
  • Copy raw_id_mappings from the input path to the generated DGL graph path for SageMaker and local.
  • Change dtype of random partition assignment pyarrow array, from int64 to uint8 for cases where num_parts < 256, uint16 otherwise.
  • Support getting SageMaker job information for both SageMaker Processing jobs, and SageMaker Training jobs, allowing us to use more instance types.
  • Make gpartition unit test verification re-usable across partition algorithm implementations.

We use this PR as a base for RangePartition.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

@thvasilo thvasilo added the ready able to trigger the CI label Jun 17, 2024
@thvasilo thvasilo self-assigned this Jun 17, 2024
@thvasilo
[GPartition/SageMaker] Improvements to partition for gpartition and S…
b94e9e3 
…ageMaker.

We use this commit as a base for range partitioning.
@thvasilo thvasilo force-pushed the gsf-partition-improvements branch from ee7cc86 to b94e9e3 Compare June 17, 2024 22:14
@thvasilo thvasilo mentioned this pull request Jun 17, 2024
Draft
jalencato
jalencato reviewed Jun 18, 2024
View reviewed changes
Copy link
Collaborator

@jalencato jalencato left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some small comments

python/graphstorm/gpartition/dist_partition_graph.py Show resolved Hide resolved
python/graphstorm/sagemaker/sagemaker_partition.py Show resolved Hide resolved
classicsong
classicsong approved these changes Jun 26, 2024
View reviewed changes
Copy link
Contributor

@classicsong classicsong left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

jalencato
jalencato approved these changes Jun 26, 2024
View reviewed changes
Copy link
Collaborator

@jalencato jalencato left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@thvasilo thvasilo merged commit f009041 into awslabs:main Jun 27, 2024
11 checks passed
@thvasilo thvasilo deleted the gsf-partition-improvements branch June 27, 2024 18:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@jalencato jalencato jalencato approved these changes

@classicsong classicsong classicsong approved these changes

Assignees

@thvasilo thvasilo

Labels
0.3.1 ready able to trigger the CI
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

DistPartition should move the raw id mappings to the partition output
3 participants
@thvasilo @classicsong @jalencato