|
32 | 32 | from .models import FOLDER_REGEX, NAME_REGEX
|
33 | 33 | from .config.user import SimvueConfiguration
|
34 | 34 | from .api.request import get_json_from_response
|
35 |
| -from .api.objects import Run, Folder, Tag, Artifact, Alert |
| 35 | +from .api.objects import Run, Folder, Tag, Artifact, Alert, FileArtifact, ObjectArtifact |
36 | 36 |
|
37 | 37 |
|
38 | 38 | CONCURRENT_DOWNLOADS = 10
|
|
42 | 42 |
|
43 | 43 |
|
44 | 44 | def _download_artifact_to_file(
|
45 |
| - artifact: Artifact, output_dir: pathlib.Path | None |
| 45 | + artifact: FileArtifact | ObjectArtifact, output_dir: pathlib.Path | None |
46 | 46 | ) -> None:
|
| 47 | + if not artifact.name: |
| 48 | + raise RuntimeError(f"Expected artifact '{artifact.id}' to have a name") |
47 | 49 | _output_file = (output_dir or pathlib.Path.cwd()).joinpath(artifact.name)
|
48 | 50 | # If this is a hierarchical structure being downloaded, need to create directories
|
49 | 51 | _output_file.parent.mkdir(parents=True, exist_ok=True)
|
@@ -475,13 +477,14 @@ def list_artifacts(
|
475 | 477 | ) # type: ignore
|
476 | 478 |
|
477 | 479 | def _retrieve_artifacts_from_server(
|
478 |
| - self, run_id: str, name: str, count: int | None = None |
479 |
| - ) -> typing.Generator[tuple[str, Artifact], None, None]: |
480 |
| - return Artifact.get( |
481 |
| - runs=json.dumps([run_id]), |
482 |
| - filters=json.dumps([f"name == {name}"]), |
483 |
| - count=count, |
484 |
| - ) # type: ignore |
| 480 | + self, run_id: str, name: str |
| 481 | + ) -> FileArtifact | ObjectArtifact | None: |
| 482 | + return Artifact.from_name( |
| 483 | + run_id=run_id, |
| 484 | + name=name, |
| 485 | + server_url=self._user_config.server.url, |
| 486 | + server_token=self._user_config.server.token, |
| 487 | + ) |
485 | 488 |
|
486 | 489 | @prettify_pydantic
|
487 | 490 | @pydantic.validate_call
|
@@ -529,12 +532,14 @@ def get_artifact(
|
529 | 532 | RuntimeError
|
530 | 533 | if retrieval of artifact from the server failed
|
531 | 534 | """
|
532 |
| - _artifact = Artifact.from_name( |
533 |
| - run_id=run_id, |
534 |
| - name=name, |
535 |
| - server_url=self._user_config.server.url, |
536 |
| - server_token=self._user_config.server.token, |
537 |
| - ) |
| 535 | + _artifact = self._retrieve_artifacts_from_server(run_id, name) |
| 536 | + |
| 537 | + if not _artifact: |
| 538 | + raise ObjectNotFoundError( |
| 539 | + obj_type="artifact", |
| 540 | + name=name, |
| 541 | + extra=f"for run '{run_id}'", |
| 542 | + ) |
538 | 543 |
|
539 | 544 | _content = b"".join(_artifact.download_content())
|
540 | 545 |
|
@@ -574,12 +579,14 @@ def get_artifact_as_file(
|
574 | 579 | if there was a failure during retrieval of information from the
|
575 | 580 | server
|
576 | 581 | """
|
577 |
| - _artifacts = self._retrieve_artifacts_from_server(run_id, name, count=1) |
| 582 | + _artifact = self._retrieve_artifacts_from_server(run_id, name) |
578 | 583 |
|
579 |
| - try: |
580 |
| - _id, _artifact = next(_artifacts) |
581 |
| - except StopIteration as e: |
582 |
| - raise ValueError(f"No artifact '{name}' found for run '{run_id}'") from e |
| 584 | + if not _artifact: |
| 585 | + raise ObjectNotFoundError( |
| 586 | + obj_type="artifact", |
| 587 | + name=name, |
| 588 | + extra=f"for run '{run_id}'", |
| 589 | + ) |
583 | 590 |
|
584 | 591 | _download_artifact_to_file(_artifact, output_dir)
|
585 | 592 |
|
|
0 commit comments