dstack 0.0.9__py3-none-any.whl → 0.20.7__py3-none-any.whl

dstack/aws/config.py

@@ -1,40 +0,0 @@
-from typing import Optional
-
-from botocore.client import BaseClient
-from rich import print
-from rich.prompt import Confirm
-
-from dstack.aws import runners
-
-
-def validate_bucket(s3_client: BaseClient, bucket_name: str, region_name: str) -> bool:
-    try:
-        response = s3_client.head_bucket(Bucket=bucket_name)
-        bucket_region = response["ResponseMetadata"]["HTTPHeaders"]["x-amz-bucket-region"]
-        if bucket_region != region_name:
-            print(f"[red bold]✗[/red bold] [red]The bucket belongs to another AWS region.")
-            return False
-    except Exception as e:
-        if hasattr(e, "response") and e.response.get("Error") and e.response["Error"].get("Code") == "403":
-            print(f"[red bold]✗[/red bold] [red]You don't have access to this bucket. "
-                  "It may belong to another AWS account.[/red]")
-            return False
-        else:
-            if hasattr(e, "response") and e.response.get("Error") and e.response["Error"].get("Code") == "404":
-                if Confirm.ask(f"[sea_green3 bold]?[/sea_green3 bold] "
-                               f"[red bold]The bucket doesn't exist. Create it?[/red bold]"):
-                    if region_name != "us-east-1":
-                        s3_client.create_bucket(Bucket=bucket_name,
-                                                CreateBucketConfiguration={"LocationConstraint": region_name})
-                    else:
-                        s3_client.create_bucket(Bucket=bucket_name)
-                else:
-                    return False
-            else:
-                raise e
-    return True
-
-
-def configure(ec2_client: BaseClient, iam_client: BaseClient, bucket_name: str, subnet_id: Optional[str]):
-    runners.instance_profile_arn(iam_client, bucket_name)
-    runners.get_security_group_id(ec2_client, bucket_name, subnet_id)

dstack/aws/jobs.py

