Skip to content

Commit

Permalink
Support for external image uploads in eecli
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 689482120
  • Loading branch information
Google Earth Engine Authors committed Oct 24, 2024
1 parent 6f51b8d commit 0c7dc76
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 43 deletions.
47 changes: 31 additions & 16 deletions python/ee/cli/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -1261,6 +1261,22 @@ def run(
utils.wait_for_tasks(task_ids, args.timeout, log_progress=args.verbose)


def _using_v1alpha(func):
"""Decorator that temporarily switches over to the v1alpha API."""

def inner(
self, args: argparse.Namespace, config: utils.CommandLineConfig
) -> None:
# pylint: disable=protected-access
original = ee._cloud_api_utils.VERSION
ee._cloud_api_utils.VERSION = 'v1alpha'
func(self, args, config)
ee._cloud_api_utils.VERSION = original
# pylint: enable=protected-access

return inner


class TaskCommand(Dispatcher):
"""Prints information about or manages long-running tasks."""

Expand Down Expand Up @@ -1432,6 +1448,20 @@ def is_tf_record(path: str) -> bool:
return manifest


class UploadExternalImageCommand(UploadImageCommand):
name = 'external_image'

@_using_v1alpha
def run(
self, args: argparse.Namespace, config: utils.CommandLineConfig
) -> None:
"""Creates an external image synchronously."""
config.ee_init()
manifest = self.manifest_from_args(args)
name = ee.data.startExternalImageIngestion(manifest, args.force)['name']
print('Created asset %s' % name)


# TODO(user): update src_files help string when secondary files
# can be uploaded.
class UploadTableCommand:
Expand Down Expand Up @@ -1621,6 +1651,7 @@ class UploadCommand(Dispatcher):

COMMANDS = [
UploadImageCommand,
UploadExternalImageCommand,
UploadTableCommand,
]

Expand Down Expand Up @@ -1937,22 +1968,6 @@ class ModelCommand(Dispatcher):
COMMANDS = [PrepareModelCommand]


def _using_v1alpha(func):
"""Decorator that temporarily switches over to the v1alpha API."""

def inner(
self, args: argparse.Namespace, config: utils.CommandLineConfig
) -> None:
# pylint: disable=protected-access
original = ee._cloud_api_utils.VERSION
ee._cloud_api_utils.VERSION = 'v1alpha'
func(self, args, config)
ee._cloud_api_utils.VERSION = original
# pylint: enable=protected-access

return inner


class ProjectConfigGetCommand:
"""Prints the current project's ProjectConfig."""

Expand Down
115 changes: 88 additions & 27 deletions python/ee/data.py
Original file line number Diff line number Diff line change
Expand Up @@ -1908,9 +1908,62 @@ def _prepare_and_run_export(
export_endpoint(project=_get_projects_path(), body=params),
num_retries=num_retries)

# TODO(user): use StrEnum when 3.11 is the min version
_INTERNAL_IMPORT = 'INTERNAL_IMPORT'
_EXTERNAL_IMPORT = 'EXTERNAL_IMPORT'


def _startIngestion(
request_id: Any,
params: Dict[str, Any],
allow_overwrite: bool = False,
import_mode: Optional[str] = _INTERNAL_IMPORT,
) -> Dict[str, Any]:
"""Starts an ingestion task or creates an external image."""
request = {
'imageManifest':
_cloud_api_utils.convert_params_to_image_manifest(params),
'overwrite':
allow_overwrite
}

# It's only safe to retry the request if there's a unique ID to make it
# idempotent.
num_retries = _max_retries if request_id else 0

image = _get_cloud_projects().image()
if import_mode == _INTERNAL_IMPORT:
import_request = image.import_(project=_get_projects_path(), body=request)
elif import_mode == _EXTERNAL_IMPORT:
import_request = image.importExternal(
project=_get_projects_path(), body=request
)
else:
raise ee_exception.EEException(
'{} is not a valid import mode'.format(import_mode)
)

result = _execute_cloud_call(
import_request,
num_retries=num_retries,
)

if import_mode == _INTERNAL_IMPORT:
return {
'id': _cloud_api_utils.convert_operation_name_to_task_id(
result['name']
),
'name': result['name'],
'started': 'OK',
}
else:
return {'name': request['imageManifest']['name']}


def startIngestion(
request_id: Any, params: Dict[str, Any], allow_overwrite: bool = False
request_id: Any,
params: Dict[str, Any],
allow_overwrite: bool = False,
) -> Dict[str, Any]:
"""Creates an image asset import task.
Expand All @@ -1923,7 +1976,7 @@ def startIngestion(
params: The object that describes the import task, which can
have these fields:
name (string) The destination asset id (e.g.,
"projects/earthengine-legacy/assets/users/foo/bar").
"projects/myproject/assets/foo/bar").
tilesets (array) A list of Google Cloud Storage source file paths
formatted like:
[{'sources': [
Expand All @@ -1942,31 +1995,39 @@ def startIngestion(
A dict with notes about the created task. This will include the ID for the
import task (under 'id'), which may be different from request_id.
"""
request = {
'imageManifest':
_cloud_api_utils.convert_params_to_image_manifest(params),
'requestId':
request_id,
'overwrite':
allow_overwrite
}
return _startIngestion(request_id, params, allow_overwrite, _INTERNAL_IMPORT)

# It's only safe to retry the request if there's a unique ID to make it
# idempotent.
num_retries = _max_retries if request_id else 0
operation = _execute_cloud_call(
_get_cloud_projects()
.image()
.import_(project=_get_projects_path(), body=request),
num_retries=num_retries,
)
return {
'id':
_cloud_api_utils.convert_operation_name_to_task_id(
operation['name']),
'name': operation['name'],
'started': 'OK',
}

def startExternalImageIngestion(
image_manifest: Dict[str, Any],
allow_overwrite: bool = False,
) -> Dict[str, Any]:
"""Creates an external image.
Args:
image_manifest: The object that describes the import task, which can
have these fields:
name (string) The destination asset id (e.g.,
"projects/myproject/assets/foo/bar").
tilesets (array) A list of Google Cloud Storage source file paths
formatted like:
[{'sources': [
{'uris': ['foo.tif', 'foo.prj']},
{'uris': ['bar.tif', 'bar.prj']},
]}]
Where path values correspond to source files' Google Cloud Storage
object names, e.g., 'gs://bucketname/filename.tif'
bands (array) An optional list of band names formatted like:
[{'id': 'R'}, {'id': 'G'}, {'id': 'B'}]
In general, this is a dict representation of an ImageManifest.
allow_overwrite: Whether the ingested image can overwrite an
existing version.
Returns:
The name of the created asset.
"""
return _startIngestion(
'unused', image_manifest, allow_overwrite, _EXTERNAL_IMPORT)


def startTableIngestion(
Expand All @@ -1983,7 +2044,7 @@ def startTableIngestion(
params: The object that describes the import task, which can
have these fields:
name (string) The destination asset id (e.g.,
"projects/earthengine-legacy/assets/users/foo/bar").
"projects/myproject/assets/foo/bar").
sources (array) A list of GCS (Google Cloud Storage) file paths
with optional character encoding formatted like this:
"sources":[{"uris":["gs://bucket/file.shp"],"charset":"UTF-8"}]
Expand Down

0 comments on commit 0c7dc76

Please sign in to comment.