datachain 0.19.1__tar.gz → 0.19.2__tar.gz

{datachain-0.19.1 → datachain-0.19.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.19.1
+Version: 0.19.2
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

datachain-0.19.2/docs/commands/auth/team.md

@@ -0,0 +1,42 @@
+# auth team
+
+Set the default team for Studio operations.
+
+## Synopsis
+
+```usage
+usage: datachain auth team [-h] [-v] [-q] [--local] [team_name]
+```
+
+## Description
+
+This command sets or shows the default team for Studio operations. By default, the team setting is saved globally, but you can use the `--local` option to set it for the current project.
+If team_name is not passed, the current team in use is shown to the user.
+
+## Arguments
+
+* `team_name` - Optional,  Name of the team to set as default
+
+## Options
+
+* `--local` - Set team locally for the current project
+* `-h`, `--help` - Show the help message and exit.
+* `-v`, `--verbose` - Be verbose.
+* `-q`, `--quiet` - Be quiet.
+
+## Examples
+
+1. Set default team for all projects:
+```bash
+datachain auth team my-team
+```
+
+2. Set default team locally for the current project:
+```bash
+datachain auth team --local my-team
+```
+
+3. Print the current default in use
+```bash
+datachain auth team
+```

{datachain-0.19.1 → datachain-0.19.2}/docs/commands/index.md

@@ -31,3 +31,5 @@ DataChain is a command-line tool for wrangling unstructured AI data at scale. Us
 	- Monitor job logs with [`datachain job logs`](job/logs.md)
 
 	- Cancel running jobs with [`datachain job cancel`](job/cancel.md)
+
+	- Check for the clusters available for jobs [`datachain job clusters`](job/clusters.md)

datachain-0.19.2/docs/commands/job/clusters.md

@@ -0,0 +1,39 @@
+# job clusters
+
+List compute clusters in Studio.
+
+## Synopsis
+
+```usage
+usage: datachain job clusters [-h] [-v] [-q] [--team TEAM]
+```
+
+## Description
+
+This command lists compute clusters available in Studio. You can specify a team to list clusters for. The command provides information about the compute resources available for running jobs.
+
+## Options
+
+* `--team TEAM` - Team to list clusters for (default: from config)
+* `-h`, `--help` - Show the help message and exit
+* `-v`, `--verbose` - Be verbose
+* `-q`, `--quiet` - Be quiet
+
+## Examples
+
+1. List all clusters for the default team:
+```bash
+datachain job clusters
+```
+
+2. List clusters for a specific team:
+```bash
+datachain job clusters --team my-team
+```
+
+
+## Notes
+
+* The command shows all compute clusters available to your team
+* Clusters represent the compute resources where your jobs can run
+* Use the `--team` option to view clusters for a different team

{datachain-0.19.1 → datachain-0.19.2}/docs/commands/job/run.md

@@ -24,6 +24,7 @@ This command runs a job in Studio using the specified query file. You can config
 * `--team TEAM` - Team to run job for (default: from config)
 * `--env-file ENV_FILE` - File with environment variables for the job
 * `--env ENV` - Environment variables in KEY=VALUE format
+* `--cluster-id CLUSTER_ID` - Compute cluster ID to run the job on
 * `--workers WORKERS` - Number of workers for the job
 * `--files FILES` - Additional files to include in the job
 * `--python-version PYTHON_VERSION` - Python version for the job (e.g., 3.9, 3.10, 3.11)
@@ -71,8 +72,17 @@ datachain job run --repository https://github.com/iterative/datachain query.py
 datachain job run --priority 2 query.py
 ```
 
+8. Run a job in a specific cluster
+```bash
+# Get the cluster id using following command
+datachain job clusters
+# Use the id  of an active clusters from above
+datachain job run --cluster-id 1 query.py
+```
+
 ## Notes
 
 * Closing the logs command (e.g., with Ctrl+C) will only stop displaying the logs but will not cancel the job execution
 * To cancel a running job, use the `datachain job cancel` command
 * The job will continue running in Studio even after you stop viewing the logs
+* You can get the list of compute clusters using `datachain job clusters` command.

{datachain-0.19.1 → datachain-0.19.2}/mkdocs.yml

@@ -97,6 +97,7 @@ nav:
                   - logs: commands/job/logs.md
                   - cancel: commands/job/cancel.md
                   - ls: commands/job/ls.md
+                  - clusters: commands/job/clusters.md
       - 📚 User Guide:
           - Overview: guide/index.md
           - 📡 Interacting with remote storage: guide/remotes.md

{datachain-0.19.1 → datachain-0.19.2}/src/datachain/cli/parser/job.py

@@ -51,6 +51,13 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         help="Environment variables in KEY=VALUE format",
     )
 
+    studio_run_parser.add_argument(
+        "--cluster-id",
+        type=int,
+        action="store",
+        help="Compute cluster ID to run the job on",
+    )
+
     studio_run_parser.add_argument(
         "--workers",
         type=int,
@@ -165,3 +172,21 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         default=None,
         help="Team to check logs for (default: from config)",
     )
+
+    studio_clusters_help = "List compute clusters in Studio"
+    studio_clusters_description = "List compute clusters in Studio."
+
+    studio_clusters_parser = jobs_subparser.add_parser(
+        "clusters",
+        parents=[parent_parser],
+        description=studio_clusters_description,
+        help=studio_clusters_help,
+        formatter_class=CustomHelpFormatter,
+    )
+
+    studio_clusters_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="Team to list clusters for (default: from config)",
+    )

{datachain-0.19.1 → datachain-0.19.2}/src/datachain/cli/parser/studio.py

@@ -89,8 +89,13 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         help="Remove the token from the local project config",
     )
 
-    auth_team_help = "Set default team for Studio operations"
-    auth_team_description = "Set the default team for Studio operations."
+    auth_team_help = "Set or show default team for Studio operations"
+    auth_team_description = (
+        "Set or show the default team for Studio operations. "
+        "This will be used globally by default. "
+        "Use --local to set the team locally for the current project. "
+        "If no team name is provided, the default team will be shown."
+    )
 
     team_parser = auth_subparser.add_parser(
         "team",
@@ -102,13 +107,15 @@ def add_auth_parser(subparsers, parent_parser) -> None:
     team_parser.add_argument(
         "team_name",
         action="store",
+        default=None,
+        nargs="?",
         help="Name of the team to set as default",
     )
     team_parser.add_argument(
-        "--global",
+        "--local",
         action="store_true",
         default=False,
-        help="Set team globally for all projects",
+        help="Set team locally for the current project",
     )
 
     auth_token_help = "View Studio authentication token"  # noqa: S105

{datachain-0.19.1 → datachain-0.19.2}/src/datachain/data_storage/schema.py

@@ -13,7 +13,16 @@ from sqlalchemy.sql import func as f
 from sqlalchemy.sql.expression import false, null, true
 
 from datachain.sql.functions import path as pathfunc
-from datachain.sql.types import Int, SQLType, UInt64
+from datachain.sql.types import (
+    JSON,
+    Boolean,
+    DateTime,
+    Int,
+    Int64,
+    SQLType,
+    String,
+    UInt64,
+)
 
 if TYPE_CHECKING:
     from sqlalchemy.engine.interfaces import Dialect
@@ -272,6 +281,19 @@ class DataTable:
             ),
         ]
 
+    @classmethod
+    def listing_columns(cls):
+        return [
+            sa.Column("file__source", String()),
+            sa.Column("file__path", String()),
+            sa.Column("file__size", Int64()),
+            sa.Column("file__version", String()),
+            sa.Column("file__etag", String()),
+            sa.Column("file__is_latest", Boolean()),
+            sa.Column("file__last_modified", DateTime()),
+            sa.Column("file__location", JSON()),
+        ]
+
     def dir_expansion(self):
         return DirExpansion(self.column)
 

datachain-0.19.2/src/datachain/lib/dc/listings.py

@@ -0,0 +1,154 @@
+from typing import (
+    TYPE_CHECKING,
+    Optional,
+)
+
+from datachain.lib.listing import LISTING_PREFIX, ls
+from datachain.lib.listing_info import ListingInfo
+from datachain.lib.settings import Settings
+from datachain.lib.signal_schema import SignalSchema
+from datachain.query import Session
+from datachain.query.dataset import DatasetQuery, QueryStep, step_result
+
+from .values import read_values
+
+if TYPE_CHECKING:
+    from typing_extensions import ParamSpec
+
+    from datachain.dataset import DatasetVersion
+    from datachain.query.dataset import StepResult
+
+    from .datachain import DataChain
+
+    P = ParamSpec("P")
+
+
+class ReadOnlyQueryStep(QueryStep):
+    """
+    This step is used to read the dataset in read-only mode.
+    It is used to avoid the need to read the table metadata from the warehouse.
+    This is useful when we want to list the files in the dataset.
+    """
+
+    def apply(self) -> "StepResult":
+        import sqlalchemy as sa
+
+        def q(*columns):
+            return sa.select(*columns)
+
+        table_name = self.catalog.warehouse.dataset_table_name(
+            self.dataset_name, self.dataset_version
+        )
+        dataset_row_cls = self.catalog.warehouse.schema.dataset_row_cls
+        table = dataset_row_cls.new_table(
+            table_name,
+            columns=(
+                [
+                    *dataset_row_cls.sys_columns(),
+                    *dataset_row_cls.listing_columns(),
+                ]
+            ),
+        )
+
+        return step_result(
+            q, table.columns, dependencies=[(self.dataset_name, self.dataset_version)]
+        )
+
+
+def listings(
+    session: Optional[Session] = None,
+    in_memory: bool = False,
+    column: str = "listing",
+    **kwargs,
+) -> "DataChain":
+    """Generate chain with list of cached listings.
+    Listing is a special kind of dataset which has directory listing data of
+    some underlying storage (e.g S3 bucket).
+
+    Example:
+        ```py
+        import datachain as dc
+        dc.listings().show()
+        ```
+    """
+    session = Session.get(session, in_memory=in_memory)
+    catalog = kwargs.get("catalog") or session.catalog
+
+    return read_values(
+        session=session,
+        in_memory=in_memory,
+        output={column: ListingInfo},
+        **{column: catalog.listings()},  # type: ignore[arg-type]
+    )
+
+
+def read_listing_dataset(
+    name: str,
+    version: Optional[str] = None,
+    path: str = "",
+    session: Optional["Session"] = None,
+    settings: Optional[dict] = None,
+) -> tuple["DataChain", "DatasetVersion"]:
+    """Read a listing dataset and return a DataChain and listing version.
+
+    Args:
+        name: Name of the dataset
+        version: Version of the dataset
+        path: Path within the listing to read. Path can have globs.
+        session: Optional Session object to use for reading
+        settings: Optional settings dictionary to use for reading
+
+    Returns:
+        tuple[DataChain, DatasetVersion]: A tuple containing:
+            - DataChain configured for listing files
+            - DatasetVersion object for the specified listing version
+
+    Example:
+        ```py
+        import datachain as dc
+        chain, listing_version = dc.read_listing_dataset(
+            "lst__s3://my-bucket/my-path", version="1.0.0", path="my-path"
+        )
+        chain.show()
+        ```
+    """
+    # Configure and return a DataChain for reading listing dataset files
+    # Uses ReadOnlyQueryStep to avoid warehouse metadata lookups
+    from datachain.lib.dc import Sys
+    from datachain.lib.file import File
+
+    from .datachain import DataChain
+
+    if not name.startswith(LISTING_PREFIX):
+        name = LISTING_PREFIX + name
+
+    session = Session.get(session)
+    dataset = session.catalog.get_dataset(name)
+    if version is None:
+        version = dataset.latest_version
+
+    query = DatasetQuery(
+        name=name,
+        session=session,
+        indexing_column_types=File._datachain_column_types,
+        fallback_to_studio=False,
+    )
+    if settings:
+        cfg = {**settings}
+        if "prefetch" not in cfg:
+            cfg["prefetch"] = 0
+        _settings = Settings(**cfg)
+    else:
+        _settings = Settings(prefetch=0)
+    signal_schema = SignalSchema({"sys": Sys, "file": File})
+
+    query.starting_step = ReadOnlyQueryStep(query.catalog, name, version)
+    query.version = version
+    # We already know that this is a listing dataset,
+    # so we can set the listing function to True
+    query.set_listing_fn(lambda: True)
+
+    chain = DataChain(query, _settings, signal_schema)
+    chain = ls(chain, path, recursive=True, column="file")
+
+    return chain, dataset.get_version(version)

{datachain-0.19.1 → datachain-0.19.2}/src/datachain/remote/studio.py

@@ -30,6 +30,7 @@ DatasetExportSignedUrls = Optional[list[str]]
 FileUploadData = Optional[dict[str, Any]]
 JobData = Optional[dict[str, Any]]
 JobListData = dict[str, Any]
+ClusterListData = dict[str, Any]
 logger = logging.getLogger("datachain")
 
 DATASET_ROWS_CHUNK_SIZE = 8192
@@ -391,6 +392,7 @@ class StudioClient:
         requirements: Optional[str] = None,
         repository: Optional[str] = None,
         priority: Optional[int] = None,
+        cluster_id: Optional[int] = None,
     ) -> Response[JobData]:
         data = {
             "query": query,
@@ -403,6 +405,7 @@ class StudioClient:
             "requirements": requirements,
             "repository": repository,
             "priority": priority,
+            "compute_cluster_id": cluster_id,
         }
         return self._send_request("datachain/job", data)
 
@@ -423,3 +426,6 @@ class StudioClient:
     ) -> Response[JobData]:
         url = f"datachain/job/{job_id}/cancel"
         return self._send_request(url, data={}, method="POST")
+
+    def get_clusters(self) -> Response[ClusterListData]:
+        return self._send_request("datachain/clusters", {}, method="GET")

{datachain-0.19.1 → datachain-0.19.2}/src/datachain/studio.py

@@ -41,6 +41,7 @@ def process_jobs_args(args: "Namespace"):
             args.req,
             args.req_file,
             args.priority,
+            args.cluster_id,
         )
 
     if args.cmd == "cancel":
@@ -51,6 +52,9 @@ def process_jobs_args(args: "Namespace"):
     if args.cmd == "ls":
         return list_jobs(args.status, args.team, args.limit)
 
+    if args.cmd == "clusters":
+        return list_clusters(args.team)
+
     raise DataChainError(f"Unknown command '{args.cmd}'.")
 
 
@@ -68,14 +72,24 @@ def process_auth_cli_args(args: "Namespace"):
         return logout(args.local)
     if args.cmd == "token":
         return token()
-
     if args.cmd == "team":
         return set_team(args)
     raise DataChainError(f"Unknown command '{args.cmd}'.")
 
 
 def set_team(args: "Namespace"):
-    level = ConfigLevel.GLOBAL if args.__dict__.get("global") else ConfigLevel.LOCAL
+    if args.team_name is None:
+        config = Config().read().get("studio", {})
+        team = config.get("team")
+        if team:
+            print(f"Default team is '{team}'")
+            return 0
+
+        raise DataChainError(
+            "No default team set. Use `datachain auth team <team_name>` to set one."
+        )
+
+    level = ConfigLevel.LOCAL if args.local else ConfigLevel.GLOBAL
     config = Config(level)
     with config.edit() as conf:
         studio_conf = conf.get("studio", {})
@@ -121,6 +135,7 @@ def login(args: "Namespace"):
     level = ConfigLevel.LOCAL if args.local else ConfigLevel.GLOBAL
     config_path = save_config(hostname, access_token, level=level)
     print(f"Authentication complete. Saved token to {config_path}.")
+    print("You can now use 'datachain auth team' to set the default team.")
     return 0
 
 
@@ -268,6 +283,7 @@ def create_job(
     req: Optional[list[str]] = None,
     req_file: Optional[str] = None,
     priority: Optional[int] = None,
+    cluster_id: Optional[int] = None,
 ):
     query_type = "PYTHON" if query_file.endswith(".py") else "SHELL"
     with open(query_file) as f:
@@ -297,6 +313,7 @@ def create_job(
         repository=repository,
         requirements=requirements,
         priority=priority,
+        cluster_id=cluster_id,
     )
     if not response.ok:
         raise DataChainError(response.message)
@@ -380,3 +397,29 @@ def show_job_logs(job_id: str, team_name: Optional[str]):
 
     client = StudioClient(team=team_name)
     show_logs_from_client(client, job_id)
+
+
+def list_clusters(team_name: Optional[str]):
+    client = StudioClient(team=team_name)
+    response = client.get_clusters()
+    if not response.ok:
+        raise DataChainError(response.message)
+
+    clusters = response.data.get("clusters", [])
+    if not clusters:
+        print("No clusters found")
+        return
+
+    rows = [
+        {
+            "ID": cluster.get("id"),
+            "Status": cluster.get("status"),
+            "Cloud Provider": cluster.get("cloud_provider"),
+            "Cloud Credentials": cluster.get("cloud_credentials"),
+            "Is Active": cluster.get("is_active"),
+            "Max Workers": cluster.get("max_workers"),
+        }
+        for cluster in clusters
+    ]
+
+    print(tabulate.tabulate(rows, headers="keys", tablefmt="grid"))

{datachain-0.19.1 → datachain-0.19.2}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.19.1
+Version: 0.19.2
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.19.1 → datachain-0.19.2}/src/datachain.egg-info/SOURCES.txt

@@ -32,6 +32,7 @@ docs/commands/auth/logout.md
 docs/commands/auth/team.md
 docs/commands/auth/token.md
 docs/commands/job/cancel.md
+docs/commands/job/clusters.md
 docs/commands/job/logs.md
 docs/commands/job/ls.md
 docs/commands/job/run.md

{datachain-0.19.1 → datachain-0.19.2}/tests/test_cli_studio.py

@@ -109,13 +109,13 @@ def test_studio_token(capsys):
 
 def test_studio_team_local():
     assert main(["auth", "team", "team_name"]) == 0
-    config = Config(ConfigLevel.LOCAL).read()
+    config = Config(ConfigLevel.GLOBAL).read()
     assert config["studio"]["team"] == "team_name"
 
 
 def test_studio_team_global():
-    assert main(["auth", "team", "team_name", "--global"]) == 0
-    config = Config(ConfigLevel.GLOBAL).read()
+    assert main(["auth", "team", "team_name", "--local"]) == 0
+    config = Config(ConfigLevel.LOCAL).read()
     assert config["studio"]["team"] == "team_name"
 
 
@@ -387,6 +387,8 @@ def test_studio_run(capsys, mocker, tmp_dir):
                     "stupidity",
                     "--repository",
                     "https://github.com/iterative/datachain",
+                    "--cluster-id",
+                    "1",
                 ]
             )
             == 0
@@ -425,4 +427,5 @@ def test_studio_run(capsys, mocker, tmp_dir):
         "team_name": "team_name",
         "repository": "https://github.com/iterative/datachain",
         "priority": 5,
+        "compute_cluster_id": 1,
     }

{datachain-0.19.1 → datachain-0.19.2}/tests/unit/lib/test_datachain.py

@@ -19,8 +19,9 @@ from datachain import Column
 from datachain.error import DatasetInvalidVersionError, DatasetVersionNotFoundError
 from datachain.lib.data_model import DataModel
 from datachain.lib.dc import C, DatasetPrepareError, Sys
+from datachain.lib.dc.listings import read_listing_dataset
 from datachain.lib.file import File
-from datachain.lib.listing import LISTING_PREFIX
+from datachain.lib.listing import LISTING_PREFIX, parse_listing_uri
 from datachain.lib.listing_info import ListingInfo
 from datachain.lib.signal_schema import (
     SignalRemoveError,
@@ -454,6 +455,56 @@ def test_listings_reindex_subpath_local_file_system(test_session, tmp_dir):
     assert dc.read_storage(subdir.as_uri(), session=test_session).count() == 1
 
 
+@pytest.mark.parametrize("version", [None, "1.0.0"])
+def test_listings_read_listing_dataset(test_session, tmp_dir, version):
+    df = pd.DataFrame(DF_DATA)
+    df.to_parquet(tmp_dir / "df.parquet")
+    uri = tmp_dir.as_uri()
+
+    ds_name, _, _ = parse_listing_uri(uri)
+    dc.read_storage(uri, session=test_session).exec()
+
+    chain, listing_version = read_listing_dataset(
+        ds_name, version=version, session=test_session
+    )
+    assert listing_version.num_objects == 1
+    assert listing_version.size > 1000
+    assert listing_version.size < 5000
+    assert listing_version.status == 4
+
+    assert chain.count() == 1
+    files = list(chain.collect("file"))
+    assert len(files) == 1
+    assert files[0].path == "df.parquet"
+    assert files[0].source == uri
+
+
+def test_listings_read_listing_dataset_with_subpath(test_session, tmp_dir):
+    subdir = tmp_dir / "subdir"
+    os.mkdir(subdir)
+
+    df = pd.DataFrame(DF_DATA)
+    df.to_parquet(tmp_dir / "df.parquet")
+    df.to_parquet(tmp_dir / "df2.parquet")
+    df.to_parquet(subdir / "df3.parquet")
+
+    ds_name, _, _ = parse_listing_uri(tmp_dir.as_uri())
+    ds_name = ds_name.removeprefix(LISTING_PREFIX)
+    dc.read_storage(tmp_dir.as_uri(), session=test_session).exec()
+
+    chain, listing_version = read_listing_dataset(
+        ds_name, path="subdir", session=test_session
+    )
+    assert listing_version.num_objects == 3
+
+    # Chain is filtered for subdir
+    assert chain.count() == 1
+    files = list(chain.collect("file"))
+    assert len(files) == 1
+    assert files[0].path == "subdir/df3.parquet"
+    assert files[0].source == tmp_dir.as_uri()
+
+
 def test_preserve_feature_schema(test_session):
     ds = dc.read_records(dc.DataChain.DEFAULT_FILE_RECORD, session=test_session)
     ds = ds.gen(

datachain-0.19.1/docs/commands/auth/team.md

@@ -1,36 +0,0 @@
-# auth team
-
-Set the default team for Studio operations.
-
-## Synopsis
-
-```usage
-usage: datachain auth team [-h] [-v] [-q] [--global] team_name
-```
-
-## Description
-
-This command sets the default team for Studio operations. By default, the team setting is project-specific, but you can use the `--global` option to set it for all projects.
-
-## Arguments
-
-* `team_name` - Name of the team to set as default
-
-## Options
-
-* `--global` - Set team globally for all projects
-* `-h`, `--help` - Show the help message and exit.
-* `-v`, `--verbose` - Be verbose.
-* `-q`, `--quiet` - Be quiet.
-
-## Examples
-
-1. Set default team for current project:
-```bash
-datachain auth team my-team
-```
-
-2. Set default team globally for all projects:
-```bash
-datachain auth team --global my-team
-```