@@ -1,245 +0,0 @@
-from typing import List, Optional, Dict, Any
-
-import yaml
-from botocore.client import BaseClient
-
-from dstack.jobs import Job, JobStatus, JobHead, Requirements, GpusRequirements, JobRefId, AppSpec, DepSpec, \
-    ArtifactSpec
-from dstack.repo import RepoData
-
-
-def serialize_job(job: Job) -> dict:
-    deps = []
-    if job.dep_specs:
-        for dep in job.dep_specs:
-            deps.append({
-                "repo_user_name": dep.repo_user_name,
-                "repo_name": dep.repo_name,
-                "run_name": dep.run_name,
-                "mount": dep.mount,
-            })
-    artifacts = []
-    if job.artifact_specs:
-        for artifact_spec in job.artifact_specs:
-            artifacts.append({
-                "path": artifact_spec.artifact_path,
-                "mount": artifact_spec.mount
-            })
-    job_data = {
-        "job_id": job.job_id,
-        "repo_user_name": job.repo_data.repo_user_name,
-        "repo_name": job.repo_data.repo_name,
-        "repo_branch": job.repo_data.repo_branch,
-        "repo_hash": job.repo_data.repo_hash,
-        "repo_diff": job.repo_data.repo_diff or '',
-        "run_name": job.run_name,
-        "workflow_name": job.workflow_name or '',
-        "provider_name": job.provider_name,
-        "status": job.status.value,
-        "submitted_at": job.submitted_at,
-        "image_name": job.image_name,
-        "commands": job.commands or [],
-        "env": job.env or {},
-        "working_dir": job.working_dir or '',
-        "artifacts": artifacts,
-        "port_count": job.port_count if job.port_count else 0,
-        "ports": [str(port) for port in job.ports] if job.ports else [],
-        "host_name": job.host_name or '',
-        "requirements": _serialize_requirements(job.requirements),
-        "deps": deps,
-        "master_job_id": job.master_job.get_id() if job.master_job else '',
-        "apps": [{
-            "port_index": a.port_index,
-            "app_name": a.app_name,
-            "url_path": a.url_path or '',
-            "url_query_params": a.url_query_params or {},
-        } for a in job.app_specs] if job.app_specs else [],
-        "runner_id": job.runner_id or '',
-        "request_id": job.request_id or '',
-        "tag_name": job.tag_name or '',
-    }
-    return job_data
-
-
-def _serialize_requirements(requirements) -> Dict[str, Any]:
-    req_data = {}
-    if requirements:
-        if requirements.cpus:
-            req_data["cpus"] = {
-                "count": requirements.cpus
-            }
-        if requirements.memory_mib:
-            req_data["memory_mib"] = requirements.memory_mib
-        if requirements.gpus:
-            req_data["gpus"] = {
-                "count": requirements.gpus.count
-            }
-            if requirements.gpus.memory_mib:
-                req_data["gpus"]["memory_mib"] = requirements.gpus.memory_mib
-            if requirements.gpus.name:
-                req_data["gpus"]["name"] = requirements.gpus.name
-        if requirements.shm_size_mib:
-            req_data["shm_size_mib"] = requirements.shm_size_mib
-        if requirements.interruptible:
-            req_data["interruptible"] = requirements.interruptible
-    return req_data
-
-
-def unserialize_job(job_data: dict) -> Job:
-    _requirements = job_data.get("requirements")
-    requirements = Requirements(
-        _requirements.get("cpus") or None,
-        _requirements.get("memory_mib") or None,
-        GpusRequirements(_requirements["gpus"].get("count") or None,
-                         _requirements["gpus"].get("memory") or None,
-                         _requirements["gpus"].get("name") or None
-                         ) if _requirements.get("gpus") else None,
-        _requirements.get("shm_size_mib") or None, _requirements.get("interruptible") or None
-    ) if _requirements else None
-    if requirements:
-        if not requirements.cpus \
-                and (not requirements.gpus or
-                     (not requirements.gpus.count
-                      and not requirements.gpus.memory_mib
-                      and not requirements.gpus.name)) \
-                and not requirements.interruptible and not not requirements.shm_size_mib:
-            requirements = None
-    dep_specs = []
-    if job_data.get("deps"):
-        for dep in job_data["deps"]:
-            if isinstance(dep, str):
-                dep_repo_user_name, dep_repo_name, dep_run_name = tuple(dep.split(","))
-                dep_spec = DepSpec(dep_repo_user_name, dep_repo_name, dep_run_name, False)
-            else:
-                dep_spec = DepSpec(dep["repo_user_name"], dep["repo_name"], dep["run_name"], dep.get("mount") is True)
-            dep_specs.append(dep_spec)
-    artifact_specs = []
-    if job_data.get("artifacts"):
-        for artifact in job_data["artifacts"]:
-            if isinstance(artifact, str):
-                artifact_spec = ArtifactSpec(artifact, False)
-            else:
-                artifact_spec = ArtifactSpec(artifact["path"], artifact.get("mount") is True)
-            artifact_specs.append(artifact_spec)
-    master_job = JobRefId(job_data["master_job_id"]) if job_data.get("master_job_id") else None
-    app_specs = ([AppSpec(a["port_index"], a["app_name"], a.get("url_path") or None, a.get("url_query_params") or None)
-                  for a in (job_data.get("apps") or [])]) or None
-    job = Job(job_data["job_id"], RepoData(job_data["repo_user_name"], job_data["repo_name"],
-                                           job_data["repo_branch"], job_data["repo_hash"],
-                                           job_data["repo_diff"] or None),
-              job_data["run_name"], job_data.get("workflow_name") or None, job_data["provider_name"],
-              JobStatus(job_data["status"]), job_data["submitted_at"], job_data["image_name"],
-              job_data.get("commands") or None, job_data["env"] or None, job_data.get("working_dir") or None,
-              artifact_specs, job_data.get("port_count") or None, job_data.get("ports") or None,
-              job_data.get("host_name") or None, requirements, dep_specs or None, master_job, app_specs,
-              job_data.get("runner_id") or None, job_data.get("request_id") or None, job_data.get("tag_name") or None)
-    return job
-
-
-def _job_head_key(job: Job):
-    prefix = f"jobs/{job.repo_data.repo_user_name}/{job.repo_data.repo_name}"
-    key = f"{prefix}/l;" \
-          f"{job.job_id};" \
-          f"{job.provider_name};" \
-          f"{job.submitted_at};" \
-          f"{job.status.value};" \
-          f"{','.join([a.artifact_path for a in (job.artifact_specs or [])])};" \
-          f"{','.join([a.app_name for a in (job.app_specs or [])])};" \
-          f"{job.tag_name or ''}"
-    return key
-
-
-def create_job(s3_client: BaseClient, bucket_name: str, job: Job, counter: List[int] = [], create_head: bool = True):
-    if len(counter) == 0:
-        counter.append(0)
-    job_id = f"{job.run_name},{job.workflow_name or ''},{counter[0]}"
-    job.set_id(job_id)
-    if create_head:
-        s3_client.put_object(Body="", Bucket=bucket_name, Key=_job_head_key(job))
-    prefix = f"jobs/{job.repo_data.repo_user_name}/{job.repo_data.repo_name}"
-    key = f"{prefix}/{job_id}.yaml"
-    s3_client.put_object(Body=yaml.dump(serialize_job(job)), Bucket=bucket_name, Key=key)
-    counter[0] += 1
-
-
-def get_job(s3_client: BaseClient, bucket_name: str, repo_user_name: str, repo_name: str, job_id: str) -> Optional[Job]:
-    prefix = f"jobs/{repo_user_name}/{repo_name}"
-    key = f"{prefix}/{job_id}.yaml"
-    try:
-        obj = s3_client.get_object(Bucket=bucket_name, Key=key)
-        job = unserialize_job(yaml.load(obj['Body'].read().decode('utf-8'), yaml.FullLoader))
-        return job
-    except Exception as e:
-        if hasattr(e, "response") and e.response.get("Error") and e.response["Error"].get("Code") == "NoSuchKey":
-            return None
-        else:
-            raise e
-
-
-def update_job(s3_client: BaseClient, bucket_name: str, job: Job):
-    prefix = f"jobs/{job.repo_data.repo_user_name}/{job.repo_data.repo_name}"
-    job_head_key_prefix = f"{prefix}/l;{job.job_id};"
-    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=job_head_key_prefix, MaxKeys=1)
-    for obj in response["Contents"]:
-        s3_client.delete_object(Bucket=bucket_name, Key=obj["Key"])
-    job_head_key = _job_head_key(job)
-    s3_client.put_object(Body="", Bucket=bucket_name, Key=job_head_key)
-    key = f"{prefix}/{job.job_id}.yaml"
-    s3_client.put_object(Body=yaml.dump(serialize_job(job)), Bucket=bucket_name, Key=key)
-
-
-def list_job_heads(s3_client: BaseClient, bucket_name: str, repo_user_name: str, repo_name: str,
-                   run_name: Optional[str] = None) -> List[JobHead]:
-    prefix = f"jobs/{repo_user_name}/{repo_name}"
-    job_head_key_prefix = f"{prefix}/l;"
-    job_head_key_run_prefix = job_head_key_prefix + run_name if run_name else job_head_key_prefix
-    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=job_head_key_run_prefix)
-    job_heads = []
-    if "Contents" in response:
-        for obj in response["Contents"]:
-            job_id, provider_name, submitted_at, status, artifacts, app_names, tag_name = tuple(
-                obj["Key"][len(job_head_key_prefix):].split(';'))
-            run_name, workflow_name, job_index = tuple(job_id.split(','))
-            job_heads.append(JobHead(job_id, repo_user_name, repo_name, run_name, workflow_name or None, provider_name,
-                                     JobStatus(status), int(submitted_at),
-                                     artifacts.split(',') if artifacts else None, tag_name or None,
-                                     app_names.split(',') or None))
-    return job_heads
-
-
-def list_job_head(s3_client: BaseClient, bucket_name: str, repo_user_name: str, repo_name: str, job_id: str) \
-        -> Optional[JobHead]:
-    prefix = f"jobs/{repo_user_name}/{repo_name}"
-    job_head_key_prefix = f"{prefix}/l;{job_id};"
-    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=job_head_key_prefix)
-    if "Contents" in response:
-        for obj in response["Contents"]:
-            provider_name, submitted_at, status, artifacts, app_names, tag_name = tuple(
-                obj["Key"][len(job_head_key_prefix):].split(';'))
-            run_name, workflow_name, job_index = tuple(job_id.split(','))
-            return JobHead(job_id, repo_user_name, repo_name, run_name, workflow_name or None, provider_name,
-                           JobStatus(status), int(submitted_at),
-                           artifacts.split(',') if artifacts else None, tag_name or None,
-                           app_names.split(',') or None)
-    return None
-
-
-def list_jobs(s3_client: BaseClient, bucket_name: str, repo_user_name: str, repo_name: str,
-              run_name: Optional[str] = None) -> List[Job]:
-    job_key_run_prefix = f"jobs/{repo_user_name}/{repo_name}/{run_name},"
-    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=job_key_run_prefix)
-    jobs = []
-    if "Contents" in response:
-        for obj in response["Contents"]:
-            job_obj = s3_client.get_object(Bucket=bucket_name, Key=obj["Key"])
-            job = unserialize_job(yaml.load(job_obj['Body'].read().decode('utf-8'), yaml.FullLoader))
-            jobs.append(job)
-    return jobs
-
-
-def delete_job_head(s3_client: BaseClient, bucket_name: str, repo_user_name: str, repo_name: str, job_id: str):
-    prefix = f"jobs/{repo_user_name}/{repo_name}"
-    job_head_key_prefix = f"{prefix}/l;{job_id};"
-    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=job_head_key_prefix, MaxKeys=1)
-    for obj in response["Contents"]:
-        s3_client.delete_object(Bucket=bucket_name, Key=obj["Key"])

dstack/aws/logs.py

@@ -1,186 +0,0 @@
-import json
-import re
-import time
-from collections import defaultdict
-from typing import Optional, Dict, List, Generator, Any, Tuple
-from urllib import parse
-
-from botocore.client import BaseClient
-
-from dstack.aws import jobs, runs
-from dstack.backend import LogEvent, LogEventSource
-from dstack.jobs import AppSpec, JobHead
-
-WAIT_N_ONCE_FINISHED = 1
-
-CHECK_STATUS_EVERY_N = 3
-
-POLL_LOGS_RATE_SECS = 1
-
-
-def _render_log_message(s3_client: BaseClient, bucket_name: str, event: Dict[str, Any],
-                        repo_user_name: str, repo_name: str,
-                        job_host_names: Dict[str, Optional[str]],
-                        job_ports: Dict[str, Optional[List[int]]],
-                        job_app_specs: Dict[str, Optional[List[AppSpec]]]) -> LogEvent:
-    message = json.loads(event["message"].strip())
-    job_id = message["job_id"]
-    log = message["log"]
-    if job_id and job_id not in job_host_names:
-        job = jobs.get_job(s3_client, bucket_name, repo_user_name, repo_name, job_id)
-        job_host_names[job_id] = job.host_name or "none" if job else "none"
-        job_ports[job_id] = job.ports if job else None
-        job_app_specs[job_id] = job.app_specs if job else None
-    host_name = job_host_names[job_id]
-    ports = job_ports[job_id]
-    app_specs = job_app_specs[job_id]
-    pat = re.compile(f'http://(localhost|0.0.0.0|{host_name}):[\\S]*[^(.+)\\s\\n\\r]')
-    if re.search(pat, log):
-        if host_name != "none" and ports and app_specs:
-            for app_spec in app_specs:
-                port = ports[app_spec.port_index]
-                url_path = app_spec.url_path or ""
-                url_query_params = app_spec.url_query_params
-                url_query = ("?" + parse.urlencode(url_query_params)) if url_query_params else ""
-                app_url = f"http://{host_name}:{port}"
-                if url_path or url_query_params:
-                    app_url += "/"
-                    if url_query_params:
-                        app_url += url_query
-                log = re.sub(pat, app_url, log)
-    return LogEvent(event["eventId"], event["timestamp"], job_id, log,
-                    LogEventSource.STDOUT if message["source"] == "stdout" else LogEventSource.STDERR)
-
-
-def _get_latest_events_and_timestamp(event_ids_per_timestamp):
-    if event_ids_per_timestamp:
-        newest_timestamp = max(event_ids_per_timestamp.keys())
-        event_ids_per_timestamp = defaultdict(
-            set, {newest_timestamp: event_ids_per_timestamp[newest_timestamp]}
-        )
-    return event_ids_per_timestamp
-
-
-def _reset_filter_log_events_params(fle_kwargs, event_ids_per_timestamp):
-    if event_ids_per_timestamp:
-        fle_kwargs['startTime'] = max(
-            event_ids_per_timestamp.keys()
-        )
-    fle_kwargs.pop('nextToken', None)
-
-
-def _filter_log_events_loop(ec2_client: BaseClient, s3_client: BaseClient, logs_client: BaseClient, bucket_name: str,
-                            repo_user_name: str, repo_name: str, job_heads: List[JobHead],
-                            filter_logs_events_kwargs: dict):
-    event_ids_per_timestamp = defaultdict(set)
-    counter = 0
-    finished_counter = 0
-    while True:
-        response = logs_client.filter_log_events(**filter_logs_events_kwargs)
-
-        for event in response["events"]:
-            if event["eventId"] not in event_ids_per_timestamp[event["timestamp"]]:
-                event_ids_per_timestamp[event["timestamp"]].add(event["eventId"])
-                yield event
-        event_ids_per_timestamp = _get_latest_events_and_timestamp(
-            event_ids_per_timestamp
-        )
-        if "nextToken" in response:
-            filter_logs_events_kwargs["nextToken"] = response["nextToken"]
-        else:
-            _reset_filter_log_events_params(
-                filter_logs_events_kwargs,
-                event_ids_per_timestamp
-            )
-            time.sleep(POLL_LOGS_RATE_SECS)
-            counter = counter + 1
-            if counter % CHECK_STATUS_EVERY_N == 0:
-                _job_heads = [jobs.get_job(s3_client, bucket_name, repo_user_name, repo_name, job_head.job_id)
-                              for job_head in job_heads]
-                run = next(iter(runs.get_run_heads(ec2_client, s3_client, bucket_name, repo_user_name, repo_name,
-                                                   _job_heads, include_request_heads=False)))
-                if run.status.is_finished():
-                    if finished_counter == WAIT_N_ONCE_FINISHED:
-                        break
-                    finished_counter += 1
-
-
-def create_log_group_if_not_exists(logs_client: BaseClient, bucket_name: str, log_group_name: str):
-    response = logs_client.describe_log_groups(logGroupNamePrefix=log_group_name)
-    if not response["logGroups"] or not any(
-            filter(lambda g: g["logGroupName"] == log_group_name, response["logGroups"])):
-        logs_client.create_log_group(
-            logGroupName=log_group_name,
-            tags={
-                'owner': 'dstack',
-                'dstack_bucket': bucket_name,
-            }
-        )
-
-
-def create_log_stream(logs_client: BaseClient, log_group_name: str, run_name: str):
-    logs_client.create_log_stream(logGroupName=log_group_name, logStreamName=run_name)
-
-
-def _filter_logs_events_kwargs(bucket_name: str, repo_user_name: str, repo_name: str, run_name: str, start_time: int,
-                               end_time: Optional[int], next_token: Optional[str]):
-    filter_logs_events_kwargs = {
-        "logGroupName": f"/dstack/jobs/{bucket_name}/{repo_user_name}/{repo_name}",
-        "logStreamNames": [run_name],
-        "startTime": start_time,
-        "interleaved": True,
-    }
-    if end_time:
-        filter_logs_events_kwargs["endTime"] = end_time
-    if next_token:
-        filter_logs_events_kwargs["nextToken"] = next_token
-    return filter_logs_events_kwargs
-
-
-def poll_logs(ec2_client: BaseClient, s3_client: BaseClient, logs_client: BaseClient, bucket_name: str,
-              repo_user_name: str, repo_name: str, job_heads: List[JobHead], start_time: int, attached: bool) \
-        -> Generator[LogEvent, None, None]:
-    run_name = job_heads[0].run_name
-    filter_logs_events_kwargs = _filter_logs_events_kwargs(bucket_name, repo_user_name, repo_name, run_name, start_time,
-                                                           end_time=None, next_token=None)
-    job_host_names = {}
-    job_ports = {}
-    job_app_specs = {}
-
-    try:
-        if attached:
-            for event in _filter_log_events_loop(ec2_client, s3_client, logs_client, bucket_name, repo_user_name,
-                                                 repo_name, job_heads, filter_logs_events_kwargs):
-                yield _render_log_message(s3_client, bucket_name, event, repo_user_name, repo_name,
-                                          job_host_names, job_ports, job_app_specs)
-        else:
-            paginator = logs_client.get_paginator("filter_log_events")
-            for page in paginator.paginate(**filter_logs_events_kwargs):
-                for event in page["events"]:
-                    yield _render_log_message(s3_client, bucket_name, event, repo_user_name,
-                                              repo_name, job_host_names, job_ports, job_app_specs)
-    except Exception as e:
-        if hasattr(e, "response") and e.response.get("Error") and e.response["Error"].get(
-                "Code") == "ResourceNotFoundException":
-            return
-        else:
-            raise e
-
-
-def query_logs(s3_client: BaseClient, logs_client: BaseClient, bucket_name: str,
-               repo_user_name: str, repo_name: str, run_name: str, start_time: int, end_time: Optional[int],
-               next_token: Optional[str],
-               job_host_names: Dict[str, Optional[str]], job_ports: Dict[str, Optional[List[int]]],
-               job_app_specs: Dict[str, Optional[List[AppSpec]]]) -> Tuple[List[LogEvent], Optional[str],
-                                                                           Dict[str, Optional[str]],
-                                                                           Dict[str, Optional[List[int]]],
-                                                                           Dict[str, Optional[List[AppSpec]]]]:
-    job_host_names = dict(job_host_names)
-    job_ports = dict(job_ports)
-    job_app_specs = dict(job_app_specs)
-    filter_logs_events_kwargs = _filter_logs_events_kwargs(bucket_name, repo_user_name, repo_name, run_name, start_time,
-                                                           end_time, next_token)
-    response = logs_client.filter_log_events(**filter_logs_events_kwargs)
-    return [_render_log_message(s3_client, bucket_name, event, repo_user_name, repo_name, job_host_names, job_ports,
-                                job_app_specs)
-            for event in response["events"]], response.get("nextToken"), job_host_names, job_ports, job_app_specs

dstack/aws/repos.py

@@ -1,137 +0,0 @@
-import json
-from typing import List, Optional
-
-from botocore.client import BaseClient
-
-from dstack.aws import runners
-from dstack.backend import RepoHead
-from dstack.repo import RepoCredentials, RepoProtocol
-
-
-def list_repo_heads(s3_client: BaseClient, bucket_name: str) -> List[RepoHead]:
-    tag_head_prefix = f"repos/l;"
-    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=tag_head_prefix)
-    repo_heads = []
-    if "Contents" in response:
-        for obj in response["Contents"]:
-            repo_user_name, repo_name, last_run_at, tags_count = tuple(obj["Key"][len(tag_head_prefix):].split(';'))
-            repo_heads.append(RepoHead(repo_user_name, repo_name, int(last_run_at) if last_run_at else None,
-                                       int(tags_count)))
-    return repo_heads
-
-
-def _get_repo_head(s3_client: BaseClient, bucket_name: str, repo_user_name: str, repo_name: str) -> Optional[RepoHead]:
-    repo_head_prefix = f"repos/l;{repo_user_name};{repo_name};"
-    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=repo_head_prefix)
-    if response.get("Contents"):
-        last_run_at, tags_count = tuple(response["Contents"][0]["Key"][len(repo_head_prefix):].split(';'))
-        return RepoHead(repo_user_name, repo_name, int(last_run_at) if last_run_at else None,
-                        int(tags_count))
-    else:
-        return None
-
-
-def _create_or_update_repo_head(s3_client: BaseClient, bucket_name: str, repo_head: RepoHead):
-    repo_head_prefix = f"repos/l;{repo_head.repo_user_name};{repo_head.repo_name};"
-    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=repo_head_prefix)
-    if "Contents" in response:
-        for obj in response["Contents"]:
-            s3_client.delete_object(Bucket=bucket_name, Key=obj["Key"])
-    repo_head_key = f"repos/l;{repo_head.repo_user_name};{repo_head.repo_name};{repo_head.last_run_at or ''};" \
-                    f"{repo_head.tags_count}"
-    s3_client.put_object(Body="", Bucket=bucket_name, Key=repo_head_key)
-
-
-def update_repo_last_run_at(s3_client: BaseClient, bucket_name: str, repo_user_name: str, repo_name: str,
-                            last_run_at: int):
-    repo_head = _get_repo_head(s3_client, bucket_name, repo_user_name, repo_name) or RepoHead(
-        repo_user_name, repo_name, None, tags_count=0)
-    repo_head.last_run_at = last_run_at
-    _create_or_update_repo_head(s3_client, bucket_name, repo_head)
-
-
-def increment_repo_tags_count(s3_client: BaseClient, bucket_name: str, repo_user_name: str, repo_name: str):
-    repo_head = _get_repo_head(s3_client, bucket_name, repo_user_name, repo_name) or RepoHead(
-        repo_user_name, repo_name, None, tags_count=0)
-    repo_head.tags_count = repo_head.tags_count + 1
-    _create_or_update_repo_head(s3_client, bucket_name, repo_head)
-
-
-def decrement_repo_tags_count(s3_client: BaseClient, bucket_name: str, repo_user_name: str, repo_name: str):
-    repo_head = _get_repo_head(s3_client, bucket_name, repo_user_name, repo_name)
-    if repo_head:
-        repo_head.tags_count = repo_head.tags_count - 1
-        _create_or_update_repo_head(s3_client, bucket_name, repo_head)
-    else:
-        raise Exception(f"No repo head is found: {repo_user_name}/{repo_name}")
-
-
-def delete_repo(s3_client: BaseClient, bucket_name: str, repo_user_name: str, repo_name: str):
-    repo_head_prefix = f"repos/l;{repo_user_name};{repo_name};"
-    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=repo_head_prefix)
-    if "Contents" in response:
-        for obj in response["Contents"]:
-            s3_client.delete_object(Bucket=bucket_name, Key=obj["Key"])
-
-
-def save_repo_credentials(sts_client: BaseClient, iam_client: BaseClient, secretsmanager_client: BaseClient,
-                          bucket_name: str, repo_user_name: str, repo_name: str,
-                          repo_credentials: RepoCredentials):
-    secret_name = f"/dstack/{bucket_name}/credentials/{repo_user_name}/{repo_name}"
-    credentials_data = {
-        "protocol": repo_credentials.protocol.value
-    }
-    if repo_credentials.protocol == RepoProtocol.HTTPS and repo_credentials.oauth_token:
-        credentials_data["oauth_token"] = repo_credentials.oauth_token
-    elif repo_credentials.protocol == RepoProtocol.SSH:
-        if repo_credentials.private_key:
-            credentials_data["private_key"] = repo_credentials.private_key
-        else:
-            raise Exception("No private key is specified")
-    try:
-        secretsmanager_client.get_secret_value(SecretId=secret_name)
-        secretsmanager_client.update_secret(
-            SecretId=secret_name,
-            SecretString=json.dumps(credentials_data)
-        )
-    except Exception as e:
-        if hasattr(e, "response") and e.response.get("Error") and e.response["Error"].get(
-                "Code") == "ResourceNotFoundException":
-            secretsmanager_client.create_secret(
-                Name=secret_name,
-                SecretString=json.dumps(credentials_data),
-                Description="Generated by dstack",
-                Tags=[
-                    {
-                        'Key': 'owner',
-                        'Value': 'dstack'
-                    },
-                    {
-                        'Key': 'dstack_bucket',
-                        'Value': bucket_name
-                    }
-                ],
-            )
-        else:
-            raise e
-    role_name = runners.role_name(iam_client, bucket_name)
-    account_id = sts_client.get_caller_identity()["Account"]
-    secretsmanager_client.put_resource_policy(
-        SecretId=secret_name,
-        ResourcePolicy=json.dumps(
-            {
-                "Version": "2012-10-17",
-                "Statement": [
-                    {
-                        "Effect": "Allow",
-                        "Principal": {"AWS": f"arn:aws:iam::{account_id}:role/{role_name}"},
-                        "Action": [
-                            "secretsmanager:GetSecretValue",
-                            "secretsmanager:ListSecrets",
-                        ],
-                        "Resource": "*"
-                    }
-                ]
-            }
-        )
-    )

dstack/aws/run_names.py

@@ -1,17 +0,0 @@
-import yaml
-from botocore.client import BaseClient
-
-
-def next_run_name_index(client: BaseClient, bucket_name: str, run_name: str) -> int:
-    count = 0
-    key = f"run-names/{run_name}.yaml"
-    try:
-        obj = client.get_object(Bucket=bucket_name, Key=key)
-        count = yaml.load(obj['Body'].read().decode('utf-8'), Loader=yaml.FullLoader)["count"]
-        client.put_object(Body=yaml.dump({"count": count + 1}), Bucket=bucket_name, Key=key)
-    except Exception as e:
-        if hasattr(e, "response") and e.response.get("Error") and e.response["Error"].get("Code") == "NoSuchKey":
-            client.put_object(Body=yaml.dump({"count": count + 1}), Bucket=bucket_name, Key=key)
-        else:
-            raise e
-    return count + 1